text_plain_output_1.png
from albumentations.augmentations.bbox_utils import denormalize_bbox, normalize_bbox

from albumentations.core.transforms_interface import DualTransform

from collections import namedtuple

import albumentations

import ast

import cv2

import matplotlib.patches as patches

import numpy as np # linear algebra

import os

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

BASE_DIR = '/kaggle/input/global-wheat-detection'
WORK_DIR = '/kaggle/working'
np.random.seed(1996)

train_df = pd.read_csv(os.path.join(BASE_DIR, 'train.csv'))
train_df[['x_min', 'y_min', 'width', 'height']] = pd.DataFrame([ast.literal_eval(x) for x in train_df.bbox.tolist()], index=train_df.index)
train_df = train_df[['image_id', 'bbox', 'source', 'x_min', 'y_min', 'width', 'height']]
train_df['area'] = train_df['width'] * train_df['height']
train_df['x_max'] = train_df['x_min'] + train_df['width']
train_df['y_max'] = train_df['y_min'] + train_df['height']
train_df = train_df.drop(['bbox', 'source'], axis=1)
train_df = train_df[['image_id', 'x_min', 'y_min', 'x_max', 'y_max', 'width', 'height', 'area']]
train_df = train_df[train_df['area'] < 100000]

image_ids = train_df['image_id'].unique()

image_id = 'c14c1e300'
image = cv2.imread(os.path.join(BASE_DIR, 'train', f'{image_id}.jpg'), cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
image /= 255.0

pascal_voc_boxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'x_max', 'y_max']].astype(np.int32).values
coco_boxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'width', 'height']].astype(np.int32).values
assert len(pascal_voc_boxes) == len(coco_boxes)
labels = np.ones((len(pascal_voc_boxes),))

def get_bbox(bboxes, col, color='white', bbox_format='pascal_voc'):
    for i in range(len(bboxes)):
        if bbox_format == 'pascal_voc':
            rect = patches.Rectangle((bboxes[i][0], bboxes[i][1]), bboxes[i][2] - bboxes[i][0], bboxes[i][3] - bboxes[i][1], linewidth=2, edgecolor=color, facecolor='none')
        else:
            rect = patches.Rectangle((bboxes[i][0], bboxes[i][1]), bboxes[i][2], bboxes[i][3], linewidth=2, edgecolor=color, facecolor='none')
        col.add_patch(rect)

aug = albumentations.Compose([albumentations.Resize(512, 512), albumentations.VerticalFlip(1)], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})

aug_result = aug(image=image, bboxes=pascal_voc_boxes, labels=labels)

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(16, 16))
get_bbox(pascal_voc_boxes, ax[0], color='red')
ax[0].title.set_text('Original Image')
ax[0].imshow(image)

get_bbox(aug_result['bboxes'], ax[1], color='red')
ax[1].title.set_text('Augmented Image')
ax[1].imshow(aug_result['image'])
plt.show()

aug = albumentations.Compose([albumentations.Resize(512, 512), albumentations.VerticalFlip(1), albumentations.Blur(p=1)], bbox_params={'format': 'coco', 'label_fields': ['labels']})

aug_result = aug(image=image, bboxes=coco_boxes, labels=labels)

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(16, 16))
get_bbox(coco_boxes, ax[0], color='red', bbox_format='coco')
ax[0].title.set_text('Original Image')
ax[0].imshow(image)

get_bbox(aug_result['bboxes'], ax[1], color='red', bbox_format='coco')
ax[1].title.set_text('Augmented Image')
ax[1].imshow(aug_result['image'])
plt.show()

class CustomCutout(DualTransform):
    """
    Custom Cutout augmentation with handling of bounding boxes 
    Note: (only supports square cutout regions)
    Author: Kaushal28
    """
    def __init__(self, fill_value=0, bbox_removal_threshold=0.5, min_cutout_size=192, max_cutout_size=512, always_apply=False, p=0.5):
        """
        Class construstor
        :param fill_value: Value to be filled in cutout (default is 0 or black color)
        :param bbox_removal_threshold: Bboxes having content cut by cutout path more than this threshold will be removed
        :param min_cutout_size: minimum size of cutout (192 x 192)
        :param max_cutout_size: maximum size of cutout (512 x 512)
        """
        super(CustomCutout, self).__init__(always_apply, p)
        self.fill_value = fill_value
        self.bbox_removal_threshold = bbox_removal_threshold
        self.min_cutout_size = min_cutout_size
        self.max_cutout_size = max_cutout_size
    def _get_cutout_position(self, img_height, img_width, cutout_size):
        """
        Randomly generates cutout position as a named tuple
        :param img_height: height of the original image
        :param img_width: width of the original image
        :param cutout_size: size of the cutout patch (square)
        :returns position of cutout patch as a named tuple
        """
        position = namedtuple('Point', 'x y')
        return position(np.random.randint(0, img_width - cutout_size + 1), np.random.randint(0, img_height - cutout_size + 1))
    def _get_cutout(self, img_height, img_width):
        """
        Creates a cutout pacth with given fill value and determines the position in the original image
        :param img_height: height of the original image
        :param img_width: width of the original image
        :returns (cutout patch, cutout size, cutout position)
        """
        cutout_size = np.random.randint(self.min_cutout_size, self.max_cutout_size + 1)
        cutout_position = self._get_cutout_position(img_height, img_width, cutout_size)
        return (np.full((cutout_size, cutout_size, 3), self.fill_value), cutout_size, cutout_position)
    def apply(self, image, **params):
        """
        Applies the cutout augmentation on the given image
        :param image: The image to be augmented
        :returns augmented image
        """
        image = image.copy()
        self.img_height, self.img_width, _ = image.shape
        cutout_arr, cutout_size, cutout_pos = self._get_cutout(self.img_height, self.img_width)
        self.image = image
        self.cutout_pos = cutout_pos
        self.cutout_size = cutout_size
        image[cutout_pos.y:cutout_pos.y + cutout_size, cutout_pos.x:cutout_size + cutout_pos.x, :] = cutout_arr
        return image
    def apply_to_bbox(self, bbox, **params):
        """
        Removes the bounding boxes which are covered by the applied cutout
        :param bbox: A single bounding box coordinates in pascal_voc format
        :returns transformed bbox's coordinates
        """
        bbox = denormalize_bbox(bbox, self.img_height, self.img_width)
        x_min, y_min, x_max, y_max = tuple(map(int, bbox))
        bbox_size = (x_max - x_min) * (y_max - y_min)
        overlapping_size = np.sum((self.image[y_min:y_max, x_min:x_max, 0] == self.fill_value) & (self.image[y_min:y_max, x_min:x_max, 1] == self.fill_value) & (self.image[y_min:y_max, x_min:x_max, 2] == self.fill_value))
        if overlapping_size / bbox_size > self.bbox_removal_threshold:
            return normalize_bbox((0, 0, 0, 0), self.img_height, self.img_width)
        return normalize_bbox(bbox, self.img_height, self.img_width)
    def get_transform_init_args_names(self):
        """
        Fetches the parameter(s) of __init__ method
        :returns: tuple of parameter(s) of __init__ method
        """
        return 'fill_value'

augmentation = albumentations.Compose([CustomCutout(p=1), albumentations.Flip(always_apply=True), albumentations.OneOf([albumentations.Blur(p=0.5), albumentations.GaussNoise(var_limit=5.0 / 255.0, p=0.5)], p=1)], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})

def get_bbox(bboxes, col, color='white'):
    for i in range(len(bboxes)):
        rect = patches.Rectangle((bboxes[i][0], bboxes[i][1]), bboxes[i][2] - bboxes[i][0], bboxes[i][3] - bboxes[i][1], linewidth=2, edgecolor=color, facecolor='none')
        col.add_patch(rect)

num_images = 5
rand_start = np.random.randint(0, len(image_ids) - 5)
fig, ax = plt.subplots(nrows=num_images, ncols=2, figsize=(16, 40))
for index, image_id in enumerate(image_ids[rand_start:rand_start + num_images]):
    image = cv2.imread(os.path.join(BASE_DIR, 'train', f'{image_id}.jpg'), cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
    image /= 255.0
    bboxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'x_max', 'y_max']].astype(np.int32).values
    labels = np.ones((len(bboxes),))
    aug_result = augmentation(image=image, bboxes=bboxes, labels=labels)
    get_bbox(bboxes, ax[index][0], color='red')
    ax[index][0].grid(False)
    ax[index][0].set_xticks([])
    ax[index][0].set_yticks([])
    ax[index][0].title.set_text('Original Image')
    ax[index][0].imshow(image)
    get_bbox(aug_result['bboxes'], ax[index][1], color='red')
    ax[index][1].grid(False)
    ax[index][1].set_xticks([])
    ax[index][1].set_yticks([])
    ax[index][1].title.set_text(f"Augmented Image: Removed bboxes: {len(bboxes) - len(aug_result['bboxes'])}")
    ax[index][1].imshow(aug_result['image'])
plt.show()