{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"scrolled":false},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport time, random, os, cv2, ast, glob, numba\nfrom numba import jit\nfrom tqdm.autonotebook import tqdm\nfrom pprint import pprint\nimport sys\n\nimport torch, torchvision\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.transforms as T\nfrom torchvision.models.detection.faster_rcnn import AnchorGenerator, FastRCNNPredictor, FasterRCNN\nfrom torchvision.models.detection import fasterrcnn_resnet50_fpn\nfrom torchvision.models import mobilenet_v2, resnet101, vgg19","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_csv = '../input/global-wheat-detection/train.csv'\ntrain_dir = '../input/global-wheat-detection/train/'\ntest_dir = '../input/global-wheat-detection/test/'\n\ndf = pd.read_csv(train_csv)\n# df","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Preprocessing Data"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"def get_dict_list():\n    \n    # Getting a key-value pair to separate training images into\n    # training data and validation data. Used for dataloading\n    # purpose.\n    train_img_dict = {}\n    val_img_dict = {}\n    img = glob.glob(train_dir + '/*.jpg')\n    n = len(img)\n\n    # 95% data for training\n    train_keys = img[:int(0.95*n)]\n    val_keys = img[int(0.95*n):]\n\n    split_dict = {}\n    for key in train_keys:\n        split_dict[key] = 'train'\n    for key in val_keys:\n        split_dict[key] = 'val'\n\n    for i in img:\n        temp = i.split('/')[-1]\n        img_name = temp.split('.')[0]\n        orig_path = train_dir + img_name.split('_')[0] + '.jpg'\n        if (split_dict[orig_path] == 'train'):\n            train_img_dict[img_name] = i\n        else:\n            val_img_dict[img_name] = i\n            \n    # Appending training and validation image paths to \n    # corresponding lists to create separate dataframes.\n    train_img_list = []\n    val_img_list = []\n\n    for i in img:\n        temp = i.split('/')[-1]\n        img_name = temp.split('.')[0]\n        orig_path = train_dir + img_name.split('_')[0] + '.jpg'\n        if (split_dict[orig_path] == 'train'):\n            train_img_list.append(img_name)\n        else:\n            val_img_list.append(img_name)\n            \n    return train_img_dict, val_img_dict, train_img_list, val_img_list","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# Function for creating a separate dataframe for training and validation images\ndef create_df(img_list):\n    \n    image_id = []\n    bbox = []\n    for i in range(len(img_list)):\n\n        for img_id, box in zip(df['image_id'].values, df['bbox'].values):\n\n            if img_list[i] == img_id:\n                image_id.append(img_id)\n                bbox.append(box)\n\n    df_new = pd.DataFrame()\n    df_new['image_id'] = image_id\n    df_new['bbox'] = bbox\n    \n    return df_new","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_img_dict, val_img_dict, train_img_list, val_img_list = get_dict_list()\n\ndf_train = create_df(train_img_list)\ndf_val = create_df(val_img_list)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Evaluation Metrics"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"# BoilerPlate code\nclass AverageMeter(object):\n    \n    # Keeps track of most recent, average, sum, and count of a metric.\n    \n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# AS PER COMPETITION METRIC\n# BoilerPlate Code\niou_thresholds = numba.typed.List()\n\nfor x in [0.5, 0.55, 0.6, 0.65, 0.7, 0.75]:\n    iou_thresholds.append(x)\n\n@jit(nopython=True)\ndef calculate_iou(gt, pr, form='pascal_voc') -> float:\n    \"\"\"Calculates the Intersection over Union.\n\n    Args:\n        gt: (np.ndarray[Union[int, float]]) coordinates of the ground-truth box\n        pr: (np.ndarray[Union[int, float]]) coordinates of the prdected box\n        form: (str) gt/pred coordinates format\n            - pascal_voc: [xmin, ymin, xmax, ymax]\n            - coco: [xmin, ymin, w, h]\n    Returns:\n        (float) Intersection over union (0.0 <= iou <= 1.0)\n    \"\"\"\n    if form == 'coco':\n        gt = gt.copy()\n        pr = pr.copy()\n\n        gt[2] = gt[0] + gt[2]\n        gt[3] = gt[1] + gt[3]\n        pr[2] = pr[0] + pr[2]\n        pr[3] = pr[1] + pr[3]\n\n    # Calculate overlap area\n    dx = min(gt[2], pr[2]) - max(gt[0], pr[0]) + 1\n    \n    if dx < 0:\n        return 0.0\n    \n    dy = min(gt[3], pr[3]) - max(gt[1], pr[1]) + 1\n\n    if dy < 0:\n        return 0.0\n\n    overlap_area = dx * dy\n\n    # Calculate union area\n    union_area = (\n            (gt[2] - gt[0] + 1) * (gt[3] - gt[1] + 1) +\n            (pr[2] - pr[0] + 1) * (pr[3] - pr[1] + 1) -\n            overlap_area\n    )\n\n    return overlap_area / union_area\n\n@jit(nopython=True)\ndef find_best_match(gts, pred, pred_idx, threshold = 0.5, form = 'pascal_voc', ious=None) -> int:\n    \"\"\"Returns the index of the 'best match' between the\n    ground-truth boxes and the prediction. The 'best match'\n    is the highest IoU. (0.0 IoUs are ignored).\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        pred: (List[Union[int, float]]) Coordinates of the predicted box\n        pred_idx: (int) Index of the current predicted box\n        threshold: (float) Threshold\n        form: (str) Format of the coordinates\n        ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n    Return:\n        (int) Index of the best match GT box (-1 if no match above threshold)\n    \"\"\"\n    best_match_iou = -np.inf\n    best_match_idx = -1\n\n    for gt_idx in range(len(gts)):\n        \n        if gts[gt_idx][0] < 0:\n            # Already matched GT-box\n            continue\n        \n        iou = -1 if ious is None else ious[gt_idx][pred_idx]\n\n        if iou < 0:\n            iou = calculate_iou(gts[gt_idx], pred, form=form)\n            \n            if ious is not None:\n                ious[gt_idx][pred_idx] = iou\n\n        if iou < threshold:\n            continue\n\n        if iou > best_match_iou:\n            best_match_iou = iou\n            best_match_idx = gt_idx\n\n    return best_match_idx\n\n@jit(nopython=True)\ndef calculate_precision(gts, preds, threshold = 0.5, form = 'coco', ious=None) -> float:\n    \"\"\"Calculates precision for GT - prediction pairs at one threshold.\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n               sorted by confidence value (descending)\n        threshold: (float) Threshold\n        form: (str) Format of the coordinates\n        ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n    Return:\n        (float) Precision\n    \"\"\"\n    n = len(preds)\n    tp = 0\n    fp = 0\n\n    # for pred_idx, pred in enumerate(preds_sorted):\n    for pred_idx in range(n):\n\n        best_match_gt_idx = find_best_match(gts, preds[pred_idx], pred_idx,\n                                            threshold=threshold, form=form, ious=ious)\n\n        if best_match_gt_idx >= 0:\n            # True positive: The predicted box matches a gt box with an IoU above the threshold.\n            tp += 1\n            # Remove the matched GT box\n            gts[best_match_gt_idx] = -1\n\n        else:\n            # No match\n            # False positive: indicates a predicted box had no associated gt box.\n            fp += 1\n\n    # False negative: indicates a gt box had no associated predicted box.\n    fn = (gts.sum(axis=1) > 0).sum()\n\n    return tp / (tp + fp + fn)\n\n@jit(nopython=True)\ndef calculate_image_precision(gts, preds, thresholds = (0.5, ), form = 'coco') -> float:\n    \"\"\"Calculates image precision.\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n               sorted by confidence value (descending)\n        thresholds: (float) Different thresholds\n        form: (str) Format of the coordinates\n\n    Return:\n        (float) Precision\n    \"\"\"\n    n_threshold = len(thresholds)\n    image_precision = 0.0\n    \n    ious = np.ones((len(gts), len(preds))) * -1\n    # ious = None\n\n    for threshold in thresholds:\n        precision_at_threshold = calculate_precision(gts.copy(), preds, threshold=threshold,\n                                                     form=form, ious=ious)\n        image_precision += precision_at_threshold / n_threshold\n\n    return image_precision\n\ndef calculate_final_score(all_predictions, score_threshold,form):\n    final_scores = []\n    for i in range(len(all_predictions)):\n        gt_boxes = all_predictions[i]['gt_boxes'].copy()\n        pred_boxes = all_predictions[i]['pred_boxes'].copy()\n        scores = all_predictions[i]['scores'].copy()\n        image_id = all_predictions[i]['image_id']\n\n        indexes = np.where(scores>score_threshold)\n        pred_boxes = pred_boxes[indexes]\n        scores = scores[indexes]\n\n        image_precision = calculate_image_precision(gt_boxes, pred_boxes,thresholds=iou_thresholds,form=form)\n        final_scores.append(image_precision)\n\n    return np.mean(final_scores)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualization Functions"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"def get_image(img_dir, dataframe, idx=None, image_id=None):\n    '''\n    Read and output the image in the form of numpy arrays.\n    Args:\n        img_dir: image directory.\n        dataframe: dataframe.\n        idx: index to get the image with ground truth boxes before training.\n        image_id: image ID, to output the image while validating.\n    Returns:\n        image in the form of numpy arrays.\n    '''\n    \n    if image_id is None:\n        img = os.path.join(img_dir, dataframe['image_id'][idx]) + '.jpg'\n    else:\n        img = os.path.join(img_dir, image_id) + '.jpg'\n                           \n    img = cv2.imread(img, cv2.IMREAD_COLOR)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    return img\n\ndef show_image(img_dir, dataframe, boxes_pred=None, show=False, idx=None, image_id=None):\n    '''\n    Function to display the image corresponding to passed arguments.\n    Args:\n        img_dir: image directory.\n        dataframe: dataframe.\n        boxes_pred: predicted box coordinates while validation.\n        show: displays the image if True.\n        idx: index to get the image with ground truth boxes before training.\n        image_id: image ID, to output the image while validating.\n    Returns:\n        if show is True then displays image with ground truth boxes before training\n        if boxes_pred is None or if boxes_pred is given then displays image with ground\n        truth and predicted boxes.\n        if show is False then returns ground truth coordinates and predicted coordinates.\n    '''\n\n    if image_id is not None:\n        image_id = image_id\n    else:\n        image_id = df['image_id'][idx]\n    img = get_image(img_dir, dataframe, idx, image_id)\n    boxes_gt = df[df['image_id'] == image_id]['bbox'].values\n    for box in boxes_gt:\n        box = ast.literal_eval(box) # https://stackoverflow.com/questions/29552950/when-to-use-ast-literal-eval/29556591\n        x, y, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])\n        w = x + w  #x_max = w\n        h = y + h  #y_max = h\n        gt_rect = cv2.rectangle(img, (x,y), (w, h), (0,255,0), 3)\n\n    if boxes_pred is not None:\n        for box in boxes_pred:\n            x, y, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])\n            # w = x + w  #x_max = w\n            # h = y + h  #y_max = h\n            pred_rect = cv2.rectangle(img, (x,y), (w, h), (255,0,0), 2)\n\n    if show:\n        if boxes_pred is None:\n            plt.figure(figsize=(8,8))\n            plt.axis('off')\n            plt.title('Image ID: '+image_id, fontdict={'color':'cyan'})\n            plt.imshow(gt_rect)\n            plt.show()\n\n        else:\n            plt.figure(figsize=(8,8))\n            plt.axis('off')\n            plt.title('Image ID: '+image_id+'       Green: Ground Truth, Box Count: '+str(len(boxes_gt))\n            +'     Red: Predicted, Box Count: '+str(len(boxes_pred)), fontdict={'color':'cyan'})\n            plt.imshow(pred_rect)\n            plt.show()\n\n    else:\n        if boxes_pred is None:\n            return gt_pred\n        else:\n            return gt_rect, pred_rect\n\ndef val_show(gts, dataframe, boxes, image_id):\n    '''\n    Function to select predicted boxes above threshold and passing that\n    as an argument to show_image function.\n    Args:\n        gts: ground truth box coordinates.\n        dataframe: dataframe\n        boxes: predicted boxes\n        image_id: corresponding image ID.\n    Returns:\n        arguments for show_image function.\n    '''\n\n    ious = np.ones((len(gts), len(boxes))) * -1\n    boxes_pred_itr = [] # for all boxes(repetition of boxes)\n    boxes_pred = [] # for unique boxes\n    for pred_idx in range(len(boxes)):\n        best_match_gt_idx = find_best_match(gts, boxes[pred_idx], pred_idx, threshold=0.5, ious=ious)\n        boxes_pred_itr.append(boxes[best_match_gt_idx])\n    \n    # for removing duplicate boxes\n    boxes_pred = list(set(tuple(box) for box in boxes_pred_itr))\n    show_image(train_dir, dataframe, boxes_pred, show=True, image_id=image_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Let's visualize a image from training directory.\nshow_image(train_dir, df, show=True, idx=365)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"class dataset(Dataset):\n    def __init__(self, df, train=True, transforms=None):\n        \n        self.df = df\n        self.train = train\n        self.image_ids = self.df['image_id'].unique()\n        self.transforms = transforms\n\n    def __getitem__(self, idx):\n        \n        image_id = self.image_ids[idx]\n        if self.train:\n            img_path = train_img_dict.get(image_id)\n        else:\n            img_path = val_img_dict.get(image_id)\n\n        img = cv2.imread(img_path, cv2.IMREAD_COLOR)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)\n        # img = np.array(Image.open(img_path).convert(\"RGB\")).astype(np.float32)\n        img = img / 255.0\n\n        boxes = np.int64(np.array([ast.literal_eval(box) for box in self.df[self.df['image_id'] == image_id]['bbox'].values]))\n        boxes[:,2] += boxes[:,0]\n        boxes[:,3] += boxes[:,1]\n\n        target = {}\n        target['boxes'] = torch.as_tensor(boxes, dtype = torch.int64)\n        target['labels'] = torch.ones((len(boxes,)), dtype = torch.int64)\n        target['iscrowd'] = torch.zeros((len(boxes,)), dtype = torch.int64)\n        target['area'] = torch.as_tensor(((boxes[:,3] - boxes[:,1]) * (boxes[:,2] - boxes[:,0])), dtype = torch.float32)\n        target['image_id'] = torch.tensor([idx])\n        \n        if self.transforms is not None:\n            img = self.transforms(img)\n            \n        return img, target, image_id\n    \n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"# https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html\ndef get_model(num_classes, backbone = None):\n    '''\n    Model function to output network according to arguments.\n    Args:\n        num_classes: number of classes(total_classes+1 for background)\n        backbone: to design network with other backbone, default backbone\n                  of faster RCNN is resnet50.\n    Returns:\n        model.\n    '''\n    \n    if backbone == 'mobile_net': \n        net = mobilenet_v2(pretrained = True)\n        backbone_ft = net.features\n        backbone_ft.out_channels = 1280\n        \n    elif backbone == 'vgg19':\n        net = vgg19(pretrained = True)\n        backbone_ft = net.features\n        backbone_ft.out_channels = 512 \n    \n    # https://stackoverflow.com/questions/58362892/resnet-18-as-backbone-in-faster-r-cnn\n    elif backbone == 'resnet101':\n        net = resnet101(pretrained = True)\n        modules = list(net.children())[:-1]\n        backbone_ft = nn.Sequential(*modules)\n        backbone_ft.out_channels = 2048\n        \n    if backbone is None:\n        \n        model = fasterrcnn_resnet50_fpn(pretrained = True)\n        in_features = model.roi_heads.box_predictor.cls_score.in_features\n        # print(in_features) = 1024\n        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n        return model\n    \n    else:\n\n        anchor_gen = AnchorGenerator(sizes=((32, 64, 128),))\n        # featmap_names = [0] gives list index out of range error.\n        roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names = ['0'],\n                                                        output_size = 7,\n                                                        sampling_ratio = 2)\n        model = FasterRCNN(backbone_ft,\n                           num_classes,\n                           rpn_anchor_generator = anchor_gen,\n                           box_roi_pool = roi_pooler)\n        \n        return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training and Validation Function"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"def train_fn(dataloader, epoch, model, optimizer, device):\n    '''\n    Training function.\n    Args:\n        dataloader: for loading training data batch-wise.\n        model: network architecture for training.\n        optimizer: optimizer used for gradient descent.\n        device: computation device for training.\n    Returns:\n        loss after every epoch.\n    '''\n    \n    model.train()  # training mode enables dropout\n    \n    loss = AverageMeter()  # loss update/reset\n    batch_time = AverageMeter()  # forward prop. + back prop. time\n    data_time = AverageMeter()  # data loading time\n    \n    start = time.time()\n    \n    loader = tqdm(dataloader, total = len(dataloader))\n    for step, (images, targets, image_id) in enumerate(loader):\n        \n        # take the list of images and targets to feed the network\n        images = [image.to(device, dtype=torch.float32) for image in images]\n        targets = [{k: v.to(device) for k,v in target.items()} for target in targets]\n        data_time.update(time.time() - start)\n\n        # forward + backward + optimize\n        loss_dict = model(images, targets)\n        # loss_dict: {'loss_classifier': tensor(0.6591, device='cuda:0', grad_fn=<NllLossBackward>),\n                    # 'loss_box_reg': tensor(0.7574, device='cuda:0', grad_fn=<DivBackward0>),\n                    # 'loss_objectness': tensor(0.6313, device='cuda:0',\n                    #                           grad_fn=<BinaryCrossEntropyWithLogitsBackward>),\n                    #  'loss_rpn_box_reg': tensor(0.1344, device='cuda:0', grad_fn=<DivBackward0>)}\n        losses = sum(loss_ind for loss_ind in loss_dict.values())\n        \n        optimizer.zero_grad()  # zero the parameter gradients\n        losses.backward()\n        optimizer.step()\n        \n        batch_time.update(time.time() - start)\n        # Update loss of after every batch.\n        loss.update(losses.item(), BATCH_SIZE)\n        \n        start = time.time()\n        \n        if step % ITER_STEP == 0:\n            \n            print('Epoch: [{0}][{1}/{2}]\\t'\n                  'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n                  'Data Time {data_time.val:.3f} ({data_time.avg:.3f})\\t'\n                  'Loss: {loss.val:.4f} ({loss.avg:.4f})\\t'.format(epoch, step, len(dataloader),\n                                                                  batch_time=batch_time,\n                                                                  data_time=data_time, loss=loss))\n        # To check the loss real-time while iterating.\n        loader.set_postfix(loss=loss.avg)\n\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def val_fn(dataloader, model, device, display_random=False, show_img_num=None):\n    '''\n    Validation function with epoch wise visualization.\n    Args:\n        dataloader: to load the data batch-wise.\n        model: trained model for validation.\n        device: device used for computation.\n        display_random: for visualiztion of random images in every epoch.\n        show_img_num: to visualize a particular image, number between 0 and\n                      batch size.\n    Returns:\n        Visualizations and a list of dictionary consisting of predicted box\n        coordinates, corresponding scores, ground truth box coordinates and\n        image ID.\n    '''\n    model.eval()\n    predictions = []\n    \n    with torch.no_grad():\n        \n        loader = tqdm(dataloader, total=len(dataloader))\n        for step, (images, targets, image_id) in enumerate(loader):\n            \n            images = [image.to(device, dtype=torch.float32) for image in images]\n            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n            output = model(images)\n\n            for i in range(len(images)):\n                # tensor.detach() creates a tensor that shares storage with tensor\n                # that does not require grad. It detaches the output from the computational\n                # graph. So no gradient will be backpropagated along this variable.\n                boxes = output[i]['boxes'].detach().cpu().numpy()\n                scores = output[i]['scores'].detach().cpu().numpy()\n\n                # boxes_itr = boxes\n                predictions.append({\n                    'pred_boxes': (boxes).astype(int),\n                    'scores': scores,\n                    'gt_boxes': (targets[i]['boxes'].cpu().numpy()).astype(int),\n                    'image_id': image_id[i],\n                })\n\n                if display_random:\n                    itr = np.random.randint(low=0, high=BATCH_SIZE-1, size=1)\n                else:\n                    itr = show_img_num\n\n                if step%15==0 and i==itr:\n\n                    gts = (targets[i]['boxes'].cpu().numpy()).astype(int)\n                    val_show(gts, df_val, boxes, image_id[i])\n                \n    return predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"# collate_fn is called with a list of data samples at each time.\n# It is expected to collate the input samples into a batch for\n# yielding from the data loader iterator.\n# https://discuss.pytorch.org/t/how-to-use-collate-fn/27181\ndef collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Engine(main function)"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def engine(device, model_path=None, init_epoch=None, resume=False):\n    '''\n    Main funtion to train and validate.\n    Args:\n        device: device for computation.\n        model_path: path of saved model.\n        init_epoch: initial epoch to resume training from.\n        resume: to resume training from last epoch.\n    Return:\n        final_score\n    '''\n    \n    final_score = []\n    best_score = 0\n    \n    # Custom DataLoaders\n    train_dataset = dataset(df_train, transforms=T.Compose([T.ToTensor()]))\n    valid_dataset = dataset(df_val, train=False, transforms=T.Compose([T.ToTensor()]))\n\n    train_loader = DataLoader(train_dataset,\n                              BATCH_SIZE,\n                              shuffle=False,\n                              num_workers=8,\n                              collate_fn=collate_fn)\n    val_loader = DataLoader(valid_dataset,\n                            BATCH_SIZE,\n                            shuffle=False,\n                            num_workers=8,\n                            collate_fn=collate_fn, )\n    \n    if resume:\n        model = torch.load(model_path)\n        init_epoch = init_epoch\n    else:\n        model = get_model(2)\n        init_epoch = 0\n    model.to(device)  # loading model on GPU\n\n    params = [p for p in model.parameters() if p.requires_grad]\n    optimizer = torch.optim.SGD(params, lr=0.002, momentum=0.9, weight_decay=0.0007)\n\n    for epoch in range(init_epoch, EPOCHS):\n        '''\n        Call the train function then validation function to take a look on how\n        model is performed in that epoch. Output of val_fn, prediction will be\n        given to evaluation metrics for getting score.\n        '''\n        train_loss = train_fn(train_loader, epoch, model, optimizer, device)\n        prediction = val_fn(val_loader, model, device, display_random=True)\n        valid_score = calculate_final_score(prediction, 0.5, 'pascal_voc')\n\n        if valid_score > best_score:\n                best_score = valid_score\n                torch.save(model.state_dict(), f'frcnn_best_{epoch}.pth')\n#                 torch.save(model, f'frcnn_best_model_epoch_{epoch}') \n        final_score.append([best_score, epoch])\n        \n    return final_score","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"%%time\n\nBATCH_SIZE = 8\nITER_STEP = 100\nEPOCHS = 10\nmodel_path = None\nINIT_EPOCH = None\nDEVICE = torch.device('cuda')\n\nfinal_score = engine(DEVICE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def get_best_epoch():\n    pprint(final_score)\n\n    # Get the epoch number for weights with best score.\n    max_score = -1\n    for score in final_score:\n        if score[0] > max_score:\n            max_score = score[0]\n            epoch = score[1]\n    print('Best model found in Epoch {}'.format(epoch))\n    \n    return epoch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epoch = get_best_epoch()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Evaluation and Submission"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"test_df = pd.read_csv('../input/global-wheat-detection/sample_submission.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# Dataset class for evaluation.\nclass eval_dataset(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        img = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)\n        img /= 255.0\n\n        if self.transforms is not None:\n            img = self.transforms(img)\n\n        return img, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# Custom dataloader for test data.\ntest_dataset = eval_dataset(test_df, test_dir, transforms=T.Compose([T.ToTensor()]))\ntest_loader = DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=8, drop_last=False, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":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,"scrolled":false},"cell_type":"code","source":"def evaluate(device, model):\n    \n    model.eval()\n    detection_threshold = 0.5\n    results = []\n\n    for images, image_ids in test_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            img = get_image(test_dir, test_df, image_id=image_id)\n\n            for box in boxes:\n                x, y, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])\n                rect = cv2.rectangle(img, (x,y), (w,h), (220, 0, 0), 3)\n\n            plt.figure(figsize=(8,8))\n            plt.axis('off')\n            plt.title('Image ID: '+image_id, fontdict={'color':'cyan'})\n            plt.imshow(rect)\n            plt.show()\n\n            result = {\n                'image_id': image_id,\n                'PredictionString': format_prediction_string(boxes, scores)\n            }\n\n            results.append(result)\n            \n    return results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"DEVICE = torch.device('cpu')\nweight_path = f'./frcnn_best_{epoch}.pth'\n\nmodel = get_model(2)\nmodel.load_state_dict(torch.load(weight_path))\n\nresults = evaluate(DEVICE, model = model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":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":"test_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"References:  \nhttps://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train  \nhttps://towardsdatascience.com/understanding-pytorch-with-an-example-a-step-by-step-tutorial-81fc5f8c4e8e"}],"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}