Refactor MAE example

Author: mathpluscode

cinema/examples/inference/mae.py

@@ -12,35 +12,22 @@
 
 
 def plot_mae_reconstruction(
-    batch: dict[str, torch.Tensor],
-    pred_dict: dict[str, torch.Tensor],
-    enc_mask_dict: dict[str, torch.Tensor],
-    patch_size_dict: dict[str, tuple[int, ...]],
-    grid_size_dict: dict[str, tuple[int, ...]],
-    sax_slices: int,
+    image_dict: dict[str, torch.Tensor],
+    reconstructed_dict: dict[str, torch.Tensor],
+    masks_dict: dict[str, torch.Tensor],
 ) -> plt.Figure:
     """Plot MAE reconstruction."""
+    sax_slices = image_dict["sax"].shape[-1]
     n_rows = sax_slices + 3
     n_cols = 4
     fig, axs = plt.subplots(n_rows, n_cols, figsize=(n_cols * 2, n_rows * 2), dpi=300)
     for i, view in enumerate(["lax_2c", "lax_3c", "lax_4c", "sax"]):
-        patches = patchify(image=batch[view], patch_size=patch_size_dict[view])
-        patches[enc_mask_dict[view]] = pred_dict[view]
-        masks = torch.zeros_like(patches)
-        masks[enc_mask_dict[view]] = 1
-        masks = unpatchify(masks, patch_size=patch_size_dict[view], grid_size=grid_size_dict[view])
-        masks = masks[0, 0]
-        reconstructed = unpatchify(
-            patches,
-            patch_size=patch_size_dict[view],
-            grid_size=grid_size_dict[view],
-        )
-        reconstructed = reconstructed[0, 0].numpy()
-        image = batch[view][0, 0].numpy()
+        masks = masks_dict[view]
+        reconstructed = reconstructed_dict[view]
+        image = image_dict[view]
         error = np.abs(reconstructed - image)
 
         if view == "sax":
-            reconstructed = reconstructed[..., :sax_slices]
             for j in range(sax_slices):
                 axs[3 + j, 0].set_ylabel(f"SAX slice {j}")
                 axs[3 + j, 0].imshow(image[..., j], cmap="gray")
@@ -68,15 +55,42 @@ def plot_mae_reconstruction(
     return fig
 
 
+def reconstruct_images(
+    batch: dict[str, torch.Tensor],
+    pred_dict: dict[str, torch.Tensor],
+    enc_mask_dict: dict[str, torch.Tensor],
+    patch_size_dict: dict[str, tuple[int, ...]],
+    grid_size_dict: dict[str, tuple[int, ...]],
+    sax_slices: int,
+) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+    """Reconstruct images from predicted patches."""
+    reconstructed_dict = {}
+    masks_dict = {}
+    for view in ["lax_2c", "lax_3c", "lax_4c", "sax"]:
+        patches = patchify(image=batch[view], patch_size=patch_size_dict[view])
+        patches[enc_mask_dict[view]] = pred_dict[view]
+        masks = torch.zeros_like(patches)
+        masks[enc_mask_dict[view]] = 1
+        masks = unpatchify(masks, patch_size=patch_size_dict[view], grid_size=grid_size_dict[view])
+        reconstructed = unpatchify(
+            patches,
+            patch_size=patch_size_dict[view],
+            grid_size=grid_size_dict[view],
+        )
+        reconstructed_dict[view] = reconstructed.detach().cpu().numpy()[0, 0]
+        masks_dict[view] = masks.detach().cpu().numpy()[0, 0]
+    reconstructed_dict["sax"] = reconstructed_dict["sax"][..., :sax_slices]
+    masks_dict["sax"] = masks_dict["sax"][..., :sax_slices]
+    return reconstructed_dict, masks_dict
+
+
 def run(device: torch.device, dtype: torch.dtype) -> None:
     """Run MAE reconstruction."""
     t = 25  # which time frame to use
 
     # load model
     model = CineMA.from_pretrained()
     model.eval()
-    patch_size_dict = model.dec_patch_size_dict
-    grid_size_dict = {k: v.patch_embed.grid_size for k, v in model.enc_down_dict.items()}
     model.to(device)
 
     # load sample data and form a batch of size 1
@@ -95,36 +109,41 @@ def run(device: torch.device, dtype: torch.dtype) -> None:
     )
     # (x, y, z, t) for SAX and (x, y, 1, t) for LAX
     exp_dir = Path(__file__).parent.parent.resolve()
-    sax_image = torch.from_numpy(
-        np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_sax.nii.gz")))
-    )
-    lax_2c_image = torch.from_numpy(
-        np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_2c.nii.gz")))
-    )
-    lax_3c_image = torch.from_numpy(
-        np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_3c.nii.gz")))
-    )
-    lax_4c_image = torch.from_numpy(
-        np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_4c.nii.gz")))
-    )
-    sax_slices = sax_image.shape[-2]
-    batch = {
-        "sax": sax_image[None, ..., t],
-        "lax_2c": lax_2c_image[None, ..., 0, t],
-        "lax_3c": lax_3c_image[None, ..., 0, t],
-        "lax_4c": lax_4c_image[None, ..., 0, t],
+    sax_image = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_sax.nii.gz")))
+    lax_2c_image = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_2c.nii.gz")))
+    lax_3c_image = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_3c.nii.gz")))
+    lax_4c_image = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_4c.nii.gz")))
+
+    image_dict = {
+        "sax": sax_image[..., t],
+        "lax_2c": lax_2c_image[..., 0, t],
+        "lax_3c": lax_3c_image[..., 0, t],
+        "lax_4c": lax_4c_image[..., 0, t],
     }
-    batch = transform(batch)
-    print(f"SAX view had originally {sax_image.shape[-2]} slices, now zero-padded to {batch['sax'].shape[-1]} slices.")  # noqa: T201
-    batch = {k: v[None, ...].to(device=device, dtype=dtype) for k, v in batch.items()}
+    batch = {k: torch.from_numpy(v[None, ...]) for k, v in image_dict.items()}
 
     # forward
+    sax_slices = batch["sax"].shape[-1]
+    batch = transform(batch)
+    batch = {k: v[None, ...].to(device=device, dtype=dtype) for k, v in batch.items()}
     with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
         _, pred_dict, enc_mask_dict, _ = model(batch, enc_mask_ratio=0.75)
+        grid_size_dict = {k: v.patch_embed.grid_size for k, v in model.enc_down_dict.items()}
+        reconstructed_dict, masks_dict = reconstruct_images(
+            batch,
+            pred_dict,
+            enc_mask_dict,
+            model.dec_patch_size_dict,
+            grid_size_dict,
+            sax_slices,
+        )
 
     # visualize
-    batch = {k: v.detach().cpu() for k, v in batch.items()}
-    fig = plot_mae_reconstruction(batch, pred_dict, enc_mask_dict, patch_size_dict, grid_size_dict, sax_slices)
+    fig = plot_mae_reconstruction(
+        image_dict,
+        reconstructed_dict,
+        masks_dict,
+    )
     fig.savefig("mae_reconstruction.png", dpi=300, bbox_inches="tight")
     plt.show(block=False)