{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Installs and Imports","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!cp '../input/weightedboxesfusion' . -r\n!pip install --no-deps './weightedboxesfusion' > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport ast\nimport numba\nimport re\nimport gc\nimport cv2\nimport time\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom glob import glob\nfrom shutil import copy\nfrom os.path import join, exists\nfrom numba import jit\nfrom typing import List, Union, Tuple\nfrom tqdm import tqdm\nfrom collections import OrderedDict\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam\nfrom torch import LongTensor as LongTensor\nfrom torch import FloatTensor as FloatTensor\nfrom torch.utils.data import Dataset, DataLoader, WeightedRandomSampler\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nimport torchvision\nfrom torchvision.models import resnet\nfrom torchvision.models._utils import IntermediateLayerGetter\nfrom torchvision.models.detection.faster_rcnn import TwoMLPHead, FastRCNNPredictor\nfrom torchvision.models.detection.rpn import AnchorGenerator, RPNHead, RegionProposalNetwork\nfrom torchvision.models.detection.roi_heads import RoIHeads\nfrom torchvision.models.detection.transform import GeneralizedRCNNTransform\nfrom torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork, LastLevelMaxPool\nfrom torchvision.ops import MultiScaleRoIAlign\nfrom torchvision.ops.misc import FrozenBatchNorm2d\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom sklearn.cluster import AgglomerativeClustering\nfrom ensemble_boxes import *\n\nnp.random.seed(123)\ntorch.manual_seed(123)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"N_CLUSTERS = 1\nPL_ROUNDS = 5\nPL_EPOCHS = 5\ndetection_threshold = 0.5\nimage_size = 1024","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_path = '/kaggle/input/global-wheat-detection'\noriginal_train_path = '/kaggle/input/original-train/original_train'\n# weights_path = '/kaggle/input/1024x1024-1/weights_e149.pth'\nweights_path = '/kaggle/input/manual-weights-7/e152.pth'\ntrain_path = join(dataset_path, 'train')\ntest_path = join(dataset_path, 'test')\ntrain_df = pd.read_csv(join(original_train_path, 'train.csv'))\ntrain_df['image_path'] = [join(original_train_path, s, i + '.jpeg') for i, s in train_df[['image_id', 'source']].values]\ntest_df = pd.read_csv(join(dataset_path, 'sample_submission.csv'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Split validation from training\n# Take 30 examples from each source\n# valid_ids = []\n# for source in pd.unique(train_df['source']):\n#     if source == 'test': continue\n#     source_ids = train_df[train_df['source'] == source]['image_id'].values\n#     valid_ids.extend(list(np.random.choice(source_ids, 50, replace=False)))\n\n# valid_ids = [p.split('/')[-1].split('.')[0] for p in glob('/kaggle/input/global-wheat-detection/test/*')[::10]]\n# print(valid_ids)\n# print(len(valid_ids))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"id_to_path = {}\nfor idx, row in train_df.iterrows():\n    image_id = row['image_id']\n    if image_id in id_to_path.keys(): continue\n    id_to_path[image_id] = row['image_path']\n\nfor idx, row in test_df.iterrows():\n    image_id = row['image_id']\n    if image_id in id_to_path.keys(): continue\n    id_to_path[image_id] = join(test_path, image_id + '.jpg')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(list(train_df.columns))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset functions","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatTrainDataset(Dataset):\n    def __init__(self, dataframe, transforms=None, test=False):\n        super().__init__()\n        self.df = dataframe\n        self.image_ids = pd.unique(dataframe['image_id'])\n        self.image_paths = [id_to_path[image_id] for image_id in self.image_ids]\n        self.length = len(self.image_ids)\n        self.transforms = transforms\n        self.test = test\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image_path = self.image_paths[index]\n        image = self.load_image(image_path) \n        boxes = self.load_boxes(image_id)\n\n        # if not self.test and random.random() > 0.5:\n        #     image, boxes = self.cutmix_image_and_boxes(image, boxes)\n\n        # there is only one class\n        labels = torch.ones((boxes.shape[0],), dtype=torch.int64)\n\n        target = {'boxes': boxes, 'labels': labels, 'image_id': torch.tensor([index])}\n\n        if self.transforms:\n            sample = {'image': image, 'bboxes': target['boxes'], 'labels': labels}\n            sample = self.transforms(**sample)\n            image, boxes = sample['image'], sample['bboxes']\n            boxes = self.filter_boxes(boxes)\n            if len(boxes):\n                target['boxes'] = torch.stack([torch.tensor(box, dtype=torch.float32) for box in zip(*boxes)]).permute(1, 0)\n            else:\n                return self.__getitem__(np.random.randint(self.length))\n                # target['boxes'] = torch.zeros((0, 4), dtype=torch.float32)\n                # target['labels'] = torch.zeros(0, dtype=torch.int64)\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.length\n\n    def load_image(self, 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\n        return image\n\n    def load_boxes(self, image_id):\n        records = self.df[self.df['image_id'] == image_id]\n        if 'x1' in records.columns and 'y1' in records.columns:\n            return records[['x', 'y', 'x1', 'y1']].values\n        else:\n            boxes = records[['x', 'y', 'w', 'h']].values\n            boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n            boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n            return boxes\n\n    def filter_boxes(self, boxes):\n        min_length = 13\n        min_area = 400\n        max_area = 145360\n        max_length_ratio = 18\n        min_length_ratio = 1.0/max_length_ratio\n\n        boxes_out = []\n        for box in boxes:\n            x, y, x1, y1 = box\n            w = x1 - x\n            if w < min_length: continue\n            h = y1 - y\n            if h < min_length: continue\n            area = w * h\n            if area < 400 or area > 145360: continue\n            length_ratio = w / h\n            if length_ratio < min_length_ratio or length_ratio > max_length_ratio: continue\n            boxes_out.append(box)\n\n        return boxes_out\n\n    def get_sample_weights(self):\n        weights = []\n        for image_id in self.image_ids:\n            w, h, source = self.df[self.df['image_id'] == image_id][['width', 'height', 'source']].values[0]\n            if w == 3072:\n                # 2x3\n                weight = 8\n            elif h == 2048:\n                # 2x2\n                weight = 5\n            elif w == 2048:\n                # 1x2\n                weight = 2.5\n            else:\n                # 1x1\n                weight = 1\n            if source == 'test': weight *= 10\n            weights.append(weight)\n        return np.array(weights)\n\n    \nclass WheatTestDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, image_size=1024, onfly=False, tta=True):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.image_size = image_size\n        self.onfly = onfly\n        self.tta = tta\n#         if onfly:\n#             self.load_all_images()\n\n#     def load_all_images(self):\n#         self.images = []\n#         for image_id in self.image_ids:\n            \n        \n    def load_image(self, image_id):\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        image = cv2.resize(image, (self.image_size, self.image_size))\n        image = image.astype(np.float32)\n        image /= 255.0\n        return image\n        \n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = self.load_image(image_id)\n        if self.tta:\n            all_images = apply_tta(image)\n            return all_images, image_id\n        else:\n            return image, image_id\n    \n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n\ndef collate_fn(batch): return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Transforms & Ensembling","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def run_wbf(predictions, image_size=1024, iou_thr=0.4, skip_box_thr=0.7, weights=None):\n    boxes = [pred['boxes'].data.cpu().numpy() for pred in predictions]\n    boxes = revert_tta(boxes, image_size)\n    boxes = [box_set/(image_size-1) for box_set in boxes]\n    scores = [pred['scores'].data.cpu().numpy() for pred in predictions]\n    labels = [np.ones(pred['scores'].shape[0]) for pred in predictions]\n    boxes, scores, labels = weighted_boxes_fusion(boxes, scores, labels, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n    boxes = boxes*(image_size-1)\n    return boxes, scores, labels\n\n# def ensemble_wbf(boxes, scores, image_size=1024, iou_thr=0.35, skip_box_thr=0.65, weights=None):\n#     boxes = [box_set/(image_size-1) for box_set in boxes]\n#     labels = [np.ones(len(score)) for score in scores]\n#     boxes, scores, labels = weighted_boxes_fusion(boxes, scores, labels, weights=None, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n#     boxes = boxes*(image_size-1)\n#     return boxes, scores, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose(\n        [\n            A.RandomCrop(1024, 1024),\n            A.RandomSizedCrop(min_max_height=(9, 800), height=1024, width=1024, p=0.5),\n            A.OneOf([\n                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.9),\n                A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.9),\n            ], p=0.9),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            # A.Resize(height=512, width=512, p=1.0),\n            ToTensorV2(p=1.0),\n        ],\n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n\n\ndef get_valid_transforms():\n    return A.Compose(\n        [\n            A.RandomCrop(1024, 1024),\n#             A.RandomSizedCrop(min_max_height=(800, 800), height=1024, width=1024, p=0.5),\n#             A.HorizontalFlip(p=0.5),\n#             A.VerticalFlip(p=0.5),\n            ToTensorV2(p=1.0),\n        ],\n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n\ndef apply_tta(image):\n    \"\"\"Apply Test Time Augmentation (TTA)\"\"\"\n    flipV_sample = flip_v_image(image)\n    flipH_sample = flip_h_image(image)\n    flipVH_sample = flip_h_image(flip_v_image(image))\n    rot1_sample = np.rot90(image)\n    rot2_sample = np.rot90(rot1_sample)\n    rot3_sample = np.rot90(rot2_sample)\n    zoom = zoom_image(image)\n    zoom_rot1 = np.rot90(zoom)\n    zoom_rot2 = np.rot90(zoom_rot1)\n    zoom_rot3 = np.rot90(zoom_rot2)\n    return [image, flipV_sample, flipH_sample, flipVH_sample, rot1_sample, rot2_sample, rot3_sample, zoom, zoom_rot1, zoom_rot2, zoom_rot3]\n\ndef revert_tta(boxes, img_size):\n    \"\"\"Undo TTA in order to ensemble predictions\"\"\"\n    sample0, flippedV, flippedH, flippedVH, rot1, rot2, rot3, zoomed, zoomed_rot1, zoomed_rot2, zoomed_rot3 = boxes\n    sample1 = flip_v_boxes(flippedV, img_size)\n    sample2 = flip_h_boxes(flippedH, img_size)\n    sample3 = flip_v_boxes(flip_h_boxes(flippedVH, img_size), img_size)\n    sample4 = rotate_boxes(rot1, img_size, 3)\n    sample5 = rotate_boxes(rot2, img_size, 2)\n    sample6 = rotate_boxes(rot3, img_size, 1)\n    sample7 = unzoom_boxes(zoomed)\n    sample8 = unzoom_boxes(rotate_boxes(zoomed_rot1, img_size, 3))\n    sample9 = unzoom_boxes(rotate_boxes(zoomed_rot2, img_size, 2))\n    sample10 = unzoom_boxes(rotate_boxes(zoomed_rot3, img_size, 1))\n    return [sample0, sample1, sample2, sample3, sample4, sample5, sample6, sample7, sample8, sample9, sample10]\n\n\n# def apply_tta(image):\n#     \"\"\"Apply Test Time Augmentation (TTA)\"\"\"\n#     transposed = np.swapaxes(image, 0, 1)\n#     flipV_sample = flip_v_image(image)\n#     flipH_sample = flip_h_image(image)\n#     flipVH_sample = flip_h_image(flip_v_image(image))\n#     rot1_sample = np.rot90(image)\n#     rot2_sample = np.rot90(rot1_sample)\n#     rot3_sample = np.rot90(rot2_sample)\n#     zoom = zoom_image(image)\n#     zoom_rot1 = np.rot90(zoom)\n#     zoom_rot2 = np.rot90(zoom_rot1)\n#     zoom_rot3 = np.rot90(zoom_rot2)\n#     flipV_zoom = flip_v_image(zoom)\n#     flipH_zoom = flip_h_image(zoom)\n#     flipVH_zoom = flip_h_image(flip_v_image(zoom))\n#     trans_zoom = np.swapaxes(zoom, 0, 1)\n#     return [image, transposed, flipV_sample, flipH_sample, flipVH_sample, rot1_sample, rot2_sample, rot3_sample, zoom, zoom_rot1, zoom_rot2, zoom_rot3, flipV_zoom, flipH_zoom, flipVH_zoom, trans_zoom]\n\n# def revert_tta(boxes, img_size):\n#     \"\"\"Undo TTA in order to ensemble predictions\"\"\"\n#     sample0, transposed, flippedV, flippedH, flippedVH, rot1, rot2, rot3, zoomed, zoomed_rot1, zoomed_rot2, zoomed_rot3, zoomed_flipV, zoomed_flipH, zoomed_flipVH, zoomed_trans = boxes\n#     sample1 = transpose_boxes(transposed)\n#     sample2 = flip_v_boxes(flippedV, img_size)\n#     sample3 = flip_h_boxes(flippedH, img_size)\n#     sample4 = flip_v_boxes(flip_h_boxes(flippedVH, img_size), img_size)\n#     sample5 = rotate_boxes(rot1, img_size, 3)\n#     sample6 = rotate_boxes(rot2, img_size, 2)\n#     sample7 = rotate_boxes(rot3, img_size, 1)\n#     sample8 = unzoom_boxes(zoomed)\n#     sample9 = unzoom_boxes(rotate_boxes(zoomed_rot1, img_size, 3))\n#     sample10 = unzoom_boxes(rotate_boxes(zoomed_rot2, img_size, 2))\n#     sample11 = unzoom_boxes(rotate_boxes(zoomed_rot3, img_size, 1))\n#     sample12 = unzoom_boxes(flip_v_boxes(zoomed_flipV, img_size))\n#     sample13 = unzoom_boxes(flip_h_boxes(zoomed_flipH, img_size))\n#     sample14 = unzoom_boxes(flip_h_boxes(flip_v_boxes(zoomed_flipVH, img_size), img_size))\n#     sample15 = unzoom_boxes(transpose_boxes(transposed))\n#     return [sample0, sample1, sample2, sample3, sample4, sample5, sample6, sample7, sample8, sample9, sample10]\n\n# def apply_tta(image):\n#     \"\"\"Apply Test Time Augmentation (TTA)\"\"\"\n#     flipV_sample = flip_v_image(image)\n#     return [image, flipV_sample]\n\n# def revert_tta(boxes, img_size):\n#     \"\"\"Undo TTA in order to ensemble predictions\"\"\"\n#     sample0, flippedV = boxes\n#     sample1 = flip_v_boxes(flippedV, img_size)\n#     return [sample0, sample1]\n\ndef zoom_image(image):\n    zoom = np.zeros_like(image)\n    zoom[100:900, 100:900] = cv2.resize(image, (800, 800))\n    return zoom.astype(np.float32)\n\ndef unzoom_boxes(boxes):\n    return (boxes - 100) * 1.28\n\ndef flip_v_image(image):\n    return np.flip(image, axis=0)\n    \ndef flip_h_image(image):\n    return np.flip(image, axis=1)\n\ndef rotate_boxes(boxes, img_size, k):\n    for _ in range(k):\n        x0 = boxes[:, 1]\n        y0 = img_size - boxes[:, 2]\n        x1 = boxes[:, 3]\n        y1 = img_size - boxes[:, 0]\n        boxes = np.stack([x0, y0, x1, y1], axis=1)\n    return boxes\n\ndef flip_v_boxes(boxes, img_size):\n    y0 = img_size - boxes[:, 3]\n    y1 = img_size - boxes[:, 1]\n    boxes[:, 1] = y0\n    boxes[:, 3] = y1\n    return boxes\n    \ndef flip_h_boxes(boxes, img_size):\n    x0 = img_size - boxes[:, 2]\n    x1 = img_size - boxes[:, 0]\n    boxes[:, 0] = x0\n    boxes[:, 2] = x1\n    return boxes\n\ndef transpose_boxes(boxes):\n    x0, y0, x1, y1 = boxes.T\n    boxes[:, 0] = y0\n    boxes[:, 1] = x0\n    boxes[:, 2] = y1\n    boxes[:, 3] = x1\n    return boxes\n\ndef tensor_transform(sample):\n    transform = A.Compose([ToTensorV2(p=1.0)], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))\n    sample = transform(**sample)\n    return sample\n\ndef resize_transform(sample, image_size=1024):\n    transform = A.Compose([A.Resize(height=image_size, width=image_size, p=1.0)], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))\n    sample = transform(**sample)\n    return sample","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Metric Definition","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"@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\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\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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Network","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_resnet():\n    backbone = resnet.__dict__['resnet50'](pretrained=False, norm_layer=FrozenBatchNorm2d)\n    for name, parameter in backbone.named_parameters():\n        if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:\n            parameter.requires_grad_(False)\n    return backbone\n\nclass MyFasterRCNN(nn.Module):\n    def __init__(self):\n        super(MyFasterRCNN, self).__init__()\n        image_mean = [0.485, 0.456, 0.406]\n        image_std = [0.229, 0.224, 0.225]\n        self.transform = GeneralizedRCNNTransform(1024, 1024, image_mean, image_std)\n        self.backbone = get_resnet()\n        return_layers = {'layer1': '0', 'layer2': '1', 'layer3': '2', 'layer4': '3'}\n        self.body = IntermediateLayerGetter(self.backbone, return_layers=return_layers)\n\n        # Feature Pyramid Network\n        out_channels = 256\n        in_channels_list = [256 * (2 ** i) for i in range(4)]\n        self.fpn = FeaturePyramidNetwork(\n            in_channels_list=in_channels_list,\n            out_channels=out_channels,\n            extra_blocks=LastLevelMaxPool())\n\n        # Regional Proposal Network\n        anchor_sizes = ((32,), (64,), (128,), (256,), (512,))\n        aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)\n        anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios)\n        head = RPNHead(out_channels, anchor_generator.num_anchors_per_location()[0])\n        self.rpn = RegionProposalNetwork(\n            anchor_generator=anchor_generator,\n            head=head,\n            fg_iou_thresh=0.7,\n            bg_iou_thresh=0.3,\n            batch_size_per_image=256,\n            positive_fraction=0.5,\n            pre_nms_top_n=dict(training=2000, testing=1000),\n            post_nms_top_n=dict(training=2000, testing=1000),\n            nms_thresh=0.7)\n\n        # RoI heads\n        representation_size = 512\n        box_roi_pool = MultiScaleRoIAlign(['0', '1', '2', '3'], 7, 2)\n        box_head = TwoMLPHead(out_channels * box_roi_pool.output_size[0] ** 2, representation_size)\n        box_predictor = FastRCNNPredictor(representation_size, num_classes=2)\n        self.roi_heads = RoIHeads(\n            box_roi_pool=box_roi_pool,\n            box_head=box_head,\n            box_predictor=box_predictor,\n            fg_iou_thresh=0.5, bg_iou_thresh=0.5,\n            batch_size_per_image=512, positive_fraction=0.25,\n            bbox_reg_weights=None,\n            score_thresh=0.05,\n            nms_thresh=0.5,\n            detections_per_img=100)\n\n    def forward(self, images, targets=None):\n        images, targets = self.transform(images, targets)\n        features = self.body(images.tensors)\n\n        fpn_features = self.fpn(features)\n        if isinstance(fpn_features, torch.Tensor): fpn_features = OrderedDict([('0', fpn_features)])\n        proposals, proposal_losses = self.rpn(images, fpn_features, targets)\n        detections, detector_losses = self.roi_heads(fpn_features, proposals, images.image_sizes, targets)\n\n        losses = {}\n        losses.update(detector_losses)\n        losses.update(proposal_losses)\n\n        return features, detections, losses\n    \n\nclass ResNetFeaturizer(nn.Module):\n    def __init__(self):\n        super(ResNetFeaturizer, self).__init__()\n        image_mean = [0.485, 0.456, 0.406]\n        image_std = [0.229, 0.224, 0.225]\n        self.transform = GeneralizedRCNNTransform(1024, 1024, image_mean, image_std)\n        self.backbone = get_resnet()\n        self.backbone.load_state_dict(torch.load('/kaggle/input/pretrained-pytorch/resnet50-19c8e357.pth'))\n\n    def forward(self, images, targets=None):\n        images, targets = self.transform(images, targets)\n        features = self.backbone(images.tensors) # images.tensors\n        return features","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model(weights_path):\n    model = MyFasterRCNN()\n    num_classes = 2  # 1 class (wheat) + background\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n    model.load_state_dict(torch.load(weights_path))\n    model = model.cuda()\n    model = model.eval()\n    return model\n\nmodel = get_model(weights_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_test_predictions():\n    for images, image_ids in test_data_loader:\n        images = [torch.tensor(image).cuda().permute(2,0,1) for image in images[0]]\n        features, predictions, loss = model(images)\n        break\n\n    sample = images[0].permute(1,2,0).cpu().numpy()\n    boxes, scores, labels = run_wbf(predictions, image_size=1024)\n    boxes = boxes.astype(np.int32)\n\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\n    for box in boxes:\n        cv2.rectangle(sample,\n                      (box[0], box[1]),\n                      (box[2], box[3]),\n                      (220, 0, 0), 2)\n\n    ax.set_axis_off()\n    ax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset = WheatTestDataset(test_df, test_path, tta=False)\ntest_data_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4, drop_last=False, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# show_test_predictions()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def cluster_test_data():\n    rn_model = ResNetFeaturizer().cuda().eval()\n    \n    feature_vecs = []\n    image_ids = []\n    for images, image_id in tqdm(test_data_loader):\n        inp = [torch.tensor(images[0]).cuda().permute(2,0,1)]\n        fv = rn_model(inp)\n        feature_vecs.append(fv.detach().cpu().numpy())\n        image_ids.append(image_id[0])\n    feature_vecs = np.squeeze(np.array(feature_vecs))\n    print(feature_vecs.shape)\n    \n    agglo = AgglomerativeClustering(n_clusters=N_CLUSTERS)\n    clusters = agglo.fit_predict(feature_vecs)\n    \n#     data = [[i,j] for i,j in ]\n#     cluster_df = pd.DataFrame(clusters, columns=['cluster_id'])\n    id_to_cluster = {k : v for k,v in zip(image_ids, clusters)}\n    \n    del rn_model\n    del inp\n    del fv\n    torch.cuda.empty_cache()\n    \n    return id_to_cluster","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"id_to_cluster = cluster_test_data()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df['cluster_id'] = [id_to_cluster[iid] for iid in test_df['image_id'].values]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_pseudo_labels(model, data_loader, combine_with_train=False):\n    model.eval()\n    data = []\n    for images, image_ids in data_loader:\n        predictions = []\n        for img in images[0]:\n            inp = [torch.tensor(img).cuda().permute(2,0,1)]\n            features, pred, loss = model(inp)\n            predictions.extend(pred)\n\n        boxes, scores, labels = run_wbf(predictions, image_size=image_size)\n        boxes = (boxes).astype(np.int32).clip(min=0, max=int(image_size-1))\n        image_id = image_ids[0]\n\n        for box in boxes:\n            x0, y0, x1, y1 = box\n            w = x1 - x0\n            h = y1 - y0\n            data.append([image_id, 1024, 1024, 'test', x0, y0, w, h, x1, y1, image_id, join(test_path, image_id + '.jpg'), id_to_cluster[image_id]])\n    s1_df = pd.DataFrame(data, columns=['image_id', 'width', 'height', 'source', 'x', 'y', 'w', 'h', 'x1', 'y1', 'image_id_orig', 'image_path', 'cluster_id'])\n    if combine_with_train:\n        s1_df = pd.concat([s1_df, train_df])\n    return s1_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_round(round_num, model):\n    for cluster_id in range(N_CLUSTERS):\n        if round_num == 0:\n            model.load_state_dict(torch.load(weights_path))\n        else:\n            model.load_state_dict(torch.load('round_%d_cluster_%d.pth' % (round_num-1, cluster_id)))\n        \n        optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0001)\n        \n        cluster_df = test_df[test_df['cluster_id'] == cluster_id]\n        cluster_dataset = WheatTestDataset(cluster_df, test_path)\n        cluster_data_loader = DataLoader(cluster_dataset, batch_size=1, shuffle=False, num_workers=4, drop_last=False, collate_fn=collate_fn)\n        cluster_df = get_pseudo_labels(model, cluster_data_loader)\n        train_dataset = WheatTrainDataset(cluster_df, get_train_transforms())\n        train_data_loader = DataLoader(train_dataset, batch_size=8, num_workers=4, drop_last=False, collate_fn=collate_fn)\n        \n        for epoch in range(PL_EPOCHS):\n            print('Round %d | Cluster %d | Epoch %d' % (round_num, cluster_id, epoch))\n            model.train()\n            for step, (images, targets, image_ids) in enumerate(train_data_loader):\n                # Load images/targets to cuda\n                images = [img.cuda() for img in images]\n                targets = [{k: v.cuda() for k, v in l.items()} for l in targets]\n                # Send images through network\n                features, detections, losses = model(images, targets)\n                # Optimizer step\n                optimizer.zero_grad()\n                loss_FS = sum(losses.values())\n                loss_FS.backward()\n                optimizer.step()\n        \n        torch.save(model.state_dict(), 'round_%d_cluster_%d.pth' % (round_num, cluster_id))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for round_num in range(PL_ROUNDS):\n    train_round(round_num, model)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Make Final Predictions","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":"detection_threshold = 0.5\nimage_size = 1024\nresults = []\n\nfor cluster_id in range(N_CLUSTERS):\n    model.load_state_dict(torch.load('round_%d_cluster_%d.pth' % (round_num, cluster_id)))\n    model.eval()\n\n    test_df_cluster = test_df[test_df['cluster_id'] == cluster_id]\n    test_dataset = WheatTestDataset(test_df_cluster, test_path)\n    test_data_loader = DataLoader(test_dataset, batch_size=1, num_workers=4, drop_last=False, collate_fn=collate_fn)\n    \n    for images, image_ids in test_data_loader:\n        image_id = image_ids[0]\n    \n        predictions = []\n        for img in images[0]:\n            inp = [torch.tensor(img).cuda().permute(2,0,1)]\n            features, pred, loss = model(inp)\n            predictions.extend(pred)\n\n        boxes, scores, labels = run_wbf(predictions, image_size=image_size)\n        boxes = boxes.astype(np.int32).clip(min=0, max=int(image_size-1))\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        results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\nsub_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"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}