{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport torch\nimport torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom albumentations.pytorch.transforms import ToTensorV2","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nWEIGHTS_FILE = \"../input/fasterrcnn-weights/fasterrcnn_weights.bin\"\n\nDETECTION_THRESHOLD = 0.66","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model():\n    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)\n    num_classes = 2  \n    # we set this to be 2: the starfish is one class and the second class is that of the background\n\n    # we need to determine the number of input features\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    #load the trained weights\n    model.load_state_dict(torch.load(WEIGHTS_FILE, map_location=DEVICE))\n    model.eval()\n\n    model = model.to(DEVICE)\n    return model\n\nmodel = get_model()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    prediction_strings = []\n    for j in zip(scores, boxes):\n        prediction_strings.append(\"{0:.10f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n    return \" \".join(prediction_strings)\n\n\ndef predict(model, pixel_array):\n    # Predictions for a single image\n    pixel_array = pixel_array.astype(np.float32) / 255.\n    tensor_img = ToTensorV2(p=1.0)(image=pixel_array)['image'].unsqueeze(0)\n    \n    # Get predictions\n    with torch.no_grad():\n        outputs = model(tensor_img.to(DEVICE))[0]\n\n    boxes = outputs['boxes'].data.cpu().numpy()\n    scores = outputs['scores'].data.cpu().numpy()\n    \n    boxes = boxes[scores >= DETECTION_THRESHOLD].astype(np.int32)\n    scores = scores[scores >= DETECTION_THRESHOLD]\n\n    boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n    boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n  \n    # the evaluation tab specified the format for submission, hence we seek to return with the same format as requested\n    return format_prediction_string(boxes, scores)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import greatbarrierreef\n\nenv = greatbarrierreef.make_env()\n\niter_test = env.iter_test() \n\nfor (pixel_array, df_pred) in iter_test:\n    df_pred['annotations'] = predict(model, pixel_array)\n    env.predict(df_pred)","metadata":{},"execution_count":null,"outputs":[]}]}