{"cells":[{"metadata":{},"cell_type":"markdown","source":"* # Faster-RCNN object detection using PyTorch\nStarter code for the competition, inspired by the excellent kernel by [Peter](https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train).","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Import required packages\nimport numpy as np\nimport os\nimport pandas as pd\nimport re\nfrom PIL import Image\nfrom PIL import ImageDraw\nimport gc\nimport warnings\nimport time\nfrom numba import jit\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\nwarnings.filterwarnings('ignore')\n\n# PyTorch libraries\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\nfrom tqdm import tqdm\n\n# OpenCV\nimport cv2\n\n# Plotting libraries\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\n%matplotlib inline\nsns.set_style('whitegrid')\nsns.set_palette('muted')\npd.set_option('display.max_columns', 500)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH = 8\nNUM_EPOCHS = 40\nSLICE_NUM_IMAGES = 0","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataloader and data splitting","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define the root directory\nroot_dir = '../input/global-wheat-detection'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load training data csv file\ntrain_data = pd.read_csv(f'{root_dir}/train.csv')\n\nprint(train_data.shape)\n# train_data.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"From the column description of `train.csv` file, we see that:\n\n- `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- `source` - source of the image file (collaborating universities which provided the data)\n\nLet's clean the dataframe a little bit by extracting the bounding box coordinates and creating separate columns.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get bounding box coordinates and make new columns in dataframe\ndef bbox_coordinates(x):\n    coords = np.array(re.findall(\"([0-9]+\\.?[0-9]*)\", x))\n    coords = list(map(float, coords))\n    if len(x) == 0:\n        coords = [-1, -1, -1, -1]\n    return coords\n\n# Sanity check\nx = '[834.0, 222.0, 56.0, 36.0]'\n# bbox_coordinates(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Add columns for bbox coordinates\ntrain_data['xmin'] = -1\ntrain_data['ymin'] = -1\ntrain_data['width'] = -1\ntrain_data['height'] = -1\n\n# join lists along `axis=0` \nprint(train_data['bbox'].apply(lambda x: bbox_coordinates(x)).shape)\ncoordinates = np.stack(train_data['bbox'].apply(lambda x: bbox_coordinates(x)))\nprint(coordinates.shape)\n\n# Updating the newly created columns\ntrain_data[['xmin', 'ymin', 'width', 'height']] = coordinates\n\n# dropping the bbox column\ntrain_data.drop(columns='bbox', inplace=True)\n# train_data.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Checking the data types of dataframe for categorical variables\n# train_data.dtypes","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Columns `image_id` and `source` are categorical variables, we can check the distribution of images based on the source.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Checking the image distribution based on `source`\n# plt.figure(figsize=(8,6))\n# sns.countplot(x='source', data=train_data, alpha=0.6)\n# plt.xlabel(\"Image source\", fontsize=15)\n# plt.ylabel(\"Count\", fontsize=15)\n# plt.title(\"Distribution of images \\nbased on the source\", fontsize=18);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Number of bounding boxes per image\n# train_data.groupby('image_id').image_id.count()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Put bounding boxes on images\ndef images_with_bbox(df, image_id):\n    boxes = df[df['image_id'] == image_id].loc[:, ['xmin', 'ymin', 'width', 'height']].values\n    image = os.path.join(root_dir, 'train', image_id) + '.jpg'\n    ### OpenCV giving TypeError: an integer is required (got type tuple)\n    ### Switching to PIL\n    # color = (255, 0, 0)\n    # frame = cv2.imread(image)\n    # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)    \n    # image = Image.fromarray(img.mul(255).permute(1,2,0).byte().numpy())\n    frame = Image.open(image).convert(\"RGB\")\n    draw = ImageDraw.Draw(frame)\n    boxes = boxes.astype(int)\n    boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n    boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\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    return frame","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Checking out a few images\ndef view_images(root_dir, image_id, df=None, show_bbox=False):\n    ncols= 4\n    nrows = min(len(image_id)//ncols, 4)\n    \n    fig, ax = plt.subplots(nrows, ncols, figsize=(20, 14))\n    \n    ax = ax.flatten()\n    \n    for i, img_id in enumerate(image_id):\n        img_name = img_id + '.jpg'\n        path = os.path.join(root_dir, 'train', img_name)\n        if show_bbox and (df is not None):\n            image = images_with_bbox(df, img_id)\n            image = np.array(image)\n        else:\n            image = Image.open(path).convert(\"RGB\")\n        ax[i].set_axis_off()\n        ax[i].imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Showing random images from the data\nimages = train_data.sample(n=8)['image_id'].values\n# view_images(root_dir, images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Showing the images with bounding boxes\n# view_images(root_dir, images, df=train_data, show_bbox=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Creating the dataset\nclass WheatDataset(Dataset):\n    def __init__(self, dataframe, data_dir, transforms=None):\n        super().__init__()\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.data = data_dir\n        self.transforms = transforms\n        # load all the images and sort them\n        self.images = sorted(os.listdir(os.path.join(data_dir)))\n        if SLICE_NUM_IMAGES != 0:\n            a = self.images.copy()\n            self.images = a[:SLICE_NUM_IMAGES]\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        boxes = boxes_info[['xmin', 'ymin', 'width', 'height']].values\n        \n        # get the xmax and ymax from width and height\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]    # xmax\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]    # ymax\n        \n        # load the image\n        path = os.path.join(self.data, image_id) + '.jpg'\n        image = Image.open(path).convert(\"RGB\")\n        \n        # convert boxes to tensor\n        boxes = torch.tensor(boxes, dtype=torch.float32)\n        \n        # create labels; there's only one class\n        labels = torch.ones((len(boxes),), dtype=torch.int64)\n        \n        # iscrowd; for background\n        iscrowd = torch.zeros((len(boxes),), dtype=torch.int64)\n        \n        # area of bounding boxes\n        area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])\n        \n        # image label\n        image_label = torch.tensor([idx])\n        \n         # storing all the attributes in a dictionary\n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = image_label\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n        \n        if self.transforms is not None:\n            image, target = self.transforms(image, target)\n            \n        return image, target, image_id\n    \n    def __len__(self):\n        return len(self.images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Sanity check to see if dataset class is working\ndata_dir = os.path.join(root_dir, 'train')\ndataset = WheatDataset(data_dir=data_dir, dataframe=train_data)\nimg, target, img_id = dataset.__getitem__(3)\n# target","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Download and Fine-tune the Faster-RCNN model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model(num_classes, model_name):\n    if model_name == 'resnet50':\n        # load the object detection model pre-trained on COCO\n        model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\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    elif model_name == 'resnet101':\n        backbone = resnet_fpn_backbone('resnet101', pretrained=True)\n        model = FasterRCNN(backbone, num_classes=2)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# helper class to keep track of loss and loss per iteration\n# source: https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train\nclass Averager:\n    def __init__(self):\n        self.current_total = 0\n        self.iterations = 0\n        \n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n        \n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n        \n    def reset(self):\n        self.current_total = 0\n        self.iterations = 0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def eli_domination(a_image):\n    print('a')\n    img = a_image\n    \n    img = np.array(img) \n    \n    ### Make smaller version of image of calculations\n    start_r = randrange(img.shape[0] / 2)\n#     end_r = start_r + randrange(200, 300)\n    end_r = start_r + randrange(2)\n\n    start_c = randrange(img.shape[1] / 2)\n#     end_c = start_r + randrange(200, 300)\n    end_c = start_r + randrange(2)\n\n    img_small = img [start_r:end_r, start_c: end_c, :]\n\n    image_work = img_small.copy()\n#     image_work = img\n    \n    ### Get Average\n    average = image_work.mean(axis=0).mean(axis=0)\n    avg = np.average(average)\n\n    ### Get dominant colors by KNN\n    pixels = np.float32(image_work.reshape(-1, 3))\n\n    n_colors = 5\n    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, .1)\n    flags = cv2.KMEANS_RANDOM_CENTERS\n\n    _, labels, palette = cv2.kmeans(pixels, n_colors, None, criteria, 10, flags)\n    _, counts = np.unique(labels, return_counts=True)\n    \n    dominant = palette[np.argmin(counts)]\n    \n    print('c')\n    ### Create new image\n    rad = avg / 1.1\n    dominant_0_min = dominant[0] - rad\n    dominant_0_max = dominant[0] + rad\n\n    dominant_1_min = dominant[1] - rad\n    dominant_1_max = dominant[1] + rad\n\n    dominant_2_min = dominant[2] - rad\n    dominant_2_max = dominant[2] + rad\n\n#     chosen_color = [0, 0, 0]\n#     chosen_color = [255, 255, 255]\n    print('d')\n    add_color = rad\n    copy_image = img.copy()\n    \n    for row in range(copy_image.shape[0]):\n      for col in range(copy_image.shape[1]):\n        if (dominant_0_max >= copy_image[row][col][0] >= dominant_0_min) \\\n        and (dominant_1_max >= copy_image[row][col][1] >= dominant_1_min) \\\n        and (dominant_2_max >= copy_image[row][col][2] >= dominant_2_min):\n          copy_image[row][col] = [min(255, copy_image[row][col][0] + add_color), min(255, copy_image[row][col][1] + add_color), min(255, copy_image[row][col][0] + add_color)]\n    print('e')\n    copy_image = Image.fromarray(copy_image)\n    return copy_image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# creating some transforms\n# source: https://github.com/pytorch/vision/blob/master/references/detection/transforms.py\n# source: https://github.com/microsoft/computervision-recipes/blob/master/utils_cv/detection/dataset.py\nimport random\nfrom random import randrange\nfrom torchvision.transforms import functional as F\nfrom torchvision.transforms import ColorJitter\n\ngammas = [0.25, 0.5, 1, 2]\n\n    \nclass Compose(object):\n    def __init__(self, transforms):\n        self.transforms = transforms\n        \n    def __call__(self, image, target):\n        for t in self.transforms:\n            image, target = t(image, target)\n        return image, target\n\nclass RandomHorizontalFlip(object):\n    \"\"\"\n    Wrapper for torchvision's HorizontalFlip\n    \"\"\"\n    def __init__(self, prob):\n        self.prob = prob\n        \n    def __call__(self, image, target):\n        if random.random() < self.prob:\n            height, width = image.shape[-2:]  # image must be a ndarray or torch tensor (PIL.Image has no attribute `.shape`)\n            image = image.flip(-1)\n            bbox = target['boxes']    # bbox coordinates MUST be of form [xmin, ymin, xmax, ymax]\n            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]\n            target['boxes'] = bbox\n        return image, target\n\nclass ColorJitterTransform(object):\n    \"\"\"\n    Wrapper for torchvision's ColorJitter\n    \"\"\"\n    def __init__(self, brightness, contrast, saturation, hue):\n        self.brightness = brightness\n        self.contrast = contrast\n        self.saturation = saturation\n        self.hue = hue\n    \n    def __call__(self, image, target):\n        image = ColorJitter(\n            brightness=self.brightness,\n            contrast=self.contrast,\n            saturation=self.saturation,\n            hue=self.hue\n        )(image)\n        return image, target\n        \nclass ToTensor(object):\n    def __call__(self, image, target):\n        image = F.to_tensor(image)   # normalizes the image and converts PIL image to torch.tensor\n        return image, target\n\nclass eli_bright(object):\n    def __init__(self, threshold, value):\n        self.threshold = threshold\n        self.value = value\n    \n    def __call__(self, image, target):\n        image[image > self.threshold] = self.value\n        \n        return image, target\n    \nclass gamma(object):\n    def __init__(self, gama, gain):\n        self.gama = gama\n        self.gain = gain\n        \n    def __call__(self, image, target):\n        image = F.adjust_gamma(image, gammas[randrange(len(gammas))], self.gain)\n            \n        return image, target\n    \nclass random_crop(object):\n    def __init__(self, nothing):\n        self.nothing = nothing\n        \n    def __call__(self, image, target):\n        image = torchvision.transforms.RandomCrop(image)\n        \n        return image, target\n    \nclass eli_dom(object):\n    def __init__(self, nothing):\n        self.nothing = nothing\n        \n    def __call__(self, image, target):\n        image = eli_domination(image)\n        \n        return image, target\n        ","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# add some image augmentation operations\ndef get_transform(train):\n    transforms = []\n    \n#     transforms.append(eli_dom(0))\n    \n    if train:\n        # needs the image to be a PIL image\n        transforms.append(ColorJitterTransform(brightness=0.2, contrast=0.2, saturation=0.4, hue=0.05))\n        \n#         transforms.append(gamma(2, 1))\n        \n#         transforms.append(random_crop(0))\n        \n    # converts a PIL image to pytorch tensor\n    transforms.append(ToTensor())\n    \n    if train:\n        # randomly flip the images, bboxes and ground truth (only during training)\n        transforms.append(RandomHorizontalFlip(0.5))  # this operation needs image to be a torch tensor\n        \n#         transforms.append(eli_bright(3, 0))\n#     if train:\n#         transforms.append(eli_bright(0.5, 0))\n\n    \n        \n    return Compose(transforms)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Use the dataset class and create train and test dataloaders\ndataset_train = WheatDataset(data_dir=data_dir, dataframe=train_data, \n                             transforms=get_transform(train=True))\ndataset_valid = WheatDataset(data_dir=data_dir, dataframe=train_data, \n                             transforms=get_transform(train=False))\n\n# split the dataset in train and valid\nindices = torch.randperm(len(dataset_train)).tolist()\n# print(int(-len(dataset_train)*0.33))\ndataset_train = torch.utils.data.Subset(dataset_train, indices[:int(-len(dataset_train)*0.15)])\ndataset_valid = torch.utils.data.Subset(dataset_valid, indices[int(-len(dataset_train)*0.15):])\n\n# creating dataloaders\ndataloader_train = DataLoader(dataset_train, batch_size=BATCH, shuffle=True, \n                                   num_workers=4, collate_fn=collate_fn)\n\ndataloader_valid = DataLoader(dataset_valid, batch_size=BATCH, shuffle=False, \n                                  num_workers=4, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# a = next(iter(dataloader_train))\n# print(a[0][3])\n# plt.imshow(a[0][7].permute(1, 2, 0))\n\n# cv2.imread(a[0][0])\n","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":"def show_bbox_single_image(images, targets):\n    image = Image.fromarray(images[0].mul(255).permute(1,2,0).cpu().byte().numpy())\n    boxes = targets[0]['boxes'].cpu().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=3)\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# sample in training set \nimages, targets, image_ids = next(iter(dataloader_train))\nimage = show_bbox_single_image(images, targets)\n# image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# sample in validation set\nimages, targets, image_id = next(iter(dataloader_valid))\nimage = show_bbox_single_image(images, targets)\n# image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 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, 'resnet101')\n\n# move the model to device\nmodel.to(device)\n\n# creating the 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# learning rate scheduler; reduce lr by 0.5 after every 3 epochs if loss reaches plateau or doesn't decrease\nlr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', \n                                                          patience=3, verbose=True, \n                                                          factor=0.5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training the model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"## Source: https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train\nloss_hist = Averager()\nitr=1\nmin_loss = np.Inf\n\n\nfor epoch in range(1, NUM_EPOCHS+1):\n    # keep track of time to run each epoch\n    epoch_start_time = time.time()\n    # reset the loss history (fresh start for each epoch)\n    loss_hist.reset()\n    # put the model in train mode\n    model.train()\n    for images, targets, image_ids in tqdm(dataloader_train):\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        losses = sum(loss for loss in loss_dict.values())\n        loss_value = losses.item()\n        \n        loss_hist.send(loss_value)\n        \n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n        \n        if itr % 500 == 0:\n            print(f\"Iteration #{itr},  loss: {loss_value:.6f}\")\n        \n        itr += 1\n    # update learning rate\n    if lr_scheduler is not None:\n        lr_scheduler.step(loss_hist.value)\n        \n    print(f\"Epoch #{epoch},  time taken: {(time.time() - epoch_start_time):.2f} seconds,  loss: {loss_hist.value:.6f}\")\n    if loss_hist.value <= min_loss:\n        print(f\"Loss decreased: {min_loss:.6f} --> {loss_hist.value:.6f}. Saving model ...\")\n        torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')\n        min_loss = loss_hist.value","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> ## Checking the model performance","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load the saved model\n# saved_model_path = '../input/themodel/fasterrcnn_resnet50_fpn.pth'\n# model.load_state_dict(torch.load(saved_model_path))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# test the model in validation set\nimages, targets, label_id = next(iter(dataloader_valid))\nimages = list(img.to(device) for img in images)\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# setting the model to evaluation mode\nmodel.eval()\ncpu_device = torch.device(\"cpu\")\noutputs = model(images)\noutputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]","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":"# checking the model output\nshow_inferences(images, outputs, 4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# saving the model\ntorch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@jit(nopython=True)\ndef calculate_iou(gt, pr, form='pascal_voc') -> float:\n    \"\"\"Calculates the Intersection over Union.\n\n    Args:\n        gt: (np.ndarray[Union[int, float]]) coordinates of the ground-truth box\n        pr: (np.ndarray[Union[int, float]]) coordinates of the prdected box\n        form: (str) gt/pred coordinates format\n            - pascal_voc: [xmin, ymin, xmax, ymax]\n            - coco: [xmin, ymin, w, h]\n    Returns:\n        (float) Intersection over union (0.0 <= iou <= 1.0)\n    \"\"\"\n    if form == 'coco':\n        gt = gt.copy()\n        pr = pr.copy()\n\n        gt[2] = gt[0] + gt[2]\n        gt[3] = gt[1] + gt[3]\n        pr[2] = pr[0] + pr[2]\n        pr[3] = pr[1] + pr[3]\n\n    # Calculate overlap area\n    dx = min(gt[2], pr[2]) - max(gt[0], pr[0]) + 1\n    \n    if dx < 0:\n        return 0.0\n    \n    dy = min(gt[3], pr[3]) - max(gt[1], pr[1]) + 1\n\n    if dy < 0:\n        return 0.0\n\n    overlap_area = dx * dy\n\n    # Calculate union area\n    union_area = (\n            (gt[2] - gt[0] + 1) * (gt[3] - gt[1] + 1) +\n            (pr[2] - pr[0] + 1) * (pr[3] - pr[1] + 1) -\n            overlap_area\n    )\n\n    return overlap_area / union_area\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@jit(nopython=True)\ndef find_best_match(gts, pred, pred_idx, threshold = 0.5, form = 'pascal_voc', ious=None) -> int:\n    \"\"\"Returns the index of the 'best match' between the\n    ground-truth boxes and the prediction. The 'best match'\n    is the highest IoU. (0.0 IoUs are ignored).\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        pred: (List[Union[int, float]]) Coordinates of the predicted box\n        pred_idx: (int) Index of the current predicted box\n        threshold: (float) Threshold\n        form: (str) Format of the coordinates\n        ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n    Return:\n        (int) Index of the best match GT box (-1 if no match above threshold)\n    \"\"\"\n    best_match_iou = -np.inf\n    best_match_idx = -1\n\n    for gt_idx in range(len(gts)):\n        \n        if gts[gt_idx][0] < 0:\n            # Already matched GT-box\n            continue\n        \n        iou = -1 if ious is None else ious[gt_idx][pred_idx]\n\n        if iou < 0:\n            iou = calculate_iou(gts[gt_idx], pred, form=form)\n            \n            if ious is not None:\n                ious[gt_idx][pred_idx] = iou\n\n        if iou < threshold:\n            continue\n\n        if iou > best_match_iou:\n            best_match_iou = iou\n            best_match_idx = gt_idx\n\n    return best_match_idx\n\n@jit(nopython=True)\ndef calculate_precision(gts, preds, threshold = 0.5, form = 'coco', ious=None) -> float:\n    \"\"\"Calculates precision for GT - prediction pairs at one threshold.\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n               sorted by confidence value (descending)\n        threshold: (float) Threshold\n        form: (str) Format of the coordinates\n        ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n    Return:\n        (float) Precision\n    \"\"\"\n    n = len(preds)\n    tp = 0\n    fp = 0\n    \n    # for pred_idx, pred in enumerate(preds_sorted):\n    for pred_idx in range(n):\n\n        best_match_gt_idx = find_best_match(gts, preds[pred_idx], pred_idx,\n                                            threshold=threshold, form=form, ious=ious)\n\n        if best_match_gt_idx >= 0:\n            # True positive: The predicted box matches a gt box with an IoU above the threshold.\n            tp += 1\n            # Remove the matched GT box\n            gts[best_match_gt_idx] = -1\n\n        else:\n            # No match\n            # False positive: indicates a predicted box had no associated gt box.\n            fp += 1\n\n    # False negative: indicates a gt box had no associated predicted box.\n    fn = (gts.sum(axis=1) > 0).sum()\n\n    return tp / (tp + fp + fn)\n\n\n@jit(nopython=True)\ndef calculate_image_precision(gts, preds, thresholds = (0.5, ), form = 'coco') -> float:\n    \"\"\"Calculates image precision.\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n               sorted by confidence value (descending)\n        thresholds: (float) Different thresholds\n        form: (str) Format of the coordinates\n\n    Return:\n        (float) Precision\n    \"\"\"\n    n_threshold = len(thresholds)\n    image_precision = 0.0\n    \n    ious = np.ones((len(gts), len(preds))) * -1\n    # ious = None\n\n    for threshold in thresholds:\n        precision_at_threshold = calculate_precision(gts.copy(), preds, threshold=threshold,\n                                                     form=form, ious=ious)\n        image_precision += precision_at_threshold / n_threshold\n\n    return image_precision\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sum_precisions = 0\nfor box_num in range(BATCH):\n    temp_precision = calculate_precision(targets[box_num]['boxes'].cpu().numpy(), outputs[box_num]['boxes'].cpu().detach().numpy(), threshold=0.5, form='coco')\n#     temp_precision = calculate_precision(targets[box_num]['boxes'].cpu().numpy(), outputs[box_num]['boxes'].cpu().detach().numpy(), threshold=0.5, form='pascal_voc')\n#     print(temp_precision)\n    sum_precisions += temp_precision\n#     sum_precisions += calculate_precision(targets[0]['boxes'].cpu().numpy(), outputs[0]['boxes'].cpu().detach().numpy(), threshold=0.5, form='pascal_voc')\n\nprecision = sum_precisions / BATCH\nprecision\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## References\n- [PyTorch vision repo](https://github.com/pytorch/vision/tree/master/references/detection)\n- [Medium post](https://towardsdatascience.com/building-your-own-object-detector-pytorch-vs-tensorflow-and-how-to-even-get-started-1d314691d4ae)","execution_count":null}],"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}