Training_&_Testing_InfNet_&_SemiInfNet.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": "Nfvxi52_PdCu"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GCH1VY7EPBeM"
      },
      "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": "lrOeyY2zPzx5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!unzip \"/content/gdrive/MyDrive/Inf-Net-master.zip\""
      ],
      "metadata": {
        "id": "Y60Fu4r3Pxh3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Install the package named \"thop\"*"
      ],
      "metadata": {
        "id": "DlXTRYWLQVp5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install thop"
      ],
      "metadata": {
        "id": "7wKqn7SMQRo-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*For recognizing the module named \"Code\"*"
      ],
      "metadata": {
        "id": "nBLraA_eQy4q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cd '/content/Inf-Net-master'"
      ],
      "metadata": {
        "id": "3KWrSufZQvXw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.path.append('/content/Inf-Net-master')"
      ],
      "metadata": {
        "id": "EaTK03_mQv99"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*For solving the error of \"ipykernel_launcher.py: error: unrecognized arguments\"*"
      ],
      "metadata": {
        "id": "GQCjBsepRp3m"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install ipykernel"
      ],
      "metadata": {
        "id": "yY9orVlvRlZU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.argv=['-f']\n",
        "del sys"
      ],
      "metadata": {
        "id": "cDl9trNnRoIt"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Training   \"MyTrain_LungInf.py\"  for Inf-Net & Semi-Inf-Net*\n",
        "\n",
        "*The following code is adjusted for Inf-Net. If you want to change it to Semi-Inf-Net, follow the [README.md](https://github.com/DengPingFan/Inf-Net/blob/master/README.md)*"
      ],
      "metadata": {
        "id": "FLBiAGX4SHdy"
      }
    },
    {
      "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",
        "1st Version: Created on 2020-05-13 (@author: Ge-Peng Ji)\n",
        "2nd Version: Fix some bugs caused by THOP on 2020-06-10 (@author: Ge-Peng Ji)\n",
        "\"\"\"\n",
        "\n",
        "import torch\n",
        "from torch.autograd import Variable\n",
        "import os\n",
        "import argparse\n",
        "from datetime import datetime\n",
        "from Code.utils.dataloader_LungInf import get_loader\n",
        "from Code.utils.utils import clip_gradient, adjust_lr, AvgMeter\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "def joint_loss(pred, mask):\n",
        "    weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)\n",
        "    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduce='none')\n",
        "    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))\n",
        "\n",
        "    pred = torch.sigmoid(pred)\n",
        "    inter = ((pred * mask)*weit).sum(dim=(2, 3))\n",
        "    union = ((pred + mask)*weit).sum(dim=(2, 3))\n",
        "    wiou = 1 - (inter + 1)/(union - inter+1)\n",
        "    return (wbce + wiou).mean()\n",
        "\n",
        "\n",
        "def train(train_loader, model, optimizer, epoch, train_save):\n",
        "    model.train()\n",
        "    # ---- multi-scale training ----\n",
        "    size_rates = [0.75, 1, 1.25]    # replace your desired scale, try larger scale for better accuracy in small object\n",
        "    loss_record1, loss_record2, loss_record3, loss_record4, loss_record5 = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()\n",
        "    for i, pack in enumerate(train_loader, start=1):\n",
        "        for rate in size_rates:\n",
        "            optimizer.zero_grad()\n",
        "            # ---- data prepare ----\n",
        "            images, gts, edges = pack\n",
        "            images = Variable(images).cuda()\n",
        "            gts = Variable(gts).cuda()\n",
        "            edges = Variable(edges).cuda()\n",
        "            # ---- rescaling the inputs (img/gt/edge) ----\n",
        "            trainsize = int(round(opt.trainsize*rate/32)*32)\n",
        "            if rate != 1:\n",
        "                images = F.upsample(images, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n",
        "                gts = F.upsample(gts, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n",
        "                edges = F.upsample(edges, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n",
        "\n",
        "            # ---- forward ----\n",
        "            lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(images)\n",
        "            # ---- loss function ----\n",
        "            loss5 = joint_loss(lateral_map_5, gts)\n",
        "            loss4 = joint_loss(lateral_map_4, gts)\n",
        "            loss3 = joint_loss(lateral_map_3, gts)\n",
        "            loss2 = joint_loss(lateral_map_2, gts)\n",
        "            loss1 = BCE(lateral_edge, edges)\n",
        "            loss = loss1 + loss2 + loss3 + loss4 + loss5\n",
        "            # ---- backward ----\n",
        "            loss.backward()\n",
        "            clip_gradient(optimizer, opt.clip)\n",
        "            optimizer.step()\n",
        "            # ---- recording loss ----\n",
        "            if rate == 1:\n",
        "                loss_record1.update(loss1.data, opt.batchsize)\n",
        "                loss_record2.update(loss2.data, opt.batchsize)\n",
        "                loss_record3.update(loss3.data, opt.batchsize)\n",
        "                loss_record4.update(loss4.data, opt.batchsize)\n",
        "                loss_record5.update(loss5.data, opt.batchsize)\n",
        "        # ---- train logging ----\n",
        "        if i % 20 == 0 or i == total_step:\n",
        "            print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], [lateral-edge: {:.4f}, '\n",
        "                  'lateral-2: {:.4f}, lateral-3: {:0.4f}, lateral-4: {:0.4f}, lateral-5: {:0.4f}]'.\n",
        "                  format(datetime.now(), epoch, opt.epoch, i, total_step, loss_record1.show(),\n",
        "                         loss_record2.show(), loss_record3.show(), loss_record4.show(), loss_record5.show()))\n",
        "    # ---- save model_lung_infection ----\n",
        "    save_path = './Snapshots/save_weights/{}/'.format(train_save)\n",
        "    os.makedirs(save_path, exist_ok=True)\n",
        "\n",
        "    if (epoch+1) % 10 == 0:\n",
        "        torch.save(model.state_dict(), save_path + 'Inf-Net-%d.pth' % (epoch+1))\n",
        "        print('[Saving Snapshot:]', save_path + 'Inf-Net-%d.pth' % (epoch+1))\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    parser = argparse.ArgumentParser()\n",
        "    # hyper-parameters\n",
        "    parser.add_argument('--epoch', type=int, default=100,\n",
        "                        help='epoch number')\n",
        "    parser.add_argument('--lr', type=float, default=1e-4,\n",
        "                        help='learning rate')\n",
        "    parser.add_argument('--batchsize', type=int, default=24,\n",
        "                        help='training batch size')\n",
        "    parser.add_argument('--trainsize', type=int, default=352,\n",
        "                        help='set the size of training sample')\n",
        "    parser.add_argument('--clip', type=float, default=0.5,\n",
        "                        help='gradient clipping margin')\n",
        "    parser.add_argument('--decay_rate', type=float, default=0.1,\n",
        "                        help='decay rate of learning rate')\n",
        "    parser.add_argument('--decay_epoch', type=int, default=50,\n",
        "                        help='every n epochs decay learning rate')\n",
        "    parser.add_argument('--is_thop', type=bool, default=False,\n",
        "                        help='whether calculate FLOPs/Params (Thop)')\n",
        "    parser.add_argument('--gpu_device', type=int, default=0,\n",
        "                        help='choose which GPU device you want to use')\n",
        "    parser.add_argument('--num_workers', type=int, default=8,\n",
        "                        help='number of workers in dataloader. In windows, set num_workers=0')\n",
        "    # model_lung_infection parameters\n",
        "    parser.add_argument('--net_channel', type=int, default=32,\n",
        "                        help='internal channel numbers in the Inf-Net, default=32, try larger for better accuracy')\n",
        "    parser.add_argument('--n_classes', type=int, default=1,\n",
        "                        help='binary segmentation when n_classes=1')\n",
        "    parser.add_argument('--backbone', type=str, default='Res2Net50',\n",
        "                        help='change different backbone, choice: VGGNet16, ResNet50, Res2Net50')\n",
        "    # training dataset\n",
        "    parser.add_argument('--train_path', type=str,\n",
        "                        default='./Dataset/COVID-SemiSeg/Dataset/TrainingSet/LungInfection-Train/Doctor-label')\n",
        "    parser.add_argument('--is_semi', type=bool, default=False,\n",
        "                        help='if True, you will turn on the mode of `Semi-Inf-Net`')\n",
        "    parser.add_argument('--is_pseudo', type=bool, default=False,\n",
        "                        help='if True, you will train the model on pseudo-label')\n",
        "    parser.add_argument('--train_save', type=str, default=None,\n",
        "                        help='If you use custom save path, please edit `--is_semi=True` and `--is_pseudo=True`')\n",
        "\n",
        "    opt = parser.parse_args()\n",
        "\n",
        "    # ---- build models ----\n",
        "    torch.cuda.set_device(opt.gpu_device)\n",
        "    # - please asign your prefer backbone in opt.\n",
        "    if opt.backbone == 'Res2Net50':\n",
        "        print('Backbone loading: Res2Net50')\n",
        "        from Code.model_lung_infection.InfNet_Res2Net import Inf_Net\n",
        "    elif opt.backbone == 'ResNet50':\n",
        "        print('Backbone loading: ResNet50')\n",
        "        from Code.model_lung_infection.InfNet_ResNet import Inf_Net\n",
        "    elif opt.backbone == 'VGGNet16':\n",
        "        print('Backbone loading: VGGNet16')\n",
        "        from Code.model_lung_infection.InfNet_VGGNet import Inf_Net\n",
        "    else:\n",
        "        raise ValueError('Invalid backbone parameters: {}'.format(opt.backbone))\n",
        "    model = Inf_Net(channel=opt.net_channel, n_class=opt.n_classes).cuda()\n",
        "\n",
        "    # ---- load pre-trained weights (mode=Semi-Inf-Net) ----\n",
        "    # - See Sec.2.3 of `README.md` to learn how to generate your own img/pseudo-label from scratch.\n",
        "    if opt.is_semi and opt.backbone == 'Res2Net50':\n",
        "        print('Loading weights from weights file trained on pseudo label')\n",
        "        model.load_state_dict(torch.load('./Snapshots/save_weights/Inf-Net_Pseduo/Inf-Net_pseudo_100.pth'))\n",
        "    else:\n",
        "        print('Not loading weights from weights file')\n",
        "\n",
        "    # weights file save path\n",
        "    if opt.is_pseudo and (not opt.is_semi):\n",
        "        train_save = 'Inf-Net_Pseudo'\n",
        "    elif (not opt.is_pseudo) and opt.is_semi:\n",
        "        train_save = 'Semi-Inf-Net'\n",
        "    elif (not opt.is_pseudo) and (not opt.is_semi):\n",
        "        train_save = 'Inf-Net'\n",
        "    else:\n",
        "        print('Use custom save path')\n",
        "        train_save = opt.train_save\n",
        "\n",
        "    # ---- calculate FLOPs and Params ----\n",
        "    if opt.is_thop:\n",
        "        from Code.utils.utils import CalParams\n",
        "        x = torch.randn(1, 3, opt.trainsize, opt.trainsize).cuda()\n",
        "        CalParams(model, x)\n",
        "\n",
        "    # ---- load training sub-modules ----\n",
        "    BCE = torch.nn.BCEWithLogitsLoss()\n",
        "\n",
        "    params = model.parameters()\n",
        "    optimizer = torch.optim.Adam(params, opt.lr)\n",
        "\n",
        "    image_root = '{}/Imgs/'.format(opt.train_path)\n",
        "    gt_root = '{}/GT/'.format(opt.train_path)\n",
        "    edge_root = '{}/Edge/'.format(opt.train_path)\n",
        "\n",
        "    train_loader = get_loader(image_root, gt_root, edge_root,\n",
        "                              batchsize=opt.batchsize, trainsize=opt.trainsize, num_workers=opt.num_workers)\n",
        "    total_step = len(train_loader)\n",
        "\n",
        "    # ---- start !! -----\n",
        "    print(\"#\"*20, \"\\nStart Training (Inf-Net-{})\\n{}\\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\".format(opt.backbone, opt), \"#\"*20)\n",
        "\n",
        "    for epoch in range(1, opt.epoch):\n",
        "        adjust_lr(optimizer, opt.lr, epoch, opt.decay_rate, opt.decay_epoch)\n",
        "        train(train_loader, model, optimizer, epoch, train_save)\n"
      ],
      "metadata": {
        "id": "NVj5QcZwQev5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Testing  \"MyTest_LungInf.py\"   for Inf-Net & Semi-Inf-Net*\n",
        "\n",
        "*The following code is adjusted for Inf-Net. If you want to change it to Semi-Inf-Net, follow the [README.md](https://github.com/DengPingFan/Inf-Net/blob/master/README.md)*"
      ],
      "metadata": {
        "id": "UnzEsna6UHe3"
      }
    },
    {
      "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 torch\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "import os\n",
        "import argparse\n",
        "from scipy import misc\n",
        "from Code.model_lung_infection.InfNet_Res2Net import Inf_Net as Network\n",
        "from Code.utils.dataloader_LungInf import test_dataset\n",
        "\n",
        "\n",
        "def inference():\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument('--testsize', type=int, default=352, help='testing size')\n",
        "    parser.add_argument('--data_path', type=str, default='./Dataset/COVID-SemiSeg/Dataset/TestingSet/LungInfection-Test/',\n",
        "                        help='Path to test data')\n",
        "    parser.add_argument('--pth_path', type=str, default='./Snapshots/save_weights/Inf-Net/Inf-Net-100.pth',\n",
        "                        help='Path to weights file. If `semi-sup`, edit it to `Semi-Inf-Net/Semi-Inf-Net-100.pth`')\n",
        "    parser.add_argument('--save_path', type=str, default='./Results/Lung infection segmentation/Inf-Net/',\n",
        "                        help='Path to save the predictions. if `semi-sup`, edit it to `Semi-Inf-Net`')\n",
        "    opt = parser.parse_args()\n",
        "\n",
        "    print(\"#\" * 20, \"\\nStart Testing (Inf-Net)\\n{}\\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@gamil.com)\\n----\\n\".format(opt), \"#\" * 20)\n",
        "\n",
        "    model = Network()\n",
        "    # model = torch.nn.DataParallel(model, device_ids=[0, 1]) # uncomment it if you have multiply GPUs.\n",
        "    model.load_state_dict(torch.load(opt.pth_path, map_location={'cuda:1':'cuda:0'}))\n",
        "    model.cuda()\n",
        "    model.eval()\n",
        "\n",
        "    image_root = '{}/Imgs/'.format(opt.data_path)\n",
        "    # gt_root = '{}/GT/'.format(opt.data_path)\n",
        "    test_loader = test_dataset(image_root, opt.testsize)\n",
        "    os.makedirs(opt.save_path, exist_ok=True)\n",
        "\n",
        "    for i in range(test_loader.size):\n",
        "        image, name = test_loader.load_data()\n",
        "\n",
        "        image = image.cuda()\n",
        "\n",
        "        lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(image)\n",
        "\n",
        "        res = lateral_map_2\n",
        "        # res = F.upsample(res, size=(ori_size[1],ori_size[0]), mode='bilinear', align_corners=False)\n",
        "        res = res.sigmoid().data.cpu().numpy().squeeze()\n",
        "        res = (res - res.min()) / (res.max() - res.min() + 1e-8)\n",
        "        #misc.imsave(opt.save_path + name, res)\n",
        "        import imageio\n",
        "        imageio.imwrite(opt.save_path + name, res)\n",
        "\n",
        "    print('Test Done!')\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    inference()\n"
      ],
      "metadata": {
        "id": "8EUSmTJwUPI2"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}