{"cells":[{"metadata":{},"cell_type":"markdown","source":"### ***A quick augmentation to squeeze out the last bit of fine-tuning in the last stretch!***","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"We have already seen the cutmix augmentation where 4 images and associated bounding boxes are joined together in a single image while training. \n\nI have experimented and found encouraging better results while using the stitch-boxes augmentation. In this we'll just pick up image patches within bounding boxes from various images and paste it over a base image.\n\nLet me know if anyone had already done this somewhere else. Was too lazy to search :p","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\nimport random\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\nfrom matplotlib import pyplot as plt\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\n\nvalid_ids = image_ids[-600:]\ntrain_ids = image_ids[:-600]\n\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\n\nvalid_df.shape, train_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Used the pytorch dataset here but it's exactly the same while doing in Keras too!","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, test=False, transforms=None):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n        self.test = test\n        \n    def fetch(self, index):\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        return image, boxes\n    \n    # the stitch boxes aug\n    \n    def stitch_boxes(self, index, imsize=1024):\n        \n        # Increase this to pick less number of boxes from each image and vice versa.\n        BOX_NUM_SELECTOR = 5\n        \n        all_boxes = []\n        \n        image, boxes = self.fetch(index)\n        for box in boxes:\n            all_boxes.append([int(x) for x in box])\n        \n        indexes = [random.randint(0, self.image_ids.shape[0] - 1) for _ in range(3)]\n        \n        for i, idx in enumerate(indexes):\n            _im, _boxes = self.fetch(idx)\n            _bi = [int(x) for x in random.sample(list(range(len(_boxes))), len(_boxes)//BOX_NUM_SELECTOR)]\n            \n            _b = _boxes[_bi]\n            for box in _b:\n                box = [int(x) for x in box]\n                all_boxes.append(box)\n                image[box[1]:box[3], box[0]:box[2]] = _im[box[1]:box[3], box[0]:box[2]]\n        \n        return image, np.array(all_boxes)\n        \n    \n    # the cutmix aug\n    \n    def load_cutmix_image_and_boxes(self, index, imsize=1024):\n        \"\"\" \n        This implementation of cutmix author:  https://www.kaggle.com/nvnnghia \n        Refactoring and adaptation: https://www.kaggle.com/shonenkov\n        \"\"\"\n        w, h = imsize, imsize\n        s = imsize // 2\n    \n        xc, yc = [int(random.uniform(imsize * 0.25, imsize * 0.75)) for _ in range(2)]  # center x, y\n        indexes = [index] + [random.randint(0, self.image_ids.shape[0] - 1) for _ in range(3)]\n    \n        result_image = np.full((imsize, imsize, 3), 1, dtype=np.float32)\n        result_boxes = []\n    \n        for i, index in enumerate(indexes):\n            image, boxes = self.fetch(index)\n            \n            \n            if i == 0:\n                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)\n                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)\n            elif i == 1:  # top right\n                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc\n                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h\n            elif i == 2:  # bottom left\n                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)\n                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)\n            elif i == 3:  # bottom right\n                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)\n                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)\n            result_image[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]\n            padw = x1a - x1b\n            padh = y1a - y1b\n    \n            boxes[:, 0] += padw\n            boxes[:, 1] += padh\n            boxes[:, 2] += padw\n            boxes[:, 3] += padh\n    \n            result_boxes.append(boxes)\n    \n        result_boxes = np.concatenate(result_boxes, 0)\n        np.clip(result_boxes[:, 0:], 0, 2 * s, out=result_boxes[:, 0:])\n        result_boxes = result_boxes.astype(np.int32)\n        result_boxes = result_boxes[np.where((result_boxes[:,2]-result_boxes[:,0])*(result_boxes[:,3]-result_boxes[:,1]) > 0)]\n        return result_image, result_boxes\n\n    def __getitem__(self, index: int):\n        \n        # randomizing augentations\n        dice = random.random()\n        if self.test == True or dice < 0.25:\n            image, boxes = self.fetch(index)\n        elif dice < 0.75:\n            image, boxes = self.stitch_boxes(index)\n        else:\n            image, boxes = self.load_cutmix_image_and_boxes(index)\n        \n        \n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # there is only one class\n        labels = torch.ones((boxes.shape[0],), dtype=torch.int64)\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((boxes.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n        image_id = 0\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Alright, let's see the image.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"traindata = WheatDataset(train_df, DIR_TRAIN, get_train_transform())\n\nim, boxes = traindata.stitch_boxes(55)\nplt.figure(figsize=(16,16))\nplt.imshow(im)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now let's visualize with boxes.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"im, boxes = traindata.stitch_boxes(111)\n\nfor box in boxes:\n    cv2.rectangle(im, (box[0], box[1]), (box[2], box[3]), (255,0,0), 1)\n\nplt.figure(figsize=(16,16))\nplt.imshow(im)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# This may be a simple image/ image with cutmix/ image with stitch boxes\nim, target, _ = traindata[99]\n\nfor box in target['boxes']:\n    cv2.rectangle(im, (box[0], box[1]), (box[2], box[3]), (255,0,0), 1)\n\nplt.figure(figsize=(16,16))\nplt.imshow(im)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"That's it! Use your own model and load the data from the dataset. Happy last minute training and fine-tuning!","execution_count":null},{"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}