main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Segmentation\n",
    "This notebook is the main file and it imports all the functions from other .py files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import argparse\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tensorboard toolbox\n",
    "from tensorboardX import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pytorch\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "import torch.optim as optim\n",
    "from torch.utils import data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# our own lib\n",
    "from lib.utils import *\n",
    "from tool.tool import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# which gpu\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some settings\n",
    "data_normalization = False\n",
    "\n",
    "# output directory\n",
    "root_dir = '/Users/weiwei/Desktop/' ##### change to your root dir\n",
    "checkpoint_dir = './checkpoint/'\n",
    "output_log_dir = './log/'\n",
    "tb_log_dir = './tb_log/'\n",
    "vis_output_path = './vis_output/'\n",
    "\n",
    "# prediction data number\n",
    "prediction_set = 'val' #['train', 'val', 'test', 'ext']\n",
    "subject_num = '1_20' \n",
    "patch_num = '88'\n",
    "case = 'whole' #['patch', 'whole']\n",
    "\n",
    "prediction_data = (prediction_set, subject_num, patch_num, case)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(args):\n",
    "\n",
    "    device = torch.device('cuda:0' if not args.no_cuda else 'cpu')\n",
    "\n",
    "    if args.mode == 'train':\n",
    "\n",
    "        ## experiment name\n",
    "        exp_name = args.model_arch + args.name\n",
    "\n",
    "        ## tensorboard log\n",
    "        tb_writer = SummaryWriter(tb_log_dir + exp_name)\n",
    "\n",
    "        ## data augmentation\n",
    "        # Todo :design which type of data aug\n",
    "        data_aug = None\n",
    "\n",
    "        ## load dataset\n",
    "        print('Load dataset...')\n",
    "        train_loader, val_loader = prepareDataset(args, root_dir, data_aug, normalize=data_normalization)\n",
    "        print(args.dataset.lower() + ' dataset loaded.')\n",
    "\n",
    "        ## load model\n",
    "        model = chooseModel(args)\n",
    "        model.to(device)\n",
    "        print(args.model_arch + ' loaded.')\n",
    "\n",
    "        ## parallel model\n",
    "        if args.gpu_num > 1:\n",
    "            model = nn.DataParallel(model)\n",
    "\n",
    "        ## loss function and optimizer\n",
    "        loss_fn = chooseLoss(args, 0)\t\n",
    "\n",
    "        optimizer = chooseOptimizer(model, args)\n",
    "\n",
    "        ## initailize statistic result\n",
    "        start_epoch = 0\n",
    "        best_per_index = 1000\n",
    "        total_tb_it = 0\n",
    "\n",
    "        ## resume training\n",
    "        # Todo\n",
    "\n",
    "        for epoch in range(start_epoch, args.num_epoch):\n",
    "\n",
    "            total_tb_it = train(args, device, model, train_loader, epoch, loss_fn, optimizer, tb_writer, total_tb_it)\n",
    "            mse_index = validate(device, model, val_loader, epoch, loss_fn, tb_writer)\n",
    "\n",
    "            state = {'epoch': epoch, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()}\n",
    "\n",
    "            if mse_index <= best_per_index:\n",
    "                best_per_index = mse_index\n",
    "                best_name = checkpoint_dir + exp_name +'_Bmodel.pkl'\n",
    "                torch.save(state, best_name)\n",
    "            else:\n",
    "                name = checkpoint_dir + exp_name +'_Emodel.pkl'\n",
    "                torch.save(state, name)\n",
    "\n",
    "        tb_writer.close()\n",
    "\n",
    "    elif args.mode == 'predict':\n",
    "\n",
    "        print('Load data...')\n",
    "        data_loader = loadData(args, root_dir, prediction_data, normalize=data_normalization)\n",
    "\n",
    "        ## load model\n",
    "        model = chooseModel(args)\n",
    "        model.load_state_dict(torch.load(args.resume_file, map_location=device)['model_state'])\n",
    "        model.to(device)\n",
    "        print(args.model_arch + ' loaded.')\n",
    "\n",
    "        ## parallel model\n",
    "        if args.gpu_num > 1:\n",
    "            model = nn.DataParallel(model)\n",
    "\n",
    "        model_name = args.resume_file.split('/')[-1].split('.')[0]\n",
    "        print(model_name)\n",
    "\n",
    "        ## loss function\n",
    "        loss_fn = chooseLoss(args, 0)\n",
    "\n",
    "        predict(args, device, model, data_loader, loss_fn, model_name)\n",
    "    else:\n",
    "        raise Exception('Unrecognized mode.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(args, device, model, train_loader, epoch, loss_fn, optimizer, tb_writer, total_tb_it):\n",
    "\n",
    "    print_freq = (len(train_loader.dataset) // args.batch_size) // 3\n",
    "\n",
    "    model.train()\n",
    "\n",
    "    for batch_count, (input_data, gt_data, input_name) in enumerate(train_loader):\n",
    "\n",
    "        #cuda\n",
    "        input_data = input_data.to(device, dtype=torch.float)\n",
    "        gt_data = gt_data.to(device, dtype=torch.float)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        output_data = model(input_data)\n",
    "\n",
    "        loss = loss_fn(output_data, gt_data)\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        per_loss = loss.item()\n",
    "\n",
    "        tb_writer.add_scalar('train/overall_loss', per_loss, total_tb_it)\n",
    "        total_tb_it += 1\n",
    "\n",
    "        if batch_count%print_freq == 0:\n",
    "            print('Epoch [%d/%d] Loss: %.8f' %(epoch, args.num_epoch, per_loss))\n",
    "\n",
    "    return total_tb_it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(device, model, val_loader, epoch, loss_fn, tb_writer):\n",
    "\n",
    "    model.eval()\n",
    "    \n",
    "    tb_loss = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "\n",
    "        for batch_count, (input_data, gt_data, input_name) in enumerate(val_loader):\n",
    "\n",
    "            #cuda\n",
    "            input_data = input_data.to(device, dtype=torch.float)\n",
    "            gt_data = gt_data.to(device, dtype=torch.float)\n",
    "\n",
    "            output_data = model(input_data)\n",
    "            \n",
    "            loss = loss_fn(output_data, gt_data)\n",
    "\n",
    "            tb_loss += loss.item()\n",
    "\n",
    "        avg_tb_loss = tb_loss / len(val_loader.dataset)\n",
    "        print('##Validate loss: %.8f' %(avg_tb_loss))\n",
    "\n",
    "        tb_writer.add_scalar('val/overall_loss', avg_tb_loss, epoch)\n",
    "\n",
    "    return avg_tb_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(args, device, model, data_loader, loss_fn, model_name):\n",
    "\n",
    "    nifti_path = root_dir + 'img_seg/dataset/original_data/'\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    if case == 'patch':\n",
    "        with torch.no_grad():\n",
    "            for batch_count, (input_data, gt_data, input_name) in enumerate(data_loader):\n",
    "\n",
    "                #cuda\n",
    "                input_data = input_data.to(device, dtype=torch.float)\n",
    "                gt_data = gt_data.to(device, dtype=torch.float)\n",
    "\n",
    "                output_data = model(input_data)\n",
    "\n",
    "                output_data = torch.squeeze(torch.sigmoid(output_data), 0).cpu().numpy()\n",
    "                if not args.no_save:\n",
    "\n",
    "                    data_path = nifti_path + 'image/' + input_name[0].split('_image')[0] + '/' +  input_name[0].split('_image')[0] + '.hdr'\n",
    "\n",
    "                    if not os.path.exists(vis_output_path + model_name):\n",
    "                        os.makedirs(vis_output_path + model_name)\n",
    "\n",
    "                    save_name = vis_output_path + model_name + '/' + input_name[0].split('_image')[0] + '_pred'\n",
    "\n",
    "                    seg_display(output_data, data_path, out_name=save_name)\n",
    "\n",
    "    elif case == 'whole':\n",
    "        with torch.no_grad():\n",
    "            for batch_count, (input_data, gt_data, input_name) in enumerate(data_loader):\n",
    "\n",
    "                #cuda\n",
    "                input_data = input_data.to(device, dtype=torch.float)\n",
    "                gt_data = gt_data.to(device, dtype=torch.float)\n",
    "\n",
    "                output_data = model(input_data)\n",
    "\n",
    "                #output_data = torch.squeeze(torch.sigmoid(gt_data), 0).cpu().numpy()\n",
    "                output_data = torch.squeeze(input_data, 0).cpu().numpy()\n",
    "\n",
    "                if not args.no_save:\n",
    "\n",
    "                    data_path = nifti_path + 'image/' + input_name[0].split('_image')[0] + '/' +  input_name[0].split('_image')[0] + '.hdr'\n",
    "\n",
    "                    if not os.path.exists(vis_output_path + model_name):\n",
    "                        os.makedirs(vis_output_path + model_name)\n",
    "\n",
    "                    save_name = vis_output_path + model_name + '/' + input_name[0].split('_image')[0] + '_pred'\n",
    "\n",
    "                    seg_display(output_data, data_path, out_name=save_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='Medical Image Segmentation')\n",
    "parser.add_argument('--mode', default='train', choices=['train', 'predict'], help='operation mode: train or predict (default: train)')\n",
    "parser.add_argument('--name', type=str, default='_test', help='the name of exp')\n",
    "parser.add_argument('--dataset', default='seg', choices=['seg'], help='dataset to use (default: seg)')\n",
    "parser.add_argument('--gpu_num', default=1, type=int, choices=[1, 2, 3, 4], help='number of gpu (default: 1)')\n",
    "parser.add_argument('--model_arch', default='vdsrr', choices=['unet', 'vdsrr'], help='network model (default: unet)')\n",
    "parser.add_argument('--num_epoch', default=200, type=int, metavar='N', help='number of total epochs to run (default: 1000)')\n",
    "parser.add_argument('--batch_size', type=int, default=20, help='batch size (default: 10)')\n",
    "parser.add_argument('--learning_rate', type=float, default=1e-3, help='learning rate (default: 1e-2)')\n",
    "parser.add_argument('--optimizer', default='adam', choices=['sgd', 'adam'], help='optimizer to use (default: adam)')\n",
    "parser.add_argument('--momentum', type=float, default=0.9, help='SGD momentum (default: 0.9)')\n",
    "parser.add_argument('--no_cuda', action='store_true', default=False, help='disables CUDA training')\n",
    "parser.add_argument('--no_save', action='store_true', default=False, help='disables saving tensors')\n",
    "parser.add_argument('--resume_file', type=str, default=None, help='the checkpoint file to resume from')\n",
    "\n",
    "\n",
    "#args = parser.parse_args()\n",
    "args, unknown = parser.parse_known_args()\n",
    "args.mode = 'train'\n",
    "args.name = '_test'\n",
    "args.dataset = 'seg'\n",
    "args.gpu_num = 1\n",
    "args.model_arch = 'vdsrr'\n",
    "args.num_epoch = 200\n",
    "args.batch_size = 20\n",
    "args.learning_rate = 0.001\n",
    "args.optimizer = 'adam'\n",
    "args.momentum = 0.9\n",
    "args.no_cuda = True\n",
    "args.no_save = False\n",
    "args.resume_file = None\n",
    "main(args)"
   ]
  },
  {
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}