{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport matplotlib.pyplot as plt\nimport cv2\nimport os\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n\n#for dirname, _, filenames in os.walk('/kaggle/input'):\n#    for filename in filenames:\n#        print(os.path.join(dirname, filename))\n\n# You can write up to 5GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom PIL import Image, ImageDraw, ImageFont\nimport torchvision\nimport torchvision.transforms as T\nimport math","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!cp -r /kaggle/input/detr-wheat/ /root/.cache/torch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ls /root/.cache/torch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DETRModel(nn.Module):\n    def __init__(self,num_classes,num_queries):\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', 'detr_resnet50', pretrained=True,force_reload=False)\n        #self.model = torch.hub.load('facebookresearch/detr', 'detr_resnet101', pretrained=True)\n        #self.model = torch.hub.load('facebookresearch/detr', 'detr_resnet101_dc5', pretrained=True)\n        self.in_features = self.model.class_embed.in_features\n        self.hidden_dim = self.model.transformer.d_model\n        \n        self.model.class_embed = nn.Linear(in_features=self.in_features,out_features=self.num_classes+1)\n        self.model.num_queries = self.num_queries\n        self.model.query_embed = nn.Embedding(self.num_queries, self.hidden_dim)\n        \n    def forward(self,images):\n        return self.model(images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = DETRModel(num_classes=1,num_queries=150)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.load_state_dict(torch.load(\"../input/detr-wheat/tf_detr_best_0_round22_best_class.pth\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n#model.to('cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"inf_transform = T.Compose([\n    T.ToTensor(),\n    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img = Image.open('../input/global-wheat-detection/test/f5a1f0358.jpg').resize((480,360))\nimw, imh = img.size\nimw, imh","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ratio = imh/imw","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img_inf = img.resize((1024,int(1024 * ratio))).convert('RGB')\nimg_inf","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img_tens = inf_transform(img_inf).unsqueeze(0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with torch.no_grad():\n    outputs = model(img_tens)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im2 = img.copy()\ndrw = ImageDraw.Draw(im2)\n\noboxes = outputs['pred_boxes'][0] * torch.Tensor([imw, imh, imw, imh])\nprob   = outputs['pred_logits'][0].softmax(1).detach().cpu().numpy()\nclasses = np.argmax(prob, axis=1)\nprob   = np.max(prob, axis=1)\n\n#print(classes)\n#print(prob)\n\n\nfor box,p,c in zip(oboxes,prob,classes):\n\n    if (p > 0.5) & (c!=1):\n        box = box.cpu()\n        print(p, box)        \n        x, y, w, h = box\n        x0, x1 = x-w//2, x+w//2\n        y0, y1 = y-h//2, y+h//2\n        drw.rectangle([x0, y0, x1, y1], outline='red', width=2)\n        drw.text((x0+4, y0+4), '%.1f'%(p*100), fill='white')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_files = os.listdir('../input/global-wheat-detection/test/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_files","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for file in test_files:\n    img = Image.open('../input/global-wheat-detection/test/'+file)\n    break\nimg","execution_count":null,"outputs":[]},{"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 =[]\nfig, ax = plt.subplots(5, 2, figsize=(30, 70))\ncount = 0\n\nfor file in test_files:\n    img = Image.open('../input/global-wheat-detection/test/'+file)\n    im_w, im_h = img.size\n    ratio = im_h / im_w\n    img = img.resize((608,int(608*ratio)))\n    img = img.resize((800,int(800*ratio))).convert('RGB')\n    img_tens = inf_transform(img).unsqueeze(0)\n    with torch.no_grad():\n        outputs = model(img_tens)\n        \n    try:\n        oboxes = outputs['pred_boxes'][0]\n        prob   = outputs['pred_logits'][0].softmax(1).detach().cpu().numpy()\n        classes = torch.from_numpy(np.argmax(prob, axis=1))\n        probs   = torch.from_numpy(np.max(prob, axis=1)).squeeze(0).squeeze(0)\n\n        pred_strings = []\n\n        mask_classes = classes !=1\n        idx_classes = torch.nonzero(mask_classes)\n\n        oboxes=oboxes[idx_classes].permute(1,0,2).squeeze(0)\n        probs=probs[idx_classes].permute(1,0).squeeze(0)\n\n        mask_probs = probs > 0.8\n        idx_probs = torch.nonzero(mask_probs)\n\n        oboxes=oboxes[idx_probs].permute(1,0,2).squeeze(0)#.detach().cpu().numpy()\n        probs=probs[idx_probs].permute(1,0).squeeze(0).detach().cpu().numpy()\n\n        oboxes=oboxes * torch.Tensor([im_w, im_h, im_w, im_h])\n        oboxes=oboxes.detach().cpu().numpy()#.astype(np.int32)\n        \n        num_boxes = oboxes.shape[0]\n        \n        transf = np.concatenate((oboxes[:, [2, 3]]/2, np.zeros((num_boxes, 2))), axis=1)     \n\n        #print(oboxes)\n        \n        oboxes = (oboxes[:, [0, 1, 2, 3]] - transf).astype(np.int32)#.clip(min=0, max=1024)\n        \n        oboxes = np.clip(oboxes, a_min = 0, a_max = [im_w, im_h, im_w, im_h]) \n        \n        #oboxes = oboxes.astype(np.int32)\n        #print(oboxes)\n        \n        #transf = oboxes[:, [2, 3]]//2\n        \n\n        print(file, len(oboxes), len(probs))\n\n    except Exception as e:\n        oboxes=np.array([])\n        probs=np.array([])\n        #print('erro')\n        #print(str(e))\n        \n       \n    image_id = file.split(\".\")[0]\n    \n    result = {'image_id': image_id,'PredictionString': format_prediction_string(oboxes, probs)}\n    results.append(result)    \n      \n    #if file == '2fd875eaa.jpg':\n    if count < 10:\n        img_ = cv2.imread('../input/global-wheat-detection/test/'+file)\n        img_ = cv2.cvtColor(img_, cv2.COLOR_BGR2RGB)\n        #print('ok')\n        #img = Image.open('../input/global-wheat-detection/test/'+file)\n        #im2 = img.copy()\n        #drw = ImageDraw.Draw(im2)\n        \n        for box,p in zip(oboxes,probs):\n            \n            x, y, w, h = box\n            #x0, x1 = x-w//2, x+w//2\n            #y0, y1 = y-h//2, y+h//2\n            \n            x0, x1 = x, x+w\n            y0, y1 = y, y+h\n            \n            #print(x0, y0, x1, y1)\n            \n            #drw.rectangle([x0, y0, x1, y1], outline='red', width=2)\n            #drw.text((x0+4, y0+4), '%.1f'%(p*100), fill='white')\n            cv2.rectangle(img_, (x0, y0), (x1, y1), (220, 0, 0), 2)\n            cv2.putText(img_, '%.1f'%(p*100), (x0, y0), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2, cv2.LINE_AA)\n            ax[count%5][count//5].imshow(img_)\n        count+=1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.sort_values(by='image_id', inplace=True)\ntest_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm -rf ./*","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)\ntest_df.head(10)","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}