{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize Tensorpack Mask-RCNN Detection Results\n",
    "\n",
    "This notebook visualizes detection results predicted by a trained [Tensorpack Mask-RCNN](https://github.com/tensorpack/tensorpack/tree/master/examples/FasterRCNN) model. \n",
    "## Load trained model\n",
    "\n",
    "First we define the system path for Python classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob\n",
    "import sys\n",
    "\n",
    "sys.path.append('/tensorpack/examples/FasterRCNN')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we initialize the ResNet FPN Mask RCNN model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from modeling.generalized_rcnn import ResNetFPNModel\n",
    "\n",
    "# create a mask r-cnn model\n",
    "mask_rcnn_model = ResNetFPNModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we specify the `model_dir` below and load the `trained model`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "\n",
    "# find best pre-trained model checkpoint\n",
    "model_dir=  f\"{os.environ['LOGDIR']}/train_log/maskrcnn/\"\n",
    "\n",
    "print(f\"Using model directory: {model_dir}\")\n",
    "model_search_path = os.path.join(model_dir, \"model-*.index\" )\n",
    "model_files = glob.glob(model_search_path)\n",
    "\n",
    "def sort_key(path):\n",
    "    index = path.rindex(\"model-\")\n",
    "    key = int(path[index+6:-6])\n",
    "    return key\n",
    "\n",
    "try:\n",
    "    model_files = sorted(model_files, key = sort_key)\n",
    "    latest_trained_model = model_files[-1]\n",
    "\n",
    "    trained_model = latest_trained_model[:-6]\n",
    "    print(f'Using model: {trained_model}')\n",
    "except:\n",
    "    print(f\"No model found in: {model_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we initialize the model configuration to match the configuration we used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from config import finalize_configs, config as cfg\n",
    "from dataset import register_coco\n",
    "\n",
    "# setup config\n",
    "cfg.MODE_FPN = True\n",
    "cfg.MODE_MASK = True\n",
    "cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS\n",
    "cfg.DATA.BASEDIR = '/efs/data/'\n",
    "register_coco(cfg.DATA.BASEDIR)\n",
    "finalize_configs(is_training=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Predictor\n",
    "\n",
    "Next we create a predictor that uses our trained model to make predictions on test images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorpack.predict.base import OfflinePredictor\n",
    "from tensorpack.tfutils.sessinit import get_model_loader\n",
    "from tensorpack.predict.config import PredictConfig\n",
    "\n",
    "# Create an inference predictor\n",
    "# PredictConfig takes a model, input tensors and output tensors\n",
    "input_tensors = mask_rcnn_model.get_inference_tensor_names()[0]\n",
    "output_tensors = mask_rcnn_model.get_inference_tensor_names()[1]\n",
    "            \n",
    "predictor = OfflinePredictor(PredictConfig(\n",
    "                model=mask_rcnn_model,\n",
    "                session_init=get_model_loader(trained_model),\n",
    "                input_names=input_tensors,\n",
    "                output_names=output_tensors))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download COCO Test 2017 dataset\n",
    "Below we download the [COCO 2017 Test dataset](http://cocodataset.org/#download) and extract the downloaded dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O /tmp/test2017.zip http://images.cocodataset.org/zips/test2017.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we extract the downloaded COCO 2017 Test dataset. This step executes quietly and may take couple of minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip -q -d /tmp/ /tmp/test2017.zip\n",
    "!rm  /tmp/test2017.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load test image\n",
    "\n",
    "Next, we find a random image to test from COCO 2017 Test dataset. You can come back to this cell when you want to load the next test image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "\n",
    "test2017_dir=os.path.join(\"/tmp\", \"test2017\")\n",
    "img_id=random.choice(os.listdir(test2017_dir))\n",
    "img_local_path = os.path.join(test2017_dir,img_id)\n",
    "print(img_local_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we load the random test image and convert the image color scheme from BGR to RBG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "img=cv2.imread(img_local_path, cv2.IMREAD_COLOR)\n",
    "print(img.shape)\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next. we show the raw image we randomly loaded from the COCO 2017 Test dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig,ax = plt.subplots(figsize=(img.shape[1]//50, img.shape[0]//50))\n",
    "ax.imshow(img.astype(int))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict\n",
    "Next, we use the predictor to predict detection results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from eval import predict_image, DetectionResult\n",
    "detection_results = predict_image(img, predictor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we draw the detection results on our image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from viz import draw_final_outputs\n",
    "final_viz = draw_final_outputs(img, detection_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we visualize the detection results drawn on the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize=(final_viz.shape[1]//50, final_viz.shape[0]//50))\n",
    "ax.imshow(final_viz.astype(int))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Go back to <b>Load test image</b> cell if you want to test more images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}