{"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":"This notebook is inspired from -\n[previous_competition on object detection](https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train/notebook)\n\nThis is a training notebook.\nThe notebook to infer the model can be found [here](https://www.kaggle.com/palash97/gbr-fasterrcnn-pytorch-inference/)","metadata":{}},{"cell_type":"markdown","source":"## Import Libraries","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport cv2\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom collections import defaultdict\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\nimport torch.nn as nn\n\nfrom torch.cuda.amp import GradScaler, autocast\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-01-29T14:57:25.690241Z","iopub.execute_input":"2022-01-29T14:57:25.690535Z","iopub.status.idle":"2022-01-29T14:57:29.071768Z","shell.execute_reply.started":"2022-01-29T14:57:25.690489Z","shell.execute_reply":"2022-01-29T14:57:29.070966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load train data","metadata":{}},{"cell_type":"code","source":"train_csv_path = '../input/tensorflow-great-barrier-reef/train.csv'\ntrain_df = pd.read_csv(train_csv_path)\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:29.073412Z","iopub.execute_input":"2022-01-29T14:57:29.073694Z","iopub.status.idle":"2022-01-29T14:57:29.134116Z","shell.execute_reply.started":"2022-01-29T14:57:29.073659Z","shell.execute_reply":"2022-01-29T14:57:29.133410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(train_df))","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:29.135202Z","iopub.execute_input":"2022-01-29T14:57:29.135460Z","iopub.status.idle":"2022-01-29T14:57:29.140618Z","shell.execute_reply.started":"2022-01-29T14:57:29.135420Z","shell.execute_reply":"2022-01-29T14:57:29.139801Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Add image path for each row in the train dataframe\n\nImage path for given video_id and video_frame is like: **video_id/video_frame.jpg**","metadata":{}},{"cell_type":"code","source":"def add_image_path(row):\n    video_id = row['video_id']\n    video_frame = row['video_frame']\n    return \"video_\" + str(video_id) + \"/\" + str(video_frame) + \".jpg\"\n\ntrain_df['image_path'] = train_df.apply(lambda x: add_image_path(x), axis=1)\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:29.142643Z","iopub.execute_input":"2022-01-29T14:57:29.143212Z","iopub.status.idle":"2022-01-29T14:57:29.558228Z","shell.execute_reply.started":"2022-01-29T14:57:29.143174Z","shell.execute_reply":"2022-01-29T14:57:29.557386Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build CustomDataset\n\nSince we are using FasterRCNN model, </br>\nDuring **training**,the model excepts **input tensors** and **targets**(list of dictionary). </br>\nDuring **inference**, the model expects only **input tensors**.\n\nMore info on input~output of FasterRCNN model - [here](https://pytorch.org/vision/main/generated/torchvision.models.detection.fasterrcnn_resnet50_fpn.html)","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, image_dir, train_df, transform):\n        self.train_df = train_df\n        self.image_dir = image_dir\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.train_df)\n\n    def __getitem__(self, index):\n        \n        ##################\n        # Read the image #\n        ##################\n        image_path = os.path.join(self.image_dir, self.train_df.loc[index, \"image_path\"])\n        image = cv2.imread(image_path, cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0    # (fasterrcnn model expects input to be in range [0-1])\n        \n        ##########################\n        # Get the bounding boxes #\n        ##########################\n        annots = self.train_df.loc[index, \"annotations\"]\n        boxes = pd.DataFrame(eval(annots), columns=['x','y','width','height']).astype(np.float32).values\n        # Shape of boxes: (num_of_bounding_boxes, 4)\n        # Columns of boxes: (x,y,w,h)\n        \n        ########################\n        # Convert xywh to xyxy #\n        ########################\n        # xyxy is nothing but (x_min, y_min, x_max, y_max) (since fasterrcnn model expects xyxy)\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]  # (x_max = x_min + w)\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]  # (y_max = y_min + h)\n        \n        ##############################\n        # Create a target dictionary #\n        ##############################\n        # (consisting of boxes and labels as its keys)\n        target = {}\n        target['boxes'] = torch.as_tensor(boxes, dtype=torch.float32)\n        target['labels'] = torch.ones((boxes.shape[0],), dtype=torch.int64)\n        # Label 1: COTS-Fish\n        # Label 0: Background (default)\n        \n        # Some extra info\n        target['image_id'] = torch.tensor([index])\n        \n        ############################\n        # Transform the input data #\n        ############################\n\n        # Before transforming check if the bboxes are are valid (i.e., not partially covering the image)\n        # The image size in the given dataset is (h=720, w=1280)\n        is_not_valid = ((boxes[:, 0] < 0).any() or\n                        (boxes[:, 1] < 0).any() or\n                        (boxes[:, 2] > 1280).any() or\n                        (boxes[:, 3] > 720).any())\n        \n        \n        if is_not_valid:\n            image = ToTensorV2()(image=image)['image']\n            return image, target\n        else:\n            # Transform\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': target['labels']\n            }\n\n            sample = self.transform(**sample)\n            image = sample['image']\n            \n            if len(boxes) > 0:\n                target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1,0)\n            \n            return image, target","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:29.559556Z","iopub.execute_input":"2022-01-29T14:57:29.559825Z","iopub.status.idle":"2022-01-29T14:57:29.575552Z","shell.execute_reply.started":"2022-01-29T14:57:29.559789Z","shell.execute_reply":"2022-01-29T14:57:29.574410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transform = A.Compose([A.Flip(0.5),\n                       ToTensorV2()\n                      ], bbox_params={'format': 'pascal_voc','label_fields': ['labels']})\n\n\nimage_dir = '../input/tensorflow-great-barrier-reef/train_images'\n\ndataset = CustomDataset(image_dir=image_dir, train_df=train_df, transform=transform)","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:29.637809Z","iopub.execute_input":"2022-01-29T14:57:29.638019Z","iopub.status.idle":"2022-01-29T14:57:29.643795Z","shell.execute_reply.started":"2022-01-29T14:57:29.637993Z","shell.execute_reply":"2022-01-29T14:57:29.642124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create PyTorch DataLoader","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 4\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_loader = DataLoader(dataset,\n                          shuffle=True,\n                          batch_size=BATCH_SIZE,\n                          collate_fn=collate_fn)","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:31.751142Z","iopub.execute_input":"2022-01-29T14:57:31.751863Z","iopub.status.idle":"2022-01-29T14:57:31.757297Z","shell.execute_reply.started":"2022-01-29T14:57:31.751821Z","shell.execute_reply":"2022-01-29T14:57:31.756446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize one sample from our train_loader","metadata":{}},{"cell_type":"code","source":"num_bboxes = 0\ntries = 10\nwhile (num_bboxes == 0):\n    images, targets = next(iter(train_loader))\n    # images and targets are of list type\n    idx = np.random.randint(0, BATCH_SIZE)\n    img = images[idx]\n    target = targets[idx]\n    num_bboxes = len(target['boxes'])\n    tries -= 1\n    if tries == 0:\n        break\n\nif num_bboxes > 0:        \n    print(img.shape)\n    print(target.keys())\n    print(target['boxes'])\n\n    img = img.permute(1,2,0).numpy()\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    if num_bboxes > 0:\n        boxes = target['boxes'].numpy()\n        for box in boxes:\n            c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))\n            cv2.rectangle(img, c1, c2,\n                      (220, 0, 0), 3)\n\n    plt.title(print(target['image_id']))\n    plt.imshow(img)\n    plt.show()\nelse:\n    print(':(')","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:33.220212Z","iopub.execute_input":"2022-01-29T14:57:33.220836Z","iopub.status.idle":"2022-01-29T14:57:34.759893Z","shell.execute_reply.started":"2022-01-29T14:57:33.220789Z","shell.execute_reply":"2022-01-29T14:57:34.755805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load the model pretrained on COCO","metadata":{}},{"cell_type":"code","source":"device = 'cuda' if torch.cuda.is_available() else 'cpu'\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\nmodel = model.to(device)","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:34.761163Z","iopub.execute_input":"2022-01-29T14:57:34.761431Z","iopub.status.idle":"2022-01-29T14:57:39.210960Z","shell.execute_reply.started":"2022-01-29T14:57:34.761392Z","shell.execute_reply":"2022-01-29T14:57:39.210123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_classes = 2  # 1 class (cots) + 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 head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes).to(device)","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:39.212579Z","iopub.execute_input":"2022-01-29T14:57:39.212837Z","iopub.status.idle":"2022-01-29T14:57:39.220456Z","shell.execute_reply.started":"2022-01-29T14:57:39.212801Z","shell.execute_reply":"2022-01-29T14:57:39.217970Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Test the model architecture on sample data ","metadata":{}},{"cell_type":"code","source":"# Model needs list of input tensors and list of targets\nimages = list(image.to(device) for image in images)\ntargets = [{k: v.long().to(device) for k, v in t.items()} for t in targets]","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:39.222466Z","iopub.execute_input":"2022-01-29T14:57:39.222745Z","iopub.status.idle":"2022-01-29T14:57:39.451059Z","shell.execute_reply.started":"2022-01-29T14:57:39.222710Z","shell.execute_reply":"2022-01-29T14:57:39.450378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loss_dict = model(images, targets)\nlosses = sum(loss for loss in loss_dict.values())\nloss = losses.item()\nloss","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:39.452358Z","iopub.execute_input":"2022-01-29T14:57:39.452745Z","iopub.status.idle":"2022-01-29T14:57:45.263124Z","shell.execute_reply.started":"2022-01-29T14:57:39.452711Z","shell.execute_reply":"2022-01-29T14:57:45.262425Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Awesome!!","metadata":{}},{"cell_type":"markdown","source":"## Let's train the model\n\nWe will be using gradient accumulation, to solve the problem of cuda out of memory (since we have set batch size to be very low).\n\nMore on gradient accumulation - [here](https://towardsdatascience.com/i-am-so-done-with-cuda-out-of-memory-c62f42947dca#:~:text=one%20by%20one.-,Gradient%20Accumulation,-This%20solution%20has)","metadata":{}},{"cell_type":"code","source":"params = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.Adam(params, lr=1e-3)","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:45.264453Z","iopub.execute_input":"2022-01-29T14:57:45.264718Z","iopub.status.idle":"2022-01-29T14:57:45.272193Z","shell.execute_reply.started":"2022-01-29T14:57:45.264683Z","shell.execute_reply":"2022-01-29T14:57:45.271547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_epoch(model, scaler, data_loader, device, optimizer, gradient_accumulations=32):\n    model.train()\n    losses = []\n    for batch_idx, data in enumerate(tqdm(data_loader)):\n        images, targets = data\n        \n        images = list(image.to(device) for image in images)\n        targets = [{k: v.long().to(device) for k, v in t.items()} for t in targets]\n        \n        with autocast():        \n            loss_dict = model(images, targets)\n            loss = sum(loss for loss in loss_dict.values())\n\n        scaler.scale(loss / gradient_accumulations).backward()\n            \n        if (batch_idx + 1) % gradient_accumulations == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            model.zero_grad()\n            \n        losses.append(loss.item())\n        \n    return np.mean(losses)","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:51.401382Z","iopub.execute_input":"2022-01-29T14:57:51.402101Z","iopub.status.idle":"2022-01-29T14:57:51.411596Z","shell.execute_reply.started":"2022-01-29T14:57:51.402043Z","shell.execute_reply":"2022-01-29T14:57:51.410874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train\nEPOCHS = 5\n\nscaler = GradScaler()\nmodel.zero_grad()\n\nfor epoch in range(EPOCHS):\n    print(f'Epoch: {epoch+1}/{EPOCHS}')\n    train_loss = train_epoch(model, scaler, train_loader, device, optimizer)\n    print(f'Train Loss: {train_loss}')","metadata":{"execution":{"iopub.status.busy":"2022-01-29T14:57:51.782885Z","iopub.execute_input":"2022-01-29T14:57:51.783436Z","iopub.status.idle":"2022-01-29T14:58:13.567381Z","shell.execute_reply.started":"2022-01-29T14:57:51.783396Z","shell.execute_reply":"2022-01-29T14:58:13.564540Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Save the model","metadata":{}},{"cell_type":"code","source":"torch.save(model.state_dict(), 'gbr_fasterrcnn_resnet50_fpn.pth.tar')","metadata":{"execution":{"iopub.status.busy":"2022-01-29T05:12:07.857974Z","iopub.execute_input":"2022-01-29T05:12:07.858691Z","iopub.status.idle":"2022-01-29T05:12:08.183128Z","shell.execute_reply.started":"2022-01-29T05:12:07.858651Z","shell.execute_reply":"2022-01-29T05:12:08.18238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}