MR_Synth_LPIPS.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "\"\"\"\n",
        "@author: Parisima\n",
        "This notebook requires predictions\n",
        "in .mat format\n",
        "and BraTS'21 test set\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "vJNkV95RZBWJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Egj20bQVLIg"
      },
      "outputs": [],
      "source": [
        "!pip install lpips"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "crBXUTyOdVWX"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow-addons"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eyy01SJcdm3p"
      },
      "outputs": [],
      "source": [
        "!pip install elasticdeform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IIk23AGRKogI"
      },
      "outputs": [],
      "source": [
        "# Access to google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hx8Up2bWWJzK"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import numpy as np\n",
        "import nibabel as nib\n",
        "import tensorflow as tf\n",
        "import os\n",
        "import pandas as pd\n",
        "import scipy.io as sio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dmBs-vH2c2Wt"
      },
      "outputs": [],
      "source": [
        "sys.path.append('/YourPath/CCL-Synthetis/')\n",
        "\n",
        "# Assuming utils and Datagen are your local modules or packages\n",
        "from utils.model_utils import modelObj\n",
        "from Synthesis.synthesis_losses import lossObj\n",
        "from Datagen.h5_pretrain_Synth_Data_Generator import DataLoaderObj\n",
        "\n",
        "import Synthesis.synth_config as cfg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q0-cdSKXW6Wx"
      },
      "outputs": [],
      "source": [
        "sys.path.append('/YourPath/multi-contrast-contrastive-learning/')\n",
        "\n",
        "from utils.utils import myCrop3D\n",
        "from utils.utils import contrastStretch\n",
        "\n",
        "def normalize_img_zmean(img, mask):\n",
        "    ''' Zero mean unit standard deviation normalization based on a mask'''\n",
        "    mask_signal = img[mask>0]\n",
        "    mean_ = mask_signal.mean()\n",
        "    std_ = mask_signal.std()\n",
        "    img = (img - mean_ )/ std_\n",
        "    return img, mean_, std_\n",
        "\n",
        "def normalize_img(img):\n",
        "    img = (img - img.min())/(img.max()-img.min())\n",
        "    return img\n",
        "\n",
        "def load_subject(datadir, subName):\n",
        "    data_suffix = ['_t1ce.nii.gz', '_t2.nii.gz', '_t1.nii.gz', '_flair.nii.gz']\n",
        "    sub_img = []\n",
        "    mask = None\n",
        "    subject_dir = os.path.join(datadir, subName)  # Correctly form the path to the subject's directory\n",
        "\n",
        "    for suffix in data_suffix:\n",
        "        img_path = os.path.join(subject_dir, subName + suffix)  # Correct path to the image file\n",
        "        img_data = nib.load(img_path).get_fdata()\n",
        "        img_data = np.rot90(img_data, -1)\n",
        "        img_data = myCrop3D(img_data, (192,192))\n",
        "\n",
        "        if mask is None:\n",
        "            mask = np.zeros(img_data.shape)\n",
        "            mask[img_data > 0] = 1\n",
        "\n",
        "        img_data = contrastStretch(img_data, mask, 0.01, 99.9)\n",
        "        img_data, mean_, std_ = normalize_img_zmean(img_data, mask) # Change to normalize_img if your model trained with this\n",
        "        sub_img.append(img_data)\n",
        "\n",
        "    sub_img = np.stack(sub_img, axis=-1)\n",
        "    sub_img = np.transpose(sub_img, (2, 0, 1, 3))  # Adjust dimensions as needed\n",
        "    sub_img = sub_img[40:120]  # Assuming your volume z-axis slice range\n",
        "\n",
        "    return sub_img\n",
        "\n",
        "#-----------------------------------------------------------------\n",
        "\n",
        "def get_data(img, contrast_idx, target_contrast_idx):\n",
        "    \"\"\"Returns tuple (input, target) correspond to sample #idx.\"\"\"\n",
        "    x_train = generate_X(img, contrast_idx)\n",
        "    y_train = generate_Y(img, target_contrast_idx)\n",
        "    return tf.identity(x_train), tf.identity(y_train)\n",
        "\n",
        "def generate_X(img, contrast_idx):\n",
        "    X = img[..., contrast_idx]\n",
        "    return X\n",
        "\n",
        "def generate_Y(img, target_contrast_idx):\n",
        "    Y = img[..., target_contrast_idx]\n",
        "    return Y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FDTubnsDXwAC"
      },
      "outputs": [],
      "source": [
        "import lpips\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sIQmm4ZjX0Ud"
      },
      "outputs": [],
      "source": [
        "# Choose the metric model\n",
        "loss_fn = lpips.LPIPS(net='alex')  # Using AlexNet\n",
        "# loss_fn = lpips.LPIPS(net='vgg')  # Using VGG"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sn3W7YxisDaK"
      },
      "outputs": [],
      "source": [
        "def load_predictions(predictions_file):\n",
        "    data = sio.loadmat(predictions_file)\n",
        "    # Assuming the data structure is known and 'predictions' is the key\n",
        "    predictions = {\n",
        "        'Baseline': data['predictions']['Baseline'][0, 0],\n",
        "        'Partial_Decoder': data['predictions']['Partial_Decoder'][0, 0],\n",
        "        'Full_Decoder': data['predictions']['Full_Decoder'][0, 0]\n",
        "    }\n",
        "    # Decode byte arrays if necessary\n",
        "    for key, value in predictions.items():\n",
        "        if isinstance(value, bytes):\n",
        "            # Example of decoding bytes to numpy array; adapt as needed\n",
        "            predictions[key] = np.frombuffer(value, dtype=np.float32).reshape((80, 192, 192, 3))\n",
        "    return predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JHmtz_FNX_uB"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import torch\n",
        "\n",
        "def prepare_image(image):\n",
        "    # Ensure the input image is a TensorFlow tensor with dtype float32\n",
        "    image = tf.convert_to_tensor(image, dtype=tf.float32)\n",
        "\n",
        "    # Normalize the image tensor to [-1, 1]\n",
        "    image_min = tf.reduce_min(image)\n",
        "    image_max = tf.reduce_max(image)\n",
        "    image = 2 * (image - image_min) / (image_max - image_min) - 1\n",
        "\n",
        "    # Check and adapt the tensor dimensions\n",
        "    # Assumption: The last dimension is channels if it's exactly 3; otherwise, we assume 1 channel.\n",
        "    if image.shape[-1] != 3:\n",
        "        # If the image does not have three channels, we tile to create three channels\n",
        "        image = tf.tile(image, [1, 1, 3])  # This expects image to be at least 3D; reshape if not\n",
        "\n",
        "    # Ensure the image tensor is in the correct shape (channels, height, width)\n",
        "    image = tf.transpose(image, [2, 0, 1])\n",
        "\n",
        "    # Convert TensorFlow tensor to a PyTorch tensor\n",
        "    image = torch.from_numpy(image.numpy()).float()\n",
        "\n",
        "    # Ensure the image tensor has a batch dimension\n",
        "    if image.dim() == 3:\n",
        "        image = image.unsqueeze(0)\n",
        "\n",
        "    return image\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3sD2tISYG12"
      },
      "outputs": [],
      "source": [
        "def compute_lpips(datadir, predictions_dir, cfg):\n",
        "    # loss_fn = lpips.LPIPS(net='alex')  # Initialize the LPIPS function\n",
        "    subject_scores = []\n",
        "    subjects = sorted([f for f in os.listdir(datadir) if f.startswith(\"BraTS2021_\")])\n",
        "\n",
        "    for subName in subjects:\n",
        "        print(f\"Processing {subName}...\")\n",
        "        img = load_subject(datadir, subName)\n",
        "        _, y_true = get_data(img, cfg.contrast_idx, cfg.target_contrast_idx)  # Ground truth images\n",
        "        y_true = y_true.numpy()\n",
        "\n",
        "        # Load predictions\n",
        "        predictions_file = f'{predictions_dir}/{subName}_predictions.mat'\n",
        "        predictions = load_predictions(predictions_file)\n",
        "\n",
        "        model_scores = {}\n",
        "        for model_name, y_pred in predictions.items():\n",
        "            scores = []\n",
        "            for i in range(y_true.shape[0]):  # Assuming y_true is numpy array\n",
        "                img0 = prepare_image(y_true[i])\n",
        "                # print(type(y_pred))\n",
        "                img1 = prepare_image(y_pred[i])\n",
        "                with torch.no_grad():\n",
        "                    score = loss_fn(img0, img1)\n",
        "                scores.append(score.item())\n",
        "            model_scores[model_name] = np.mean(scores)\n",
        "\n",
        "        subject_scores.append({\n",
        "            'Subject_ID': subName,\n",
        "            'Baseline_Avg_LPIPS': model_scores['Baseline'],\n",
        "            'Full_Decoder_Avg_LPIPS': model_scores['Full_Decoder'],\n",
        "            'Partial_Decoder_Avg_LPIPS': model_scores['Partial_Decoder']\n",
        "        })\n",
        "\n",
        "    return pd.DataFrame(subject_scores)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FyIjysLNZF32"
      },
      "outputs": [],
      "source": [
        "datadir = '/YourPath/BraTS2021_Test/'\n",
        "predictions_dir = '/YourPath/predictions/'\n",
        "\n",
        "results_df = compute_lpips(datadir, predictions_dir, cfg)\n",
        "results_df.to_csv('SaveDir.csv', index=False)\n",
        "print('CSV file has been saved.')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}