train_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change these flags to train a specific model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_RESNET = False\n",
    "TRAIN_UNODE = True\n",
    "TRAIN_UNET = False\n",
    "\n",
    "def get_title():\n",
    "    if TRAIN_UNODE: return 'U-NODE'\n",
    "    elif TRAIN_RESNET: return 'RESNET'\n",
    "    elif TRAIN_UNET: return 'UNET'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.utils.data\n",
    "\n",
    "import PIL\n",
    "import skimage.measure\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm, tqdm_notebook\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "from models import ConvODEUNet, ConvResUNet, ODEBlock, Unet\n",
    "from dataloader import GLaSDataLoader\n",
    "from train_utils import plot_losses\n",
    "\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Commands to download the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('Warwick QU Dataset (Released 2016_07_08)'):\n",
    "    !wget https://warwick.ac.uk/fac/sci/dcs/research/tia/glascontest/download/warwick_qu_dataset_released_2016_07_08.zip\n",
    "    !unzip warwick_qu_dataset_released_2016_07_08.zip     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "val_set_idx = torch.LongTensor(10).random_(0, 85)\n",
    "train_set_idx = torch.arange(0, 85)\n",
    "\n",
    "overlapping = (train_set_idx[..., None] == val_set_idx).any(-1)\n",
    "train_set_idx = torch.masked_select(train_set_idx, ~overlapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = GLaSDataLoader((352, 512), dataset_repeat=1, images=train_set_idx)\n",
    "valset = GLaSDataLoader((352, 512), dataset_repeat=1, images=val_set_idx, validation=True)\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=1, shuffle=True, num_workers=10)\n",
    "valloader = torch.utils.data.DataLoader(valset, batch_size=1, shuffle=False, num_workers=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting train data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=5, ncols=6, figsize=(24, 15))\n",
    "\n",
    "for y in range(5):\n",
    "    for x in range(3):\n",
    "        sample = trainset[y]\n",
    "        ax[y, x * 2].imshow(sample[0].numpy().transpose(1,2,0))\n",
    "        ax[y, x * 2 + 1].imshow(sample[1][0])\n",
    "        ax[y, x * 2].axis('off')\n",
    "        ax[y, x * 2 + 1].axis('off')\n",
    "\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(24, 15))\n",
    "\n",
    "sample = trainset[0]\n",
    "ax[1].imshow(sample[1][0].numpy())\n",
    "ax[2].imshow(sample[1].sum(dim=0))\n",
    "ax[0].imshow(sample[0].numpy().transpose(1,2,0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=5, ncols=6, figsize=(24, 15))\n",
    "\n",
    "for y in range(5):\n",
    "    for x in range(3):\n",
    "        sample = valset[y]\n",
    "        ax[y, x * 2].imshow(sample[0].numpy().transpose(1,2,0))\n",
    "        ax[y, x * 2 + 1].imshow(sample[1][1])\n",
    "        ax[y, x * 2].axis('off')\n",
    "        ax[y, x * 2 + 1].axis('off')\n",
    "\n",
    "plt.show(); "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "\n",
    "if TRAIN_UNODE:\n",
    "    net = ConvODEUNet(num_filters=16, output_dim=2, time_dependent=True, \n",
    "                      non_linearity='lrelu', adjoint=True, tol=1e-3)\n",
    "    net.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TRAIN_RESNET:\n",
    "    net = ConvResUNet(num_filters=16, output_dim=2, non_linearity='lrelu')\n",
    "    net.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TRAIN_UNET:\n",
    "    net = Unet(depth=5, num_filters=64, output_dim=2).cuda()\n",
    "    net.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in net.modules():\n",
    "    if isinstance(m, torch.nn.Conv2d):\n",
    "        torch.nn.init.kaiming_normal_(m.weight)\n",
    "        torch.nn.init.constant_(m.bias, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_parameters(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.BCEWithLogitsLoss()\n",
    "val_criterion = torch.nn.BCEWithLogitsLoss()\n",
    "\n",
    "if TRAIN_UNET:\n",
    "    cross_entropy = torch.nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def criterion(conf, labels):\n",
    "        out_shape = conf.shape[2:4]\n",
    "        label_shape = labels.shape[2:4]\n",
    "\n",
    "        w = (label_shape[1] - out_shape[1]) // 2\n",
    "        h = (label_shape[1] - out_shape[1]) // 2\n",
    "        dh, dw = out_shape[0:2]\n",
    "\n",
    "        conf_loss_ce = cross_entropy(conf, labels[:, :, h:h+dh, w:w+dw])\n",
    "\n",
    "        return conf_loss_ce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = []\n",
    "val_losses = []\n",
    "nfe = [[],[],[],[],[],[],[],[],[]] if TRAIN_UNODE else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accumulate_batch = 8  # mini-batch size by gradient accumulation\n",
    "accumulated = 0\n",
    "\n",
    "if TRAIN_RESNET: filename = 'best_border_resnet_model.pt'\n",
    "elif TRAIN_UNODE: filename = 'best_border_unode_model.pt'\n",
    "elif TRAIN_UNET: filename = 'best_border_unet_model.pt'\n",
    "\n",
    "def run(lr=1e-3, epochs=100):\n",
    "    accumulated = 0\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        \n",
    "        # training loop with gradient accumulation\n",
    "        running_loss = 0.0\n",
    "        optimizer.zero_grad()\n",
    "        for data in tqdm(trainloader):\n",
    "            inputs, labels = data[0].cuda(), data[1].cuda()\n",
    "            outputs = net(inputs)\n",
    "            loss = criterion(outputs, labels) / accumulate_batch\n",
    "            loss.backward()\n",
    "            accumulated += 1\n",
    "            if accumulated == accumulate_batch:\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "                accumulated = 0\n",
    "\n",
    "            running_loss += loss.item() * accumulate_batch\n",
    "\n",
    "        losses.append(running_loss / len(trainloader))\n",
    "        \n",
    "        # validation loop\n",
    "        with torch.no_grad():\n",
    "            running_loss = 0.0\n",
    "            for data in valloader:\n",
    "                inputs, labels = data[0].cuda(), data[1].cuda()\n",
    "                outputs = net(inputs)\n",
    "                loss = criterion(outputs, labels)\n",
    "                running_loss += loss.item()\n",
    "\n",
    "            val_losses.append(running_loss / len(valloader))\n",
    "            if np.argmin(val_losses) == len(val_losses) - 1 and loss < 0.4:\n",
    "                torch.save(net, filename)\n",
    "                \n",
    "            clear_output(wait=True)\n",
    "            plot_losses(inputs, outputs, losses, val_losses, get_title(), nfe, net=net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TRAIN_UNODE or TRAIN_RESNET: lr = 1e-3 \n",
    "else: lr = 1e-4\n",
    "\n",
    "run(lr, 600 - len(losses))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load best model\n",
    "net = torch.load(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    running_loss = 0.0\n",
    "    for data in tqdm(valloader):\n",
    "        inputs, labels = data[0].cuda(), data[1].cuda()\n",
    "        outputs = net(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        running_loss += loss.item()\n",
    "\n",
    "    print(\"Check validation loss:\", running_loss / len(valloader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize results on validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from inference_utils import inference_image, postprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=5, ncols=3, figsize=(4*3,3*5))\n",
    "\n",
    "ax[0, 0].set_title('Image')\n",
    "ax[0, 1].set_title('Ground-truth')\n",
    "ax[0, 2].set_title(get_title())\n",
    "\n",
    "for col in range(3):\n",
    "    for row in range(5):\n",
    "        index = val_set_idx[row]\n",
    "        image = PIL.Image.open(f'Warwick QU Dataset (Released 2016_07_08)/train_{index}.bmp')\n",
    "        gt = PIL.Image.open(f'Warwick QU Dataset (Released 2016_07_08)/train_{index}_anno.bmp')\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            result, input_image = inference_image(net, image, shouldpad=TRAIN_UNET)\n",
    "            result = postprocess(result, gt)\n",
    "        if col == 0:\n",
    "            ax[row, col].imshow(image)\n",
    "        elif col == 1:\n",
    "            ax[row, col].imshow(np.array(gt) > 0)\n",
    "        else:\n",
    "            ax[row, col].imshow(image)\n",
    "            ax[row, col].imshow(result, alpha=0.5)\n",
    "                \n",
    "        ax[row, col].set_axis_off()\n",
    "\n",
    "plt.show(); "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate metrics on test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from metrics import ObjectDice, ObjectHausdorff, F1score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_RESNET = False\n",
    "TEST_UNODE = True\n",
    "TEST_UNET = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TEST_UNODE: net = torch.load('best_border_unode_model.pt')\n",
    "elif TEST_RESNET: net = torch.load('best_border_resnet_model.pt')\n",
    "elif TEST_UNET: net = torch.load('best_border_unet_model.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# USE THIS TO LOAD WEIGHTS FROM PAPER\n",
    "#\n",
    "\n",
    "# if TEST_UNODE: \n",
    "#     net = ConvODEUNet(num_filters=16, output_dim=2, time_dependent=True, \n",
    "#                       non_linearity='lrelu', adjoint=True, tol=1e-3)\n",
    "#     net = net.cuda()\n",
    "#     state_dict = torch.load('best_border_unode_paper.pt')\n",
    "#     net.load_state_dict(state_dict)\n",
    "    \n",
    "# if TEST_RESNET:\n",
    "#     net = ConvResUNet(num_filters=16, output_dim=2, non_linearity='lrelu')\n",
    "#     net = net.cuda()\n",
    "#     state_dict = torch.load('best_border_resnet_paper.pt')\n",
    "#     net.load_state_dict(state_dict)\n",
    "\n",
    "# if TEST_UNET:\n",
    "#     net = Unet(depth=5, num_filters=64, output_dim=2).cuda()\n",
    "#     net = net.cuda()\n",
    "#     state_dict = torch.load('best_border_unet_paper.pt')\n",
    "#     net.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dice, hausdorff, f1, dice_full = 0, 0, 0, 0\n",
    "\n",
    "if TEST_UNODE: folder = 'results_unode'\n",
    "elif TEST_UNET: folder = 'results_unet'\n",
    "elif TEST_RESNET: folder = 'results_resnet'\n",
    "    \n",
    "images = []\n",
    "for index in np.arange(1, 81):\n",
    "    if index < 61: images.append(f'testA_{index}_anno.bmp')\n",
    "    else: images.append(f'testB_{index - 60}_anno.bmp')\n",
    "        \n",
    "for i, fname in tqdm_notebook(enumerate(images), total=80):\n",
    "    gt = PIL.Image.open(f'Warwick QU Dataset (Released 2016_07_08)/' + fname)\n",
    "    image = PIL.Image.open(f'Warwick QU Dataset (Released 2016_07_08)/' + fname.replace('_anno', ''))\n",
    "    \n",
    "    result, resized = inference_image(net, image, shouldpad=TEST_UNET)\n",
    "    result = postprocess(result, gt)\n",
    "\n",
    "    gt = skimage.measure.label(np.array(gt))\n",
    "    \n",
    "    f1_img = F1score(result, gt)\n",
    "    hausdorff_img = ObjectHausdorff(result, gt)\n",
    "    dice_img = ObjectDice(result, gt)\n",
    "    \n",
    "    f1 += f1_img\n",
    "    hausdorff += hausdorff_img\n",
    "    dice += dice_img\n",
    "\n",
    "    if i == 59:        \n",
    "        diceA = dice \n",
    "        hausdorffA = hausdorff \n",
    "        f1A = f1\n",
    "\n",
    "    print(i, f1_img, hausdorff_img, dice_img)\n",
    "\n",
    "diceB = dice - diceA\n",
    "hausdorffB = hausdorff - hausdorffA\n",
    "f1B = f1 - f1A\n",
    "\n",
    "print('ObjectDice:', dice / 80, 'A', diceA / 60, 'B', diceB / 20)\n",
    "print('Hausdorff:', hausdorff / 80, 'A', hausdorffA / 60, 'B', hausdorffB / 20)\n",
    "print('F1:', f1 / 80, 'A', f1A / 60, 'B', f1B / 20)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}