{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-output":true,"_kg_hide-input":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":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import sys\n#sys.path.insert(0, \"../input/effdetpipeline/timm-efficientdet-pytorch/timm-efficientdet-pytorch\")\nsys.path.insert(0, \"../input/omegaconf\")\nsys.path.insert(0, \"../input/weightedboxesfusion\")\nsys.path.insert(0, \"../input/timm-efficientdet-pytorch/\")\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\nfrom effdet.efficientdet import HeadNet","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_valid_transforms():\n    return A.Compose([\n            A.Resize(height=1024, width=1024, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_ROOT_PATH = \"../input/wheatspike/new/new\"\n\nclass DatasetRetriever(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)[:1024,:1024,:]\n        #print(image.shape)\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":{"trusted":true},"cell_type":"code","source":"dataset = DatasetRetriever(\n    image_ids=np.array([path.split('/')[-1][:-4] for path in glob(f'{DATA_ROOT_PATH}/*.jpg')]),\n    transforms=get_valid_transforms()\n)\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndata_loader = DataLoader(\n    dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=2,\n    drop_last=False,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_net(checkpoint_path):\n    config = get_efficientdet_config('tf_efficientdet_d6')\n    net = EfficientDet(config, pretrained_backbone=False)\n\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(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()\n\nnet = load_net('../input/effdet5fold/stac.bin')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Custom TTA API\n\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":{"trusted":true},"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 TTARotate180(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 2, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 2, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        boxes[:, [0,1,2,3]] = self.image_size - boxes[:, [2,3,0,1]]\n        return boxes\n    \nclass TTARotate270(BaseWheatTTA):\n    \"\"\" author: @shonenkov \"\"\"\n    \n    def augment(self, image):\n        return torch.rot90(image, 3, (1, 2))\n\n    def batch_augment(self, images):\n        return torch.rot90(images, 3, (2, 3))\n    \n    def deaugment_boxes(self, boxes):\n        res_boxes = boxes.copy()\n        res_boxes[:, [0,2]] = boxes[:, [1,3]]\n        res_boxes[:, [1,3]] = self.image_size - 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":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"def process_det(index, det, score_threshold=0.25):\n    boxes = det[index].detach().cpu().numpy()[:,:4]    \n    scores = det[index].detach().cpu().numpy()[:,4]\n    boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n    boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\n    boxes = (boxes).clip(min=0, max=1023).astype(int)\n    indexes = np.where(scores>score_threshold)\n    boxes = boxes[indexes]\n    scores = scores[indexes]\n    return boxes, scores","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# you can try own combinations:\ntransform = TTACompose([\n    TTARotate90(),\n    TTAVerticalFlip(),\n])\n\nfig, ax = plt.subplots(1, 3, figsize=(16, 6))\n\nimage, image_id = dataset[5]\n\nnumpy_image = image.permute(1,2,0).cpu().numpy().copy()\n\nax[0].imshow(numpy_image);\nax[0].set_title('original')\n\ntta_image = transform.augment(image)\ntta_image_numpy = tta_image.permute(1,2,0).cpu().numpy().copy()\n\ndet = net(tta_image.unsqueeze(0).float().cuda(), torch.tensor([1]).float().cuda())\nboxes, scores = process_det(0, det)\n\nfor box in boxes:\n    cv2.rectangle(tta_image_numpy, (box[0], box[1]), (box[2],  box[3]), (0, 1, 0), 2)\n\nax[1].imshow(tta_image_numpy);\nax[1].set_title('tta')\n    \nboxes = transform.deaugment_boxes(boxes)\n\nfor box in boxes:\n    cv2.rectangle(numpy_image, (box[0], box[1]), (box[2],  box[3]), (0, 1, 0), 2)\n    \nax[2].imshow(numpy_image);\nax[2].set_title('deaugment predictions');","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(), TTARotate180(), TTARotate270(), 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, score_threshold=0.37):\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.541, skip_box_thr=0.406, 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.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":{"trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nfor j, (images, image_ids) in enumerate(data_loader):\n    break\n\npredictions = make_tta_predictions(images)\n\ni = 1\nsample = images[i].permute(1,2,0).cpu().numpy()\n\nboxes, scores, labels = run_wbf(predictions, image_index=i)\nboxes = boxes.round().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), 1)\n\nax.set_axis_off()\nax.imshow(sample);","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":"results = []\n\nfor images, image_ids in data_loader:\n    predictions = make_tta_predictions(images)\n    for i, image in enumerate(images):\n        boxes, scores, labels = run_wbf(predictions, image_index=i)\n        boxes = (boxes).round().astype(np.int32).clip(min=0, max=1023)\n        image_id = image_ids[i]\n\n        score_threshold = 0.125\n        indexes = np.where(scores>score_threshold)\n        boxes = boxes[indexes]\n        scores = scores[indexes]\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()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"a = 0\nfor i in range(10):\n    b = (len(test_df.iloc[i].PredictionString.split(' '))/5)\n    print(b)\n    a+=b\nprint('====')\nprint(a/10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfig, ax = plt.subplots(1, 1, figsize=(30, 30))\n#fig.axes.get_xaxis().set_visible(False)\n#fig.axes.get_yaxis().set_visible(False)\nplt.axis('off')\nDIR_TEST = \"../input/wheatspike/SPIKE Dataset/positive\"\n\nnum = 8\n\nimg_pth = os.path.join(DIR_TEST, test_df.image_id[num]+'.jpg')\n#print(img_pth)\nimage = cv2.imread(img_pth)\nb = [float(i) for i in test_df.PredictionString[num].split(' ')]\nboxes = []\ni = 0\nwhile i < len(b):\n    #print(b[i])\n    boxes.append([int(b[i+1]),int(b[i+2]),int(b[i+3]),int(b[i+4]),float(b[i])])\n    #print(boxes)\n    i+=5\nboxes = np.array(boxes)\nfor box in boxes.astype(float):\n    cv2.rectangle(image,(int(box[0]), int(box[1])),(int(box[0])+int(box[2]),  int(box[1])+int(box[3])),(0, 32, 255), 3)\n    #print(box[4])\n    cv2.putText(image, str(box[4]), (int(box[0]), int(box[1])+20), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)\nax.imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\n\nfig, ax = plt.subplots(3, 3, figsize=(30, 30),gridspec_kw = {'wspace':0, 'hspace':0})\nplt.axis('off')\n\n#fig.axes.get_xaxis().set_visible(False)\n#fig.axes.get_yaxis().set_visible(False)\nDIR_TEST = \"../input/wheatspike/SPIKE Dataset/positive\"\n\nfor i1 in range(3):\n    for j1 in range(3):\n        num = j1*3+i1\n\n        img_pth = os.path.join(DIR_TEST, test_df.image_id[num]+'.jpg')\n        #print(img_pth)\n        image = cv2.imread(img_pth)\n        b = [float(i) for i in test_df.PredictionString[num].split(' ')]\n        boxes = []\n        i = 0\n        while i < len(b):\n            boxes.append([int(b[i+1]),int(b[i+2]),int(b[i+3]),int(b[i+4]),float(b[i])])\n            i+=5\n        boxes = np.array(boxes)\n        for box in boxes.astype(float):\n            cv2.rectangle(image,(int(box[0]), int(box[1])),(int(box[0])+int(box[2]),  int(box[1])+int(box[3])),(0, 32, 255), 3)\n            #print(box[4])\n            cv2.putText(image, str(box[4]), (int(box[0]), int(box[1])+20), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)\n        ax[i1][j1].imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}