AURA_net.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "AURA-net.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "pEkOmPPFk1KH",
        "py246hQJk9WN",
        "Km09asDimU9R",
        "kC_633Vsl_Yx",
        "yZOgXF0In6qa"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pEkOmPPFk1KH"
      },
      "source": [
        "###Utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CxlLrgFDkr4M"
      },
      "source": [
        "import os\n",
        "import random\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.ndimage.interpolation import map_coordinates\n",
        "from scipy.ndimage.filters import gaussian_filter\n",
        "from torchvision.datasets import ImageFolder\n",
        "from torchvision import transforms\n",
        "import torch.nn as nn\n",
        "import torch\n",
        "from torch.optim import Adam, SGD\n",
        "from torch.autograd import Variable\n",
        "from torch.optim import lr_scheduler\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader\n",
        "from matplotlib.backends.backend_agg import FigureCanvasAgg\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision.transforms import Compose\n",
        "from torch.utils.data.sampler import SubsetRandomSampler\n",
        "from torchvision import transforms\n",
        "from torchvision import models\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchvision import models\n",
        "from torch.autograd import Variable as V\n",
        "import torch.optim as optim\n",
        "from torch.optim import lr_scheduler\n",
        "import glob\n",
        "from collections import OrderedDict\n",
        "from sklearn.preprocessing import MinMaxScaler\n",
        "\n",
        "import PIL.Image as Image\n",
        "\n",
        "from PIL import Image\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QfnTUZLfk3pG"
      },
      "source": [
        "!pip install numpy libtiff"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JP2JS-4Dk5Iu"
      },
      "source": [
        "!pip install albumentations"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FVBPoevIk6m-"
      },
      "source": [
        "from libtiff import TIFF\n",
        "from util import tif_to_nparray,create_dir,to_numpy,resize_my_images,load_image,load_set,save_to_tif"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "py246hQJk9WN"
      },
      "source": [
        "###Dataloading"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9-ll_0n6k_TO"
      },
      "source": [
        "from albumentations import (\n",
        "    Resize,\n",
        "    PadIfNeeded,\n",
        "    HorizontalFlip,\n",
        "    VerticalFlip,    \n",
        "    CenterCrop,    \n",
        "    Crop,\n",
        "    Compose,\n",
        "    Transpose,\n",
        "    RandomRotate90,\n",
        "    ElasticTransform,\n",
        "    GridDistortion, \n",
        "    OpticalDistortion,\n",
        "    RandomSizedCrop,\n",
        "    OneOf,\n",
        "    CLAHE,\n",
        "    RandomBrightnessContrast,    \n",
        "    RandomGamma    \n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jHsdBy7-lC02"
      },
      "source": [
        "original_height=512#put the size of you're image\n",
        "original_width=512#put the size of you're image\n",
        "#You can modify this code by adding/removing transformations\n",
        "transform_train=Compose([\n",
        "    VerticalFlip(p=0.5),              \n",
        "    RandomRotate90(p=0.5),\n",
        "    OneOf([\n",
        "        ElasticTransform(p=0.8, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),\n",
        "        GridDistortion(p=0.2),\n",
        "        OpticalDistortion(p=0.4, distort_limit=2, shift_limit=0.5)                  \n",
        "        ], p=0.8)\n",
        "         ])\n",
        "transform_val=Compose([\n",
        "   Resize(512,512)])#depends of your image sizes\n",
        "transform_nolastic=Compose([\n",
        "    VerticalFlip(p=0.5),              \n",
        "    RandomRotate90(p=0.5)])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ggMNVVkGlEt3"
      },
      "source": [
        "class Embryo_elastic(Dataset):\n",
        "    def __init__(self,img_fol,mask_fol,transform=None):\n",
        "        self.img_fol=img_fol\n",
        "        self.mask_fol=mask_fol\n",
        "        self.transform=transform\n",
        "    def __getitem__(self, idx):\n",
        "        image=load_set(self.img_fol,is_mask=False)[0][idx]\n",
        "        mask=load_set(self.mask_fol,is_mask=True)[0][idx]\n",
        "        if self.transform:\n",
        "            augmented = self.transform(image=image, mask=mask)\n",
        "            image=augmented['image']\n",
        "            mask=augmented['mask']\n",
        "            #change the normalize img by the normalization you want to do\n",
        "            normalize_img= transforms.Compose([transforms.ToTensor()])\n",
        "            image=normalize_img(image)\n",
        "            image=image.permute(0,2,1) \n",
        "            transform_to_tensor = transforms.Compose([transforms.ToTensor()])\n",
        "            mask=transform_to_tensor(mask)\n",
        "\n",
        "\n",
        "        else:\n",
        "            #same here\n",
        "            normalize_img= transforms.Compose([transforms.ToTensor()])\n",
        "            image=normalize_img(image)\n",
        "            image=image.permute(0,2,1) \n",
        "            transform_to_tensor = transforms.Compose([transforms.ToTensor()]) \n",
        "            mask=transform_to_tensor(mask)                        \n",
        "        \n",
        "        return image, mask\n",
        " \n",
        "\n",
        "    def __len__(self):\n",
        "        return len(load_set(self.mask_fol,is_mask=True)[1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3aoZqyFslf8_"
      },
      "source": [
        "def get_emb_elastic_loader(path_img,path_mask, validation_split=.20,  shuffle_dataset=True):\n",
        "    dataset = Embryo_elastic(path_img,path_mask)  # instantiating the data set.\n",
        "    dataset_size = len(dataset)\n",
        "    indices = list(range(dataset_size))\n",
        "\n",
        "    split_val = int(np.floor(validation_split * dataset_size))\n",
        "\n",
        "    if shuffle_dataset:\n",
        "        np.random.shuffle(indices)\n",
        "\n",
        "    train_indices = indices[split_val :]\n",
        "    val_indices = indices[: split_val]\n",
        "    \n",
        "    train_sampler = SubsetRandomSampler(train_indices)\n",
        "    valid_sampler = SubsetRandomSampler(val_indices)\n",
        "    \n",
        "    \n",
        "    dataset_train=Embryo_elastic(path_img,path_mask,transform=transform_train)\n",
        "    dataset_val=Embryo_elastic(path_img,path_mask,transform=transform_val)\n",
        "    loader = {\n",
        "        'train': DataLoader(dataset_train, batch_size=4, sampler=train_sampler),\n",
        "        'val': DataLoader(dataset_val, batch_size=1, sampler=valid_sampler),\n",
        "    }\n",
        "    return loader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UqCJks9xloIv"
      },
      "source": [
        "dataloader=get_emb_elastic_loader(path_to_img,path_to_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4EvtTdRFmGGW"
      },
      "source": [
        "#Dataloader for testing\n",
        "class Embryo_elastic_for_test(Dataset):\n",
        "    def __init__(self,img_fol,mask_fol):\n",
        "        self.img_fol=img_fol\n",
        "        self.mask_fol=mask_fol\n",
        "    def __getitem__(self, idx):\n",
        "        image=load_set(self.img_fol,is_mask=False)[0][idx]\n",
        "        mask=load_set(self.mask_fol,is_mask=True)[0][idx]\n",
        "        #add normalization if you want\n",
        "        normalize_img= transforms.Compose([transforms.ToTensor()])\n",
        "        image=normalize_img(image)\n",
        "        image=image.permute(0,2,1) \n",
        "        transform_to_tensor = transforms.Compose([transforms.ToTensor()]) \n",
        "        mask=transform_to_tensor(mask) \n",
        "                               \n",
        "        return image,mask\n",
        " \n",
        "\n",
        "    def __len__(self):\n",
        "        return len(load_set(self.img_fol,is_mask=False)[1])\n",
        "\n",
        "\n",
        "\n",
        "def get_dataset_test(path_img_test,path_mask_test):\n",
        "  dataset_test=Embryo_elastic_for_test(path_img_test,path_mask_test)\n",
        "  loader={\n",
        "      'test':DataLoader(dataset_test,batch_size=1)\n",
        "  }\n",
        "  return loader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vcUlmcnAkt00"
      },
      "source": [
        "dataloader_test=get_dataset_test(path_to_test_img,path_to_test_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Km09asDimU9R"
      },
      "source": [
        "###Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NwnMLDqemWzt"
      },
      "source": [
        "from collections import defaultdict\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ak4qCJXImdF2"
      },
      "source": [
        "def calc_loss(pred, target, metrics, bce_weight=0.5):#you can use the weights you want for the bce_weights\n",
        "    bce = torch.nn.functional.binary_cross_entropy_with_logits(pred, target)\n",
        "    \n",
        "    pred = torch.sigmoid(pred)\n",
        "    dice = dice_loss(pred, target)\n",
        "\n",
        "    AC=active_contour_loss(target,pred)\n",
        "\n",
        "    loss_bce_dice = bce * bce_weight + dice * (1 - bce_weight)\n",
        "    \n",
        "    loss=0.75*loss_bce_dice + 0.25*AC#put whatever weights you want\n",
        "\n",
        "    metrics['bce'] += bce.data.cpu().numpy() * target.size(0)\n",
        "    metrics['dice'] += dice.data.cpu().numpy() * target.size(0)\n",
        "    metrics['AC'] += AC*target.size(0)\n",
        "    metrics['loss_bce_dice'] += loss_bce_dice.data.cpu().numpy() * target.size(0)\n",
        "    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)\n",
        "\n",
        "    return loss\n",
        "\n",
        "def active_contour_loss(y_true, y_pred):\n",
        "  '''\n",
        "  y_true, y_pred: tensor of shape (B, C, H, W), where y_true[:,:,region_in_contour] == 1, y_true[:,:,region_out_contour] == 0.\n",
        "  weight: scalar, length term weight.\n",
        "  '''\n",
        "  # length term\n",
        "  delta_r = y_pred[:,:,1:,:] - y_pred[:,:,:-1,:] # horizontal gradient (B, C, H-1, W) \n",
        "  delta_c = y_pred[:,:,:,1:] - y_pred[:,:,:,:-1] # vertical gradient   (B, C, H,   W-1)\n",
        "  \n",
        "  delta_r    = delta_r[:,:,1:,:-2]**2  # (B, C, H-2, W-2)\n",
        "  delta_c    = delta_c[:,:,:-2,1:]**2  # (B, C, H-2, W-2)\n",
        "  delta_pred = torch.abs(delta_r + delta_c) \n",
        "\n",
        "  epsilon = 1e-8 # where is a parameter to avoid square root is zero in practice.\n",
        "  lenth = torch.mean(torch.sqrt(delta_pred + epsilon)) # eq.(11) in the paper, mean is used instead of sum.\n",
        "  \n",
        "  # region term\n",
        "  C_in  = torch.ones_like(y_pred)\n",
        "  C_out = torch.zeros_like(y_pred)\n",
        "\n",
        "  region_in  = torch.mean( y_pred     * (y_true - C_in )**2 ) # equ.(12) in the paper, mean is used instead of sum.\n",
        "  region_out = torch.mean( (1-y_pred) * (y_true - C_out)**2 ) \n",
        "  region = region_in + region_out\n",
        "  \n",
        "  loss =  0.2*lenth + 0.8*region\n",
        "\n",
        "  return loss\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def compute_metrics(metrics, epoch_samples):\n",
        "    computed_metrics = {}\n",
        "    for k in metrics.keys():\n",
        "        computed_metrics[k] = metrics[k] / epoch_samples\n",
        "    return computed_metrics\n",
        "\n",
        "\n",
        "def print_metrics(computed_metrics, phase):\n",
        "    outputs = []\n",
        "    for k in computed_metrics.keys():\n",
        "        outputs.append(\"{}:{:4f}\".format(k, computed_metrics[k]))\n",
        "\n",
        "    print(\"\\t{}-> {}\".format(phase.ljust(5), \"|\".join(outputs)))\n",
        "\n",
        "\n",
        "def dice_loss(pred, target, smooth=1.):\n",
        "    pred = pred.contiguous()\n",
        "    target = target.contiguous()\n",
        "\n",
        "    intersection = (pred * target).sum(dim=2).sum(dim=2)\n",
        "\n",
        "    loss = (1 - ((2. * intersection + smooth) /\n",
        "                 (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))\n",
        "\n",
        "    return loss.mean()\n",
        "\n",
        "\n",
        "\n",
        "def normalise_mask_set(mask, threshold):\n",
        "  \n",
        "  mask[mask > threshold] = 1\n",
        "  mask[mask <= threshold] = 0\n",
        "  return mask \n",
        "\n",
        "\n",
        "def normalise_mask(mask, threshold=0.5):\n",
        "  \n",
        "  mask[mask > threshold] = 1\n",
        "  mask[mask <= threshold] = 0\n",
        "  return mask    \n",
        "\n",
        "\n",
        "def metrics_line(data):\n",
        "    phases = list(data.keys())\n",
        "    metrics = list(data[phases[0]][0].keys())\n",
        "\n",
        "    i = 0\n",
        "    fig, axs = plt.subplots(1, len(metrics))\n",
        "    fig.set_figheight(6)\n",
        "    fig.set_figwidth(6 * len(metrics))\n",
        "    for metric in metrics:\n",
        "        for phase in phases:\n",
        "            axs[i].plot([i[metric] for i in data[phase]], label=phase)\n",
        "        axs[i].set_title(metric)\n",
        "        i += 1\n",
        "\n",
        "    plt.legend()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CahDFInmm1GO"
      },
      "source": [
        "from collections import defaultdict\n",
        "import copy\n",
        "import time\n",
        "\n",
        "import torch\n",
        "\n",
        "\n",
        "class Trainers(object):\n",
        "\n",
        "    def __init__(self, model, optimizer=None, scheduler=None):\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        self.device = torch.device(\n",
        "            'cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "        self.model = model.to(self.device)\n",
        "\n",
        "        self.optimizer = optimizer\n",
        "        if self.optimizer == None:\n",
        "            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)\n",
        "\n",
        "        self.scheduler = scheduler\n",
        "        if self.scheduler == None:\n",
        "            self.scheduler = torch.optim.lr_scheduler.StepLR(\n",
        "                self.optimizer, step_size=5, gamma=0.1)\n",
        "\n",
        "    def train_model(self, dataloaders, num_epochs=25):\n",
        "        best_model_wts = copy.deepcopy(self.model.state_dict())\n",
        "        best_loss = 1e10\n",
        "        epochs_metrics = {\n",
        "            'train': [],\n",
        "            'val': []\n",
        "        }\n",
        "\n",
        "        for epoch in range(num_epochs):\n",
        "            print('Epoch {}/{}:'.format(epoch+1, num_epochs))\n",
        "\n",
        "            since = time.time()\n",
        "\n",
        "\n",
        "            # Each epoch has a training and validation phase\n",
        "            for phase in ['train', 'val']:\n",
        "                if phase == 'train':\n",
        "                    for param_group in self.optimizer.param_groups:\n",
        "                        print(\"\\tlearning rate: {:.2e}\".format(\n",
        "                            param_group['lr']))\n",
        "\n",
        "                    self.model.train()  # Set model to training mode\n",
        "                else:\n",
        "                    self.model.eval()   # Set model to evaluate mode\n",
        "\n",
        "                metrics = defaultdict(float)\n",
        "                epoch_samples = 0\n",
        "\n",
        "                for inputs, labels in dataloaders[phase]:\n",
        "                    \n",
        "                    inputs = inputs.to(self.device)\n",
        "                    labels = labels.to(self.device)\n",
        "                    inputs=inputs.permute(0,1,3,2)#put the mask and inputs in the same settings\n",
        "                    \n",
        "\n",
        "                    # zero the parameter gradients\n",
        "                    self.optimizer.zero_grad()\n",
        "\n",
        "                    # forward\n",
        "                    # track history if only in train\n",
        "                    with torch.set_grad_enabled(phase == 'train'):\n",
        "                        \n",
        "                        outputs = self.model(inputs)\n",
        "                        \n",
        "                        loss = calc_loss(outputs, labels,metrics)\n",
        "                        \n",
        "                        # backward + optimize only if in training phase\n",
        "                        if phase == 'train':\n",
        "                            loss.backward()\n",
        "                            self.optimizer.step()\n",
        "                            \n",
        "\n",
        "                    # statistics\n",
        "                    epoch_samples += inputs.size(0)\n",
        "\n",
        "                computed_metrics = compute_metrics(metrics, epoch_samples)\n",
        "                print_metrics(computed_metrics, phase)\n",
        "                epochs_metrics[phase].append(computed_metrics)\n",
        "                epoch_loss = metrics['loss'] / epoch_samples\n",
        "\n",
        "                if phase == 'train':\n",
        "                    self.scheduler.step()\n",
        "\n",
        "                # deep copy the model\n",
        "                if phase == 'val' and epoch_loss < best_loss:\n",
        "                    print(\"\\tSaving best model, epoch loss {:4f} < best loss {:4f}\".format(\n",
        "                        epoch_loss, best_loss))\n",
        "                    best_loss = epoch_loss\n",
        "                    best_model_wts = copy.deepcopy(self.model.state_dict())\n",
        "\n",
        "            time_elapsed = time.time() - since\n",
        "            print('\\t{:.0f}m {:.0f}s'.format(\n",
        "                time_elapsed // 60, time_elapsed % 60))\n",
        "            print('-' * 10)\n",
        "\n",
        "        print('Best val loss: {:4f}'.format(best_loss))\n",
        "\n",
        "        # load best model weights\n",
        "        self.model.load_state_dict(best_model_wts)\n",
        "\n",
        "        metrics_line(epochs_metrics)\n",
        "   \n",
        "    \n",
        "    \n",
        "    def predict(self, X):\n",
        "        \n",
        "        self.model.eval()\n",
        "        X=X.permute(0,1,3,2)\n",
        "        inputs = X.to(self.device)\n",
        "        pred = self.model(inputs)\n",
        "        \n",
        "        avant_norm = pred.data.cpu().numpy()\n",
        "        \n",
        "       \n",
        "        \n",
        "\n",
        "        return avant_norm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kC_633Vsl_Yx"
      },
      "source": [
        "###AURA-net"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PakVkssLmBOl"
      },
      "source": [
        "from __future__ import print_function, division\n",
        "import torch.nn as nn\n",
        "import torch.utils.data\n",
        "import torch\n",
        "from torchvision import models\n",
        "\n",
        "def convrelu(in_channels, out_channels, kernel, padding):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
        "        nn.ReLU(inplace=True),\n",
        "    )\n",
        "\n",
        "class conv_block(nn.Module):\n",
        "    \"\"\"\n",
        "    Convolution Block \n",
        "    \"\"\"\n",
        "    def __init__(self, in_ch, out_ch):\n",
        "        super(conv_block, self).__init__()\n",
        "        \n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "            nn.BatchNorm2d(out_ch),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "            nn.BatchNorm2d(out_ch),\n",
        "            nn.ReLU(inplace=True))\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        x = self.conv(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class up_conv(nn.Module):\n",
        "    \"\"\"\n",
        "    Up Convolution Block\n",
        "    \"\"\"\n",
        "    def __init__(self, in_ch, out_ch):\n",
        "        super(up_conv, self).__init__()\n",
        "        self.up = nn.Sequential(\n",
        "            nn.Upsample(scale_factor=2),\n",
        "            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "            nn.BatchNorm2d(out_ch),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.up(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "\n",
        "class Attention_block(nn.Module):\n",
        "    \"\"\"\n",
        "    Attention Block\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, F_g, F_l, F_int):\n",
        "        super(Attention_block, self).__init__()\n",
        "\n",
        "        self.W_g = nn.Sequential(\n",
        "            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),\n",
        "            nn.BatchNorm2d(F_int)\n",
        "        )\n",
        "\n",
        "        self.W_x = nn.Sequential(\n",
        "            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),\n",
        "            nn.BatchNorm2d(F_int)\n",
        "        )\n",
        "\n",
        "        self.psi = nn.Sequential(\n",
        "            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),\n",
        "            nn.BatchNorm2d(1),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, g, x):\n",
        "        g1 = self.W_g(g)\n",
        "        x1 = self.W_x(x)\n",
        "        psi = self.relu(g1 + x1)\n",
        "        psi = self.psi(psi)\n",
        "        out = x * psi\n",
        "        return out\n",
        "\n",
        "\n",
        "def convrelu(in_channels, out_channels, kernel, padding):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
        "        nn.ReLU(inplace=True),\n",
        "    )\n",
        "\n",
        "\n",
        "class AURA-net(nn.Module):\n",
        "    \n",
        "    def __init__(self, img_ch=3, output_ch=1):\n",
        "        super(AURA-net, self).__init__()\n",
        "\n",
        "        n1 = 64\n",
        "        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]\n",
        "\n",
        "        self.base_model = models.resnet18(pretrained=True)\n",
        "        \n",
        "\n",
        "        self.base_layers = list(self.base_model.children())\n",
        "        \n",
        "        # size=(N, 64, x.H/2, x.W/2)\n",
        "        \n",
        "\n",
        "        self.layer0 = nn.Sequential(*self.base_layers[:3])\n",
        "     \n",
        "        self.layer1 = nn.Sequential(*self.base_layers[3:5])\n",
        "       \n",
        "        self.layer2 = self.base_layers[5]  # size=(N, 128, x.H/8, x.W/8)\n",
        "       \n",
        "        self.layer3 = self.base_layers[6]  # size=(N, 256, x.H/16, x.W/16)\n",
        "        \n",
        "        self.layer4 = self.base_layers[7]  # size=(N, 512, x.H/32, x.W/32)\n",
        "\n",
        "        \n",
        "      \n",
        "        self.fix=nn.Upsample(scale_factor=4)\n",
        "        \n",
        "        \n",
        "        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "\n",
        "        self.Conv1 = conv_block(img_ch, filters[0])\n",
        "        self.Conv2 = conv_block(filters[0], filters[1])\n",
        "        self.Conv3 = conv_block(filters[1], filters[2])\n",
        "        self.Conv4 = conv_block(filters[2], filters[3])\n",
        "        self.Conv5 = conv_block(filters[3], filters[4])\n",
        "\n",
        "        self.Up5 = up_conv(filters[4], filters[3])\n",
        "        self.Att5 = Attention_block(F_g=filters[3], F_l=filters[3], F_int=filters[2])\n",
        "        self.Up_conv5 = conv_block(filters[4], filters[3])\n",
        "\n",
        "        self.Up4 = up_conv(filters[3], filters[2])\n",
        "        self.Att4 = Attention_block(F_g=filters[2], F_l=filters[2], F_int=filters[1])\n",
        "        self.Up_conv4 = conv_block(filters[3], filters[2])\n",
        "\n",
        "        self.Up3 = up_conv(filters[2], filters[1])\n",
        "        self.Att3 = Attention_block(F_g=filters[1], F_l=filters[1], F_int=filters[0])\n",
        "        self.Up_conv3 = conv_block(filters[2], filters[1])\n",
        "\n",
        "        self.Up2 = up_conv(filters[1], filters[0])\n",
        "        self.Att2 = Attention_block(F_g=filters[0], F_l=filters[0], F_int=32)\n",
        "        self.Up_conv2 = conv_block(filters[1], filters[0])\n",
        "\n",
        "        self.Up1 = up_conv(filters[0], filters[0])\n",
        "        self.Att1 = Attention_block(F_g=filters[0], F_l=filters[0],  F_int=32)\n",
        "        self.Up_conv1 = conv_block(filters[1], filters[0])\n",
        "\n",
        "        self.Up0 = up_conv(filters[0], filters[0])\n",
        "        self.Att0 = Attention_block(F_g=filters[0], F_l=filters[0],  F_int=32)\n",
        "        self.Up_conv0 = conv_block(filters[0], filters[0])\n",
        "\n",
        "        self.conv_original_size0 = convrelu(3, 64, 3, 1)\n",
        "        self.conv_original_size1 = convrelu(64, 64, 3, 1)\n",
        "        self.conv_original_size2 = convrelu(128, 64, 3, 1)\n",
        "\n",
        "        self.Conv = nn.Conv2d(filters[0], output_ch, kernel_size=1, stride=1, padding=0)\n",
        "        \n",
        "        #self.active = torch.nn.Sigmoid()\n",
        "        \n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        x_original = self.conv_original_size0(x)\n",
        "        x_original = self.conv_original_size1(x_original)\n",
        "        \n",
        "       \n",
        "        \n",
        "        e0 = self.layer0(x)\n",
        "        \n",
        "        \n",
        "      \n",
        "        e1= self.layer1(e0)\n",
        "        \n",
        "        \n",
        "        e2 = self.layer2(e1)\n",
        "    \n",
        "        \n",
        "        \n",
        "        e3 = self.layer3(e2)\n",
        "        \n",
        "        \n",
        "        \n",
        "        e4 = self.layer4(e3)\n",
        "        \n",
        "       \n",
        "        e5 = self.Maxpool4(e4)\n",
        "        e5 = self.Conv5(e5)\n",
        "        \n",
        "        \n",
        "        d5 = self.Up5(e5)\n",
        "        x4 = self.Att5(g=d5, x=e4)\n",
        "        d5 = torch.cat((x4, d5), dim=1)\n",
        "        d5 = self.Up_conv5(d5)\n",
        "       \n",
        "\n",
        "        d4 = self.Up4(d5)\n",
        "        x3 = self.Att4(g=d4, x=e3)\n",
        "        d4 = torch.cat((x3, d4), dim=1)\n",
        "        d4 = self.Up_conv4(d4)\n",
        "        \n",
        "        \n",
        "        d3 = self.Up3(d4)\n",
        "        x2 = self.Att3(g=d3, x=e2)\n",
        "        d3 = torch.cat((x2, d3), dim=1)\n",
        "        d3 = self.Up_conv3(d3)\n",
        "        \n",
        "        \n",
        "        d2 = self.Up2(d3)\n",
        "        x1 = self.Att2(g=d2, x=e1)\n",
        "        d2 = torch.cat((x1, d2), dim=1)\n",
        "        d2 = self.Up_conv2(d2)\n",
        "        \n",
        "\n",
        "        d1 = self.Up1(d2)\n",
        "        x0 = self.Att1(g=d1, x=e0)\n",
        "        d1 = torch.cat((x0, d1), dim=1)\n",
        "        d1 = self.Up_conv1(d1)\n",
        "        \n",
        "\n",
        "        d0 = self.Up0(d1)\n",
        "        x=torch.cat((x_original,d0), dim=1)\n",
        "        d0=self.conv_original_size2(x)\n",
        "        \n",
        "\n",
        "        out = self.Conv(d0)\n",
        "\n",
        "\n",
        "\n",
        "      #  out = self.active(out)\n",
        "\n",
        "        return out\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yZOgXF0In6qa"
      },
      "source": [
        "###Training\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K7zPN7Ftobvy"
      },
      "source": [
        "learning_rate=\n",
        "step_size=\n",
        "gamma=\n",
        "num_epochs="
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mCSK4r8Qn8Km"
      },
      "source": [
        "model=AURA-net()\n",
        "optimizer_func = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "scheduler = lr_scheduler.StepLR(optimizer_func, step_size=step_size, gamma=gamma)\n",
        "trainer = Trainers(model, optimizer=optimizer_func, scheduler=scheduler)\n",
        "trainer.train_model(dataloader, num_epochs=num_epochs)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "INkZBhH2ool2"
      },
      "source": [
        "#to plot all images in the validation set, and you can do the same with the test set using dataloader_test['test'].\n",
        "for epoch, (images, masks) in enumerate(dataloader['val']):\n",
        "        truth=masks\n",
        "        img=images\n",
        "        proba_prediction= trainer.predict(images)\n",
        "        print('###########################################################')\n",
        "        print('###########################################################')\n",
        "        print('###########################################################')\n",
        "        tmp=proba_prediction.squeeze()\n",
        "        normalizedImg = np.zeros((512, 512))#put the size of your images here\n",
        "        normalizedImg = cv2.normalize(tmp,  normalizedImg, 0, 1, cv2.NORM_MINMAX)\n",
        "        mask_pred=normalise_mask_set(normalizedImg,0.5)\n",
        "        \n",
        "        plot_side_by_side(to_numpy(truth.squeeze()),mask_pred.squeeze())\n",
        "        plot_side_by_side(to_numpy(img.permute(0,3,2,1).squeeze()),proba_prediction.squeeze())\n",
        "       \n",
        "        print('###########################################################')\n",
        "        print('###########################################################')\n",
        "        print('###########################################################')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}