{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Train and Inference for FasterRCNN with Augmentation\n\nFor EDA and Beginner Friendly code checkout [this](https://www.kaggle.com/yashchoudhary/gwd-eda-and-starter-code-beginner-friendly) notebook.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Setup","execution_count":null},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2, ToTensor\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\nfrom matplotlib import pyplot as plt\ngodimg = 0\nDIR_INPUT = '/kaggle/input'\nDIR_TRAIN = f'{DIR_INPUT}/global-wheat-detection/train'\nDIR_TEST = f'{DIR_INPUT}/global-wheat-detection/test'","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/global-wheat-detection/train.csv')\ntest_df = pd.read_csv(f'{DIR_INPUT}/global-wheat-detection/sample_submission.csv')\ntrain_df.shape\n\ntrain_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['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\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)\n\nimage_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]\n\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Creating the Train and Test Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.utils import shuffle\nimport random\nclass WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n        \n        self.df = dataframe\n        self.image_ids = dataframe['image_id'].unique()\n        self.image_ids = shuffle(self.image_ids)\n        self.labels = [np.zeros((0, 5), dtype=np.float32)] * len(self.image_ids)\n        self.img_size = 1024\n        im_w = 1024\n        im_h = 1024\n        for i, img_id in enumerate(self.image_ids):\n            records = self.df[self.df['image_id'] == img_id]\n            boxes = records[['x', 'y', 'w', 'h']].values\n            boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n            boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n            boxesyolo = []\n            for box in boxes:\n                x1, y1, x2, y2 = box\n                xc, yc, w, h = 0.5*x1/im_w+0.5*x2/im_w, 0.5*y1/im_h+0.5*y2/im_h, abs(x2/im_w-x1/im_w), abs(y2/im_h-y1/im_h)\n                boxesyolo.append([0, xc, yc, w, h])\n            self.labels[i] = np.array(boxesyolo)\n        \n        self.image_dir = image_dir\n        self.transforms = transforms\n        \n        self.mosaic = False\n        self.augment = True\n\n    def __getitem__(self, index: int):\n\n        #img, labels = load_mosaic(self, index)\n        self.mosaic = True\n        if random.randint(0,1) ==0:\n            self.mosaic = False\n        if self.mosaic:\n            # Load mosaic\n            img, labels = load_mosaic(self, index)\n            shapes = None\n\n        else:\n            # Load image\n            img, (h0, w0), (h, w) = load_image(self, index)\n\n            # Letterbox\n            shape = self.img_size  # final letterboxed shape\n            img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)\n            shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling\n\n            # Load labels\n            labels = []\n            x = self.labels[index]\n            if x.size > 0:\n                # Normalized xywh to pixel xyxy format\n                labels = x.copy()\n                labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0]  # pad width\n                labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1]  # pad height\n                labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0]\n                labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1]\n        \n        if self.augment:\n            # Augment imagespace\n            if not self.mosaic:\n                img, labels = random_affine(img, labels,\n                                            degrees=0,\n                                            translate=0,\n                                            scale=0,\n                                            shear=0)\n\n            # Augment colorspace\n            augment_hsv(img, hgain=0.0138, sgain= 0.678, vgain=0.36)\n        \n        d = {}\n        d['boxes'] = torch.from_numpy(labels[:,1:].astype(np.float32))\n        d['labels'] = torch.ones((labels[:,0].shape[0],), dtype=torch.int64)\n        \n        \n        return torch.from_numpy(torch.from_numpy(img).permute(2, 0, 1).numpy().astype(np.float32) / 255.0), d\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    \n    \nclass WheatTestDataset(Dataset):\n\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, index: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\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 self.transforms:\n            sample = {\n                'image': image,\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n\n        return image, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentation functions\nThese augmentations are inspired by [CutMix](https://arxiv.org/abs/1905.04899) and [MixMatch](https://arxiv.org/abs/1905.02249).<br>\n[Source](https://github.com/ultralytics/yolov3)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_image(self, index):\n    # loads 1 image from dataset, returns img, original hw, resized hw\n    image_id = self.image_ids[index]\n    imgpath = f'{DIR_INPUT}/global-wheat-detection/train'\n    img = cv2.imread(f'{imgpath}/{image_id}.jpg', cv2.IMREAD_COLOR)\n    \n    assert img is not None, 'Image Not Found ' + imgpath\n    h0, w0 = img.shape[:2]  # orig hw\n    return img, (h0, w0), img.shape[:2]  # img, hw_original, hw_resized\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_mosaic(self, index):\n    # loads images in a mosaic\n\n    labels4 = []\n    s = self.img_size\n    xc, yc = [int(random.uniform(s * 0.5, s * 1.5)) for _ in range(2)]  # mosaic center x, y\n    indices = [index] + [random.randint(0, len(self.labels) - 1) for _ in range(3)]  # 3 additional image indices\n    for i, index in enumerate(indices):\n        # Load image\n        img, _, (h, w) = load_image(self, index)\n\n        # place img in img4\n        if i == 0:  # top left\n            img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles\n            x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)\n            x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)\n        elif i == 1:  # top right\n            x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc\n            x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h\n        elif i == 2:  # bottom left\n            x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)\n            x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)\n        elif i == 3:  # bottom right\n            x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)\n            x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)\n\n        img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]\n        padw = x1a - x1b\n        padh = y1a - y1b\n\n        # Labels\n        x = self.labels[index]\n        labels = x.copy()\n        if x.size > 0:  # Normalized xywh to pixel xyxy format\n            labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw\n            labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + padh\n            labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padw\n            labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + padh\n        labels4.append(labels)\n\n    # Concat/clip labels\n    if len(labels4):\n        labels4 = np.concatenate(labels4, 0)\n        # np.clip(labels4[:, 1:] - s / 2, 0, s, out=labels4[:, 1:])  # use with center crop\n        np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:])  # use with random_affine\n\n    # Augment\n    # img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)]  # center crop (WARNING, requires box pruning)\n    img4, labels4 = random_affine(img4, labels4,\n                                  degrees=1.98 * 2,\n                                  translate=0.05 * 2,\n                                  scale=0.05 * 2,\n                                  shear=0.641 * 2,\n                                  border=-s // 2)  # border to remove\n\n    return img4, labels4\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def random_affine(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, border=0):\n    # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))\n    # https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4\n\n    if targets is None:  # targets = [cls, xyxy]\n        targets = []\n    height = img.shape[0] + border * 2\n    width = img.shape[1] + border * 2\n\n    # Rotation and Scale\n    R = np.eye(3)\n    a = random.uniform(-degrees, degrees)\n    # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations\n    s = random.uniform(1 - scale, 1 + scale)\n    R[:2] = cv2.getRotationMatrix2D(angle=a, center=(img.shape[1] / 2, img.shape[0] / 2), scale=s)\n\n    # Translation\n    T = np.eye(3)\n    T[0, 2] = random.uniform(-translate, translate) * img.shape[0] + border  # x translation (pixels)\n    T[1, 2] = random.uniform(-translate, translate) * img.shape[1] + border  # y translation (pixels)\n\n    # Shear\n    S = np.eye(3)\n    S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)\n    S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)\n\n    # Combined rotation matrix\n    M = S @ T @ R  # ORDER IS IMPORTANT HERE!!\n    if (border != 0) or (M != np.eye(3)).any():  # image changed\n        img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(114, 114, 114))\n\n    # Transform label coordinates\n    n = len(targets)\n    if n:\n        # warp points\n        xy = np.ones((n * 4, 3))\n        xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2)  # x1y1, x2y2, x1y2, x2y1\n        xy = (xy @ M.T)[:, :2].reshape(n, 8)\n\n        # create new boxes\n        x = xy[:, [0, 2, 4, 6]]\n        y = xy[:, [1, 3, 5, 7]]\n        xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T\n\n        # # apply angle-based reduction of bounding boxes\n        # radians = a * math.pi / 180\n        # reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5\n        # x = (xy[:, 2] + xy[:, 0]) / 2\n        # y = (xy[:, 3] + xy[:, 1]) / 2\n        # w = (xy[:, 2] - xy[:, 0]) * reduction\n        # h = (xy[:, 3] - xy[:, 1]) * reduction\n        # xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T\n\n        # reject warped points outside of image\n        xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)\n        xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)\n        w = xy[:, 2] - xy[:, 0]\n        h = xy[:, 3] - xy[:, 1]\n        area = w * h\n        area0 = (targets[:, 3] - targets[:, 1]) * (targets[:, 4] - targets[:, 2])\n        ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16))  # aspect ratio\n        i = (w > 4) & (h > 4) & (area / (area0 * s + 1e-16) > 0.2) & (ar < 10)\n\n        targets = targets[i]\n        targets[:, 1:5] = xy[i]\n\n    return img, targets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):\n    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains\n    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))\n    dtype = img.dtype  # uint8\n\n    x = np.arange(0, 256, dtype=np.int16)\n    lut_hue = ((x * r[0]) % 180).astype(dtype)\n    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)\n    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)\n\n    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)\n    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):\n    # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232\n    shape = img.shape[:2]  # current shape [height, width]\n    if isinstance(new_shape, int):\n        new_shape = (new_shape, new_shape)\n\n    # Scale ratio (new / old)\n    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])\n    if not scaleup:  # only scale down, do not scale up (for better test mAP)\n        r = min(r, 1.0)\n\n    # Compute padding\n    ratio = r, r  # width, height ratios\n    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))\n    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding\n    if auto:  # minimum rectangle\n        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding\n    elif scaleFill:  # stretch\n        dw, dh = 0.0, 0.0\n        new_unpad = new_shape\n        ratio = new_shape[0] / shape[1], new_shape[1] / shape[0]  # width, height ratios\n\n    dw /= 2  # divide padding into 2 sides\n    dh /= 2\n\n    if shape[::-1] != new_unpad:  # resize\n        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)\n    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))\n    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))\n    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border\n    return img, ratio, (dw, dh)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Setting-up Dataloaders","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import random, math\n# Albumentations\ndef get_train_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\n\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\n\ndef get_test_transform():\n    return A.Compose([\n        # A.Resize(512, 512),\n        ToTensorV2(p=1.0)\n    ])\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())\ntest_dataset = WheatTestDataset(test_df, DIR_TEST, get_test_transform())\n\n#train_dataset = WheatDataset(train_df, DIR_TRAIN, get_augumentation(phase='train'))\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=16,\n    shuffle=True,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=4,\n    drop_last=False,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training Code\nWe will be using pretrained pseudo labeling weights for better results.\n<br>\n<br>\nWe will be training for multiple epochs and saving the weights when the validation loss is the least. This will help us in getting the best weights file and avoid overfitting.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WEIGHTS_FILE = \"../input/fork-of-fasterrcnn-pseudo-labeling/fasterrcnn_resnet50_fpn2nd.pth\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_classes = 2  # 1 class (wheat) + 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)\n\nmodel.load_state_dict(torch.load(WEIGHTS_FILE))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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":{"trusted":true},"cell_type":"code","source":"images, targets = 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]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = targets[2]['boxes'].cpu().numpy().astype(np.int32)\nsample = images[2].permute(1,2,0).cpu().numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, 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":"## Hyperparameters","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"model.to(device)\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)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\n\n# optimizer = torch.optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=0.00001)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5)\n\n#lr_scheduler = None\n\nnum_epochs = 10 #Increase it for better results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loss_hist = Averager()\nval_loss_hist = Averager()\nitr = 1\nleast_loss = float('inf')\nfor epoch in range(num_epochs):\n    loss_hist.reset()\n    val_loss_hist.reset()\n    \n    for images, targets in train_data_loader:\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        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    #Validation Step\n    for images, targets in valid_data_loader:\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        val_loss_dict = model(images, targets)\n\n        val_losses = sum(loss for loss in val_loss_dict.values())\n        val_loss_value = val_losses.item()\n\n        val_loss_hist.send(val_loss_value)\n\n    \n    # update the learning rate\n    if val_loss_hist.value<least_loss:\n        least_loss = val_loss_hist.value\n        lval=int(least_loss*1000)/1000\n        torch.save(model.state_dict(), f'fasterrcnn_custom_test_ep{epoch}_loss{lval}.pth')\n        torch.save(model.state_dict(), 'best_weights.pth')\n        \n    else:\n        if lr_scheduler is not None:\n            lr_scheduler.step()\n    print(f\"Epoch #{epoch} train_loss: {loss_hist.value} val_loss: {val_loss_hist.value}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference Code","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WEIGHTS_FILE = 'best_weights.pth'\nmodel.load_state_dict(torch.load(WEIGHTS_FILE))\nmodel.eval()\ncpu_device = torch.device(\"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"detection_threshold = 0.5\nresults = []\n\nfor images, image_ids in test_data_loader:\n\n    images = list(image.to(device) for image in images)\n    outputs = model(images)\n\n    for i, image in enumerate(images):\n\n        boxes = outputs[i]['boxes'].data.cpu().numpy()\n        scores = outputs[i]['scores'].data.cpu().numpy()\n        \n        boxes = boxes[scores >= detection_threshold].astype(np.int32)\n        scores = scores[scores >= detection_threshold]\n        image_id = image_ids[i]\n        \n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n        \n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores)\n        }\n\n        \n        results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample = images[1].permute(1,2,0).cpu().numpy()\nboxes = outputs[1]['boxes'].data.cpu().numpy()\nscores = outputs[1]['scores'].data.cpu().numpy()\n\nboxes = boxes[scores >= detection_threshold].astype(np.int32)\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), 2)\n    \nax.set_axis_off()\nax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"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}