{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\nimport sys\nsys.path.insert(0, \"../input/weightedboxesfusion\")\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport torch\nfrom torchvision import transforms\nfrom matplotlib import pyplot as plt\nfrom torchvision import transforms\nimport torchvision \nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torch.utils.data import Dataset,DataLoader\nimport glob\nimport albumentations as A\nimport cv2\nfrom albumentations.pytorch.transforms import ToTensorV2\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)\nfrom ensemble_boxes import *\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n# import os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 5GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_resnet101_model(checkpoint_path):\n    num_classes = 2  # 1 class (wheat) + background\n    backbone = resnet_fpn_backbone('resnet101', pretrained= False)\n    model_faster = FasterRCNN(backbone, num_classes)\n    # replace the pre-trained head with a new one\n    in_features = model_faster.roi_heads.box_predictor.cls_score.in_features\n    model_faster.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n    checkpoint = torch.load(checkpoint_path)\n    model_faster.load_state_dict(checkpoint['model_state_dict'])\n    model_faster.to(device);\n    model_faster.eval()\n    return model_faster\n\ndef load_resnet50_model(checkpoint_path):\n    num_classes = 2  # 1 class (wheat) + background\n    backbone = resnet_fpn_backbone('resnet50', pretrained= False)\n    model_faster = FasterRCNN(backbone, num_classes)\n    # replace the pre-trained head with a new one\n    in_features = model_faster.roi_heads.box_predictor.cls_score.in_features\n    model_faster.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n    model_faster.to(device);\n    checkpoint = torch.load(checkpoint_path)\n    model_faster.load_state_dict(checkpoint['model_state_dict'])\n    model_faster.to(device);\n    model_faster.eval()\n    return model_faster\n\n\ndef load_resnet152_model(checkpoint_path):\n    num_classes = 2  # 1 class (wheat) + background\n    backbone = resnet_fpn_backbone('resnet152', pretrained= False)\n    model_faster = FasterRCNN(backbone, num_classes)\n    # replace the pre-trained head with a new one\n    in_features = model_faster.roi_heads.box_predictor.cls_score.in_features\n    model_faster.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n    model_faster.to(device);\n    checkpoint = torch.load(checkpoint_path)\n    model_faster.load_state_dict(checkpoint['model_state_dict'])\n    model_faster.to(device);\n    model_faster.eval()\n    return model_faster\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"resnet101_33e = \"../input/resnet10133e-1024/cp33.pt\"\nresnet50 = \"../input/resnet50-40e/resnet50_40e.pt\"\nresnet101_25e = \"../input/resnet101mymodel/cp25.pt\"\n# resnet101_50e = \"../input/resnet101-50e/cp49.pt\"\nresnet152_20e = \"../input/resnet152-20e/cp20.pt\"\nresnet152_19e =\"../input/resnet152-19e/cp19 (1).pt\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"models = [\n    load_resnet50_model(resnet50),\n#     load_resnet101_model(resnet101_33e),\n#     load_resnet101_model(resnet101_25e),\n#     load_resnet101_model(resnet101_50e),\n    load_resnet152_model(resnet152_20e),\n    load_resnet152_model(resnet152_19e)\n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_ROOT_PATH = '../input/global-wheat-detection/test/'\n\nclass 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]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_valid_transforms():\n    return A.Compose([\n            ToTensorV2(p=1.0),\n        ], p=1.0)","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":"dataset = DatasetRetriever(np.array([path.split('/')[-1][:-4] for path in glob.glob(f'{DATA_ROOT_PATH}/*.jpg')]),get_valid_transforms())\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndata_loader = DataLoader(\n    dataset,\n    batch_size=1,\n    shuffle=False,\n    num_workers=2,\n    drop_last=False,\n    collate_fn=collate_fn\n)","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[:, [1,3]]\n        res_boxes[:, [1,3]] = boxes[:, [2,0]]\n        return res_boxes\n    \nclass TTARotate180(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 2, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 2, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,1,2,3]] = self.image_size - boxes[:, [2,3,0,1]]\n        return boxes\n    \nclass TTARotate270(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 3, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 3, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = boxes[:, [1,3]]\n        res_boxes[:, [1,3]] = self.image_size - 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":"def make_tta_predictions(model,images, score_threshold=0.35):\n  model.eval()\n  model.to(device)\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 = model(tta_transform.batch_augment(images.clone()))\n\n          for i in range(images.shape[0]):\n              boxes = det[i]['boxes'].detach().cpu().numpy()[:,:4]    \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\ndef make_ensemble_predictions(images):\n    images = list(image.to(device) for image in images)    \n    result = []\n    for model in models:\n        predictions = make_tta_predictions(model, images)\n        with torch.no_grad():\n            for i in range(len(images)):\n                boxes, scores, labels = run_wbf(predictions, image_index=i)\n                outputs = {'boxes': boxes,'labels': labels, 'scores':scores }\n                result.append([outputs])\n                torch.cuda.empty_cache()\n\n    return result\n\n\n# def make_ensemble_predictions(images):\n#     images = list(image.to(device) for image in images)    \n#     result = []\n#     for model in models:\n#         with torch.no_grad():\n#             outputs = model(images)\n#             result.append(outputs)\n#             torch.cuda.empty_cache()\n#     return result\n\n\ndef run_wbf(predictions, image_index, image_size=1024, iou_thr=0.45, skip_box_thr=0.35,score_threshold = 0.35, weights=None,conf_type='max'):\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=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr, conf_type=conf_type)\n    indexes = np.where(scores > score_threshold)[0]\n    boxes = boxes[indexes]    \n    scores = scores[indexes]\n    boxes = boxes*(image_size-1)\n    return boxes, scores, labels","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(), TTARotate180(), TTARotate270(), None]):\n    tta_transforms.append(TTACompose([tta_transform for tta_transform in tta_combination if tta_transform]))\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nfor j, (images, image_ids) in enumerate(data_loader):\n    break\n\npredictions = make_ensemble_predictions(images)\n\ni = 0\nsample = images[i].permute(1,2,0).cpu().numpy()\nboxes, scores, labels = run_wbf(predictions, image_index=i, iou_thr=0.45, skip_box_thr=0.35, score_threshold = 0.35, weights=None,conf_type='max')\n# print(len(boxes),len(scores))\nboxes = boxes.round().astype(np.int32).clip(min=0, max=1023)\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor j,box in enumerate(boxes):\n    cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (0, 0, 0), 3)\n    score=str(float(\"{:.3f}\".format(scores[j])))\n    \n#     cv2.putText(\n#               sample,\n#               score,\n#               org=(int(box[0]), int(box[1] )), # bottom left\n#               fontFace=cv2.FONT_HERSHEY_PLAIN,\n#               fontScale=3,\n#               color=(1,1, 1),\n#               thickness=2\n#             )\nax.set_axis_off()\nax.imshow(sample);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n    \nvalidation_image_precisions = []\niou_thresholds = [x for x in np.arange(0.5, 0.76, 0.05)]\n# model.eval()\n# model.to(device)\n\nresults = []\nts = 0.147\n\nfor images, image_ids in data_loader:    \n    predictions = make_ensemble_predictions(images)\n    for i, image in enumerate(images):\n#         boxes, scores, labels = run_wbf(predictions, image_index=i, iou_thr=0.35, skip_box_thr=0.35, score_threshold = 0.147, weights=None,conf_type='avg')# res50, res152(20) \n        boxes, scores, labels = run_wbf(predictions, image_index=i, iou_thr=0.35, skip_box_thr=0.35, score_threshold = 0.35, weights=None,conf_type='avg')\n        boxes = boxes.round().astype(np.int32).clip(min=0, max=1023)\n        image_id = image_ids[i]\n\n        \n        \n#         plot the images\n#         sample = images[i].permute(1,2,0).cpu().numpy()\n#         fig, ax = plt.subplots(1, 1, figsize=(15,15))\n#         indexes = np.where(scores > ts)[0]\n\n#         for j,box in enumerate(boxes):\n#             if j not  in indexes:\n#                 c = (1,0,0)\n#             else:\n#                 c = (0,0,0)\n                \n#             cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), c, 3)\n#             score=str(float(\"{:.3f}\".format(scores[j])))\n\n#             cv2.putText(\n#               sample,\n#               score,\n#               org=(int(box[0]), int(box[1] )), # bottom left\n#               fontFace=cv2.FONT_HERSHEY_PLAIN,\n#               fontScale=3,\n#               color=(1,1, 1),\n#               thickness=2\n#             )\n#         ax.set_axis_off()\n#         ax.imshow(sample)\n        \n\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        results.append(result)\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head(10)","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}