{"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\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\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":{},"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)\ntrain_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]'\nbbox_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)\ntrain_data.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Checking the data types of dataframe for categorical variables\ntrain_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`\nplt.figure(figsize=(8,6))\nsns.countplot(x='source', data=train_data, alpha=0.6)\nplt.xlabel(\"Image source\", fontsize=15)\nplt.ylabel(\"Count\", fontsize=15)\nplt.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\ntrain_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\nview_images(root_dir, images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Showing the images with bounding boxes\nview_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, data_dir, dataframe, transforms=None):\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        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)\ntarget","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):\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    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":"# 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 torchvision.transforms import functional as F\nfrom torchvision.transforms import ColorJitter\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","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    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    # 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    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()\ndataset_train = torch.utils.data.Subset(dataset_train, indices[:-700])\ndataset_valid = torch.utils.data.Subset(dataset_valid, indices[-700:])\n\n# creating dataloaders\ndataloader_train = DataLoader(dataset_train, batch_size=8, shuffle=True, \n                                   num_workers=4, collate_fn=collate_fn)\n\ndataloader_valid = DataLoader(dataset_valid, batch_size=4, shuffle=False, \n                                  num_workers=4, collate_fn=collate_fn)","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)\nimage","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)\nimage","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)\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\nnum_epochs = 20\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 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\nsaved_model_path = 'fasterrcnn_resnet50_fpn.pth'\nmodel.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]","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\n# torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')","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}