{"cells":[{"metadata":{},"cell_type":"markdown","source":"# This notebook is the inference for DETR model\n\nIf you want to know how to train it, I suggest you look at this\n* [End to End Object Detection with Transformers:DETR](https://www.kaggle.com/tanulsingh077/end-to-end-object-detection-with-transformers-detr/notebook)\n\nAlso, I've created `detrmodels` dataset cloned from [facebookresearch](https://github.com/facebookresearch/detr/) including its pretrained weights","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Import Libs","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport sys\nfrom pathlib import Path\n\nimport albumentations as A\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom albumentations.pytorch import ToTensorV2\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\n\n\nsys.path.append('../input/detrmodels/facebookresearch_detr_master/')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Copy Pretrained Weights","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# copy pretrained weights to the folder PyTorch will search by default\nPath('/root/.cache/torch/hub/').mkdir(exist_ok=True, parents=True)\nPath('/root/.cache/torch/hub/checkpoints/').mkdir(exist_ok=True, parents=True)\n\ndetr_path = '/root/.cache/torch/hub/checkpoints/detr-r50-e632da11.pth'\nresnet50_pretrained = '/root/.cache/torch/hub/checkpoints/resnet50-19c8e357.pth'\ndetr_hub = '/root/.cache/torch/hub/facebookresearch_detr_master'\n\n!cp ../input/detrmodels/detr-r50-e632da11.pth {detr_path}\n!cp ../input/detrmodels/resnet50-19c8e357.pth {resnet50_pretrained}\n!cp -R ../input/detrmodels/facebookresearch_detr_master {detr_hub}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TEST = f'{DIR_INPUT}/test'\n\nDIR_WEIGHTS = '../input/detrglobalwheatbestmodels/'\nWEIGHTS_FILE = f'{DIR_WEIGHTS}/detr_best_0.pth'\n\nbatch_size = 5","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\n\ndef get_valid_transforms():\n    return A.Compose([A.Resize(height=512, width=512, p=1.0),\n                      ToTensorV2(p=1.0)], \n                      p=1.0, \n                    )\n\n\nclass WheatDataset(Dataset):\n    def __init__(self,image_ids,dataframe,transforms=None):\n        self.image_ids = image_ids\n        self.df = dataframe\n        self.transforms = transforms\n        \n        \n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    \n    def __getitem__(self,index):\n        image_id = self.image_ids[index]\n        \n        image = cv2.imread(f'{DIR_TEST}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        \n        if self.transforms:\n            sample = {\n                'image': image,\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n    \n        \n        return image, image_id","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class DETRModel(nn.Module):\n    def __init__(self, num_classes, num_queries, model_name='detr_resnet50'):\n        super(DETRModel, self).__init__()\n        self.num_classes = num_classes\n        self.num_queries = num_queries\n\n        self.model = torch.hub.load('facebookresearch/detr', model_name, pretrained=True)\n        self.in_features = self.model.class_embed.in_features\n\n        self.model.class_embed = nn.Linear(in_features=self.in_features,\n                                           out_features=self.num_classes)\n        self.model.num_queries = self.num_queries\n\n    def forward(self, images):\n        return self.model(images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.read_csv(f'{DIR_INPUT}/sample_submission.csv')\ntest_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = test_df.image_id.unique(); image_ids","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nnum_classes = 2\nnum_queries = 100\nmodel = DETRModel(num_classes=num_classes, num_queries=num_queries, model_name='detr_resnet50')\nmodel = model.to(device)\nmodel.load_state_dict(torch.load(WEIGHTS_FILE))\nmodel.eval()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predict and Visualise bboxes","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"clf_gt = []\n\ntest_dataset = WheatDataset(\n    image_ids=image_ids,\n    dataframe=test_df,\n    transforms=get_valid_transforms()\n)\n\ndata_test_loader = DataLoader(\n        test_dataset, batch_size=batch_size, \n        collate_fn=collate_fn,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"rows = 2\ncolumns = 5\nconfidence_thrsh = 0.5\nfig, ax = plt.subplots(rows, columns, figsize=(15, 8))\nfor i, (images, image_ids) in enumerate(tqdm(data_test_loader)):\n    with torch.no_grad():\n        images = list(image.to(device, dtype=torch.float) for image in images)\n        outputs = model(images)\n    outputs = [{k: v.to('cpu') for k, v in outputs.items()}][0]\n    \n    row = i % rows\n    for j, (image_name, bboxes, logits) in enumerate(zip(image_ids, outputs['pred_boxes'], outputs['pred_logits'])):\n        column = j % columns\n\n        oboxes = bboxes.detach().cpu().numpy()\n        oboxes = np.array([\n            np.array(box).astype(np.int32) \n            for box in A.augmentations.bbox_utils.denormalize_bboxes(oboxes,512,512)\n        ])\n        prob   = logits.softmax(1).detach().cpu().numpy()[:, 0]\n        # scale boxes \n        oboxes = (oboxes*2).astype(np.int32).clip(min=0, max=1023)\n        \n        # Comment this out of you dont want to plot bboxes #\n        sample = cv2.imread(f'{DIR_TEST}/{image_name}.jpg', cv2.IMREAD_COLOR)\n        sample = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB)\n        for box,p in zip(oboxes,prob):\n            if p > confidence_thrsh:\n                color = (0,0,220)\n                cv2.rectangle(sample,\n                      (box[0], box[1]),\n                      (box[2]+box[0], box[3]+box[1]),\n                      color, 3)\n    \n        ax[row, column].set_axis_off()\n        ax[row, column].imshow(sample)\n        # ================================================= #\n        \n        clf_gt.append({\n                'image_id': image_name,\n                'PredictionString': ' '.join(\n                    str(round(confidence,4)) \n                    + ' '\n                    + ' '.join(str(int(round(float(x)))) for x in box) \n                    for box, confidence in zip(oboxes, prob)\n                    if confidence > confidence_thrsh\n                )\n                ,\n            })","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission_df = pd.DataFrame(clf_gt)\nsubmission_df['PredictionString'] = submission_df['PredictionString'].fillna('')\nsubmission_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission_df","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}