{"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_minor":4,"nbformat":4,"cells":[{"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":{}},{"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'","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:52.644852Z","iopub.execute_input":"2022-07-24T14:11:52.645296Z","iopub.status.idle":"2022-07-24T14:11:53.504044Z","shell.execute_reply.started":"2022-07-24T14:11:52.645254Z","shell.execute_reply":"2022-07-24T14:11:53.50318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:53.507833Z","iopub.execute_input":"2022-07-24T14:11:53.508128Z","iopub.status.idle":"2022-07-24T14:11:53.762465Z","shell.execute_reply.started":"2022-07-24T14:11:53.508102Z","shell.execute_reply":"2022-07-24T14:11:53.761643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- image_id - the unique image ID\n- width, height - the width and height of the images\n- bbox - a bounding box, formatted as a Python-style list of [xmin, ymin, width, height]\n- etc.","metadata":{}},{"cell_type":"markdown","source":"Split`bbox` to `x` `y` `w` `h`:","metadata":{}},{"cell_type":"code","source":"train_df[['x','y','w','h']] = 0\ntrain_df[['x','y','w','h']] = np.stack(train_df['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=','))).astype(np.float)\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:53.763644Z","iopub.execute_input":"2022-07-24T14:11:53.763959Z","iopub.status.idle":"2022-07-24T14:11:54.592472Z","shell.execute_reply.started":"2022-07-24T14:11:53.763924Z","shell.execute_reply":"2022-07-24T14:11:54.591568Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Analysis Dataset\n\nThe size of train data:","metadata":{}},{"cell_type":"code","source":"train_df.shape","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:54.593782Z","iopub.execute_input":"2022-07-24T14:11:54.594115Z","iopub.status.idle":"2022-07-24T14:11:54.603239Z","shell.execute_reply.started":"2022-07-24T14:11:54.59408Z","shell.execute_reply":"2022-07-24T14:11:54.602307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The number of unique `image_id` in train data:","metadata":{}},{"cell_type":"code","source":"train_df['image_id'].nunique()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:54.607023Z","iopub.execute_input":"2022-07-24T14:11:54.607273Z","iopub.status.idle":"2022-07-24T14:11:54.626084Z","shell.execute_reply.started":"2022-07-24T14:11:54.607249Z","shell.execute_reply":"2022-07-24T14:11:54.625274Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(os.listdir(DIR_TRAIN))","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:54.629251Z","iopub.execute_input":"2022-07-24T14:11:54.629605Z","iopub.status.idle":"2022-07-24T14:11:54.745357Z","shell.execute_reply.started":"2022-07-24T14:11:54.629571Z","shell.execute_reply":"2022-07-24T14:11:54.744454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"There are `3422-3373=49` images which do not have any wheat heads in it (without annotation). ","metadata":{}},{"cell_type":"markdown","source":"Distribution of size of train images:","metadata":{}},{"cell_type":"code","source":"train_df['width'].unique(), train_df['height'].unique()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:54.746804Z","iopub.execute_input":"2022-07-24T14:11:54.747141Z","iopub.status.idle":"2022-07-24T14:11:54.759278Z","shell.execute_reply.started":"2022-07-24T14:11:54.747107Z","shell.execute_reply":"2022-07-24T14:11:54.757776Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Distribution of number of bounding boxes:","metadata":{}},{"cell_type":"code","source":"counts = train_df['image_id'].value_counts()\nprint(f'number of boxes, range [{min(counts)}, {max(counts)}]')\nsns.displot(counts, kde=False)\nplt.xlabel('boxes')\nplt.ylabel('images')\nplt.title('boxes vs. images')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:54.760339Z","iopub.execute_input":"2022-07-24T14:11:54.760758Z","iopub.status.idle":"2022-07-24T14:11:55.081069Z","shell.execute_reply.started":"2022-07-24T14:11:54.760718Z","shell.execute_reply":"2022-07-24T14:11:55.080082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Distribution of size of bounding boxes:","metadata":{}},{"cell_type":"code","source":"train_df['cx'] = train_df['x'] + train_df['w'] / 2\ntrain_df['cy'] = train_df['y'] + train_df['h'] / 2\n\nax = plt.subplots(1, 4, figsize=(16, 4), tight_layout=True)[1].ravel()\nax[0].set_title('x vs. y')\nax[0].set_xlim(0, 1024)\nax[0].set_ylim(0, 1024)\nax[1].set_title('cx vs. cy')\nax[1].set_xlim(0, 1024)\nax[1].set_ylim(0, 1024)\nax[2].set_title('w vs. h')\nax[3].set_title('area size')\nsns.histplot(data=train_df, x='x', y='y', ax=ax[0], bins=50, pmax=0.9)\nsns.histplot(data=train_df, x='cx', y='cy', ax=ax[1], bins=50, pmax=0.9)\nsns.histplot(data=train_df, x='w', y='h', ax=ax[2], bins=50, pmax=0.9)\nsns.histplot(train_df['w'] * train_df['h'], ax=ax[3], bins=50, kde=False)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:55.082346Z","iopub.execute_input":"2022-07-24T14:11:55.082693Z","iopub.status.idle":"2022-07-24T14:11:55.982107Z","shell.execute_reply.started":"2022-07-24T14:11:55.082658Z","shell.execute_reply":"2022-07-24T14:11:55.981208Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Split dataset to `train:valid=8:2`:","metadata":{}},{"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\n\nsplit_len = round(len(image_ids)*0.8)\n\ntrain_ids = image_ids[:split_len]\nvalid_ids = image_ids[split_len:]\n\ntrain = train_df[train_df['image_id'].isin(train_ids)]\nvalid = train_df[train_df['image_id'].isin(valid_ids)]\n\ntrain.shape, valid.shape","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:55.983316Z","iopub.execute_input":"2022-07-24T14:11:55.983819Z","iopub.status.idle":"2022-07-24T14:11:56.039811Z","shell.execute_reply.started":"2022-07-24T14:11:55.983781Z","shell.execute_reply":"2022-07-24T14:11:56.039082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Show Images","metadata":{}},{"cell_type":"code","source":"def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):\n  figsize = (num_cols * scale, num_rows * scale)\n  _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)\n  axes = axes.flatten()\n  for i, (ax, img) in enumerate(zip(axes, imgs)):\n    ax.imshow(img)\n    ax.axes.get_xaxis().set_visible(False)\n    ax.axes.get_yaxis().set_visible(False)\n    if titles and len(titles) > i:\n      ax.set_title(titles[i])\n  return axes\n\ndef show_bboxes(axes, bboxes, labels=None, colors=None):\n  def _make_list(obj, default_values=None):\n    if obj is None:\n      obj = default_values\n    elif not isinstance(obj, (list, tuple)):\n      obj = [obj]\n    return obj\n\n  labels = _make_list(labels)\n  colors = _make_list(colors, ['b', 'g', 'r', 'm', 'c'])\n  for i, bbox in enumerate(bboxes):\n    color = colors[i % len(colors)]\n    rect = plt.Rectangle(\n      xy=(bbox[0], bbox[1]),\n      width=bbox[2] - bbox[0],\n      height=bbox[3] - bbox[1],\n      fill=False,\n      edgecolor=color,\n      linewidth=2)\n    axes.add_patch(rect)\n    if labels and len(labels) > i:\n      text_color = 'k' if color == 'w' else 'w'\n      axes.text(rect.xy[0], rect.xy[1], labels[i], va='center',\n                ha='center', fontsize=9, color=text_color,\n                bbox=dict(facecolor=color, lw=0))\n\n# https://github.com/d2l-ai/d2l-en/blob/master/d2l/torch.py","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:56.042286Z","iopub.execute_input":"2022-07-24T14:11:56.04295Z","iopub.status.idle":"2022-07-24T14:11:56.055426Z","shell.execute_reply.started":"2022-07-24T14:11:56.042888Z","shell.execute_reply":"2022-07-24T14:11:56.054691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Show images without `bboxes`:","metadata":{}},{"cell_type":"code","source":"num_rows, num_cols = 2, 4\nimgs = [plt.imread(f'{DIR_TRAIN}/{n}.jpg') for n in train_df['image_id'].unique()[:num_rows*num_cols]]\nshow_images(imgs, num_rows, num_cols, scale=4)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:56.058303Z","iopub.execute_input":"2022-07-24T14:11:56.0586Z","iopub.status.idle":"2022-07-24T14:11:57.506574Z","shell.execute_reply.started":"2022-07-24T14:11:56.058576Z","shell.execute_reply":"2022-07-24T14:11:57.505552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Show images with `bboxes`:","metadata":{}},{"cell_type":"code","source":"num_rows, num_cols = 1, 2\nids = train_df['image_id'].unique()[:num_rows*num_cols]\nimgs = [plt.imread(f'{DIR_TRAIN}/{n}.jpg') for n in ids]\naxes = show_images(imgs, num_rows, num_cols, scale=8)\nfor ax, id in zip(axes, ids):\n  datas = train_df[train_df['image_id'] == id]\n  bboxes = [(d['x'], d['y'], d['x']+d['w'], d['y']+d['h']) for _, d in datas.iterrows()]\n  show_bboxes(ax, bboxes, labels=None, colors=['w'])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:11:57.50775Z","iopub.execute_input":"2022-07-24T14:11:57.508061Z","iopub.status.idle":"2022-07-24T14:11:58.251693Z","shell.execute_reply.started":"2022-07-24T14:11:57.508032Z","shell.execute_reply":"2022-07-24T14:11:58.250735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create Dataset\n\nThe dataset should inherit from the standard [torch.utils.data.Dataset](https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset) class, and implement `__len__` and `__getitem__`.\n\nThe only specificity that we require is that the dataset `__getitem__` should return:\n\n* image: a `numpy.ndarray` image\n* target: a dict containing the following fields\n    * `boxes` (`FloatTensor[N, 4]`): the coordinates of the `N` bounding boxes in `[x0, y0, x1, y1]` format, ranging from `0` to `W` and `0` to `H`\n    * `labels` (`Int64Tensor[N]`): the label for each bounding box\n    * `image_id` (`Int64Tensor[1]`): an image identifier. It should be unique between all the images in the dataset, and is used during evaluation\n    * `area` (`Tensor[N]`): The area of the bounding box. This is used during evaluation with the COCO metric, to separate the metric scores between small, medium and large boxes.\n    * `iscrowd` (`UInt8Tensor[N]`): instances with `iscrowd=True` will be ignored during evaluation.","metadata":{}},{"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 Wheat(Dataset):\n\n  def __init__(self, dataframe, image_dir, transforms=None):\n    super().__init__()\n    self.image_ids = dataframe['image_id'].unique()\n    self.df = dataframe\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    records = self.df[self.df['image_id'] == image_id]\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    boxes = records[['x', 'y', 'w', 'h']].values\n\n    area = boxes[:, 2] * boxes[:, 3]\n    area = torch.as_tensor(area, dtype=torch.float32)\n\n    boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n    boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n\n    # there is only one class\n    labels = torch.ones((records.shape[0],), dtype=torch.int64)\n    # suppose all instances are not crowd\n    iscrowd = torch.zeros((records.shape[0],), dtype=torch.uint8)\n\n    target = {}\n    target['boxes'] = boxes\n    target['labels'] = labels\n    target['image_id'] = torch.tensor([idx])\n    target['area'] = area\n    target['iscrowd'] = iscrowd\n\n    if self.transforms:\n      sample = {\n        'image': image,\n        'bboxes': target['boxes'],\n        'labels': labels,\n      }\n      sample = self.transforms(**sample)\n      image = sample['image']\n      target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n    return image, target, image_id\n\n  def __len__(self) -> int:\n    return len(self.image_ids)\n\n  # albumentations\n  #  https://github.com/albumentations-team/albumentations\n\n  @staticmethod\n  def get_train_transform():\n    return A.Compose([\n      A.Flip(0.5),\n      A.ToGray (p=1),\n    ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\n  @staticmethod\n  def get_valid_transform():\n    return A.Compose([\n      ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:13:52.732819Z","iopub.execute_input":"2022-07-24T14:13:52.733224Z","iopub.status.idle":"2022-07-24T14:13:52.748258Z","shell.execute_reply.started":"2022-07-24T14:13:52.733171Z","shell.execute_reply":"2022-07-24T14:13:52.747448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Create `train` `valid` datasets:","metadata":{}},{"cell_type":"code","source":"train_dataset = Wheat(train, DIR_TRAIN, Wheat.get_train_transform())\nvalid_dataset = Wheat(valid, DIR_TRAIN, Wheat.get_valid_transform())","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:13:55.332555Z","iopub.execute_input":"2022-07-24T14:13:55.33288Z","iopub.status.idle":"2022-07-24T14:13:55.349964Z","shell.execute_reply.started":"2022-07-24T14:13:55.332844Z","shell.execute_reply":"2022-07-24T14:13:55.349225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Show images with `transforms`:","metadata":{}},{"cell_type":"code","source":"datas = [train_dataset[i] for i in range(2)]\nimgs = [d[0].permute(1, 2, 0).numpy() for d in datas]\naxes = show_images(imgs, 1, 2, scale=8)\nfor ax, (image, target, image_id) in zip(axes, datas):\n  show_bboxes(ax, target['boxes'], labels=None, colors=['w'])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:13:57.150654Z","iopub.execute_input":"2022-07-24T14:13:57.15102Z","iopub.status.idle":"2022-07-24T14:13:57.977197Z","shell.execute_reply.started":"2022-07-24T14:13:57.150986Z","shell.execute_reply":"2022-07-24T14:13:57.976437Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare Model","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:04.617764Z","iopub.execute_input":"2022-07-24T14:14:04.618098Z","iopub.status.idle":"2022-07-24T14:14:04.623286Z","shell.execute_reply.started":"2022-07-24T14:14:04.618068Z","shell.execute_reply":"2022-07-24T14:14:04.622407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create Model\n\nCreate a Faster R-CNN model pre-trained on COCO:","metadata":{}},{"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:06.472085Z","iopub.execute_input":"2022-07-24T14:14:06.472487Z","iopub.status.idle":"2022-07-24T14:14:09.705718Z","shell.execute_reply.started":"2022-07-24T14:14:06.472455Z","shell.execute_reply":"2022-07-24T14:14:09.704902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Finetune Model","metadata":{}},{"cell_type":"markdown","source":"Print the last layer:","metadata":{}},{"cell_type":"code","source":"print(model.roi_heads.box_predictor)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:11.304786Z","iopub.execute_input":"2022-07-24T14:14:11.305121Z","iopub.status.idle":"2022-07-24T14:14:11.310105Z","shell.execute_reply.started":"2022-07-24T14:14:11.30509Z","shell.execute_reply":"2022-07-24T14:14:11.308897Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Finetune the last layer:","metadata":{}},{"cell_type":"code","source":"num_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)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:13.834949Z","iopub.execute_input":"2022-07-24T14:14:13.835384Z","iopub.status.idle":"2022-07-24T14:14:13.845928Z","shell.execute_reply.started":"2022-07-24T14:14:13.835323Z","shell.execute_reply":"2022-07-24T14:14:13.845112Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Print the last layer again:","metadata":{}},{"cell_type":"code","source":"print(model.roi_heads.box_predictor)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:15.335258Z","iopub.execute_input":"2022-07-24T14:14:15.335621Z","iopub.status.idle":"2022-07-24T14:14:15.341016Z","shell.execute_reply.started":"2022-07-24T14:14:15.335588Z","shell.execute_reply":"2022-07-24T14:14:15.339867Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train Model","metadata":{}},{"cell_type":"markdown","source":"### Create DataLoader","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import DataLoader\n\ndef collate_fn(batch):\n  return tuple(zip(*batch))\n\ntrain_data_loader = DataLoader(\n  train_dataset,\n  batch_size=16,\n  shuffle=False,\n  num_workers=4,\n  collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n  valid_dataset,\n  batch_size=16,\n  shuffle=False,\n  num_workers=4,\n  collate_fn=collate_fn\n)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:20.291415Z","iopub.execute_input":"2022-07-24T14:14:20.29174Z","iopub.status.idle":"2022-07-24T14:14:20.297665Z","shell.execute_reply.started":"2022-07-24T14:14:20.291711Z","shell.execute_reply":"2022-07-24T14:14:20.296567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create Params ","metadata":{}},{"cell_type":"code","source":"# train on the GPU or on the CPU, if a GPU is not available\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n# move model to the right device\nmodel.to(device)\n\n# create an optimizer\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n\n# create a learning rate scheduler\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = None\n\n# train it for 10 epochs\nnum_epochs = 10","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:22.699992Z","iopub.execute_input":"2022-07-24T14:14:22.70032Z","iopub.status.idle":"2022-07-24T14:14:27.176441Z","shell.execute_reply.started":"2022-07-24T14:14:22.700291Z","shell.execute_reply":"2022-07-24T14:14:27.175667Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Let's Train","metadata":{}},{"cell_type":"code","source":"import time\n#from tqdm import tqdm \nfrom tqdm.notebook import tqdm as tqdm\n\nitr = 1\n\ntotal_train_loss = []\ntotal_valid_loss = []\n\nlosses_value = 0\n\nfor epoch in range(num_epochs):\n\n  start_time = time.time()\n\n  # train ------------------------------\n\n  model.train()\n  train_loss = []\n  \n  pbar = tqdm(train_data_loader, desc='let\\'s train')\n  for images, targets, image_ids in pbar:\n    \n    images = list(image.to(device) for image in images)\n    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n    loss_dict = model(images, targets)\n\n    losses = sum(loss for loss in loss_dict.values())\n    losses_value = losses.item()\n    train_loss.append(losses_value)   \n\n    optimizer.zero_grad()\n    losses.backward()\n    optimizer.step()\n\n    pbar.set_description(f\"Epoch: {epoch+1}, Batch: {itr}, Loss: {losses_value}\")\n    itr += 1\n  \n  epoch_train_loss = np.mean(train_loss)\n  total_train_loss.append(epoch_train_loss)\n\n  # update the learning rate\n  if lr_scheduler is not None:\n    lr_scheduler.step()\n\n  # valid ------------------------------\n\n  with torch.no_grad():\n    valid_loss = []\n\n    for images, targets, image_ids in valid_data_loader:\n      images = list(image.to(device) for image in images)\n      targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n      loss_dict = model(images, targets)\n\n      losses = sum(loss for loss in loss_dict.values())\n      loss_value = losses.item()\n      valid_loss.append(loss_value)\n        \n  epoch_valid_loss = np.mean(valid_loss)\n  total_valid_loss.append(epoch_valid_loss)    \n  \n  # print ------------------------------\n\n  print(f\"Epoch Completed: {epoch+1}/{num_epochs}, Time: {time.time()-start_time}, \"\n        f\"Train Loss: {epoch_train_loss}, Valid Loss: {epoch_valid_loss}\")","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:14:30.21318Z","iopub.execute_input":"2022-07-24T14:14:30.213578Z","iopub.status.idle":"2022-07-24T15:10:48.475607Z","shell.execute_reply.started":"2022-07-24T14:14:30.213543Z","shell.execute_reply":"2022-07-24T15:10:48.474467Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Show Loss","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(8, 5))\nsns.set_style(style=\"whitegrid\")\nsns.lineplot(x=range(1, len(total_train_loss)+1), y=total_train_loss, label=\"Train Loss\")\nsns.lineplot(x=range(1, len(total_train_loss)+1), y=total_valid_loss, label=\"Valid Loss\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Loss\")\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T15:18:41.447647Z","iopub.execute_input":"2022-07-24T15:18:41.447991Z","iopub.status.idle":"2022-07-24T15:18:41.7028Z","shell.execute_reply.started":"2022-07-24T15:18:41.447959Z","shell.execute_reply":"2022-07-24T15:18:41.702053Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{}},{"cell_type":"code","source":"test_df = pd.read_csv(f'{DIR_INPUT}/sample_submission.csv')\ntest_df.tail()","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.223789Z","iopub.status.idle":"2022-07-24T14:12:01.224511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create Dataset","metadata":{}},{"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    ])","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.225606Z","iopub.status.idle":"2022-07-24T14:12:01.226305Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create DataLoader","metadata":{}},{"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())","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.227455Z","iopub.status.idle":"2022-07-24T14:12:01.228166Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.229237Z","iopub.status.idle":"2022-07-24T14:12:01.229964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Model","metadata":{}},{"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)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.231137Z","iopub.status.idle":"2022-07-24T14:12:01.231993Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Let's Inference","metadata":{}},{"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))","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.233048Z","iopub.status.idle":"2022-07-24T14:12:01.233783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Show Outputs","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.234851Z","iopub.status.idle":"2022-07-24T14:12:01.23558Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save Results","metadata":{}},{"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]","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.236647Z","iopub.status.idle":"2022-07-24T14:12:01.237347Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.238487Z","iopub.status.idle":"2022-07-24T14:12:01.23919Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2022-07-24T14:12:01.240251Z","iopub.status.idle":"2022-07-24T14:12:01.241053Z"},"trusted":true},"execution_count":null,"outputs":[]}]}