{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"code","source":"import random as rnd\nimport numpy as np\nimport cv2 as cv\nimport torch\n\n\n__all__ = [\n    'Mosaic',\n    'MixUp',\n    'RandomAugmentation'\n]\n\n\nclass Mosaic:\n    m_num_required_elements = 5\n\n    @staticmethod\n    def get_area(box: np.ndarray) -> float:\n        \"\"\"\n        ...\n\n        Args:\n            box (np.ndarray): coordinates of the box    [pascal-voc form]\n\n        Returns:\n            (float) Area of the box.\n        \"\"\"\n        return (box[2] - box[0]) * (box[3] - box[1])\n\n    @staticmethod\n    def get_overlap_percentage(b_1: np.ndarray, b_2: np.ndarray) -> tuple:\n        \"\"\"\n        Calculates the percentage of the overlap area with respect to box 1 and the overlap 'directions'.\n\n        Args:\n            b_1 (np.ndarray): coordinates of the first box  [pascal-voc form]\n            b_2 (np.ndarray): coordinates of the second box [pascal-voc form]\n\n        Returns:\n            (float) Percentage of the overlap area with respect to box 1.\n            (tuple) Overlap 'directions'.       (up, down, left, right)\n        \"\"\"\n\n        # In case a bbox coordinate is on the edge of an image (--> make it fit in the array)\n        if b_1[2] > 1023:\n            b_1[2] = 1023\n        if b_1[3] > 1023:\n            b_1[3] = 1023\n\n        # Calculate overlap area\n        d_x = min(b_1[2], b_2[2]) - max(b_1[0], b_2[0])\n        d_y = min(b_1[3], b_2[3]) - max(b_1[1], b_2[1])\n\n        if d_x < 0 != d_y < 0:\n            return 0.0, (False, False, False, False)  # not overlapping\n\n        overlap_area = d_x * d_y\n        box_1_area = Mosaic.get_area(box=b_1)\n\n        # Calculate overlap 'directions'        (up, down, left, right)\n        overlap_directions = (\n            (min(b_1[1], b_2[1]) == b_1[1]),\n            (max(b_1[3], b_2[3]) == b_1[3]),\n            (min(b_1[0], b_2[0]) == b_1[0]),\n            (max(b_1[2], b_2[2]) == b_1[2])\n        )\n\n        return (overlap_area / box_1_area), overlap_directions\n\n    @staticmethod\n    def apply(imgs: tuple, bboxes: tuple, p: float, p_mixup: float = 0.8) -> (np.ndarray, list):\n        if rnd.random() > p:\n            i = rnd.randint(0, 3)\n            return imgs[i], bboxes[i]\n\n        h, w, c = imgs[0].shape\n        h -= 1\n        w -= 1\n        shape = (h, w)\n\n        # Apply random MixUp\n        img_0, img_1, img_2, img_3, img_4 = imgs\n        bboxes_0, bboxes_1, bboxes_2, bboxes_3, bboxes_4 = bboxes\n        img_0, bboxes_0 = MixUp.apply(imgs=(img_0, img_1), bboxes=(bboxes_0, bboxes_1), p=p_mixup, p_chaos=0.0)\n        imgs = (img_0, img_2, img_3, img_4)\n        bboxes = (bboxes_0, bboxes_2, bboxes_3, bboxes_4)\n\n        rnd_x_scale = rnd.uniform(0.2, 0.8)\n        rnd_y_scale = rnd.uniform(0.2, 0.8)\n\n        img = Mosaic.apply_to_imgs(imgs, shape, rnd_x_scale, rnd_y_scale)\n        bboxes = Mosaic.apply_to_bboxes(bboxes, shape, rnd_x_scale, rnd_y_scale)\n\n        return img, bboxes\n\n    @staticmethod\n    def apply_to_imgs(imgs: tuple, shape: tuple, rnd_x_scale: float, rnd_y_scale: float) -> np.ndarray:\n        img_1, img_2, img_3, img_4 = imgs\n        h, w = shape\n\n        # Concat the four input images to upper and lower image pairs\n        img_lower = np.concatenate((img_3[int(rnd_y_scale*h):h, :int(rnd_x_scale*w)],\n                                    img_4[int(rnd_y_scale*h):h, int(rnd_x_scale*w):w]),\n                                   axis=1)\n        img_upper = np.concatenate((img_1[:int(rnd_y_scale*h), :int(rnd_x_scale*w)],\n                                    img_2[:int(rnd_y_scale*h), int(rnd_x_scale*w):w]),\n                                   axis=1)\n\n        # Concat upper and lower image pairs to final image\n        img = np.concatenate((img_upper, img_lower), axis=0)\n\n        return img\n\n    @staticmethod\n    def apply_to_bboxes(bboxes: tuple, shape: tuple, rnd_x_scale: float, rnd_y_scale: float) -> list:\n        h, w = shape\n        d_x = int(rnd_x_scale*w)\n        d_y = int(rnd_y_scale*h)\n        img_areas = [[0, 0, d_x, d_y], [d_x, 0, w, d_y], [0, d_y, d_x, h], [d_x, d_y, w, h]]    # pascal-voc form; [left_upper, right_upper, left_lower, right_lower]\n        overlap_threshold = 0.0\n        p_area_threshold = 0.25\n\n        keep_bboxes = []\n        for bboxes_i in range(len(bboxes)):\n            img_area = img_areas[bboxes_i]\n            for bbox in bboxes[bboxes_i]:\n                op, overlap_directions = Mosaic.get_overlap_percentage(b_1=bbox, b_2=img_area)\n                if op > overlap_threshold:\n                    area_0 = Mosaic.get_area(box=bbox)\n                    if op < 1.0:\n                        x_mode = bboxes_i % 2       # 0 for upper_left and lower_left; 1 for upper_right and lower_right\n                        y_mode = int(bboxes_i > 1)  # 0 for upper_left and upper_right; 1 for lower_left and lower_right\n                        if (overlap_directions[2] and x_mode == 1) or (overlap_directions[3] and x_mode == 0):\n                            bbox[(2-((bboxes_i % 2)*2))] = d_x\n                        if (overlap_directions[0] and y_mode == 1) or (overlap_directions[1] and y_mode == 0):\n                            bbox[(3-(int(bboxes_i > 1)*2))] = d_y\n                        p_area = Mosaic.get_area(box=bbox) / area_0\n                        if p_area > p_area_threshold:\n                            keep_bboxes.append((*bbox.tolist(), ))\n                    else:\n                        keep_bboxes.append((*bbox.tolist(), ))\n\n        return keep_bboxes\n\n\nclass MixUp:\n    m_num_required_elements = 3\n\n    @staticmethod\n    def apply(imgs: tuple, bboxes: tuple, p: float, p_chaos=0.3) -> (np.ndarray, list):\n        if rnd.random() > p:\n            i = rnd.randint(0, 1)\n            return imgs[i], bboxes[i]\n\n        if rnd.random() > p_chaos:\n            n_imgs = 2\n        else:\n            n_imgs = 3\n        return_img = MixUp.apply_to_imgs(imgs=imgs, n_imgs=n_imgs)\n        return_bboxes = MixUp.apply_to_bboxes(bboxes=bboxes, n_imgs=n_imgs)\n\n        return return_img, return_bboxes\n\n    @staticmethod\n    def apply_to_imgs(imgs: tuple, n_imgs: int) -> np.ndarray:\n        if n_imgs == 3:\n            img_1, img_2, img_3 = imgs\n        else:\n            img_1, img_2 = imgs[ :2]\n\n        #l = np.random.beta(a=16.0, b=16.0)\n        l = 0.5\n        new_img = cv.addWeighted(src1=img_1, alpha=l, src2=img_2, beta=(1-l), gamma=0.0)\n\n        if n_imgs == 3:\n            l = (l*2) / 3\n            new_img = cv.addWeighted(src1=new_img, alpha=l, src2=img_3, beta=(1-l), gamma=0.0)\n\n        return new_img\n\n    @staticmethod\n    def apply_to_bboxes(bboxes: tuple, n_imgs: int) -> list:\n        if n_imgs == 3:\n            bboxes_1, bboxes_2, bboxes_3 = bboxes\n        else:\n            bboxes_1, bboxes_2 = bboxes[ :2]\n\n        if n_imgs == 3:\n            cated_bboxes = torch.cat(tensors=(bboxes_1, bboxes_2, bboxes_3), dim=0)\n        else:\n            cated_bboxes = torch.cat(tensors=(bboxes_1, bboxes_2), dim=0)\n\n        return cated_bboxes\n\n\nclass RandomAugmentation:\n    m_augmenatation_distribution = [0.5, 0.5]\n    m_augmentation_methods = [\n        Mosaic,\n        MixUp\n    ]\n\n    @staticmethod\n    def get_random_augmentation(augmentation_distribution=m_augmenatation_distribution) -> object:\n        augmentation_method = rnd.choices(population=RandomAugmentation.m_augmentation_methods,\n                                          weights=augmentation_distribution, k=1)[0]\n\n        return augmentation_method\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\nimport torch\nfrom ast import literal_eval\nimport cv2 as cv\nimport numpy as np\nimport random as rnd\n\n\nclass WheatDetectionDataset:\n    def __init__(self, root, transforms, mode):\n        self.root = root\n        self.transforms = transforms\n        self.mode = mode\n\n        self.imgs = list(os.listdir(os.path.join(root, mode)))\n        if mode is 'train':\n            self.preds = pd.read_csv(os.path.join(root, 'train.csv'))\n\n    def get_img(self, index):\n        img_path = os.path.join(self.root, self.mode, self.imgs[index])\n        img = cv.imread(img_path, cv.IMREAD_COLOR)\n        img = cv.cvtColor(img, cv.COLOR_BGR2RGB).astype(np.float32)\n        #img = cv.GaussianBlur(img, ksize=(5, 5), sigmaX=0)\n        img /= 255.0\n\n        return img\n\n    def __getitem__(self, index, mode=0):\n        # Load image\n        img = self.get_img(index=index)\n\n        if self.mode == 'train':\n            # Load prediction boxes\n            img_boxes = self.preds.loc[self.preds['image_id'] == self.imgs[index].split('.')[0]]['bbox']\n            num_boxes = len(img_boxes)\n            boxes = []\n            # Box coords in pascal-voc form\n            for box in img_boxes:\n                coords = literal_eval(box)\n                x_min = float(coords[0])\n                y_min = float(coords[1])\n                x_max = float(x_min + coords[2])\n                y_max = float(y_min + coords[3])\n                boxes.append([x_min, y_min, x_max, y_max])\n            img_boxes = torch.as_tensor(boxes)\n            \n            img_id = torch.tensor([index])\n            \n            target = {}\n            target['boxes'] = img_boxes\n            #target['labels'] = None\n            target['image_id'] = img_id\n            #target['area'] = None\n            #target['iscrowd'] = None\n            \n            if self.transforms is not None and mode == 0:\n                sample = {\n                    'image': img,\n                    'bboxes': target['boxes'],\n                    'labels': None\n                }\n                \n                # Apply own augmentation\n                augmentation_method = RandomAugmentation.get_random_augmentation(augmentation_distribution=[0.5, 0.5])\n                imgs = []\n                bboxes = []\n                for i in range(augmentation_method.m_num_required_elements):\n                    apply_img, apply_target = self.__getitem__(index=rnd.randint(0, (len(self.imgs) - 1)), mode=1)\n                    imgs.append(apply_img)\n                    bboxes.append(apply_target['boxes'])\n                imgs = tuple(imgs)\n                bboxes = tuple(bboxes)\n                \n                sample['image'], sample['bboxes'] = augmentation_method.apply(\n                    imgs=imgs,\n                    bboxes=bboxes,\n                    p=0.9\n                )\n                \n                sample = self.transforms(**sample)\n                img = sample['image']\n                num_boxes = len(sample['bboxes'])\n                target['boxes'] = torch.tensor(sample['bboxes'], dtype=torch.float32)\n\n            try:\n                area = (img_boxes[:, 3] - img_boxes[:, 1]) * (img_boxes[:, 2] - img_boxes[:, 0])\n            except IndexError:\n                area = 0\n\n            # set up labels\n            labels = torch.ones((num_boxes, ), dtype=torch.int64)\n\n            is_crowd = torch.zeros((num_boxes, ), dtype=torch.int64)\n\n            target['area'] = torch.as_tensor(area, dtype=torch.float32)\n            target['labels'] = labels\n            target['is_crowd'] = is_crowd\n\n            return img, target\n\n        else:\n            if self.transforms is not None:\n                img = self.transforms(img)\n\n            return img, self.imgs[index]\n\n    def __len__(self) -> int:\n        return len(self.imgs)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# using: http://openaccess.thecvf.com/content_CVPR_2019/papers/He_Bag_of_Tricks_for_Image_Classification_with_Convolutional_Neural_Networks_CVPR_2019_paper.pdf\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nimport torchvision.transforms as T\nimport torchvision\nimport torch\nimport numpy as np\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport matplotlib.pyplot as plt\nimport cv2 as cv\n\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n\ndef get_transform(train):\n    transforms = []\n    if train:\n        transforms = [T.ToPILImage()] + augmentation.get_augmentation(n=2)\n    transforms.append(T.ToTensor())\n    #transforms.append(T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))\n    return T.Compose(transforms)\n\n\ndef get_train_transform():\n    return A.Compose([\n        A.RandomContrast(p=0.1),\n        A.RandomBrightness(p=0.1),\n        A.RandomGamma(p=0.1),\n        A.GaussianBlur(p=0.1),\n        A.HueSaturationValue(p=0.1),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\n\n# using: https://www.kaggle.com/maherdeebcv/test-pytorch-faster-r-cnn-with-resnet152-backbone\ndef fasterrcnn_resnet_fpn(pretrained=False, num_classes=2, pretrained_backbone=False, **kwargs):\n    if pretrained:\n        pretrained_backbone = False\n    \n    backbone = resnet_fpn_backbone(backbone_name='resnet152', pretrained=pretrained_backbone)\n    model = FasterRCNN(backbone=backbone, num_classes=num_classes, **kwargs)\n\n    return model\n\n\ndef get_model(num_classes):\n    #model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n    model = fasterrcnn_resnet_fpn(pretrained=False, pretrained_backbone=True)\n\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    return model\n\n\ndef load_checkpoint(checkpoint_path, model, optim):\n    checkpoint = torch.load(checkpoint_path)\n    model.load_state_dict(checkpoint['model_state_dict'])\n    optim.load_state_dict(checkpoint['optimizer_state_dict'])\n\n    return model, optim, checkpoint['epoch']\n\n\ndef main():\n    device = torch.device('cuda')\n    \n    continue_training = True\n\n    num_classes = 2     # 1 class (wheat) + background\n    #dataset_train = WheatDetectionDataset(root='data', transforms=get_transform(train=True), mode='train')\n    dataset_train = WheatDetectionDataset(root='../input/global-wheat-detection', transforms=get_train_transform(), mode='train')\n\n    # define dataset loaders\n    data_loader_train = torch.utils.data.DataLoader(\n        dataset_train, batch_size=8, shuffle=True, num_workers=0, collate_fn=collate_fn\n    )\n\n    # get model\n    model = get_model(num_classes).to(device)\n    #model = torch.load('models/wheat_model_10.0.pth', map_location=device).train()\n    print('Model loaded successfully.')\n\n    i_lr = 0.005      # initial learning rate\n    num_batches_warmup = 325\n    params = [p for p in model.parameters() if p.requires_grad]\n    optim = torch.optim.SGD(params=params, lr=i_lr, momentum=0.9, weight_decay=0.0005)\n    \n    start_epoch = 0\n    \n    if continue_training:\n        model, optim, start_epoch = load_checkpoint(checkpoint_path='../input/wheat-detection-model/wheat_model_checkpoint_11.4.pth', \n                                                    model=model, optim=optim)\n        print('Model checkpoint loaded successfully.')\n    \n    num_epochs = start_epoch + 3\n\n    len_dataloader = len(data_loader_train)\n    \n    for epoch in range(start_epoch, num_epochs):\n        model.train()\n        i = 0\n\n        loss_hist = []\n        for imgs, targets in data_loader_train:\n            i += 1\n\n            imgs = [img.to(device) for img in imgs]\n            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n            try:\n                \"\"\"test_img = imgs[0].detach().permute(1, 2, 0).to('cpu').numpy()\n                bboxes = targets[0]['boxes'].detach().to('cpu')\n                for it, bbox in enumerate(bboxes):\n                    bbox = bbox.numpy()\n                    p_1 = (int(bbox[0]), int(bbox[1]))\n                    p_2 = (int(bbox[2]), int(bbox[3]))\n                    test_img = cv.rectangle(img=test_img, pt1=p_1, pt2=p_2, color=(0, 0, 255))\n                #print(bboxes)\n                #print(test_img.shape)\n                cv.imshow('TEST_0', test_img)\n                cv.waitKey()\"\"\"\n\n                loss_dict = model(imgs, targets)\n                losses = sum(loss for loss in loss_dict.values())\n\n                optim.zero_grad()\n                losses.backward()\n                optim.step()\n\n                loss_hist.append(losses.__float__())\n\n                # adjust learning rate\n                current_batch = (epoch * len_dataloader) + i\n                if current_batch <= num_batches_warmup:\n                    # learning rate warmup\n                    # starting with a too big learning rate may result in something unwanted\n                    lr = current_batch * (i_lr / num_batches_warmup)\n                else:\n                    # cosine learning rate decay\n                    # (smoother than step learning rate decay)\n                    lr = i_lr * 0.5 * (1 + np.cos(((current_batch-num_batches_warmup) * np.pi) / (len_dataloader * num_epochs)))\n                for param_group in optim.param_groups:\n                    param_group['lr'] = lr\n\n                if i % 214 == 0:\n                    print(f'Iteration: {i}/{len_dataloader}, Loss: {np.mean(loss_hist)}, Epoch: {epoch}')\n                    loss_hist = []\n\n            except IndexError:\n                continue\n\n    torch.save(model, '/kaggle/working/wheat_model_11.5.pth')\n    torch.save({\n        'epoch': num_epochs,\n        'model_state_dict': model.state_dict(),\n        'optimizer_state_dict': optim.state_dict(),\n    }, '/kaggle/working/wheat_model_checkpoint_11.5.pth')\n\n\nmain()\n","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}