dense_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 minconv(input):\n",
    "    device = input.device\n",
    "    \n",
    "    disp1d = disp.view(1, disp_width, disp_width, disp_width, 3)[0,0,0,:,0]\n",
    "    regular1d = (disp1d.view(1,-1) - disp1d.view(-1,1)) ** 2\n",
    "    output = torch.min( input.view(-1, disp_width, 1, disp_width, disp_width) + regular1d.view(1, disp_width, disp_width, 1, 1), 1)[0]\n",
    "    \n",
    "    disp1d = disp.view(1, disp_width, disp_width, disp_width, 3)[0,0,:,0,1]\n",
    "    regular1d = (disp1d.view(1,-1) - disp1d.view(-1,1)) ** 2\n",
    "    output = torch.min(output.view(-1, disp_width, disp_width, 1, disp_width) + regular1d.view(1, 1, disp_width, disp_width, 1), 2)[0]\n",
    "    \n",
    "    disp1d = disp.view(1, disp_width, disp_width, disp_width, 3)[0,:,0,0,2]\n",
    "    regular1d = (disp1d.view(1,-1) - disp1d.view(-1,1)) ** 2\n",
    "    output = torch.min(output.view(-1, disp_width, disp_width, disp_width, 1) + regular1d.view(1, 1, 1, disp_width, disp_width), 3)[0]\n",
    "    \n",
    "    output = output - (torch.min(output.view(-1, disp_width**3), 1)[0]).view(-1, 1, 1, 1)\n",
    "\n",
    "    return output.view_as(input)\n",
    "\n",
    "def dense_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",
    "    candidates = disp.view(1, 1, -1, 3).expand(-1, N, -1, -1)\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 = torch.zeros(1, N, k1, device=device)\n",
    "    for j in range(int(math.ceil(N/128))):\n",
    "        j1 = j*128\n",
    "        j2 = min((j+1)*128,N) \n",
    "        candidates_cost[:, j1:j2, :] = alpha*(feat_kpts_fixed[:,j1:j2, :].unsqueeze(2)-feat_kpts_moving[:,j1:j2,:,:]).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 = minconv(multi_data_cost)\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",
    "def get_disp(disp_step, disp_radius):\n",
    "    disp = torch.stack(torch.meshgrid(torch.arange(- disp_step * disp_radius, disp_step * disp_radius + 1, disp_step, device=device),\n",
    "                                      torch.arange(- disp_step * disp_radius, disp_step * disp_radius + 1, disp_step, device=device),\n",
    "                                      torch.arange(- disp_step * disp_radius, disp_step * disp_radius + 1, disp_step, device=device)), dim=3).view(1, -1, 3)\n",
    "    disp = flow_pt(disp, (D, H, W)).view(1, 1, -1, 1, 3)\n",
    "    return disp\n",
    "\n",
    "disp_radius = 5\n",
    "disp_width = disp_radius*2+1\n",
    "disp_step = 4\n",
    "disp = get_disp(disp_step, disp_radius)\n",
    "k1 = disp_width ** 3\n",
    "        \n",
    "dice_initial = 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",
    "\n",
    "with torch.no_grad():\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 = dense_lbp(kpts_fixed, feat_kpts_fixed, feat_moving, sps)\n",
    "\n",
    "        torch.cuda.synchronize()\n",
    "        t1 = time.time()\n",
    "\n",
    "        torch.cuda.empty_cache()\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",
    "        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, :])\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",
    "        del feat_fixed\n",
    "        del feat_moving\n",
    "        del flow\n",
    "        del jac_det\n",
    "        del dense_flow\n",
    "        del seg_moving_warped\n",
    "        torch.cuda.empty_cache()\n",
    "        torch.cuda.empty_cache()\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
}