text_plain_output_1.png
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_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()