{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0,'../input/weightedboxesfusion')\n\nimport cv2\nimport gc\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport torch\nimport torchvision\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm.notebook import tqdm\nfrom matplotlib import pyplot as plt\nfrom glob import glob\nfrom ensemble_boxes import *\nfrom itertools import product\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor,FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\n%matplotlib inline\nDIR = '../input/global-wheat-detection/'\nDATA_ROOT_PATH = DIR + \"test\"\n\nMODELS_DIR_PATH = '../input/fasterrcnn-resnet50-fpn-best' # //\n\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"from albumentations import (Compose, Resize,Normalize )\nfrom albumentations.pytorch.transforms import ToTensorV2\n\ndef get_test_transform():\n    return Compose([\n        Resize(height=512, width=512, p=1.0),\n        ToTensorV2(p=1.0)  \n    ])\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef load_model(path):\n    model_type = path.split('/')[-1].split('.')[0].split('_')[0]\n    backbone = resnet_fpn_backbone(model_type, pretrained=False)\n\n    model = FasterRCNN(backbone, 2)                    \n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)\n\n    model_dict = torch.load(path)\n    model.load_state_dict(model_dict)\n\n    del model_dict\n    gc.collect()\n    \n    model.to(device)\n    model.eval()\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"models = [load_model(f'{MODELS_DIR_PATH}/resnet50_0.pth'),\n          load_model(f'{MODELS_DIR_PATH}/resnet50_1.pth'),\n          load_model(f'{MODELS_DIR_PATH}/resnet50_2.pth'),\n          load_model(f'{MODELS_DIR_PATH}/resnet101_3.pth'),\n          load_model(f'{MODELS_DIR_PATH}/resnet152_4.pth')]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DatasetRetriever(Dataset):\n\n    def __init__(self, image_ids, transforms=None):\n        super().__init__()\n        self.image_ids = image_ids\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{DATA_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        if self.transforms:\n            sample = {'image': image}\n            sample = self.transforms(**sample)\n            image = sample['image']\n        return image, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = DatasetRetriever(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{DATA_ROOT_PATH}/*.jpg')]),\n    transforms=get_test_transform()\n)\ndata_loader = DataLoader(\n    dataset,\n    batch_size=1,\n    shuffle=False,\n    num_workers=4,\n    drop_last=False,\n    collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BaseWheatTTA:\n    \"\"\" author: @shonenkov \"\"\"\n    image_size = 512\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[:, [1,3]]\n        res_boxes[:, [1,3]] = boxes[:, [2,0]]\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":"\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]))\n\ndef make_tta_predictions(net, images, score_threshold=0.4):\n    with torch.no_grad():\n        images = torch.stack(images).float().cuda()\n        predictions = []\n        for tta_transform in tta_transforms:\n            result = []\n            det = net(tta_transform.batch_augment(images.clone()))\n\n            for i in range(images.shape[0]):\n                boxes = det[i]['boxes'].detach().cpu().numpy()\n                scores = det[i]['scores'].detach().cpu().numpy()\n                indexes = np.where(scores > score_threshold)[0]\n                boxes = boxes[indexes]\n                # boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n                # boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\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\n\n\n\n\ndef make_ensemble_predictions(images):\n    images = list(image.to(device) for image in images)    \n    result = []\n    for net in models:\n        outputs = net(images)\n        result.append(outputs)\n    return result\n\ndef run_wbf(predictions, image_index, image_size=512, iou_thr=0.55, skip_box_thr=0.7, weights=None,conf_type='avg'):\n    boxes = [prediction[image_index]['boxes']/(image_size) for prediction in predictions]\n    scores = [prediction[image_index]['scores'] for prediction in predictions]\n    labels = [np.ones(prediction[image_index]['scores'].shape[0]) for prediction in predictions]\n    boxes, scores, labels = weighted_boxes_fusion(boxes, scores, labels, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr,conf_type=conf_type)\n    boxes = boxes*(image_size)\n    return boxes, scores, labels\n\ndef TTA_make_ensemble_predictions(images):\n    images = list(image.to(device) for image in images)    \n    result = []\n    for net in models:\n        predictions = make_tta_predictions(net, images)\n        for i in range(len(images)):\n            boxes, scores, labels = run_wbf(predictions, image_index=i, iou_thr=0.5, skip_box_thr=0.7)\n            outputs = {'boxes':boxes,\n                       'labels':labels,\n                       'scores':scores}\n            result.append([outputs])\n    return result\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Submission","execution_count":null},{"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    return \" \".join(pred_strings)\n\ndef visualize_results(sample, boxes, scores, score_thresh=0.33,  show_error=False):\n        sample = cv2.resize(sample,(1024,1024))\n        indexes = np.where(scores > score_thresh)[0]\n        if not show_error:\n            boxes = boxes[indexes]\n\n        fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n        colors = [(0, 0, 0),(1, 0, 0)]\n        for index, (box, score) in enumerate(zip(boxes,scores)):\n            \n            if index not in indexes:\n                c = colors[1]\n            else:\n                c = colors[0]\n            cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]),  c, 2)\n            cv2.putText(sample, f'{score:.3f}', (box[0]+20, box[1]+20),\n                        cv2.FONT_HERSHEY_COMPLEX,  \n                        1, (1, 1, 1), 2, cv2.LINE_AA) \n            \n        # ax.set_axis_off()\n        ax.imshow(sample);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nTHRESH = 0.35\nresults = []\nfor j, (images, image_ids) in tqdm(enumerate(data_loader), total=len(data_loader), desc='Prediction:'):\n    predictions = TTA_make_ensemble_predictions(images)\n    for i in range(len(images)):\n        boxes, scores, labels = run_wbf(predictions, image_index=i, iou_thr=0.5, skip_box_thr=0.7, conf_type='max')\n        boxes = (boxes*2).astype(np.int32).clip(min=0, max=1023)\n\n#         visualize_results(images[i].permute(1,2,0).cpu().numpy(), boxes, scores, THRESH, show_error=True)\n\n        indexes = np.where(scores > THRESH)[0]\n        boxes = boxes[indexes]\n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n\n        result = {\n            'image_id': image_ids[i],\n            'PredictionString': format_prediction_string(boxes, scores)\n        }\n        results.append(result)\ntest_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head()\n","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}