{"cells":[{"metadata":{},"cell_type":"markdown","source":"# RetinaNet in PyTorch - Global Wheat Detection 🌾\n\nIn this notebook, I will show how to train RetinaNet for object detection using PyTorch. You can find more details [here](https://arxiv.org/pdf/1708.02002.pdf).  \n  \nI am using [this](https://github.com/yhenon/pytorch-retinanet) implementation of RetinaNet in PyTorch by [yhenon](https://github.com/yhenon).\n","execution_count":null},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"### Cloning Github Repository \n!git clone https://github.com/yhenon/pytorch-retinanet.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Copying RetinaNet Folder to root dir so we can import it easily\n!cp -r /kaggle/working/pytorch-retinanet/retinanet ./","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install pycocotools","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport re\nimport cv2\nimport time\nimport numpy as np\nimport pandas as pd\n\n\nimport torch\nimport torch.optim as optim\nimport torchvision.transforms as T\nfrom torchvision.utils import make_grid \nfrom torch.utils.data import DataLoader, Dataset\n\nfrom retinanet import model\nfrom retinanet.dataloader import collater, Resizer, Augmenter, Normalizer, UnNormalizer\n\nimport seaborn as sns\nfrom matplotlib import pyplot as plt\n%matplotlib inline\n\nDIR = \"../input/global-wheat-detection/\"\nDIR_TRAIN = DIR + \"train\"\nDIR_TEST = DIR + \"test\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Exploring Dataset 📊","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"### Loading Dataset\ndf = pd.read_csv(DIR + \"train.csv\")\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Converting bbox list from original df to some appropriate form","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"### Converting bbox list in appropriate form\n\ndf['x'] = -1\ndf['y'] = -1\ndf['w'] = -1\ndf['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ndf[['x', 'y', 'w', 'h']] = np.stack(df['bbox'].apply(lambda x: expand_bbox(x)))\ndf.drop(columns=['bbox'], inplace=True)\ndf['x'] = df['x'].astype(np.float)\ndf['y'] = df['y'].astype(np.float)\ndf['w'] = df['w'].astype(np.float)\ndf['h'] = df['h'].astype(np.float)\n\ndf.head()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Null Values, Unique Images, etc.","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"### Null Values, Unique Images, etc.\n\nunq_values = df[\"image_id\"].unique()\nprint(\"Total Records: \", len(df))\nprint(\"Unique Images: \",len(unq_values))\n\nnull_values = df.isnull().sum(axis = 0)\nprint(\"\\n> Null Values in each column <\")\nprint(null_values)\n\n","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"### Data Sources\n\nsources = df[\"source\"].unique()\nprint(\"Total Sources: \",len(sources))\nprint(\"\\n> Sources <\\n\",sources)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Visualizing Source Distribution\n\nplt.figure(figsize=(14,8))\nplt.title('Source Distribution', fontsize= 20)\nsns.countplot(x = \"source\", data = df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Splitting Train Dataset into train - val (80:20)\n\nimages = df['image_id'].unique()\nvalid_imgs = images[-674:]\ntrain_imgs = images[:-674]\n\nvalid_df = df[df['image_id'].isin(valid_imgs)]\ntrain_df = df[df['image_id'].isin(train_imgs)]\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualize Random Images with BBox 🕵️‍","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"### Function to plot image\n\ndef plot_img(image_name):\n    \n    fig, ax = plt.subplots(1, 2, figsize = (10, 10))\n    ax = ax.flatten()\n    \n    records = df[df['image_id'] == image_name]\n    img_path = os.path.join(DIR_TRAIN, image_name + \".jpg\")\n    \n    image = cv2.imread(img_path, cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n    image /= 255.0\n    image2 = image\n    \n    ax[0].set_title('Original Image')\n    ax[0].imshow(image)\n    \n    for idx, row in records.iterrows():\n        box = row[['x', 'y', 'w', 'h']].values\n        xmin = box[0]\n        ymin = box[1]\n        width = box[2]\n        height = box[3]\n        \n        cv2.rectangle(image2, (int(xmin),int(ymin)), (int(xmin + width),int(ymin + height)), (255,0,0), 3)\n    \n    ax[1].set_title('Image with Bondary Box')\n    ax[1].imshow(image2)\n\n    plt.show()\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Pass any image id as parameter\n\nplot_img(\"0126b7d11\")\nplot_img(\"00333207f\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Preparing Dataset for Training 📂","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"### Creating targets for model using Dataset Class\n\nclass GWD(Dataset):\n\n    def __init__(self, dataframe, image_dir, mode = \"train\", transforms = None):\n        \n        super().__init__()\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.mode = mode\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n\n        # Retriving image id and records from df\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        # Loading Image\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        # If mode is set to train, then only we create targets\n        if self.mode == \"train\" or self.mode == \"valid\":\n\n            # Converting xmin, ymin, w, h to x1, y1, x2, y2\n            boxes = np.zeros((records.shape[0], 5))\n            boxes[:, 0:4] = records[['x', 'y', 'w', 'h']].values\n            boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n            boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n            boxes[:, 4] = 1 # This is for label, as we have only 1 class, it is always 1\n            \n            # Applying Transforms\n            sample = {'img': image, 'annot': boxes}\n                \n            if self.transforms:\n                sample = self.transforms(sample)\n\n            return sample\n        \n        elif self.mode == \"test\":\n            \n            # We just need to apply transoforms and return image\n            if self.transforms:\n                \n                sample = {'img' : image}\n                sample = self.transforms(sample)\n                \n            return sample\n        \n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Preparing Datasets and Dataloaders for Training \n\n# Dataset Object\ntrain_dataset = GWD(train_df, DIR_TRAIN, mode = \"train\", transforms = T.Compose([Augmenter(), Normalizer(), Resizer()]))\nvalid_dataset = GWD(valid_df, DIR_TRAIN, mode = \"valid\", transforms = T.Compose([Normalizer(), Resizer()]))\n\n# DataLoaders\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size = 8,\n    shuffle = True,\n    num_workers = 4,\n    collate_fn = collater\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size = 8,\n    shuffle = True,\n    num_workers = 4,\n    collate_fn = collater\n)\n\n\ntest_data_loader = DataLoader(\n    valid_dataset,\n    batch_size = 1,\n    shuffle = True,\n    num_workers = 4,\n    collate_fn = collater\n)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Model - RetinaNet 🔨","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"### Utilize GPU if available\n\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\ntorch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### I am using Pre-trained Resnet50 as backbone\n\nretinanet = model.resnet50(num_classes = 2, pretrained = True)\n\n# Loading Pre-trained model - if you load pre-trained model, comment above line.\n#retinanet = torch.load(\"path_to_.pt_file\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Preparing model for training\n\n# Defininig Optimizer\noptimizer = torch.optim.Adam(retinanet.parameters(), lr = 0.0001)\n\n# Learning Rate Scheduler\n#lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5, gamma=0.5)\n\nretinanet.to(device)\n\n#No of epochs\nepochs = 15\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Now comes everbody's favorite part 😋, let's train it!\nI have defined functions to just improve the readability of the code, model and other parameters are defined outside.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"### One Epoch - Train\n\ndef train_one_epoch(epoch_num, train_data_loader):\n    \n    print(\"Epoch - {} Started\".format(epoch_num))\n    st = time.time()\n    \n    retinanet.train()\n    \n    epoch_loss = []\n\n    for iter_num, data in enumerate(train_data_loader):\n                \n        # Reseting gradients after each iter\n        optimizer.zero_grad()\n            \n        # Forward\n        classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot'].cuda().float()])\n                \n        # Calculating Loss\n        classification_loss = classification_loss.mean()\n        regression_loss = regression_loss.mean()\n\n        loss = classification_loss + regression_loss\n\n        if bool(loss == 0):\n            continue\n                \n        # Calculating Gradients\n        loss.backward()\n\n        # Gradient Clipping\n        torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)\n                \n        # Updating Weights\n        optimizer.step()\n\n        #Epoch Loss\n        epoch_loss.append(float(loss))\n\n            \n        print(\n            'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(\n                epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(epoch_loss)))\n\n        del classification_loss\n        del regression_loss\n        \n    # Update the learning rate\n    #if lr_scheduler is not None:\n        #lr_scheduler.step()\n        \n    et = time.time()\n    print(\"\\n Total Time - {}\\n\".format(int(et - st)))\n    \n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### One Epoch - Valid\n\ndef valid_one_epoch(epoch_num, valid_data_loader):\n    \n    print(\"Epoch - {} Started\".format(epoch_num))\n    st = time.time()\n    \n    epoch_loss = []\n\n    for iter_num, data in enumerate(valid_data_loader):\n                \n        with torch.no_grad():\n            \n            # Forward\n            classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot'].cuda().float()])\n\n            # Calculating Loss\n            classification_loss = classification_loss.mean()\n            regression_loss = regression_loss.mean()\n            loss = classification_loss + regression_loss\n\n            #Epoch Loss\n            epoch_loss.append(float(loss))\n\n            print(\n                'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(\n                    epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(epoch_loss)))\n\n            del classification_loss\n            del regression_loss\n        \n    et = time.time()\n    print(\"\\n Total Time - {}\\n\".format(int(et - st)))\n    \n    # Save Model after each epoch\n    torch.save(retinanet, \"retinanet_gwd.pt\")\n    \n        ","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true,"collapsed":true},"cell_type":"code","source":"### Training Loop\nfor epoch in range(epochs):\n    \n    # Call train function\n    train_one_epoch(epoch, train_data_loader)\n    \n    # Call valid function\n    valid_one_epoch(epoch, valid_data_loader)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Sample Results\nretinanet.eval()\nunnormalize = UnNormalizer()\n\nfor iter_num, data in enumerate(test_data_loader):\n    \n    # Getting Predictions\n    scores, classification, transformed_anchors = retinanet(data['img'].cuda().float())\n    \n    idxs = np.where(scores.cpu()>0.5)\n    img = np.array(255 * unnormalize(data['img'][0, :, :, :])).copy()\n    \n    img[img<0] = 0\n    img[img>255] = 255\n\n    img = np.transpose(img, (1, 2, 0))\n\n    img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_BGR2RGB)\n    \n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    \n    for j in range(idxs[0].shape[0]):\n        bbox = transformed_anchors[idxs[0][j], :]\n        x1 = int(bbox[0])\n        y1 = int(bbox[1])\n        x2 = int(bbox[2])\n        y2 = int(bbox[3])\n\n        cv2.rectangle(img, (x1, y1), (x2, y2), color = (0, 0, 255), thickness = 2)\n        \n    ax.imshow(img)\n    \n    break\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Hope you liked it 😜, comment below your suggestions / feedback. I will upload another notebook for inference soon.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}