{"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)\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\nimport os\nfor 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 20GB 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":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import cv2\nimport os\nimport re\n\n# Imports for image transforms\n# Albumentations bounding box augmentation docs: https://albumentations.ai/docs/getting_started/bounding_boxes_augmentation/\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\n# Torch imports\nimport torch\nimport torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\nfrom matplotlib import pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"INPUT_DIR = '/kaggle/input/global-wheat-detection'\n# PRETRAINED_DIR = '/kaggle/input/wheat-dataset-project'\nPRETRAINED_DIR = '/kaggle/input/wheat-frcnn-bayesian/'\nOUTPUT_DIR = '/kaggle/output/'\nTRAIN_DIR = f'{INPUT_DIR}/train'\nTEST_DIR = f'{INPUT_DIR}/test'\n# MODEL_LOC = f'{PRETRAINED_DIR}/fasterrcnn_resnet50_fpn_TRAINED.pth'\nMODEL_LOC = f'{PRETRAINED_DIR}/fasterrcnn_resnet50_fpn_BAYESOPT.pth'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model Inference\n* Create TestDataset class (similar to WheatDataset but no bboxes, transforms only change to tensor)\n* Create dataset using TestDataset clas\n* Create dataloader\n* Loop over images, image_ids in dataloader\n    * Within each iteration, get outputs by calling model(images)\n    * Loop over i, image in enum(images)\n        * Get boxes and scores from outputs for element i\n        * Threshold boxes and scores\n        * Get boxes from \\[xmin ymin xmax ymax\\] form into \\[x y w h\\] form\n        * make result dict for image id and prediction string (in competition format) and append to result list\n* Sample from outputs as before (with score threshold on boxes) to display prediction"},{"metadata":{"trusted":true},"cell_type":"code","source":"def test_transform():\n    return A.Compose([ToTensorV2(p=1.0)])\n\nclass TestDataset(Dataset):\n    \n    def __init__(self, df, directory, transforms=None):\n        super().__init__()\n        \n        self.image_ids = df['image_id'].unique()\n        self.df = df\n        self.dir = directory\n        self.transforms = transforms\n        \n    def __len__(self):\n        return int(self.image_ids.shape[0])\n    \n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{self.dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        # cv2 reads images into BGR format, must convert to RGB for f-RCNN\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        # f-RCNN requires images in [C,W,H] form with values in [0,1]\n        image /= 255.0\n        \n        if self.transforms:\n            dataToTransform = {'image': image}\n            transData = self.transforms(**dataToTransform)\n            image = transData['image']\n        \n        return image, image_id","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nnum_classes = 2\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\npretrained_state_dict = torch.load(MODEL_LOC)\nmodel.load_state_dict(pretrained_state_dict)\nmodel.eval()\nmodel.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.read_csv(f'{INPUT_DIR}/sample_submission.csv')\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset = TestDataset(test_df, TEST_DIR, test_transform())\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ntest_dl = DataLoader(dataset=test_dataset, batch_size=4, num_workers=4, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"THRESHOLD = .5\nres = []\n\ndef getPredString(outputTup):\n    formatted_strings = []\n    for tup in outputTup:\n        score = tup[0]\n        x, y, w, h = tup[1]\n        box_string = f'{score} {x} {y} {w} {h}'\n        formatted_strings.append(box_string)\n    return \" \".join(formatted_strings)\n    \n    \n\nfor imgs, img_ids in test_dl:\n    imgs = list(image.to(device) for image in imgs)\n    model_outputs = model(imgs)\n    \n    for i,img in enumerate(imgs):\n        scores = model_outputs[i]['scores'].data.cpu().numpy()\n        bboxes = model_outputs[i]['boxes'].data.cpu().numpy()\n        bboxes = bboxes[scores >= THRESHOLD].astype(np.int32)\n        scores = scores[scores >= THRESHOLD]\n        \n        this_id = img_ids[i]\n        \n        bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]\n        bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]\n        \n        res.append({'image_id': this_id, 'PredictionString': getPredString(zip(scores, bboxes))})\n\ntest_df = pd.DataFrame(res, columns=['image_id', 'PredictionString'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample = imgs[1].permute(1,2,0).cpu().numpy()\nscores = model_outputs[1]['scores'].data.cpu().numpy()\nbboxes = model_outputs[1]['boxes'].data.cpu().numpy()\nbboxes = bboxes[scores >= THRESHOLD].astype(np.int32)\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in bboxes:\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":"test_df.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}