add visualization for image registration

mu_map/vis/registration.py

+import argparse
+
+import cv2 as cv
+import numpy as np
+import SimpleITK as sitk
+
+from mu_map.file.dicom import load_dcm
+from mu_map.file.convert import dicom_to_nrrd, nrrd_to_dicom
+from mu_map.file.nrrd import (
+    NRRD,
+    load_nrrd,
+    change_spacing,
+    pad_crop,
+)
+
+from mu_map.vis.util import to_grayscale
+
+
+def print_info(nrrd: NRRD, title: str):
+    """
+    Print information (origin, direction, spacing and size) about an image in NRRD format.
+
+    Parameter
+    ---------
+    title: str
+    nrrd: NRRD
+    """
+    print(f"{title}:")
+    print(f" -    Origin: {nrrd.header.GetOrigin()}")
+    print(f" - Direction: {nrrd.header.GetDirection()}")
+    print(f" -   Spacing: {nrrd.header.GetSpacing()}")
+    print(f" -      Size: {nrrd.header.GetSize()}")
+
+
+def clip_to_percentage(img: np.ndarray, percentage: float) -> np.ndarray:
+    """
+    Clip an image to a percentage of its maximum value.
+
+    Parameter
+    ---------
+    img: np.ndarray
+    percentage: float
+        in range of [0, 1]
+
+    Returns
+    -------
+    np.ndarray
+    """
+    return np.clip(img, a_min=img.min(), a_max=img.max() * percentage)
+
+
+def window_img(img: np.ndarray, w_center: float, w_width: float) -> np.ndarray:
+    """
+    Window an image by clipping to (w_center - w_width, w_center + w_width)
+
+    Parameter
+    ---------
+    img: np.ndarray
+    w_center: float
+    w_width: float
+
+    Returns
+    -------
+    np.ndarray
+    """
+    return np.clip(img, w_center - w_width, w_center + w_width)
+
+
+parser = argparse.ArgumentParser(
+    description="Visualize the registration of a CT/MuMap to a SPECT reconstruction",
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+)
+parser.add_argument(
+    "--mu_map",
+    type=str,
+    required=True,
+    help="the file/directory containing the CT/MuMap",
+)
+parser.add_argument(
+    "--recon",
+    type=str,
+    required=True,
+    help="the file containing the SPECT reconstruction",
+)
+parser.add_argument(
+    "--verbose",
+    action="store_true",
+    help="print additional information about the images",
+)
+parser.add_argument(
+    "--percentage",
+    type=float,
+    default=0.9,
+    help="normalize the reconstruction by clipping above this percentage of its maximum value",
+)
+parser.add_argument(
+    "--window",
+    type=int,
+    nargs=2,
+    help="windowing for the CT/MuMap as (center, width) - e.g., 200 600",
+)
+parser.add_argument(
+    "--save",
+    type=str,
+    help="save the resulting animation as a GIF file with the given name",
+)
+args = parser.parse_args()
+
+recon = load_nrrd(args.recon)
+mu_map = load_nrrd(args.mu_map)
+
+if args.verbose:
+    print("μ Map:")
+    print_info(mu_map)
+    print()
+    print("Recon:")
+    print_info(recon)
+    print()
+
+recon = sitk.Resample(recon.header, mu_map.header, sitk.Transform(), sitk.sitkLinear)
+recon.SetOrigin(mu_map.header.GetOrigin())
+recon = NRRD(recon, sitk.GetArrayFromImage(recon))
+
+if args.verbose:
+    print("Recon - Resampled:")
+    print_info(recon)
+    print()
+
+
+img_recon = to_grayscale(clip_to_percentage(recon.image, args.percentage))
+
+img_mumap = mu_map.image
+w_center = 0 if args.window is None else args.window[0]
+w_width = (
+    max(img_mumap.max(), -img_mumap.min()) if args.window is None else args.window[1]
+)
+img_mumap = to_grayscale(window_img(img_mumap, w_center, w_width))
+
+imgs_recon = [
+    cv.applyColorMap(img_recon[i], cv.COLORMAP_INFERNO)
+    for i in range(img_recon.shape[0])
+]
+imgs_mumap = [
+    img_mumap[i].repeat(3).reshape((*img_mumap.shape[1:], 3))
+    for i in range(img_mumap.shape[0])
+]
+imgs = zip(imgs_mumap, imgs_recon)
+imgs = map(lambda _imgs: cv.addWeighted(_imgs[0], 0.8, _imgs[1], 0.8, 0.0), imgs)
+imgs = map(lambda img: cv.resize(img, (512, 512)), imgs)
+imgs = list(imgs)
+
+if args.save:
+    import os
+
+    _, ext = os.path.splitext(args.save)
+    if ext == ".gif":
+        import imageio
+
+        imageio.mimsave(
+            args.save,
+            list(
+                map(
+                    lambda img: cv.cvtColor(
+                        cv.resize(img, (256, 256)), cv.COLOR_BGR2RGB
+                    ),
+                    imgs,
+                )
+            ),
+            duration=100,
+        )
+    else:
+        recon_dcm = load_dcm(args.recon)
+        nrrd_to_dicom(recon, ref=recon_dcm).write(args.save)
+
+_slice = 0
+while True:
+    img = imgs[_slice]
+    _slice = (_slice + 1) % img_recon.shape[0]
+
+    cv.imshow("Slice", img)
+    if cv.waitKey(100) == ord("q"):
+        break