{"cells":[{"metadata":{},"cell_type":"markdown","source":"# TorchVision Faster R-CNN Finetuning\n\n## Prepare Dataset"},{"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Add Data\n\n[Global Wheat Detection](https://www.kaggle.com/c/global-wheat-detection)\n\n- train.csv - the training data\n- sample_submission.csv - a sample submission file in the correct format\n- train.zip - training images\n- test.zip - test images"},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Read Data\n\nRead the `train.csv` data:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"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":"\nSplit`bbox` to `x` `y` `w` `h`:"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Analysis Dataset\n\nThe size of train data:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The number of unique `image_id` in train data:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['image_id'].nunique()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The number of images in train dir:"},{"metadata":{"trusted":true},"cell_type":"code","source":"len(os.listdir(DIR_TRAIN))","execution_count":null,"outputs":[]},{"metadata":{},"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":"\nDistribution of size of train images:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['width'].unique(), train_df['height'].unique()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Distribution of number of bounding boxes:"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Distribution of size of bounding boxes:"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Split dataset to `train:valid=8:2`:"},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Show Images"},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Show images without `bboxes`:"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Show images with `bboxes`:"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"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":{"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 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      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']})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Create `train` `valid` datasets:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = Wheat(train, DIR_TRAIN, Wheat.get_train_transform())\nvalid_dataset = Wheat(valid, DIR_TRAIN, Wheat.get_valid_transform())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Show images with `transforms`:"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare 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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create Model\n\nCreate a Faster R-CNN model pre-trained on COCO:"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Finetune Model"},{"metadata":{},"cell_type":"markdown","source":"Print the last layer:"},{"metadata":{"trusted":true},"cell_type":"code","source":"print(model.roi_heads.box_predictor)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"Finetune the last layer:"},{"metadata":{"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Print the last layer again:"},{"metadata":{"trusted":true},"cell_type":"code","source":"print(model.roi_heads.box_predictor)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Model"},{"metadata":{},"cell_type":"markdown","source":"### Create DataLoader"},{"metadata":{"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create Params "},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Let's Train"},{"metadata":{"trusted":true},"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}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Show Loss"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Save Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Let's Inference\n\nSee [TorchVision Faster R-CNN Inference](https://www.kaggle.com/gocoding/torchvision-faster-r-cnn-inference)"},{"metadata":{},"cell_type":"markdown","source":"## References\n\n- [TorchVision Instance Segmentation Finetuning Tutorial](https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html)\n- [Kaggle: Global Wheat Detection](https://www.kaggle.com/c/global-wheat-detection)\n  - [Pytorch Starter - FasterRCNN Train](https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train)\n  - [Global Wheat Detection: Starter EDA](https://www.kaggle.com/kaushal2896/global-wheat-detection-starter-eda)"}],"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}