brain_seg_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "brain-seg.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fOQKH5ig2Sph",
        "outputId": "cb1f48fa-06f5-4fb8-dab9-a403c7da045f"
      },
      "source": [
        "!git clone https://github.com/aqntks/brain-seg"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'brain-seg'...\n",
            "remote: Enumerating objects: 77, done.\u001b[K\n",
            "remote: Counting objects: 100% (77/77), done.\u001b[K\n",
            "remote: Compressing objects: 100% (54/54), done.\u001b[K\n",
            "remote: Total 77 (delta 32), reused 61 (delta 18), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (77/77), done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IMrEFWDc2pDO",
        "outputId": "2a0fa9f6-0700-416a-b5c6-8213daed2af5"
      },
      "source": [
        "%cd brain-seg\n",
        "!mkdir data"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "mkdir: cannot create directory ‘data’: File exists\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9X10rByz2v03",
        "outputId": "4e722a6a-abe1-44e6-ba4c-03dec96790a5"
      },
      "source": [
        "!pip install -r requirements.txt"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torch>=1.5 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 1)) (1.9.0+cu111)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 2)) (1.19.5)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 3)) (1.1.5)\n",
            "Requirement already satisfied: nibabel in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 4)) (3.0.2)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 5)) (3.2.2)\n",
            "Collecting monai\n",
            "  Downloading monai-0.7.0-202109240007-py3-none-any.whl (650 kB)\n",
            "\u001b[K     |████████████████████████████████| 650 kB 5.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm>=4.47.0 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 7)) (4.62.3)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.5->-r requirements.txt (line 1)) (3.7.4.3)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas->-r requirements.txt (line 3)) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->-r requirements.txt (line 3)) (2018.9)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas->-r requirements.txt (line 3)) (1.15.0)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->-r requirements.txt (line 5)) (2.4.7)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->-r requirements.txt (line 5)) (1.3.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->-r requirements.txt (line 5)) (0.10.0)\n",
            "Installing collected packages: monai\n",
            "Successfully installed monai-0.7.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BXp22lvK3TO3"
      },
      "source": [
        "import os\n",
        "import time\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "from collections import OrderedDict\n",
        "\n",
        "from monai.data import decollate_batch\n",
        "from monai.losses import DiceLoss\n",
        "from monai.inferers import sliding_window_inference\n",
        "from monai.metrics import DiceMetric\n",
        "from monai.networks.nets import SegResNet\n",
        "from monai.transforms import (\n",
        "    Activations,\n",
        "    AsDiscrete,\n",
        "    Compose,\n",
        "    EnsureType,\n",
        ")\n",
        "\n",
        "import data_loader\n",
        "from utils import pre_visualize, train_graph"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cHpWyUio5Mui"
      },
      "source": [
        "train_ds, val_ds, train_loader, val_loader = data_loader.brats_brain('data/')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QufiLZbj6daU"
      },
      "source": [
        "pre_visualize(val_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x_sNW62x6hNE"
      },
      "source": [
        "device = torch.device(\"cuda:0\")\n",
        "\n",
        "max_epochs = 300\n",
        "val_interval = 1\n",
        "VAL_AMP = True"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xfuYBu6J29eH"
      },
      "source": [
        "class UNet3d(nn.Module):\n",
        "    def __init__(self, in_channels=3, out_channels=1, init_features=32):\n",
        "        super(UNet3d, self).__init__()\n",
        "\n",
        "        features = init_features\n",
        "        self.encoder1 = UNet3d._block(in_channels, features, name=\"enc1\")  # 1/32/32\n",
        "        self.pool1 = nn.MaxPool3d((2, 2, 2))\n",
        "        self.encoder2 = UNet3d._block(features, features * 2, name=\"enc2\")  # 32/64/64\n",
        "        self.pool2 = nn.MaxPool3d((2, 2, 2))\n",
        "        self.encoder3 = UNet3d._block(features * 2, features * 4, name=\"enc3\")  # 64/128/128\n",
        "        self.pool3 = nn.MaxPool3d((2, 2, 2))\n",
        "        self.encoder4 = UNet3d._block(features * 4, features * 8, name=\"enc4\")  # 128/256/256\n",
        "        self.pool4 = nn.MaxPool3d((2, 2, 2))\n",
        "\n",
        "        self.bottleneck = UNet3d._block(features * 8, features * 16, name=\"bottleneck\")  # 256/512/512\n",
        "\n",
        "        self.upconv4 = nn.ConvTranspose3d(features * 16, features * 8, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=True)\n",
        "        self.decoder4 = UNet3d._block((features * 8) * 2, features * 8, name=\"dec4\")\n",
        "        self.upconv3 = nn.ConvTranspose3d(features * 8, features * 4, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=True)\n",
        "        self.decoder3 = UNet3d._block((features * 4) * 2, features * 4, name=\"dec3\")\n",
        "        self.upconv2 = nn.ConvTranspose3d(features * 4, features * 2, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=True)\n",
        "        self.decoder2 = UNet3d._block((features * 2) * 2, features * 2, name=\"dec2\")\n",
        "        self.upconv1 = nn.ConvTranspose3d(features * 2, features, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=True)\n",
        "        self.decoder1 = UNet3d._block(features * 2, features, name=\"dec1\")\n",
        "\n",
        "        self.conv = nn.Conv3d(in_channels=features, out_channels=out_channels, kernel_size=(1, 1, 1), padding=(0, 0, 0), bias=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        enc1 = self.encoder1(x)\n",
        "        enc2 = self.encoder2(self.pool1(enc1))\n",
        "        enc3 = self.encoder3(self.pool2(enc2))\n",
        "        enc4 = self.encoder4(self.pool3(enc3))\n",
        "\n",
        "        bottleneck = self.bottleneck(self.pool4(enc4))\n",
        "\n",
        "        dec4 = self.upconv4(bottleneck)\n",
        "        dec4 = torch.cat((dec4, enc4), dim=1)\n",
        "        dec4 = self.decoder4(dec4)\n",
        "        dec3 = self.upconv3(dec4)\n",
        "        dec3 = torch.cat((dec3, enc3), dim=1)\n",
        "        dec3 = self.decoder3(dec3)\n",
        "        dec2 = self.upconv2(dec3)\n",
        "        dec2 = torch.cat((dec2, enc2), dim=1)\n",
        "        dec2 = self.decoder2(dec2)\n",
        "        dec1 = self.upconv1(dec2)\n",
        "        dec1 = torch.cat((dec1, enc1), dim=1)\n",
        "        dec1 = self.decoder1(dec1)\n",
        "\n",
        "        return torch.sigmoid(self.conv(dec1))\n",
        "\n",
        "    @staticmethod\n",
        "    def _block(in_channels, features, name):\n",
        "        return nn.Sequential(\n",
        "            OrderedDict(\n",
        "                [\n",
        "                    (\n",
        "                        name + \"conv1\",\n",
        "                        nn.Conv3d(\n",
        "                            in_channels=in_channels,\n",
        "                            out_channels=features,\n",
        "                            kernel_size=(3, 3, 3),\n",
        "                            stride=(1, 1, 1),\n",
        "                            padding=(1, 1, 1),\n",
        "                            bias=True,\n",
        "                        ),\n",
        "                    ),\n",
        "                    (name + \"norm1\", nn.BatchNorm3d(num_features=features)),\n",
        "                    (name + \"relu1\", nn.ReLU(inplace=True)),\n",
        "                    (\n",
        "                        name + \"conv2\",\n",
        "                        nn.Conv3d(\n",
        "                            in_channels=features,\n",
        "                            out_channels=features,\n",
        "                            kernel_size=(3, 3, 3),\n",
        "                            stride=(1, 1, 1),\n",
        "                            padding=(1, 1, 1),\n",
        "                            bias=True,\n",
        "                        ),\n",
        "                    ),\n",
        "                    (name + \"norm2\", nn.BatchNorm3d(num_features=features)),\n",
        "                    (name + \"relu2\", nn.ReLU(inplace=True)),\n",
        "                ]\n",
        "            )\n",
        "        )"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NVI9Ekos6a-A"
      },
      "source": [
        "model = UNet3d(in_channels=4, out_channels=3).to(device)\n",
        "\n",
        "loss_function = DiceLoss(smooth_nr=0, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, sigmoid=True)\n",
        "optimizer = torch.optim.Adam(model.parameters(), 1e-4, weight_decay=1e-5)\n",
        "lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)\n",
        "\n",
        "dice_metric = DiceMetric(include_background=True, reduction=\"mean\")\n",
        "dice_metric_batch = DiceMetric(include_background=True, reduction=\"mean_batch\")\n",
        "\n",
        "post_trans = Compose(\n",
        "    [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)]\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T8HZVCRd3K7O"
      },
      "source": [
        "def train(train_ds, val_ds, train_loader, val_loader):\n",
        "    def inference(input):\n",
        "\n",
        "        def _compute(input):\n",
        "            return sliding_window_inference(\n",
        "                inputs=input,\n",
        "                roi_size=(240, 240, 160),\n",
        "                sw_batch_size=1,\n",
        "                predictor=model,\n",
        "                overlap=0.5,\n",
        "            )\n",
        "\n",
        "        if VAL_AMP:\n",
        "            with torch.cuda.amp.autocast():\n",
        "                return _compute(input)\n",
        "        else:\n",
        "            return _compute(input)\n",
        "\n",
        "    scaler = torch.cuda.amp.GradScaler()\n",
        "    torch.backends.cudnn.benchmark = True\n",
        "\n",
        "    # train\n",
        "    best_metric = -1\n",
        "    best_metric_epoch = -1\n",
        "    best_metrics_epochs_and_time = [[], [], []]\n",
        "    epoch_loss_values = []\n",
        "    metric_values = []\n",
        "    metric_values_tc = []\n",
        "    metric_values_wt = []\n",
        "    metric_values_et = []\n",
        "\n",
        "    total_start = time.time()\n",
        "    for epoch in range(max_epochs):\n",
        "        epoch_start = time.time()\n",
        "        print(\"-\" * 10)\n",
        "        print(f\"epoch {epoch + 1}/{max_epochs}\")\n",
        "        model.train()\n",
        "        epoch_loss = 0\n",
        "        step = 0\n",
        "        for batch_data in train_loader:\n",
        "            step_start = time.time()\n",
        "            step += 1\n",
        "            inputs, labels = (\n",
        "                batch_data[\"image\"].to(device),\n",
        "                batch_data[\"label\"].to(device),\n",
        "            )\n",
        "            optimizer.zero_grad()\n",
        "            with torch.cuda.amp.autocast():\n",
        "                outputs = model(inputs)\n",
        "                loss = loss_function(outputs, labels)\n",
        "            scaler.scale(loss).backward()\n",
        "            scaler.step(optimizer)\n",
        "            scaler.update()\n",
        "            epoch_loss += loss.item()\n",
        "            print(\n",
        "                f\"{step}/{len(train_ds) // train_loader.batch_size}\"\n",
        "                f\", train_loss: {loss.item():.4f}\"\n",
        "                f\", step time: {(time.time() - step_start):.4f}\"\n",
        "            )\n",
        "        lr_scheduler.step()\n",
        "        epoch_loss /= step\n",
        "        epoch_loss_values.append(epoch_loss)\n",
        "        print(f\"epoch {epoch + 1} average loss: {epoch_loss:.4f}\")\n",
        "\n",
        "        if (epoch + 1) % val_interval == 0:\n",
        "            model.eval()\n",
        "            with torch.no_grad():\n",
        "\n",
        "                for val_data in val_loader:\n",
        "                    val_inputs, val_labels = (\n",
        "                        val_data[\"image\"].to(device),\n",
        "                        val_data[\"label\"].to(device),\n",
        "                    )\n",
        "                    val_outputs = inference(val_inputs)\n",
        "                    val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]\n",
        "                    dice_metric(y_pred=val_outputs, y=val_labels)\n",
        "                    dice_metric_batch(y_pred=val_outputs, y=val_labels)\n",
        "\n",
        "                metric = dice_metric.aggregate().item()\n",
        "                metric_values.append(metric)\n",
        "                metric_batch = dice_metric_batch.aggregate()\n",
        "                metric_tc = metric_batch[0].item()\n",
        "                metric_values_tc.append(metric_tc)\n",
        "                metric_wt = metric_batch[1].item()\n",
        "                metric_values_wt.append(metric_wt)\n",
        "                metric_et = metric_batch[2].item()\n",
        "                metric_values_et.append(metric_et)\n",
        "                dice_metric.reset()\n",
        "                dice_metric_batch.reset()\n",
        "\n",
        "                if metric > best_metric:\n",
        "                    best_metric = metric\n",
        "                    best_metric_epoch = epoch + 1\n",
        "                    best_metrics_epochs_and_time[0].append(best_metric)\n",
        "                    best_metrics_epochs_and_time[1].append(best_metric_epoch)\n",
        "                    best_metrics_epochs_and_time[2].append(time.time() - total_start)\n",
        "                    torch.save(\n",
        "                        model.state_dict(),\n",
        "                        os.path.join(\"weights/\", \"best_metric_model.pth\"),\n",
        "                    )\n",
        "                    print(\"saved new best metric model\")\n",
        "                print(\n",
        "                    f\"current epoch: {epoch + 1} current mean dice: {metric:.4f}\"\n",
        "                    f\" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}\"\n",
        "                    f\"\\nbest mean dice: {best_metric:.4f}\"\n",
        "                    f\" at epoch: {best_metric_epoch}\"\n",
        "                )\n",
        "        print(f\"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}\")\n",
        "    total_time = time.time() - total_start\n",
        "\n",
        "    print(f\"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}, total time: {total_time}.\")\n",
        "\n",
        "    train_graph(epoch_loss_values, val_interval, metric_values, metric_values_tc, metric_values_wt, metric_values_et)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4LzhZ2r63uI2"
      },
      "source": [
        "train(train_ds, val_ds, train_loader, val_loader)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}