diff --git a/mu_map/scripts/create_video.py b/mu_map/scripts/create_video.py new file mode 100644 index 0000000000000000000000000000000000000000..665d72f0f0a8f239ddd50087beeb48a359df6074 --- /dev/null +++ b/mu_map/scripts/create_video.py @@ -0,0 +1,67 @@ +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() diff --git a/res/polar_maps.png b/res/polar_maps.png new file mode 100644 index 0000000000000000000000000000000000000000..0b60bb8d35522ac24131f63f82e1a7a7797c9582 Binary files /dev/null and b/res/polar_maps.png differ