add script to visualize and test the alignment of mu maps and reconstructions

mu_map/vis/alignment.py

+import argparse
+import math
+import os
+
+import cv2 as cv
+import numpy as np
+import pandas as pd
+
+from mu_map.data.prepare import headers
+from mu_map.file.dicom import load_dcm_img
+
+parser = argparse.ArgumentParser(
+    description="Visualize different slice-alignment types of reconstructions and attenuation maps",
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+)
+parser.add_argument("--recon", type=str, help="the reconstruction")
+parser.add_argument("--mu_map", type=str, help="the attenuation map")
+parser.add_argument(
+    "--dataset_dir",
+    type=str,
+    help="if provided images are loaded from the dataset with the given id instead of using the --recon and --mu_map parameters",
+)
+parser.add_argument(
+    "--images_dir",
+    default="images",
+    type=str,
+    help="sub-directory of the dataset containing the dicom images",
+)
+parser.add_argument(
+    "--csv",
+    default="meta.csv",
+    type=str,
+    help="csv file in the dataset containing meta information",
+)
+parser.add_argument(
+    "--id", type=int, default=1, help="the id of the image in the dataset to display"
+)
+parser.add_argument(
+    "--alignment",
+    choices=["first", "center", "center_left", "center_right", "last"],
+    default="center_left",
+    help="choose how the reconstruction and mu map are aligned - at their first, center or last slice",
+)
+parser.add_argument(
+    "-d",
+    "--direction",
+    type=int,
+    default=0,
+    help="make sure that slices are ordered in this direction according to the `SpacingBetweenSlices` parameter - 0 = no re-ordering, 1 = in positive direction, -1 = in negative direction",
+)
+parser.add_argument(
+    "--ac", action="store_true", help="use the attenuation corrected reconstruction"
+)
+args = parser.parse_args()
+
+
+if args.dataset_dir and args.id:
+    meta = pd.read_csv(os.path.join(args.dataset_dir, args.csv))
+    row = meta[meta[headers.id] == args.id].iloc[0]
+
+    file_recon = headers.file_recon_ac_nsc if args.ac else headers.file_recon_nac_nsc
+    args.recon = os.path.join(args.dataset_dir, args.images_dir, row[file_recon])
+    args.mu_map = os.path.join(
+        args.dataset_dir, args.images_dir, row[headers.file_mu_map]
+    )
+
+
+recon = load_dcm_img(args.recon, direction=args.direction)
+mu_map = load_dcm_img(args.mu_map, direction=args.direction)
+
+diff_slices = recon.shape[0] - mu_map.shape[0]
+diff_slices_h = diff_slices / 2
+if args.alignment == "center_left" or args.alignment == "center":
+    diff_slices_h = math.ceil(diff_slices_h)
+    args.alignment = "center"
+elif args.alignment == "center_right":
+    diff_slices_h = math.floor(diff_slices_h)
+    args.alignment = "center"
+
+
+def display_slice(img, _slice: int):
+    img_slice = (img[_slice] - img.min()) / (img.max() - img.min())
+    img_slice = (255 * img_slice).astype(np.uint8)
+    img_slice = cv.resize(img_slice, (512, 512))
+    return img_slice
+
+
+def combine(img_recon, img_mu_map, slice_recon, slice_mu_map):
+    img_slice_recon = display_slice(img_recon, slice_recon)
+    img_slice_mu_map = display_slice(img_mu_map, slice_mu_map)
+
+    img_slice_recon = cv.cvtColor(img_slice_recon, cv.COLOR_GRAY2BGR)
+    img_slice_mu_map = cv.cvtColor(img_slice_mu_map, cv.COLOR_GRAY2BGR)
+
+    img_com_1 = img_slice_mu_map.copy()
+    img_com_2 = cv.applyColorMap(img_slice_recon, cv.COLORMAP_INFERNO)
+    img_com = cv.addWeighted(img_com_1, 0.8, img_com_2, 0.4, 0.0)
+
+    img_slice_recon = cv.putText(
+        img_slice_recon,
+        f"{1 + slice_recon:02d}/{img_recon.shape[0]}",
+        (0, 30),
+        cv.FONT_HERSHEY_SIMPLEX,
+        1,
+        255,
+        3,
+    )
+    img_slice_mu_map = cv.putText(
+        img_slice_mu_map,
+        f"{1 + slice_mu_map:02d}/{img_mu_map.shape[0]}",
+        (0, 30),
+        cv.FONT_HERSHEY_SIMPLEX,
+        1,
+        255,
+        3,
+    )
+
+    space = np.full((img_com.shape[0], 10, 3), 239, np.uint8)
+    return np.hstack((img_slice_mu_map, img_com, img_slice_recon))
+
+
+wname_align = args.alignment[0].upper() + args.alignment[1:]
+wname_correction = "AC" if args.ac else "nAC"
+wname = f"{wname_align} Alignment - {wname_correction}"
+cv.namedWindow(wname, cv.WINDOW_NORMAL)
+cv.resizeWindow(wname, 1600, 900)
+
+timeout = 0
+slice_recon = 0
+while True:
+    if args.alignment == "first":
+        slice_mu_map = min(slice_recon, mu_map.shape[0] - 1)
+    elif args.alignment == "center":
+        slice_mu_map = min(max(0, slice_recon - diff_slices_h), mu_map.shape[0] - 1)
+    elif args.alignment == "last":
+        slice_mu_map = min(max(0, slice_recon - diff_slices), mu_map.shape[0] - 1)
+    else:
+        raise ValueError(f"Unknown alignment {args.alignment}")
+
+    img = combine(recon, mu_map, slice_recon, slice_mu_map)
+    cv.imshow(wname, img)
+    key = cv.waitKey(timeout)
+    if key == ord("q"):
+        break
+    elif key == ord("p"):
+        timeout = 0 if timeout > 0 else 100
+    elif key == ord("a"):
+        slice_recon = (slice_recon - 1) % recon.shape[0]
+        continue
+
+    slice_recon = (slice_recon + 1) % recon.shape[0]