add polar maps and script to create video of attenuation maps

mu_map/scripts/create_video.py

+import os
+
+import cv2 as cv
+import numpy as np
+import torch
+
+from mu_map.dataset.default import MuMapDataset
+from mu_map.dataset.transform import PadCropTranform, SequenceTransform
+from mu_map.models.unet import UNet
+from mu_map.util import to_grayscale, COLOR_WHITE
+from mu_map.random_search.cgan import load_params
+from mu_map.vis.slices import join_images
+
+torch.set_grad_enabled(False)
+
+random_search_iter_dir = "cgan_random_search/001"
+
+params = load_params(os.path.join(random_search_iter_dir, "params.json"))
+print(params["normalization"])
+print(params["generator_features"])
+dataset = MuMapDataset(
+    "data/second/",
+    transform_normalization=SequenceTransform(
+        [params["normalization"], PadCropTranform(dim=3, size=32)]
+    ),
+    split_name="test",
+    scatter_correction=False,
+)
+recon, mu_map_ct = dataset[1]
+mu_map_ct = mu_map_ct.squeeze().numpy()
+
+device = torch.device("cpu")
+model = UNet(features=params["generator_features"])
+model.load_state_dict(
+    torch.load(
+        os.path.join(random_search_iter_dir, "snapshots/val_min_generator.pth"), map_location=device
+    )
+)
+
+mu_map_dl = model(recon.unsqueeze(dim=0)).squeeze().numpy()
+mu_map_dl = np.clip(mu_map_dl, 0, mu_map_ct.max())
+
+volumes = [mu_map_ct, mu_map_dl, np.abs(mu_map_dl - mu_map_ct)]
+min_val = 0
+max_val = max(mu_map_ct.max(), mu_map_dl.max())
+print(mu_map_ct.max(), mu_map_dl.max())
+
+fourcc = cv.VideoWriter_fourcc(*"mp4v")
+frame_size = (512, 3 * 512 + 2 * 10)
+print(f"Frame size {frame_size}")
+# video_writer = cv.VideoWriter("mu_map_comparison.mp4", fourcc, 25, frame_size, isColor=False)
+video_writer = cv.VideoWriter("mu_map_comparison.mp4", cv.VideoWriter_fourcc(*"mp4v"), 25, frame_size[::-1], isColor=False)
+for i in range(mu_map_ct.shape[0]):
+    images = map(lambda volume: volume[i], volumes)
+    images = map(
+        lambda img: to_grayscale(img, min_val=min_val, max_val=max_val), images
+    )
+    images = map(lambda img: cv.resize(img, (512, 512)), images)
+    images = list(images)
+
+    txt = f"{str(i):{len(str(mu_map_ct.shape[0]))}}/{mu_map_ct.shape[0]}"
+    cv.putText(images[0], txt, (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, 255, 3)
+
+    image = join_images(images)
+    for i in range(5):
+        video_writer.write(image)
+video_writer.release()