{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport os\nimport re\n\nfrom PIL import Image\nfrom PIL import ImageDraw\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\n\nfrom torch.utils.data import DataLoader, Dataset\n\nfrom matplotlib import pyplot as plt\n\nweights_dir = '../input/wheat-head-detection'\n\nsaved_model_path = os.path.join(weights_dir,'fasterrcnn_resnet50_fpn.pth')\nsaved_model_path","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Creating the test dataset\nfrom torchvision.transforms import functional as F\n\nclass WheatDatasetTest(Dataset):\n    def __init__(self, data_dir, dataframe, transforms=True):\n        super().__init__()\n        self.data = data_dir\n        self.df = dataframe\n        self.transforms = transforms\n        # load all the images and sort them\n        self.images = sorted(os.listdir(os.path.join(data_dir)))\n        \n    def __getitem__(self, idx:int):\n        # getting the image id for the given image\n        image_id = self.images[idx].split('.')[0]\n        \n        # get the bounding box info for the given image\n        boxes_info = (self.df[self.df['image_id'] == image_id])\n\n        # load the image\n        path = os.path.join(self.data, image_id) + '.jpg'\n        image = Image.open(path).convert(\"RGB\")\n        \n        if self.transforms:\n            image= F.to_tensor(image)   # normalizes and returns `torch.tensor`\n            \n        return image, image_id\n    \n    def __len__(self):\n        return len(self.images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define the root directory\nroot_dir = '../input/global-wheat-detection'\n\n# create test dataframe\ntest_df = pd.read_csv(os.path.join(root_dir, 'sample_submission.csv'))\ntest_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# location of test images\ntest_data_dir = os.path.join(root_dir, 'test')\ntest_data_dir","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# merge all the images in a batch\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# Use the dataset class and create train and test dataloaders\ndataset_test = WheatDatasetTest(data_dir=test_data_dir, dataframe=test_df)\n\n# creating dataloader\ndataloader_test = DataLoader(dataset_test, batch_size=10, shuffle=False, \n                                   num_workers=4, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model(num_classes):\n    # load the object detection model pre-trained on COCO\n    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, \n                                                                 pretrained_backbone=False)\n    \n    # get the input features in the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    \n    # replace the input features of pretrained head with the num_classes\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if torch.cuda.is_available():\n    print(torch.cuda.get_device_name())\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# get the model and load the saved model\n# our dataset has only two classes: wheat heads and background\nnum_classes = 2\n\n# get the model using helper function\nmodel = get_model(num_classes)\n\n# Load the trained weights\nmodel.load_state_dict(torch.load(saved_model_path))\nmodel.eval()\n\n# move the model to device\nmodel.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    # source: https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train\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":"detection_threshold = 0.45\nresults = []\n\nfor images, image_ids in dataloader_test:\n\n    images = list(image.to(device) for image in images)\n    outputs = model(images)\n\n    for i, image in enumerate(images):\n\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        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":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Showing inference on model output\ndef show_inferences(images, targets, num_images):\n    ncols = 2\n    nrows = min(num_images//ncols, 2)\n    \n    fig, ax = plt.subplots(nrows, ncols, figsize=(20, 14))\n    \n    ax = ax.flatten()  \n    for idx in range(num_images):\n        image = Image.fromarray(images[idx].mul(255).permute(1,2,0).cpu().detach().byte().numpy())\n        boxes = targets[idx]['boxes'].cpu().detach().numpy().astype(np.int64)\n\n        draw = ImageDraw.Draw(image)\n\n        for box in boxes:\n            coord1 = (box[0], box[1])\n            coord2 = (box[2], box[3])\n            draw.rectangle([coord1, coord2], outline=(220,20,60), width=10)\n        ax[idx].set_axis_off()\n        ax[idx].imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_inferences(images, outputs, 4)","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}