{"cells":[{"metadata":{},"cell_type":"markdown","source":"## General information\n\nThis is an inference kernel of my training kernel: https://www.kaggle.com/artgor/object-detection-with-pytorch-lightning/","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"from __future__ import print_function\nfrom PIL import Image\nfrom albumentations.core.composition import Compose\nfrom collections import defaultdict, deque\nfrom itertools import product\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import Dataset\nfrom torchvision.transforms import functional as F\nfrom typing import Any, Dict, List\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport ast\nimport copy\nimport cv2\nimport datetime\nimport errno\nimport importlib\nimport json\nimport numpy as np\nimport os\nimport pandas as pd\nimport pickle\nimport random\nimport shutil\nimport tempfile\nimport time\nimport torch\nimport torch._six\nimport torch.distributed as dist\nimport torch.utils.data\nimport torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nimport matplotlib.pyplot as plt\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Helper code","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def set_seed(seed: int = 666):\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.backends.cudnn.benchmark = False\n    torch.backends.cudnn.deterministic = True\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    \ndef get_test_dataset():\n    \"\"\"\n    Get test dataset\n\n    Args:\n        cfg:\n\n    Returns:\n\n    \"\"\"\n\n    test_img_dir = f'/kaggle/input/global-wheat-detection/test'\n\n    valid_augs = A.Compose([ToTensorV2(always_apply=True, p=1.0)],\n                           p=1.0,\n                           bbox_params={'format': 'pascal_voc', 'label_fields': ['labels'], 'min_area': 0.0, 'min_visibility': 0.0},\n                           keypoint_params=None,\n                           additional_targets={})\n\n    test_dataset = WheatDataset(None,\n                                 'test',\n                                 test_img_dir,\n                                 valid_augs)\n\n    return test_dataset\n\nclass WheatDataset(Dataset):\n\n    def __init__(self,\n                 dataframe: pd.DataFrame = None,\n                 mode: str = 'train',\n                 image_dir: str = '',\n                 transforms: Compose = None):\n        \"\"\"\n        Prepare data for wheat competition.\n\n        Args:\n            dataframe: dataframe with image id and bboxes\n            mode: train/val/test\n            image_dir: path to images\n            transforms: albumentations\n        \"\"\"\n        self.image_dir = image_dir\n        self.df = dataframe\n        self.mode = mode\n        self.image_ids = os.listdir(self.image_dir) if self.df is None else self.df['image_id'].unique()\n        self.transforms = transforms\n\n    def __getitem__(self, idx: int):\n        image_id = self.image_ids[idx].split('.')[0]\n        # print(image_id)\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\n        # normalization.\n        # TO DO: refactor preprocessing\n        image /= 255.0\n\n        # test dataset must have some values so that transforms work.\n        target = {'labels': torch.as_tensor([[0]], dtype=torch.float32),\n                  'boxes': torch.as_tensor([[0, 0, 0, 0]], dtype=torch.float32)}\n\n        # for train and valid test create target dict.\n        if self.mode != 'test':\n            image_data = self.df.loc[self.df['image_id'] == image_id]\n            boxes = image_data[['x', 'y', 'x1', 'y1']].values\n\n            areas = image_data['area'].values\n            areas = torch.as_tensor(areas, dtype=torch.float32)\n\n            # there is only one class\n            labels = torch.ones((image_data.shape[0],), dtype=torch.int64)\n            iscrowd = torch.zeros((image_data.shape[0],), dtype=torch.int64)\n\n            target['boxes'] = boxes\n            target['labels'] = labels\n            target['image_id'] = torch.tensor([idx])\n            target['area'] = areas\n            target['iscrowd'] = iscrowd\n\n            if self.transforms:\n                image_dict = {\n                    'image': image,\n                    'bboxes': target['boxes'],\n                    'labels': labels\n                }\n                image_dict = self.transforms(**image_dict)\n                image = image_dict['image']\n                target['boxes'] = torch.as_tensor(image_dict['bboxes'], dtype=torch.float32)\n\n        else:\n            image_dict = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': target['labels']\n            }\n            image = self.transforms(**image_dict)['image']\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return len(self.image_ids)\n    \n    \ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n\ndef format_prediction_string(boxes, scores):\n    pred_strings = []\n    for s, b in zip(scores, boxes.astype(int)):\n        pred_strings.append(f'{s:.4f} {b[0]} {b[1]} {b[2] - b[0]} {b[3] - b[1]}')\n\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"set_seed(42)\n\ntest_dataset = get_test_dataset()\n\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)\ndevice = 'cuda'\nnum_classes = 2\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n# Load the trained weights\nmodel.load_state_dict(torch.load('/kaggle/input/object-detection-with-pytorch-lightning/fasterrcnn_resnet50_fpn.pth'))\nmodel.cuda()\nmodel.eval()\n\ntest_loader = torch.utils.data.DataLoader(test_dataset,\n                                          batch_size=4,\n                                          num_workers=2,\n                                          shuffle=False,\n                                          collate_fn=collate_fn)\ndetection_threshold = 0.5\nresults = []\n\nfor images, _, image_ids in test_loader:\n\n    images = list(image.to(device) for image in images)\n    outputs = model(images)\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\n        boxes = boxes[scores >= detection_threshold].astype(np.int32)\n        scores = scores[scores >= detection_threshold]\n        image_id = image_ids[i]\n\n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores)\n        }\n\n        results.append(result)","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":{"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}