{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"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')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = train_df[train_df['class_id'] != 14].reset_index(drop=True)\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['image_path'] = '../input/vinbigdata-512-image-dataset/vinbigdata/train/'+train_df.image_id+'.png'\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## GROUP KFOLD"},{"metadata":{"trusted":true},"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()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.groupby('fold')['image_id'].agg(lambda x: x.nunique()).reset_index()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df[train_df['image_id'] == '9a5094b2563a1ef3ff50dc5c7ff71345']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dt = VBDDataset(train_df, train_dir)\ndt[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Albumentations\ndef get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Testing Sample"},{"metadata":{"_kg_hide-input":true,"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"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    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for fold in range(num_folds):\n    run_fold(fold)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualize Model"},{"metadata":{"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{},"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":{"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}