{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Execute EfficientDet TTA with ODAch\nThis is a example to conduct TTA, using a tta-tool **odach** with EfficientDet.\n\nOdach page is bellow, please check it out for details!\n\nhttps://github.com/kentaroy47/ODA-Object-Detection-ttA"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","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# Import odach\nsys.path.insert(0, \"../input/odachkaggle/ODA-Object-Detection-ttA-main/\")\n\nimport odach as oda\n\nfrom ensemble_boxes import *\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=512, width=512, 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/global-wheat-detection/test'\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)\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=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 effdet"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_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=512\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/wheat-effdet5-fold0-best-checkpoint/fold0-best-all-states.bin')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Wrap effdet outputs for odach"},{"metadata":{"trusted":true},"cell_type":"code","source":"class wrap_effdet:\n    def __init__(self, model):\n        self.model = model\n    \n    def __call__(self, img, score_threshold=0.22):\n        det = self.model(img, torch.tensor([1]*images.shape[0]).float().cuda())\n        predictions = []\n        for i in range(img.shape[0]):\n            boxes = det[i][:,:4]  \n            scores = det[i][:,4]\n            npscore = scores.detach().cpu().numpy()\n            indexes = np.where(npscore > score_threshold)[0]\n            boxes = boxes[indexes]\n            boxes[:, 2] = boxes[:, 2] + boxes[:, 0]\n            boxes[:, 3] = boxes[:, 3] + boxes[:, 1]\n            # clamp boxes\n            boxes = boxes.clamp(0, 511)\n            predictions.append({\n                'boxes': boxes[indexes], #/(img.shape[2]-1),\n                'scores': scores[indexes],\n                \"labels\": torch.from_numpy(np.ones_like(npscore[indexes])).cuda()\n            })\n        return predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# wrap effdet\neffdet = wrap_effdet(net.eval())\n# Declare TTA variations\ntta = [oda.HorizontalFlip(), oda.VerticalFlip(), oda.Rotate90()]\n# Declare scales to tta\nscale = [1]\n# wrap model and tta\ntta_model = oda.TTAWrapper(effdet, tta, scale)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference on a image"},{"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\nimages2 = images.cuda().float()\nboxes, scores, labels = tta_model(images2)\n\n# Visualize\ni = 0\nsample = images[i].permute(1,2,0).cpu().numpy()\nsample = cv2.UMat(sample).get()\nboxes = (boxes*(images.shape[2]-1)).astype(np.int32).clip(min=0, max=511)\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (1,0,0), 2)\n    \nax.set_axis_off()\nax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Submit!"},{"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    images = images.cuda()\n    i=0\n    boxes, scores, labels = tta_model(images)\n    boxes = (boxes*1023).astype(np.int32).clip(min=0, max=1023)\n    image_id = image_ids[i]\n    \n    if len(data_loader)<30:\n        i = 0\n        fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n        sample = images[i].permute(1,2,0).cpu().numpy()\n        sample = cv2.UMat(sample).get()\n        for box in boxes:\n            box = (box/2).astype(np.int32)\n            cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (1, 0, 0), 2)\n\n        ax.set_axis_off()\n        ax.imshow(sample);\n\n    boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n    boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n    \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":"","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}