test_main_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'\n",
    "\n",
    "import cv2\n",
    "import glob\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as dset\n",
    "import torch.nn.functional as F\n",
    "import torchvision.utils as vutils\n",
    "import pickle\n",
    "from PIL import ImageFile\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "\n",
    "from sklearn.metrics import jaccard_score\n",
    "from unet import UNet\n",
    "from sklearn.model_selection import train_test_split\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from sklearn.metrics import recall_score\n",
    "from sklearn.metrics import confusion_matrix\n",
    "\n",
    "from albumentations import (\n",
    "    HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n",
    "    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n",
    "    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine,\n",
    "    IAASharpen, IAAEmboss, RandomContrast, RandomBrightness, Flip, OneOf, Compose,\n",
    "    RandomCrop, Normalize, Resize\n",
    ")\n",
    "\n",
    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "def show(img):\n",
    "    npimg = img.detach().numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "img_size = 128\n",
    "IMG_HEIGHT, IMG_WIDTH = 128, 128\n",
    "batchSize = 32\n",
    "use_cuda = torch.cuda.is_available()\n",
    "  \n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print (device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_cuda:\n",
    "#     net = UNet(n_channels=3, n_classes=1).cuda()\n",
    "    net = nn.DataParallel(UNet(n_channels=4, n_classes=1)).cuda()\n",
    "else:\n",
    "    net = UNet(n_channels=3, n_classes=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model\n",
    "### Note: you should change the path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# net.load_state_dict(torch.load('./results/patch_withPred_ver2/net_025.pth'))\n",
    "net.load_state_dict(torch.load('/home/quang/working/fundus_segmentation/model_paper/results/patch_withPred_ver2/net_025.pth'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict whole image\n",
    "### Note: you should change the path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8196\n"
     ]
    }
   ],
   "source": [
    "IMG_SIZE = 1200\n",
    "PATCH_SIZE = 128\n",
    "STEP = 64\n",
    "\n",
    "# INPUT_DIR_FULL_IMAGE = '/home/quang/working/fundus_segmentation/data/segmentation_doctor/data_full/'\n",
    "# INPUT_DIR_MASK_PRED_FULL_IMAGE = '/home/quang/working/fundus_segmentation/data/segmentation_doctor/data_full_mask_pred/'\n",
    "\n",
    "INPUT_DIR_FULL_IMAGE = '/home/quang/working/Face-Aging-CAAE/data/data_amd_resize_1200_8196/'\n",
    "INPUT_DIR_MASK_PRED_FULL_IMAGE = '/home/quang/working/fundus_segmentation/data/segmentation_doctor/prediction_data_amd_512_8196/'\n",
    "\n",
    "# OUTPUT_DIR_MASK_PRED_FULL_IMAGE = './predictions/main_model/'\n",
    "OUTPUT_DIR_MASK_PRED_FULL_IMAGE = '/home/quang/working/fundus_segmentation/data/segmentation_doctor/prediction_data_amd_1200_8196/'\n",
    "\n",
    "list_paths_img = glob.glob(INPUT_DIR_FULL_IMAGE + '*.jpg')\n",
    "\n",
    "list_name_f = [x.split('/')[-1].split('.')[0] for x in list_paths_img]\n",
    "print (len(list_name_f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_album_FULL_IMAGE = Compose([\n",
    "    Resize(int(PATCH_SIZE), int(PATCH_SIZE)), \n",
    "], additional_targets = {'image0': 'image', 'mask_pred': 'mask'})\n",
    "\n",
    "test_transform_torch_FULL_IMAGE = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "    ])\n",
    "\n",
    "test_transform_torch_mask_pred_FULL_IMAGE = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5,), (0.5,))\n",
    "    ])\n",
    "\n",
    "toTensor = transforms.ToTensor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/8196 [00:00<?, ?it/s]/home/quang/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py:1350: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
      "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n",
      "100%|██████████| 8196/8196 [4:22:36<00:00,  1.91s/it]  \n"
     ]
    }
   ],
   "source": [
    "net.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    for temp_name in tqdm(list_name_f):\n",
    "        img = cv2.imread(INPUT_DIR_FULL_IMAGE + temp_name + '.jpg')\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #??????\n",
    "        img_rescaled = cv2.resize(img, (IMG_SIZE, IMG_SIZE), interpolation = cv2.INTER_CUBIC)\n",
    "\n",
    "        mask_pred = cv2.imread(INPUT_DIR_MASK_PRED_FULL_IMAGE + temp_name + '.jpg', 0)\n",
    "        mask_pred_rescaled = cv2.resize(mask_pred, (IMG_SIZE, IMG_SIZE), interpolation = cv2.INTER_CUBIC)\n",
    "        mask_pred_rescaled = np.expand_dims(mask_pred_rescaled, axis=-1)\n",
    "\n",
    "        img_gray = cv2.cvtColor(img_rescaled, cv2.COLOR_BGR2GRAY)\n",
    "        img_gray_mask = np.where(img_gray > 10, 1, 0)\n",
    "        circle = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.uint8)\n",
    "        cv2.circle(circle, (int(IMG_SIZE/2), int(IMG_SIZE/2)), int(IMG_SIZE/2), 1, thickness=-1)\n",
    "\n",
    "        img_gray_mask = img_gray_mask*circle\n",
    "\n",
    "        mask_predict = np.zeros((IMG_SIZE,IMG_SIZE))\n",
    "        mask_predict_count = np.zeros((IMG_SIZE,IMG_SIZE))\n",
    "\n",
    "        for i in range(0, IMG_SIZE-PATCH_SIZE, STEP):\n",
    "            for j in range(0, IMG_SIZE-PATCH_SIZE, STEP):\n",
    "                if img_gray_mask[i:i+PATCH_SIZE, j:j+PATCH_SIZE].sum() > PATCH_SIZE*PATCH_SIZE/2 or True:\n",
    "                    img_patch = img_rescaled[i:i+PATCH_SIZE, j:j+PATCH_SIZE,:]\n",
    "                    mask_pred_patch = mask_pred_rescaled[i:i+PATCH_SIZE, j:j+PATCH_SIZE]\n",
    "\n",
    "                    img_auged = test_album_FULL_IMAGE(image=img_patch, mask_pred=mask_pred_patch)\n",
    "\n",
    "                    image_tensor = img_auged['image']\n",
    "                    mask_pred_aug = img_auged['mask_pred']\n",
    "                    image_tensor = test_transform_torch_FULL_IMAGE(image_tensor)\n",
    "                    mask_pred_tensor = test_transform_torch_mask_pred_FULL_IMAGE(mask_pred_aug).float()\n",
    "                    image_tensor = image_tensor.cuda()\n",
    "                    mask_pred_tensor = mask_pred_tensor.cuda()\n",
    "                    data_catted = torch.cat([image_tensor, mask_pred_tensor], 0)\n",
    "                    masks_pred = net(torch.unsqueeze(data_catted, 0))\n",
    "                    mask_predict[i:i+PATCH_SIZE, j:j+PATCH_SIZE] += masks_pred.cpu().numpy()[0,0]\n",
    "                    mask_predict_count[i:i+PATCH_SIZE, j:j+PATCH_SIZE] += 1\n",
    "\n",
    "        mask_predict_count_temp = np.where(mask_predict_count == 0, 1, mask_predict_count)\n",
    "        mask_predict_avg = mask_predict/mask_predict_count_temp\n",
    "        mask_predict_avg_binary = np.where(mask_predict_avg > 0.5, 1, 0)\n",
    "\n",
    "        cv2.imwrite(OUTPUT_DIR_MASK_PRED_FULL_IMAGE + temp_name + '.png', mask_predict_avg_binary*255)"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}