R2U_Net_Training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import R2U_Net\n",
    "u_net = R2U_Net()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DataSetClass import RetinaDataset \n",
    "tr = torch.utils.data.DataLoader(RetinaDataset('data/training/','data/manual/'), batch_size=8, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "optimizer = optim.Adam(u_net.parameters())\n",
    "# criterion = F.BCELoss()\n",
    "# criterion = nn.CrossEntropyLoss()\n",
    "criterion = nn.functional.mse_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1,     1) item-loss: 0.160448119\n",
      "(1,    11) item-loss: 0.132301971\n",
      "(1,    21) item-loss: 0.131982177\n",
      "(1,    31) item-loss: 0.120753273\n",
      "(1,    41) item-loss: 0.122668535\n",
      "[1] epoch-loss: 6.238536499\n",
      "(2,     1) item-loss: 0.118186601\n",
      "(2,    11) item-loss: 0.091019548\n",
      "(2,    21) item-loss: 0.093979277\n",
      "(2,    31) item-loss: 0.101368628\n",
      "(2,    41) item-loss: 0.075588480\n",
      "[2] epoch-loss: 4.983094707\n",
      "(3,     1) item-loss: 0.079567812\n",
      "(3,    11) item-loss: 0.090454176\n",
      "(3,    21) item-loss: 0.076946020\n",
      "(3,    31) item-loss: 0.076976016\n",
      "(3,    41) item-loss: 0.092525363\n",
      "[3] epoch-loss: 4.611879103\n",
      "(4,     1) item-loss: 0.098241948\n",
      "(4,    11) item-loss: 0.087108977\n",
      "(4,    21) item-loss: 0.070236750\n",
      "(4,    31) item-loss: 0.084769100\n",
      "(4,    41) item-loss: 0.081391819\n",
      "[4] epoch-loss: 4.473543867\n",
      "(5,     1) item-loss: 0.097281545\n",
      "(5,    11) item-loss: 0.086161196\n",
      "(5,    21) item-loss: 0.079491332\n",
      "(5,    31) item-loss: 0.109092072\n",
      "(5,    41) item-loss: 0.083614387\n",
      "[5] epoch-loss: 4.370078616\n",
      "(6,     1) item-loss: 0.099583343\n",
      "(6,    11) item-loss: 0.090123028\n",
      "(6,    21) item-loss: 0.099926263\n",
      "(6,    31) item-loss: 0.078724906\n",
      "(6,    41) item-loss: 0.084340855\n",
      "[6] epoch-loss: 4.313848421\n",
      "(7,     1) item-loss: 0.084514566\n",
      "(7,    11) item-loss: 0.067235291\n",
      "(7,    21) item-loss: 0.095349453\n",
      "(7,    31) item-loss: 0.093388766\n",
      "(7,    41) item-loss: 0.065173149\n",
      "[7] epoch-loss: 4.237556592\n",
      "(8,     1) item-loss: 0.095340289\n",
      "(8,    11) item-loss: 0.100034885\n",
      "(8,    21) item-loss: 0.092104390\n",
      "(8,    31) item-loss: 0.064848050\n",
      "(8,    41) item-loss: 0.067972213\n",
      "[8] epoch-loss: 4.244976271\n",
      "(9,     1) item-loss: 0.073186815\n",
      "(9,    11) item-loss: 0.063353777\n",
      "(9,    21) item-loss: 0.054400876\n",
      "(9,    31) item-loss: 0.070045143\n",
      "(9,    41) item-loss: 0.064166404\n",
      "[9] epoch-loss: 4.176835198\n",
      "(10,     1) item-loss: 0.096050285\n",
      "(10,    11) item-loss: 0.105491966\n",
      "(10,    21) item-loss: 0.092738539\n",
      "(10,    31) item-loss: 0.093130559\n",
      "(10,    41) item-loss: 0.090923518\n",
      "[10] epoch-loss: 4.146363873\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "loss_history = []\n",
    "for epoch in range(10):\n",
    "    train_loss = 0.0\n",
    "    for i,data in enumerate(tr):\n",
    "        x, y = data\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = u_net(x)\n",
    "        loss = criterion(y_pred, y)\n",
    "        x = loss.item()\n",
    "        train_loss += x\n",
    "        if (i%10==0):\n",
    "            print('(%d, %5d) item-loss: %.9f'%(epoch + 1, i + 1, x))\n",
    "        loss.backward()\n",
    "        optimizer.step()   \n",
    "    print('[%d] epoch-loss: %.9f'%(epoch + 1, train_loss))\n",
    "    loss_history.append(train_loss)\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.plot(loss_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##saving the model\n",
    "PATH = 'model/retina_unet_256_16_msef2.pth'\n",
    "torch.save(u_net.state_dict(), PATH)\n",
    "##loading the model\n",
    "model = U_Net()\n",
    "model.load_state_dict(torch.load(PATH))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Showing 5 random segmentations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torchvision\n",
    "\n",
    "for i in range(2):\n",
    "    idx = np.random.randint(100)\n",
    "    im1 = Image.open('data/training/' + str(idx)  + '_orig.jpg')\n",
    "    im2 = Image.open('data/training/' + str(idx)  + '_rot.jpg')\n",
    "    im3 = Image.open('data/training/' + str(idx)  + '_noise.jpg')\n",
    "\n",
    "    s1 = torchvision.transforms.ToTensor()(im1).unsqueeze(0)\n",
    "    s2 = torchvision.transforms.ToTensor()(im2).unsqueeze(0)\n",
    "    s3 = torchvision.transforms.ToTensor()(im3).unsqueeze(0)\n",
    "\n",
    "    pred1 = model(s1).cpu().detach().numpy()\n",
    "    pred2 = model(s2).cpu().detach().numpy()\n",
    "    pred3 = model(s3).cpu().detach().numpy()\n",
    "\n",
    "    f = plt.figure()\n",
    "    f.add_subplot(2,3, 1)\n",
    "    plt.imshow(im1)\n",
    "    f.add_subplot(2,3, 2)\n",
    "    plt.imshow(im2)\n",
    "    f.add_subplot(2,3, 3)\n",
    "    plt.imshow(im3)\n",
    "    f.add_subplot(2,3, 4)\n",
    "    plt.imshow(pred1[0][0],cmap='gray')\n",
    "    f.add_subplot(2,3, 5)\n",
    "    plt.imshow(pred2[0][0],cmap='gray')\n",
    "    f.add_subplot(2,3, 6)\n",
    "    plt.imshow(pred3[0][0],cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(10):\n",
    "    train_loss = 0.0\n",
    "    for i,data in enumerate(tr):\n",
    "        x, y = data\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = u_net(x)\n",
    "        loss = F.mse_loss(y_pred, y)\n",
    "        x = loss.item()\n",
    "        train_loss += x\n",
    "        if (i%10==0):\n",
    "            print('(%d, %5d) item-loss: %.9f'%(epoch + 1, i + 1, x))\n",
    "        loss.backward()\n",
    "        optimizer.step()   \n",
    "    print('[%d] epoch-loss: %.9f'%(epoch + 1, train_loss))\n",
    "    loss_history.append(train_loss)\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(10):\n",
    "    train_loss = 0.0\n",
    "    for i,data in enumerate(tr):\n",
    "        x, y = data\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = u_net(x)\n",
    "        loss = F.mse_loss(y_pred, y)\n",
    "        x = loss.item()\n",
    "        train_loss += x\n",
    "        if (i%10==0):\n",
    "            print('(%d, %5d) item-loss: %.9f'%(epoch + 1, i + 1, x))\n",
    "        loss.backward()\n",
    "        optimizer.step()   \n",
    "    print('[%d] epoch-loss: %.9f'%(epoch + 1, train_loss))\n",
    "    loss_history.append(train_loss)\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.plot(loss_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##saving the model\n",
    "PATH = 'model/retina_unet_256_16_msef2.pth'\n",
    "# torch.save(u_net.state_dict(), PATH)\n",
    "##loading the model\n",
    "model = U_Net()\n",
    "model.load_state_dict(torch.load(PATH))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Showing 2 random segmentations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torchvision\n",
    "\n",
    "for i in range(2):\n",
    "    idx = np.random.randint(100)\n",
    "    im1 = Image.open('data/training/' + str(idx)  + '_orig.jpg')\n",
    "    im2 = Image.open('data/training/' + str(idx)  + '_rot.jpg')\n",
    "    im3 = Image.open('data/training/' + str(idx)  + '_noise.jpg')\n",
    "\n",
    "    s1 = torchvision.transforms.ToTensor()(im1).unsqueeze(0)\n",
    "    s2 = torchvision.transforms.ToTensor()(im2).unsqueeze(0)\n",
    "    s3 = torchvision.transforms.ToTensor()(im3).unsqueeze(0)\n",
    "\n",
    "    pred1 = model(s1).cpu().detach().numpy()\n",
    "    pred2 = model(s2).cpu().detach().numpy()\n",
    "    pred3 = model(s3).cpu().detach().numpy()\n",
    "\n",
    "    f = plt.figure()\n",
    "    f.add_subplot(2,3, 1)\n",
    "    plt.imshow(im1)\n",
    "    f.add_subplot(2,3, 2)\n",
    "    plt.imshow(im2)\n",
    "    f.add_subplot(2,3, 3)\n",
    "    plt.imshow(im3)\n",
    "    f.add_subplot(2,3, 4)\n",
    "    plt.imshow(pred1[0][0],cmap='gray')\n",
    "    f.add_subplot(2,3, 5)\n",
    "    plt.imshow(pred2[0][0],cmap='gray')\n",
    "    f.add_subplot(2,3, 6)\n",
    "    plt.imshow(pred3[0][0],cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Results on Test Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Showing 2 random segmentations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torchvision\n",
    "\n",
    "for i in range(2):\n",
    "    idx = np.random.randint(20)\n",
    "    im1 = Image.open('data/testing/' + str(idx)  + '_orig.jpg')\n",
    "    im2 = Image.open('data/testing/' + str(idx)  + '_rot.jpg')\n",
    "    im3 = Image.open('data/testing/' + str(idx)  + '_noise.jpg')\n",
    "    im4 = Image.open('data/testing_manual/' + str(idx)  + '_orig.jpg')\n",
    "    im5 = Image.open('data/testing_manual/' + str(idx)  + '_rot.jpg')\n",
    "    im6 = Image.open('data/testing_manual/' + str(idx)  + '_noise.jpg')\n",
    "\n",
    "    s1 = torchvision.transforms.ToTensor()(im1).unsqueeze(0)\n",
    "    s2 = torchvision.transforms.ToTensor()(im2).unsqueeze(0)\n",
    "    s3 = torchvision.transforms.ToTensor()(im3).unsqueeze(0)\n",
    "\n",
    "    pred1 = model(s1).cpu().detach().numpy()\n",
    "    pred2 = model(s2).cpu().detach().numpy()\n",
    "    pred3 = model(s3).cpu().detach().numpy()\n",
    "\n",
    "    f = plt.figure()\n",
    "    f.add_subplot(3,3, 1)\n",
    "    plt.imshow(im1)\n",
    "    f.add_subplot(3,3, 2)\n",
    "    plt.imshow(im2)\n",
    "    f.add_subplot(3,3, 3)\n",
    "    plt.imshow(im3)\n",
    "    f.add_subplot(3,3, 4)\n",
    "    plt.imshow(pred1[0][0],cmap='gray')\n",
    "    f.add_subplot(3,3, 5)\n",
    "    plt.imshow(pred2[0][0],cmap='gray')\n",
    "    f.add_subplot(3,3, 6)\n",
    "    plt.imshow(pred3[0][0],cmap='gray')\n",
    "    f.add_subplot(3,3, 7)\n",
    "    plt.imshow(im4,cmap='gray')\n",
    "    f.add_subplot(3,3, 8)\n",
    "    plt.imshow(im5,cmap='gray')\n",
    "    f.add_subplot(3,3, 9)\n",
    "    plt.imshow(im6,cmap='gray')"
   ]
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}