{"cells":[{"metadata":{"trusted":false},"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"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\nimport ensemble_boxes \nimport torch\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nfrom torch.utils.data import Dataset,DataLoader\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport cv2\nimport gc\nfrom matplotlib import pyplot as plt\nfrom effdet import get_efficientdet_config, EfficientDet, DetBenchEval, DetBenchTrain\nfrom effdet.efficientdet import HeadNet\nimport os\nfrom datetime import datetime\nimport time\nimport random\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Seed","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"# SEED = 42\n\n# def 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\n# seed_everything(SEED) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference Settings","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"# plabel\nPLABEL_OPT = True\nplabel_epoch = 2\nplabel_bs = 2\nplabel_lr = 0.0001\nplabel_score_threshold = 0.3\n\nif PLABEL_OPT and len(os.listdir('../input/global-wheat-detection/test/'))>10:\n    PLABEL_OPT = True\nelse:\n    PLABEL_OPT = False\n\n# TTA\nTTA_OPT = True\n\n# WBF\nWBF_OPT = True\nWBF_iou_thr = 0.44\nWBF_skip_box_thr = 0.43\n\n# OOF\noof_score_threshold=0.3\n\n# val data\nval_data_type = [0]    # [0]  /  [0,1,2,3,4]  /  None\n\n# model path\nmodel_path = [\n    '',\n]\n\n# img_size\nimg_size = 1024\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data Augmentation for train val test","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    # as same as training code\n    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.1, sat_shift_limit=0.3,\n                                         val_shift_limit=0.3, p=0.9),\n                    A.RandomBrightnessContrast(brightness_limit=0.4,\n                                               contrast_limit=0.3, 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.Resize(height=img_size, width=img_size, 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\n    def get_valid_transforms():\n        return A.Compose(\n            [\n                A.Resize(height=img_size, width=img_size, 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        )","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def get_test_transforms():\n    return A.Compose([\n            A.Resize(height=img_size, width=img_size, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load test data ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Test Dataset","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"TEST_ROOT_PATH = '../input/global-wheat-detection/test'\n\nclass DatasetTest(Dataset):\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'{TEST_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]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset and Dataloader\nNote that batchsize of test dataloader must be 1","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"test_dataset = DatasetTest(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{TEST_ROOT_PATH}/*.jpg')]),\n    transforms=get_test_transforms()\n)\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# batchsize must be 1\ntest_dataloader = DataLoader(\n    test_dataset,\n    batch_size=1,\n    shuffle=False,\n    num_workers=4,\n    drop_last=False,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load model","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"def load_test_net(checkpoint_path):\n    config = get_efficientdet_config('tf_efficientdet_d5')\n    net = EfficientDet(config, pretrained_backbone=False)\n\n    config.num_classes = 1\n    config.image_size=img_size\n    net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n\n    checkpoint = torch.load(checkpoint_path)\n    net.load_state_dict(checkpoint['model_state_dict'])\n\n    del checkpoint\n    gc.collect()\n\n    net = DetBenchEval(net, config)\n    net.eval();\n    return net.cuda()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    def load_train_net(ckpt_path):    \n        config = get_efficientdet_config('tf_efficientdet_d5')\n        net = EfficientDet(config, pretrained_backbone=False)\n\n        config.num_classes = 1\n        config.image_size = img_size\n        net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n        checkpoint = torch.load(ckpt_path)\n        net.load_state_dict(checkpoint['model_state_dict'])\n\n        return DetBenchTrain(net, config)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# TTA  \nIdea is simple: \n- `augment` make tta for one image\n- `batch_augment` make tta for batch of images\n- `deaugment_boxes` return tta predicted boxes in back to original state of image\n\nAlso we are interested in `Compose` with combinations of tta :)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Function Definition","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if TTA_OPT:\n    class BaseWheatTTA:\n        \"\"\" author: @shonenkov \"\"\"\n        image_size = img_size\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\n    class 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\n    class 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\n    class 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\n    class 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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### TTA compose","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if TTA_OPT:\n    from itertools import product\n\n    tta_transforms = []\n    for 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]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Inference function using TTA","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if TTA_OPT:\n    def make_tta_predictions(images, models, score_threshold):\n        with torch.no_grad():\n            images = torch.stack(images).float().cuda()\n            assert images.shape[0] == 1\n            predictions = []\n            boxes_All = []\n            scores_All = []\n            for tta_transform in tta_transforms:\n\n                for net in models:  # model ensemble\n                    \n                    det = net(tta_transform.batch_augment(images.clone()), torch.tensor([1]*images.shape[0]).float().cuda())  # send to model\n\n                    boxes = det[0].detach().cpu().numpy()[:,:4]    \n                    scores = det[0].detach().cpu().numpy()[:,4]\n\n                    boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n                    boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\n                    boxes = tta_transform.deaugment_boxes(boxes.copy())    # de-TTA: back to org\n\n                    indexes = np.where(scores > score_threshold)[0]  # select via simple threshold\n                    boxes = boxes[indexes]\n                    scores = scores[indexes]\n                    \n                    if len(boxes_All) == 0:\n                        boxes_All = boxes\n                        scores_All = scores\n                    else:\n                        boxes_All = np.concatenate((boxes_All,boxes),0)\n                        scores_All = np.concatenate((scores_All,scores),0)\n\n            result = {\n                'boxes': boxes_All,\n                'scores': scores_All,\n            }\n            predictions.append([result])   \n            \n        return predictions","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Inference function without TTA","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"def make_ensemble_predictions(images, models, score_threshold):\n    with torch.no_grad():\n        images = torch.stack(images).float().cuda()\n        assert images.shape[0] == 1\n        # loop\n        predictions = []\n        boxes_All = []\n        scores_All = []\n        for net in models:  # model ensemble\n\n            det = net(images.clone(), torch.tensor([1]*images.shape[0]).float().cuda())  # send to model\n\n            boxes = det[0].detach().cpu().numpy()[:,:4]    \n            scores = det[0].detach().cpu().numpy()[:,4]\n\n            boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n            boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\n\n            indexes = np.where(scores > score_threshold)[0]  # select via simple threshold\n            boxes = boxes[indexes]\n            scores = scores[indexes]\n\n            if len(boxes_All) == 0:\n                boxes_All = boxes\n                scores_All = scores\n            else:\n                boxes_All = np.concatenate((boxes_All,boxes),0)\n                scores_All = np.concatenate((scores_All,scores),0)\n\n        result = {\n            'boxes': boxes_All,\n            'scores': scores_All,\n        }\n        predictions.append([result])   \n\n    return predictions","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# WBF \nTo fuse all predictions for different TTAs and different models\n### Function Definition","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if WBF_OPT:\n    def run_wbf(predictions, image_size, iou_thr, skip_box_thr, weights=None):\n        assert len(predictions) == 1\n        assert len(predictions[0]) == 1\n        boxes = [(predictions[0][0]['boxes']/(image_size-1)).tolist()]\n        scores = [(predictions[0][0]['scores']).tolist()]\n        labels = [(np.ones(predictions[0][0]['scores'].shape[0]).astype(int)).tolist()]\n        boxes, scores, labels = ensemble_boxes.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\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def run_without_wbf(predictions):\n    assert len(predictions) == 1\n    assert len(predictions[0]) == 1\n    boxes = predictions[0][0]['boxes']\n    scores = predictions[0][0]['scores']\n    labels = np.ones(predictions[0][0]['scores'].shape[0]).astype(int)\n    return boxes, scores, labels\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Plabel ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Plabel Training Settings ","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    class TrainGlobalConfig:\n        num_workers = 0\n        batch_size = plabel_bs \n        n_epochs = plabel_epoch    \n        lr = plabel_lr\n\n        folder = 'plabel_model'\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 = torch.optim.lr_scheduler.CosineAnnealingLR\n        scheduler_params = dict(\n            T_max=plabel_epoch,\n            eta_min=0,\n        )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Init model","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"# remember to rewrite with your ckpt path\nif PLABEL_OPT:\n    test_models = []\n    for p in model_path:\n        test_models.append(load_test_net(p))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"###  Generate plabel for test data","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    results_plabel = []\n    for images, image_ids in test_dataloader:\n        # get predictions\n        if TTA_OPT:\n            predictions = make_tta_predictions(images=images, models=test_models, score_threshold=plabel_score_threshold)\n        else:\n            predictions = make_ensemble_predictions(images=images, models=test_models, score_threshold=plabel_score_threshold)\n        \n        # next step\n        for i, image in enumerate(images):\n            assert i == 0\n            image_id = image_ids[i]\n            image_ = cv2.imread(f'{TEST_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n            h,w,_ = np.shape(image_)\n            \n            # WBF\n            if WBF_OPT:\n                boxes, scores, labels = run_wbf(predictions, image_size=img_size, iou_thr=WBF_iou_thr, skip_box_thr=WBF_skip_box_thr)\n            else:\n                boxes, scores, labels = run_without_wbf(predictions)\n\n            # resize bbox\n            rate = float(h / img_size)\n            boxes = (boxes*rate).astype(np.int32).clip(min=0, max=h-1)\n            \n            # back to xywh\n            boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n            boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n\n            # make plabel\n            for box in boxes:\n                result_p = {\n                    'image_id': image_id,\n                    'width':w,\n                    'height':h,\n                    'source':'usask_1',\n                    'x':box[0],\n                    'y':box[1],\n                    'w':box[2],\n                    'h':box[3],\n                }\n                results_plabel.append(result_p)\n    del test_models\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"###  Put test data info and plabel to pd","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    results_df = pd.DataFrame(results_plabel, columns=['image_id', 'width','height','source','x','y','w','h'])\n    results_df.head()\n    results_df['image_id'] = results_df['image_id'].apply(lambda x: TEST_ROOT_PATH+'/'+ x+'.jpg')   # 考虑改一下以区分test和train数据","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load train data via .csv file","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    # as same as training code\n    marking = pd.read_csv('../input/global-wheat-detection/train.csv')\n\n    bboxs = np.stack(marking['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\n    for i, column in enumerate(['x', 'y', 'w', 'h']):\n        marking[column] = bboxs[:,i]\n    marking.drop(columns=['bbox'], inplace=True)\n    marking.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Concat train and test in pd","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    TRAIN_ROOT_PATH = '../input/global-wheat-detection/train'\n    marking['image_id'] = marking['image_id'].apply(lambda x: TRAIN_ROOT_PATH+'/'+ x+'.jpg')     # 考虑改一下以区分test和train数据\n    train_data_plabel = pd.concat([results_df,marking], axis=0)     # concat \"train pd\"  and \"test pd\" ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Split k-fold for the whole pd","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n\n    df_folds = train_data_plabel[['image_id']].copy()\n    df_folds.loc[:, 'bbox_count'] = 1\n    df_folds = df_folds.groupby('image_id').count()\n    df_folds.loc[:, 'source'] = marking[['image_id', 'source']].groupby('image_id').min()['source']  \n    df_folds.loc[:, 'stratify_group'] = np.char.add(\n        df_folds['source'].values.astype(str),\n        df_folds['bbox_count'].apply(lambda x: f'_{x // 15}').values.astype(str)\n    )\n    df_folds.loc[:, 'fold'] = 0\n\n    for fold_number, (train_index, val_index) in enumerate(skf.split(X=df_folds.index, y=df_folds['stratify_group'])):\n        df_folds.loc[df_folds.iloc[val_index].index, 'fold'] = fold_number","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset class for training","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    class DatasetRetriever(Dataset):\n        def __init__(self, marking, image_ids, transforms=None, test=False):\n            super().__init__()\n\n            self.image_ids = image_ids\n            self.marking = marking\n            self.transforms = transforms\n            self.test = test\n\n        def __getitem__(self, index: int):\n            image_id = self.image_ids[index]\n\n            # mixup & cutmix\n            random_num = random.random()\n            if self.test or random_num < 0.4:\n                image, boxes = self.load_image_and_boxes(index)\n            elif random_num >= 0.4 and random_num < 0.7:\n                image, boxes = self.load_mixup_image_and_boxes(index)\n            else:\n                image, boxes = self.load_cutmix_image_and_boxes(index)\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['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            # image = cv2.imread(f'{TRAIN_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n            image = cv2.imread(image_id, cv2.IMREAD_COLOR)\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\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_mixup_image_and_boxes(self, index):\n            image, boxes = self.load_image_and_boxes(index)   # 加载图片和 bbox\n            r_image, r_boxes = self.load_image_and_boxes(random.randint(0, self.image_ids.shape[0] - 1))   # 随机加载另外一张图片和 bbox\n            mixup_image = (image + r_image) / 2    # 进行 mixup 图片的融合，这里简单的利用 0.5 权重\n            mixup_boxes = np.concatenate((boxes, r_boxes), 0)     # 进行 mixup bbox的融合\n            return mixup_image, mixup_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 // 2\n\n            xc, yc = [int(random.uniform(imsize * 0.25, imsize * 0.75)) for _ in range(2)]  # center x, y\n            indexes = [index] + [random.randint(0, self.image_ids.shape[0] - 1) for _ in range(3)]\n\n            result_image = np.full((imsize, imsize, 3), 1, dtype=np.float32)\n            result_boxes = []\n\n            for i, index in enumerate(indexes):\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]) > 0)]\n            return result_image, result_boxes","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load train and val data for plabel training","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    def get_dataset_dataloader(fold_number=None):\n        if fold_number is not None:\n            train_dataset = DatasetRetriever(\n                image_ids=df_folds[df_folds['fold'] != fold_number].index.values,\n                marking=train_data_plabel,\n                transforms=get_train_transforms(),\n                test=False,\n            )\n            validation_dataset = DatasetRetriever(\n                image_ids=df_folds[df_folds['fold'] == fold_number].index.values,\n                marking=train_data_plabel,\n                transforms=get_valid_transforms(),\n                test=True,\n            )\n            train_loader = torch.utils.data.DataLoader(\n                train_dataset,\n                batch_size=TrainGlobalConfig.batch_size,\n                shuffle=True,\n                pin_memory=False,\n                drop_last=False,\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                drop_last=False,\n                pin_memory=False,\n                collate_fn=collate_fn,\n            )\n            return train_dataset, validation_dataset, train_loader, val_loader\n        else:\n            train_dataset = DatasetRetriever(\n                image_ids=df_folds.index.values,\n                marking=train_data_plabel,\n                transforms=get_train_transforms(),\n                test=False,\n            )\n            train_loader = torch.utils.data.DataLoader(\n                train_dataset,\n                batch_size=TrainGlobalConfig.batch_size,\n                shuffle=True,\n                pin_memory=False,\n                drop_last=False,\n                num_workers=TrainGlobalConfig.num_workers,\n                collate_fn=collate_fn,\n            )\n            return train_dataset, None, train_loader, None","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training Functions","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    import warnings\n    warnings.filterwarnings(\"ignore\")\n\n    class Fitter:\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\n            param_optimizer = list(self.model.named_parameters())\n            no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n            optimizer_grouped_parameters = [\n                {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n                {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n            ] \n\n            self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=config.lr)\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, model_idx):\n            self.best_summary_loss = 10**5\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                # train\n                t = time.time()\n                summary_loss = self.train_one_epoch(train_loader)\n                self.log(f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n                if validation_loader is None:\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.bin')\n\n                # val\n                if validation_loader is not None:\n                    t = time.time()\n                    summary_loss = self.validation(validation_loader)\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(model_idx)}.bin')\n                \n                # modify lr\n                if self.config.validation_scheduler:\n                    self.scheduler.step()\n\n                self.epoch += 1\n\n        def validation(self, val_loader):\n            self.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)\n                    batch_size = images.shape[0]\n                    images = images.to(self.device).float()\n                    boxes = [target['boxes'].to(self.device).float() for target in targets]\n                    labels = [target['labels'].to(self.device).float() for target in targets]\n\n                    loss, _, _ = self.model(images, boxes, labels)\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            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)\n                images = images.to(self.device).float()\n                batch_size = images.shape[0]\n                boxes = [target['boxes'].to(self.device).float() for target in targets]\n                labels = [target['labels'].to(self.device).float() for target in targets]\n\n                self.optimizer.zero_grad()\n\n                loss, _, _ = self.model(images, boxes, labels)\n\n                loss.backward()\n\n                summary_loss.update(loss.detach().item(), batch_size)\n\n                self.optimizer.step()\n\n                if self.config.step_scheduler:\n                    self.scheduler.step()\n\n            return summary_loss\n\n        def save(self, path):\n            self.model.eval()\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')","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    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","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    def run_training(models):\n        device = torch.device('cuda:0')\n        # load data\n        if val_data_type is None:   \n            assert len(models) == 1   # only one model\n            train_dataset, validation_dataset, train_loader, val_loader = get_dataset_dataloader(fold_number=None)\n            fitter = Fitter(model=models[0].to(device), device=device, config=TrainGlobalConfig)\n            fitter.fit(train_loader, val_loader, None)\n        else:\n            i = 0\n            for fold_num in val_data_type:\n                assert len(val_data_type) == len(models)\n                train_dataset, validation_dataset, train_loader, val_loader = get_dataset_dataloader(fold_number=fold_num)\n                fitter = Fitter(model=models[i].to(device), device=device, config=TrainGlobalConfig)\n                fitter.fit(train_loader, val_loader, fold_num)\n                i = i + 1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Run plabel training (fine-tune)","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"if PLABEL_OPT:\n    train_models = []\n    for p in model_path:\n        train_models.append(load_train_net(p))\n\n    run_training(train_models) \n    \n    del train_models\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\n\n\n# Final Inference","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Format for Preditions","execution_count":null},{"metadata":{"trusted":false},"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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load final model","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"# remenber to rewrite\nif PLABEL_OPT:\n    if val_data_type is None:\n        final_models = [\n            load_test_net(f'plabel_model/best_checkpoint.bin')\n        ]\n        print('use pl')\n    else:\n        final_models = []\n        for i in val_data_type:\n            final_models.append(load_test_net(f'plabel_model/best_checkpoint-{str(i)}.bin'))\nelse:\n    final_models = []\n    for p in model_path:\n        final_models.append(load_test_net(p))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### inference","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Output ","execution_count":null},{"metadata":{"trusted":false},"cell_type":"code","source":"results = []\n\nfor images, image_ids in test_dataloader:\n    # TTA \n    if TTA_OPT:\n        predictions = make_tta_predictions(images=images, models=final_models, score_threshold=oof_score_threshold)\n    else:\n        predictions = make_ensemble_predictions(images=images, models=final_models, score_threshold=oof_score_threshold)\n    \n    # next step\n    for i, image in enumerate(images):\n        assert i == 0\n        # image id\n        image_id = image_ids[i]\n        # WBF\n        if WBF_OPT:\n            boxes, scores, labels = run_wbf(predictions, image_size=img_size, iou_thr=WBF_iou_thr, skip_box_thr=WBF_skip_box_thr)\n        else:\n            boxes, scores, labels = run_without_wbf(predictions)\n        # resize bbox\n        rate = float(h / img_size)\n        print('rate', rate)\n        boxes = (boxes*rate).astype(np.int32).clip(min=0, max=h-1)\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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head()","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}