{"cells":[{"metadata":{},"cell_type":"markdown","source":"# TorchVision Faster R-CNN Inference"},{"metadata":{},"cell_type":"markdown","source":"## Prepare Data\n\n- [Global Wheat Detection](https://www.kaggle.com/c/global-wheat-detection)\n- [TorchVision Faster R-CNN Finetuning](https://www.kaggle.com/gocoding/torchvision-faster-r-cnn-finetuning)"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\n\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nfrom matplotlib import pyplot as plt\nfrom PIL import Image\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\nDIR_WEIGHTS = '/kaggle/input/torchvision-faster-r-cnn-finetuning'\nWEIGHTS_FILE = f'{DIR_WEIGHTS}/fasterrcnn_resnet50_fpn.pth'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Submission File\n\nThe submission format requires a space delimited set of bounding boxes. For example:\n\n`ce4833752,0.5 0 0 100 100`\n\nindicates that image `ce4833752` has a bounding box with a confidence of `0.5`, at `x` == 0 and `y` == 0, with a `width` and `height` of 100."},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.read_csv(f'{DIR_INPUT}/sample_submission.csv')\ntest_df.tail()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"import cv2 as cv\nimport numpy as np\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nfrom torch.utils.data import Dataset\n\n\nclass WheatTest(Dataset):\n\n  def __init__(self, image_ids, image_dir, transforms=None):\n    super().__init__()\n    self.image_ids = image_ids\n    self.image_dir = image_dir\n    self.transforms = transforms\n\n  def __getitem__(self, idx: int):\n    image_id = self.image_ids[idx]\n\n    image = cv.imread(f'{self.image_dir}/{image_id}.jpg', cv.IMREAD_COLOR)\n    image = cv.cvtColor(image, cv.COLOR_BGR2RGB).astype(np.float32)\n    image /= 255.0\n\n    if self.transforms:\n      sample = {\n        'image': image,\n      }\n      sample = self.transforms(**sample)\n      image = sample['image']\n\n    return image, image_id\n\n  def __len__(self) -> int:\n    return len(self.image_ids)\n\n  @staticmethod\n  def get_test_transform():\n    return A.Compose([\n      ToTensorV2(p=1.0)\n    ])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create DataLoader"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_image_ids(p):\n  import glob\n  image_ids = []\n  for p in glob.glob(f'{p}/*.jpg'):\n    n, _ = os.path.splitext(os.path.basename(p))\n    image_ids.append(n)\n  return image_ids\n\n# try more images for submission\n#test_dataset = WheatTest(get_image_ids(DIR_TRAIN), DIR_TRAIN, WheatTest.get_test_transform())\n\ntest_dataset = WheatTest(test_df[\"image_id\"].unique(), DIR_TEST, WheatTest.get_test_transform())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import DataLoader\n\ndef collate_fn(batch):\n  return tuple(zip(*batch))\n\ntest_data_loader = DataLoader(\n  test_dataset,\n  batch_size=8, # GPU not enough if inference 16 images\n  shuffle=False,\n  num_workers=4,\n  drop_last=False,\n  collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n# create a Faster R-CNN model without pre-trained\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)\n\nnum_classes = 2 # wheat or not(background)\n\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# replace the pre-trained model's head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n# load the trained weights\nmodel.load_state_dict(torch.load(WEIGHTS_FILE, map_location=device))\nmodel.eval()\n\n# move model to the right device\n_ = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Let's Inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"score_threshold = 0.7\nimage_outputs = []\n\nfor images, image_ids in test_data_loader:\n  images = list(image.to(device) for image in images)\n  outputs = model(images)\n    \n  for image_id, output in zip(image_ids, outputs):\n    boxes = output['boxes'].data.cpu().numpy()\n    scores = output['scores'].data.cpu().numpy()\n    \n    mask = scores >= score_threshold\n    boxes = boxes[mask].astype(np.int32)\n    scores = scores[mask]\n\n    image_outputs.append((image_id, boxes, scores))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Show Outputs"},{"metadata":{"trusted":true},"cell_type":"code","source":"num_rows, num_cols = 1, 2\nscale = 16\nfigsize = (num_rows * scale, num_cols * scale)\n_, axes = plt.subplots(num_rows, num_cols, figsize=figsize)\naxes = axes.flatten()\n\nfig_n = num_rows * num_cols\nfig_i = 0\n\nfor i in range(1, 1+fig_n):\n  image, image_id = test_dataset[i]\n  _, boxes, scores = image_outputs[i]\n\n  sample = image.permute(1, 2, 0).cpu().numpy()\n\n  for box in boxes:\n    cv.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (220, 0, 0), 2)\n\n  axes[fig_i].imshow(sample)\n  fig_i += 1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Save Results"},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n  pred_strings = []\n  for score, box in zip(scores, boxes):\n    pred_strings.append(round(score, 4))\n    pred_strings.extend(box)\n  return ' '.join(map(str, pred_strings))\n\nresults = []\n\nfor image_id, boxes, scores in image_outputs:\n  #boxes = boxes_.copy()\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\nresults[0]","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":"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}