Blur Labs Open Imaging Models

import torch
import deepinv as dinv
from blurlabs.models import RAM2

device = dinv.utils.get_device()
model = RAM2().to(device)  # downloads the pretrained weights

x = dinv.utils.load_example("butterfly.png", img_size=256).to(device) # dummy image

physics = dinv.physics.Blur(
    filter=dinv.physics.functional.blur.gaussian_blur(sigma=(2.0, 2.0)),
    noise_model=dinv.physics.GaussianNoise(0.01),
    device=device,
)

y = physics(x) # simulate measurements

with torch.no_grad():
    x_hat = model(y, physics)

dinv.utils.plot([x, y, x_hat], ["Original", "Measurements", "RAM 2"])

import torch
import deepinv as dinv
from blurlabs import RAM2

device = dinv.utils.get_device()
torch.manual_seed(0)

model = RAM2().to(device).eval()
x = dinv.utils.load_example("CBSD_0010.png", img_size=(200, 200), device=device)

physics_generator = (
        dinv.physics.generator.MotionBlurGenerator((31, 31), l=2.0, sigma=2.4, device=device)
        + dinv.physics.generator.SigmaGenerator(sigma_min=0.001, sigma_max=0.2, device=device)
)

# Draw noise and blur parameters at random
parameters = physics_generator.step(seed=0)

physics = dinv.physics.BlurFFT(
    img_size=x.shape[1:],
    filter=parameters["filter"],
    noise_model=dinv.physics.GaussianNoise(parameters["sigma"]),
).to(device)

# Generate simulated measurements
y = physics(x)

with torch.no_grad():
    x_lin = physics.A_adjoint(y)
    x_hat = model(y, physics)

psnr = dinv.metric.PSNR()
linear_psnr = psnr(x_lin, x).item()
ram2_psnr = psnr(x_hat, x).item()

dinv.utils.plot(
    {"Measurement (adjoint)": x_lin, "RAM 2": x_hat, "Ground truth": x},
    subtitles=[f"{linear_psnr:.2f} dB", f"{ram2_psnr:.2f} dB", "Reference",],
    figsize=(6, 4),
)

import torch
import deepinv as dinv
from blurlabs import RAM2

device = dinv.utils.get_device()
torch.manual_seed(0)

model = RAM2().to(device).eval()
x = dinv.utils.load_example("JAX_018_011_RGB.tif", device=device)[..., :300, :300]

physics = dinv.physics.Denoising(noise_model=dinv.physics.PoissonGaussianNoise(sigma=0.1, gain=0.1))

y = physics(x)

with torch.no_grad():
    x_hat = model(y, physics)

psnr = dinv.metric.PSNR()
measurement_psnr = psnr(y, x).item()
ram2_psnr = psnr(x_hat, x).item()

dinv.utils.plot(
    {"Noisy measurement": y, "RAM 2": x_hat, "Ground truth": x,},
    subtitles=[f"{measurement_psnr:.2f} dB", f"{ram2_psnr:.2f} dB", "Reference",],
    figsize=(6, 4),
)

import torch
import deepinv as dinv
from blurlabs import RAM2

device = dinv.utils.get_device()
torch.manual_seed(0)

model = RAM2().to(device).eval()
x = dinv.utils.load_example("CT100_256x256_0.pt", device=device)

physics = dinv.physics.Tomography(img_width=x.shape[-1], angles=10, normalize=True, device=device)

y = physics(x)

with torch.no_grad():
    x_fbp = physics.A_dagger(y, fbp=True)
    x_hat = model(y, physics)

psnr = dinv.metric.PSNR()
fbp_psnr = psnr(x_fbp, x).item()
ram2_psnr = psnr(x_hat, x).item()

dinv.utils.plot(
    {"FBP": x_fbp, "RAM 2": x_hat, "Ground truth": x},
    subtitles=[f"{fbp_psnr:.2f} dB", f"{ram2_psnr:.2f} dB", "Reference"],
    figsize=(6, 4),
)

import torch
from torch.utils.data import DataLoader

import deepinv as dinv
from blurlabs import RAM2

device = dinv.utils.get_device()
torch.manual_seed(0)

cache = dinv.utils.get_cache_home()
filename = "file_brain_AXT2_200_2000341.h5"
measurement_root = cache / "fastmri_brain_multicoil_test"
target_root = cache / "fastmri_brain_multicoil_test_full"

dinv.datasets.download_archive(dinv.utils.get_image_url("demo_fastmri_brain_multicoil_test_file_brain_AXT2_200_2000341.h5"), measurement_root / filename,)
dinv.datasets.download_archive(dinv.utils.get_image_url("demo_fastmri_brain_multicoil_test_full_file_brain_AXT2_200_2000341.h5"), target_root / filename,)

dataset = dinv.datasets.FastMRISliceDataset(
    measurement_root,
    target_root=target_root,
    slice_index="middle",
    transform=dinv.datasets.MRISliceTransform(),
    use_dict_output=True,
)

batch = next(iter(DataLoader(dataset)))
x = batch["x"].to(device)
y = batch["y"].to(device)
mask = batch["params"]["mask"].to(device)

# DeepInverse estimates these maps with SigPy's ESPIRiT implementation.
coil_maps = dinv.physics.MultiCoilMRI.estimate_coil_maps(y, calib_size=24, espirit_crop=0.85).to(device)

physics = dinv.physics.MultiCoilMRI(img_size=y.shape[-2:], mask=mask, coil_maps=coil_maps, device=device)

model = RAM2().to(device).eval()

with torch.no_grad():
    x_zf = physics.A_adjoint(y)
    physics.phase_correct_maps(x_zf)
    x_zf = physics.A_adjoint(y)

    # Conjugate-gradient SENSE reconstruction with the same calibrated physics.
    x_sense = physics.A_dagger(y)

    # Raw FastMRI k-space is below RAM 2's minimum noise scale. Normalize the
    # measurement for inference, then restore the acquisition's original scale.
    input_scale = x_zf.abs().amax().item()
    y_normalized = y / input_scale

    # Estimate the noise level in the measurement domain
    noise_sigma = dinv.models.WaveletNoiseEstimator()(y_normalized).item()

    x_hat = model(y_normalized, physics, sigma=noise_sigma) * input_scale

# Match the FastMRI target convention: crop, then take the RSS magnitude.
x_zf = physics.rss(physics.crop(x_zf, shape=x.shape), multicoil=False)
x_sense = physics.rss(physics.crop(x_sense, shape=x.shape), multicoil=False)
x_hat = physics.rss(physics.crop(x_hat, shape=x.shape), multicoil=False)

psnr = dinv.metric.PSNR(max_pixel=None)

zf_psnr = psnr(x_zf, x).item()
sense_psnr = psnr(x_sense, x).item()
ram2_psnr = psnr(x_hat, x).item()

print(f"Estimated sigma: {noise_sigma:.4f}")

dinv.utils.plot(
    {"Zero-filled": x_zf, "CG-SENSE": x_sense, "RAM 2": x_hat, "Ground truth": x},
    subtitles=[f"{zf_psnr:.2f} dB", f"{sense_psnr:.2f} dB", f"{ram2_psnr:.2f} dB", "Reference"],
    figsize=(8, 4),
)