{"cells":[{"metadata":{},"cell_type":"markdown","source":"* The script was based on Alex Shonenkov's awesome https://www.kaggle.com/shonenkov/training-efficientdet notebook, all credits to him!\n* The following is an implementation for the fasterrcnn model, with image and bounding boxes augmentation taken from Alex's notebook","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\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\nimport os\nfrom datetime import datetime\nimport time\nimport random\nimport cv2\nimport pandas as pd\nimport numpy as np\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nfrom glob import glob\n\nSEED = 42\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n\nseed_everything(SEED)\n\nDIR = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR}/train'\nDIR_TEST = f'{DIR}/test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR}/train.csv')\nprint(train_df.shape)\ntrain_df.sample(5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"bboxs = np.stack(train_df['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\nfor i, column in enumerate(['x', 'y', 'w', 'h']):\n    train_df[column] = bboxs[:,i]\ntrain_df.drop(columns=['bbox'], inplace=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n\n# df_folds = train_df[['image_id']].copy()\n# df_folds.loc[:, 'bbox_count'] = 1\n# df_folds = df_folds.groupby('image_id').count()\n# df_folds.loc[:, 'source'] = train_df[['image_id', 'source']].groupby('image_id').min()['source']\n# df_folds.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# df_folds.loc[:, 'stratify_group'] = np.char.add(\n#     df_folds['source'].values.astype(str),\n#     df_folds['bbox_count'].apply(lambda x: f'_{x // 15}').values.astype(str)\n# )\n# df_folds.loc[:, 'fold'] = 0\n\n# for fold_number, (train_index, val_index) in enumerate(skf.split(X=df_folds.index, y=df_folds['stratify_group'])):\n#     df_folds.loc[df_folds.iloc[val_index].index, 'fold'] = fold_number\n\n# df_folds.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Albumentations","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose(\n        [\n            A.OneOf([\n                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit= 0.2, \n                                     val_shift_limit=0.2, p=0.9),\n                A.RandomBrightnessContrast(brightness_limit=0.2, \n                                           contrast_limit=0.2, p=0.9),\n            ],p=0.9),\n            A.ToGray(p=0.01),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.5),\n            ToTensorV2(p=1.0),\n        ], \n        p=1.0, \n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0, \n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n\ndef get_valid_transforms():\n    return A.Compose(\n        [\n            ToTensorV2(p=1.0),\n        ], \n        p=1.0, \n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0, \n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_ROOT_PATH = '../input/global-wheat-detection/train'\nclass DatasetRetriever(Dataset):\n\n    def __init__(self, marking, image_ids, transforms=None, test=False):\n        super().__init__()\n        \n        self.image_ids = image_ids\n        self.df = marking\n        self.transforms = transforms\n        self.test = test\n\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        if self.test or random.random() > 0.5:\n            image, boxes, area, iscrowd = self.load_image_and_boxes(index)\n        else:\n            image, boxes, area, iscrowd = self.load_cutmix_image_and_boxes(index)\n\n        # there is only one class\n        labels = torch.ones((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            for i in range(10):\n                sample = self.transforms(**{\n                    'image': image,\n                    'bboxes': target['boxes'],\n                    'labels': labels\n                })\n                if len(sample['bboxes']) > 0:\n                    image = sample['image']\n                    target['boxes'] = torch.as_tensor(target['boxes'][:,1:], dtype = torch.float32)\n                    target['boxes'] = torch.stack(tuple(map(torch.Tensor, zip(*sample['bboxes'])))).permute(1, 0)\n                    break\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n    def load_image_and_boxes(self, index):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{DIR_TRAIN}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        records = self.df[self.df['image_id'] == image_id]\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        area = (boxes[:,3]-boxes[:,1])*(boxes[:,2]-boxes[:,0])\n        area = torch.as_tensor(area, dtype = torch.float32)\n        iscrowd = torch.zeros((records.shape[0]), dtype=torch.int64)\n        return image, boxes, area, iscrowd\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,_, iscrowd = self.load_image_and_boxes(index)\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        area = (result_boxes[:,3]-result_boxes[:,1])*(result_boxes[:,2]-result_boxes[:,0])\n        area = torch.as_tensor(area, dtype = torch.float32)\n        return result_image, result_boxes, area, iscrowd","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fold_number = 0\n\ntrain_dataset = DatasetRetriever(\n    image_ids=np.array(train_df['image_id'].unique()),\n    marking=train_df,\n    transforms=get_train_transforms(),\n    test=False,\n)\n\n# validation_dataset = DatasetRetriever(\n#     image_ids=df_folds[df_folds['fold'] == fold_number].index.values,\n#     marking=train_df,\n#     transforms=get_valid_transforms(),\n#     test=True,\n# )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images = []\nbxs = []\nimage_ids = []\nfor i in range(16):\n    imgs, tgts, img_ids = train_dataset[i]\n    boxes = tgts['boxes'].cpu().numpy().astype(np.int32)\n    images.append(imgs)\n    bxs.append(boxes)\n    image_ids.append(img_ids)\n\ndef trim_axs(axs, N):\n    \"\"\"little helper to massage the axs list to have correct length...\"\"\"\n    axs = axs.flat\n    for ax in axs[N:]:\n        ax.remove()\n    return axs[:N]\n\nfig, axs = plt.subplots(3,3,figsize = (23,18), constrained_layout=True, tight_layout = True)\naxs = trim_axs(axs, 8)\nfor i, ax in enumerate(axs):\n    single_boxes = bxs[i]\n    single_image = images[i].permute(1,2,0).cpu().numpy() \n    for box in single_boxes:\n        cv2.rectangle(single_image,\n                  (int(box[0]), int(box[1])),\n                  (int(box[2]), int(box[3])),\n                  220,3)\n\n    ax.set_axis_off()\n    ax.imshow(single_image)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Initialize model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nclass Fitter:\n    \n    def __init__(self, model, device, config):\n        self.config = config\n        self.epoch = 0\n\n        self.base_dir = f'./{config.folder}'\n        if not os.path.exists(self.base_dir):\n            os.makedirs(self.base_dir)\n        \n        self.log_path = f'{self.base_dir}/log.txt'\n        self.best_summary_loss = 10**5\n\n        self.model = model\n        self.device = device\n\n        param_optimizer = list(self.model.named_parameters())\n        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n        optimizer_grouped_parameters = [\n            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n        ] \n\n        self.optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=config.lr)\n        self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)\n        self.log(f'Fitter prepared. Device is {self.device}')\n\n    def fit(self, train_loader):\n        for e in range(self.config.n_epochs):\n            if self.config.verbose:\n                lr = self.optimizer.param_groups[0]['lr']\n                timestamp = datetime.utcnow().isoformat()\n                self.log(f'\\n{timestamp}\\nLR: {lr}')\n\n            t = time.time()\n            summary_loss = self.train_one_epoch(train_loader)\n\n            self.log(f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n\n            if summary_loss.avg < self.best_summary_loss:\n                self.best_summary_loss = summary_loss.avg\n                self.model.eval()\n                self.save(f'{self.base_dir}/best-checkpoint-{str(self.epoch).zfill(3)}epoch.bin')\n                for path in sorted(glob(f'{self.base_dir}/best-checkpoint-*epoch.bin'))[:-3]:\n                    os.remove(path)\n\n            if self.config.validation_scheduler:\n                self.scheduler.step(metrics=summary_loss.avg)\n\n            self.epoch += 1\n\n\n    def train_one_epoch(self, train_loader):\n        self.model.train()\n        summary_loss = AverageMeter()\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(train_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    print(\n                        f'Train Step {step}/{len(train_loader)}, ' + \\\n                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}', end='\\r'\n                    )\n            \n            images = list(image.to(self.device) for image in images)\n\n            batch_size = (len(images))\n            targets = [{k:v.to(self.device) for k,v in t.items()} for t in targets]\n\n\n            self.optimizer.zero_grad()\n            \n            loss_dict = self.model(images, targets)\n            losses = sum(loss for loss in loss_dict.values())\n\n            losses.backward()\n\n            summary_loss.update(losses.detach().item(), batch_size)\n\n            self.optimizer.step()\n\n            if self.config.step_scheduler:\n                self.scheduler.step()\n\n        return summary_loss\n    \n    def save(self, path):\n\n        torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')\n\n    def log(self, message):\n        if self.config.verbose:\n            print(message)\n        with open(self.log_path, 'a+') as logger:\n            logger.write(f'{message}\\n')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TrainGlobalConfig:\n    num_workers = 4\n    batch_size = 8 \n    n_epochs = 18\n    lr = 0.00003\n\n    folder = 'fasterrcnn-cutmix-augmix'\n\n    # -------------------\n    verbose = True\n    verbose_step = 1\n    # -------------------\n\n    # --------------------\n    step_scheduler = False  # do scheduler.step after optimizer.step\n    validation_scheduler = True  # do scheduler.step after validation stage loss\n    \n\n#     SchedulerClass = torch.optim.lr_scheduler.OneCycleLR\n#     scheduler_params = dict(\n#         max_lr=0.01,\n#         epochs=n_epochs,\n#         steps_per_epoch=int(len(train_dataset) / batch_size),\n#         pct_start=0.1,\n#         anneal_strategy='cos', \n#         final_div_factor=10**5\n#     )\n    \n    SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau\n    scheduler_params = dict(\n        mode='min',\n        factor=0.7,\n        patience=1,\n        verbose=False, \n        threshold=0.0001,\n        threshold_mode='abs',\n        cooldown=0, \n        min_lr=1e-8,\n        eps=1e-08\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ndef run_training():\n    device = torch.device('cuda')\n    model.to(device)\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=RandomSampler(train_dataset),\n        pin_memory=False,\n        drop_last=True,\n        num_workers=TrainGlobalConfig.num_workers,\n        collate_fn=collate_fn,\n    )\n\n    fitter = Fitter(model=model, device=device, config=TrainGlobalConfig)\n    fitter.fit(train_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained = True)\n\nc = 2 #number of classes\n\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, c)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"run_training()","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}