CRF_VOC_2012.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.optim\n",
    "import torch.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import torchvision\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "import torch.optim as optim\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "\n",
    "from include.data import ObjectSegmentationDataset\n",
    "from include.hybrid_net import SegmentationModel\n",
    "\n",
    "from include.Utility_Functions import Validate_IOU\n",
    "from torch.nn.modules.upsampling import Upsample\n",
    "\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from inferno.extensions.criteria.set_similarity_measures import SorensenDiceLoss\n",
    "\n",
    "from include.Utility_Functions import IOU\n",
    "from include import crf\n",
    "import time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loaders for Training and Validation\n",
    "\n",
    "# It is important to normalise the Input images acording to\n",
    "# https://pytorch.org/docs/stable/torchvision/models.html\n",
    "\n",
    "norm_transform = transforms.Normalize(mean=[0.45752926384356873,0.4377081874543038,0.40432555437277296],std=[0.2664644516691583,0.2634024345463397,0.2775109722016356])\n",
    "transform = transforms.Compose([transforms.ToTensor(),norm_transform])\n",
    "number_of_classes = 21\n",
    "crf_dataset = ObjectSegmentationDataset(\n",
    "    src_image_dir=\"/home/snajder/AML/data/VOCdevkit/VOC2012/VAL/Data\",\n",
    "    seg_image_dir=\"/home/snajder/AML/data/VOCdevkit/VOC2012/VAL/Ground_Truth\",\n",
    "    num_classes=number_of_classes,\n",
    "    transform=transform,\n",
    "    rescale=False,\n",
    "    gt_one_hot=False,\n",
    "    augment=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = SegmentationModel(num_classes=20)\n",
    "net_file=\"IOU_Best_params.pth\"\n",
    "checkpoint = torch.load(net_file)\n",
    "net.load_state_dict(checkpoint['state_dict'])\n",
    "net = net.cuda()\n",
    "net.eval()\n",
    "pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs=1\n",
    "\n",
    "dataset_len = len(crf_dataset)\n",
    "indices = list(range(dataset_len))\n",
    "num_train = 100\n",
    "np.random.seed(43)\n",
    "train_idx = np.random.choice(indices, size=num_train, replace=False)\n",
    "test_idx = list(set(indices) - set(train_idx))\n",
    "np.random.seed(int(time.time()))\n",
    "Train_Loader = torch.utils.data.DataLoader(crf_dataset, batch_size=1, num_workers=4, sampler=SubsetRandomSampler(train_idx))\n",
    "Test_Loader = torch.utils.data.DataLoader(crf_dataset, batch_size=1, num_workers=4, sampler=SubsetRandomSampler(test_idx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_improvement = crf.gridSearchCRFParameters(net, Train_Loader, norm_transform,20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "best_params = np.unravel_index(np.argmax(total_improvement), total_improvement.shape)\n",
    "print(\"Best improvement\", total_improvement[best_params])\n",
    "print(\"Improvement with published parameters\", total_improvement[2,5,2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Reject prediction with really low IOU (not used)\n",
    "IOU_threshold = 0\n",
    "val_set_size = 0\n",
    "TP = np.zeros(20)\n",
    "FP = np.zeros(20)\n",
    "FN = np.zeros(20)\n",
    "TP_afterCRF= np.zeros(20)\n",
    "FP_afterCRF = np.zeros(20)\n",
    "FN_afterCRF = np.zeros(20)\n",
    "for data in Test_Loader:\n",
    "        src_img, seg_img, seg_img_ds, src_img_raw = data\n",
    "\n",
    "        Input = Variable(src_img, requires_grad=False).float().cuda()\n",
    "        Target = Variable(seg_img.long(), requires_grad=False).cuda()\n",
    "\n",
    "        Output = net(Input)\n",
    "        upsampler = nn.Upsample(size=(src_img.shape[2],src_img.shape[3]), mode='bilinear', align_corners=True)\n",
    "        Output_upsampled = upsampler(Output)\n",
    "        out=nn.functional.softmax(Output_upsampled[:,:,:,:], dim=1).cpu().detach().numpy()\n",
    "        Result=np.argmax(out, axis=1)\n",
    "        for i in range(Output.shape[0]):\n",
    "            GT=Target[i,:,:].cpu().numpy()\n",
    "            for n in range(20):\n",
    "                A = GT == n\n",
    "                B = Result[i] == n\n",
    "                TP[n] += np.sum(B & A)\n",
    "                FP[n] += np.sum(B & ~A)\n",
    "                FN[n] += np.sum(~B & A)\n",
    "            \n",
    "    \n",
    "            val_set_size += 1\n",
    "            # now apply CRF\n",
    "            unary = out[i].reshape((20,-1))\n",
    "            unary = - np.log(unary)\n",
    "            image = (src_img_raw[i]*255).byte().cpu().detach().numpy()\n",
    "            after_crf = crf.applyDenseCRF(unary, image,70,3,3,5,3)\n",
    "            after_crf = np.array(after_crf).reshape(20,image.shape[0],image.shape[1])\n",
    "            Result_after_crf = np.argmax(after_crf, axis=0)\n",
    "            for n in range(20):\n",
    "                A = GT == n\n",
    "                B = Result_after_crf == n\n",
    "                TP_afterCRF[n] += np.sum(B & A)\n",
    "                FP_afterCRF[n] += np.sum(B & ~A)\n",
    "                FN_afterCRF[n] += np.sum(~B & A)\n",
    "        \n",
    "        IOU = TP / ( TP + FP + FN)\n",
    "        IOU_afterCRF = TP_afterCRF / ( TP_afterCRF + FP_afterCRF + FN_afterCRF)\n",
    "        if val_set_size > 0:\n",
    "            print(\"Image %d, mean IOU w/o CRF: %f, mean IOU with CRF: %f\" % (val_set_size, np.mean(IOU), np.mean(IOU_afterCRF) ))\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "en = enumerate(Test_Loader, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\n",
    "test_input, test_gt, test_gt_ds, test_input_raw = next(en)[1]\n",
    "test_input_raw = test_input_raw[0].clone()\n",
    "test_gt = test_gt[0]\n",
    "test_gt_ds = test_gt_ds[0]\n",
    "test_input = test_input.cuda()\n",
    "m = nn.Upsample(size=(test_input.shape[2],test_input.shape[3]), mode='bilinear', align_corners=True)\n",
    "test_out = net(test_input)[0]\n",
    "test_out = nn.functional.softmax(test_out[:,:,:], dim=0)\n",
    "test_out = test_out[:,:,:]\n",
    "test_out_upsampled = m(test_out.unsqueeze(0)).cpu()\n",
    "_, test_out_labels = test_out_upsampled[0].max(dim=0)\n",
    "\n",
    "plt.imshow(test_input_raw)\n",
    "plt.show()\n",
    "\n",
    "unary_ori = test_out_upsampled[0].cpu().detach().numpy()\n",
    "unary = unary_ori.reshape((20,-1))\n",
    "unary = - np.log(unary)\n",
    "image = (test_input_raw*255).byte().cpu().detach().numpy()\n",
    "#image = np.transpose(image, axes=(1,2,0)).copy()\n",
    "after_crf = crf.applyDenseCRF(unary, image,70,5,3,5,3)\n",
    "#after_crf = np.log(after_crf)\n",
    "after_crf = np.array(after_crf).reshape(20,image.shape[0],image.shape[1])\n",
    "unary = unary.reshape(20,image.shape[0],image.shape[1])\n",
    "\n",
    "\n",
    "plt.imshow(test_gt.numpy())\n",
    "plt.show()\n",
    "plt.imshow(np.argmax(unary_ori,axis=0))\n",
    "plt.show()\n",
    "plt.imshow(np.argmax(after_crf,axis=0))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prints probability maps for classes\n",
    "for part in range(20):\n",
    "    print(\"Part \",part)\n",
    "    plt.imshow(unary_ori[part])\n",
    "    plt.show()\n",
    "    plt.imshow(after_crf[part])\n",
    "    plt.show()\n"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}