{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Global Wheat Detection -- FasterRCNN with ResNet50 backbone","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport cv2\n\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm.autonotebook import tqdm\n\nimport pathlib\nimport os\n\nIS_KAGGLE_ENV = True\nDIR_INPUT = pathlib.Path('/kaggle/input/global-wheat-detection')\nDIR_WEIGHTS = pathlib.Path('/kaggle/input/resnet50-weights-imagenet-pth')\nFILENAME_WEIGHTS = 'resnet50-19c8e357.pth'\n\nif not IS_KAGGLE_ENV:  # My local machine.\n    DIR_INPUT = pathlib.Path('.')\n    DIR_WEIGHTS = pathlib.Path(f'{str(pathlib.Path.home())}/.cache/torch/hub/checkpoints')\n\nos.listdir(DIR_INPUT)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = pd.read_csv(DIR_INPUT / 'train.csv')\ndf_train.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Feature Transformation -- Group by `image_id`","execution_count":null},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"def df_to_list_of_dict_dataset(df: pd.DataFrame) -> list:\n    \"\"\"\n    Transform the training DataFrame to list of dict.\n    Each sample contains following keys mainly:\n        `image_id`: Unique image id locates on single image.\n        `bboxes`: Multiple bounding boxes for the image.\n        `labels`: Same number of bboxes labeled with ones,\n            for this task is just to predict where foregrounds locate at.\n    \"\"\"\n    ds_dict = {}\n    for idx, line in df.iterrows():\n        if line['image_id'] not in ds_dict:\n            ds_dict[line['image_id']] = {\n                'image_id': line['image_id'],\n                'width': float(line['width']),\n                'height': float(line['height']),\n                'bboxes': [eval(line['bbox'])],\n            }\n        else:\n            ds_dict[line['image_id']]['bboxes'].append(eval(line['bbox']))\n\n    ds_list = [sample for sample in ds_dict.values()]\n    for sample in ds_list:\n        sample['bboxes'] = np.asarray(sample['bboxes'], dtype='int64')\n        sample['labels'] = np.ones(shape=(len(sample['bboxes']),), dtype='int64')\n    return ds_list\n\n\nds_list_trn = df_to_list_of_dict_dataset(df_train)\nds_list_trn[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare Dataset and Augmentations for training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDatasetTrain(Dataset):\n    \n    def __init__(self, py_data: list, img_dir: pathlib.Path, transforms=None):\n        self.py_data = py_data\n        self.img_dir = img_dir\n        self.transforms = transforms\n    \n    def __len__(self):\n        return len(self.py_data)\n    \n    def __getitem__(self, item):\n        if torch.is_tensor(item):\n            item = item.tolist()\n        \n        sample = self.py_data[item]\n        image = cv2.imread(str(self.img_dir / f\"{sample['image_id']}.jpg\"))[..., ::-1].copy()\n        \n        if self.transforms is not None:\n            sample_to_transform = {'image': image, 'bboxes': sample['bboxes'],\n                                   'labels': sample['labels']}\n            sample_transformed = self.transforms(image=sample_to_transform['image'],\n                                                 bboxes=sample_to_transform['bboxes'],\n                                                 labels=sample_to_transform['labels'])\n            image = sample_transformed['image'].to(torch.float32) / 255.\n            boxes = torch.tensor(sample_transformed['bboxes'], dtype=torch.int64)\n            boxes[:, [2, 3]] += boxes[:, [0, 1]]  # from coco format to pascal_voc format\n            labels = torch.tensor(sample_transformed['labels'], dtype=torch.int64)\n            target = {'boxes': boxes, 'labels': labels}\n        else:\n            image = torch.from_numpy(image.transpose(2, 0, 1)\n                                     .astype('float32')) / 255.\n            boxes = torch.tensor(sample['bboxes'], dtype=torch.int64)\n            boxes[:, [2, 3]] += boxes[:, [0, 1]]\n            labels = torch.tensor(sample['labels'], dtype=torch.int64)\n            target = {'boxes': boxes, 'labels': labels}\n        return image, target\n\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n\nimage_w, image_h = 1024, 1024\n# Using A.RandomResizedCrop instance may lost all the bboxes in some sample,\n# while method FasterRCNN.forward() requires at least one bbox for each sample.\n# See base class at:\n# https://github.com/pytorch/vision/blob/v0.7.0/torchvision/models/detection/generalized_rcnn.py#L64\ntransform = A.Compose([\n    A.RandomBrightnessContrast(p=0.5),\n    A.Blur(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.HorizontalFlip(p=0.5),\n    A.RandomRotate90(p=0.5),\n    A.RandomSizedBBoxSafeCrop(image_h, image_w, erosion_rate=0.05, p=0.5),\n    ToTensorV2(),\n], bbox_params=A.BboxParams(format='coco', label_fields=['labels']))\n\ndataset_trn = WheatDatasetTrain(ds_list_trn, DIR_INPUT / 'train', transform)\nloader_trn = DataLoader(dataset_trn, batch_size=8, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"def show_augmented(transform, sample_dict, img_dir: pathlib.Path):\n    img = cv2.imread(str(img_dir / f\"{sample_dict['image_id']}.jpg\"))[..., ::-1].copy()\n    augmented = transform(image=img, bboxes=sample_dict['bboxes'], labels=sample_dict['labels'])\n    img = augmented['image'].numpy().transpose(1, 2, 0)\n    img = img[..., ::-1].astype('uint8')\n    bboxes = np.asarray(augmented['bboxes'], dtype='int64')\n    bboxes[:, [2, 3]] += bboxes[:, [0, 1]]\n    bboxes = bboxes.tolist()\n    for bbox in bboxes:\n        cv2.rectangle(img, tuple(bbox[:2]), tuple(bbox[2:]), (255, 0, 0), 2)\n    plt.figure(figsize=(6, 6))\n    plt.imshow(img)\n    plt.show()\n\n\nfor _ in range(3):\n    show_augmented(transform, ds_list_trn[0], DIR_INPUT / 'train')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Detector build: FasterRCNN with ResNet50 backbone having pre-trained on ImageNet\n\nWe use ResNet50 as backbone followed by Feature Pyramid Network(FPN),\nand Region Proposal Network(RPN) with default AnchorGenerator(scales=(32, 64, 128, 256, 512), ratios=(0.5, 1, 2))\nto produce Region of Intrests(RoI) filtered by Non Max Suppression(nms) with above 0.7 IoU threshold.\nAfter RoIAlign, the predictor predicts class score and bounding boxes.\n\nThe Predictor accepts bounding boxes in pascal_voc format ONLY.\n\nFor details, see: [FasterRCNN in torchvision](https://github.com/pytorch/vision/blob/master/torchvision/models/detection/faster_rcnn.py)","execution_count":null},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"backbone = torchvision.models.detection.backbone_utils.resnet_fpn_backbone('resnet50', pretrained=False)\n\n# Pretrained weights are better, although it's pre-trained on ImageNet.\nmissing, unexpected = backbone.body.load_state_dict(torch.load(str(DIR_WEIGHTS / FILENAME_WEIGHTS)), strict=False)\nprint(f\"Missing: {missing}\\nUnexpected in loaded state_dict: {unexpected}\")\n\nmodel = torchvision.models.detection.FasterRCNN(backbone, num_classes=2)  # Including the background\nprint(model)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train the model","execution_count":null},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"def train_n_epochs(model, optimizer, loader, device, lr_scheduler=None, n_epochs=10):\n    model = model.to(device)\n    model.train()\n    losses = []\n    for epoch in range(n_epochs):\n        loss_epoch = 0.\n        loader_len = len(loader)\n        for img_list, target in tqdm(loader):\n            img_list = [img.to(device) for img in img_list]\n            target = list(target)\n            for i in range(len(target)):\n                target[i]['boxes'] = target[i]['boxes'].to(device)\n                target[i]['labels'] = target[i]['labels'].to(device)\n            losses_batch = model(img_list, target)\n            losses_reduce = sum(l for l in losses_batch.values())\n\n            optimizer.zero_grad()\n            losses_reduce.backward()\n            optimizer.step()\n\n            loss_epoch += losses_reduce.item()\n        if lr_scheduler is not None:\n            lr_scheduler.step()\n        loss_epoch = np.round(loss_epoch / loader_len, decimals=5)\n        losses.append(loss_epoch)\n        if loss_epoch <= min(losses):  # Save best only.\n            torch.save(model.state_dict(), 'model.pt')\n            print(f'Model is serialized in {repr(\"model.pt\")}')\n        print(f\"Loss: {loss_epoch}\")\n            \n    model.eval()\n    return model\n\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nmodel = train_n_epochs(model, optimizer, loader_trn, device, lr_scheduler, 10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImgDatasetTest(Dataset):\n    \n    def __init__(self, img_dir: pathlib.Path):\n        self.img_dir = img_dir\n        self.img_names = [fname for fname in os.listdir(self.img_dir)\n                          if os.path.splitext(fname)[-1].lower() in ['.jpg', '.png']]\n        \n    def __len__(self):\n        return len(self.img_names)\n    \n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        img_fname = self.img_names[idx]\n        img_id, ext = os.path.splitext(img_fname)\n        img = cv2.imread(str(self.img_dir / img_fname))[..., ::-1].astype('float32') / 255.\n        img_tensor = torch.from_numpy(img.transpose(2, 0, 1))\n        return img_tensor, img_id\n\n\ndataset_test = ImgDatasetTest(DIR_INPUT / 'test')\nloader_test = DataLoader(dataset_test, batch_size=8)\nmodel.load_state_dict(torch.load('model.pt'))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Submission","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_result_string(score, box):\n    return f\"{score:.4f} {' '.join(str(num) for num in box)}\"\n\n\npred_threshold = 0.3\nresults = []\nwith torch.no_grad():\n    for imgs, img_ids in loader_test:\n        imgs = imgs.to(device)\n        preds = model(imgs)\n        for pred, img_id in zip(preds, img_ids):\n            sample_pred = {'image_id': img_id, 'PredictionString': ''}\n            pred['boxes'] = pred['boxes'].data.cpu().numpy()\n            pred['scores'] = pred['scores'].data.cpu().numpy()\n            boxes = pred['boxes'][pred['scores'] >= pred_threshold].astype('int64')\n            boxes[:, [2, 3]] -= boxes[:, [0, 1]]\n            scores = pred['scores'][pred['scores'] >= pred_threshold]\n            sample_pred['PredictionString'] += ' '.join(\n                format_result_string(score, box)\n                for score, box in zip(scores, boxes)\n            )\n            results.append(sample_pred)\n\nsubmission = pd.DataFrame(results)\nsubmission.to_csv('submission.csv', index=False)\nsubmission","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Plot predicted bboxes by submission","execution_count":null},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"def plot_prediction_from_submission(df: pd.DataFrame, img_dir: pathlib.Path, ext='.jpg'):\n    for idx, line in df.iterrows():\n        img = cv2.imread(str(img_dir / (line['image_id'] + ext)))\n        probas_bboxes = np.array(line['PredictionString'].split(' ')).reshape(-1, 5).astype('float64')\n        probas = probas_bboxes[:, 0]\n        bboxes = probas_bboxes[:, 1:].astype('int64')\n        bboxes[:, [2, 3]] += bboxes[:, [0, 1]]\n        for proba, bbox in zip(probas, bboxes):\n            # The lower probability predicted, the color of bbox will be more blue,\n            # otherwise it will be more red.\n            cv2.rectangle(img, tuple(bbox[:2]), tuple(bbox[2:]),\n                          (int(255 - 255 * proba), 0, int(255 * proba)), 3)\n        img = img[..., ::-1].copy()\n        plt.figure(figsize=(6, 6))\n        plt.imshow(img)\n        plt.show()\n\n\nplot_prediction_from_submission(submission, DIR_INPUT / 'test')","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}