04 - Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/twloehfelm/SAR2020/blob/master/04%20-%20Segmentation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SkHz0FVN6F8u"
   },
   "source": [
    "<table width=\"100%\">\n",
    "    <tr>\n",
    "        <td valign=\"top\"><img src=\"https://cdn.ymaws.com/www.abdominalradiology.org/graphics/logo.jpg\"/></td>\n",
    "        <td valign=\"middle\" align=\"right\"><h1>SAR 2020<br/>AI Masters Class</h1></td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td align=\"center\" colspan=2><h1>Segmentation</h1></td>\n",
    "    </tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2O-42dYnoR3Q"
   },
   "source": [
    "We are going to build a liver segmentation tool using Facebook's detectron2 object identification algorithm trained on 18 CTs from the Combined Healthy Abdominal Organ Segmentation (CHAOS) Challenge. We'll test it on 2 CTs from the same challenge that are not included in the training set.\n",
    "<br/><br/><br/>\n",
    "**References to the original dataset and related publications:**\n",
    "\n",
    "> <sub><sup>A.E. Kavur, M. A. Selver, O. Dicle, M. Barış,  N.S. Gezer. CHAOS - Combined (CT-MR) Healthy Abdominal Organ Segmentation Challenge Data (Version v1.03) [Data set]. Apr.  2019. Zenodo. http://doi.org/10.5281/zenodo.3362844 </sup></sub>\n",
    "\n",
    "> <sub><sup>A.E. Kavur, N.S. Gezer, M. Barış, P.-H. Conze, V. Groza, D.D. Pham, et al. \"CHAOS Challenge - Combined (CT-MR) Healthy Abdominal Organ Segmentation\", arXiv pre-print, Jan. 2020. https://arxiv.org/abs/2001.06535</sup></sub>\n",
    "\n",
    "> <sub><sup>A.E. Kavur, N.S. Gezer, M. Barış, Y.Şahin, S. Özkan, B. Baydar, et al.  \"Comparison of semi-automatic and deep learning-based automatic methods for liver segmentation in living liver transplant donors\", Diagnostic and  Interventional  Radiology,  vol. 26, pp. 11–21, Jan. 2020. https://doi.org/10.5152/dir.2019.19025</sup></sub>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vM54r6jlKTII"
   },
   "source": [
    "# Install required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9_FzH13EjseR"
   },
   "outputs": [],
   "source": [
    "# install dependencies: (use cu100 because colab is on CUDA 10.0)\n",
    "!pip install -U torch==1.4+cu100 torchvision==0.5+cu100 -f https://download.pytorch.org/whl/torch_stable.html > /dev/null\n",
    "!pip install cython pyyaml==5.1 > /dev/null\n",
    "!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI' > /dev/null\n",
    "import torch, torchvision\n",
    "torch.__version__\n",
    "!gcc --version\n",
    "\n",
    "#install detectron2:\n",
    "!pip3 install -q detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu100/index.html\n",
    "\n",
    "#install pydicom\n",
    "!pip3 install -q pydicom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4Qg7zSVOulkb"
   },
   "outputs": [],
   "source": [
    "# Download liver seg training data to root directory\n",
    "!wget -q https://www.dropbox.com/s/ajg922tunqptd2z/chaos_train.zip\n",
    "!unzip chaos_train.zip > /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZyAvNCJMmvFF"
   },
   "outputs": [],
   "source": [
    "## You may need to restart your runtime prior to this, to let your installation take effect\n",
    "# Some basic setup:\n",
    "# Setup detectron2 logger\n",
    "import detectron2\n",
    "from detectron2.utils.logger import setup_logger\n",
    "setup_logger()\n",
    "\n",
    "# import some common libraries\n",
    "import numpy as np\n",
    "import cv2\n",
    "import random\n",
    "\n",
    "# import some common detectron2 utilities\n",
    "from detectron2 import model_zoo\n",
    "from detectron2.engine import DefaultPredictor, DefaultTrainer, default_argument_parser, default_setup, launch\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.utils.visualizer import Visualizer, ColorMode\n",
    "\n",
    "from detectron2.structures import BoxMode\n",
    "from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_train_loader, build_detection_test_loader\n",
    "from detectron2.data import transforms as T\n",
    "from detectron2.data import detection_utils as utils\n",
    "\n",
    "# Imports for liver seg\n",
    "import os\n",
    "from pathlib import Path\n",
    "import json\n",
    "import pydicom\n",
    "import pycocotools\n",
    "import png\n",
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "\n",
    "try:\n",
    "  import google.colab\n",
    "  IN_COLAB = True\n",
    "except:\n",
    "  IN_COLAB = False\n",
    "\n",
    "if IN_COLAB:\n",
    "  # opencv is pre-installed on colab\n",
    "  # colab cannot use the standard imshow due to some html/web limitation\n",
    "  from google.colab.patches import cv2_imshow as imshow\n",
    "else:\n",
    "  #!pip install opencv-python\n",
    "  from matplotlib.pyplot import imshow\n",
    "  %matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qi2amjKT9Trx"
   },
   "outputs": [],
   "source": [
    "!rm -rf ./__MACOSX\n",
    "!rm -rf ./chaos_train.zip\n",
    "!rm -rf ./sample_data\n",
    "ROOT_PATH = Path('/content')\n",
    "chaos = ROOT_PATH/'chaos_train'\n",
    "pts_dir = chaos/'CT'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b2bjrfb2LDeo"
   },
   "source": [
    "# Prepare the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JFR6B1Kh99Bb"
   },
   "outputs": [],
   "source": [
    "# Create lists of DICOMs and segmentation files for the training and validation datasets\n",
    "dicoms_train = list()\n",
    "segs_train = list()\n",
    "dicoms_val = list()\n",
    "segs_val = list()\n",
    "\n",
    "# Randomly split patients 90/10 into training/validation. 20 studies, so 18 train and 2 val\n",
    "pts = [x for x in pts_dir.iterdir() if x.is_dir()]\n",
    "random.seed(716)\n",
    "random.shuffle(pts)\n",
    "train_pts = pts[:18]\n",
    "val_pts = pts[18:]\n",
    "\n",
    "# For each patient, add DICOMs and segmentation files to the respective lists\n",
    "# N.B. DICOM and seg file names must sort in the same order\n",
    "for p in train_pts:\n",
    "  dicoms_train += sorted([x for x in (p/'DICOM_anon').iterdir() if x.is_file()])\n",
    "  segs_train += sorted([x for x in (p/'Ground').iterdir() if x.is_file()])\n",
    "for p in val_pts:\n",
    "  dicoms_val += sorted([x for x in (p/'DICOM_anon').iterdir() if x.is_file()])\n",
    "  segs_val += sorted([x for x in (p/'Ground').iterdir() if x.is_file()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yfsEnRXLkU2N"
   },
   "outputs": [],
   "source": [
    "def mapper(dataset_dict):\n",
    "  \"\"\"\n",
    "  Args:\n",
    "      dataset_dict (dict): Metadata of one image, in detectron2 Dataset format.\n",
    "\n",
    "  Returns:\n",
    "      dict: a format that builtin models in detectron2 accept\n",
    "  \"\"\"\n",
    "  # Implement a mapper, similar to the default DatasetMapper, but with your own customizations\n",
    "  dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below\n",
    "  ds = pydicom.dcmread(dataset_dict[\"file_name\"])\n",
    "  image = ds.pixel_array\n",
    "  # Convert pixel values to Hounsfield units\n",
    "  image = image*ds.RescaleSlope + ds.RescaleIntercept\n",
    "  image, transforms = T.apply_transform_gens([T.RandomBrightness(0.8, 1.2), T.RandomContrast(0.8, 1.2)], image)\n",
    "  dataset_dict[\"image\"] = torch.as_tensor(image.astype(\"float32\"))\n",
    "\n",
    "  annos = [\n",
    "      utils.transform_instance_annotations(obj, transforms, image.shape[:2])\n",
    "      for obj in dataset_dict.pop(\"annotations\")\n",
    "      if obj.get(\"iscrowd\", 0) == 0\n",
    "  ]\n",
    "  instances = utils.annotations_to_instances(annos, image.shape[:2], mask_format='bitmask')\n",
    "  dataset_dict[\"instances\"] = utils.filter_empty_instances(instances)\n",
    "  return dataset_dict\n",
    "\n",
    "class LiverTrainer(DefaultTrainer):\n",
    "    @classmethod\n",
    "    def build_test_loader(cls, cfg, dataset_name):\n",
    "        return build_detection_test_loader(cfg, dataset_name, mapper=mapper)\n",
    "\n",
    "    @classmethod\n",
    "    def build_train_loader(cls, cfg):\n",
    "        return build_detection_train_loader(cfg, mapper=mapper)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c3TFNA-47FYP"
   },
   "outputs": [],
   "source": [
    "def load_mask(image_path):\n",
    "    \"\"\"Load masks for the given image.\n",
    "    Mask is a binary true/false map of the same size as the image.\n",
    "    args:\n",
    "      image_path = complete path to segmentation file\n",
    "    returns:\n",
    "      boolean array of the liver segmentation mask \n",
    "    \"\"\"\n",
    "    mask = plt.imread(str(image_path))\n",
    "    return mask.astype(np.bool)\n",
    "\n",
    "def bbox(img):\n",
    "    \"\"\"Generates a bounding box from a segmentation mask\n",
    "    Bounding box is given as coordinates of upper left and lower right corner\n",
    "    args:\n",
    "      img = boolean array of a segmentation mask\n",
    "    returns:\n",
    "      coordinates of upper left and lower right corner \n",
    "    \"\"\"\n",
    "    try:\n",
    "      x,y = np.where(img)\n",
    "    except ValueError:\n",
    "      return None\n",
    "    if x.size != 0:\n",
    "      bbox = y.min(), x.min(), y.max(), x.max()\n",
    "    else:\n",
    "      bbox = None\n",
    "    return bbox\n",
    "\n",
    "# from fastai2 medical imaging\n",
    "def windowed(px, w, l):\n",
    "    \"\"\"Windows a pixel_array of Houndfield units\n",
    "    args:\n",
    "      px = pixel array in Houndfield units\n",
    "      w = window width (HU range)\n",
    "      l = window level (center point)\n",
    "    returns:\n",
    "      pixel_array convered to the given window/level\n",
    "    \"\"\"\n",
    "    if type(w) == pydicom.multival.MultiValue:\n",
    "      w = w[0]\n",
    "    if type(l) == pydicom.multival.MultiValue:\n",
    "      l = l[0]\n",
    "    px_min = l - w//2\n",
    "    px_max = l + w//2\n",
    "    px[px<px_min] = px_min\n",
    "    px[px>px_max] = px_max\n",
    "    return (px-px_min) / (px_max-px_min)\n",
    "\n",
    "# Used in the DatasetCatalog.register call\n",
    "def get_liver_dicts(train_or_val):\n",
    "  \"\"\"Builds a dataset_dict for detectron2\n",
    "    args:\n",
    "      train_or_val = string 'train' or 'val' indicating whether to return\n",
    "        the training or validation dataset_dict\n",
    "    returns:\n",
    "      dataset_dict with each element of the training or validation dataset\n",
    "    \"\"\"\n",
    "  if train_or_val == \"train\":\n",
    "    dicoms = dicoms_train\n",
    "    segs = segs_train\n",
    "  elif train_or_val == \"val\":\n",
    "    dicoms = dicoms_val\n",
    "    segs = segs_val\n",
    "\n",
    "  dataset_dicts = []\n",
    "  for idx, d in enumerate(dicoms):\n",
    "    record = {}\n",
    "    \n",
    "    filename = str(d)\n",
    "    ds = pydicom.dcmread(filename)\n",
    "    height, width = ds.Rows, ds.Columns\n",
    "\n",
    "    # Mininum required fields for each element in the dict\n",
    "    record[\"file_name\"] = filename  # Full path to image file\n",
    "    record[\"image_id\"] = idx        # Index of file (unique serial number)\n",
    "    record[\"height\"] = height       # Image dimension\n",
    "    record[\"width\"] = width         # Image dimension\n",
    "\n",
    "    try:\n",
    "      mask = load_mask(str(segs[idx]))\n",
    "    except IndexError:\n",
    "      mask = None\n",
    "\n",
    "    # Add list of segmentation object(s)\n",
    "    objs = []\n",
    "    if bbox(mask) is not None:\n",
    "      obj = {\n",
    "          \"bbox\": bbox(mask),\n",
    "          \"bbox_mode\": BoxMode.XYXY_ABS,\n",
    "          \"segmentation\": pycocotools.mask.encode(np.asarray(mask, order=\"F\")), # Convert binary mask to RLE format\n",
    "          \"category_id\": 0,\n",
    "          \"is_crowd\": 0\n",
    "      }\n",
    "      objs.append(obj)\n",
    "    record[\"annotations\"] = objs\n",
    "    dataset_dicts.append(record)\n",
    "  return dataset_dicts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RgYw3tuOAs2Q"
   },
   "outputs": [],
   "source": [
    "# Clear existing DatasetCatalog and then register the training and validation datasets\n",
    "DatasetCatalog.clear()  \n",
    "for d in [\"train\", \"val\"]:\n",
    "    DatasetCatalog.register(\"liver_\" + d, lambda d=d: get_liver_dicts(d))\n",
    "    MetadataCatalog.get(\"liver_\" + d).set(thing_classes=[\"liver\"])\n",
    "liver_metadata = MetadataCatalog.get(\"liver_train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6ljbWTX0Wi8E"
   },
   "source": [
    "# Verify Data Loading\n",
    "To verify the data loading is correct, let's visualize the annotations of randomly selected samples in the training set:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Tp1Ft4x-Kx23"
   },
   "outputs": [],
   "source": [
    "dataset_dicts = get_liver_dicts(\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show examples of training data: mask + DICOM\n",
    "for d in random.sample(dataset_dicts, 3):\n",
    "  try:\n",
    "    # Decode the binary mask from the RLE string\n",
    "    mask = pycocotools.mask.decode(d[\"annotations\"][0][\"segmentation\"])\n",
    "    # Extract pixel array from DICOM and convert to default W/L\n",
    "    ds=pydicom.dcmread(d[\"file_name\"])\n",
    "    im=ds.pixel_array\n",
    "    im=im*ds.RescaleSlope + ds.RescaleIntercept\n",
    "    im = windowed(im, ds.WindowWidth, ds.WindowCenter)\n",
    "    \n",
    "    #Plot the mask and windowed DICOM\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(10, 5), sharey=True)\n",
    "    axs[0].imshow(mask, cmap=plt.cm.binary_r)\n",
    "    axs[1].imshow(im, cmap='gray')\n",
    "    plt.show()\n",
    "  except IndexError:\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UkNbUzUOLYf0"
   },
   "outputs": [],
   "source": [
    "# Same as previous, but use detectron's default display format with overlaid mask\n",
    "for d in random.sample(dataset_dicts, 3):\n",
    "    ds=pydicom.dcmread(d[\"file_name\"])\n",
    "    im=ds.pixel_array\n",
    "    im=im*ds.RescaleSlope + ds.RescaleIntercept\n",
    "    im = windowed(im, ds.WindowWidth, ds.WindowCenter)\n",
    "    im = np.stack((im,) * 3, -1)\n",
    "    im=im*255\n",
    "    v = Visualizer(im[:, :, ::-1],\n",
    "                   metadata=liver_metadata, \n",
    "                   scale=0.8\n",
    "    )\n",
    "    v = v.draw_dataset_dict(d)\n",
    "    imshow(v.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XW5EakiuV2Ar"
   },
   "source": [
    "# Train\n",
    "\n",
    "Now, let's fine-tune a coco-pretrained R50-FPN Mask R-CNN model on the liver dataset. It takes ~2 minutes to train 300 iterations on Colab.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7unkuuiqLdqd"
   },
   "outputs": [],
   "source": [
    "cfg = get_cfg()\n",
    "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"))\n",
    "cfg.DATASETS.TRAIN = (\"liver_train\",)\n",
    "cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS = False\n",
    "cfg.DATASETS.TEST = ()\n",
    "cfg.DATALOADER.NUM_WORKERS = 2\n",
    "cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\")  # Let training initialize from model zoo\n",
    "\n",
    "# These three SOLVER parameters are probably the best places to start tweaking to modify performance\n",
    "cfg.SOLVER.IMS_PER_BATCH = 8\n",
    "cfg.SOLVER.BASE_LR = 0.001  # Can experiment with differnt base learning rates\n",
    "cfg.SOLVER.MAX_ITER = 500    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset\n",
    "\n",
    "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this simple dataset (default: 512)\n",
    "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (liver)\n",
    "\n",
    "cfg.INPUT.FORMAT = \"F\" #32-bit single channel floating point pixels\n",
    "cfg.INPUT.MASK_FORMAT = \"bitmask\" # Needed to change this from the default \"polygons\"\n",
    "\n",
    "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
    "trainer = LiverTrainer(cfg)\n",
    "trainer.resume_or_load(resume=False)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hBXeH8UXFcqU"
   },
   "outputs": [],
   "source": [
    "# Look at training curves in tensorboard:\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "h0JeEVP7MS8A"
   },
   "outputs": [],
   "source": [
    "%reload_ext tensorboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0e4vdDIOXyxF"
   },
   "source": [
    "# Inference & evaluation using the trained model\n",
    "Now, let's run inference with the trained model on the validation dataset. First, let's create a predictor using the model we just trained:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CYHrysR9MgV6"
   },
   "outputs": [],
   "source": [
    "from detectron2.modeling import build_model\n",
    "from detectron2.checkpoint import DetectionCheckpointer\n",
    "\n",
    "class LiverPredictor:\n",
    "    \"\"\"\n",
    "    Create a simple end-to-end predictor with the given config that runs on\n",
    "    single device for a single input image.\n",
    "    Compared to using the model directly, this class does the following additions:\n",
    "    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.\n",
    "    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.\n",
    "    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.\n",
    "    4. Take one input image and produce a single output, instead of a batch.\n",
    "    If you'd like to do anything more fancy, please refer to its source code\n",
    "    as examples to build and use the model manually.\n",
    "    Attributes:\n",
    "        metadata (Metadata): the metadata of the underlying dataset, obtained from\n",
    "            cfg.DATASETS.TEST.\n",
    "    Examples:\n",
    "    .. code-block:: python\n",
    "        pred = DefaultPredictor(cfg)\n",
    "        inputs = cv2.imread(\"input.jpg\")\n",
    "        outputs = pred(inputs)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, cfg):\n",
    "        self.cfg = cfg.clone()  # cfg can be modified by model\n",
    "        self.model = build_model(self.cfg)\n",
    "        self.model.eval()\n",
    "        self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])\n",
    "\n",
    "        checkpointer = DetectionCheckpointer(self.model)\n",
    "        checkpointer.load(cfg.MODEL.WEIGHTS)\n",
    "\n",
    "        self.transform_gen = T.ResizeShortestEdge(\n",
    "            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST\n",
    "        )\n",
    "\n",
    "        self.input_format = cfg.INPUT.FORMAT\n",
    "        \n",
    "    def __call__(self, original_image):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            original_image (np.ndarray): a single channel image.\n",
    "        Returns:\n",
    "            predictions (dict):\n",
    "                the output of the model for one image only.\n",
    "                See :doc:`/tutorials/models` for details about the format.\n",
    "        \"\"\"\n",
    "        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258\n",
    "            height, width = original_image.shape[:2]\n",
    "            image = original_image\n",
    "            image = torch.as_tensor(image.astype(\"float32\"))\n",
    "\n",
    "            inputs = {\"image\": image, \"height\": height, \"width\": width}\n",
    "            predictions = self.model([inputs])[0]\n",
    "            return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ya5nEuMELeq8"
   },
   "outputs": [],
   "source": [
    "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model\n",
    "cfg.DATASETS.TEST = (\"liver_val\", )\n",
    "predictor = LiverPredictor(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qWq1XHfDWiXO"
   },
   "source": [
    "Then, we randomly select several samples to visualize the prediction results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "U5LhISJqWXgM"
   },
   "outputs": [],
   "source": [
    "dataset_dicts = get_liver_dicts(\"val\")\n",
    "for d in random.sample(dataset_dicts, 10): \n",
    "    ds=pydicom.dcmread(d[\"file_name\"])\n",
    "    im=ds.pixel_array\n",
    "    im=im*ds.RescaleSlope + ds.RescaleIntercept\n",
    "    outputs = predictor(im)\n",
    "    im = windowed(im, ds.WindowWidth, ds.WindowCenter)\n",
    "    im = np.stack((im,) * 3, -1)\n",
    "    im=im*255\n",
    "    v = Visualizer(im[:, :, ::-1],\n",
    "                   metadata=liver_metadata, \n",
    "                   scale=0.8, \n",
    "                   instance_mode=ColorMode.IMAGE\n",
    "    )\n",
    "    v = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n",
    "    imshow(v.get_image()[:, :, ::-1])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "vM54r6jlKTII",
    "b2bjrfb2LDeo",
    "6ljbWTX0Wi8E",
    "XW5EakiuV2Ar",
    "0e4vdDIOXyxF"
   ],
   "include_colab_link": true,
   "name": "SAR 2020 - Segmentation.ipynb",
   "provenance": []
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}