{"cells":[{"metadata":{"trusted":true},"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":true},"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.optim.lr_scheduler import _LRScheduler\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\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":true},"cell_type":"code","source":"SEED = 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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Settings","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# path\nTRAIN_DATA_PATH = '../input/global-wheat-detection/train/'\nTRAIN_CSV_PATH = '../input/global-wheat-detection/train.csv'\nTEST_DATA_PATH = '../input/global-wheat-detection/test/'\n\n# Plabel\nif len(os.listdir('../input/global-wheat-detection/test/'))>11:\n    PL_OPT = True\nelse:\n    PL_OPT = False\n\nwarmup_opt = True\nwarmup_epoch = 1\nPL_lr_sche = 'plat'  # plat\ncos_lr_min = 1e-8\nPL_lr = 0.0001\nPL_batchsize = 1\nPL_epoch = 5\nPL_thr = 0.25\n\n# img size\nimg_size = 1024\n\n# OOF\nOOF_thr = 0.3\n\n# WBF\nWBF_iou_thr = 0.44\nWBF_skip_thr = 0.43\n\n# model path\npath1 = [\n    '../input/212121/21_all - 0.39332 train - epoch89.bin'\n]\n\npath2 = [\n    '../input/292929/29_0 - 0.40681 train - 0.37596 val - epoch94.bin',\n]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# transform","execution_count":null},{"metadata":{"trusted":true},"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.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\ndef 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    )\ndef 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 data","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# load train data\nmarking = pd.read_csv(TRAIN_CSV_PATH)\nbboxs = np.stack(marking['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\nfor i, column in enumerate(['x', 'y', 'w', 'h']):\n    marking[column] = bboxs[:,i]\nmarking.drop(columns=['bbox'], inplace=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# load test data\nclass DatasetT(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'{TEST_DATA_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]\n\n    \ndataset = DatasetT(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{TEST_DATA_PATH}/*.jpg')]),\n    transforms=get_test_transforms()\n)\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndata_loader = DataLoader(\n    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":{"trusted":true},"cell_type":"code","source":"class DatasetRetriever(Dataset):\n\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        if self.test or random.random() > 0.5:\n            image, boxes = self.load_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(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_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 model","execution_count":null},{"metadata":{"trusted":true},"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":true},"cell_type":"code","source":"from effdet import get_efficientdet_config, EfficientDet, DetBenchTrain\nfrom effdet.efficientdet import HeadNet\n\ndef load_train_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    checkpoint = torch.load(checkpoint_path)\n    net.load_state_dict(checkpoint['model_state_dict'])\n    \n    del checkpoint\n    gc.collect()\n    \n    return DetBenchTrain(net, config)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# TTA","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"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\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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"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":"# WBF","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# plabel settings","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class TrainGlobalConfig:\n    num_workers = 2\n    batch_size = PL_batchsize \n    n_epochs = PL_epoch # n_epochs = 40\n    lr = PL_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    if PL_lr_sche == 'cos':\n        SchedulerClass = torch.optim.lr_scheduler.CosineAnnealingLR\n        scheduler_params = dict(\n            T_max=PL_epoch,\n            eta_min=cos_lr_min,\n        )\n    elif PL_lr_sche == 'plat':\n        SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau\n        scheduler_params = dict(\n            mode='min',\n            factor=0.5,\n            patience=1,\n            verbose=False, \n            threshold=0.0001,\n            threshold_mode='abs',\n            cooldown=0, \n            min_lr=1e-8,\n            eps=1e-08\n        )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# warmup","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WarmUp(_LRScheduler):\n    \"\"\"warmup_training learning rate scheduler\n    Args:\n        optimizer: optimzier(e.g. SGD)\n        total_iters: totoal_iters of warmup phase\n    \"\"\"\n\n    def __init__(self, optimizer, total_iters, last_epoch=-1):\n        self.total_iters = total_iters\n        super(WarmUp, self).__init__(optimizer, last_epoch)\n\n    def get_lr(self):\n        \"\"\"we will use the first m batches, and set the learning\n        rate to base_lr * m / total_iters\n        \"\"\"\n        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# plabel train code ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")\n\nclass Fitter:\n    \n    def __init__(self, model, device, config, train_loader_length):\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        # warm up\n        if warmup_opt:\n            # iter_per_epoch = math.ceil(len(train_loader) / self.accumulate)\n            iter_per_epoch = train_loader_length\n            total_iters = iter_per_epoch * warmup_epoch\n            self.WarmUp_scheduler = WarmUp(self.optimizer, total_iters)\n        \n        self.log(f'Fitter prepared. Device is {self.device}')\n\n    def fit(self, train_loader, validation_loader):\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            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-checkpoint1.bin')\n\n            # lr sche\n            if warmup_opt and (self.epoch+1) <= warmup_epoch:\n                pass\n            else:\n                if self.config.validation_scheduler:\n                    if PL_lr_sche == 'plat':\n                        self.scheduler.step(metrics=summary_loss.avg)\n                    elif PL_lr_sche == 'cos':\n                        self.scheduler.step()\n            \n#             t = time.time()\n#             summary_loss = self.validation(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\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 warmup_opt and (self.epoch+1) <= warmup_epoch:\n                self.WarmUp_scheduler.step()\n            \n#             if self.config.step_scheduler:\n#                 self.scheduler.step()\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":true},"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# run plabel","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"if PL_OPT:\n    # ---------------------------------- get pred label ---------------------------------------\n    test_models = []\n    for p in path2:\n        test_models.append(load_test_net(p))\n    \n    results_plabel = []\n    for images, image_ids in data_loader:\n        predictions = make_tta_predictions(images, test_models, PL_thr)\n        \n        for i, image in enumerate(images):\n            assert i == 0\n            image_id = image_ids[i]\n            image_ = cv2.imread(f'{TEST_DATA_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n            h,w,_ = np.shape(image_)\n            \n            boxes, scores, labels = run_wbf(predictions, img_size, WBF_iou_thr, WBF_skip_thr)\n            \n            if img_size == 512:\n                boxes = (boxes*2).astype(np.int32).clip(min=0, max=1023)\n            elif img_size == 1024:\n                boxes = boxes.astype(np.int32).clip(min=0, max=1023)\n            \n            image_id = image_ids[i]\n\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()\n    # ------------------------------------ process data --------------------------------------\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_DATA_PATH+'/'+ x+'.jpg')\n\n    marking['image_id'] = marking['image_id'].apply(lambda x: TRAIN_DATA_PATH+'/'+ x+'.jpg')\n    # 把测试集的 plabel 和 训练集拼接到一起\n    train_data_plabel = pd.concat([results_df,marking], axis=0)\n    # 分层交叉验证\n    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\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    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\n    \n    # ------------------------------ dataloader -----------------------------------------\n    fold_number = 0\n\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        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    # ------------------------------- get train model --------------------------------\n    train_models = []\n    for p in path1:\n        train_models.append(load_train_net(p))\n    \n    device = torch.device('cuda:0')\n    \n    fitter = Fitter(model=train_models[0].to(device), device=device, config=TrainGlobalConfig, train_loader_length=len(train_loader))\n    fitter.fit(train_loader, val_loader)\n    \n    del train_models\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# inference","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if PL_OPT:\n    final_models = [\n        load_test_net(f'plabel_model/last-checkpoint1.bin')\n    ]\nelse:\n    final_models = []\n    for p in path1:\n        final_models.append(load_test_net(p))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results = []\n\nfor images, image_ids in data_loader:\n    predictions = make_tta_predictions(images, final_models, OOF_thr)\n    \n    for i, image in enumerate(images):\n        assert i == 0\n        boxes, scores, labels = run_wbf(predictions, img_size, WBF_iou_thr, WBF_skip_thr)\n        \n        if img_size == 512:\n            boxes = (boxes*2).astype(np.int32).clip(min=0, max=1023)\n        elif img_size == 1024:\n            boxes = boxes.astype(np.int32).clip(min=0, max=1023)\n        \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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head(10)","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}