{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!cp -rp ../input/nvidiaapex/apex-master/* ./","metadata":{"_kg_hide-input":true,"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./","metadata":{"_kg_hide-input":true,"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --no-deps '../input/timm-package/timm-0.1.26-py3-none-any.whl' > /dev/null\n!pip install --no-deps '../input/pycocotools/pycocotools-2.0-cp37-cp37m-linux_x86_64.whl' > /dev/null","metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_kg_hide-output":true,"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.insert(0, \"../input/timm-efficientdet-pytorch\")\nsys.path.insert(0, \"../input/omegaconf\")\nsys.path.insert(0, \"../input/weightedboxesfusion\")\n\nfrom ensemble_boxes import *\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nimport random\nfrom torch.utils.data import Dataset,DataLoader\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom datetime import datetime\nimport time\nimport cv2\nimport gc\nimport os\nfrom matplotlib import pyplot as plt\nfrom effdet import get_efficientdet_config, EfficientDet, DetBenchTrain,DetBenchEval\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nfrom sklearn.model_selection import StratifiedKFold, train_test_split\nfrom effdet.efficientdet import HeadNet\nimport effdet.anchors as anchors\nfrom timm.scheduler import CosineLRScheduler\nimport math\nfrom copy import deepcopy\nfrom apex import amp\nimport torch.nn as nn\n\nSEED = 42\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n\nseed_everything(SEED)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trained_weight_path = '../input/efficidentdet7fullsize/best-checkpoint-054epoch.bin'","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TTA + Create Pseudo label","metadata":{}},{"cell_type":"code","source":"class BaseWheatTTA:\n    \"\"\" author: @shonenkov \"\"\"\n    image_size = 1024\n\n    def augment(self, image):\n        raise NotImplementedError\n    \n    def batch_augment(self, images):\n        raise NotImplementedError\n    \n    def deaugment_boxes(self, boxes):\n        raise NotImplementedError\n\nclass TTAHorizontalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n\n    def augment(self, image):\n        return image.flip(1)\n    \n    def batch_augment(self, images):\n        return images.flip(2)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [1,3]] = self.image_size - boxes[:, [3,1]]\n        return boxes\n\nclass TTAVerticalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return image.flip(2)\n    \n    def batch_augment(self, images):\n        return images.flip(3)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,2]] = self.image_size - boxes[:, [2,0]]\n        return boxes\n    \nclass TTARotate90(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 1, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 1, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = self.image_size - boxes[:, [1,3]]\n        res_boxes[:, [1,3]] = boxes[:, [2,0]]\n        return res_boxes\n\nclass TTACompose(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    def __init__(self, transforms):\n        self.transforms = transforms\n        \n    def augment(self, image):\n        for transform in self.transforms:\n            image = transform.augment(image)\n        return image\n    \n    def batch_augment(self, images):\n        for transform in self.transforms:\n            images = transform.batch_augment(images)\n        return images\n    \n    def prepare_boxes(self, boxes):\n        result_boxes = boxes.copy()\n        result_boxes[:,0] = np.min(boxes[:, [0,2]], axis=1)\n        result_boxes[:,2] = np.max(boxes[:, [0,2]], axis=1)\n        result_boxes[:,1] = np.min(boxes[:, [1,3]], axis=1)\n        result_boxes[:,3] = np.max(boxes[:, [1,3]], axis=1)\n        return result_boxes\n    \n    def deaugment_boxes(self, boxes):\n        for transform in self.transforms[::-1]:\n            boxes = transform.deaugment_boxes(boxes)\n        return self.prepare_boxes(boxes)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from itertools import product\n\ntta_transforms = []\nfor tta_combination in product([TTAHorizontalFlip(), None], \n                               [TTAVerticalFlip(), None],\n                               [TTARotate90(), None]):\n    tta_transforms.append(TTACompose([tta_transform for tta_transform in tta_combination if tta_transform]))","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_tta_predictions(images, nets, score_threshold=0.25):\n    all_predictions = []\n    images = torch.stack(images).float().cuda()\n    for fold_num, net in enumerate(nets):\n        with torch.no_grad():\n            predictions = []\n            for tta_transform in tta_transforms:\n                result = []\n                det = net(tta_transform.batch_augment(images.clone()), torch.tensor([1]*images.shape[0]).float().cuda())\n\n                for i in range(images.shape[0]):\n                    boxes = det[i].detach().cpu().numpy()[:,:4]    \n                    scores = det[i].detach().cpu().numpy()[:,4]\n                    indexes = np.where(scores > score_threshold)[0]\n                    boxes = boxes[indexes]\n                    boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n                    boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\n                    boxes = tta_transform.deaugment_boxes(boxes.copy())\n                    result.append({\n                        'boxes': boxes,\n                        'scores': scores[indexes],\n                    })\n                predictions.append(result)\n            all_predictions.append(predictions)\n            \n    return all_predictions\n\ndef run_wbf(all_predictions, image_index, image_size=1024, iou_thr=0.44, skip_box_thr=0.43, weights=None):\n    for i in range(len(nets)):\n        predictions = all_predictions[i]\n        if i == 0:\n            boxes = [(prediction[image_index]['boxes']/(image_size-1)).tolist()  for prediction in predictions]\n            scores = [prediction[image_index]['scores'].tolist()  for prediction in predictions]\n            labels = [np.ones(prediction[image_index]['scores'].shape[0]).tolist() for prediction in predictions]\n        else:\n            boxes += [(prediction[image_index]['boxes']/(image_size-1)).tolist()  for prediction in predictions]\n            scores += [prediction[image_index]['scores'].tolist()  for prediction in predictions]\n            labels += [np.ones(prediction[image_index]['scores'].shape[0]).tolist() for prediction in predictions]\n            \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","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_model_efficientdet_test(compound_coef:int, weight_path=None,\n                              num_classes=1, image_size=1024, pretrained=False, model_ema=False):\n\n    config = get_efficientdet_config(f'tf_efficientdet_d{compound_coef}')\n    model = EfficientDet(config, pretrained_backbone=pretrained)\n    config.num_classes = num_classes\n    config.image_size = image_size\n    model.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n    \n    checkpoint = torch.load(weight_path)\n    #ck = fix_model_state_dict(checkpoint['state_dict_ema'])\n    if model_ema:\n        model.load_state_dict(checkpoint['state_dict_ema'])\n    else:\n        model.load_state_dict(checkpoint['model_state_dict'])\n\n    \n    del checkpoint\n    gc.collect()\n\n    model = DetBenchEval(model, config)\n    anchors.MAX_DETECTIONS_PER_IMAGE=250\n    model.eval();\n    \n    return model.cuda()\n\n\nnets = [create_model_efficientdet_test(compound_coef=4, \n                                  weight_path='../input/efficidentdet4-fullsize/best-checkpoint-058epoch.bin', \n                                  pretrained=False,\n                                  model_ema=True),\n        create_model_efficientdet_test(compound_coef=5, \n                                 weight_path='../input/efficientdet5-fullsize/best-checkpoint-052epoch.bin', \n                                 pretrained=False),\n        create_model_efficientdet_test(compound_coef=7, \n                                weight_path=trained_weight_path, \n                                pretrained=False,\n                                model_ema=True)]\n#before_pesudo_label_net = create_model_efficientdet_test(compound_coef=7, \n                                #weight_path=trained_weight_path, \n                                #pretrained=False,\n                                #model_ema=True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_test_transforms():\n    return A.Compose([\n            A.Resize(height=1024, width=1024, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.0)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_ROOT_PATH = '../input/global-wheat-detection/test'\n\nclass DatasetRetrieverTest(Dataset):\n\n    def __init__(self, image_ids, transforms=None):\n        super().__init__()\n        self.image_ids = image_ids\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{DATA_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        if self.transforms:\n            sample = {'image': image}\n            sample = self.transforms(**sample)\n            image = sample['image']\n        return image, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset = DatasetRetrieverTest(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{DATA_ROOT_PATH}/*.jpg')]),\n    transforms=get_test_transforms()\n)\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=2,\n    shuffle=False,\n    num_workers=4,\n    drop_last=False,\n    collate_fn=collate_fn\n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict_test(data_loader,nets):\n    results = []\n    testdf_psuedo = []\n\n    for images, image_ids in data_loader:\n        predictions = make_tta_predictions(images=images, nets=nets)\n        for i, image in enumerate(images):\n            boxes, scores, labels = run_wbf(predictions, image_index=i)\n            boxes = boxes.astype(np.int32).clip(min=0, max=1023)\n            image_id = image_ids[i]\n\n            boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n            boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n\n            for box in boxes:\n                result = {\n                    'image_id': image_id,\n                    'source': 'BULLBULL',\n                    'width':1024,\n                    'height':1024,\n                    'x': box[0],\n                    'y': box[1],\n                    'w': box[2],\n                    'h': box[3]        \n                }\n                testdf_psuedo.append(result)\n\n    test_df_pseudo = pd.DataFrame(testdf_psuedo, columns=['image_id', 'width', 'height', 'source', 'x', 'y', 'w', 'h'])\n    return test_df_pseudo","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df_pseudo = predict_test(test_data_loader, nets)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df_pseudo.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create Training Data","metadata":{}},{"cell_type":"code","source":"def generate_train(test_df_pseudo, fold_=1):\n    print(\"Loading training set...\")\n    marking = pd.read_csv('../input/wheat-cleaned-dataset/clean_train_v2.csv')\n    marking = marking[marking['area']> 300]\n    df_with_pseudo = pd.concat([marking, test_df_pseudo])\n\n    index = list(set(marking.image_id))\n    fold = 0\n    val_index = index[len(index)*fold//5:len(index)*(fold+1)//5]\n    train_index = index[:len(index)*fold//5] + index[len(index)*(fold+1)//5:]\n\n    train_df_pseudo_fold = test_df_pseudo['image_id'].unique()\n\n    train_fold = np.concatenate((train_index, train_df_pseudo_fold))\n    print('train dataset size:', len(train_fold))\n\n    \n    val_fold = np.array(val_index)\n    print('val dataset size:', len(val_fold))\n    \n    return df_with_pseudo, train_fold, val_fold","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_with_pesudo, train_fold, val_fold = generate_train(test_df_pseudo)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose(\n        [\n            #A.RandomSizedCrop(min_max_height=(800, 800), height=1024, width=1024, p=0.5),\n            A.OneOf([\n                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit= 0.2, \n                                     val_shift_limit=0.2, p=0.9),\n                A.RandomBrightnessContrast(brightness_limit=0.2, \n                                           contrast_limit=0.2, p=0.9),\n            ],p=0.9),\n            A.ToGray(p=0.01),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            \n            A.Resize(height=1024, width=1024, p=1),\n            #A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.5),\n            ToTensorV2(p=1.0),\n        ], \n        p=1.0, \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 get_valid_transforms():\n    return A.Compose(\n        [\n            A.Resize(height=1024, width=1024, p=1.0),\n            ToTensorV2(p=1.0),\n        ], \n        p=1.0, \n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0, \n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def randomsizedCroped(image, boxes, min_max_height=(800, 800), height=1024, width=1024, w2h_ratio=1.0):\n    \"\"\"\n    https://github.com/albumentations-team/albumentations/\n    \"\"\"\n    crop_height = random.randint(min_max_height[0], min_max_height[1])\n    h_start = random.random()\n    w_start = random.random()\n    crop_width = int(crop_height * w2h_ratio)\n    image_ = image.copy()\n    boxes_ = boxes.copy()\n\n    def get_random_crop_coords(height, width, crop_height, crop_width, h_start, w_start):\n        y1 = int((height - crop_height) * h_start)\n        y2 = y1 + crop_height\n        x1 = int((width - crop_width) * w_start)\n        x2 = x1 + crop_width\n        return x1, y1, x2, y2\n\n    x1, y1, x2, y2 = get_random_crop_coords(height, width, crop_height, crop_width, h_start, w_start)\n    image_ = image_[y1:y2, x1:x2]\n    image_ = cv2.resize(image_, (height, width))\n\n    ratio = [width / crop_width, height / crop_height]\n    # shift box\n    boxes_[:, 0] = boxes_[:, 0] - x1\n    boxes_[:, 1] = boxes_[:, 1] - y1\n    boxes_[:, 2] = boxes_[:, 2] - x1\n    boxes_[:, 3] = boxes_[:, 3] - y1\n\n    # resize box\n    boxes_[:, 0] = ratio[0] * boxes_[:, 0]\n    boxes_[:, 1] = ratio[1] * boxes_[:, 1]\n    boxes_[:, 2] = ratio[0] * boxes_[:, 2]\n    boxes_[:, 3] = ratio[1] * boxes_[:, 3]\n    boxes_[:, [0, 2]] = boxes_[:, [0, 2]].clip(min=0, max=width - 1)\n    boxes_[:, [1, 3]] = boxes_[:, [1, 3]].clip(min=0, max=height - 1)\n\n    boxes_ = boxes_.astype(np.int32)\n    boxes_ = boxes_[np.where((boxes_[:, 2] - boxes_[:, 0]) * (boxes_[:, 3] - boxes_[:, 1]) > 0.0)]\n\n    if len(boxes_) == 0:\n        return image, boxes\n    else:\n        return image_, boxes_","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def cutout(image, labels):\n    # https://arxiv.org/abs/1708.04552\n    # https://github.com/hysts/pytorch_cutout/blob/master/dataloader.py\n    # https://towardsdatascience.com/when-conventional-wisdom-fails-revisiting-data-augmentation-for-self-driving-cars-4831998c5509\n    # https://github.com/ultralytics/yolov5/blob/master/utils/datasets.py\n    h, w = image.shape[:2]\n    image_copy = image.copy()\n    labels_copy = labels.copy()\n\n    def bbox_ioa(box1, box2):\n        # Returns the intersection over box2 area given box1, box2. box1 is 4, box2 is nx4. boxes are x1y1x2y2\n        box2 = box2.transpose()\n\n        # Get the coordinates of bounding boxes\n        b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]\n        b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]\n\n        # Intersection area\n        inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \\\n                     (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)\n\n        # box2 area\n        box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + 1e-16\n\n        # Intersection over box2 area\n        return inter_area / box2_area\n\n    # create random masks\n    scales = [0.125] * 8 + [0.0625] * 12 + [0.03125] * 16 # image size fraction\n    for s in scales:\n        y = random.randint(0, h)\n        x = random.randint(0, w)\n\n        y1 = np.clip(y - h*s // 2, 0, h)\n        y2 = np.clip(y1 + h*s, 0, h)\n        x1 = np.clip(x - w*s // 2, 0, w)\n        x2 = np.clip(x1 + w*s, 0, w)\n\n        # apply random color mask\n        image[int(y1):int(y2), int(x1):int(x2)] = 0\n\n        # return unobscured labels\n        if len(labels) and s > 0.03:\n            box = np.array([x1, y1, x2, y2], dtype=np.float32)\n            ioa = bbox_ioa(box, labels[:, 0:4])  # intersection over area\n            labels = labels[ioa < 0.60]  # remove >60% obscured labels\n\n    if len(labels) == 0:\n        return image_copy, labels_copy\n\n    return image, labels","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def random_affine(img, targets=(), degrees=0, translate=0, scale=0.5, shear=0, border=[-1024 // 2, -1024 // 2]):\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    \n    img_, target_ = img.copy(), targets.copy()\n    \n    height = img.shape[0] + border[0] * 2  # shape(h,w,c)\n    width = img.shape[1] + border[1] * 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    # s = 2 ** random.uniform(-scale, 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[1] + border[1]  # x translation (pixels)\n    T[1, 2] = random.uniform(-translate, translate) * img.shape[0] + border[0]  # 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] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed\n        img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, \n                             borderValue=(114/255, 114/255, 114/255))\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[:, [0, 1, 2, 3, 0, 3, 2, 1]].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[:, 2] - targets[:, 0]) * (targets[:, 3] - targets[:, 1])\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[:, 0:4] = xy[i]\n    \n    if len(targets) == 0:\n        return img_, target_\n    \n    return img, targets","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TRAIN_ROOT_PATH = '../input/global-wheat-detection/train'\n\nclass DatasetRetriever(Dataset):\n\n    def __init__(self, marking, image_ids, transforms=None, test=False, mixup=True):\n        super().__init__()\n\n        self.image_ids = image_ids\n        self.marking = marking\n        self.transforms = transforms\n        self.test = test\n        self.mixup = mixup\n        self.mosaic_border = [-1024//2, -1024//2]\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        \n        origin_image = False\n        if self.test or random.random() > 0.5:\n            image, boxes = self.load_image_and_boxes(index)\n            origin_image = True\n        else:\n            image, boxes = self.load_cutmix_image_and_boxes(index)\n            \n        #if self.mixup and origin_image == True and (random.random() > 0.5) and (self.test == False):\n            #random_index = random.choice(np.delete(np.arange(self.image_ids.shape[0]), index))\n            #mix_image, mix_boxes = self.load_image_and_boxes(random_index)\n            #factor = 0.5\n            #image = factor*image + (1-factor)*mix_image\n            #boxes = np.concatenate([boxes, mix_boxes], 0)\n        #if (random.random() < 0.5) and (self.test == False):\n            #image, boxes = randomsizedCroped(image, boxes)               \n        if (random.random() < 0.9) and (self.test == False):\n            image, boxes = cutout(image, boxes)\n\n        # there is only one class\n        labels = torch.ones((boxes.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n\n        if self.transforms:\n            for i in range(10):\n                sample = self.transforms(**{\n                    'image': image,\n                    'bboxes': target['boxes'],\n                    'labels': labels\n                })\n                if len(sample['bboxes']) > 0:\n                    image = sample['image']\n                    target['labels'] = torch.ones((len(sample['bboxes']),), dtype=torch.int64)\n                    target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n                    target['boxes'][:,[0,1,2,3]] = target['boxes'][:,[1,0,3,2]]  #yxyx: be warning\n                    break\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n    def load_image_and_boxes(self, index):\n        image_id = self.image_ids[index]\n        if os.path.exists(f'{TRAIN_ROOT_PATH}/{image_id}.jpg'):\n            image = cv2.imread(f'{TRAIN_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        else:\n            image = cv2.imread(f'{DATA_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR) \n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image = cv2.resize(image, (1024, 1024))\n        image /= 255.0\n        records = self.marking[self.marking['image_id'] == image_id]\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        return image, boxes\n\n    def load_cutmix_image_and_boxes(self, index, imsize=1024):\n        \"\"\" \n        This implementation of cutmix author:  https://www.kaggle.com/nvnnghia \n        Refactoring and adaptation: https://www.kaggle.com/shonenkov\n        \"\"\"\n        w, h = imsize, imsize\n        s = imsize\n\n        xc, yc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border]\n        indexes = [index] + [random.randint(0, self.image_ids.shape[0] - 1) for _ in range(3)]\n\n        result_image = np.full((imsize*2, imsize*2, 3), 114/255, dtype=np.float32)\n        result_boxes = []\n\n        for i, index in enumerate(indexes):\n            \n            image, boxes = self.load_image_and_boxes(index)\n            if i == 0:\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            result_image[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]\n            padw = x1a - x1b\n            padh = y1a - y1b\n\n            boxes[:, 0] += padw\n            boxes[:, 1] += padh\n            boxes[:, 2] += padw\n            boxes[:, 3] += padh\n\n            result_boxes.append(boxes)\n\n        result_boxes = np.concatenate(result_boxes, 0)\n        np.clip(result_boxes[:, 0:], 0, 2 * s, out=result_boxes[:, 0:])\n        result_boxes = result_boxes.astype(np.int32)\n        result_boxes = result_boxes[np.where((result_boxes[:, 2] - result_boxes[:, 0]) * (result_boxes[:, 3] - result_boxes[:, 1]) > 300.0)]\n        result_boxes = result_boxes[np.where((result_boxes[:, 2] - result_boxes[:, 0]) >= 10.0)]\n        result_boxes = result_boxes[np.where((result_boxes[:, 3] - result_boxes[:, 1]) >= 10.0)]\n        result_image, result_boxes = random_affine(result_image, targets=result_boxes)\n\n        return result_image, result_boxes","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = DatasetRetriever(\n    image_ids=train_fold,\n    marking=df_with_pesudo,\n    transforms=get_train_transforms(),\n    test=False,\n)\n\nvalidation_dataset = DatasetRetriever(\n    image_ids=val_fold,\n    marking=df_with_pesudo,\n    transforms=get_valid_transforms(),\n    test=True,\n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image, target, image_id = train_dataset[0]\nboxes = target['boxes'].cpu().numpy().astype(np.int32)\n\nnumpy_image = image.permute(1,2,0).cpu().numpy()\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(numpy_image, (box[1], box[0]), (box[3],  box[2]), (0, 1, 0), 2)\n    \nax.set_axis_off()\nax.imshow(numpy_image);","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\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","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TrainGlobalConfig:\n    num_workers = 2\n    batch_size = 2\n    use_amp = True #use nvidia apex\n    use_ema = True #use model ema\n    n_epochs = 8\n    total_batch_size = 8\n    clip_grad = True\n    accumulate_steps = int(total_batch_size / batch_size) #real batch_size = 16, \n    lr = 0.0085\n    folder = 'effdet7'\n\n    # -------------------\n    verbose = True\n    verbose_step = 1\n    # -------------------\n\n    # --------------------\n    step_scheduler = False  # do scheduler.step after optimizer.step\n    validation_scheduler = True  # do scheduler.step after validation stage loss\n\n    SchedulerClass = CosineLRScheduler\n    scheduler_params = dict(\n        t_initial=n_epochs, \n        t_mul=1.0, \n        lr_min = 1e-8,\n    )\n    # --------------------","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def is_parallel(model):\n    # is model is parallel with DP or DDP\n    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)\n\nclass ModelEMA:\n    \"\"\" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models\n    Keep a moving average of everything in the model state_dict (parameters and buffers).\n    This is intended to allow functionality like\n    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage\n    A smoothed version of the weights is necessary for some training schemes to perform well.\n    This class is sensitive where it is initialized in the sequence of model init,\n    GPU assignment and distributed training wrappers.\n    \"\"\"\n\n    def __init__(self, model, decay=0.9999, updates=0):\n        # Create EMA\n        self.ema = deepcopy(model.module if is_parallel(model) else model).eval()  # FP32 EMA\n        # if next(model.parameters()).device.type != 'cpu':\n        #     self.ema.half()  # FP16 EMA\n        self.updates = updates  # number of EMA updates\n        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)\n        for p in self.ema.parameters():\n            p.requires_grad_(False)\n\n    def update(self, model):\n        # Update EMA parameters\n        with torch.no_grad():\n            self.updates += 1\n            d = self.decay(self.updates)\n\n            msd = model.module.state_dict() if is_parallel(model) else model.state_dict()  # model state_dict\n            for k, v in self.ema.state_dict().items():\n                if v.dtype.is_floating_point:\n                    v *= d\n                    v += (1. - d) * msd[k].detach()\n\n    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):\n        # Update EMA attributes\n        copy_attr(self.ema, model, include, exclude)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nclass Fitter:\n    \n    def __init__(self, model, device, config):\n        self.config = config\n        self.epoch = 0\n\n        self.base_dir = f'./{config.folder}'\n        if not os.path.exists(self.base_dir):\n            os.makedirs(self.base_dir)\n        \n        self.log_path = f'{self.base_dir}/log.txt'\n        self.best_summary_loss = 10**5\n        \n        self.model = model\n        self.device = device\n        self.use_amp = config.use_amp\n        self.accumulate_steps = config.accumulate_steps\n        self.clip_grad = config.clip_grad\n        self.max_norm = 10.0\n        self.use_ema = config.use_ema\n        self.model_ema = None\n        self.model_ema_decay = 0.999\n        \n        pg0, pg1, pg2 = [], [], []  # optimizer parameter groups\n        for k, v in self.model.named_parameters():\n            if v.requires_grad:\n                if '.bias' in k:\n                    pg2.append(v)  # biases\n                elif '.weight' in k and '.bn' not in k:\n                    pg1.append(v)  # apply weight decay\n                else:\n                    pg0.append(v)  # all else\n        self.optimizer = torch.optim.SGD(pg0, lr=config.lr, momentum=0.9, nesterov=True)\n        self.optimizer.add_param_group({'params': pg1, 'weight_decay': 5e-4})\n        self.optimizer.add_param_group({'params': pg2})  \n        \n        self.model.to(self.device)\n        \n        if self.use_amp:\n            self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level=\"O1\",verbosity=0)\n        if self.use_ema:\n            self.model_ema = ModelEMA(self.model, decay=self.model_ema_decay)\n        self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)\n        self.log(f'Fitter prepared. Device is {self.device}')\n\n    def fit(self, train_loader, validation_loader):\n        \n        best_epoch = 0\n        for e in range(self.config.n_epochs):\n            if self.config.verbose:\n                lr = self.optimizer.param_groups[0]['lr']\n                timestamp = datetime.utcnow().isoformat()\n                self.log(f'\\n{timestamp}\\nLR: {lr}')\n\n            t = time.time()\n            summary_loss = self.train_one_epoch(train_loader)\n\n            self.log(f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n            self.save(f'{self.base_dir}/last-checkpoint.bin')\n\n            t = time.time()\n            if self.use_ema:\n                summary_loss = self.validation(self.model_ema.ema.module if hasattr(self.model_ema.ema, 'module') else self.model_ema.ema, \n                                               validation_loader)\n            else:\n                summary_loss = self.validation(self.model, validation_loader)\n\n            self.log(f'[RESULT]: Val. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n            if summary_loss.avg < self.best_summary_loss:\n                self.best_summary_loss = summary_loss.avg\n                self.model.eval()\n                self.save(f'{self.base_dir}/best-checkpoint-{str(self.epoch).zfill(3)}epoch.bin')\n                for path in sorted(glob(f'{self.base_dir}/best-checkpoint-*epoch.bin'))[:-3]:\n                    os.remove(path)\n                best_epoch = e\n\n            if self.config.validation_scheduler:\n                self.scheduler.step(e+1)\n\n            self.epoch += 1\n            \n        return best_epoch\n            \n    def validation(self, model, val_loader):\n        model.eval()\n        summary_loss = AverageMeter()\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(val_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    print(\n                        f'Val Step {step}/{len(val_loader)}, ' + \\\n                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}', end='\\r'\n                    )\n            with torch.no_grad():\n                images = torch.stack(images).to(self.device).float()\n                batch_size = images.shape[0]\n                target_input = {'bbox':[target['boxes'].to(self.device).float() for target in targets],\n                                'cls': [target['labels'].to(self.device).float() for target in targets]} \n                loss = model(images, target_input['bbox'], target_input['cls'])\n                loss, class_loss, box_loss = loss\n                summary_loss.update(loss.detach().item(), batch_size)\n\n        return summary_loss\n\n    def train_one_epoch(self, train_loader):\n        self.model.train()\n        summary_loss = AverageMeter()\n        t = time.time()\n        self.optimizer.zero_grad() #very important\n        for step, (images, targets, image_ids) in enumerate(train_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    print(\n                        f'Train Step {step}/{len(train_loader)}, ' + \\\n                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}', end='\\r'\n                    )\n            \n            images = torch.stack(images).to(self.device).float()\n            batch_size = images.shape[0]\n            target_input = {'bbox':[target['boxes'].to(self.device).float() for target in targets],\n                            'cls': [target['labels'].to(self.device).float() for target in targets]} \n\n            loss = self.model(images, target_input['bbox'], target_input['cls'])\n            loss, class_loss, box_loss = loss\n            \n            if self.use_amp:\n                with amp.scale_loss(loss, self.optimizer) as scaled_loss:\n                    scaled_loss.backward()\n            else:\n                loss.backward()\n            if (step+1) % self.accumulate_steps == 0: # Wait for several backward steps\n                if self.clip_grad:\n                    if self.use_amp:\n                        torch.nn.utils.clip_grad_norm_(amp.master_params(self.optimizer), self.max_norm)\n                    else:\n                        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_norm)\n                \n                self.optimizer.step()  # Now we can do an optimizer step\n                self.optimizer.zero_grad()\n                if self.model_ema is not None:\n                    self.model_ema.update(self.model)\n\n            summary_loss.update(loss.detach().item(), batch_size)\n\n            if self.config.step_scheduler:\n                self.scheduler.step()\n\n            \n        return summary_loss\n    \n    def save(self, path):\n        self.model.eval()\n        \n        if self.use_ema:\n            torch.save({\n                'model_state_dict': self.model.model.state_dict(),\n                'state_dict_ema': self.model_ema.ema.module.state_dict() if hasattr(self.model_ema.ema, 'module') else self.model_ema.ema.model.state_dict(),\n                'optimizer_state_dict': self.optimizer.state_dict(),\n                'scheduler_state_dict': self.scheduler.state_dict(),\n                'best_summary_loss': self.best_summary_loss,\n                'epoch': self.epoch,\n            }, path)\n        else:\n            torch.save({\n                'model_state_dict': self.model.model.state_dict(),\n                'optimizer_state_dict': self.optimizer.state_dict(),\n                'scheduler_state_dict': self.scheduler.state_dict(),\n                'best_summary_loss': self.best_summary_loss,\n                'epoch': self.epoch,\n            }, path)\n\n    def load(self, path):\n        checkpoint = torch.load(path)\n        self.model.model.load_state_dict(checkpoint['model_state_dict'])\n        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n        self.best_summary_loss = checkpoint['best_summary_loss']\n        self.epoch = checkpoint['epoch'] + 1\n        \n    def log(self, message):\n        if self.config.verbose:\n            print(message)\n        with open(self.log_path, 'a+') as logger:\n            logger.write(f'{message}\\n')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ndef run_training():\n    device = torch.device('cuda:0')\n    net.to(device)\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=RandomSampler(train_dataset),\n        pin_memory=False,\n        drop_last=True,\n        num_workers=TrainGlobalConfig.num_workers,\n        collate_fn=collate_fn,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        validation_dataset, \n        batch_size=TrainGlobalConfig.batch_size,\n        num_workers=TrainGlobalConfig.num_workers,\n        shuffle=False,\n        sampler=SequentialSampler(validation_dataset),\n        pin_memory=False,\n        collate_fn=collate_fn,\n    )\n\n    fitter = Fitter(model=net, device=device, config=TrainGlobalConfig)\n    best_epoch = fitter.fit(train_loader, val_loader)\n    \n    return best_epoch","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_net():\n    config = get_efficientdet_config('tf_efficientdet_d7')\n    net = EfficientDet(config, pretrained_backbone=False)\n    config.num_classes = 1\n    config.image_size = 1024\n    net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n    \n    checkpoint = torch.load(trained_weight_path)\n    net.load_state_dict(checkpoint['state_dict_ema'])\n    \n    return DetBenchTrain(net, config)\n\nnet = get_net()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_epoch = None \nif len(os.listdir(\"../input/global-wheat-detection/test/\")) > 11:\n    best_epoch = run_training()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference ","metadata":{}},{"cell_type":"code","source":"if best_epoch is not None:\n    print(\"load retrained model {}\".format(best_epoch))\n    net = create_model_efficientdet_test(compound_coef=7, \n                                     weight_path=f'./effdet7/best-checkpoint-{str(best_epoch).zfill(3)}epoch.bin', \n                                     image_size=1024,\n                                     model_ema=True)\nelse:\n    net = create_model_efficientdet_test(compound_coef=7, \n                                     weight_path=trained_weight_path, \n                                     image_size=1024,\n                                     model_ema=True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BaseWheatTTA:\n    \"\"\" author: @shonenkov \"\"\"\n    image_size = 1024\n\n    def augment(self, image):\n        raise NotImplementedError\n    \n    def batch_augment(self, images):\n        raise NotImplementedError\n    \n    def deaugment_boxes(self, boxes):\n        raise NotImplementedError\n\nclass TTAHorizontalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n\n    def augment(self, image):\n        return image.flip(1)\n    \n    def batch_augment(self, images):\n        return images.flip(2)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [1,3]] = self.image_size - boxes[:, [3,1]]\n        return boxes\n\nclass TTAVerticalFlip(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return image.flip(2)\n    \n    def batch_augment(self, images):\n        return images.flip(3)\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,2]] = self.image_size - boxes[:, [2,0]]\n        return boxes\n    \nclass TTARotate90(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 1, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 1, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = self.image_size - boxes[:, [1,3]]\n        res_boxes[:, [1,3]] = boxes[:, [2,0]]\n        return res_boxes\n\nclass TTACompose(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    def __init__(self, transforms):\n        self.transforms = transforms\n        \n    def augment(self, image):\n        for transform in self.transforms:\n            image = transform.augment(image)\n        return image\n    \n    def batch_augment(self, images):\n        for transform in self.transforms:\n            images = transform.batch_augment(images)\n        return images\n    \n    def prepare_boxes(self, boxes):\n        result_boxes = boxes.copy()\n        result_boxes[:,0] = np.min(boxes[:, [0,2]], axis=1)\n        result_boxes[:,2] = np.max(boxes[:, [0,2]], axis=1)\n        result_boxes[:,1] = np.min(boxes[:, [1,3]], axis=1)\n        result_boxes[:,3] = np.max(boxes[:, [1,3]], axis=1)\n        return result_boxes\n    \n    def deaugment_boxes(self, boxes):\n        for transform in self.transforms[::-1]:\n            boxes = transform.deaugment_boxes(boxes)\n        return self.prepare_boxes(boxes)\n\nfrom itertools import product\n\ntta_transforms = []\nfor tta_combination in product([TTAHorizontalFlip(), None], \n                               [TTAVerticalFlip(), None],\n                               [TTARotate90(), None]):\n    tta_transforms.append(TTACompose([tta_transform for tta_transform in tta_combination if tta_transform]))\n\ndef make_tta_predictions(images, net, score_threshold=0.25):\n    with torch.no_grad():\n        images = torch.stack(images).float().cuda()\n        predictions = []\n        for tta_transform in tta_transforms:\n            result = []\n            det = net(tta_transform.batch_augment(images.clone()), torch.tensor([1]*images.shape[0]).float().cuda())\n\n            for i in range(images.shape[0]):\n                boxes = det[i].detach().cpu().numpy()[:,:4]    \n                scores = det[i].detach().cpu().numpy()[:,4]\n                indexes = np.where(scores > score_threshold)[0]\n                boxes = boxes[indexes]\n                boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n                boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\n                boxes = tta_transform.deaugment_boxes(boxes.copy())\n                result.append({\n                    'boxes': boxes,\n                    'scores': scores[indexes],\n                })\n            predictions.append(result)\n    return predictions\n\ndef run_wbf(predictions, image_index, image_size=1024, iou_thr=0.44, skip_box_thr=0.43, weights=None):\n    boxes = [(prediction[image_index]['boxes']/(image_size-1)).tolist() for prediction in predictions]\n    scores = [prediction[image_index]['scores'].tolist() for prediction in predictions]\n    labels = [np.ones(prediction[image_index]['scores'].shape[0]).astype(int).tolist() for prediction in predictions]\n    boxes, scores, labels = ensemble_boxes_wbf.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","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_test_transforms():\n    return A.Compose([\n            A.Resize(height=1024, width=1024, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.0)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset = DatasetRetrieverTest(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{DATA_ROOT_PATH}/*.jpg')]),\n    transforms=get_test_transforms()\n)\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=2,\n    shuffle=False,\n    num_workers=4,\n    drop_last=False,\n    collate_fn=collate_fn\n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for j, (images, image_ids) in enumerate(test_data_loader):\n    break\n\npredictions = make_tta_predictions(images, net=net)\n\ni = 1\nsample = images[i].permute(1,2,0).cpu().numpy()\n\nboxes, scores, labels = run_wbf(predictions, image_index=i)\nboxes = boxes.astype(np.int32).clip(min=0, max=1023)\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (1, 0, 0), 3)\n    \nax.set_axis_off()\nax.imshow(sample);","metadata":{},"execution_count":null,"outputs":[]},{"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    return \" \".join(pred_strings)\nresults = []\n\nfor images, image_ids in test_data_loader:\n    predictions = make_tta_predictions(images, net=net)\n    for i, image in enumerate(images):\n        boxes, scores, labels = run_wbf(predictions, image_index=i)\n        boxes = boxes.astype(np.int32).clip(min=0, max=1023)\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        results.append(result)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head()","metadata":{},"execution_count":null,"outputs":[]}]}