focal_phi_loss_mendeley_vanilla_unet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "focal_phi_loss_mendeley_vanilla_unet.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lCx9-Tf21c0_"
      },
      "source": [
        "## **Install the required packages and dependencies**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JkPXhdd31Zgk"
      },
      "source": [
        "#check cuda version \n",
        "!nvcc --version \n",
        "# install pytorch according to the cuda version from the pytorch website\n",
        "!pip install torch==1.7.0+cu101 torchvision==0.8.1+cu101 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "# install mmcv\n",
        "!pip install mmcv-full==1.1.4\n",
        "#clone github repository\n",
        "!git clone https://github.com/rubeea/focal_phi_loss_mmsegmentation.git \n",
        "#install all requirements\n",
        "!pip install -qr focal_phi_loss_mmsegmentation/requirements.txt \n",
        "%cd /content/focal_phi_loss_mmsegmentation\n",
        "#install mmsegmentation in develop mode\n",
        "!python setup.py develop "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZhbqjmm8IIw"
      },
      "source": [
        "# **Make data directories and upload the data in them**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sn1Qu9g0kJap"
      },
      "source": [
        "#declare global variables\r\n",
        "root_dir= '/content/focal_phi_loss_mmsegmentation/'\r\n",
        "data_root= \"/content/focal_phi_loss_mmsegmentation/data/\"\r\n",
        "dataset= \"mendeley\""
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aQcS320P22zn"
      },
      "source": [
        "import os\n",
        "\n",
        "#script to extract data from zip files and create the directory structure for mendeley\n",
        "!python /content/focal_phi_loss_mmsegmentation/tools/convert_datasets/mendeley.py /content/focal_phi_loss_mmsegmentation/train_imgs.zip /content/focal_phi_loss_mmsegmentation/train_gt.zip /content/focal_phi_loss_mmsegmentation/val_imgs.zip /content/focal_phi_loss_mmsegmentation/val_gt.zip \n",
        "\n",
        "#check file counts\n",
        "train_imgs= data_root + dataset+'/img_dir/train'\n",
        "val_imgs= data_root + dataset +'/img_dir/val'\n",
        "train_gt= data_root + dataset+'/ann_dir/train'\n",
        "val_gt= data_root + dataset+'/ann_dir/val'\n",
        "\n",
        "#print the file count\n",
        "print(len(os.listdir(train_imgs)))\n",
        "print(len(os.listdir(val_imgs)))\n",
        "print(len(os.listdir(train_gt)))\n",
        "print(len(os.listdir(val_gt)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CL9RLxEPabM_"
      },
      "source": [
        "# **Visualize training and mask images**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZayX5trmabNC"
      },
      "source": [
        "import mmcv\n",
        "import matplotlib.pyplot as plt\n",
        "from google.colab.patches import cv2_imshow\n",
        "import numpy as np\n",
        "import os\n",
        "import glob\n",
        "\n",
        "\n",
        "PATH_TO_TRAIN_IMAGES_DIR= data_root+dataset+\"/img_dir/train\"\n",
        "PATH_TO_LABEL_IMAGES_DIR= data_root+dataset+\"/ann_dir/train\"\n",
        "TRAIN_IMAGE_PATHS = glob.glob(os.path.join(PATH_TO_TRAIN_IMAGES_DIR, \"*.jpg\"))\n",
        "LABEL_PATHS= glob.glob(os.path.join(PATH_TO_LABEL_IMAGES_DIR, \"*.png\"))\n",
        "\n",
        "TRAIN_IMAGE_PATHS.sort()\n",
        "LABEL_PATHS.sort()\n",
        "\n",
        "palette = [[120, 120, 120], [6, 230, 230]] #dataset palette\n",
        "\n",
        "display_num = 5\n",
        "r_choices = np.random.choice(len(TRAIN_IMAGE_PATHS), display_num)\n",
        "\n",
        "#visualize any 5 random images and their mask images\n",
        "plt.figure(figsize=(10, 15))\n",
        "for i in range(0, display_num * 2, 2):\n",
        "  img_num = r_choices[i // 2]\n",
        "  img = mmcv.imread(TRAIN_IMAGE_PATHS[img_num])\n",
        "  label = mmcv.imread(LABEL_PATHS[img_num])\n",
        "\n",
        "  plt.subplot(display_num, 2, i + 1)\n",
        "  plt.imshow(img)\n",
        "  plt.title(\"Input image\")\n",
        "\n",
        "  plt.subplot(display_num, 2, i + 2)\n",
        "  plt.imshow(label)\n",
        "  plt.title(\"Ground Truth\")\n",
        "\n",
        "\n",
        "plt.suptitle(\"Examples of Images and their Masks\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yUVtmn3Iq3WA"
      },
      "source": [
        "#**Create a config file**\n",
        "In the next step, we need to modify the config for the training. To accelerate the process, we finetune the model from trained weights."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wwnj9tRzqX_A"
      },
      "source": [
        "#config for ACU-Net\n",
        "from mmcv import Config\n",
        "cfg = Config.fromfile('/content/focal_phi_loss_mmsegmentation/configs/unet/fcn_vanilla_unet_s5-d16_256x256_40k_hrf.py')"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1y2oV5w97jQo"
      },
      "source": [
        "Since the given config is used to train fcn vanilla U-Net on hrf dataset, we need to modify it accordingly for our new dataset.  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eyKnYC1Z7iCV"
      },
      "source": [
        "from mmseg.apis import set_random_seed\n",
        "import torch\n",
        "\n",
        "cfg_data_root = data_root+dataset\n",
        "train_img_dir = 'img_dir/train'\n",
        "train_ann_dir = 'ann_dir/train'\n",
        "val_img_dir = 'img_dir/val'\n",
        "val_ann_dir = 'ann_dir/val'\n",
        "\n",
        "# Since we use ony one GPU, BN is used instead of SyncBN\n",
        "cfg.norm_cfg = dict(type='BN', requires_grad=True)\n",
        "cfg.model.backbone.norm_cfg = cfg.norm_cfg\n",
        "\n",
        "cfg.model.decode_head.norm_cfg = cfg.norm_cfg\n",
        "\n",
        "# modify num classes of the model in decode head\n",
        "cfg.model.decode_head.num_classes = 2\n",
        "\n",
        "# Modify dataset type and path\n",
        "cfg.dataset_type = 'MendeleyDataset'\n",
        "cfg.data_root = cfg_data_root\n",
        "\n",
        "#batch size\n",
        "cfg.data.samples_per_gpu = 2\n",
        "cfg.data.workers_per_gpu=2\n",
        "\n",
        "\n",
        "#Balanced CE loss\n",
        "#assign class weights to tackle the imbalanced nature of the dataset in balanced CE loss\n",
        "# weights = [0.02, 0.98]\n",
        "# cfg.model.decode_head.loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, \n",
        "#                                             loss_weight=1.0, class_weight= weights)\n",
        "\n",
        "#Dice Loss\n",
        "# cfg.model.decode_head.loss_decode=dict(type='TverskyLoss',  alpha=0.5, beta=0.5, gamma=1.0, use_focal=False, loss_weight=1.0)\n",
        "\n",
        "#Tversky Loss\n",
        "# cfg.model.decode_head.loss_decode=dict(type='TverskyLoss',  alpha=0.3, beta=0.7, gamma=1.0, use_focal=False, loss_weight=1.0)\n",
        "\n",
        "#Focal Tversky Loss\n",
        "# cfg.model.decode_head.loss_decode=dict(type='TverskyLoss',  alpha=0.3, beta=0.7, gamma=0.75, use_focal=True, loss_weight=1.0)\n",
        "\n",
        "#MCC loss\n",
        "# cfg.model.decode_head.loss_decode=dict(type='PhiLoss', loss_weight=1.0, gamma=1.0)\n",
        "\n",
        "#Focal Phi Loss\n",
        "cfg.model.decode_head.loss_decode=dict(type='PhiLoss', loss_weight=1.0, gamma=0.5)\n",
        "\n",
        "#normalizing the dataset\n",
        "cfg.img_norm_cfg = dict(\n",
        "    mean= [142.393, 137.978, 83.073], std= [23.228, 20.046, 21.623], to_rgb=True) #for mendeley dataset\n",
        "\n",
        "cfg.crop_size = (256, 256)\n",
        "\n",
        "#dataset config\n",
        "cfg.train_pipeline = [\n",
        "    dict(type='LoadImageFromFile'),\n",
        "    dict(type='LoadAnnotations'),\n",
        "    dict(type='Resize', img_scale=(512, 512), ratio_range=(0.5, 1.5)), #for mendeley dataset\n",
        "    dict(type='RandomRotate', prob=0.5, degree=(45.0,315.0),pad_val=0,seg_pad_val=255), #for pldu and mendeley dataset\n",
        "    dict(type='RandomCrop', crop_size=cfg.crop_size, cat_max_ratio=0.75),\n",
        "    dict(type='RandomFlip', flip_ratio=0.5),\n",
        "    dict(type='PhotoMetricDistortion'),\n",
        "    dict(type='Normalize', **cfg.img_norm_cfg),\n",
        "    dict(type='Pad', size=cfg.crop_size, pad_val=0, seg_pad_val=255),\n",
        "    dict(type='DefaultFormatBundle'),\n",
        "    dict(type='Collect', keys=['img', 'gt_semantic_seg']),\n",
        "]\n",
        "\n",
        "cfg.test_pipeline = [\n",
        "    dict(type='LoadImageFromFile'),\n",
        "    dict(\n",
        "        type='MultiScaleFlipAug',\n",
        "        img_scale= (512, 512), #for mendeley dataset\n",
        "        flip=False,\n",
        "        transforms=[\n",
        "            dict(type='Resize', keep_ratio=True),\n",
        "            dict(type='RandomFlip'),\n",
        "            dict(type='Normalize', **cfg.img_norm_cfg),\n",
        "            dict(type='ImageToTensor', keys=['img']),\n",
        "            dict(type='Collect', keys=['img']),\n",
        "        ])\n",
        "]\n",
        "\n",
        "\n",
        "cfg.data.train.type = cfg.dataset_type\n",
        "cfg.data.train.data_root = cfg.data_root\n",
        "cfg.data.train.img_dir = train_img_dir\n",
        "cfg.data.train.ann_dir = train_ann_dir\n",
        "cfg.data.train.pipeline = cfg.train_pipeline\n",
        "cfg.data.train.split = None\n",
        "\n",
        "cfg.data.val.type = cfg.dataset_type\n",
        "cfg.data.val.data_root = cfg.data_root\n",
        "cfg.data.val.img_dir = val_img_dir\n",
        "cfg.data.val.ann_dir = val_ann_dir\n",
        "cfg.data.val.pipeline = cfg.test_pipeline\n",
        "cfg.data.val.split = None\n",
        "\n",
        "cfg.data.test.type = cfg.dataset_type\n",
        "cfg.data.test.data_root = cfg.data_root\n",
        "cfg.data.test.img_dir = val_img_dir\n",
        "cfg.data.test.ann_dir = val_ann_dir\n",
        "cfg.data.test.pipeline = cfg.test_pipeline\n",
        "cfg.data.test.split = None\n",
        "\n",
        "\n",
        "# Set up working dir to save files and logs.\n",
        "cfg.work_dir = root_dir+'/work_dirs/tutorial'\n",
        "\n",
        "cfg.total_iters = 160\n",
        "cfg.log_config.interval = 50\n",
        "cfg.evaluation.interval = 200 #validation at every 200 iterations\n",
        "cfg.checkpoint_config.interval = 200\n",
        "\n",
        "# Set seed to facilitate reproducing the result\n",
        "cfg.seed = 0\n",
        "set_random_seed(0, deterministic=False)\n",
        "cfg.gpu_ids = range(1)\n",
        "\n",
        "#for Dice Loss, Tversky Loss and Focal Tversky Loss on Mendeley dataset use the following learning rates with Adam:\n",
        "# lr = 5e-5\n",
        "\n",
        "cfg.optimizer = dict(type='Adam', lr=1e-3, weight_decay=0.0001,\n",
        "                     paramwise_cfg = dict(\n",
        "                        custom_keys={\n",
        "                            'head': dict(lr_mult=10.)\n",
        "                        }\n",
        "                        ))\n",
        "\n",
        "\n",
        "# Let's have a look at the final config used for training\n",
        "print(f'Config:\\n{cfg.pretty_text}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YXvx9Ll9EvNS"
      },
      "source": [
        "# **Train the model according to the config**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jYKoSfdMF12B"
      },
      "source": [
        "from mmseg.datasets import build_dataset\n",
        "from mmseg.models import build_segmentor\n",
        "from mmseg.apis import train_segmentor\n",
        "from mmseg.utils import collect_env, get_root_logger\n",
        "import mmcv\n",
        "import os.path as osp\n",
        "\n",
        "# Build the dataset\n",
        "datasets = [build_dataset(cfg.data.train)]\n",
        "\n",
        "\n",
        "meta = dict()\n",
        "    # log env info\n",
        "env_info_dict = collect_env()\n",
        "env_info = '\\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])\n",
        "meta['env_info'] = env_info\n",
        "\n",
        "meta['seed'] = cfg.seed\n",
        "meta['exp_name'] = osp.basename(cfg.filename)\n",
        "cfg.checkpoint_config.meta = dict(\n",
        "            CLASSES=datasets[0].CLASSES,\n",
        "            PALETTE=datasets[0].PALETTE)\n",
        "\n",
        "# Build the detector\n",
        "model = build_segmentor(\n",
        "    cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)\n",
        "# Add an attribute for visualization convenience\n",
        "model.CLASSES = datasets[0].CLASSES\n",
        "\n",
        "# Create work_dir\n",
        "mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))\n",
        "train_segmentor(model, datasets, cfg, distributed=False, validate=True, \n",
        "                meta=meta)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFKytoqKGp92"
      },
      "source": [
        "# **Inference with Trained Model**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KwPcnvRXeiVo"
      },
      "source": [
        "#create direcrtory for storing results\r\n",
        "!mkdir /content/focal_phi_loss_mmsegmentation/work_dirs/results\r\n",
        "!mkdir /content/focal_phi_loss_mmsegmentation/work_dirs/results/masks\r\n",
        "!mkdir /content/focal_phi_loss_mmsegmentation/work_dirs/results/overlays"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P1lhqp-aIfln"
      },
      "source": [
        "from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot\n",
        "import mmcv\n",
        "import matplotlib.pyplot as plt\n",
        "from google.colab.patches import cv2_imshow\n",
        "import numpy as np\n",
        "import os\n",
        "import glob\n",
        "import zipfile\n",
        "\n",
        "PATH_TO_TEST_IMAGES_DIR= data_root+dataset+\"/img_dir/val\"\n",
        "PATH_TO_LABEL_IMAGES_DIR= data_root+dataset+\"/ann_dir/val\"\n",
        "TEST_IMAGE_PATHS = glob.glob(os.path.join(PATH_TO_TEST_IMAGES_DIR, \"*.jpg\"))\n",
        "LABEL_PATHS= glob.glob(os.path.join(PATH_TO_LABEL_IMAGES_DIR, \"*.png\"))\n",
        "\n",
        "TEST_IMAGE_PATHS.sort()\n",
        "LABEL_PATHS.sort()\n",
        "\n",
        "checkpoint= root_dir+'/work_dirs/tutorial/iter_19200.pth' #checkpoint path\n",
        "model = init_segmentor(cfg, checkpoint, device='cuda:0')\n",
        "palette = [[120, 120, 120], [6, 230, 230]] #dataset palette\n",
        "print(len(TEST_IMAGE_PATHS))\n",
        "\n",
        "for i in range(0, len(TEST_IMAGE_PATHS)):\n",
        "  img = mmcv.imread(TEST_IMAGE_PATHS[i])\n",
        "  label = mmcv.imread(LABEL_PATHS[i])\n",
        "  result = inference_segmentor(model, img)\n",
        "  img_name= LABEL_PATHS[i].split(\"/\")[7] #FOR MASK\n",
        "  # img_name= TEST_IMAGE_PATHS[i].split(\"/\")[7] #for overlay\n",
        "\n",
        "  #prediction mask\n",
        "  arr= np.array(result)\n",
        "  arr=np.squeeze(arr)\n",
        "\n",
        "  #overlay image\n",
        "  # overlay = model.show_result(img, result, palette=palette, show=False)\n",
        "\n",
        "  # plt.subplot(1, 7 , 1)\n",
        "  # plt.imshow(img)\n",
        "  # plt.title(\"Input image\")\n",
        "\n",
        "  # plt.subplot(1, 7,  3)\n",
        "  # plt.imshow(label)\n",
        "  # plt.title(\"Ground Truth\")\n",
        "\n",
        "  # plt.subplot(1, 7,  5)\n",
        "  plt.imshow(arr)\n",
        "  plt.title(\"Prediction\")\n",
        "\n",
        "  # plt.subplot(1, 7,  7)\n",
        "  # plt.imshow(mmcv.bgr2rgb(overlay))\n",
        "  # plt.title(\"Overlay\")\n",
        "\n",
        "  # show_result_pyplot(model, img, result,palette)\n",
        "\n",
        "  image_file_name= root_dir+\"/work_dirs/results/masks/\"+img_name\n",
        "  plt.savefig(image_file_name)\n",
        "  print(\"saving \"+img_name)\n",
        "  i= i+1\n",
        "\n",
        "# plt.figure(figsize=(8, 6))\n",
        "# show_result_pyplot(model, img, result,palette)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PCvbTJmBdNyB"
      },
      "source": [
        "# **Downloading Inference Results**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4nOek-cqitGi"
      },
      "source": [
        "from google.colab import files\n",
        "#mask results\n",
        "!zip -r /content/focal_phi_loss_mmsegmentation/work_dirs/results/masks /content/focal_phi_loss_mmsegmentation/work_dirs/results/masks\n",
        "files.download(\"/content/focal_phi_loss_mmsegmentation/work_dirs/results/masks.zip\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "80ZSKSCenKFS"
      },
      "source": [
        "#overlay results\n",
        "!zip -r /content/focal_phi_loss_mmsegmentation/work_dirs/overlays /content/focal_phi_loss_mmsegmentation/work_dirs/overlays\n",
        "files.download(\"/content/focal_phi_loss_mmsegmentation/work_dirs/overlays.zip\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}