{"cells":[{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"pip install --no-deps ../input/wtfml/","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport ast\n\nimport pandas as pd\nimport numpy as np\n\nimport albumentations\nimport argparse\nimport torch\nimport torchvision\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom sklearn import metrics\nfrom sklearn.model_selection import train_test_split\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n\nfrom wtfml.engine import RCNNEngine\nfrom wtfml.data_loaders.image import RCNNLoader\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nfrom PIL import Image\nfrom PIL import ImageFile\nimport matplotlib.patches as patches","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    # function taken from: https://www.kaggle.com/arunmohan003/fasterrcnn-using-pytorch-baseline\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\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n\nclass Model(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.base_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(\n            pretrained=False, \n            pretrained_backbone=False\n        )\n        in_features = self.base_model.roi_heads.box_predictor.cls_score.in_features\n        self.base_model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)\n\n    def forward(self, images, targets):\n        if targets is None:\n            return self.base_model(images, targets)\n        else:\n            output = self.base_model(images, targets)\n            if isinstance(output, list):\n                return output\n            loss = sum(loss for loss in output.values())\n            return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class args:\n\tdata_path = \"../input/global-wheat-detection\"\n\tepochs = 30\n\tdevice = \"cuda\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Model()\nmodel.to(args.device)\n\nmodel.load_state_dict(\n    torch.load('../input/wtfwheat/model.bin')\n)\n\nmean = (0., 0., 0.)\nstd = (1, 1, 1)\naug = albumentations.Compose(\n    [albumentations.Normalize(mean, std, max_pixel_value=255.0, always_apply=True)]\n)\n\ntest_df = pd.read_csv(os.path.join(args.data_path, \"sample_submission.csv\"))\ntest_df.loc[:, \"bbox\"] = [\"[0, 0, 10, 10]\"] * len(test_df)\ntest_df.bbox = test_df.bbox.apply(ast.literal_eval)\ntest_df = test_df.groupby('image_id')['bbox'].apply(list).reset_index(name='bboxes')\ntest_image_ids = test_df.image_id.values\n\nimages = test_df.image_id.values.tolist()\nimages = [os.path.join(args.data_path, \"test\", i + \".jpg\") for i in images]\ntargets = test_df.bboxes.values\n\ntest_dataset = RCNNLoader(\n    image_paths=images, \n    bounding_boxes=targets, \n    augmentations=aug\n)\ntest_loader = torch.utils.data.DataLoader(\n    test_dataset, batch_size=1, shuffle=False, num_workers=4, collate_fn=collate_fn\n)\n\nprediction_strings = []\npredictions = RCNNEngine.predict(test_loader, model, device=args.device)\nfor p in predictions:\n    boxes = p['boxes'].numpy()\n    scores = p['scores'].numpy()\n\n    boxes = boxes[scores >= 0.5].astype(np.int32)\n    scores = scores[scores >= 0.5]\n\n    boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n    boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n    prediction_strings.append(format_prediction_string(boxes, scores))\n\nsample = pd.read_csv(os.path.join(args.data_path, \"sample_submission.csv\"))\nsample.loc[:, \"PredictionString\"] = prediction_strings\nsample.loc[:, \"image_id\"] = test_image_ids\nsample.to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for idx in range(len(sample)):\n    npimage = np.array(Image.open(os.path.join(args.data_path, \"test\", sample.image_id.values[idx] + \".jpg\")))\n    boxes = sample.PredictionString.values[idx]\n    boxes = [x for i, x in enumerate(boxes.split()) if i % 5 != 0]\n    boxes = np.array(boxes).reshape(-1, 4).astype(float)\n\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    for b in boxes:\n        ax.add_patch(\n         patches.Rectangle(\n            (b[0], b[1]),\n            b[2],\n            b[3],\n            fill=False,\n            color='red')\n        )\n    ax.set_axis_off()\n    ax.imshow(npimage)","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}