robreg

import os
import sys
from pathlib import Path
from urllib.request import urlretrieve

# Allow unsupported MPS ops to fall back to CPU so the notebook can still test MPS.
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

import nibabel as nib
import torch

REPO_ROOT = Path.cwd().resolve()
if not (REPO_ROOT / "neuroreg").exists():
    REPO_ROOT = REPO_ROOT.parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from neuroreg import robreg
from neuroreg.image import map
from neuroreg.image.map import map_r2r
from neuroreg.transforms import LTA, affine_dist
from neuroreg.transforms.matrices import get_affine

DATA_DIR = Path(".")
orig_path = DATA_DIR / "140_orig.mgz"
trg_path = DATA_DIR / "140_trg.mgz"
reg_path = DATA_DIR / "140_reg.mgz"
lta_path = DATA_DIR / "140.lta"
gt_lta_path = DATA_DIR / "140_ground_truth.lta"


device = "cpu"
if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
    print("MPS detected; unsupported ops will fall back to CPU via PYTORCH_ENABLE_MPS_FALLBACK=1.")

print(f"Running on device: {device}")

# Download a test image.
url = "https://surfer.nmr.mgh.harvard.edu/pub/data/tutorial_data/buckner_data/tutorial_subjs/140/mri/orig.mgz"
if not orig_path.exists():
    urlretrieve(url, orig_path)
else:
    print(f"File {orig_path} already exists, skipping download")

# Load the image and create a small synthetic rigid transform.
img = nib.load(str(orig_path))
idata = torch.from_numpy(img.get_fdata()).float()
img_affine = torch.from_numpy(img.affine).float()

translation_mm = torch.tensor([3.1, 0.6, 1.2])
rotation_rad = torch.tensor([0.08, 0.04, 0.02])
v2v_gt = get_affine(translation=translation_mm, rotvec=rotation_rad).float()
r2r_gt = img_affine @ v2v_gt @ torch.linalg.inv(img_affine)

print(f"Applied Vox-to-Vox matrix:\n{v2v_gt}")
print(f"Ground-truth RAS-to-RAS matrix:\n{r2r_gt}")

# Map the image and save the moved volume.
mapped_data = map(idata, transform=v2v_gt, is_torch_mat=False)
target_img = nib.MGHImage(mapped_data.squeeze().numpy(), img.affine, img.header)
target_img.to_filename(str(trg_path))

# Save the ground-truth transform as an LTA.
LTA.from_matrix(r2r_gt.numpy(), str(orig_path), img, str(trg_path), target_img).write(str(gt_lta_path))
print(f"Saved synthetic target image to: {trg_path}")
print(f"Saved ground-truth LTA to: {gt_lta_path}")

# Register the synthetic pair.
# By default, robreg returns an RAS-to-RAS transform.
Mr2r = robreg(
    str(orig_path),
    str(trg_path),
    device=device,
)
Mr2r_cpu = Mr2r.detach().cpu()

recovered_data = map_r2r(
    idata,
    Mr2r_cpu.float(),
    source_affine=img_affine,
    target_affine=img_affine,
    target_shape=tuple(int(v) for v in img.shape[:3]),
    mode="bilinear",
)
recovered_img = nib.MGHImage(recovered_data.detach().cpu().numpy(), img.affine, img.header)
recovered_img.to_filename(str(reg_path))
LTA.from_matrix(Mr2r_cpu.numpy(), str(orig_path), img, str(trg_path), target_img).write(str(lta_path))

recovered_lta = LTA.read(str(lta_path))
affine_distance = float(affine_dist(Mr2r_cpu.float(), r2r_gt.float(), radius=100.0))
saved_lta_distance = float(
    affine_dist(torch.from_numpy(recovered_lta.r2r()).float(), r2r_gt.float(), radius=100.0)
)

print(f"\nRecovered RAS-to-RAS matrix:\n{Mr2r_cpu}")
print(f"\nSaved LTA RAS-to-RAS matrix:\n{torch.from_numpy(recovered_lta.r2r()).float()}")
print(f"\nAffine distance to the ground-truth transform (r=100 mm): {affine_distance:.4f}")
print(f"Affine distance for the saved LTA (r=100 mm): {saved_lta_distance:.4f}")
print(f"Mapped output saved to: {reg_path}")
print(f"Recovered LTA saved to: {lta_path}")