{"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":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\nimport time\n\nfrom sklearn.model_selection import GroupKFold\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\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\nfrom matplotlib import pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dir = '../input/vinbigdata-512-image-dataset/vinbigdata/train'\ntest_dir = '../input/vinbigdata-512-image-dataset/vinbigdata/test'\ntrain_df = pd.read_csv('../input/vinbigdata-512-image-dataset/vinbigdata/train.csv')","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = train_df[train_df['class_id'] != 14].reset_index(drop=True)\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df['image_path'] = '../input/vinbigdata-512-image-dataset/vinbigdata/train/'+train_df.image_id+'.png'\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## GROUP KFOLD","metadata":{}},{"cell_type":"code","source":"gkf  = GroupKFold(n_splits = 5)\ntrain_df['fold'] = -1\nfor fold, (train_idx, val_idx) in enumerate(gkf.split(train_df, groups = train_df.image_id.tolist())):\n    train_df.loc[val_idx, 'fold'] = fold\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.groupby('fold')['image_id'].agg(lambda x: x.nunique()).reset_index()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMG_SIZE = 512\ntrain_df['xmin'] = (train_df['x_min']/train_df['width'])*IMG_SIZE\ntrain_df['ymin'] = (train_df['y_min']/train_df['height'])*IMG_SIZE\ntrain_df['xmax'] = (train_df['x_max']/train_df['width'])*IMG_SIZE\ntrain_df['ymax'] = (train_df['y_max']/train_df['height'])*IMG_SIZE","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"assert train_df['xmin'].all() <= IMG_SIZE\nassert train_df['ymin'].all() <= IMG_SIZE\nassert train_df['xmax'].all() <= IMG_SIZE\nassert train_df['ymax'].all() <= IMG_SIZE","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df[train_df['image_id'] == '9a5094b2563a1ef3ff50dc5c7ff71345']","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class_dict = dict(set(zip(train_df.class_id, train_df.class_name)))\nclasses = []\nfor key in sorted(class_dict.keys()): \n    classes.append(class_dict[key])\n\nclasses = ['_'] + classes   # adding background\nclasses","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class VBDDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\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):\n\n        image_id = self.image_ids[idx]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.png', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        boxes = records[['xmin', 'ymin', 'xmax', 'ymax']].values\n        \n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n        # all the labels are shifted by 1 to accomodate background\n        labels = torch.squeeze(torch.as_tensor((records.class_id.values+1,), dtype=torch.int64))\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        # target['masks'] = None\n        target['image_id'] = torch.tensor([idx])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\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            \n            target['boxes'] = torch.as_tensor(sample['bboxes'])\n\n        return image, target, image_id\n\n    def __len__(self):\n        return self.image_ids.shape[0]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dt = VBDDataset(train_df, train_dir)\ndt[0]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Albumentations\ndef get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        A.ShiftScaleRotate(rotate_limit=10, p=0.5),\n        A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=1.0),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        A.ShiftScaleRotate(rotate_limit=10, p=0.5),\n        A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=1.0),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_classes = 15  # 14 classes + 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)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# A Class for keeping track of average\nclass Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.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.0\n        self.iterations = 0.0","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Testing Sample","metadata":{}},{"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = VBDDataset(train_df, train_dir, get_train_transform())\nvalid_dataset = VBDDataset(train_df, train_dir, get_valid_transform())\n\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=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nimages, targets, image_ids = next(iter(train_data_loader))\nimages = list(image.to(device) for image in images)\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\nboxes = targets[2]['boxes'].cpu().numpy().astype(np.int32)\nsample = images[2].permute(1, 2, 0).cpu().numpy()\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2], box[3]),\n                  (220, 0, 0), 3)\n    \nax.set_axis_off()\nax.imshow(sample)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"def get_dataloaders(df, trn_idx, val_idx):\n    \n    train_ = df.loc[trn_idx,:].reset_index(drop=True)\n    valid_ = df.loc[val_idx,:].reset_index(drop=True)\n        \n    def collate_fn(batch):\n        return tuple(zip(*batch))\n\n    train_dataset = VBDDataset(train_, train_dir, get_train_transform())\n    valid_dataset = VBDDataset(valid_, train_dir, get_valid_transform())\n\n\n    train_data_loader = DataLoader(\n        train_dataset,\n        batch_size=16,\n        shuffle=False,\n        num_workers=4,\n        collate_fn=collate_fn\n    )\n\n    valid_data_loader = DataLoader(\n        valid_dataset,\n        batch_size=8,\n        shuffle=False,\n        num_workers=4,\n        collate_fn=collate_fn\n    )\n    \n    return train_data_loader, valid_data_loader\n\n\n\ndef train_model(model, dataloader, device, epochs, optimizer, lr_scheduler, fold):\n    \n    best_loss = 1e10\n    loss_hist = Averager()\n    itr = 1\n    all_losses = []\n\n    for epoch in range(epochs):\n        loss_hist.reset() \n    \n        for images, targets, image_ids in dataloader:\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            loss_value = losses.item()\n\n            loss_hist.send(loss_value)\n            all_losses.append(loss_value)\n            \n            optimizer.zero_grad()\n            losses.backward()\n            optimizer.step()\n\n            if itr % 50 == 0:\n                print(f\"Iteration #{itr} loss: {loss_value}\")\n\n            itr += 1\n        \n        # saving the model based on training loss for now. - later can be moved to validation\n        if loss_hist.value < best_loss:\n            best_loss = loss_hist.value\n            torch.save(model.state_dict(), f'fasterrcnn_model_{fold}.pt')\n\n        # update the learning rate\n        if lr_scheduler is not None:\n            lr_scheduler.step()\n\n        print(f\"Epoch #{epoch} loss: {loss_hist.value}\\n\")\n        \n    return all_losses\n        \n        \ndef validate_model(model, dataloader, device):\n    print(\"\\n Starting Validation ... \")\n    loss_hist = Averager()\n    itr = 1\n\n    loss_hist.reset() \n\n    for images, targets, image_ids in dataloader:\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        loss_value = losses.item()\n\n        loss_hist.send(loss_value)\n\n        if itr % 50 == 0:\n            print(f\"Iteration #{itr} loss: {loss_value}\")\n\n        itr += 1\n\n    print(f\"\\nFinal loss: {loss_hist.value}\")\n\n\n    ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_fold(fold):\n    print(f\"Starting fold {fold}\")\n    start = time.time()\n    trn_idx = train_df[train_df['fold'] != fold].index\n    val_idx = train_df[train_df['fold'] == fold].index\n    \n    \n    trainloader, valloader = get_dataloaders(train_df, trn_idx, val_idx)\n    \n    loss_hist = train_model(model, trainloader, device, epochs, optimizer, lr_scheduler, fold)\n    \n    # plot training loss\n    plt.figure(figsize=(8,5))\n    plt.plot(loss_hist)\n    plt.title(\"Training Loss Statistic\", size=17)\n    plt.xlabel(\"Iteration\", size=15)\n    plt.ylabel(\"Loss Value\", size=15)\n    plt.show()\n    \n    validate_model(model, valloader, device)\n    \n    print(f\"Completed Fold {fold} in {round(time.time()-start, 2)} seconds\")\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.to(device)\n\n# set params for model\nparams = [p for p in model.parameters() if p.requires_grad]\n\n# set optimizer\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n\n# set lr scheduler\nlr_scheduler = None\n\n# set epochs\nepochs = 20\n\n# set folds\nnum_folds = 1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for fold in range(num_folds):\n    run_fold(fold)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize Model","metadata":{}},{"cell_type":"code","source":"len(valid_data_loader)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model0 = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\nmodel0.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\nmodel0.load_state_dict(torch.load(model_dir + 'fasterrcnn_model_4.pt'))\nmodel0.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_output(model, num_imgs, score_threshold=0.9):\n    img_cnt = 0\n    for images, targets, image_ids in iter(valid_data_loader):\n        if img_cnt == num_imgs:\n            break\n        img_cnt += 1\n\n        print(f\"----------{img_cnt}-----------\")\n\n        images = list(img.to(device) for img in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n        boxes = targets[1][\"boxes\"].cpu().numpy().astype(np.int32)\n        sample = images[1].permute(1, 2, 0).cpu().numpy()\n        clss = targets[1][\"labels\"].cpu().numpy().astype(np.int32)\n\n        for c in range(len(clss)):\n          print(clss[c], boxes[c])\n        # print(\"boxes:\")\n        # print(boxes)\n        # print(\"classes:\")\n        # print(clss)\n\n        # clss_boxes_dict = {}\n        # ## first box\n        # for i in range(len(clss)):\n        #   if clss[i] not in clss_boxes_dict.keys():\n        #     clss_boxes_dict[clss[i]] = boxes[i]\n        # print(clss_boxes_dict)\n\n        # clss_boxes_dict = {}\n        # ## average box\n        # for i in range(len(clss)):\n        #   clas = clss[i]\n        #   curr_box = boxes[i]\n\n        #   if clas not in clss_boxes_dict.keys():\n        #     clss_boxes_dict[clas] = (curr_box, 1)\n\n        #   else:\n        #     curr_sum, curr_len = clss_boxes_dict[clas]\n        #     clss_boxes_dict[clas] = (curr_sum + curr_box, curr_len + 1)\n\n        # for clas in clss_boxes_dict.keys():\n        #   clss_boxes_dict[clas] = (clss_boxes_dict[clas][0] / clss_boxes_dict[clas][1]).astype(np.int32)\n        # print(clss_boxes_dict)\n\n        model.eval()\n        cpu_device = torch.device(\"cpu\")\n\n        outputs = model(images)\n        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]\n\n        # outputs_class_boxes prepares labels and boxes to be plotted\n\n        outputs_class_boxes = {}\n        for output_dict in outputs:\n            boxes, labels, scores = (\n                output_dict[\"boxes\"],\n                output_dict[\"labels\"],\n                output_dict[\"scores\"],\n            )\n            for s in range(len(scores)):\n                box, label, score = (\n                    boxes[s].detach().numpy().astype(np.int32),\n                    labels[s].item(),\n                    scores[s].item(),\n                )\n                if scores[s] > score_threshold:\n                    if labels[s] not in outputs_class_boxes.keys():\n                        outputs_class_boxes[label] = [box]\n                    else:\n                        print(f\"adding more boxes for label {label}\")\n                        outputs_class_boxes[label] += [box]\n                        \n        for clas in sorted(outputs_class_boxes.keys()):\n            print(clas, outputs_class_boxes[clas])\n\n        fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\n        for box, clas in zip(boxes, clss):\n            cv2.putText(\n                sample,\n                #f\"{classes[clas]}\",\n                f\"{clas}\",\n                (box[0], box[1] + 20),\n                cv2.FONT_HERSHEY_DUPLEX,\n                1.0,\n                (127, 0, 127),\n                1,\n            )\n            cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (220, 0, 0), 1)\n\n        for clas in outputs_class_boxes.keys():\n            for box in outputs_class_boxes[clas]:\n                cv2.putText(\n                    sample,\n                    #f\"{classes[clas]}\",\n                    f\"{clas}\",\n                    #((box[0] + box[2]) // 2, (box[1] + box[3]) // 2),\n                    (box[0], box[3]),\n                    cv2.FONT_HERSHEY_DUPLEX,\n                    1.0,\n                    (0, 0, 255),\n                    1,\n                )\n                cv2.rectangle(\n                    sample, (box[0], box[1]), (box[2], box[3]), (0, 255, 127), 2\n                )\n\n        ax.set_axis_off()\n        ax.imshow(sample)\n\n        for _ in range(3):\n          print()\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_output(model0, 5, 0.95)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"images, targets, image_ids = next(iter(valid_data_loader))\n\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\nboxes = targets[1]['boxes'].cpu().numpy().astype(np.int32)\nsample = images[1].permute(1,2,0).cpu().numpy()\nclss = targets[1]['labels'].cpu().numpy().astype(np.int32)\n\nmodel.eval()\ncpu_device = torch.device(\"cpu\")\n\noutputs = model(images)\noutputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box, clas in zip(boxes, clss):\n    cv2.putText(sample, f\"{classes[clas]}\", (box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1)\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2], box[3]),\n                  (220, 0, 0), 1)\n    \nax.set_axis_off()\nax.imshow(sample)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Acknowledgements\nNotebook Heavily inspired by this Notebook - https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train/notebook.\n\n**If the kernel helps you in any way, kindly Upvote**","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}