{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.utils.data\nimport torchvision\nimport cv2\nfrom PIL import Image\nfrom matplotlib import pyplot as plt\n#from pycocotools.coco import COCO\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom torch.utils.data import DataLoader","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# path to your own data and coco file\ntrain_data_dir = \"data/upload\"\ntrain_coco = \"object_detaction__wheat_ml/abc.json\"\n\nModel_path = \"../input/trained-model/fasterrcnn_resnet50.pth\"\n\n# Batch size\ntrain_batch_size = 1\n\n# Params for dataloader\ntrain_shuffle_dl = True\nnum_workers_dl = 4\n\n# Params for training\n\n# Two classes; Only target class or background\nnum_classes = 2\nnum_epochs = 2\n\nlr = 0.005\nmomentum = 0.9\nweight_decay = 0.005\n\nTEST_FOLDER = \"../input/global-wheat-detection/test\"\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\n\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n\ntest_df = pd.read_csv(f'{DIR_INPUT}/sample_submission.csv')\ntest_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class myOwnDataset(torch.utils.data.Dataset):\n    \n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\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\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# In my case, just added ToTensor\ndef get_transform():\n    return A.Compose([\n        # A.Resize(512, 512),\n        ToTensorV2(p=1.0)\n    ])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntest_dataset = myOwnDataset(test_df, DIR_TEST, get_transform())\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=4,\n    drop_last=False,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# collate_fn needs for batch\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n\ndef get_model_instance_segmentation(num_classes):\n    # load an instance segmentation model pre-trained pre-trained on COCO\n    model =  torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False,\n                                                                  pretrained_backbone=False)\n    # get number of input features for the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    # replace the pre-trained head with a new one\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = get_model_instance_segmentation(num_classes)\nmodel.load_state_dict(torch.load(Model_path))\nmodel.eval()\nx = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def prediction_sting(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)\n\ndef predict(filepath, image_id, detection_threshold=0.3):\n    with open(filepath, mode='rb') as f:\n        image = Image.open(f).convert('RGB')\n    image_height, image_width = image.size\n    \n    trasformer = get_transform()    \n    trasform_image = trasformer(image)\n    images = torch.stack([trasform_image])\n    #with torch.no_grad():\n    results = model(images)\n    boxes = results[0]['boxes'].data.cpu().numpy()\n    scores = results[0]['scores'].data.cpu().numpy()\n    #print(boxes)\n    #print(scores)\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': prediction_sting(boxes, scores)\n    }\n    return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"detection_threshold = 0.40\n\nresults = []\nfor images, image_ids in test_data_loader:\n\n    images = list(image.to(device) for image in images)\n    outputs = model(images)\n\n    for i, image in enumerate(images):\n\n        boxes = outputs[i]['boxes'].data.cpu().numpy()\n        scores = outputs[i]['scores'].data.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': prediction_sting(boxes, scores)\n        }\n\n        \n        results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"all_result = []\nfor dirname, _, filenames in os.walk(TEST_FOLDER):\n    for filename in filenames:\n        image_path = os.path.join(dirname, filename)\n        image_id = filename.split(\".\")[0]\n        result = \"\"\n        #result = predict(image_path, image_id)\n        all_result.append(result)\n#print(all_result)        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])\ntest_df.to_csv('submission.csv', index=False)\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample = images[1].permute(1,2,0).cpu().numpy()\nboxes = outputs[1]['boxes'].data.cpu().numpy()\nscores = outputs[1]['scores'].data.cpu().numpy()\n\nboxes = boxes[scores >= detection_threshold].astype(np.int32)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2], box[3]),\n                  (220, 0, 0), 2)\n    \nax.set_axis_off()\nax.imshow(sample)","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}