{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"/kaggle/input/weightedboxesfusion\")\n\nimport ensemble_boxes\n\nimport pandas as pd\nimport numpy as np\nimport cv2\nimport os, re\n\nimport torch\nimport torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom torch.utils.data import DataLoader, Dataset\n\nfrom matplotlib import pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"DATA_DIR = \"/kaggle/input/global-wheat-detection\"\nMODELS_IN_DIR = \"/kaggle/input/frcnnfoldthree\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.read_csv(os.path.join(DATA_DIR, \"sample_submission.csv\"))\ntest_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n    def __init__(self, dataframe, image_dir, 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\n    def __len__(self) -> int:\n        return len(self.image_ids)\n\n    def __getitem__(self, idx: int):\n        image_id = self.image_ids[idx]\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#         # change the shape from [h,w,c] to [c,h,w]  \n#         image = torch.from_numpy(image).permute(2,0,1)\n\n        records = self.df[self.df['image_id'] == image_id]\n    \n        if self.transforms:\n            sample = {\"image\": image}\n            sample = self.transforms(**sample)\n            image = sample['image']\n\n        return image, image_id\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model():\n    \"\"\"\n    https://stackoverflow.com/questions/58362892/resnet-18-as-backbone-in-faster-r-cnn\n    \"\"\"\n    backbone = resnet_fpn_backbone('resnet101', pretrained=False)\n    model = FasterRCNN(backbone, num_classes=2)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nmodel = get_model()\nmodel.load_state_dict(torch.load(os.path.join(MODELS_IN_DIR, \"best_model.pth\")))\nmodel.eval()\n\nmodel.to(DEVICE)\n1 == 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_test_transforms():\n    return A.Compose([\n            A.Resize(height=1024, width=1024, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntest_dataset = WheatDataset(test_df, os.path.join(DATA_DIR, \"test\"), get_test_transforms())\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=1,\n    drop_last=False,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BaseWheatTTA:\n    \"\"\" author: @shonenkov \"\"\"\n    image_size = 1024\n\n    def augment(self, image):\n        raise NotImplementedError\n    \n    def batch_augment(self, images):\n        raise NotImplementedError\n    \n    def deaugment_boxes(self, boxes):\n        raise NotImplementedError\n\nclass TTAHorizontalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n\n    def augment(self, image):\n        return image.flip(1)\n    \n    def batch_augment(self, images):\n        return images.flip(2)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [1,3]] = self.image_size - boxes[:, [3,1]]\n        return boxes\n\nclass TTAVerticalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return image.flip(2)\n    \n    def batch_augment(self, images):\n        return images.flip(3)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,2]] = self.image_size - boxes[:, [2,0]]\n        return boxes\n    \nclass TTARotate90(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 1, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 1, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = self.image_size - boxes[:, [3,1]] \n        res_boxes[:, [1,3]] = boxes[:, [0,2]]\n        return res_boxes\n\nclass TTACompose(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    def __init__(self, transforms):\n        self.transforms = transforms\n        \n    def augment(self, image):\n        for transform in self.transforms:\n            image = transform.augment(image)\n        return image\n    \n    def batch_augment(self, images):\n        for transform in self.transforms:\n            images = transform.batch_augment(images)\n        return images\n    \n    def prepare_boxes(self, boxes):\n        result_boxes = boxes.copy()\n        result_boxes[:,0] = np.min(boxes[:, [0,2]], axis=1)\n        result_boxes[:,2] = np.max(boxes[:, [0,2]], axis=1)\n        result_boxes[:,1] = np.min(boxes[:, [1,3]], axis=1)\n        result_boxes[:,3] = np.max(boxes[:, [1,3]], axis=1)\n        return result_boxes\n    \n    def deaugment_boxes(self, boxes):\n        for transform in self.transforms[::-1]:\n            boxes = transform.deaugment_boxes(boxes)\n        return self.prepare_boxes(boxes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from itertools import product\n\ntta_transforms = []\nfor tta_combination in product([TTAHorizontalFlip(), None], \n                               [TTAVerticalFlip(), None],\n                               [TTARotate90(), None]):\n    tta_transforms.append(TTACompose([tta_transform for tta_transform in tta_combination if tta_transform]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_tta_predictions(images, score_threshold=0.5):\n    with torch.no_grad():\n        images = torch.stack(images).float().to(DEVICE)\n        predictions = []\n        for tta_transform in tta_transforms:\n            result = []\n            outputs = model(tta_transform.batch_augment(images.clone()))\n\n            for i, image in enumerate(images):\n                boxes = outputs[i]['boxes'].data.cpu().numpy()   \n                scores = outputs[i]['scores'].data.cpu().numpy()\n                indexes = np.where(scores > score_threshold)[0]\n                boxes = boxes[indexes]\n                boxes = tta_transform.deaugment_boxes(boxes.copy())\n                result.append({\n                    'boxes': boxes,\n                    'scores': scores[indexes],\n                })\n            predictions.append(result)\n    return predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run_wbf(predictions, image_index, image_size=1024, iou_thr=0.5, skip_box_thr=0.43, weights=None):\n    boxes = [(prediction[image_index]['boxes']/(image_size-1)).tolist() for prediction in predictions]\n    scores = [prediction[image_index]['scores'].tolist() for prediction in predictions]\n    labels = [np.ones(prediction[image_index]['scores'].shape[0]).astype(int).tolist() for prediction in predictions]\n    boxes, scores, labels = ensemble_boxes.ensemble_boxes_wbf.weighted_boxes_fusion(boxes, scores, labels, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n    boxes = boxes*(image_size-1)\n    return boxes, scores, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results = []\n\nfor images, image_ids in test_data_loader:\n\n    predictions = make_tta_predictions(images)\n\n    for i, image in enumerate(images):\n\n        boxes, scores, labels = run_wbf(predictions, image_index=i)\n        boxes = boxes.round().astype(np.int32).clip(min=0, max=1023)\n        image_id = image_ids[i]\n        \n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n        \n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores)\n        }\n\n        \n        results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results[0:2]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)","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}