{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"# cloning detr's github repository\n!git clone https://github.com/facebookresearch/detr.git","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Importing our Dependencies","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np \nimport pandas as pd \nfrom datetime import datetime\nimport time\nimport random\nimport sys\nimport numba  # python and numpy code acceleartor\nfrom tqdm.autonotebook import tqdm\n\n#Torch\nimport torch\nimport torch. nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\n\n#CV\nimport cv2\n\n################# DETR FUCNTIONS FOR LOSS######################## \n# losses used by detr\nsys.path.append('../input/detrfiles/results/')\nsys.path.append('../input/detrfiles/results/detr/')\n\n\nfrom detr.models.matcher import HungarianMatcher\nfrom detr.models.detr import SetCriterion\n################################################################\n\n#Albumenatations\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom albumentations.pytorch.transforms import ToTensorV2\n\n#Glob\nfrom glob import glob\n\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom PIL import Image, ImageDraw\nfrom IPython.display import display","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# returns imgages in the form of a batch\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# transforms to be applied on the test images\ndef get_test_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\ndef change_to_xmin_ymin_wh(bboxes,size):\n    \"\"\"A utility function to convert bounding boxes output by detr to (x_min, y_min, w, h) format\"\"\"\n    \n    box = np.array(bboxes)\n    xy , wh = box[:,[0,1]] , box[:,[2,3]]\n    xy = xy-(wh/2)\n    xywh = np.concatenate((xy,wh),axis=1)\n    \n    # denormalise bbox\n    xywh = xywh*size\n    return xywh","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Creating Test Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_TEST = \"../input/global-wheat-detection/test\"\n\nclass WheatTestDataset(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        # reading and processing image\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        # applying transforms\n        if self.transforms:\n            sample = {\n                'image': image,\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\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(\"../input/global-wheat-detection/sample_submission.csv\")\nimage_ids = test_df.image_id.unique()\nprint(test_df.shape)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# path to trained weights\nWEIGHTS_FILE = \"../input/detrfiles/detr_best_3.pth\"\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nnum_classes = 2\nnum_queries = 100\n\n# instantiating the model\nmodel = DETRModel(num_classes=num_classes, num_queries=num_queries, model_name='detr_resnet50')\nmodel = model.to(device)\n\n# loading trained weights for predictions\nmodel.load_state_dict(torch.load(WEIGHTS_FILE))\nmodel.eval()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"clf_gt = []\n\n# creating the test dataset\ntest_dataset = WheatTestDataset(\n    image_ids=image_ids,\n    dataframe=test_df,\n    transforms=get_test_transforms()\n)\n\n# loading the test dataset batch-wise\ndata_test_loader = DataLoader(\n    test_dataset, \n    batch_size=5, \n    collate_fn=collate_fn,\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Making predictions","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"bounding_boxes = {}\n\nconfidence_thrsh = 0.5\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   \n    outputs = {k: v.to('cpu') for k, v in outputs.items()} # dictionary of outputs\n     \n    for j, (image_name, bboxes, logits) in enumerate(zip(image_ids, outputs['pred_boxes'], outputs['pred_logits'])):\n\n        # denormalizing and scaling  boxes \n        oboxes = bboxes.detach().cpu()\n        \n        # converting to required format\n        oboxes = change_to_xmin_ymin_wh(oboxes,512)     \n        \n        # rescaling the bounding boxes\n        oboxes = (oboxes*2).astype(np.int32).clip(min=0,max=1023)\n        \n        # applying softmax on logits\n        prob   = logits.softmax(1).detach().cpu().numpy()[:, 0]\n        \n        # creating a dictionary for submission\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            })\n        \n        temp = []\n        \n        # considering bounding boxes with confidence greater than a threshold\n        for box, confidence in zip(oboxes, prob):\n            if confidence > confidence_thrsh:\n                temp.append(box)\n                \n        bounding_boxes[image_name]=temp","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)\nsubmission_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def display_bboxes(img, bboxes):\n    draw = ImageDraw.Draw(box_img)\n    for box in bboxes:\n        x,y,w,h = box\n        x0,x1 = x,x+w\n        y0,y1 = y,y+h\n\n        draw.rectangle([x0, y0, x1, y1], outline = \"blue\", width = 2)\n    display(box_img)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualizing the predicted Bounding Boxes","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"image_id = image_ids[0]\nboxes = bounding_boxes[image_id]\nimg = Image.open(\"../input/global-wheat-detection/test/\"+image_id+\".jpg\")\nbox_img = img.copy()\ndisplay_bboxes(box_img, boxes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_id = image_ids[1]\nboxes = bounding_boxes[image_id]\nimg = Image.open(\"../input/global-wheat-detection/test/\"+image_id+\".jpg\")\nbox_img = img.copy()\ndisplay_bboxes(box_img, boxes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_id = image_ids[2]\nboxes = bounding_boxes[image_id]\nimg = Image.open(\"../input/global-wheat-detection/test/\"+image_id+\".jpg\")\nbox_img = img.copy()\ndisplay_bboxes(box_img, boxes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_id = image_ids[3]\nboxes = bounding_boxes[image_id]\nimg = Image.open(\"../input/global-wheat-detection/test/\"+image_id+\".jpg\")\nbox_img = img.copy()\ndisplay_bboxes(box_img, boxes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_id = image_ids[4]\nboxes = bounding_boxes[image_id]\nimg = Image.open(\"../input/global-wheat-detection/test/\"+image_id+\".jpg\")\nbox_img = img.copy()\ndisplay_bboxes(box_img, boxes)","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}