Training_&_Testing_SemiInfNet_+_MulticlassUNet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "**Pooyan Rezaeipour Lasaki**\n",
        "\n",
        "e-mails:\n",
        "rezaeipourpooyan@gmail.com &\n",
        "pooyan_rezaeipour@elec.iust.ac.ir\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "NBzk2gT4cbWc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Mohsen Safaei**\n",
        "\n",
        "e-mails: mfsafaei78@gmail.com & mo_safaei@elec.iust.ac.ir"
      ],
      "metadata": {
        "id": "4sF28w9ooSmx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Saeed Chamani**\n",
        "\n",
        "e-mails: saeed.chamani10@gmail.com\n",
        "& saeed_chamani@elec.iust.ac.ir\n"
      ],
      "metadata": {
        "id": "2DbUtbJ2ocOP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "*Biomedical Engineering Department, School of Electrical Engineering, \"Iran University of Science and Technology\", Tehran*"
      ],
      "metadata": {
        "id": "6Bz-krJCoYBE"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "*To find the files which you have uploaded in your \"Google Drive\"*"
      ],
      "metadata": {
        "id": "Q3zewxs-bj_v"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ixa4_nkha_gh"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Unzipping the mentioned file which is named \"Inf-Net-master\"*"
      ],
      "metadata": {
        "id": "_K4pW3h1bnkX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!unzip \"/content/gdrive/MyDrive/Inf-Net-master.zip\""
      ],
      "metadata": {
        "id": "1YG1syJcbfBI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*For recognizing the module named \"Code\"*"
      ],
      "metadata": {
        "id": "nT6Yvu7_cT-7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cd '/content/Inf-Net-master'"
      ],
      "metadata": {
        "id": "q9A5tTYGb55F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.path.append('/content/Inf-Net-master')"
      ],
      "metadata": {
        "id": "_CjMZsNTb9TS"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Training  \"MyTrain_MulClsLungInf_UNet\"   for Semi-Inf-Net + Multi-class UNet*"
      ],
      "metadata": {
        "id": "IF1yiJRNcuBa"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Note** that if you run this file uninterruptedly after Semi-Inf-net(second method), you do not face a problem because the folder named \"Snapshots\" was created automatically.\n",
        "\n",
        " But for running this file seperately(which we do below), you must create a folder in \"Snapshots\" which is named \"save_weights\",then create a folder in \"save_weights\" which is named \"{}\".\n",
        "\n",
        "Also change the last line which is \" save_path='Semi-Inf-Net_UNet' \" to \" save_path='{}' \""
      ],
      "metadata": {
        "id": "QkjXlXO8dJJJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# -*- coding: utf-8 -*-\n",
        "\n",
        "\"\"\"Preview\n",
        "Code for 'Inf-Net: Automatic COVID-19 Lung Infection Segmentation from CT Scans'\n",
        "submit to Transactions on Medical Imaging, 2020.\n",
        "\n",
        "First Version: Created on 2020-05-13 (@author: Ge-Peng Ji)\n",
        "\"\"\"\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import torch.optim as optim\n",
        "from Code.utils.dataloader_MulClsLungInf_UNet import LungDataset\n",
        "from torchvision import transforms\n",
        "# from LungData import test_dataloader, train_dataloader  # pls change batch_size\n",
        "from torch.utils.data import DataLoader\n",
        "from Code.model_lung_infection.InfNet_UNet import *\n",
        "\n",
        "\n",
        "def train(epo_num, num_classes, input_channels, batch_size, lr, save_path):\n",
        "    train_dataset = LungDataset(\n",
        "        imgs_path='./Dataset/COVID-SemiSeg/Dataset/TrainingSet/MultiClassInfection-Train/Imgs/',\n",
        "        # NOTES: prior is borrowed from the object-level label of train split\n",
        "        pseudo_path='./Dataset/COVID-SemiSeg/Dataset/TrainingSet/MultiClassInfection-Train/Prior/',\n",
        "        label_path='./Dataset/COVID-SemiSeg/Dataset/TrainingSet/MultiClassInfection-Train/GT/',\n",
        "        transform=transforms.Compose([\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]))\n",
        "    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\n",
        "\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "    lung_model = Inf_Net_UNet(input_channels, num_classes)  # input_channels=3， n_class=3\n",
        "    print(lung_model)\n",
        "    lung_model = lung_model.to(device)\n",
        "\n",
        "    criterion = nn.BCELoss().to(device)\n",
        "    optimizer = optim.SGD(lung_model.parameters(), lr=lr, momentum=0.7)\n",
        "\n",
        "    print(\"#\" * 20, \"\\nStart Training (Inf-Net)\\nThis code is written for 'Inf-Net: Automatic COVID-19 Lung \"\n",
        "                    \"Infection Segmentation from CT Scans', 2020, TMI.\\n\"\n",
        "                    \"----\\nPlease cite the paper if you use this code and dataset. \"\n",
        "                    \"And any questions feel free to contact me \"\n",
        "                    \"via E-mail (gepengai.ji@gmail.com)\\n----\\n\", \"#\" * 20)\n",
        "\n",
        "    for epo in range(epo_num):\n",
        "\n",
        "        train_loss = 0\n",
        "        lung_model.train()\n",
        "\n",
        "        for index, (img, pseudo, img_mask, _) in enumerate(train_dataloader):\n",
        "\n",
        "            img = img.to(device)\n",
        "            pseudo = pseudo.to(device)\n",
        "            img_mask = img_mask.to(device)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            output = lung_model(torch.cat((img, pseudo), dim=1))\n",
        "\n",
        "            output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])\n",
        "            loss = criterion(output, img_mask)\n",
        "\n",
        "            loss.backward()\n",
        "            iter_loss = loss.item()\n",
        "            train_loss += iter_loss\n",
        "            optimizer.step()\n",
        "\n",
        "            if np.mod(index, 20) == 0:\n",
        "                print('Epoch: {}/{}, Step: {}/{}, Train loss is {}'.format(epo, epo_num, index, len(train_dataloader), iter_loss))\n",
        "\n",
        "        os.makedirs('./checkpoints//UNet_Multi-Class-Semi', exist_ok=True)\n",
        "        if np.mod(epo+1, 10) == 0:\n",
        "            torch.save(lung_model.state_dict(),\n",
        "                       './Snapshots/save_weights/{}/unet_model_{}.pkl'.format(save_path, epo+1))\n",
        "            print('Saving checkpoints: unet_model_{}.pkl'.format(epo+1))\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    train(epo_num=200,\n",
        "          num_classes=3,\n",
        "          input_channels=6,\n",
        "          batch_size=16,\n",
        "          lr=1e-2,\n",
        "          save_path='{}')\n"
      ],
      "metadata": {
        "id": "qsEQFhNebvFK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Testing \"MyTest_MulClsLungInf_UNet\"   for Semi-Inf-Net + Multi-class UNet*"
      ],
      "metadata": {
        "id": "dww4tHi5hGdS"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For running this file seperately from second method(which we do below), change the one before the last line which is \" snapshot_dir='./Snapshots/save_weights/Semi-Inf-Net_UNet/unet_model_200.pkl' \" to  \n",
        " \" snapshot_dir='./Snapshots/save_weights/{}/unet_model_200.pkl' \""
      ],
      "metadata": {
        "id": "Wq3I-cmghukX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# -*- coding: utf-8 -*-\n",
        "\n",
        "\"\"\"Preview\n",
        "Code for 'Inf-Net: Automatic COVID-19 Lung Infection Segmentation from CT Scans'\n",
        "submit to Transactions on Medical Imaging, 2020.\n",
        "\n",
        "First Version: Created on 2020-05-13 (@author: Ge-Peng Ji)\n",
        "\"\"\"\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "from Code.utils.dataloader_MulClsLungInf_UNet import LungDataset\n",
        "from torchvision import transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from Code.model_lung_infection.InfNet_UNet import *  # use U-Net for multi-class segmentation\n",
        "from scipy import misc\n",
        "from Code.utils.split_class import split_class\n",
        "import shutil\n",
        "\n",
        "\n",
        "def inference(num_classes, input_channels, snapshot_dir, save_path):\n",
        "    test_dataset = LungDataset(\n",
        "        imgs_path='./Dataset/COVID-SemiSeg/Dataset/TestingSet/MultiClassInfection-Test/Imgs/',\n",
        "        pseudo_path='./Dataset/COVID-SemiSeg/Results/Lung infection segmentation/Semi-Inf-Net/',  # NOTES: generated from `Semi-Inf-Net`\n",
        "        label_path='./Dataset/COVID-SemiSeg/Dataset/TestingSet/MultiClassInfection-Test/GT/',\n",
        "        transform=transforms.Compose([\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize(mean=[0.485, 0.456, 0.406], \n",
        "                                 std=[0.229, 0.224, 0.225])]),\n",
        "        is_test=True\n",
        "    )\n",
        "    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=0)\n",
        "\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "    lung_model = Inf_Net_UNet(input_channels, num_classes).cuda()\n",
        "    print(lung_model)\n",
        "    lung_model.load_state_dict(torch.load(snapshot_dir))\n",
        "    lung_model.eval()\n",
        "\n",
        "    for index, (img, pseudo, img_mask, name) in enumerate(test_dataloader):\n",
        "        img = img.to(device)\n",
        "        pseudo = pseudo.to(device)\n",
        "        img_mask = img_mask.to(device)\n",
        "\n",
        "        output = lung_model(torch.cat((img, pseudo), dim=1))\n",
        "        output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])\n",
        "        b, _, w, h = output.size()\n",
        "        _, _, w_gt, h_gt = img_mask.size()\n",
        "\n",
        "        # output b*n_class*h*w -- > b*h*w\n",
        "        pred = output.cpu().permute(0, 2, 3, 1).contiguous().view(-1, num_classes).max(1)[1].view(b, w, h).numpy().squeeze()\n",
        "        print('Class numbers of prediction in total:', np.unique(pred))\n",
        "        # pred = misc.imresize(pred, size=(w_gt, h_gt))\n",
        "        os.makedirs(save_path, exist_ok=True)\n",
        "        #misc.imsave(save_path + name[0].replace('.jpg', '.png'), pred)\n",
        "        import imageio\n",
        "        imageio.imwrite(save_path + name[0].replace('.jpg', '.png'), pred)\n",
        "        split_class(save_path, name[0].replace('.jpg', '.png'), w_gt, h_gt)\n",
        "\n",
        "    shutil.rmtree(save_path)\n",
        "    print('Test done!')\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    inference(num_classes=3,\n",
        "              input_channels=6,\n",
        "              snapshot_dir='./Snapshots/save_weights/{}/unet_model_200.pkl',\n",
        "              save_path='./Results/Multi-class lung infection segmentation/class_12/'\n",
        "              )\n"
      ],
      "metadata": {
        "id": "o9b4KsLBhRqi"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}