source: https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/

segmentation_brain_MRI_using_PyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# type of problem? - instance segmentation problem\n",
    "# input data? - tiff format image data (2D - height, width)\n",
    "# output data? - 2D\n",
    "# Goal - An accurate model to predict brain tumors from a new FLAIR image.\n",
    "\n",
    "# This segmentation is using FLAIR MRI .tiff images from TCIA(not .nii)\n",
    "# These files are easy compared to .nii files at this point.\n",
    "# One interesting thing would be to check whether fastai can beat this DICE score. \n",
    "\n",
    "# sources:\n",
    "## https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/\n",
    "## https://www.kaggle.com/gbstanford/brain-segmentation-pytorch/edit\n",
    "## https://github.com/gireeshkbogu/brain-segmentation-pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "\n",
    "from collections import OrderedDict\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.backends.backend_agg import FigureCanvasAgg\n",
    "from tqdm import tqdm\n",
    "from skimage.exposure import rescale_intensity\n",
    "from skimage.io import imread, imsave\n",
    "from skimage.transform import resize, rescale, rotate\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision.transforms import Compose\n",
    "\n",
    "from torchvision import models # this is missing in original code and throws UNet related error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_sample(x):\n",
    "    volume, mask = x\n",
    "    volume[volume < np.max(volume) * 0.1] = 0\n",
    "    z_projection = np.max(np.max(np.max(volume, axis=-1), axis=-1), axis=-1)\n",
    "    z_nonzero = np.nonzero(z_projection)\n",
    "    z_min = np.min(z_nonzero)\n",
    "    z_max = np.max(z_nonzero) + 1\n",
    "    y_projection = np.max(np.max(np.max(volume, axis=0), axis=-1), axis=-1)\n",
    "    y_nonzero = np.nonzero(y_projection)\n",
    "    y_min = np.min(y_nonzero)\n",
    "    y_max = np.max(y_nonzero) + 1\n",
    "    x_projection = np.max(np.max(np.max(volume, axis=0), axis=0), axis=-1)\n",
    "    x_nonzero = np.nonzero(x_projection)\n",
    "    x_min = np.min(x_nonzero)\n",
    "    x_max = np.max(x_nonzero) + 1\n",
    "    return (\n",
    "        volume[z_min:z_max, y_min:y_max, x_min:x_max],\n",
    "        mask[z_min:z_max, y_min:y_max, x_min:x_max],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sample(x):\n",
    "    volume, mask = x\n",
    "    a = volume.shape[1]\n",
    "    b = volume.shape[2]\n",
    "    if a == b:\n",
    "        return volume, mask\n",
    "    diff = (max(a, b) - min(a, b)) / 2.0\n",
    "    if a > b:\n",
    "        padding = ((0, 0), (0, 0), (int(np.floor(diff)), int(np.ceil(diff))))\n",
    "    else:\n",
    "        padding = ((0, 0), (int(np.floor(diff)), int(np.ceil(diff))), (0, 0))\n",
    "    mask = np.pad(mask, padding, mode=\"constant\", constant_values=0)\n",
    "    padding = padding + ((0, 0),)\n",
    "    volume = np.pad(volume, padding, mode=\"constant\", constant_values=0)\n",
    "    return volume, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def resize_sample(x, size=256):\n",
    "    volume, mask = x\n",
    "    v_shape = volume.shape\n",
    "    out_shape = (v_shape[0], size, size)\n",
    "    mask = resize(\n",
    "        mask,\n",
    "        output_shape=out_shape,\n",
    "        order=0,\n",
    "        mode=\"constant\",\n",
    "        cval=0,\n",
    "        anti_aliasing=False,\n",
    "    )\n",
    "    out_shape = out_shape + (v_shape[3],)\n",
    "    volume = resize(\n",
    "        volume,\n",
    "        output_shape=out_shape,\n",
    "        order=2,\n",
    "        mode=\"constant\",\n",
    "        cval=0,\n",
    "        anti_aliasing=False,\n",
    "    )\n",
    "    return volume, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_volume(volume):\n",
    "    p10 = np.percentile(volume, 10)\n",
    "    p99 = np.percentile(volume, 99)\n",
    "    volume = rescale_intensity(volume, in_range=(p10, p99))\n",
    "    m = np.mean(volume, axis=(0, 1, 2))\n",
    "    s = np.std(volume, axis=(0, 1, 2))\n",
    "    volume = (volume - m) / s\n",
    "    return volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BrainSegmentationDataset(Dataset):\n",
    "    \"\"\"Brain MRI dataset for FLAIR abnormality segmentation\"\"\"\n",
    "\n",
    "    in_channels = 3\n",
    "    out_channels = 1\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        images_dir,\n",
    "        transform=None,\n",
    "        image_size=256,\n",
    "        subset=\"train\",\n",
    "        random_sampling=True,\n",
    "        seed=42,\n",
    "    ):\n",
    "        assert subset in [\"all\", \"train\", \"validation\"]\n",
    "\n",
    "        # read images\n",
    "        volumes = {}\n",
    "        masks = {}\n",
    "        print(\"reading {} images...\".format(subset))\n",
    "        for (dirpath, dirnames, filenames) in os.walk(images_dir):\n",
    "            image_slices = []\n",
    "            mask_slices = []\n",
    "            for filename in sorted(\n",
    "                filter(lambda f: \".tif\" in f, filenames),\n",
    "                key=lambda x: int(x.split(\".\")[-2].split(\"_\")[4]),\n",
    "            ):\n",
    "                filepath = os.path.join(dirpath, filename)\n",
    "                if \"mask\" in filename:\n",
    "                    mask_slices.append(imread(filepath, as_gray=True))\n",
    "                else:\n",
    "                    image_slices.append(imread(filepath))\n",
    "            if len(image_slices) > 0:\n",
    "                patient_id = dirpath.split(\"/\")[-1]\n",
    "                volumes[patient_id] = np.array(image_slices[1:-1])\n",
    "                masks[patient_id] = np.array(mask_slices[1:-1])\n",
    "\n",
    "        self.patients = sorted(volumes)\n",
    "\n",
    "        # select cases to subset\n",
    "        if not subset == \"all\":\n",
    "            random.seed(seed)\n",
    "            validation_patients = random.sample(self.patients, k=10)\n",
    "            if subset == \"validation\":\n",
    "                self.patients = validation_patients\n",
    "            else:\n",
    "                self.patients = sorted(\n",
    "                    list(set(self.patients).difference(validation_patients))\n",
    "                )\n",
    "\n",
    "        print(\"preprocessing {} volumes...\".format(subset))\n",
    "        # create list of tuples (volume, mask)\n",
    "        self.volumes = [(volumes[k], masks[k]) for k in self.patients]\n",
    "\n",
    "        print(\"cropping {} volumes...\".format(subset))\n",
    "        # crop to smallest enclosing volume\n",
    "        self.volumes = [crop_sample(v) for v in self.volumes]\n",
    "\n",
    "        print(\"padding {} volumes...\".format(subset))\n",
    "        # pad to square\n",
    "        self.volumes = [pad_sample(v) for v in self.volumes]\n",
    "\n",
    "        print(\"resizing {} volumes...\".format(subset))\n",
    "        # resize\n",
    "        self.volumes = [resize_sample(v, size=image_size) for v in self.volumes]\n",
    "\n",
    "        print(\"normalizing {} volumes...\".format(subset))\n",
    "        # normalize channel-wise\n",
    "        self.volumes = [(normalize_volume(v), m) for v, m in self.volumes]\n",
    "\n",
    "        # probabilities for sampling slices based on masks\n",
    "        self.slice_weights = [m.sum(axis=-1).sum(axis=-1) for v, m in self.volumes]\n",
    "        self.slice_weights = [\n",
    "            (s + (s.sum() * 0.1 / len(s))) / (s.sum() * 1.1) for s in self.slice_weights\n",
    "        ]\n",
    "\n",
    "        # add channel dimension to masks\n",
    "        self.volumes = [(v, m[..., np.newaxis]) for (v, m) in self.volumes]\n",
    "\n",
    "        print(\"done creating {} dataset\".format(subset))\n",
    "\n",
    "        # create global index for patient and slice (idx -> (p_idx, s_idx))\n",
    "        num_slices = [v.shape[0] for v, m in self.volumes]\n",
    "        self.patient_slice_index = list(\n",
    "            zip(\n",
    "                sum([[i] * num_slices[i] for i in range(len(num_slices))], []),\n",
    "                sum([list(range(x)) for x in num_slices], []),\n",
    "            )\n",
    "        )\n",
    "\n",
    "        self.random_sampling = random_sampling\n",
    "\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.patient_slice_index)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        patient = self.patient_slice_index[idx][0]\n",
    "        slice_n = self.patient_slice_index[idx][1]\n",
    "\n",
    "        if self.random_sampling:\n",
    "            patient = np.random.randint(len(self.volumes))\n",
    "            slice_n = np.random.choice(\n",
    "                range(self.volumes[patient][0].shape[0]), p=self.slice_weights[patient]\n",
    "            )\n",
    "\n",
    "        v, m = self.volumes[patient]\n",
    "        image = v[slice_n]\n",
    "        mask = m[slice_n]\n",
    "\n",
    "        if self.transform is not None:\n",
    "            image, mask = self.transform((image, mask))\n",
    "\n",
    "        # fix dimensions (C, H, W)\n",
    "        image = image.transpose(2, 0, 1)\n",
    "        mask = mask.transpose(2, 0, 1)\n",
    "\n",
    "        image_tensor = torch.from_numpy(image.astype(np.float32))\n",
    "        mask_tensor = torch.from_numpy(mask.astype(np.float32))\n",
    "\n",
    "        # return tensors\n",
    "        return image_tensor, mask_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transforms(scale=None, angle=None, flip_prob=None):\n",
    "    transform_list = []\n",
    "\n",
    "    if scale is not None:\n",
    "        transform_list.append(Scale(scale))\n",
    "    if angle is not None:\n",
    "        transform_list.append(Rotate(angle))\n",
    "    if flip_prob is not None:\n",
    "        transform_list.append(HorizontalFlip(flip_prob))\n",
    "\n",
    "    return Compose(transform_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Scale(object):\n",
    "\n",
    "    def __init__(self, scale):\n",
    "        self.scale = scale\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image, mask = sample\n",
    "\n",
    "        img_size = image.shape[0]\n",
    "\n",
    "        scale = np.random.uniform(low=1.0 - self.scale, high=1.0 + self.scale)\n",
    "\n",
    "        image = rescale(\n",
    "            image,\n",
    "            (scale, scale),\n",
    "            multichannel=True,\n",
    "            preserve_range=True,\n",
    "            mode=\"constant\",\n",
    "            anti_aliasing=False,\n",
    "        )\n",
    "        mask = rescale(\n",
    "            mask,\n",
    "            (scale, scale),\n",
    "            order=0,\n",
    "            multichannel=True,\n",
    "            preserve_range=True,\n",
    "            mode=\"constant\",\n",
    "            anti_aliasing=False,\n",
    "        )\n",
    "\n",
    "        if scale < 1.0:\n",
    "            diff = (img_size - image.shape[0]) / 2.0\n",
    "            padding = ((int(np.floor(diff)), int(np.ceil(diff))),) * 2 + ((0, 0),)\n",
    "            image = np.pad(image, padding, mode=\"constant\", constant_values=0)\n",
    "            mask = np.pad(mask, padding, mode=\"constant\", constant_values=0)\n",
    "        else:\n",
    "            x_min = (image.shape[0] - img_size) // 2\n",
    "            x_max = x_min + img_size\n",
    "            image = image[x_min:x_max, x_min:x_max, ...]\n",
    "            mask = mask[x_min:x_max, x_min:x_max, ...]\n",
    "\n",
    "        return image, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Rotate(object):\n",
    "\n",
    "    def __init__(self, angle):\n",
    "        self.angle = angle\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image, mask = sample\n",
    "\n",
    "        angle = np.random.uniform(low=-self.angle, high=self.angle)\n",
    "        image = rotate(image, angle, resize=False, preserve_range=True, mode=\"constant\")\n",
    "        mask = rotate(\n",
    "            mask, angle, resize=False, order=0, preserve_range=True, mode=\"constant\"\n",
    "        )\n",
    "        return image, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HorizontalFlip(object):\n",
    "\n",
    "    def __init__(self, flip_prob):\n",
    "        self.flip_prob = flip_prob\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        image, mask = sample\n",
    "\n",
    "        if np.random.rand() > self.flip_prob:\n",
    "            return image, mask\n",
    "\n",
    "        image = np.fliplr(image).copy()\n",
    "        mask = np.fliplr(mask).copy()\n",
    "\n",
    "        return image, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "\n",
    "    def __init__(self, in_channels=3, out_channels=1, init_features=32):\n",
    "        super(UNet, self).__init__()\n",
    "\n",
    "        features = init_features\n",
    "        self.encoder1 = UNet._block(in_channels, features, name=\"enc1\")\n",
    "        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.encoder2 = UNet._block(features, features * 2, name=\"enc2\")\n",
    "        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.encoder3 = UNet._block(features * 2, features * 4, name=\"enc3\")\n",
    "        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.encoder4 = UNet._block(features * 4, features * 8, name=\"enc4\")\n",
    "        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        self.bottleneck = UNet._block(features * 8, features * 16, name=\"bottleneck\")\n",
    "\n",
    "        self.upconv4 = nn.ConvTranspose2d(\n",
    "            features * 16, features * 8, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder4 = UNet._block((features * 8) * 2, features * 8, name=\"dec4\")\n",
    "        self.upconv3 = nn.ConvTranspose2d(\n",
    "            features * 8, features * 4, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder3 = UNet._block((features * 4) * 2, features * 4, name=\"dec3\")\n",
    "        self.upconv2 = nn.ConvTranspose2d(\n",
    "            features * 4, features * 2, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder2 = UNet._block((features * 2) * 2, features * 2, name=\"dec2\")\n",
    "        self.upconv1 = nn.ConvTranspose2d(\n",
    "            features * 2, features, kernel_size=2, stride=2\n",
    "        )\n",
    "        self.decoder1 = UNet._block(features * 2, features, name=\"dec1\")\n",
    "\n",
    "        self.conv = nn.Conv2d(\n",
    "            in_channels=features, out_channels=out_channels, kernel_size=1\n",
    "        )\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",
    "        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.Conv2d(\n",
    "                            in_channels=in_channels,\n",
    "                            out_channels=features,\n",
    "                            kernel_size=3,\n",
    "                            padding=1,\n",
    "                            bias=False,\n",
    "                        ),\n",
    "                    ),\n",
    "                    (name + \"norm1\", nn.BatchNorm2d(num_features=features)),\n",
    "                    (name + \"relu1\", nn.ReLU(inplace=True)),\n",
    "                    (\n",
    "                        name + \"conv2\",\n",
    "                        nn.Conv2d(\n",
    "                            in_channels=features,\n",
    "                            out_channels=features,\n",
    "                            kernel_size=3,\n",
    "                            padding=1,\n",
    "                            bias=False,\n",
    "                        ),\n",
    "                    ),\n",
    "                    (name + \"norm2\", nn.BatchNorm2d(num_features=features)),\n",
    "                    (name + \"relu2\", nn.ReLU(inplace=True)),\n",
    "                ]\n",
    "            )\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiceLoss(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(DiceLoss, self).__init__()\n",
    "        self.smooth = 1.0\n",
    "\n",
    "    def forward(self, y_pred, y_true):\n",
    "        assert y_pred.size() == y_true.size()\n",
    "        y_pred = y_pred[:, 0].contiguous().view(-1)\n",
    "        y_true = y_true[:, 0].contiguous().view(-1)\n",
    "        intersection = (y_pred * y_true).sum()\n",
    "        dsc = (2. * intersection + self.smooth) / (\n",
    "            y_pred.sum() + y_true.sum() + self.smooth\n",
    "        )\n",
    "        return 1. - dsc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_images(x, y_true, y_pred, channel=1):\n",
    "    images = []\n",
    "    x_np = x[:, channel].cpu().numpy()\n",
    "    y_true_np = y_true[:, 0].cpu().numpy()\n",
    "    y_pred_np = y_pred[:, 0].cpu().numpy()\n",
    "    for i in range(x_np.shape[0]):\n",
    "        image = gray2rgb(np.squeeze(x_np[i]))\n",
    "        image = outline(image, y_pred_np[i], color=[255, 0, 0])\n",
    "        image = outline(image, y_true_np[i], color=[0, 255, 0])\n",
    "        images.append(image)\n",
    "    return images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gray2rgb(image):\n",
    "    w, h = image.shape\n",
    "    image += np.abs(np.min(image))\n",
    "    image_max = np.abs(np.max(image))\n",
    "    if image_max > 0:\n",
    "        image /= image_max\n",
    "    ret = np.empty((w, h, 3), dtype=np.uint8)\n",
    "    ret[:, :, 2] = ret[:, :, 1] = ret[:, :, 0] = image * 255\n",
    "    return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def outline(image, mask, color):\n",
    "    mask = np.round(mask)\n",
    "    yy, xx = np.nonzero(mask)\n",
    "    for y, x in zip(yy, xx):\n",
    "        if 0.0 < np.mean(mask[max(0, y - 1) : y + 2, max(0, x - 1) : x + 2]) < 1.0:\n",
    "            image[max(0, y) : y + 1, max(0, x) : x + 1] = color\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_loaders(batch_size, workers, image_size, aug_scale, aug_angle):\n",
    "    dataset_train, dataset_valid = datasets(\"/labs/mpsnyder/gbogu17/mri/lgg_mri_segmentation/kaggle_3m\", image_size, aug_scale, aug_angle)\n",
    "\n",
    "    def worker_init(worker_id):\n",
    "        np.random.seed(42 + worker_id)\n",
    "\n",
    "    loader_train = DataLoader(\n",
    "        dataset_train,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        drop_last=True,\n",
    "        num_workers=workers,\n",
    "        worker_init_fn=worker_init,\n",
    "    )\n",
    "    loader_valid = DataLoader(\n",
    "        dataset_valid,\n",
    "        batch_size=batch_size,\n",
    "        drop_last=False,\n",
    "        num_workers=workers,\n",
    "        worker_init_fn=worker_init,\n",
    "    )\n",
    "\n",
    "    return loader_train, loader_valid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def datasets(images, image_size, aug_scale, aug_angle):\n",
    "    train = BrainSegmentationDataset(\n",
    "        images_dir=images,\n",
    "        subset=\"train\",\n",
    "        image_size=image_size,\n",
    "        transform=transforms(scale=aug_scale, angle=aug_angle, flip_prob=0.5),\n",
    "    )\n",
    "    valid = BrainSegmentationDataset(\n",
    "        images_dir=images,\n",
    "        subset=\"validation\",\n",
    "        image_size=image_size,\n",
    "        random_sampling=False,\n",
    "    )\n",
    "    return train, valid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dsc(y_pred, y_true):\n",
    "    y_pred = np.round(y_pred).astype(int)\n",
    "    y_true = np.round(y_true).astype(int)\n",
    "    return np.sum(y_pred[y_true == 1]) * 2.0 / (np.sum(y_pred) + np.sum(y_true))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dsc_distribution(volumes):\n",
    "    dsc_dict = {}\n",
    "    for p in volumes:\n",
    "        y_pred = volumes[p][1]\n",
    "        y_true = volumes[p][2]\n",
    "        dsc_dict[p] = dsc(y_pred, y_true)\n",
    "    return dsc_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dsc_per_volume(validation_pred, validation_true, patient_slice_index):\n",
    "    dsc_list = []\n",
    "    num_slices = np.bincount([p[0] for p in patient_slice_index])\n",
    "    index = 0\n",
    "    for p in range(len(num_slices)):\n",
    "        y_pred = np.array(validation_pred[index : index + num_slices[p]])\n",
    "        y_true = np.array(validation_true[index : index + num_slices[p]])\n",
    "        dsc_list.append(dsc(y_pred, y_true))\n",
    "        index += num_slices[p]\n",
    "    return dsc_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def postprocess_per_volume(\n",
    "    input_list, pred_list, true_list, patient_slice_index, patients\n",
    "):\n",
    "    volumes = {}\n",
    "    num_slices = np.bincount([p[0] for p in patient_slice_index])\n",
    "    index = 0\n",
    "    for p in range(len(num_slices)):\n",
    "        volume_in = np.array(input_list[index : index + num_slices[p]])\n",
    "        volume_pred = np.round(\n",
    "            np.array(pred_list[index : index + num_slices[p]])\n",
    "        ).astype(int)\n",
    "        volume_true = np.array(true_list[index : index + num_slices[p]])\n",
    "        volumes[patients[p]] = (volume_in, volume_pred, volume_true)\n",
    "        index += num_slices[p]\n",
    "    return volumes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_loss_summary(loss, step, prefix=\"\"):\n",
    "    print(\"epoch {} | {}: {}\".format(step + 1, prefix + \"loss\", np.mean(loss)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_scalar_summary(tag, value, step):\n",
    "    print(\"epoch {} | {}: {}\".format(step + 1, tag, value))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_dsc(dsc_dist):\n",
    "    y_positions = np.arange(len(dsc_dist))\n",
    "    dsc_dist = sorted(dsc_dist.items(), key=lambda x: x[1])\n",
    "    values = [x[1] for x in dsc_dist]\n",
    "    labels = [x[0] for x in dsc_dist]\n",
    "    labels = [\"_\".join(l.split(\"_\")[1:-1]) for l in labels]\n",
    "    fig = plt.figure(figsize=(12, 8))\n",
    "    canvas = FigureCanvasAgg(fig)\n",
    "    plt.barh(y_positions, values, align=\"center\", color=\"skyblue\")\n",
    "    plt.yticks(y_positions, labels)\n",
    "    plt.xticks(np.arange(0.0, 1.0, 0.1))\n",
    "    plt.xlim([0.0, 1.0])\n",
    "    plt.gca().axvline(np.mean(values), color=\"tomato\", linewidth=2)\n",
    "    plt.gca().axvline(np.median(values), color=\"forestgreen\", linewidth=2)\n",
    "    plt.xlabel(\"Dice coefficient\", fontsize=\"x-large\")\n",
    "    plt.gca().xaxis.grid(color=\"silver\", alpha=0.5, linestyle=\"--\", linewidth=1)\n",
    "    plt.tight_layout()\n",
    "    canvas.draw()\n",
    "    plt.close()\n",
    "    s, (width, height) = canvas.print_to_buffer()\n",
    "    return np.fromstring(s, np.uint8).reshape((height, width, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch_size = 16 [reduced because of CUDA out of memory error]\n",
    "batch_size = 4\n",
    "# epochs = 50 [reduced to 5 to finish the process faster but not optimal]\n",
    "epochs = 5\n",
    "lr = 0.0001\n",
    "workers = 2\n",
    "weights = \"./\"\n",
    "image_size = 224\n",
    "aug_scale = 0.05\n",
    "aug_angle = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_validate():\n",
    "    device = torch.device(\"cpu\" if not torch.cuda.is_available() else \"cuda:0\")\n",
    "    \n",
    "    loader_train, loader_valid = data_loaders(batch_size, workers, image_size, aug_scale, aug_angle)\n",
    "    loaders = {\"train\": loader_train, \"valid\": loader_valid}\n",
    "    \n",
    "    unet = UNet(in_channels=BrainSegmentationDataset.in_channels, out_channels=BrainSegmentationDataset.out_channels)\n",
    "    unet.to(device)\n",
    "    \n",
    "    dsc_loss = DiceLoss()\n",
    "    best_validation_dsc = 0.0\n",
    "    \n",
    "    optimizer = optim.Adam(unet.parameters(), lr=lr)\n",
    "    \n",
    "    loss_train = []\n",
    "    loss_valid = []\n",
    "    \n",
    "    step = 0\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        for phase in [\"train\", \"valid\"]:\n",
    "            if phase == \"train\":\n",
    "                unet.train()\n",
    "            else:\n",
    "                unet.eval()\n",
    "    \n",
    "            validation_pred = []\n",
    "            validation_true = []\n",
    "    \n",
    "            for i, data in enumerate(loaders[phase]):\n",
    "                if phase == \"train\":\n",
    "                    step += 1\n",
    "    \n",
    "                x, y_true = data\n",
    "                x, y_true = x.to(device), y_true.to(device)\n",
    "    \n",
    "                optimizer.zero_grad()\n",
    "    \n",
    "                with torch.set_grad_enabled(phase == \"train\"):\n",
    "                    y_pred = unet(x)\n",
    "    \n",
    "                    loss = dsc_loss(y_pred, y_true)\n",
    "    \n",
    "                    if phase == \"valid\":\n",
    "                        loss_valid.append(loss.item())\n",
    "                        y_pred_np = y_pred.detach().cpu().numpy()\n",
    "                        validation_pred.extend(\n",
    "                            [y_pred_np[s] for s in range(y_pred_np.shape[0])]\n",
    "                        )\n",
    "                        y_true_np = y_true.detach().cpu().numpy()\n",
    "                        validation_true.extend(\n",
    "                            [y_true_np[s] for s in range(y_true_np.shape[0])]\n",
    "                        )\n",
    "                        \n",
    "                    if phase == \"train\":\n",
    "                        loss_train.append(loss.item())\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "    \n",
    "            if phase == \"train\":\n",
    "                log_loss_summary(loss_train, epoch)\n",
    "                loss_train = []\n",
    "\n",
    "            if phase == \"valid\":\n",
    "                log_loss_summary(loss_valid, epoch, prefix=\"val_\")\n",
    "                mean_dsc = np.mean(\n",
    "                    dsc_per_volume(\n",
    "                        validation_pred,\n",
    "                        validation_true,\n",
    "                        loader_valid.dataset.patient_slice_index,\n",
    "                    )\n",
    "                )\n",
    "                log_scalar_summary(\"val_dsc\", mean_dsc, epoch)\n",
    "                if mean_dsc > best_validation_dsc:\n",
    "                    best_validation_dsc = mean_dsc\n",
    "                    torch.save(unet.state_dict(), os.path.join(weights, \"unet.pt\"))\n",
    "                loss_valid = []\n",
    "    \n",
    "    print(\"\\nBest validation mean DSC: {:4f}\\n\".format(best_validation_dsc))\n",
    "    \n",
    "    state_dict = torch.load(os.path.join(weights, \"unet.pt\"))\n",
    "    unet.load_state_dict(state_dict)\n",
    "    unet.eval()\n",
    "    \n",
    "    input_list = []\n",
    "    pred_list = []\n",
    "    true_list = []\n",
    "    \n",
    "    for i, data in enumerate(loader_valid):\n",
    "        x, y_true = data\n",
    "        x, y_true = x.to(device), y_true.to(device)\n",
    "        with torch.set_grad_enabled(False):\n",
    "            y_pred = unet(x)\n",
    "            y_pred_np = y_pred.detach().cpu().numpy()\n",
    "            pred_list.extend([y_pred_np[s] for s in range(y_pred_np.shape[0])])\n",
    "            y_true_np = y_true.detach().cpu().numpy()\n",
    "            true_list.extend([y_true_np[s] for s in range(y_true_np.shape[0])])\n",
    "            x_np = x.detach().cpu().numpy()\n",
    "            input_list.extend([x_np[s] for s in range(x_np.shape[0])])\n",
    "            \n",
    "    volumes = postprocess_per_volume(\n",
    "        input_list,\n",
    "        pred_list,\n",
    "        true_list,\n",
    "        loader_valid.dataset.patient_slice_index,\n",
    "        loader_valid.dataset.patients,\n",
    "    )\n",
    "    \n",
    "    dsc_dist = dsc_distribution(volumes)\n",
    "\n",
    "    dsc_dist_plot = plot_dsc(dsc_dist)\n",
    "    imsave(\"./dsc.png\", dsc_dist_plot)\n",
    "\n",
    "    for p in volumes:\n",
    "        x = volumes[p][0]\n",
    "        y_pred = volumes[p][1]\n",
    "        y_true = volumes[p][2]\n",
    "        for s in range(x.shape[0]):\n",
    "            image = gray2rgb(x[s, 1])  # channel 1 is for FLAIR\n",
    "            image = outline(image, y_pred[s, 0], color=[255, 0, 0])\n",
    "            image = outline(image, y_true[s, 0], color=[0, 255, 0])\n",
    "            filename = \"{}-{}.png\".format(p, str(s).zfill(2))\n",
    "            filepath = os.path.join(\"./\", filename)\n",
    "            imsave(filepath, image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reading train images...\n",
      "preprocessing train volumes...\n",
      "cropping train volumes...\n",
      "padding train volumes...\n",
      "resizing train volumes...\n",
      "normalizing train volumes...\n",
      "done creating train dataset\n",
      "reading validation images...\n",
      "preprocessing validation volumes...\n",
      "cropping validation volumes...\n",
      "padding validation volumes...\n",
      "resizing validation volumes...\n",
      "normalizing validation volumes...\n",
      "done creating validation dataset\n",
      "epoch 1 | loss: 0.8003252012123593\n",
      "epoch 1 | val_loss: 0.8957568256251783\n",
      "epoch 1 | val_dsc: 0.33690554734963557\n",
      "epoch 2 | loss: 0.5626050973538872\n",
      "epoch 2 | val_loss: 0.7640607184674366\n",
      "epoch 2 | val_dsc: 0.7338745620930366\n",
      "epoch 3 | loss: 0.28837961887570024\n",
      "epoch 3 | val_loss: 0.6644515215632427\n",
      "epoch 3 | val_dsc: 0.7918927802227225\n",
      "epoch 4 | loss: 0.1743763219967163\n",
      "epoch 4 | val_loss: 0.6254718842276608\n",
      "epoch 4 | val_dsc: 0.8280030150396515\n",
      "epoch 5 | loss: 0.13543601852240894\n",
      "epoch 5 | val_loss: 0.6106225941554609\n",
      "epoch 5 | val_dsc: 0.8151660227317103\n",
      "\n",
      "Best validation mean DSC: 0.828003\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_validate()"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}