{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport os\nimport glob\nimport cv2\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.transforms as transforms\nfrom PIL import Image\nfrom torch.utils.data import dataloader,dataset\nimport albumentations as A\nfrom albumentations.pytorch import ToTensor \n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nimport warnings\nwarnings.filterwarnings(\"ignore\")\ntest_path = '../input/global-wheat-detection/test'\nmodel_path = '../input/modelweight'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.read_csv('../input/global-wheat-detection/sample_submission.csv')\ntest_df.shape","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\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class test_dataset(torch.utils.data.Dataset):\n  \n  def __init__(self, image_list, transforms=None):\n    super().__init__()\n    self.images = image_list\n    self.transforms = transforms\n    \n\n  def __len__(self):\n    return len(self.images)\n\n  def __getitem__(self,item):\n    image = self.images[item]\n    \n    base_name = os.path.basename(image)\n    image_name_split = os.path.splitext(base_name)\n    image_id = image_name_split[0] \n    \n    image = cv2.imread(image, cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)\n    image = image/255.0\n    image = np.transpose(image,(2,1,0))\n    \n    #if self.transforms:\n      #sample = {'image':image}\n    #  image = self.transforms(image)\n\n    #image = sample['image']\n    #image = self.transforms(image)\n\n    return torch.as_tensor(image,dtype=torch.float32),image_id","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images_list= glob.glob(os.path.join(test_path,'*.jpg'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#transform = A.Compose([ToTensor()])\ntest_data=test_dataset(images_list)\ntest_loader = torch.utils.data.DataLoader(test_data,batch_size=4,shuffle=False,collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False,pretrained_backbone=False)\nnum_classes = 2\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.load_state_dict(torch.load(os.path.join(model_path,'resnet50_GWD_6.pth')))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nmodel = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\ndetection_threshold = 0.5\nresults = []\n\nfor images,image_ids in test_loader:\n\n    images_12 = list(image.to(device) for image in images)\n    with torch.no_grad():\n        outputs = model(images_12)\n    #outputs = model(images)\n\n        for i, image in enumerate(images_12):\n\n            boxes = outputs[i]['boxes'].cpu().numpy()\n            scores = outputs[i]['scores'].cpu().numpy()\n        \n            boxes = boxes[scores >= detection_threshold].astype(np.int32)\n            scores = scores[scores >= detection_threshold]\n            image_id = image_ids[i]\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\n        \n            results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df['PredictionString']=results\ntest_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.head(10\n            )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = outputs[0]['boxes'].data.cpu().numpy()\nscores = outputs[0]['scores'].data.cpu().numpy()\n\nboxes = boxes[scores >= detection_threshold].astype(np.int32)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nimage = cv2.imread(os.path.join(test_path,'f5a1f0358'+'.jpg'),cv2.IMREAD_COLOR)\nimage = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in boxes:\n  image = cv2.rectangle(image,(i[0],i[1]),(i[2],i[3]),(220,0,0),3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig,ax = plt.subplots(1,1,figsize=(15,15))\nax.imshow(image)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)","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}