UNet_Neuro.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/travislatchman/Neuroimage-Registration-and-Synthesis/blob/main/UNet_Neuro.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installations"
      ],
      "metadata": {
        "id": "Q32lmgVmL-MC"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QW0EUVHtQOqt",
        "outputId": "142cbe27-0ff8-4323-e787-270d1a9904cd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting nilearn\n",
            "  Downloading nilearn-0.10.1-py3-none-any.whl (10.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.3/10.3 MB\u001b[0m \u001b[31m90.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: joblib>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from nilearn) (1.2.0)\n",
            "Requirement already satisfied: lxml in /usr/local/lib/python3.10/dist-packages (from nilearn) (4.9.2)\n",
            "Collecting nibabel>=3.2.0 (from nilearn)\n",
            "  Downloading nibabel-5.1.0-py3-none-any.whl (3.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.3/3.3 MB\u001b[0m \u001b[31m85.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.19.0 in /usr/local/lib/python3.10/dist-packages (from nilearn) (1.22.4)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from nilearn) (23.1)\n",
            "Requirement already satisfied: pandas>=1.1.5 in /usr/local/lib/python3.10/dist-packages (from nilearn) (1.5.3)\n",
            "Requirement already satisfied: requests>=2.25.0 in /usr/local/lib/python3.10/dist-packages (from nilearn) (2.27.1)\n",
            "Requirement already satisfied: scikit-learn>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from nilearn) (1.2.2)\n",
            "Requirement already satisfied: scipy>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from nilearn) (1.10.1)\n",
            "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->nilearn) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->nilearn) (2022.7.1)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.25.0->nilearn) (1.26.15)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.25.0->nilearn) (2022.12.7)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests>=2.25.0->nilearn) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.25.0->nilearn) (3.4)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=1.0.0->nilearn) (3.1.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.1->pandas>=1.1.5->nilearn) (1.16.0)\n",
            "Installing collected packages: nibabel, nilearn\n",
            "  Attempting uninstall: nibabel\n",
            "    Found existing installation: nibabel 3.0.2\n",
            "    Uninstalling nibabel-3.0.2:\n",
            "      Successfully uninstalled nibabel-3.0.2\n",
            "Successfully installed nibabel-5.1.0 nilearn-0.10.1\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: nibabel in /usr/local/lib/python3.10/dist-packages (5.1.0)\n",
            "Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.10/dist-packages (from nibabel) (1.22.4)\n",
            "Requirement already satisfied: packaging>=17 in /usr/local/lib/python3.10/dist-packages (from nibabel) (23.1)\n"
          ]
        }
      ],
      "source": [
        "!pip install nilearn\n",
        "!pip install nibabel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HBYSQ831xUwZ",
        "outputId": "0cfdfd7c-f2ce-434a-a512-e1544289d777"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting torchviz\n",
            "  Downloading torchviz-0.0.2.tar.gz (4.9 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from torchviz) (2.0.0+cu118)\n",
            "Requirement already satisfied: graphviz in /usr/local/lib/python3.10/dist-packages (from torchviz) (0.20.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->torchviz) (3.12.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch->torchviz) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->torchviz) (1.11.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->torchviz) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->torchviz) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch->torchviz) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch->torchviz) (3.25.2)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch->torchviz) (16.0.3)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->torchviz) (2.1.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->torchviz) (1.3.0)\n",
            "Building wheels for collected packages: torchviz\n",
            "  Building wheel for torchviz (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for torchviz: filename=torchviz-0.0.2-py3-none-any.whl size=4131 sha256=acb9f0edca4df405f1ceb13c9bd5a6c6fc5fa0e9f71f49b0526de5e848e5a919\n",
            "  Stored in directory: /root/.cache/pip/wheels/4c/97/88/a02973217949e0db0c9f4346d154085f4725f99c4f15a87094\n",
            "Successfully built torchviz\n",
            "Installing collected packages: torchviz\n",
            "Successfully installed torchviz-0.0.2\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting onnx\n",
            "  Downloading onnx-1.14.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (14.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m14.6/14.6 MB\u001b[0m \u001b[31m103.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from onnx) (1.22.4)\n",
            "Requirement already satisfied: protobuf>=3.20.2 in /usr/local/lib/python3.10/dist-packages (from onnx) (3.20.3)\n",
            "Requirement already satisfied: typing-extensions>=3.6.2.1 in /usr/local/lib/python3.10/dist-packages (from onnx) (4.5.0)\n",
            "Installing collected packages: onnx\n",
            "Successfully installed onnx-1.14.0\n"
          ]
        }
      ],
      "source": [
        "!pip install torchviz\n",
        "!pip install onnx\n",
        "!pip install torch torchvision\n",
        "!pip install netron\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lpXCwqlIqWP5"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import transforms\n",
        "import nibabel as nib\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from skimage.metrics import mean_squared_error as mse, peak_signal_noise_ratio as psnr, structural_similarity as ssim\n",
        "import csv\n",
        "from torchviz import make_dot\n",
        "import torch.onnx\n",
        "import random\n",
        "import time\n",
        "import copy\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_SsbKTHpQLMU",
        "outputId": "f000f172-f5b3-4c84-b497-8ac9d5bbe699"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R_eHHYokbqGI",
        "outputId": "b1ce90da-55da-4c77-f20f-807b590c7a38"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7fea3201eef0>"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "seed = 42\n",
        "np.random.seed(seed)\n",
        "torch.manual_seed(seed)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aR_ahDk4sLBi"
      },
      "source": [
        "## MedIA_Project2_registered folder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "446OARZCrVbz"
      },
      "outputs": [],
      "source": [
        "class MedIARegisteredDataset(Dataset):\n",
        "    \"\"\"\n",
        "    A class used to represent a Medical Image Analysis (MedIA) Registered Dataset.\n",
        "\n",
        "    ...\n",
        "\n",
        "    Attributes\n",
        "    ----------\n",
        "    root_dir : str\n",
        "        The path to the root directory containing the patient folders.\n",
        "    transform : callable, optional\n",
        "        An optional transform function to apply to the images.\n",
        "    patient_ids : list\n",
        "        A list of patient IDs used to identify the corresponding patient folders.\n",
        "\n",
        "    Methods\n",
        "    -------\n",
        "    __init__(self, root_dir, patient_ids, transform=None):\n",
        "        Initializes the MedIARegisteredDataset object.\n",
        "    __len__(self):\n",
        "        Returns the number of patients in the dataset.\n",
        "    __getitem__(self, idx):\n",
        "        Loads and returns a tuple of the input images (T1w and T2w) and target images (FA and ADC) for the patient\n",
        "        at the given index.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, root_dir, patient_ids, transform=None):\n",
        "        \"\"\"\n",
        "        Parameters\n",
        "        ----------\n",
        "        root_dir : str\n",
        "            The path to the root directory containing the patient folders.\n",
        "        patient_ids : list\n",
        "            A list of patient IDs used to identify the corresponding patient folders.\n",
        "        transform : callable, optional\n",
        "            An optional transform function to apply to the images.\n",
        "        \"\"\"\n",
        "        self.root_dir = root_dir\n",
        "        self.transform = transform\n",
        "        self.patient_ids = patient_ids\n",
        "\n",
        "    def __len__(self):\n",
        "        \"\"\"\n",
        "        Determines the number of patients in the dataset.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        int\n",
        "            The number of patients in the dataset.\n",
        "        \"\"\"\n",
        "        return len(self.patient_ids)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        \"\"\"\n",
        "        Loads and returns a tuple of the input images (T1w and T2w) and target images (FA and ADC) for the patient\n",
        "        at the given index.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        idx : int\n",
        "            The index of the patient.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        tuple\n",
        "            A tuple containing two torch.Tensor objects:\n",
        "            - input_imgs: a torch.Tensor containing the stacked T1w and T2w images.\n",
        "            - target_imgs: a torch.Tensor containing the stacked FA and ADC images.\n",
        "        \"\"\"\n",
        "        patient_id = self.patient_ids[idx]\n",
        "        patient_folder = os.path.join(self.root_dir, patient_id)\n",
        "\n",
        "        # Load the input images (T1w and T2w)\n",
        "        t1w_path = os.path.join(patient_folder, 'T1w_1mm.nii.gz')\n",
        "        t2w_path = os.path.join(patient_folder, 'T2w_1mm_noalign.nii.gz')\n",
        "        t1w_img = nib.load(t1w_path).get_fdata()\n",
        "        t2w_img = nib.load(t2w_path).get_fdata()\n",
        "\n",
        "        # Stack the input images along the channel dimension\n",
        "        input_imgs = np.stack([t1w_img, t2w_img], axis=0)\n",
        "\n",
        "        # Load the target images (FA and ADC)\n",
        "        fa_path = os.path.join(patient_folder, 'FA_deformed.nii.gz')\n",
        "        adc_path = os.path.join(patient_folder, 'ADC_deformed.nii.gz')\n",
        "        fa_img = nib.load(fa_path).get_fdata()\n",
        "        adc_img = nib.load(adc_path).get_fdata()\n",
        "\n",
        "        # Stack the target images along the channel dimension\n",
        "        target_imgs = np.stack([fa_img, adc_img], axis=0)\n",
        "\n",
        "        if self.transform:\n",
        "            input_imgs, target_imgs = self.transform(input_imgs, target_imgs)\n",
        "\n",
        "        return torch.tensor(input_imgs).float(), torch.tensor(target_imgs).float()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AOcyu8aozqIE"
      },
      "outputs": [],
      "source": [
        "#root_dir = r\"G:\\My Drive\\Project 2 Neuroimage Synthesis\\MedIA_Project2_registered\"\n",
        "root_dir = \"/content/drive/My Drive/Project 2 Neuroimage Synthesis/MedIA_Project2_registered\"\n",
        "\n",
        "# Define patient_ids here\n",
        "patient_ids = [f\"{i:03d}\" for i in range(1, 201) if i != 163]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "30f1s7ZBsnN2"
      },
      "source": [
        "## split the patient IDs into training, validation, and testing sets (80/10/10 split)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wRWB5h_issZt"
      },
      "outputs": [],
      "source": [
        "# #patient_ids = [str(i).zfill(3) for i in range(1, 201) if i != 163]\n",
        "rng = np.random.default_rng(seed)\n",
        "rng.shuffle(patient_ids)\n",
        "\n",
        "train_size = int(0.8 * len(patient_ids))\n",
        "val_size = int(0.1 * len(patient_ids))\n",
        "test_size = len(patient_ids) - train_size - val_size\n",
        "\n",
        "train_ids = patient_ids[:train_size]\n",
        "val_ids = patient_ids[train_size:train_size+val_size]\n",
        "test_ids = patient_ids[train_size+val_size:]\n",
        "\n",
        "# separate instances of the dataset for each split\n",
        "train_dataset = MedIARegisteredDataset(root_dir, patient_ids=train_ids)\n",
        "val_dataset = MedIARegisteredDataset(root_dir, patient_ids=val_ids)\n",
        "test_dataset = MedIARegisteredDataset(root_dir, patient_ids=test_ids)\n",
        "\n",
        "# create data loaders for each dataset\n",
        "\n",
        "batch_size = 2\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YErk6lRVqiZk"
      },
      "source": [
        "##UNet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O_zD42Lu_tcU"
      },
      "outputs": [],
      "source": [
        "class ConvBlock(nn.Module):\n",
        "    \"\"\"\n",
        "    A class used to represent a convolution block for a 3D U-Net.\n",
        "\n",
        "    ...\n",
        "\n",
        "    Attributes\n",
        "    ----------\n",
        "    conv : torch.nn.Module\n",
        "        The convolution layer.\n",
        "    relu : torch.nn.Module\n",
        "        The ReLU activation function.\n",
        "\n",
        "    Methods\n",
        "    -------\n",
        "    __init__(self, in_channels, out_channels):\n",
        "        Initializes the ConvBlock object.\n",
        "    forward(self, x):\n",
        "        Performs the forward pass of the convolution block.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "      \"\"\"\n",
        "        Parameters\n",
        "        ----------\n",
        "        in_channels : int\n",
        "            The number of input channels.\n",
        "        out_channels : int\n",
        "            The number of output channels.\n",
        "        \"\"\"\n",
        "        super(ConvBlock, self).__init__()\n",
        "        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1)\n",
        "        self.relu = nn.ReLU(inplace=False)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        Performs the forward pass of the convolution block.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        x : torch.Tensor\n",
        "            The input tensor.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        torch.Tensor\n",
        "            The output tensor after applying the convolution and the ReLU activation function.\n",
        "        \"\"\"\n",
        "        return self.relu(self.conv(x))\n",
        "\n",
        "\n",
        "class UNet(nn.Module):\n",
        "    \"\"\"\n",
        "    A class used to represent a 3D U-Net.\n",
        "\n",
        "    ...\n",
        "\n",
        "    Attributes\n",
        "    ----------\n",
        "    enc1, enc2, enc3, enc4 : ConvBlock\n",
        "        The encoder blocks.\n",
        "    pool : torch.nn.Module\n",
        "        The max pooling layer.\n",
        "    middle : ConvBlock\n",
        "        The middle block.\n",
        "    up4, up3, up2, up1 : torch.nn.Module\n",
        "        The up-convolution layers.\n",
        "    dec4, dec3, dec2, dec1 : ConvBlock\n",
        "        The decoder blocks.\n",
        "    output : torch.nn.Module\n",
        "        The output convolution layer.\n",
        "\n",
        "    Methods\n",
        "    -------\n",
        "    __init__(self, in_channels, out_channels):\n",
        "        Initializes the UNet object.\n",
        "    forward(self, x):\n",
        "        Performs the forward pass of the U-Net.\n",
        "    up_concat(self, up, encode, decode):\n",
        "        Performs the up-sampling and concatenation operation.\n",
        "    \"\"\"\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        \"\"\"\n",
        "        Parameters\n",
        "        ----------\n",
        "        in_channels : int\n",
        "            The number of input channels.\n",
        "        out_channels : int\n",
        "            The number of output channels.\n",
        "        \"\"\"\n",
        "        super(UNet, self).__init__()\n",
        "\n",
        "        self.enc1 = ConvBlock(in_channels, 32)\n",
        "        self.enc2 = ConvBlock(32, 64)\n",
        "        self.enc3 = ConvBlock(64, 128)\n",
        "        self.enc4 = ConvBlock(128, 256)\n",
        "\n",
        "        self.pool = nn.MaxPool3d(2)\n",
        "\n",
        "        self.middle = ConvBlock(256, 512)\n",
        "\n",
        "        self.up4 = nn.ConvTranspose3d(512, 256, kernel_size=2, stride=2)\n",
        "        self.dec4 = ConvBlock(512, 256)\n",
        "\n",
        "        self.up3 = nn.ConvTranspose3d(256, 128, kernel_size=2, stride=2)\n",
        "        self.dec3 = ConvBlock(256, 128)\n",
        "\n",
        "        self.up2 = nn.ConvTranspose3d(128, 64, kernel_size=2, stride=2)\n",
        "        self.dec2 = ConvBlock(128, 64)\n",
        "\n",
        "        self.up1 = nn.ConvTranspose3d(64, 32, kernel_size=2, stride=2)\n",
        "        self.dec1 = ConvBlock(64, 32)\n",
        "\n",
        "        self.output = nn.Conv3d(32, out_channels, kernel_size=1)\n",
        "\n",
        "    def forward(self, x):\n",
        "      \"\"\"\n",
        "        Performs the forward pass of the U-Net.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        x : torch.Tensor\n",
        "            The input tensor.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        torch.Tensor\n",
        "            The output tensor after passing through the U-Net.\n",
        "        \"\"\"\n",
        "        enc1 = self.enc1(x)\n",
        "        enc2 = self.enc2(self.pool(enc1))\n",
        "        enc3 = self.enc3(self.pool(enc2))\n",
        "        enc4 = self.enc4(self.pool(enc3))\n",
        "\n",
        "        middle = self.middle(self.pool(enc4))\n",
        "\n",
        "        up4 = self.up_concat(self.up4, enc4, middle)\n",
        "        dec4 = self.dec4(up4)\n",
        "\n",
        "        up3 = self.up_concat(self.up3, enc3, dec4)\n",
        "        dec3 = self.dec3(up3)\n",
        "\n",
        "        up2 = self.up_concat(self.up2, enc2, dec3)\n",
        "        dec2 = self.dec2(up2)\n",
        "\n",
        "        up1 = self.up_concat(self.up1, enc1, dec2)\n",
        "        dec1 = self.dec1(up1)\n",
        "\n",
        "        output = self.output(dec1)\n",
        "\n",
        "        return output\n",
        "\n",
        "    def up_concat(self, up, encode, decode):\n",
        "       \"\"\"\n",
        "        Performs the up-sampling and concatenation operation.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        up : torch.nn.Module\n",
        "            The up-convolution layer.\n",
        "        encode : torch.Tensor\n",
        "            The output tensor from the corresponding encoder block.\n",
        "        decode : torch.Tensor\n",
        "            The output tensor from the previous decoder block.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        torch.Tensor\n",
        "            The merged tensor after up-sampling and concatenation.\n",
        "        \"\"\"\n",
        "        up = up(decode)\n",
        "        up = F.interpolate(up, size=encode.size()[2:], mode='trilinear', align_corners=False)\n",
        "        merge = torch.cat([encode, up], dim=1)\n",
        "        return merge\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a model instance\n",
        "model = UNet(in_channels=1, out_channels=2)\n",
        "\n",
        "# Create a dummy input tensor\n",
        "x = torch.randn(1, 1, 64, 64, 64)\n",
        "\n",
        "# Export the model to an ONNX file\n",
        "torch.onnx.export(model, x, \"/content/drive/My Drive/Project 2 Neuroimage Synthesis/Team2/code/Code - Project 2 Data/unet3d.onnx\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dYl8Y3qDS78L",
        "outputId": "166e2738-929d-4c1c-a40d-0411ce266ab3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============= Diagnostic Run torch.onnx.export version 2.0.0+cu118 =============\n",
            "verbose: False, log level: Level.ERROR\n",
            "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Instantiate model\n",
        "model = UNet(2, 2).to(device)\n",
        "\n",
        "# Define loss function and optimizer\n",
        "criterion = nn.MSELoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)"
      ],
      "metadata": {
        "id": "n_dAuVPj3ZEZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KPHeeXiPt9dZ"
      },
      "source": [
        "## Train the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7u5KlE7Ydc8P"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Instantiate model\n",
        "model = UNet(2, 2).to(device)\n",
        "\n",
        "# Define loss function and optimizer\n",
        "criterion = nn.MSELoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
        "\n",
        "num_epochs = 100\n",
        "train_losses = []\n",
        "val_losses = []\n",
        "patience = 10\n",
        "early_stopping_counter = 0\n",
        "# model_save_path = \"model.pt\"\n",
        "# model_save_path = r\"G:\\My Drive\\Project 2 Neuroimage Synthesis\\model.pt\"\n",
        "# model_save_path = \"/content/drive/My Drive/Project 2 Neuroimage Synthesis/model.pt\"\n",
        "model_save_path = \"/content/drive/My Drive/Project 2 Neuroimage Synthesis/Team2/code/Code - Project 2 Data/unet.pt\"\n",
        "best_validation_loss = float('inf')\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DePsGVsYt_PE",
        "outputId": "354a9fe1-049e-426d-f4c5-252b328fded4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/100, Train Loss: 10.3701, Val Loss: 1.9945, Epoch Duration: 1220.2573 seconds, Total Time: 1220.2578 seconds\n",
            "Epoch 2/100, Train Loss: 1.3925, Val Loss: 0.9439, Epoch Duration: 129.4910 seconds, Total Time: 1349.8435 seconds\n",
            "Epoch 3/100, Train Loss: 0.8151, Val Loss: 0.6636, Epoch Duration: 129.4560 seconds, Total Time: 1479.4026 seconds\n",
            "Epoch 4/100, Train Loss: 0.4843, Val Loss: 0.0232, Epoch Duration: 129.2185 seconds, Total Time: 1608.7228 seconds\n",
            "Epoch 5/100, Train Loss: 0.0168, Val Loss: 0.0159, Epoch Duration: 130.1015 seconds, Total Time: 1738.9285 seconds\n",
            "Epoch 6/100, Train Loss: 0.0154, Val Loss: 0.0153, Epoch Duration: 129.4336 seconds, Total Time: 1868.4747 seconds\n",
            "Epoch 7/100, Train Loss: 0.0149, Val Loss: 0.0149, Epoch Duration: 129.6653 seconds, Total Time: 1998.2429 seconds\n",
            "Epoch 8/100, Train Loss: 0.0146, Val Loss: 0.0146, Epoch Duration: 129.3313 seconds, Total Time: 2127.6777 seconds\n",
            "Epoch 9/100, Train Loss: 0.0143, Val Loss: 0.0143, Epoch Duration: 129.9111 seconds, Total Time: 2257.7013 seconds\n",
            "Epoch 10/100, Train Loss: 0.0140, Val Loss: 0.0140, Epoch Duration: 129.6334 seconds, Total Time: 2387.4411 seconds\n",
            "Epoch 11/100, Train Loss: 0.0137, Val Loss: 0.0137, Epoch Duration: 129.6644 seconds, Total Time: 2517.2185 seconds\n",
            "Epoch 12/100, Train Loss: 0.0134, Val Loss: 0.0134, Epoch Duration: 129.8185 seconds, Total Time: 2647.1474 seconds\n",
            "Epoch 13/100, Train Loss: 0.0131, Val Loss: 0.0131, Epoch Duration: 129.7310 seconds, Total Time: 2776.9861 seconds\n",
            "Epoch 14/100, Train Loss: 0.0128, Val Loss: 0.0128, Epoch Duration: 129.5855 seconds, Total Time: 2906.6798 seconds\n",
            "Epoch 15/100, Train Loss: 0.0125, Val Loss: 0.0125, Epoch Duration: 129.7747 seconds, Total Time: 3036.5615 seconds\n",
            "Epoch 16/100, Train Loss: 0.0123, Val Loss: 0.0122, Epoch Duration: 129.8790 seconds, Total Time: 3166.5640 seconds\n",
            "Epoch 17/100, Train Loss: 0.0120, Val Loss: 0.0120, Epoch Duration: 130.1974 seconds, Total Time: 3296.8787 seconds\n",
            "Epoch 18/100, Train Loss: 0.0117, Val Loss: 0.0117, Epoch Duration: 129.7782 seconds, Total Time: 3426.7664 seconds\n",
            "Epoch 19/100, Train Loss: 0.0115, Val Loss: 0.0114, Epoch Duration: 130.3288 seconds, Total Time: 3557.2071 seconds\n",
            "Epoch 20/100, Train Loss: 0.0112, Val Loss: 0.0112, Epoch Duration: 129.7144 seconds, Total Time: 3687.0348 seconds\n",
            "Epoch 21/100, Train Loss: 0.0110, Val Loss: 0.0111, Epoch Duration: 130.1557 seconds, Total Time: 3817.2972 seconds\n",
            "Epoch 22/100, Train Loss: 0.0107, Val Loss: 0.0108, Epoch Duration: 130.1134 seconds, Total Time: 3947.5227 seconds\n",
            "Epoch 23/100, Train Loss: 0.0105, Val Loss: 0.0105, Epoch Duration: 129.7720 seconds, Total Time: 4077.4029 seconds\n",
            "Epoch 24/100, Train Loss: 0.0103, Val Loss: 0.0103, Epoch Duration: 130.3190 seconds, Total Time: 4207.8291 seconds\n",
            "Epoch 25/100, Train Loss: 0.0101, Val Loss: 0.0101, Epoch Duration: 130.5105 seconds, Total Time: 4338.4465 seconds\n",
            "Epoch 26/100, Train Loss: 0.0099, Val Loss: 0.0099, Epoch Duration: 130.3891 seconds, Total Time: 4468.9480 seconds\n",
            "Epoch 27/100, Train Loss: 0.0097, Val Loss: 0.0097, Epoch Duration: 130.1814 seconds, Total Time: 4599.2443 seconds\n",
            "Epoch 28/100, Train Loss: 0.0095, Val Loss: 0.0095, Epoch Duration: 128.2911 seconds, Total Time: 4727.6379 seconds\n",
            "Epoch 29/100, Train Loss: 0.0094, Val Loss: 0.0095, Epoch Duration: 127.3909 seconds, Total Time: 4855.1355 seconds\n",
            "Epoch 30/100, Train Loss: 0.0093, Val Loss: 0.0092, Epoch Duration: 130.6968 seconds, Total Time: 4985.9453 seconds\n",
            "Epoch 31/100, Train Loss: 0.0092, Val Loss: 0.0091, Epoch Duration: 130.5248 seconds, Total Time: 5116.5755 seconds\n",
            "Epoch 32/100, Train Loss: 0.0090, Val Loss: 0.0090, Epoch Duration: 130.2976 seconds, Total Time: 5246.9784 seconds\n",
            "Epoch 33/100, Train Loss: 0.0089, Val Loss: 0.0090, Epoch Duration: 130.5997 seconds, Total Time: 5377.6849 seconds\n",
            "Epoch 34/100, Train Loss: 0.0088, Val Loss: 0.0087, Epoch Duration: 130.1569 seconds, Total Time: 5507.9474 seconds\n",
            "Epoch 35/100, Train Loss: 0.0088, Val Loss: 0.0087, Epoch Duration: 130.5420 seconds, Total Time: 5638.5981 seconds\n",
            "Epoch 36/100, Train Loss: 0.0087, Val Loss: 0.0086, Epoch Duration: 130.5418 seconds, Total Time: 5769.1401 seconds\n",
            "Epoch 37/100, Train Loss: 0.0085, Val Loss: 0.0085, Epoch Duration: 130.4151 seconds, Total Time: 5899.6624 seconds\n",
            "Epoch 38/100, Train Loss: 0.0084, Val Loss: 0.0083, Epoch Duration: 130.0391 seconds, Total Time: 6029.8092 seconds\n",
            "Epoch 39/100, Train Loss: 0.0083, Val Loss: 0.0082, Epoch Duration: 130.4061 seconds, Total Time: 6160.3240 seconds\n",
            "Epoch 40/100, Train Loss: 0.0082, Val Loss: 0.0081, Epoch Duration: 130.2916 seconds, Total Time: 6290.7233 seconds\n",
            "Epoch 41/100, Train Loss: 0.0080, Val Loss: 0.0082, Epoch Duration: 130.1177 seconds, Total Time: 6420.9495 seconds\n",
            "Epoch 42/100, Train Loss: 0.0080, Val Loss: 0.0079, Epoch Duration: 130.2835 seconds, Total Time: 6551.2333 seconds\n",
            "Epoch 43/100, Train Loss: 0.0078, Val Loss: 0.0078, Epoch Duration: 130.4183 seconds, Total Time: 6681.7669 seconds\n",
            "Epoch 44/100, Train Loss: 0.0077, Val Loss: 0.0079, Epoch Duration: 129.9600 seconds, Total Time: 6811.8362 seconds\n",
            "Epoch 45/100, Train Loss: 0.0077, Val Loss: 0.0077, Epoch Duration: 130.2257 seconds, Total Time: 6942.0622 seconds\n",
            "Epoch 46/100, Train Loss: 0.0076, Val Loss: 0.0075, Epoch Duration: 130.0550 seconds, Total Time: 7072.2295 seconds\n",
            "Epoch 47/100, Train Loss: 0.0074, Val Loss: 0.0074, Epoch Duration: 126.6128 seconds, Total Time: 7198.9509 seconds\n",
            "Epoch 48/100, Train Loss: 0.0073, Val Loss: 0.0074, Epoch Duration: 130.2888 seconds, Total Time: 7329.3507 seconds\n",
            "Epoch 49/100, Train Loss: 0.0073, Val Loss: 0.0076, Epoch Duration: 130.2228 seconds, Total Time: 7459.6916 seconds\n",
            "Epoch 50/100, Train Loss: 0.0072, Val Loss: 0.0073, Epoch Duration: 129.9594 seconds, Total Time: 7589.6513 seconds\n",
            "Epoch 51/100, Train Loss: 0.0071, Val Loss: 0.0070, Epoch Duration: 129.9363 seconds, Total Time: 7719.6992 seconds\n",
            "Epoch 52/100, Train Loss: 0.0070, Val Loss: 0.0069, Epoch Duration: 130.4322 seconds, Total Time: 7850.2603 seconds\n",
            "Epoch 53/100, Train Loss: 0.0069, Val Loss: 0.0069, Epoch Duration: 130.0574 seconds, Total Time: 7980.4326 seconds\n",
            "Epoch 54/100, Train Loss: 0.0070, Val Loss: 0.0068, Epoch Duration: 130.1510 seconds, Total Time: 8110.6959 seconds\n",
            "Epoch 55/100, Train Loss: 0.0067, Val Loss: 0.0066, Epoch Duration: 130.1064 seconds, Total Time: 8240.9245 seconds\n",
            "Epoch 56/100, Train Loss: 0.0067, Val Loss: 0.0066, Epoch Duration: 130.5289 seconds, Total Time: 8371.5621 seconds\n",
            "Epoch 57/100, Train Loss: 0.0066, Val Loss: 0.0065, Epoch Duration: 130.0366 seconds, Total Time: 8501.7035 seconds\n",
            "Epoch 58/100, Train Loss: 0.0065, Val Loss: 0.0064, Epoch Duration: 130.2155 seconds, Total Time: 8632.0268 seconds\n",
            "Epoch 59/100, Train Loss: 0.0064, Val Loss: 0.0063, Epoch Duration: 130.4154 seconds, Total Time: 8762.5489 seconds\n",
            "Epoch 60/100, Train Loss: 0.0062, Val Loss: 0.0066, Epoch Duration: 130.0411 seconds, Total Time: 8892.7013 seconds\n",
            "Epoch 61/100, Train Loss: 0.0062, Val Loss: 0.0062, Epoch Duration: 130.3061 seconds, Total Time: 9023.0077 seconds\n",
            "Epoch 62/100, Train Loss: 0.0061, Val Loss: 0.0061, Epoch Duration: 129.7811 seconds, Total Time: 9152.8987 seconds\n",
            "Epoch 63/100, Train Loss: 0.0060, Val Loss: 0.0060, Epoch Duration: 130.2610 seconds, Total Time: 9283.2764 seconds\n",
            "Epoch 64/100, Train Loss: 0.0059, Val Loss: 0.0060, Epoch Duration: 131.2207 seconds, Total Time: 9414.6026 seconds\n",
            "Epoch 65/100, Train Loss: 0.0059, Val Loss: 0.0058, Epoch Duration: 130.7966 seconds, Total Time: 9545.5151 seconds\n",
            "Epoch 66/100, Train Loss: 0.0058, Val Loss: 0.0057, Epoch Duration: 130.7328 seconds, Total Time: 9676.3673 seconds\n",
            "Epoch 67/100, Train Loss: 0.0057, Val Loss: 0.0058, Epoch Duration: 130.3534 seconds, Total Time: 9806.8298 seconds\n",
            "Epoch 68/100, Train Loss: 0.0057, Val Loss: 0.0056, Epoch Duration: 130.1757 seconds, Total Time: 9937.0064 seconds\n",
            "Epoch 69/100, Train Loss: 0.0056, Val Loss: 0.0056, Epoch Duration: 130.6421 seconds, Total Time: 10067.7725 seconds\n",
            "Epoch 70/100, Train Loss: 0.0055, Val Loss: 0.0056, Epoch Duration: 130.4875 seconds, Total Time: 10198.2604 seconds\n",
            "Epoch 71/100, Train Loss: 0.0055, Val Loss: 0.0055, Epoch Duration: 130.0070 seconds, Total Time: 10328.3817 seconds\n",
            "Epoch 72/100, Train Loss: 0.0055, Val Loss: 0.0056, Epoch Duration: 130.2859 seconds, Total Time: 10458.7793 seconds\n",
            "Epoch 73/100, Train Loss: 0.0054, Val Loss: 0.0054, Epoch Duration: 130.1560 seconds, Total Time: 10588.9355 seconds\n",
            "Epoch 74/100, Train Loss: 0.0054, Val Loss: 0.0054, Epoch Duration: 129.9026 seconds, Total Time: 10718.9439 seconds\n",
            "Epoch 75/100, Train Loss: 0.0054, Val Loss: 0.0054, Epoch Duration: 130.3278 seconds, Total Time: 10849.2723 seconds\n",
            "Epoch 76/100, Train Loss: 0.0054, Val Loss: 0.0053, Epoch Duration: 129.8320 seconds, Total Time: 10979.2082 seconds\n",
            "Epoch 77/100, Train Loss: 0.0053, Val Loss: 0.0054, Epoch Duration: 130.3177 seconds, Total Time: 11109.6321 seconds\n",
            "Epoch 78/100, Train Loss: 0.0052, Val Loss: 0.0052, Epoch Duration: 129.9251 seconds, Total Time: 11239.5583 seconds\n",
            "Epoch 79/100, Train Loss: 0.0052, Val Loss: 0.0052, Epoch Duration: 130.2505 seconds, Total Time: 11369.9214 seconds\n",
            "Epoch 80/100, Train Loss: 0.0052, Val Loss: 0.0053, Epoch Duration: 130.5388 seconds, Total Time: 11500.5699 seconds\n",
            "Epoch 81/100, Train Loss: 0.0051, Val Loss: 0.0052, Epoch Duration: 130.8653 seconds, Total Time: 11631.4356 seconds\n",
            "Epoch 82/100, Train Loss: 0.0051, Val Loss: 0.0052, Epoch Duration: 130.5566 seconds, Total Time: 11762.0971 seconds\n",
            "Epoch 83/100, Train Loss: 0.0051, Val Loss: 0.0051, Epoch Duration: 130.3197 seconds, Total Time: 11892.5246 seconds\n",
            "Epoch 84/100, Train Loss: 0.0051, Val Loss: 0.0052, Epoch Duration: 130.7968 seconds, Total Time: 12023.4417 seconds\n",
            "Epoch 85/100, Train Loss: 0.0050, Val Loss: 0.0050, Epoch Duration: 130.2317 seconds, Total Time: 12153.6742 seconds\n",
            "Epoch 86/100, Train Loss: 0.0050, Val Loss: 0.0051, Epoch Duration: 130.5767 seconds, Total Time: 12284.3607 seconds\n",
            "Epoch 87/100, Train Loss: 0.0051, Val Loss: 0.0050, Epoch Duration: 130.4009 seconds, Total Time: 12414.7620 seconds\n",
            "Epoch 88/100, Train Loss: 0.0049, Val Loss: 0.0049, Epoch Duration: 130.5009 seconds, Total Time: 12545.3763 seconds\n",
            "Epoch 89/100, Train Loss: 0.0048, Val Loss: 0.0050, Epoch Duration: 130.5803 seconds, Total Time: 12676.0663 seconds\n",
            "Epoch 90/100, Train Loss: 0.0048, Val Loss: 0.0048, Epoch Duration: 130.6060 seconds, Total Time: 12806.6726 seconds\n",
            "Epoch 91/100, Train Loss: 0.0048, Val Loss: 0.0048, Epoch Duration: 130.3802 seconds, Total Time: 12937.1637 seconds\n",
            "Epoch 92/100, Train Loss: 0.0047, Val Loss: 0.0047, Epoch Duration: 130.3569 seconds, Total Time: 13067.6329 seconds\n",
            "Epoch 93/100, Train Loss: 0.0046, Val Loss: 0.0047, Epoch Duration: 130.9098 seconds, Total Time: 13198.6615 seconds\n",
            "Epoch 94/100, Train Loss: 0.0046, Val Loss: 0.0047, Epoch Duration: 130.2773 seconds, Total Time: 13329.0557 seconds\n",
            "Epoch 95/100, Train Loss: 0.0045, Val Loss: 0.0045, Epoch Duration: 130.2561 seconds, Total Time: 13459.3127 seconds\n",
            "Epoch 96/100, Train Loss: 0.0044, Val Loss: 0.0044, Epoch Duration: 130.2236 seconds, Total Time: 13589.6435 seconds\n",
            "Epoch 97/100, Train Loss: 0.0044, Val Loss: 0.0045, Epoch Duration: 130.2607 seconds, Total Time: 13720.0106 seconds\n",
            "Epoch 98/100, Train Loss: 0.0043, Val Loss: 0.0043, Epoch Duration: 130.1291 seconds, Total Time: 13850.1402 seconds\n",
            "Epoch 99/100, Train Loss: 0.0041, Val Loss: 0.0041, Epoch Duration: 129.8278 seconds, Total Time: 13980.0826 seconds\n",
            "Epoch 100/100, Train Loss: 0.0041, Val Loss: 0.0042, Epoch Duration: 130.6320 seconds, Total Time: 14110.8221 seconds\n"
          ]
        }
      ],
      "source": [
        "train_losses = []\n",
        "val_losses = []\n",
        "\n",
        "best_validation_loss = float('inf')\n",
        "early_stopping_counter = 0\n",
        "patience = 10\n",
        "\n",
        "start_time = time.time()\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    torch.autograd.set_detect_anomaly(True)\n",
        "    model.train()\n",
        "\n",
        "    epoch_start_time = time.time()\n",
        "\n",
        "    train_loss = 0.0\n",
        "    for i, (input_imgs, target_imgs) in enumerate(train_loader):\n",
        "        input_imgs = input_imgs.to(device)\n",
        "        target_imgs = F.interpolate(target_imgs, size=[182, 218, 182], mode='trilinear', align_corners=False)\n",
        "        #target_imgs = F.interpolate(target_imgs, size=(182, 218, 182), mode='trilinear', align_corners=False)\n",
        "        target_imgs = target_imgs.to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(input_imgs)\n",
        "        loss = criterion(outputs, target_imgs)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        train_loss += loss.item() * input_imgs.size(0)\n",
        "\n",
        "    train_loss /= len(train_loader.dataset)\n",
        "    train_losses.append(train_loss)\n",
        "\n",
        "    model.eval()\n",
        "    val_loss = 0.0\n",
        "    with torch.no_grad():\n",
        "        for inputs, targets in val_loader:\n",
        "            inputs = inputs.to(device)\n",
        "            targets = F.interpolate(targets, size=[182, 218, 182], mode='trilinear', align_corners=False)\n",
        "            targets = targets.to(device)\n",
        "\n",
        "            outputs = model(inputs)\n",
        "            loss = criterion(outputs, targets)\n",
        "            val_loss += loss.item() * inputs.size(0)\n",
        "\n",
        "\n",
        "    val_loss /= len(val_loader.dataset)\n",
        "    val_losses.append(val_loss)\n",
        "\n",
        "    epoch_end_time = time.time()\n",
        "    epoch_duration = epoch_end_time - epoch_start_time\n",
        "    total_training_time = time.time() - start_time\n",
        "\n",
        "    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Epoch Duration: {epoch_duration:.4f} seconds, Total Time: {total_training_time:.4f} seconds')\n",
        "\n",
        "    if val_loss < best_validation_loss:\n",
        "        best_validation_loss = val_loss\n",
        "        early_stopping_counter = 0\n",
        "        torch.save(model.state_dict(), model_save_path)\n",
        "    else:\n",
        "        early_stopping_counter += 1\n",
        "        if early_stopping_counter >= patience:\n",
        "            print(\"Early stopping triggered.\")\n",
        "            break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRS-P8cIrAEW"
      },
      "source": [
        "### save model state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TYhTh951YiK5"
      },
      "outputs": [],
      "source": [
        "torch.save({\n",
        "            'epoch': epoch,\n",
        "            'model_state_dict': model.state_dict(),\n",
        "            'optimizer_state_dict': optimizer.state_dict(),\n",
        "            'loss': loss,\n",
        "            }, \"unet.pt\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eMcgQG-dzPZV"
      },
      "source": [
        "## Train/Val Loss Plot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 482
        },
        "id": "OBeNThFdpoUd",
        "outputId": "50371cad-787f-49b0-c9f9-ab0729ff3b66"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1000x500 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAA0kAAAHACAYAAAB3WSN5AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA8uElEQVR4nO3deXiU1d3/8c8kM5ksZGPLIgGCjbIFjGwFXPABWVRksyqNFdQ+/KwRRatVVJCliqJSHpfiUkVtBVwqiChaRAVUZEexILggRFmCIglrQjL374+QYQYChCTMOQnv13XNlZn7nuU7MzeQD+f+nuNyHMcRAAAAAECSFGa6AAAAAACwCSEJAAAAAAIQkgAAAAAgACEJAAAAAAIQkgAAAAAgACEJAAAAAAIQkgAAAAAgACEJAAAAAAK4TRdwqvl8Pm3ZskWxsbFyuVymywEAAABgiOM42r17t1JTUxUWduzxolofkrZs2aK0tDTTZQAAAACwRG5urho1anTM/bU+JMXGxkoq/SDi4uIMVwMAAADAlIKCAqWlpfkzwrHU+pBUdopdXFwcIQkAAADACdtwmLgBAAAAAAIQkgAAAAAgACEJAAAAAALU+p4kAAAA4Hgcx1FxcbFKSkpMl4IqCg8Pl9vtrvLSP4QkAAAAnLaKioq0detW7du3z3QpqCbR0dFKSUlRREREpZ+DkAQAAIDTks/n08aNGxUeHq7U1FRFRERUeQQC5jiOo6KiIu3YsUMbN25URkbGcReMPR5CEgAAAE5LRUVF8vl8SktLU3R0tOlyUA2ioqLk8Xi0adMmFRUVKTIyslLPw8QNAAAAOK1VdrQBdqqO75MjAgAAAAACEJIAAAAAIAAhCQAAADjNNW3aVJMnTzZdhjUISQAAAEAN4XK5jnsZM2ZMpZ532bJlGjZsWJVq69atm0aMGFGl57AFs9sBAAAANcTWrVv911999VWNHj1a69ev92+rU6eO/7rjOCopKZHbfeJf+Rs0aFC9hdZwjCSF0FXPLNbFkxZo0y97TZcCAACAIziOo31FxUYujuNUqMbk5GT/JT4+Xi6Xy3/766+/VmxsrObOnat27drJ6/Xqk08+0Xfffad+/fopKSlJderUUYcOHfTBBx8EPe+Rp9u5XC794x//0IABAxQdHa2MjAzNnj27Sp/vv//9b7Vq1Uper1dNmzbVY489FrT/73//uzIyMhQZGamkpCRdccUV/n1vvPGGMjMzFRUVpXr16qlHjx7au/fU/U7NSFIIfbdjj37eU6R9RSWmSwEAAMAR9h8sUcvR7xt57bXjeik6onp+Nb/77rv16KOPqlmzZkpMTFRubq4uueQSPfDAA/J6vXr55ZfVt29frV+/Xo0bNz7m84wdO1YTJ07UI488oieeeELZ2dnatGmT6tate9I1rVixQldeeaXGjBmjq666Sp999pluuukm1atXT0OHDtXy5ct1yy236J///Ke6dOminTt3atGiRZJKR88GDx6siRMnasCAAdq9e7cWLVpU4WBZGYSkEPKElw7cHSzxGa4EAAAAtdW4ceN08cUX+2/XrVtXbdu29d8eP368Zs6cqdmzZ+vmm28+5vMMHTpUgwcPliQ9+OCDevzxx7V06VL17t37pGuaNGmSunfvrlGjRkmSzjrrLK1du1aPPPKIhg4dqs2bNysmJkaXXXaZYmNj1aRJE2VlZUkqDUnFxcUaOHCgmjRpIknKzMw86RpOBiEphAhJAAAA9oryhGvtuF7GXru6tG/fPuj2nj17NGbMGL3zzjv+wLF//35t3rz5uM/Tpk0b//WYmBjFxcUpLy+vUjWtW7dO/fr1C9rWtWtXTZ48WSUlJbr44ovVpEkTNWvWTL1791bv3r39p/q1bdtW3bt3V2Zmpnr16qWePXvqiiuuUGJiYqVqqQh6kkLIE+6SJBUVn7qhQQAAAFSOy+VSdITbyMXlclXb+4iJiQm6fccdd2jmzJl68MEHtWjRIq1evVqZmZkqKio67vN4PJ6jPh+f79T8Z39sbKxWrlyp6dOnKyUlRaNHj1bbtm21a9cuhYeHa968eZo7d65atmypJ554QmeffbY2btx4SmqRCEkhxUgSAAAAQu3TTz/V0KFDNWDAAGVmZio5OVk//PBDSGto0aKFPv3006PqOuussxQeXjqK5na71aNHD02cOFFffvmlfvjhB3344YeSSgNa165dNXbsWK1atUoRERGaOXPmKauX0+1CKMJNSAIAAEBoZWRk6M0331Tfvn3lcrk0atSoUzYitGPHDq1evTpoW0pKiv785z+rQ4cOGj9+vK666iotXrxYTz75pP7+979LkubMmaPvv/9eF1xwgRITE/Xuu+/K5/Pp7LPP1pIlSzR//nz17NlTDRs21JIlS7Rjxw61aNHilLwHiZAUUowkAQAAINQmTZqk66+/Xl26dFH9+vV11113qaCg4JS81rRp0zRt2rSgbePHj9d9992n1157TaNHj9b48eOVkpKicePGaejQoZKkhIQEvfnmmxozZowOHDigjIwMTZ8+Xa1atdK6deu0cOFCTZ48WQUFBWrSpIkee+wx9enT55S8B0lyOady7jwLFBQUKD4+Xvn5+YqLizNay9XPLtbn3+/U44OzdHnbVKO1AAAAnO4OHDigjRs3Kj09XZGRkabLQTU53vda0WxgtCdp4cKF6tu3r1JTU+VyuTRr1qyg/Y7jaPTo0UpJSVFUVJR69Oihb775xkyx1cA/klTMSBIAAABgK6Mhae/evWrbtq2eeuqpcvdPnDhRjz/+uJ5++mktWbJEMTEx6tWrlw4cOBDiSqtHBKfbAQAAANYz2pPUp0+fY55L6DiOJk+erPvuu88/p/rLL7+spKQkzZo1S1dffXUoS60W9CQBAAAA9rN2CvCNGzdq27Zt6tGjh39bfHy8OnXqpMWLFxusrPI8h2a3Kyqp1W1gAAAAQI1m7ex227ZtkyQlJSUFbU9KSvLvK09hYaEKCwv9t0/VzB2VUbaYLCNJAAAAgL2sHUmqrAkTJig+Pt5/SUtLM12SXwQTNwAAAADWszYkJScnS5K2b98etH379u3+feUZOXKk8vPz/Zfc3NxTWufJoCcJAAAAsJ+1ISk9PV3JycmaP3++f1tBQYGWLFmizp07H/NxXq9XcXFxQRdblIUkepIAAAAAexntSdqzZ4++/fZb/+2NGzdq9erVqlu3rho3bqwRI0bor3/9qzIyMpSenq5Ro0YpNTVV/fv3N1d0FXjc9CQBAADAvG7duumcc87R5MmTTZdiJaMjScuXL1dWVpaysrIkSbfffruysrI0evRoSdJf/vIXDR8+XMOGDVOHDh20Z88evffeezV2RWTWSQIAAEBV9O3bV7179y5336JFi+RyufTll19W+XVefPFFJSQkVPl5aiqjI0ndunWT4xz71DOXy6Vx48Zp3LhxIazq1CEkAQAAoCpuuOEGDRo0SD/++KMaNWoUtG/q1Klq37692rRpY6i62sPanqTayL9OUjE9SQAAADh5l112mRo0aKAXX3wxaPuePXv0+uuv64YbbtAvv/yiwYMH64wzzlB0dLQyMzM1ffr0aq1j8+bN6tevn+rUqaO4uDhdeeWVQROuffHFF7rooosUGxuruLg4tWvXTsuXL5ckbdq0SX379lViYqJiYmLUqlUrvfvuu9VaX1VZu05SbcTsdgAAABZzHOngPjOv7YmWXK4T3s3tduvaa6/Viy++qHvvvVeuQ495/fXXVVJSosGDB2vPnj1q166d7rrrLsXFxemdd97RH/7wB5155pnq2LFjlUv1+Xz+gLRgwQIVFxcrJydHV111lT7++GNJUnZ2trKysjRlyhSFh4dr9erV8ng8kqScnBwVFRVp4cKFiomJ0dq1a1WnTp0q11WdCEkhFMFisgAAAPY6uE96MNXMa9+zRYqIqdBdr7/+ej3yyCNasGCBunXrJqn0VLtBgwb51wq94447/PcfPny43n//fb322mvVEpLmz5+vNWvWaOPGjf41SV9++WW1atVKy5YtU4cOHbR582bdeeedat68uSQpIyPD//jNmzdr0KBByszMlCQ1a9asyjVVN063CyFGkgAAAFBVzZs3V5cuXfTCCy9Ikr799lstWrRIN9xwgySppKRE48ePV2ZmpurWras6dero/fff1+bNm6vl9detW6e0tDR/QJKkli1bKiEhQevWrZNUOiHbH//4R/Xo0UMPPfSQvvvuO/99b7nlFv31r39V165ddf/991fLRBPVjZGkEGKdJAAAAIt5oktHdEy99km44YYbNHz4cD311FOaOnWqzjzzTF144YWSpEceeUT/93//p8mTJyszM1MxMTEaMWKEioqKTkXl5RozZox+//vf65133tHcuXN1//33a8aMGRowYID++Mc/qlevXnrnnXf0n//8RxMmTNBjjz2m4cOHh6y+E2EkKYTKJm44WMxIEgAAgHVcrtJT3kxcKtCPFOjKK69UWFiYpk2bppdfflnXX3+9vz/p008/Vb9+/XTNNdeobdu2atasmTZs2FBtH1OLFi2Um5ur3Nxc/7a1a9dq165datmypX/bWWedpdtuu03/+c9/NHDgQE2dOtW/Ly0tTTfeeKPefPNN/fnPf9Zzzz1XbfVVB0aSQoieJAAAAFSHOnXq6KqrrtLIkSNVUFCgoUOH+vdlZGTojTfe0GeffabExERNmjRJ27dvDwowFVFSUqLVq1cHbfN6verRo4cyMzOVnZ2tyZMnq7i4WDfddJMuvPBCtW/fXvv379edd96pK664Qunp6frxxx+1bNkyDRo0SJI0YsQI9enTR2eddZZ+/fVXffTRR2rRokVVP5JqRUgKIXqSAAAAUF1uuOEGPf/887rkkkuUmnp4won77rtP33//vXr16qXo6GgNGzZM/fv3V35+/kk9/549e5SVlRW07cwzz9S3336rt956S8OHD9cFF1ygsLAw9e7dW0888YQkKTw8XL/88ouuvfZabd++XfXr19fAgQM1duxYSaXhKycnRz/++KPi4uLUu3dv/e1vf6vip1G9XM7xVnOtBQoKChQfH6/8/HzFxcUZrWXhhh269oWlapESp7m3nm+0FgAAgNPdgQMHtHHjRqWnpysyMtJ0Oagmx/teK5oN6EkKIUaSAAAAAPsRkkIowk1PEgAAAGA7QlII+UeSmN0OAAAAsBYhKYRYJwkAAACwHyEphOhJAgAAAOxHSAqhCEISAACAdWr5ZM+nner4PglJIeRh4gYAAABreDweSdK+ffsMV4LqVPZ9ln2/lcFisiF0+HQ7R47jyOVyGa4IAADg9BUeHq6EhATl5eVJkqKjo/n9rAZzHEf79u1TXl6eEhISFB4eXunnIiSFUFlIkkqDUtmU4AAAADAjOTlZkvxBCTVfQkKC/3utLEJSCEUEhSSfItyc7QgAAGCSy+VSSkqKGjZsqIMHD5ouB1Xk8XiqNIJUhpAUQp7wwyNH9CUBAADYIzw8vFp+uUbtwFBGCIWHuVR2mmsRIQkAAACwEiEphFwuV9DkDQAAAADsQ0gKMf9aScWMJAEAAAA2IiSFWFlfEj1JAAAAgJ0ISSFWdrodPUkAAACAnQhJIUZPEgAAAGA3QlKIla2NxOl2AAAAgJ0ISSHm70li4gYAAADASoSkEKMnCQAAALAbISnE6EkCAAAA7EZICjH/OkmMJAEAAABWIiSFmMfNOkkAAACAzQhJIebvSWLiBgAAAMBKhKQQoycJAAAAsBshKcToSQIAAADsRkgKMf86SYQkAAAAwEqEpBBjnSQAAADAboSkEPO4D51uV0xPEgAAAGAjQlKI0ZMEAAAA2I2QFGL0JAEAAAB2IySFGD1JAAAAgN0ISSHm4XQ7AAAAwGqEpBCLYOIGAAAAwGqEpBCjJwkAAACwGyEpxOhJAgAAAOxGSAoxepIAAAAAuxGSQuzwOkn0JAEAAAA2IiSFmMdNTxIAAABgM0JSiPl7kooJSQAAAICNCEkhRk8SAAAAYDdCUojRkwQAAADYjZAUYowkAQAAAHYjJIVY2WKyrJMEAAAA2ImQFGIeNyNJAAAAgM0ISSHm70kqpicJAAAAsBEhKcToSQIAAADsRkgKMXqSAAAAALtZHZJKSko0atQopaenKyoqSmeeeabGjx8vx6m5p6oxkgQAAADYzW26gON5+OGHNWXKFL300ktq1aqVli9fruuuu07x8fG65ZZbTJdXKRFu1kkCAAAAbGZ1SPrss8/Ur18/XXrppZKkpk2bavr06Vq6dKnhyiqvbCSpxOeoxOcoPMxluCIAAAAAgaw+3a5Lly6aP3++NmzYIEn64osv9Mknn6hPnz7HfExhYaEKCgqCLjYp60mSOOUOAAAAsJHVI0l33323CgoK1Lx5c4WHh6ukpEQPPPCAsrOzj/mYCRMmaOzYsSGs8uSUjSRJpSEp0hNusBoAAAAAR7J6JOm1117TK6+8omnTpmnlypV66aWX9Oijj+qll1465mNGjhyp/Px8/yU3NzeEFZ9YcEiiLwkAAACwjdUjSXfeeafuvvtuXX311ZKkzMxMbdq0SRMmTNCQIUPKfYzX65XX6w1lmSclPMyl8DCXSnwOp9sBAAAAFrJ6JGnfvn0KCwsuMTw8XD5fzQ4X/rWSimv2+wAAAABqI6tHkvr27asHHnhAjRs3VqtWrbRq1SpNmjRJ119/venSqsQTHqYDB32MJAEAAAAWsjokPfHEExo1apRuuukm5eXlKTU1Vf/v//0/jR492nRpVRIRzlpJAAAAgK2sDkmxsbGaPHmyJk+ebLqUauXxhyRGkgAAAADbWN2TVFt53Id6kghJAAAAgHUISQb4R5KYuAEAAACwDiHJAHqSAAAAAHsRkgygJwkAAACwFyHJAP86SYQkAAAAwDqEJAPKRpJYTBYAAACwDyHJgAg3p9sBAAAAtiIkGRBBTxIAAABgLUKSAf7T7ZjdDgAAALAOIckAj5t1kgAAAABbEZIMKJvdjtPtAAAAAPsQkgygJwkAAACwFyHJAHqSAAAAAHsRkgzwMJIEAAAAWIuQZIDHfagniYkbAAAAAOsQkgygJwkAAACwFyHJAHqSAAAAAHsRkgygJwkAAACwFyHJANZJAgAAAOxFSDIgws1IEgAAAGArQpIB/p6kYnqSAAAAANsQkgygJwkAAACwFyHJAHqSAAAAAHsRkgxgnSQAAADAXoQkA1gnCQAAALAXIckAT9nsdsWMJAEAAAC2ISQZQE8SAAAAYC9CkgH0JAEAAAD2IiQZcHgKcHqSAAAAANsQkgw4PHEDI0kAAACAbQhJBkS46UkCAAAAbEVIMsB/uh2z2wEAAADWISQZQE8SAAAAYC9CkgGBPUmOQ1ACAAAAbEJIMqBsCnBJKvYRkgAAAACbEJIM8ByauEFi8gYAAADANoQkAzwBI0kHixlJAgAAAGxCSDLAHXZ4JIm1kgAAAAC7EJIMcLlc/r4kTrcDAAAA7EJIMsQTzoKyAAAAgI0ISYZ43IwkAQAAADYiJBniXyuJiRsAAAAAqxCSDKEnCQAAALATIckQepIAAAAAOxGSDPGfbkdIAgAAAKxCSDLE4z/djp4kAAAAwCaEJEP8s9sVM5IEAAAA2ISQZEgEPUkAAACAlQhJhtCTBAAAANiJkGQIPUkAAACAnQhJhnhYJwkAAACwEiHJkAg3PUkAAACAjQhJhvh7kpjdDgAAALAKIckQepIAAAAAOxGSDKEnCQAAALATIckQ1kkCAAAA7ERIMoR1kgAAAAA7WR+SfvrpJ11zzTWqV6+eoqKilJmZqeXLl5suq8o87kOn2xXTkwQAAADYxG26gOP59ddf1bVrV1100UWaO3euGjRooG+++UaJiYmmS6syepIAAAAAO1kdkh5++GGlpaVp6tSp/m3p6ekGK6o+9CQBAAAAdrL6dLvZs2erffv2+t3vfqeGDRsqKytLzz333HEfU1hYqIKCgqCLjehJAgAAAOxkdUj6/vvvNWXKFGVkZOj999/Xn/70J91yyy166aWXjvmYCRMmKD4+3n9JS0sLYcUVxzpJAAAAgJ2sDkk+n0/nnnuuHnzwQWVlZWnYsGH63//9Xz399NPHfMzIkSOVn5/vv+Tm5oaw4oo7PHEDI0kAAACATawOSSkpKWrZsmXQthYtWmjz5s3HfIzX61VcXFzQxUb0JAEAAAB2sjokde3aVevXrw/atmHDBjVp0sRQRdWHniQAAADATlaHpNtuu02ff/65HnzwQX377beaNm2ann32WeXk5JgurcqYAhwAAACwk9UhqUOHDpo5c6amT5+u1q1ba/z48Zo8ebKys7NNl1ZlTNwAAAAA2MnqdZIk6bLLLtNll11muoxqF+GmJwkAAACwkdUjSbWZvyeJ2e0AAAAAqxCSDKEnCQAAALATIckQepIAAAAAOxGSDIlgJAkAAACwEiHJEA8TNwAAAABWIiQZwsQNAAAAgJ0ISYZE0JMEAAAAWImQZAiz2wEAAAB2IiQZ4gkv7Ukq9jny+RhNAgAAAGxRqZCUm5urH3/80X976dKlGjFihJ599tlqK6y287gPf/QHfYwmAQAAALaoVEj6/e9/r48++kiStG3bNl188cVaunSp7r33Xo0bN65aC6ytynqSJPqSAAAAAJtUKiR99dVX6tixoyTptddeU+vWrfXZZ5/plVde0Ysvvlid9dVansCQxAx3AAAAgDUqFZIOHjwor9crSfrggw90+eWXS5KaN2+urVu3Vl91tVh4mEthpW1JTN4AAAAAWKRSIalVq1Z6+umntWjRIs2bN0+9e/eWJG3ZskX16tWr1gJrM/9aSYQkAAAAwBqVCkkPP/ywnnnmGXXr1k2DBw9W27ZtJUmzZ8/2n4aHE4tws1YSAAAAYBt3ZR7UrVs3/fzzzyooKFBiYqJ/+7BhwxQdHV1txdV2EayVBAAAAFinUiNJ+/fvV2FhoT8gbdq0SZMnT9b69evVsGHDai2wNvOfbsfEDQAAAIA1KhWS+vXrp5dfflmStGvXLnXq1EmPPfaY+vfvrylTplRrgbWZx106cwMjSQAAAIA9KhWSVq5cqfPPP1+S9MYbbygpKUmbNm3Syy+/rMcff7xaC6zNPOH0JAEAAAC2qVRI2rdvn2JjYyVJ//nPfzRw4ECFhYXpt7/9rTZt2lStBdZm9CQBAAAA9qlUSPrNb36jWbNmKTc3V++//7569uwpScrLy1NcXFy1FlibMQU4AAAAYJ9KhaTRo0frjjvuUNOmTdWxY0d17txZUumoUlZWVrUWWJt5wg/1JDFxAwAAAGCNSk0BfsUVV+i8887T1q1b/WskSVL37t01YMCAaiuutqMnCQAAALBPpUKSJCUnJys5OVk//vijJKlRo0YsJHuSDi8my0gSAAAAYItKnW7n8/k0btw4xcfHq0mTJmrSpIkSEhI0fvx4+Xz8wl9R9CQBAAAA9qnUSNK9996r559/Xg899JC6du0qSfrkk080ZswYHThwQA888EC1Fllb+XuSCEkAAACANSoVkl566SX94x//0OWXX+7f1qZNG51xxhm66aabCEkV5O9JYuIGAAAAwBqVOt1u586dat68+VHbmzdvrp07d1a5qNNFBBM3AAAAANapVEhq27atnnzyyaO2P/nkk2rTpk2Vizpd0JMEAAAA2KdSp9tNnDhRl156qT744AP/GkmLFy9Wbm6u3n333WotsDbzuOlJAgAAAGxTqZGkCy+8UBs2bNCAAQO0a9cu7dq1SwMHDtR///tf/fOf/6zuGmutw+skEZIAAAAAW1R6naTU1NSjJmj44osv9Pzzz+vZZ5+tcmGnA3qSAAAAAPtUaiQJ1cPfk8TsdgAAAIA1CEkGcbodAAAAYB9CkkFM3AAAAADY56R6kgYOHHjc/bt27apKLacdepIAAAAA+5xUSIqPjz/h/muvvbZKBZ1OWCcJAAAAsM9JhaSpU6eeqjpOS0zcAAAAANiHniSDPOH0JAEAAAC2ISQZFOFmdjsAAADANoQkg/xTgBczcQMAAABgC0KSQUzcAAAAANiHkGQQPUkAAACAfQhJBh1eJ4mQBAAAANiCkGSQx81isgAAAIBtCEkGsU4SAAAAYB9CkkH0JAEAAAD2ISQZRE8SAAAAYB9CkkH+dZLoSQIAAACsQUgyqGziBtZJAgAAAOxBSDIosCfJcRhNAgAAAGxASDKorCfJcaQSHyEJAAAAsAEhyaCyniSJviQAAADAFoQkgwJDEn1JAAAAgB0ISQaV9SRJTAMOAAAA2IKQZJDL5WJBWQAAAMAyNSokPfTQQ3K5XBoxYoTpUqqNf62kYnqSAAAAABvUmJC0bNkyPfPMM2rTpo3pUqpVWUiiJwkAAACwQ40ISXv27FF2draee+45JSYmmi6nWvlHkghJAAAAgBVqREjKycnRpZdeqh49epzwvoWFhSooKAi62CyCniQAAADAKm7TBZzIjBkztHLlSi1btqxC958wYYLGjh17iquqPh43I0kAAACATaweScrNzdWtt96qV155RZGRkRV6zMiRI5Wfn++/5ObmnuIqq8bfk8TEDQAAAIAVrB5JWrFihfLy8nTuuef6t5WUlGjhwoV68sknVVhYqPDw8KDHeL1eeb3eUJdaafQkAQAAAHaxOiR1795da9asCdp23XXXqXnz5rrrrruOCkg1ET1JAAAAgF2sDkmxsbFq3bp10LaYmBjVq1fvqO01FSNJAAAAgF2s7kk6HRxeJ4meJAAAAMAGVo8klefjjz82XUK18s9uV8xIEgAAAGADRpIMoycJAAAAsAshyTB6kgAAAAC7EJIMoycJAAAAsAshyTBGkgAAAAC7EJIMi3Af6kli4gYAAADACoQkwxhJAgAAAOxCSDKMniQAAADALoQkwxhJAgAAAOxCSDKMdZIAAAAAuxCSDGMkCQAAALALIckwj/tQT1IxPUkAAACADQhJhjGSBAAAANiFkGQYPUkAAACAXQhJhjGSBAAAANiFkGQY6yQBAAAAdiEkGVY2ccPBYkaSAAAAABsQkgyjJwkAAACwCyHJMHqSAAAAALsQkgyjJwkAAACwCyHJMEaSAAAAALsQkgyLcNOTBAAAANiEkGSYfySJ2e0AAAAAKxCSDItw05MEAAAA2ISQZBg9SQAAAIBdCEmGRRCSAAAAAKsQkgxjJAkAAACwCyHJME942ex2jhyHviQAAADANEKSYR734a/gIJM3AAAAAMYRkgwr60mSOOUOAAAAsAEhyTAPIQkAAACwCiHJsPAwl8JK25JUREgCAAAAjCMkWeDwDHf0JAEAAACmEZIs4F8rqZiRJAAAAMA0QpIFyma4oycJAAAAMI+QZIGytZLoSQIAAADMIyRZgJ4kAAAAwB6EJAv4e5IYSQIAAACMIyRZwMPEDQAAAIA1CEkW8LjpSQIAAABsQUiyAD1JAAAAgD0ISRbw0JMEAAAAWIOQZAEmbgAAAADsQUiygH+dJCZuAAAAAIwjJFmAniQAAADAHoQkC3jcnG4HAAAA2IKQZAF6kgAAAAB7EJIs4O9JIiQBAAAAxhGSLODvSSqmJwkAAAAwjZBkAdZJAgAAAOxBSLJABBM3AAAAANYgJFmAniQAAADAHoQkC3C6HQAAAGAPQpIFmLgBAAAAsAchyQKskwQAAADYg5BkAXqSAAAAAHsQkizgYXY7AAAAwBqEJAscnriBniQAAADANKtD0oQJE9ShQwfFxsaqYcOG6t+/v9avX2+6rGpHTxIAAABgD6tD0oIFC5STk6PPP/9c8+bN08GDB9WzZ0/t3bvXdGnVqmwkqaiYkAQAAACY5jZdwPG89957QbdffPFFNWzYUCtWrNAFF1xgqKoqcByp4CcpvlHQ5rKJGxhJAgAAAMyzOiQdKT8/X5JUt27dY96nsLBQhYWF/tsFBQWnvK4KKS6U/tZK2rtDuuMbqU5D/67DEzfQkwQAAACYZvXpdoF8Pp9GjBihrl27qnXr1se834QJExQfH++/pKWlhbDK43B7pajE0utbVgftoicJAAAAsEeNCUk5OTn66quvNGPGjOPeb+TIkcrPz/dfcnNzQ1RhBaScU/pz6+qgzf6eJEISAAAAYFyNON3u5ptv1pw5c7Rw4UI1atTouPf1er3yer0hquwkpWZJa147aiSJniQAAADAHlaHJMdxNHz4cM2cOVMff/yx0tPTTZdUNannlP7csipos3+dpGJ6kgAAAADTrA5JOTk5mjZtmt566y3FxsZq27ZtkqT4+HhFRUUZrq4SkttIckm7t0h78vyTN0S46UkCAAAAbGF1T9KUKVOUn5+vbt26KSUlxX959dVXTZdWOd46Uv2zSq8HnHJHTxIAAABgD6tHkhynFp5+lnqO9PP60lPuzuopiZ4kAAAAwCZWjyTVSqlZpT8DZrg7PAV4LQyFAAAAQA1DSAq1smnAAyZvKDvdrsTnqMRHUAIAAABMIiSFWnKmSidv2Crt3i5J8rgPfw2ccgcAAACYRUgKNW8dqcHZpdcPnXJX1pMkMXkDAAAAYBohyYQjTrnzhAWMJBUTkgAAAACTCEkmlE3ecGga8LAwl9xhZTPc0ZMEAAAAmERIMiH1nNKfATPcecJZUBYAAACwASHJhORMyRV2aPKGbZIO9yXRkwQAAACYRUgyISJGqn9o8oZDp9xFuBlJAgAAAGxASDKl7JS7sskbyk63K6YnCQAAADCJkGRK2Qx3/mnAS78KTrcDAAAAzCIkmXLEDHdlPUmcbgcAAACYRUgypWzyhj3bpIKtzG4HAAAAWIKQZEpEtNSgeen1rasVH+WRJG38ea/BogAAAAAQkkwq60vaskoXt0ySJM1c9ZO5egAAAAAQkozyz3C3Wpefk6rwMJdWbd6l73fsMVoWAAAAcDojJJlUNnnD1tVqWMer8zPqS5JmMZoEAAAAGENIMimp9aHJG7ZLu7dqQNYZkqQ3V/0kn4/1kgAAAAATCEkmBU7esGW1erZMVh2vWz/+ul/LN/1qtjYAAADgNEVIMi3glLuoiHD1aZ0sSZq56keDRQEAAACnL0KSaQEz3EnSwHMbSZLmfLlVBw6WGCoKAAAAOH0RkkwrG0naslpyHHVKr6szEqK0+0Cx5q/LM1oaAAAAcDoiJJmW3FpyhUt786SCLQoLc6l/Vqok6c2VnHIHAAAAhBohyTRP1OHJG7auliQNyCo95W7Bhh36eU+hocIAAACA0xMhyQaBp9xJ+k3DOmrbKF7FPkdvf7HFXF0AAADAaYiQZIPUc0p/Hpq8QZJ/zaSZLCwLAAAAhBQhyQZlM9xtXS05pYvI9m2bKneYS1/+mK9v83YbKw0AAAA43RCSbOCfvGGHVFA6clSvjlfdzm4gSXpzJaNJAAAAQKgQkmzgiZKSWpVe/+xJ/+ayCRxmrfpJPp9jojIAAADgtENIssVF95T+XDJFWve2JKl7i4aKjXRrS/4Bfb7xF4PFAQAAAKcPQpItzu4jdb659PqsHOnXHxTpCddlbVIkccodAAAAECqEJJv0GCM16iAV5kuvXycVF2nguaWn3M1ds1X7i0rM1gcAAACcBghJNgn3SFdMlSITpC0rpXmj1L5JotLqRmlvUYneXPWj6QoBAACAWo+QZJuENGnA06XXlzwt17q3NbRLuiRp4nvrlbf7gMHiAAAAgNqPkGSjs/tIXYaXXn/rZg1p7qj1GXHK339QY99ea7Y2AAAAoJYjJNmq+/1So45SYb7c/75OD/c7W+FhLr3z5VbNX7fddHUAAABArUVIslW4R7riBSkqUdq6Wq2+elQ3nFd62t19s77SnsJiwwUCAAAAtRMhyWYJadKAZ0qvL31Gfz5zq9LqRmlr/gE9+v56s7UBAAAAtRQhyXZn9ZLOHSJJ8v53hh4ckClJemnxD1q5+VeTlQEAAAC1EiGpJjjn96U/18/V+elxGnjuGXIcaeS/16io2Ge2NgAAAKCWISTVBI06SrEpUmGB9N1Huu/SlqobE6H123fr2YXfma4OAAAAqFUISTVBWJjU4vLS62vfUt2YCI2+rKUk6fH53+q7HXsMFgcAAADULoSkmqJlv9Kf69+RiovU75xUXXBWAxWV+DTyzTXy+Ryz9QEAAAC1BCGppmj8WymmoXQgX9q4UC6XSw/0b60oT7iWbtypN1f9ZLpCAAAAoFYgJNUUYeFSy7JT7mZKktLqRuuW7hmSpCkff8toEgAAAFANCEk1Sdkpd1+/I5UclCRd89vGivW69d2OvfpofZ7B4gAAAIDagZBUkzTuIkXXl/b/Kv2wSJIUG+nR7zs1liQ9t+h7k9UBAAAAtQIhqSYJd0st+pZe/+8s/+ahXZvKHebS59/v1Jof883UBgAAANQShKSaxn/K3RyppFiSlBIfpcvapEhiNAkAAACoKkJSTdP0fCmqrrTvF2nTp/7Nfzy/mSTpnTVb9dOu/aaqAwAAAGo8QlJNE+6WWlxWen3tW/7Nrc+IV5cz66nE52jqJxsNFQcAAADUfISkmqjslLt1b0u+Ev/m/72gdDRpxrJcFRw4aKIyAAAAoMYjJNVE6RdKkQnS3jxp82L/5m5nNVBGwzraU1isGUs3m6sPAAAAqMEISTVRuEdqfvQpdy6XS388P12SNPXTH3SwxGeiOgAAAKBGIyTVVGWn3K2dLfkOh6F+55yh+nW82pp/QO98udVQcQAAAEDNRUiqqZp1k7zx0p5tUu4S/+ZIT7iGdG4iqXQ6cMdxDBUIAAAA1Ew1IiQ99dRTatq0qSIjI9WpUyctXbrUdEnmuSOk5peUXg845U6SrvltE0V6wvTfLQVa/N0vBooDAAAAai7rQ9Krr76q22+/Xffff79Wrlyptm3bqlevXsrLyzNdmnn+We6CT7lLjInQ79qlSWJxWQAAAOBkuRzLz8fq1KmTOnTooCeffFKS5PP5lJaWpuHDh+vuu+8+4eMLCgoUHx+v/Px8xcXFnepyQ+vgAemR30hFu6Xz75Diz5DcUZInUtv3uXT7zPXa70Tojl7NFR/lkSvMJZfLJZcrTC6p9GeYK+AJXXK5Sn9KpRNBHN4VFrTNFfSwsKPu71LgYwPvHHj/st3lvE7A4w/XeDjTu456Tlc5Nw/VGnbk/wW4Al47eN9x34NLcnT041yH9gd9JK5j11re+y3vfRz3cUfW6r8e8P2UfVdBn9vxXyP4+Y/92ZzMY0/uccc5bgAAAKqootnAHcKaTlpRUZFWrFihkSNH+reFhYWpR48eWrx4cbmPKSwsVGFhof92QUHBKa/TGE+kdHYfac1r0qJHg3YlSXol4tCNj0NdGBBaPic4UDlB1ysetgLvW7X/PTr2ax6vnuO95vEfV7lAWdnHneh5jvs+jht+T1TP4Wd2HeNVDtfiOqoWx7+tou+7ct9V+fWUf7s6VOQ5yz6vY93T8f8M/DMQ6mOuuo7HgOvW/WdLZT+36nm2I5/nWN/Hsf58lf+cRz5HaD/zU/Fn6lSp7PFYtfdYtc+n/H9PK/eceYlZ6jT8pSrVE0pWh6Sff/5ZJSUlSkpKCtqelJSkr7/+utzHTJgwQWPHjg1FeXb4n/skb6x0YFfpyFLxfv/PwgP7tHNXgRzHF/AX3qFD3HGC/hI88i/EE90+cnvwc5Upb1959zvi/s7R9zlereU/d/m/ELiO+ON+rNc5+v0fFuayevD1tHT874TvCzht8ccfNjnNj8eCvQ1Nl3BSrA5JlTFy5Ejdfvvt/tsFBQVKS0szWNEplthEumxSubu8klJCWw0kOf7+MEeOIykg9DmH/oZ0HEdynODIVs79Dj/n0f//d/Q+5/BjnYDXOdZrHOdM2+M97tCGcms5utYTPE+5+5zAjYe2VLzWIEd9bseqrPzXrfhjAx967PXJjlvrUfsq+l0dp5jj3PnEZ1of731U7DVO5vvXEZ9b8J8NlXP65YnGRI6s5chtJ1hH7nifeQUfd9T9HJ+OVfeR/zFzrL8bjn7K4/8HwdH/cx3w31jOka9d0T8DJ3McV+hRx39cpbsCKv8b6Qn/rFZkRODIv0dPwW/ILjnHHMk58tVcx/k75igVfY/lPa4yQv16pQ+uxJ4TvWQV3n9lHnYyz1vu93/iz9w58u/KwP/IPuoYL0fA9xoTV/+Er2cTq0NS/fr1FR4eru3btwdt3759u5KTk8t9jNfrldfrDUV5QLkCe6BqzkkAAAAAKGP17HYRERFq166d5s+f79/m8/k0f/58de7c2WBlAAAAAGorq0eSJOn222/XkCFD1L59e3Xs2FGTJ0/W3r17dd1115kuDQAAAEAtZH1Iuuqqq7Rjxw6NHj1a27Zt0znnnKP33nvvqMkcAAAAAKA6WL9OUlXV6nWSAAAAAFRYRbOB1T1JAAAAABBqhCQAAAAACEBIAgAAAIAAhCQAAAAACEBIAgAAAIAAhCQAAAAACEBIAgAAAIAAhCQAAAAACEBIAgAAAIAAhCQAAAAACOA2XcCp5jiOJKmgoMBwJQAAAABMKssEZRnhWGp9SNq9e7ckKS0tzXAlAAAAAGywe/duxcfHH3O/yzlRjKrhfD6ftmzZotjYWLlcLqO1FBQUKC0tTbm5uYqLizNaC2oOjhtUFscOKoPjBpXBcYPKCvWx4ziOdu/erdTUVIWFHbvzqNaPJIWFhalRo0amywgSFxfHXyA4aRw3qCyOHVQGxw0qg+MGlRXKY+d4I0hlmLgBAAAAAAIQkgAAAAAgACEphLxer+6//355vV7TpaAG4bhBZXHsoDI4blAZHDeoLFuPnVo/cQMAAAAAnAxGkgAAAAAgACEJAAAAAAIQkgAAAAAgACEJAAAAAAIQkkLoqaeeUtOmTRUZGalOnTpp6dKlpkuCRSZMmKAOHTooNjZWDRs2VP/+/bV+/fqg+xw4cEA5OTmqV6+e6tSpo0GDBmn79u2GKoaNHnroIblcLo0YMcK/jeMG5fnpp590zTXXqF69eoqKilJmZqaWL1/u3+84jkaPHq2UlBRFRUWpR48e+uabbwxWDBuUlJRo1KhRSk9PV1RUlM4880yNHz9egfOAcexg4cKF6tu3r1JTU+VyuTRr1qyg/RU5Rnbu3Kns7GzFxcUpISFBN9xwg/bs2ROy90BICpFXX31Vt99+u+6//36tXLlSbdu2Va9evZSXl2e6NFhiwYIFysnJ0eeff6558+bp4MGD6tmzp/bu3eu/z2233aa3335br7/+uhYsWKAtW7Zo4MCBBquGTZYtW6ZnnnlGbdq0CdrOcYMj/frrr+ratas8Ho/mzp2rtWvX6rHHHlNiYqL/PhMnTtTjjz+up59+WkuWLFFMTIx69eqlAwcOGKwcpj388MOaMmWKnnzySa1bt04PP/ywJk6cqCeeeMJ/H44d7N27V23bttVTTz1V7v6KHCPZ2dn673//q3nz5mnOnDlauHChhg0bFqq3IDkIiY4dOzo5OTn+2yUlJU5qaqozYcIEg1XBZnl5eY4kZ8GCBY7jOM6uXbscj8fjvP766/77rFu3zpHkLF682FSZsMTu3budjIwMZ968ec6FF17o3HrrrY7jcNygfHfddZdz3nnnHXO/z+dzkpOTnUceecS/bdeuXY7X63WmT58eihJhqUsvvdS5/vrrg7YNHDjQyc7OdhyHYwdHk+TMnDnTf7six8jatWsdSc6yZcv895k7d67jcrmcn376KSR1M5IUAkVFRVqxYoV69Ojh3xYWFqYePXpo8eLFBiuDzfLz8yVJdevWlSStWLFCBw8eDDqOmjdvrsaNG3McQTk5Obr00kuDjg+J4wblmz17ttq3b6/f/e53atiwobKysvTcc8/592/cuFHbtm0LOm7i4+PVqVMnjpvTXJcuXTR//nxt2LBBkvTFF1/ok08+UZ8+fSRx7ODEKnKMLF68WAkJCWrfvr3/Pj169FBYWJiWLFkSkjrdIXmV09zPP/+skpISJSUlBW1PSkrS119/bagq2Mzn82nEiBHq2rWrWrduLUnatm2bIiIilJCQEHTfpKQkbdu2zUCVsMWMGTO0cuVKLVu27Kh9HDcoz/fff68pU6bo9ttv1z333KNly5bplltuUUREhIYMGeI/Nsr7d4vj5vR29913q6CgQM2bN1d4eLhKSkr0wAMPKDs7W5I4dnBCFTlGtm3bpoYNGwbtd7vdqlu3bsiOI0ISYKGcnBx99dVX+uSTT0yXAsvl5ubq1ltv1bx58xQZGWm6HNQQPp9P7du314MPPihJysrK0ldffaWnn35aQ4YMMVwdbPbaa6/plVde0bRp09SqVSutXr1aI0aMUGpqKscOahVOtwuB+vXrKzw8/KjZpLZv367k5GRDVcFWN998s+bMmaOPPvpIjRo18m9PTk5WUVGRdu3aFXR/jqPT24oVK5SXl6dzzz1XbrdbbrdbCxYs0OOPPy63262kpCSOGxwlJSVFLVu2DNrWokULbd68WZL8xwb/buFId955p+6++25dffXVyszM1B/+8AfddtttmjBhgiSOHZxYRY6R5OTkoyY3Ky4u1s6dO0N2HBGSQiAiIkLt2rXT/Pnz/dt8Pp/mz5+vzp07G6wMNnEcRzfffLNmzpypDz/8UOnp6UH727VrJ4/HE3QcrV+/Xps3b+Y4Oo11795da9as0erVq/2X9u3bKzs723+d4wZH6tq161FLDGzYsEFNmjSRJKWnpys5OTnouCkoKNCSJUs4bk5z+/btU1hY8K+P4eHh8vl8kjh2cGIVOUY6d+6sXbt2acWKFf77fPjhh/L5fOrUqVNoCg3J9BBwZsyY4Xi9XufFF1901q5d6wwbNsxJSEhwtm3bZro0WOJPf/qTEx8f73z88cfO1q1b/Zd9+/b573PjjTc6jRs3dj788ENn+fLlTufOnZ3OnTsbrBo2CpzdznE4bnC0pUuXOm6323nggQecb775xnnllVec6Oho51//+pf/Pg899JCTkJDgvPXWW86XX37p9OvXz0lPT3f2799vsHKYNmTIEOeMM85w5syZ42zcuNF58803nfr16zt/+ctf/Pfh2MHu3budVatWOatWrXIkOZMmTXJWrVrlbNq0yXGcih0jvXv3drKyspwlS5Y4n3zyiZORkeEMHjw4ZO+BkBRCTzzxhNO4cWMnIiLC6dixo/P555+bLgkWkVTuZerUqf777N+/37npppucxMREJzo62hkwYICzdetWc0XDSkeGJI4blOftt992Wrdu7Xi9Xqd58+bOs88+G7Tf5/M5o0aNcpKSkhyv1+t0797dWb9+vaFqYYuCggLn1ltvdRo3buxERkY6zZo1c+69916nsLDQfx+OHXz00Ufl/k4zZMgQx3Eqdoz88ssvzuDBg506deo4cXFxznXXXefs3r07ZO/B5TgBSyQDAAAAwGmOniQAAAAACEBIAgAAAIAAhCQAAAAACEBIAgAAAIAAhCQAAAAACEBIAgAAAIAAhCQAAAAACEBIAgAggMvl0qxZs0yXAQAwiJAEALDG0KFD5XK5jrr07t3bdGkAgNOI23QBAAAE6t27t6ZOnRq0zev1GqoGAHA6YiQJAGAVr9er5OTkoEtiYqKk0lPhpkyZoj59+igqKkrNmjXTG2+8EfT4NWvW6H/+538UFRWlevXqadiwYdqzZ0/QfV544QW1atVKXq9XKSkpuvnmm4P2//zzzxowYICio6OVkZGh2bNn+/f9+uuvys7OVoMGDRQVFaWMjIyjQh0AoGYjJAEAapRRo0Zp0KBB+uKLL5Sdna2rr75a69atkyTt3btXvXr1UmJiopYtW6bXX39dH3zwQVAImjJlinJycjRs2DCtWbNGs2fP1m9+85ug1xg7dqyuvPJKffnll7rkkkuUnZ2tnTt3+l9/7dq1mjt3rtatW6cpU6aofv36ofsAAACnnMtxHMd0EQAASKU9Sf/6178UGRkZtP2ee+7RPffcI5fLpRtvvFFTpkzx7/vtb3+rc889V3//+9/13HPP6a677lJubq5iYmIkSe+++6769u2rLVu2KCkpSWeccYauu+46/fWvfy23BpfLpfvuu0/jx4+XVBq86tSpo7lz56p37966/PLLVb9+fb3wwgun6FMAAJhGTxIAwCoXXXRRUAiSpLp16/qvd+7cOWhf586dtXr1aknSunXr1LZtW39AkqSuXbvK5/Np/fr1crlc2rJli7p3737cGtq0aeO/HhMTo7i4OOXl5UmS/vSnP2nQoEFauXKlevbsqf79+6tLly6Veq8AADsRkgAAVomJiTnq9LfqEhUVVaH7eTyeoNsul0s+n0+S1KdPH23atEnvvvuu5s2bp+7duysnJ0ePPvpotdcLADCDniQAQI3y+eefH3W7RYsWkqQWLVroiy++0N69e/37P/30U4WFhenss89WbGysmjZtqvnz51ephgYNGmjIkCH617/+pcmTJ+vZZ5+t0vMBAOzCSBIAwCqFhYXatm1b0Da32+2fHOH1119X+/btdd555+mVV17R0qVL9fzzz0uSsrOzdf/992vIkCEaM2aMduzYoeHDh+sPf/iDkpKSJEljxozRjTfeqIYNG6pPnz7avXu3Pv30Uw0fPrxC9Y0ePVrt2rVTq1atVFhYqDlz5vhDGgCgdiAkAQCs8t577yklJSVo29lnn62vv/5aUunMczNmzNBNN92klJQUTZ8+XS1btpQkRUdH6/3339ett96qDh06KDo6WoMGDdKkSZP8zzVkyBAdOHBAf/vb33THHXeofv36uuKKKypcX0REhEaOHKkffvhBUVFROv/88zVjxoxqeOcAAFswux0AoMZwuVyaOXOm+vfvb7oUAEAtRk8SAAAAAAQgJAEAAABAAHqSAAA1BmeIAwBCgZEkAAAAAAhASAIAAACAAIQkAAAAAAhASAIAAACAAIQkAAAAAAhASAIAAACAAIQkAAAAAAhASAIAAACAAIQkAAAAAAjw/wER7gVY0bRBQwAAAABJRU5ErkJggg==\n"
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Total training time: 14138.3713 seconds\n"
          ]
        }
      ],
      "source": [
        "# Plot training and validation loss\n",
        "plt.figure(figsize=(10, 5))\n",
        "plt.plot(train_losses, label='Train Loss')\n",
        "plt.plot(val_losses, label='Val Loss')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Loss')\n",
        "plt.legend()\n",
        "#plt.savefig(\"/content/drive/My Drive/Project 2 Neuroimage Synthesis/loss_plot.png\")\n",
        "plt.show()\n",
        "\n",
        "# Total training time\n",
        "end_time = time.time()\n",
        "total_time = end_time - start_time\n",
        "print(f'Total training time: {total_time:.4f} seconds')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Reload Model"
      ],
      "metadata": {
        "id": "YA79CeFcEqdJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the model\n",
        "model_path = \"/content/drive/My Drive/Project 2 Neuroimage Synthesis/unet.pt\"\n",
        "\n",
        "# model_path = \"/content/drive/My Drive/Project 2 Neuroimage Synthesis/Team2/code/Code - unet.pt\"\n",
        "model.load_state_dict(torch.load(model_path))\n",
        "\n",
        "# Switch to eval mode\n",
        "model.eval()\n"
      ],
      "metadata": {
        "id": "5tmJhuDK78Ek",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fcb9e34b-a1cb-411a-ff38-25370a1f580d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "UNet(\n",
              "  (enc1): ConvBlock(\n",
              "    (conv): Conv3d(2, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (enc2): ConvBlock(\n",
              "    (conv): Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (enc3): ConvBlock(\n",
              "    (conv): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (enc4): ConvBlock(\n",
              "    (conv): Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (pool): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "  (middle): ConvBlock(\n",
              "    (conv): Conv3d(256, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (up4): ConvTranspose3d(512, 256, kernel_size=(2, 2, 2), stride=(2, 2, 2))\n",
              "  (dec4): ConvBlock(\n",
              "    (conv): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (up3): ConvTranspose3d(256, 128, kernel_size=(2, 2, 2), stride=(2, 2, 2))\n",
              "  (dec3): ConvBlock(\n",
              "    (conv): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (up2): ConvTranspose3d(128, 64, kernel_size=(2, 2, 2), stride=(2, 2, 2))\n",
              "  (dec2): ConvBlock(\n",
              "    (conv): Conv3d(128, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (up1): ConvTranspose3d(64, 32, kernel_size=(2, 2, 2), stride=(2, 2, 2))\n",
              "  (dec1): ConvBlock(\n",
              "    (conv): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
              "    (relu): ReLU()\n",
              "  )\n",
              "  (output): Conv3d(32, 2, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lg43aZtYuEHE"
      },
      "source": [
        "## Test - Produce Synthesized Maps of FA and ADC"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import os\n",
        "import nibabel as nib\n",
        "from nibabel.processing import resample_to_output\n",
        "import csv\n",
        "from skimage.metrics import structural_similarity as ssim\n",
        "from skimage.metrics import peak_signal_noise_ratio as psnr\n",
        "from skimage.metrics import mean_squared_error as mse"
      ],
      "metadata": {
        "id": "BLSGNE995hep"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j6C0trMP6uNF",
        "outputId": "47e727ae-d7da-4ad8-c6be-f926a84e1cf7"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['103',\n",
              " '177',\n",
              " '140',\n",
              " '184',\n",
              " '042',\n",
              " '111',\n",
              " '023',\n",
              " '020',\n",
              " '165',\n",
              " '065',\n",
              " '048',\n",
              " '114',\n",
              " '141',\n",
              " '183',\n",
              " '015',\n",
              " '142',\n",
              " '196',\n",
              " '037',\n",
              " '098',\n",
              " '195',\n",
              " '137']"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "source": [
        "test_ids"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from nilearn.image import resample_img\n",
        "\n",
        "# Define the new voxel size for resampling\n",
        "new_voxel_size = (1.25, 1.25, 1.25)\n",
        "\n",
        "successful_patients = []\n",
        "\n",
        "with torch.no_grad():\n",
        "    for i, (input_imgs, target_imgs) in enumerate(test_loader):\n",
        "        patient_id = test_ids[i]\n",
        "        print(f\"Processing patient {patient_id}\")\n",
        "\n",
        "        # Save the original images for visualization\n",
        "        original_input_imgs = input_imgs.clone()\n",
        "\n",
        "        # Move tensors to the device\n",
        "        input_imgs = input_imgs.to(device)\n",
        "        target_imgs = target_imgs.to(device)\n",
        "\n",
        "        # Generate output maps from input images\n",
        "        outputs = model(input_imgs)\n",
        "\n",
        "        # Check the number of outputs\n",
        "        if len(outputs) == 2:\n",
        "            # Unpack the outputs\n",
        "            fake_fa, fake_adc = outputs\n",
        "\n",
        "            # Convert tensors to Nifti images\n",
        "            fake_fa_nifti = nib.Nifti1Image(fake_fa[0].cpu().numpy(), np.eye(4))\n",
        "            fake_adc_nifti = nib.Nifti1Image(fake_adc[0].cpu().numpy(), np.eye(4))\n",
        "\n",
        "            # Resample the images to the new voxel size\n",
        "            fake_fa_nifti = resample_img(fake_fa_nifti, target_affine=np.diag(new_voxel_size))\n",
        "            fake_adc_nifti = resample_img(fake_adc_nifti, target_affine=np.diag(new_voxel_size))\n",
        "\n",
        "            # Save the synthesized FA and ADC maps\n",
        "            patient_id = test_ids[i]\n",
        "            patient_folder = os.path.join(root_dir, patient_id)\n",
        "\n",
        "            fa_syn_path = os.path.join(patient_folder, 'FA_syn.nii.gz')\n",
        "            adc_syn_path = os.path.join(patient_folder, 'ADC_syn.nii.gz')\n",
        "\n",
        "            nib.save(fake_fa_nifti, fa_syn_path)\n",
        "            nib.save(fake_adc_nifti, adc_syn_path)\n",
        "\n",
        "            # Add the successful patient_id to the list\n",
        "            successful_patients.append(patient_id)\n",
        "        elif len(outputs) == 1:\n",
        "            print(f\"Only one output for patient {test_ids[i]}, skipping.\")\n",
        "        else:\n",
        "            print(f\"Unexpected number of outputs for patient {test_ids[i]}, skipping.\")\n",
        "\n",
        "print(\"Successfully saved patients:\", successful_patients)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rE3yAHqb6PXa",
        "outputId": "a0f36e10-8d9f-45d3-ea07-0e1dbeec4d60"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Processing patient 103\n",
            "Processing patient 177\n",
            "Processing patient 140\n",
            "Processing patient 184\n",
            "Processing patient 042\n",
            "Processing patient 111\n",
            "Processing patient 023\n",
            "Processing patient 020\n",
            "Processing patient 165\n",
            "Processing patient 065\n",
            "Processing patient 048\n",
            "Only one output for patient 048, skipping.\n",
            "Successfully saved patients: ['103', '177', '140', '184', '042', '111', '023', '020', '165', '065']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RzrTJHVf8i_i"
      },
      "source": [
        "## Save Metrics (MSE, PSNR, SSIM) to a CSV comparing synthesized FA and ADC maps to original FA and ADC for each patient in test set."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import csv\n",
        "import os\n",
        "import numpy as np\n",
        "import nibabel as nib\n",
        "from skimage.metrics import mean_squared_error as mse, peak_signal_noise_ratio as psnr, structural_similarity as ssim\n",
        "from sklearn.metrics import mean_absolute_error as mae\n",
        "from sklearn.metrics import mean_absolute_error\n",
        "\n",
        "from nilearn.image import resample_to_img\n"
      ],
      "metadata": {
        "id": "iW-ofTGwqDBF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the headers for the CSV file\n",
        "headers = ['Patient_ID', 'MSE_FA', 'PSNR_FA', 'SSIM_FA', 'MAE_FA', 'MSE_ADC', 'PSNR_ADC', 'SSIM_ADC', 'MAE_ADC']\n",
        "\n",
        "\n",
        "# Initialize running sums for calculating the averages\n",
        "sum_mse_fa, sum_psnr_fa, sum_ssim_fa = 0, 0, 0\n",
        "sum_mse_adc, sum_psnr_adc, sum_ssim_adc = 0, 0, 0\n",
        "sum_mae_fa, sum_mae_adc = 0, 0\n",
        "\n",
        "\n",
        "\n",
        "# Open the CSV file for writing\n",
        "with open('attention_metrics.csv', 'w', newline='') as csvfile:\n",
        "    csv_writer = csv.writer(csvfile)\n",
        "    csv_writer.writerow(headers)\n",
        "\n",
        "    for patient_id in successful_patients:\n",
        "        patient_folder = os.path.join(root_dir, patient_id)\n",
        "\n",
        "        fa_syn_path = os.path.join(patient_folder, 'FA_syn.nii.gz')\n",
        "        adc_syn_path = os.path.join(patient_folder, 'ADC_syn.nii.gz')\n",
        "        fa_deformed_path = os.path.join(patient_folder, 'FA_deformed.nii.gz')\n",
        "        adc_deformed_path = os.path.join(patient_folder, 'ADC_deformed.nii.gz')\n",
        "\n",
        "        fa_syn_nifti = nib.load(fa_syn_path)\n",
        "        fa_syn = fa_syn_nifti.get_fdata()\n",
        "        adc_syn_nifti = nib.load(adc_syn_path)\n",
        "        adc_syn = adc_syn_nifti.get_fdata()\n",
        "\n",
        "        fa_deformed_nifti = nib.load(fa_deformed_path)\n",
        "        fa_deformed = fa_deformed_nifti.get_fdata()\n",
        "        adc_deformed_nifti = nib.load(adc_deformed_path)\n",
        "        adc_deformed = adc_deformed_nifti.get_fdata()\n",
        "\n",
        "        # Resample the synthesized FA and ADC images to match the deformed FA and ADC images\n",
        "        fa_syn_nifti = resample_to_img(fa_syn_nifti, fa_deformed_nifti)\n",
        "        fa_syn = fa_syn_nifti.get_fdata()\n",
        "        adc_syn_nifti = resample_to_img(adc_syn_nifti, adc_deformed_nifti)\n",
        "        adc_syn = adc_syn_nifti.get_fdata()\n",
        "\n",
        "        mae_fa = mean_absolute_error(fa_syn.flatten(), fa_deformed.flatten())\n",
        "        mse_fa = mse(fa_syn, fa_deformed)\n",
        "        psnr_fa = psnr(fa_syn, fa_deformed, data_range=fa_syn.max() - fa_syn.min())\n",
        "        ssim_fa = ssim(fa_syn, fa_deformed)\n",
        "\n",
        "        mae_adc = mean_absolute_error(adc_syn.flatten(), adc_deformed.flatten())\n",
        "        mse_adc = mse(adc_syn, adc_deformed)\n",
        "        psnr_adc = psnr(adc_syn, adc_deformed, data_range=adc_syn.max() - adc_syn.min())\n",
        "        ssim_adc = ssim(adc_syn, adc_deformed)\n",
        "\n",
        "        # Write the metrics to the CSV file\n",
        "        csv_writer.writerow([patient_id, mse_fa, psnr_fa, ssim_fa, mae_fa, mse_adc, psnr_adc, ssim_adc, mae_adc])\n",
        "\n",
        "        # Update the running sums\n",
        "        sum_mse_fa += mse_fa\n",
        "        sum_psnr_fa += psnr_fa\n",
        "        sum_ssim_fa += ssim_fa\n",
        "        sum_mse_adc += mse_adc\n",
        "        sum_psnr_adc += psnr_adc\n",
        "        sum_ssim_adc += ssim_adc\n",
        "\n",
        "        sum_mae_fa += mae_fa\n",
        "        sum_mae_adc += mae_adc\n",
        "\n",
        "\n",
        "\n",
        "# Calculate the averages\n",
        "num_successful_patients = len(successful_patients)\n",
        "avg_mse_fa = sum_mse_fa / num_successful_patients\n",
        "avg_psnr_fa = sum_psnr_fa / num_successful_patients\n",
        "avg_ssim_fa = sum_ssim_fa / num_successful_patients\n",
        "avg_mse_adc = sum_mse_adc / num_successful_patients\n",
        "avg_psnr_adc = sum_psnr_adc / num_successful_patients\n",
        "avg_ssim_adc = sum_ssim_adc / num_successful_patients\n",
        "avg_mae_fa = sum_mae_fa / num_successful_patients\n",
        "avg_mae_adc = sum_mae_adc / num_successful_patients\n",
        "\n",
        "\n",
        "# Output the averages\n",
        "print(\"Average MSE FA:\", avg_mse_fa)\n",
        "print(\"Average PSNR FA:\", avg_psnr_fa)\n",
        "print(\"Average SSIM FA:\", avg_ssim_fa)\n",
        "print(\"Average MSE ADC:\", avg_mse_adc)\n",
        "print(\"Average PSNR ADC:\", avg_psnr_adc)\n",
        "print(\"Average SSIM ADC:\", avg_ssim_adc)\n",
        "print(\"Average MAE FA:\", avg_mae_fa)\n",
        "print(\"Average MAE ADC:\", avg_mae_adc)\n",
        "\n",
        "\n",
        "print(\"Metrics saved to metrics.csv\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0mJO3TiurAYV",
        "outputId": "67fe0abb-f27a-4f5b-dfc5-881c4dbc8805"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average MSE FA: 0.032549113481524825\n",
            "Average PSNR FA: 14.829087339733443\n",
            "Average SSIM FA: 0.6989486501670562\n",
            "Average MSE ADC: 0.003323685569761844\n",
            "Average PSNR ADC: 24.715863547509294\n",
            "Average SSIM ADC: 0.9539950789508278\n",
            "Average MAE FA: 0.07557693652318702\n",
            "Average MAE ADC: 0.011682469192816108\n",
            "Metrics saved to metrics.csv\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}