app.py

import io
import os
import sys

import gradio as gr
import numpy as np
import spaces
import torch

# from huggingface_hub import hf_hub_download
from huggingface_hub import snapshot_download
from PIL import Image, ImageDraw, ImageFont

zero = torch.Tensor([0]).cuda()

# Set the working directory to the root directory
# root_dir = os.path.abspath("..")
# os.chdir(root_dir)
# sys.path.insert(0, root_dir)

# download dataset & weights
snapshot_download(repo_id="armeet/fastmri-tiny", repo_type="dataset", local_dir=".")


device = "cuda"
# dataset_path = "/global/homes/p/peterwg/pscratch/datasets/mri_knee_dummy"
dataset_path = "dataset"

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F

import fastmri
from fastmri.datasets import SliceDatasetLMDB, SliceSample
from fastmri.subsample import create_mask_for_mask_type
from models.lightning.no_varnet_module import NOVarnetModule
from models.lightning.varnet_module import VarNetModule

acceleration_to_fractions = {
    1: 1,
    2: 0.16,
    4: 0.08,
    6: 0.06,
    8: 0.04,
    16: 0.02,
    32: 0.01,
}


def create_mask_fn(center_fraction, acceleration):
    mask_fn = create_mask_for_mask_type(
        "equispaced_fraction",
        [center_fraction],
        [acceleration],
    )
    return mask_fn


mask_4x = create_mask_fn(acceleration_to_fractions[4], 4)
mask_6x = create_mask_fn(acceleration_to_fractions[6], 6)
mask_8x = create_mask_fn(acceleration_to_fractions[8], 8)
mask_16x = create_mask_fn(acceleration_to_fractions[16], 16)

val_dataset_4x = SliceDatasetLMDB(
    "knee",
    partition="val",
    mask_fns=[mask_4x],
    complex=False,
    root=dataset_path,
    crop_shape=(320, 320),
    coils=15,
)

val_dataset_6x = SliceDatasetLMDB(
    "knee",
    partition="val",
    mask_fns=[mask_6x],
    complex=False,
    root=dataset_path,
    crop_shape=(320, 320),
    coils=15,
)

val_dataset_8x = SliceDatasetLMDB(
    "knee",
    partition="val",
    mask_fns=[mask_8x],
    complex=False,
    root=dataset_path,
    crop_shape=(320, 320),
    coils=15,
)
val_dataset_16x = SliceDatasetLMDB(
    "knee",
    partition="val",
    mask_fns=[mask_16x],
    complex=False,
    root=dataset_path,
    crop_shape=(320, 320),
    coils=15,
)

vn = VarNetModule.load_from_checkpoint(
    "vn.ckpt",
)
no = NOVarnetModule.load_from_checkpoint(
    "no.ckpt",
)
no.eval()
vn.eval()

bright_samples = [42, 69, 80, 137, 139, 226, 229]


def v(x):
    return x.detach().cpu().numpy().squeeze()


def viz(x, cmap="gray", vmin=0, vmax=1):
    processed_data = v(x)
    fig, ax = plt.subplots()
    ax.imshow(processed_data, cmap=cmap, vmin=vmin, vmax=vmax)
    ax.axis("off")  # Turn off axes
    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)  # Adjust margins
    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
    buf.seek(0)  # Rewind the buffer to the beginning
    plt.show()
    try:
        img = Image.open(buf)
        img_array = np.array(img)
    except Exception as e:
        print(f"Error converting image buffer to NumPy array: {e}")
        img_array = None
    finally:
        plt.close(fig)
        buf.close()
    return img_array


@spaces.GPU
def forward(model, idx, rate):
    if rate == 4:
        dataset = val_dataset_4x
    elif rate == 6:
        dataset = val_dataset_6x
    elif rate == 8:
        dataset = val_dataset_8x
    elif rate == 16:
        dataset = val_dataset_16x
    else:
        raise ValueError("Invalid rate")

    sample = dataset[idx]
    mask, k, target = (
        sample.mask.to(device),
        sample.masked_kspace.to(device),
        sample.target.to(device),
    )
    pred = model(k.unsqueeze(0), mask.unsqueeze(0), None)

    return mask, k, target, pred[0]


@spaces.GPU
def update_interface(sample_id, sample_rate):
    n = [None] * 6
    if sample_id is None or sample_rate is None or sample_id not in bright_samples:
        return n

    mask, k, target, pred_vn = forward(vn, sample_id, sample_rate)
    _, _, _, pred_no = forward(no, sample_id, sample_rate)

    k = viz(mask[0, :, :, 0], cmap="gray", vmin=0, vmax=1)
    target_res = viz(target, cmap="gray", vmin=None, vmax=None)

    pred_no_res = viz(pred_no, cmap="gray", vmin=None, vmax=None)
    pred_vn_res = viz(pred_vn, cmap="gray", vmin=None, vmax=None)

    diff_no_res = viz(torch.abs(pred_no - target), cmap=None, vmin=None, vmax=None)
    diff_vn_res = viz(torch.abs(pred_vn - target), cmap=None, vmin=None, vmax=None)

    return k, target_res, pred_no_res, pred_vn_res, diff_no_res, diff_vn_res


with gr.Blocks(theme=gr.themes.Monochrome(), fill_width=True) as demo:
    gr.Markdown(
        "# A Unified Model for Compressed Sensing MRI Across Undersampling Patterns [CPVR 2025]"
    )
    gr.Markdown("""  
> Armeet Singh Jatyani, Jiayun Wang, Aditi Chandrashekar, Zihui Wu, Miguel Liu-Schiaffini, Bahareh Tolooshams, Anima Anandkumar  
                """)
    gr.Markdown(
        "[![arXiv](https://img.shields.io/badge/arXiv-2410.16290-b31b1b.svg?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2410.16290)"
    )
    gr.Markdown(
        "[![](https://img.shields.io/badge/Blog-armeet.ca%2Fnomri-yellow?style=flat-square)](https://armeet.ca/nomri)"
    )

    gr.Markdown(
        "This demo showcases the performance of our unified model for compressed sensing MRI across different acceleration rates."
    )
    gr.Markdown(
        "We recommend trying samples with a 16x acceleration pattern first, as reconstruction differences are easy to observe."
    )

    gr.Markdown(
        "At lower acceleration rates (4x or 6x), the difference in reconstruction quality is difficult to discern. At higher acceleration rates, look for blurring, repeating, or distortion, especially near edges and in backgrounds. We provide difference images to help identify reconstruction errors."
    )
    with gr.Row():
        dropdown_sample = gr.Dropdown(
            choices=bright_samples,
            label="Select a Sample",
            info="Choose one of the available samples.",
            filterable=False,
            value=229,
        )
    with gr.Row():
        dropdown_rate = gr.Radio(
            choices=[16, 8, 6, 4],
            value=16,
            label="Select an Acceleration Rate",
            info="Ex: 4x means the model is trained to reconstruct from 4x undersampled k-space data",
            # filterable=False,
        )

    with gr.Row():
        with gr.Column():
            gr.Label("Undersampling Mask")
            k = gr.Image(label=None, interactive=False)
        with gr.Column():
            gr.Label("Ground Truth")
            target = gr.Image(label=None, interactive=False)
        with gr.Column():
            gr.Label("NO (ours)")
            pred_no = gr.Image(label="Reconstruction (ours)", interactive=False)
        with gr.Column():
            gr.Label("VN (existing)")
            pred_vn = gr.Image(label="Reconstruction (existing)", interactive=False)
    with gr.Row():
        with gr.Column():
            pass
        with gr.Column():
            pass
        with gr.Column():
            diff_no = gr.Image(label="| Recon - GT | (ours)", interactive=False)
        with gr.Column():
            diff_vn = gr.Image(label="| Recon - GT | (existing)", interactive=False)

    gr.Markdown("""
```
@inproceedings{jatyani2025nomri,
  author    = {Armeet Singh Jatyani* and Jiayun Wang* and Aditi Chandrashekar and Zihui Wu and Miguel Liu-Schiaffini and Bahareh Tolooshams and Anima Anandkumar},
  title     = {A Unified Model for Compressed Sensing MRI Across Undersampling Patterns},
  booktitle = {Conference on Computer Vision and Pattern Recognition (CVPR) Proceedings},
  abbr      = {CVPR},
  year      = {2025}
}
```
                """)

    update_inputs = [dropdown_sample, dropdown_rate]
    update_outputs = [k, target, pred_no, pred_vn, diff_no, diff_vn]

    dropdown_sample.change(
        fn=update_interface, inputs=update_inputs, outputs=update_outputs
    )
    dropdown_rate.change(
        fn=update_interface, inputs=update_inputs, outputs=update_outputs
    )

if __name__ == "__main__":
    # demo.launch(share=True)
    demo.launch()