iter_lbp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import os\n",
    "from scipy.sparse import csr_matrix, csgraph\n",
    "from skimage.morphology import convex_hull_image\n",
    "import struct\n",
    "import time\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import warnings\n",
    "\n",
    "from evaluation import *\n",
    "from mindssc import *\n",
    "from thin_plate_spline import *\n",
    "from utils import *\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '4'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/path/to/data/l2r_2020/task_03/'\n",
    "device = 'cuda'\n",
    "\n",
    "pairs = torch.tensor(np.genfromtxt(os.path.join(data_dir, 'pairs_val.csv'), delimiter=',')[1:], dtype=torch.long)\n",
    "cases = pairs.unique()\n",
    "ind = torch.zeros(cases.max()+1, dtype=torch.long)\n",
    "ind[cases] = torch.arange(cases.shape[0])\n",
    "\n",
    "D, H, W = 192, 160, 256\n",
    "mind_delta = 3\n",
    "mind_sigma = 3\n",
    "mind_patch_step = 3\n",
    "mind_patch_radius = 3\n",
    "mind_stride = 2\n",
    "N = 2048\n",
    "d = 2\n",
    "k = 10\n",
    "k1 = 12\n",
    "alpha = 2.5\n",
    "alpha1 = 100\n",
    "sps = torch.linspace(.3,.03,20)\n",
    "lbp_iter = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def load_case(case, plot=False):\n",
    "    img_path = os.path.join(data_dir, 'Training/img', 'img{:04d}.nii.gz'.format(case))\n",
    "    img = torch.from_numpy(nib.load(img_path).get_fdata().astype(np.float32))\n",
    "    \n",
    "    seg_path = os.path.join(data_dir, 'Training/label', 'label{:04d}.nii.gz'.format(case))\n",
    "    seg = torch.from_numpy(nib.load(seg_path).get_fdata().astype(np.int64))\n",
    "    \n",
    "    mask = seg > 0\n",
    "    mask = F.avg_pool3d(mask.view(1, 1, D, H, W).float().cuda(), 17, stride=1, padding=8)[0, 0].cpu() > 0.001\n",
    "    for i in range(D):\n",
    "        mask[i, :, :] = torch.from_numpy(convex_hull_image(mask[i,:,:]))\n",
    "    for i in range(H):\n",
    "        mask[:, i, :] = torch.from_numpy(convex_hull_image(mask[:,i,:]))\n",
    "    for i in range(W):\n",
    "        mask[:, :, i] = torch.from_numpy(convex_hull_image(mask[:,:,i]))\n",
    "\n",
    "    if plot:\n",
    "        cmap = plt.get_cmap('Set1')\n",
    "        plt.figure(figsize=(24,8))\n",
    "        plt.subplot(131)\n",
    "        plt.imshow(img[D//2, :, :], cmap='gray')\n",
    "        seg_plot = cmap(seg[D//2, :, :]/13.)\n",
    "        seg_plot[:, :, 3] = seg[D//2, :, :] != 0\n",
    "        plt.imshow(seg_plot, alpha=0.5)\n",
    "        plt.imshow(mask[D//2, :, :], alpha=0.1)\n",
    "        plt.axis('off')\n",
    "        \n",
    "        plt.subplot(132)\n",
    "        plt.imshow(img[:, H//2, :], cmap='gray')\n",
    "        seg_plot = cmap(seg[:, H//2, :]/13.)\n",
    "        seg_plot[:, :, 3] = seg[:, H//2, :] != 0\n",
    "        plt.imshow(seg_plot, alpha=0.5)\n",
    "        plt.imshow(mask[:, H//2, :], alpha=0.1)\n",
    "        plt.axis('off')\n",
    "        \n",
    "        plt.subplot(133)\n",
    "        plt.imshow(img[:, :, W//2], cmap='gray')\n",
    "        seg_plot = cmap(seg[:, :, W//2]/13.)\n",
    "        seg_plot[:, :, 3] = seg[:, :, W//2] != 0\n",
    "        plt.imshow(seg_plot, alpha=0.5)\n",
    "        plt.imshow(mask[:, :, W//2], alpha=0.1)\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "        \n",
    "    return img, seg, mask\n",
    "  \n",
    "grid = F.affine_grid(torch.eye(3, 4, device='cuda').unsqueeze(0), (1, 1, D//mind_stride, H//mind_stride, W//mind_stride))\n",
    "mind_patch = torch.stack(torch.meshgrid(torch.arange(0, 2 * mind_patch_radius + 1, mind_patch_step, device='cuda'),\n",
    "                                        torch.arange(0, 2 * mind_patch_radius + 1, mind_patch_step, device='cuda'),\n",
    "                                        torch.arange(0, 2 * mind_patch_radius + 1, mind_patch_step, device='cuda')), dim=3).view(1, -1, 3) - mind_patch_radius\n",
    "mind_patch = flow_pt(mind_patch, (D//mind_stride, H//mind_stride, W//mind_stride), align_corners=True).view(1, 1, -1, 1, 3)\n",
    "    \n",
    "imgs = torch.zeros(len(cases), 1, D, H, W).pin_memory()\n",
    "segs = torch.zeros(len(cases), 1, D, H, W, dtype=torch.long).pin_memory()\n",
    "masks = torch.zeros(len(cases), 1, D, H, W, dtype=torch.bool).pin_memory()\n",
    "feats = torch.zeros(len(cases), 12*mind_patch.shape[2], D//mind_stride, H//mind_stride, W//mind_stride).pin_memory()\n",
    "for i, case in enumerate(cases):\n",
    "    print('Case {}'.format(case))\n",
    "    img, seg, mask = load_case(case, plot=True)\n",
    "\n",
    "    mind = F.grid_sample(mindssc(img.view(1, 1, D, H, W).cuda(), mind_delta, mind_sigma), grid.view(1, 1, 1, -1, 3) + mind_patch).view(12*mind_patch.shape[2], D//mind_stride, H//mind_stride, W//mind_stride).cpu()\n",
    "    \n",
    "    imgs[i, 0, :, :, :] = img\n",
    "    segs[i, 0, :, :, :] = seg\n",
    "    masks[i, 0, :, :, :] = mask\n",
    "    feats[i, :, :, : ,:] = mind\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def iter_lbp(kpts_fixed, feat_kpts_fixed, feat_moving, sps):\n",
    "    edges, edges_reverse_idx = lbp_graph(kpts_fixed, k)\n",
    "    N_edges = edges.shape[0]\n",
    "\n",
    "    flow = torch.zeros_like(kpts_fixed) \n",
    "    for sp in sps:\n",
    "        candidates = flow.view(1, -1, 1, 3) + (torch.rand((1, N, k1, 3), device=device) - 0.5) * sp\n",
    "        candidates_edges0 = candidates[0, edges[:, 0], :, :]\n",
    "        candidates_edges1 = candidates[0, edges[:, 1], :, :]\n",
    "\n",
    "        feat_kpts_moving = F.grid_sample(feat_moving, kpts_fixed.view(1, 1, -1, 1, 3) + candidates.view(1, 1, -1, k1, 3)).view(1, -1, N, k1).permute(0, 2, 3, 1)\n",
    "        candidates_cost = alpha*(feat_kpts_fixed.unsqueeze(2)-feat_kpts_moving).pow(2).mean(3)\n",
    "\n",
    "        messages = torch.zeros((N_edges, k1), device=device)\n",
    "        temp_messages = torch.zeros((N, k1), device=device)\n",
    "        for _ in range(lbp_iter):\n",
    "            multi_data_cost = torch.gather(temp_messages + candidates_cost.view(-1, k1), 0, edges[:,0].view(-1, 1).expand(-1, k1))\n",
    "\n",
    "            reverse_message = torch.gather(messages, 0, edges_reverse_idx.view(-1, 1).expand(-1, k1))\n",
    "            multi_data_cost -= reverse_message\n",
    "\n",
    "            messages = sparse_minconv(multi_data_cost, candidates_edges0, candidates_edges1)\n",
    "\n",
    "            torch.fill_(temp_messages, 0)\n",
    "            temp_messages.scatter_add_(0, edges[:,1].view(-1, 1).expand(-1, k1), messages)\n",
    "\n",
    "        marginals = temp_messages + candidates_cost.view(-1, k1)\n",
    "\n",
    "        flow = torch.sum(F.softmax(-alpha1 * marginals, 1).unsqueeze(2) * candidates.view(N, k1, 3), 1).view(1, N, 3)\n",
    "\n",
    "    return flow\n",
    "        \n",
    "dice_initial = torch.zeros((len(pairs), 13), device=device)\n",
    "dice_affine = torch.zeros((len(pairs), 13), device=device)\n",
    "dice_affine1 = torch.zeros((len(pairs), 13), device=device)\n",
    "dice = torch.zeros((len(pairs), 13), device=device)\n",
    "sd_log_j = torch.zeros((len(pairs),), device=device)\n",
    "runtimes = torch.zeros((len(pairs),), device=device)\n",
    "runtimes1 = torch.zeros((len(pairs),), device=device)\n",
    "\n",
    "torch.manual_seed(30)\n",
    "for i, pair in enumerate(pairs):\n",
    "    idx_fixed = ind[pair[0]]\n",
    "    idx_moving = ind[pair[1]]\n",
    "\n",
    "    img_fixed = imgs[idx_fixed:idx_fixed+1].to(device, non_blocking=True)\n",
    "    seg_fixed = segs[idx_fixed:idx_fixed+1].to(device, non_blocking=True)\n",
    "    mask_fixed = masks[idx_fixed:idx_fixed+1].to(device, non_blocking=True)\n",
    "    feat_fixed = feats[idx_fixed:idx_fixed+1].to(device, non_blocking=True)\n",
    "\n",
    "    img_moving = imgs[idx_moving:idx_moving+1].to(device, non_blocking=True)\n",
    "    seg_moving = segs[idx_moving:idx_moving+1].to(device, non_blocking=True)\n",
    "    mask_moving = masks[idx_moving:idx_moving+1].to(device, non_blocking=True)\n",
    "    feat_moving = feats[idx_moving:idx_moving+1].to(device, non_blocking=True)\n",
    "    \n",
    "    kpts_fixed = random_kpts(mask_fixed, d, num_points=N)\n",
    "    feat_kpts_fixed = F.grid_sample(feat_fixed, kpts_fixed.view(1, 1, 1, -1, 3)).view(1, -1, N).permute(0, 2, 1)\n",
    "\n",
    "    torch.cuda.synchronize()\n",
    "    t0 = time.time()\n",
    "\n",
    "    flow = iter_lbp(kpts_fixed, feat_kpts_fixed, feat_moving, sps)\n",
    "\n",
    "    torch.cuda.synchronize()\n",
    "    t1 = time.time()\n",
    "\n",
    "    dense_flow = thin_plate_dense(kpts_fixed.cuda(), flow.cuda(), (D, H, W), 3, 0.001)\n",
    "    seg_moving_warped = F.grid_sample(seg_moving.float().cuda(), F.affine_grid(torch.eye(3,4,device='cuda').unsqueeze(0), (1,1,D,H,W)) + dense_flow, mode='nearest').to(device)\n",
    "    img_moving_warped = F.grid_sample(img_moving+3024.0005, F.affine_grid(torch.eye(3,4,device='cuda').unsqueeze(0), (1,1,D,H,W)) + dense_flow, mode='bilinear').to(device)-3024.0005\n",
    "    jac_det = jacobian_determinant(flow_world(dense_flow.view(1, -1, 3), (D//2, H//2, W//2)).view_as(dense_flow)[:, ::2, ::2, ::2, :]).to(device)\n",
    "    \n",
    "    dice_initial[i] = dice_coeff(seg_moving, seg_fixed, 13)\n",
    "    dice[i] = dice_coeff(seg_moving_warped, seg_fixed, 13)\n",
    "    sd_log_j[i] = torch.log((jac_det + 3).clamp_(0.000000001, 1000000000)).std()\n",
    "    runtimes[i] = t1-t0\n",
    "\n",
    "    print('Fixed: {}, Moving: {}'.format(pair[0], pair[1]))\n",
    "    print('Initial Dice: {:.2f}'.format(dice_initial[i].mean()))\n",
    "    print('Dice: {:.2f}'.format(dice[i].mean()))\n",
    "    print('SDlogJ: {:.2f}'.format(sd_log_j[i]))\n",
    "    print('Runtime: {:.2f} s'.format(runtimes[i]))\n",
    "    print()\n",
    "\n",
    "print('---')\n",
    "print('Mean Initial Dice: {:.3f}'.format(dice_initial.mean()))\n",
    "print('Mean Dice: {:.3f}'.format(dice.mean()))\n",
    "print('Mean SDlogJ: {:.3f}'.format(sd_log_j.mean()))\n",
    "print('Mean Runtime: {:.3f} s'.format(runtimes.mean()))\n"
   ]
  },
  {
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}