{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Start \n\n","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"%ls\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport ast\nfrom collections import namedtuple\nimport random\nimport collections\nimport uuid\nfrom glob import glob\nfrom datetime import datetime\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nfrom tqdm import tqdm\nfrom PIL import Image\n\nimport joblib\nfrom joblib import Parallel, delayed\n\nimport cv2\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom albumentations.core.transforms_interface import DualTransform\nfrom albumentations.augmentations.bbox_utils import denormalize_bbox, normalize_bbox\n\nfrom sklearn.model_selection import StratifiedKFold\n\nimport torch\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.utils.data as data_utils\n\nfrom matplotlib import pyplot as plt\nimport matplotlib.patches as patches\nfrom matplotlib.image import imsave","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Constants\nBASE_DIR = '/kaggle/input/global-wheat-detection'\nWORK_DIR = '/kaggle/working'\n\n# Set seed for numpy for reproducibility\n#np.random.seed(1996)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%cp -r /kaggle/input/global-wheat-detection/train /kaggle/working\n%cp /kaggle/input/global-wheat-detection/train.csv /kaggle/working","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# START Puzzle Augmentation\n1. divide an image into 4 pieces(puzzles) with bbox\n2. make pool of puzzles of all original images (in train_data)\n3. sample 4 pieces & merge (repeat k-times)\n4. save as a dataframe & jpg images","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def adjust_bbox_in_image(\n    i_w,\n    i_h,\n    bbox,\n    min_bbox_size=20):\n    \"\"\"crop bbox if it cover the edge of the cropped image\"\"\"\n    \n    x, y, w, h = bbox\n    if x < 0:\n        w += x\n        x = 0\n    if y < 0:\n        h += y\n        y = 0\n    if i_w < x+w:\n        w -= (x+w-i_w)\n    if i_h < y+h:\n        h -= (y+h-i_h)\n    \n    on_border = x < 3 or y < 3 or i_w-3 < (x+w) or i_h-3 < (y+h)\n    under_min_size = w < min_bbox_size or h < min_bbox_size\n    if on_border and under_min_size:\n        return None\n    \n    return (x, y, w, h)\n\n\ndef make_puzzles(\n    image_id,\n    bboxes,\n    min_bbox_size=20,\n    image_root='/kaggle/working/train/'):\n    \"\"\"divide given image into 4 pieces with bboxes\"\"\"\n\n    img_path = os.path.join(image_root, '{}.jpg'.format(image_id))\n    \n    image = cv2.imread(img_path)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    \n    row, col, _ = image.shape\n    \n    y_div = int(row/2)\n    x_div = int(col/2)\n    lt = image[:y_div, :x_div] # left top\n    rt = image[:y_div, x_div:] # right top\n    lb = image[y_div:, :x_div] # left bottom\n    rb = image[y_div:, x_div:] # right bottom\n    \n    lt_bboxes, rt_bboxes, lb_bboxes, rb_bboxes = [], [], [], []\n    \n    for bbox in bboxes:\n        x, y, w, h = bbox\n        \n        # Quandrant-2\n        lt_x, lt_y = x, y\n        if lt_y < y_div and lt_x < x_div:\n            _bbox = adjust_bbox_in_image(col/2, row/2, bbox, min_bbox_size=min_bbox_size)\n            if _bbox:\n                lt_bboxes.append(_bbox)\n        \n        # Quandrant-1\n        rt_x, rt_y = x+w-1, y\n        if rt_y < y_div and x_div <= rt_x:\n            _bbox = (bbox[0]-col/2, bbox[1], bbox[2], bbox[3])\n            _bbox = adjust_bbox_in_image(col/2, row/2, _bbox, min_bbox_size=min_bbox_size)\n            if _bbox:\n                rt_bboxes.append(_bbox)\n            \n        # Quandrant-3\n        lb_x, lb_y = x, y+h-1\n        if y_div <= lb_y and lb_x < x_div:\n            _bbox = (bbox[0], bbox[1]-row/2, bbox[2], bbox[3])\n            _bbox = adjust_bbox_in_image(col/2, row/2, _bbox, min_bbox_size=min_bbox_size)\n            if _bbox:\n                lb_bboxes.append(_bbox)\n        \n        # Quandrant-4\n        rb_x, rb_y = x+w-1, y+h-1\n        if y_div <= rb_y and x_div <= rb_x:\n            _bbox = (bbox[0]-col/2, bbox[1]-row/2, bbox[2], bbox[3])\n            _bbox = adjust_bbox_in_image(col/2, row/2, _bbox, min_bbox_size=min_bbox_size)\n            if _bbox:\n                rb_bboxes.append(_bbox)\n    \n    puzzle_bbox_pairs = [\n        (lt, lt_bboxes),\n        (rt, rt_bboxes),\n        (lb, lb_bboxes),\n        (rb, rb_bboxes)\n    ]\n    \n    return puzzle_bbox_pairs\n\n\ndef merge_random_4_puzzles(puzzle_bbox_pairs):\n    lt_img, lt_bboxes = puzzle_bbox_pairs[0]\n    rt_img, rt_bboxes = puzzle_bbox_pairs[1]\n    lb_img, lb_bboxes = puzzle_bbox_pairs[2]\n    rb_img, rb_bboxes = puzzle_bbox_pairs[3]\n    \n    row, col, ch = lt_img.shape\n    x_div = col\n    y_div = row\n    \n    empty_img = np.zeros((row*2, col*2, ch), np.uint8)\n    \n    empty_img[:y_div,:x_div,:] = lt_img\n    empty_img[:y_div,x_div:,:] = rt_img\n    empty_img[y_div:,:x_div,:] = lb_img\n    empty_img[y_div:,x_div:,:] = rb_img\n    \n    _lt_bboxes = lt_bboxes[:]\n    _rt_bboxes = rt_bboxes[:]\n    for i, bbox in enumerate(_rt_bboxes):\n        x, y, w, h = bbox\n        _rt_bboxes[i] = (x+x_div, y, w, h)\n    \n    _lb_bboxes = lb_bboxes[:]\n    for i, bbox in enumerate(_lb_bboxes):\n        x, y, w, h = bbox\n        _lb_bboxes[i] = (x, y+y_div, w, h)\n    \n    _rb_bboxes = rb_bboxes[:]\n    for i, bbox in enumerate(_rb_bboxes):\n        x, y, w, h = bbox\n        _rb_bboxes[i] = (x+x_div, y+y_div, w, h)\n        \n    merged_bbox = _lt_bboxes + _rt_bboxes + _lb_bboxes + _rb_bboxes\n    \n    return (empty_img, merged_bbox)\n    \n\ndef visualize_4_image_bbox(puzzle_bbox_pairs):\n    fig, ax = plt.subplots(2, 2, figsize=(12, 12))\n    ax = ax.flatten()\n\n    labels = ['left top', 'right top', 'left bot', 'right bot']\n\n    for i in range(4):\n        image, bboxes = puzzle_bbox_pairs[i]\n        for row in bboxes:\n            x, y, w, h = (int(n) for n in row)\n            cv2.rectangle(image,\n                          (x, y),\n                          (x+w, y+h),\n                          (220, 0, 0), 3)\n        ax[i].set_axis_off()\n        ax[i].imshow(image)\n        ax[i].set_title(labels[i], color='yellow')\n\n        \ndef visualize_image_bbox(image, bboxes):\n    fig, ax = plt.subplots(1, 1, figsize=(12, 12))\n\n    for row in bboxes:\n        x, y, w, h = (int(n) for n in row)\n        cv2.rectangle(image,\n                      (x, y),\n                      (x+w, y+h),\n                      (220, 0, 0), 3)\n    ax.imshow(image)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Make puzzle pool\n\ndf = pd.read_csv('train.csv')\n\nbboxs = np.stack(df['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\nfor i, column in enumerate(['x', 'y', 'w', 'h']):\n    df[column] = bboxs[:,i]\ndf.drop(columns=['bbox'], inplace=True)\ndf.head()\n\n\npool = []\nimage_ids = set(df.image_id.values)\n\nfor i, image_id in enumerate(image_ids):\n    \n    filtered = df[df['image_id'] == image_id]\n    bboxes = filtered[['x', 'y', 'w', 'h']].values\n\n    puzzles = make_puzzles(image_id, bboxes, min_bbox_size=20)\n    \n    pool += puzzles","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#%rm -r merged_puzzles\n#%rm merged_puzzles.csv\n\nk = 5000 # NUM NEW IMAGES\nAugData = collections.namedtuple('AugData', 'image_id,x_min,y_min,x_max,y_max,width,height,area,source')\n\naug_data = []\n\nos.makedirs('./merged_puzzles')\n\nfor i in range(k):\n    random.shuffle(pool)\n    a = ([bbox for img, bbox in pool[:4]])\n    merged_image, merged_bboxes = merge_random_4_puzzles(pool[:4])\n    ih, iw, ch = merged_image.shape\n    image_id = str(uuid.uuid4())\n    for bbox in merged_bboxes:\n        x, y, w, h = bbox\n        \n        #aug_data.append(AugData(image_id=image_id, width=iw, height=ih, source='aug', x=x, y=y, w=w, h=h))\n        \n        aug_data.append(AugData(image_id=image_id, x_min=x, y_min=y, x_max=x+w, y_max=y+h, width=w, height=h, area=w*h, source='aug'))\n\n    merged_image = cv2.cvtColor(merged_image, cv2.COLOR_RGB2BGR)\n    cv2.imwrite('merged_puzzles/{}.jpg'.format(image_id), merged_image)         ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SAVE new data","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"aug_df = pd.DataFrame(data=aug_data)\naug_df.to_csv('merged_puzzles.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dfvv = pd.read_csv('merged_puzzles.csv')\ntrain_dfvv.head()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"final_train_df = pd.read_csv(os.path.join('merged_puzzles.csv'))\n\nnum_images_final = final_train_df['image_id'].unique()\nprint(f'Total number of training images in csv: {len(num_images_final)}')\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# combined train.csv + merged_puzzles.csv","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"extension = 'csv'\nall_filenames = [i for i in glob('*.{}'.format(extension))]\nall_filenames","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#combine all files in the list\ncombined_csv = pd.concat([pd.read_csv(f) for f in all_filenames ])\n#export to csv\ncombined_csv.to_csv( \"combined_csv.csv\", index=False, encoding='utf-8-sig')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!zip -r -qq merged_puzzles.zip merged_puzzles\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Start work with traing data","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"final_train_df = pd.read_csv('combined_csv.csv')\n\nprint(len(df_train['image_id'].unique()))\nprint((df_train.shape[0])/len(df_train['image_id'].unique()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Image_id v/s # of bounding boxes\")\nprint(df_train['image_id'].value_counts())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Height\")\nprint(df_train['height'].value_counts())\nprint(\"Width\")\ndf_train['width'].value_counts()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv('combined_csv.csv')\n\n# Let's expand the bounding box coordinates and calculate the area of all the bboxes\ntrain_df[['x_min','y_min', 'width', 'height']] = pd.DataFrame([ast.literal_eval(x) for x in train_df.bbox.tolist()], index= train_df.index)\ntrain_df = train_df[['image_id', 'bbox', 'source', 'x_min', 'y_min', 'width', 'height']]\ntrain_df['area'] = train_df['width'] * train_df['height']\ntrain_df['x_max'] = train_df['x_min'] + train_df['width']\ntrain_df['y_max'] = train_df['y_min'] + train_df['height']\ntrain_df = train_df.drop(['bbox'], axis=1)\ntrain_df = train_df[['image_id', 'x_min', 'y_min', 'x_max', 'y_max', 'width', 'height', 'area', 'source']]\n\n# There are some buggy annonations in training images having huge bounding boxes. Let's remove those bboxes\ntrain_df = train_df[train_df['area'] < 100000]\n\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(train_df.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\nprint(f'Total number of training images: {len(image_ids)}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Read the image on which data augmentaion is to be performed\nimage_id = 'c14c1e300'\nimage = cv2.imread(os.path.join(WORK_DIR, 'train', f'{image_id}.jpg'), cv2.IMREAD_COLOR)\nimage = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\nimage /= 255.0\nplt.figure(figsize = (10, 10))\nplt.imshow(image)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"There are two major formats of bounding boxes:\n\n1. **pascal_voc**, which is [x_min, y_min, x_max, y_max]\n2. **COCO**, which is [x_min, y_min, width, height]\n\nWe'll see how to perform image augmentations for both the formats. Let's first start with **pascal_voc** format.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"\npascal_voc_boxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'x_max', 'y_max']].astype(np.int32).values\ncoco_boxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'width', 'height']].astype(np.int32).values\nassert(len(pascal_voc_boxes) == len(coco_boxes))\nlabels = np.ones((len(pascal_voc_boxes), ))\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# class WheatDataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n    \n    def __init__(self, df):\n        self.df = df\n        self.image_ids = self.df['image_id'].unique()\n\n    def __len__(self):\n        return len(self.image_ids)\n    \n    def __getitem__(self, index):\n        image_id = self.image_ids[index]\n        image = cv2.imread(os.path.join(WORK_DIR, 'train', f'{image_id}.jpg'), cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0  # Normalize\n        \n        # Get bbox coordinates for each wheat head(s)\n        bboxes_df = self.df[self.df['image_id'] == image_id]\n        boxes, areas = [], []\n        n_objects = len(bboxes_df)  # Number of wheat heads in the given image\n\n        for i in range(n_objects):\n            x_min = bboxes_df.iloc[i]['x_min']\n            x_max = bboxes_df.iloc[i]['x_max']\n            y_min = bboxes_df.iloc[i]['y_min']\n            y_max = bboxes_df.iloc[i]['y_max']\n\n            boxes.append([x_min, y_min, x_max, y_max])\n            areas.append(bboxes_df.iloc[i]['area'])\n\n        return {\n            'image_id': image_id,\n            'image': image,\n            'boxes': boxes,\n            'area': areas,\n        }\n    \n    \ndef collate_fn(batch):\n    images, bboxes, areas, image_ids = ([] for _ in range(4))\n    for data in batch:\n        images.append(data['image'])\n        bboxes.append(data['boxes'])\n        areas.append(data['area'])\n        image_ids.append(data['image_id'])\n\n    return np.array(images), np.array(bboxes), np.array(areas), np.array(image_ids)     ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH_SIZE = 16\n\ntrain_dataset = WheatDataset(train_df)\ntrain_loader = data_utils.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# class CustomCutout - AUGMENTATION square cutout regions","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomCutout(DualTransform):\n    \"\"\"\n    Custom Cutout augmentation with handling of bounding boxes \n    Note: (only supports square cutout regions)\n\n    \"\"\"\n    \n    def __init__(\n        self,\n        fill_value=0,\n        bbox_removal_threshold=0.50,\n        min_cutout_size=120,  # SIZE\n        max_cutout_size=512,\n        always_apply=False,\n        p=0.5\n    ):\n        \"\"\"\n        Class construstor\n        \n        :param fill_value: Value to be filled in cutout (default is 0 or black color)\n        :param bbox_removal_threshold: Bboxes having content cut by cutout path more than this threshold will be removed\n        :param min_cutout_size: minimum size of cutout (192 x 192)\n        :param max_cutout_size: maximum size of cutout (512 x 512)\n        \"\"\"\n        super(CustomCutout, self).__init__(always_apply, p)  # Initialize parent class\n        self.fill_value = fill_value\n        self.bbox_removal_threshold = bbox_removal_threshold\n        self.min_cutout_size = min_cutout_size\n        self.max_cutout_size = max_cutout_size\n        \n    def _get_cutout_position(self, img_height, img_width, cutout_size):\n        \"\"\"\n        Randomly generates cutout position as a named tuple\n        \n        :param img_height: height of the original image\n        :param img_width: width of the original image\n        :param cutout_size: size of the cutout patch (square)\n        :returns position of cutout patch as a named tuple\n        \"\"\"\n        position = namedtuple('Point', 'x y')\n        return position(\n            np.random.randint(0, img_width - cutout_size + 1),\n            np.random.randint(0, img_height - cutout_size + 1)\n        )\n        \n    def _get_cutout(self, img_height, img_width):\n        \"\"\"\n        Creates a cutout pacth with given fill value and determines the position in the original image\n        \n        :param img_height: height of the original image\n        :param img_width: width of the original image\n        :returns (cutout patch, cutout size, cutout position)\n        \"\"\"\n        cutout_size = np.random.randint(self.min_cutout_size, self.max_cutout_size + 1)\n        cutout_position = self._get_cutout_position(img_height, img_width, cutout_size)\n        return np.full((cutout_size, cutout_size, 3), self.fill_value), cutout_size, cutout_position\n        \n    def apply(self, image, **params):\n        \"\"\"\n        Applies the cutout augmentation on the given image\n        \n        :param image: The image to be augmented\n        :returns augmented image\n        \"\"\"\n        image = image.copy()  # Don't change the original image\n        self.img_height, self.img_width, _ = image.shape\n        cutout_arr, cutout_size, cutout_pos = self._get_cutout(self.img_height, self.img_width)\n        \n        # Set to instance variables to use this later\n        self.image = image\n        self.cutout_pos = cutout_pos\n        self.cutout_size = cutout_size\n        \n        image[cutout_pos.y:cutout_pos.y+cutout_size, cutout_pos.x:cutout_size+cutout_pos.x, :] = cutout_arr\n        return image\n    \n    def apply_to_bbox(self, bbox, **params):\n        \"\"\"\n        Removes the bounding boxes which are covered by the applied cutout\n        \n        :param bbox: A single bounding box coordinates in pascal_voc format\n        :returns transformed bbox's coordinates\n        \"\"\"\n\n        # Denormalize the bbox coordinates\n        bbox = denormalize_bbox(bbox, self.img_height, self.img_width)\n        x_min, y_min, x_max, y_max = tuple(map(int, bbox))\n\n        bbox_size = (x_max - x_min) * (y_max - y_min)  # width * height\n        overlapping_size = np.sum(\n            (self.image[y_min:y_max, x_min:x_max, 0] == self.fill_value) &\n            (self.image[y_min:y_max, x_min:x_max, 1] == self.fill_value) &\n            (self.image[y_min:y_max, x_min:x_max, 2] == self.fill_value)\n        )\n\n        # Remove the bbox if it has more than some threshold of content is inside the cutout patch\n        if overlapping_size / bbox_size > self.bbox_removal_threshold:\n            return normalize_bbox((0, 0, 0, 0), self.img_height, self.img_width)\n\n        return normalize_bbox(bbox, self.img_height, self.img_width)\n\n    def get_transform_init_args_names(self):\n        \"\"\"\n        Fetches the parameter(s) of __init__ method\n        :returns: tuple of parameter(s) of __init__ method\n        \"\"\"\n        return ('fill_value', 'bbox_removal_threshold', 'min_cutout_size', 'max_cutout_size', 'always_apply', 'p')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Add augmentation albumentations","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"#rm\nfirst_version = '''\naugmentation0 = albumentations.Compose([\n    CustomCutout(p=0.5),\n    A.Flip(p=0.60),\n    A.RandomRotate90(p=0.5),\n    A.RandomBrightness(limit=0.3, p=0.60),\n    A.OneOf([  # One of blur or adding gauss noise\n        A.Blur(p=0.50),  # Blurs the image\n        A.GaussNoise(var_limit=5.0 / 255.0, p=0.50)  # Adds Gauss noise to image\n    ], p=0.5)\n], bbox_params = {\n    'format': 'pascal_voc',\n    'label_fields': ['labels']\n})\n\n'''","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### determining the augmentations used ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# CustomCutout(p=.5) # function call\n\nflip = A.Flip(p=.6)\nrot90 = A.RandomRotate90(p=.5)\n\nbr_contr = A.RandomBrightnessContrast(brightness_limit=.3, contrast_limit=.3, p=.5)\nbrigh = A.RandomBrightness(limit=.3, p=.6)\ncontrast = A.RandomContrast(limit=.3, p=.6)\n\nblur = A.Blur(p=.3)\nnoise = A.GaussNoise(var_limit=5.0 / 255.0, p=.3)\n\ndef oneof(arr=[blur, noise], p=.5):\n    return A.OneOf( arr, p )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"required_aug = [\n    [\n        CustomCutout(p=.5), \n        flip, \n        rot90, \n        brigh, \n        oneof()], \n    [\n        CustomCutout(p=.5), \n        flip, \n        rot90, \n        contrast, \n        oneof()]\n]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualization augmented images","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_bbox(bboxes, col, color='white'):\n    for i in range(len(bboxes)):\n        # Create a Rectangle patch\n        rect = patches.Rectangle(\n            (bboxes[i][0], bboxes[i][1]),\n            bboxes[i][2] - bboxes[i][0], \n            bboxes[i][3] - bboxes[i][1], \n            linewidth=2, \n            edgecolor=color, \n            facecolor='none')\n\n        # Add the patch to the Axes\n        col.add_patch(rect)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = A.Compose(\n    required_aug[0], \n    bbox_params = {\n    'format': 'pascal_voc',\n    'label_fields': ['labels']\n})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_images = 5\nrand_start = np.random.randint(0, len(image_ids) - 5)\nfig, ax = plt.subplots(nrows=num_images, ncols=2, figsize=(16, 40))\n\nfor index, image_id in enumerate(image_ids[rand_start : rand_start + num_images]):\n    # Read the image from image id\n    image = cv2.imread(os.path.join(WORK_DIR, 'train', f'{image_id}.jpg'), cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n    image /= 255.0  # Normalize\n    \n    # Get the bboxes details and apply all the augmentations\n    bboxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'x_max', 'y_max']].astype(np.int32).values\n    labels = np.ones((len(bboxes), ))  # As we have only one class (wheat heads)\n        \n    aug_result = transform(image=image, bboxes=bboxes, labels=labels)\n\n    get_bbox(bboxes, ax[index][0], color='red')\n    ax[index][0].grid(False)\n    ax[index][0].set_xticks([])\n    ax[index][0].set_yticks([])\n    ax[index][0].title.set_text('Original Image')\n    ax[index][0].imshow(image)\n\n    get_bbox(aug_result['bboxes'], ax[index][1], color='red')\n    ax[index][1].grid(False)\n    ax[index][1].set_xticks([])\n    ax[index][1].set_yticks([])\n    ax[index][1].title.set_text(f'Augmented Image: Removed bboxes: {len(bboxes) - len(aug_result[\"bboxes\"])}')\n    ax[index][1].imshow(aug_result['image'])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"required_aug = [\n    [\n        CustomCutout(p=.5), \n        flip, \n        rot90, \n        brigh, \n        oneof()\n    ], \n    [\n        CustomCutout(p=.5), \n        flip, \n        rot90, \n        contrast, \n        oneof()\n    ]\n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"coefficient = len(required_aug)\ncoefficient","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_dataset(index, image_id, coefficient = 0):\n    # Read the image from image id\n    image = cv2.imread(os.path.join(WORK_DIR, 'train', f'{image_id}.jpg'), cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n    # Get the bboxes details and apply all the augmentations\n    bboxes = train_df[train_df['image_id'] == image_id][['x_min', 'y_min', 'x_max', 'y_max']].astype(np.int32).values\n    source = train_df[train_df['image_id'] == image_id]['source'].unique()[0]\n    labels = np.ones((len(bboxes), ))  # As we have only one class (wheat heads)\n    \n    some_aug = required_aug[coefficient]\n    \n    transform = A.Compose(\n        some_aug, \n        bbox_params = {\n        'format': 'pascal_voc',\n        'label_fields': ['labels']})\n    \n    aug_result = transform(image=image, bboxes=bboxes, labels=labels)\n    name_img_aug = f'{image_id}_aug_{coefficient}'\n\n    aug_image = aug_result['image']\n    aug_bboxes = aug_result['bboxes']\n    \n    Image.fromarray(image).save(os.path.join(WORK_DIR, 'train', f'{image_id}.jpg'))\n    Image.fromarray(aug_image).save(os.path.join(WORK_DIR, 'train', f'{name_img_aug}.jpg'))\n\n    image_metadata = []\n    for bbox in aug_bboxes:\n        bbox = tuple(map(int, bbox))\n        image_metadata.append({\n            'image_id': name_img_aug,\n            'x_min': bbox[0],\n            'y_min': bbox[1],\n            'x_max': bbox[2],\n            'y_max': bbox[3],\n            'width': bbox[2] - bbox[0],\n            'height': bbox[3] - bbox[1],\n            'area': (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]),\n            'source': source\n        })\n    return image_metadata","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#%rm train.csv\n#%rm -r train\n\nif not os.path.isdir('train'):\n    os.mkdir('train')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for el_k in range(coefficient):\n    image_metadata = Parallel(n_jobs=8)(delayed(create_dataset)(index, image_id, el_k) for index, image_id in tqdm(enumerate(image_ids), total=len(image_ids)))\n    image_metadata = [item for sublist in image_metadata for item in sublist]\n    coefficient -= 1\n    aug_train_df = pd.DataFrame(image_metadata)\n    train_df = pd.concat([train_df, aug_train_df]).reset_index(drop=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(aug_train_df.shape)\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Save result","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Add a new column to store kfold indices\ntrain_df.loc[:, 'kfold'] = -1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_source = train_df[['image_id', 'source']].drop_duplicates()\n\n# get lists for image_ids and sources\nimage_ids = image_source['image_id'].to_numpy()\nsources = image_source['source'].to_numpy()\n\nskf = StratifiedKFold(n_splits=5, shuffle=True, random_state=1996)\nsplit = skf.split(image_ids, sources) # second arguement is what we are stratifying by\n\nfor fold, (train_idx, val_idx) in enumerate(split):\n    translated_val_idx = train_df[train_df['image_id'].isin(image_ids[val_idx])].index.values\n    print(len(translated_val_idx))\n    train_df.loc[translated_val_idx, 'kfold'] = fold\n    \ntrain_df.to_csv('train.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Check result","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"final_train_df = pd.read_csv(os.path.join('train.csv'))\n\nnum_images_final = final_train_df['image_id'].unique()\nprint(f'Total number of training images: {len(num_images_final)}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_files = len(os.listdir('./train/'))\nnum_files","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}