{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install --no-index ../input/ensemble-boxes/ensemble_boxes-1.0.4-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os, json, time, random, operator, functools\n\nfrom tqdm.notebook import tqdm\nfrom IPython import display\nfrom ipywidgets import Output\n\nimport pandas as pd\nimport numpy as np\nfrom skimage import io\nfrom PIL import Image, ImageDraw \nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.utils\nimport torchvision \nfrom torchvision import transforms\nimport torch.nn as nn\nimport torch.optim as optim\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport ensemble_boxes","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"path = '../input/global-wheat-detection'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\ncpu = torch.device(\"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class wheatDataset(object):\n    def __init__(self, root, mode, transform = None):\n        self.root = root\n        self.transform = transform\n        self.mode = mode\n\n        self.imgs = sorted(list(os.listdir(os.path.join(root, mode))))\n        \n        if mode == \"train\":\n          self.table = pd.read_csv(os.path.join(root, \"train.csv\"))\n          self.table.bbox = self.table.bbox.map(lambda x: json.loads(x))\n          self.table = self.table.groupby(['image_id'])\n          self.ID = self.table.groups.keys()\n        #root - train.csv\n        #     |- train - image\n\n    def __getitem__(self, idx):\n        # load images\n        image_name = self.imgs[idx]\n        image_id = image_name[:-4]\n\n        if (self.mode == 'train'): \n          if (not(image_id in self.ID)):\n              return self.__getitem__(random.randint(0,len(self.imgs)-1))##avoid no wheat\n        \n        img_path = os.path.join(self.root, self.mode, image_name)\n        img = np.asarray(Image.open(img_path).convert(\"RGB\"))\n\n        if (self.mode == 'test'):\n          return transforms.ToTensor()(img)\n\n        # get bounding box\n        boxes_list = np.array(self.table.get_group(image_id).bbox.to_list())\n        boxes_list[:,2] += boxes_list[:,0]#x2 = x1 + box width\n        boxes_list[:,3] += boxes_list[:,1]#y2 = y1 + box hight\n        \n\n#         # convert everything into a torch.Tensor\n#         boxes = torch.as_tensor(boxes_list, dtype=torch.float32)\n#         # there is only one class\n#         label = torch.ones((len(boxes_list),), dtype=torch.int64)\n#         image_id = torch.tensor([idx])\n        \n        # not crowd\n#         iscrowd = torch.zeros((1,), dtype=torch.int64)\n        boxes = boxes_list\n        label = np.ones((len(boxes_list),), dtype=np.int64)\n        \n        target = {}\n        if self.transform is not None:\n            transformed = self.transform(image = img, bboxes=boxes, labels = label)\n            while(len(transformed['bboxes']) == 0):\n                transformed = self.transform(image = img, bboxes=boxes, labels = label)\n            \n            img = transforms.ToTensor()(transformed['image'])\n            target[\"boxes\"] = torch.as_tensor(np.asarray(transformed['bboxes']), dtype=torch.float32)\n            target[\"labels\"] = torch.as_tensor(np.asarray(transformed['labels']), dtype=torch.int64)\n        else:\n            img = transforms.ToTensor()(img)\n            target[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.float32)\n            target[\"labels\"] = torch.as_tensor(label, dtype=torch.int64)\n            \n        \n        return img, target\n\n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained_backbone=False)\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nnum_classes = 2  # wheat + background\nmodel.roi_heads.box_predictor = FastRCNNPredictor(model.roi_heads.box_predictor.cls_score.in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.load_state_dict(torch.load(\"../input/fastrcnnwheat/Trained-wheat-fastrcnn\"))\nmodel.to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# test","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_test = wheatDataset(path, \"test\")\ndata_loader = torch.utils.data.DataLoader(\n        dataset_test, batch_size=10, shuffle=False, num_workers=5, pin_memory = True,\n        collate_fn=lambda batch: batch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def test_numpy():\n  result = list()\n  for images in tqdm(data_loader):\n    images = list(image.to(device) for image in images)\n    model.eval()\n    predictions = model(images)\n    \n    images = list(image.to(cpu).detach().numpy() for image in images)\n    predictions = list({k:v.to(cpu).detach().numpy() for k, v in p.items()} for p in predictions)\n    \n    result.append(predictions)\n#     for x,y in zip(images, predictions):\n#       drawBondingBox_test(x,y['boxes'])\n  return result\n\nresult = test_numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result = functools.reduce(operator.add, result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"IDs = list(map(lambda n: n[:-4], sorted(list(os.listdir(os.path.join(path, 'test'))))))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def drawBondingBox_test(image_name, listOfbox):\n    img_path = os.path.join(path, 'test', image_name + '.jpg')\n    img = Image.open(img_path).convert(\"RGB\")\n    draw = ImageDraw.Draw(img) \n    for b in listOfbox:\n        x0,y0,x1,y1 = b\n        draw.rectangle((x0,y0,x1,y1), fill=None, outline=\"#FA6E1A\", width=4)\n    plt.figure(figsize = (10,10))\n    plt.imshow(img, interpolation='nearest')\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n = 0\nsubmission = pd.DataFrame(columns=['image_id', 'PredictionString'])\nfor predictions in result:\n    y = predictions\n    # print(IDs[n])\n    \n    boxes = [(b/1024).tolist() for b in y['boxes']]\n    scores = y['scores'].tolist()\n    labels = y['labels'].astype(int).tolist()\n    \n    \n    boxes, scores, labels = ensemble_boxes.ensemble_boxes_wbf.weighted_boxes_fusion([boxes], [scores], [labels], weights=None, iou_thr=0.44, skip_box_thr=0.43)\n    boxes = (np.array(boxes)*1024).astype(int)\n#     drawBondingBox_test(IDs[n], boxes)\n    PredictionString = \"\"\n    for b, s in zip(boxes, scores):\n        PredictionString += \"{:.4f} {:.0f} {:.0f} {:.0f} {:.0f} \".format(s, b[0], b[1], b[2]-b[0], b[3]-b[1])\n    PredictionString = PredictionString[:-1]\n    submission.loc[n] = [IDs[n], PredictionString]\n    n+=1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission.to_csv(\"submission.csv\",index=False)","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}