Refactor

app.py

@@ -1,12 +1,12 @@
 import numpy as np
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import matplotlib.pyplot as plt
 import SimpleITK as sitk  # noqa: N813
 import torch
 from monai.transforms import Compose, ScaleIntensityd, SpatialPadd
 from cinema import ConvUNetR
 from pathlib import Path
+from examples.inference.segmentation_sax import plot_segmentations, plot_volume_changes
 import spaces
 
 # cache directories
@@ -25,7 +25,6 @@ def inferece(
     # set device and dtype
     dtype, device = torch.float32, torch.device("cpu")
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
         device = torch.device("cuda")
         if torch.cuda.is_bf16_supported():
             dtype = torch.bfloat16
@@ -81,13 +80,7 @@ def run_inference(trained_dataset, seed, image_id, t_step, progress=gr.Progress(
     transform = Compose(
         [
             ScaleIntensityd(keys=view),
-            SpatialPadd(
-                keys=view,
-                spatial_size=(192, 192, 16),
-                method="end",
-                lazy=True,
-                allow_missing_keys=True,
-            ),
+            SpatialPadd(keys=view, spatial_size=(192, 192, 16), method="end"),
         ]
     )
 
@@ -96,50 +89,8 @@ def run_inference(trained_dataset, seed, image_id, t_step, progress=gr.Progress(
     labels = inferece(images, view, transform, model, progress)
 
     progress(1, desc="Plotting results...")
-    # Create segmentation visualization
-    n_slices, n_frames = labels.shape[-2:]
-    fig1, axs = plt.subplots(n_frames, n_slices, figsize=(n_slices, n_frames), dpi=300)
-    for t in range(n_frames):
-        for z in range(n_slices):
-            axs[t, z].imshow(images[..., z, t], cmap="gray")
-            axs[t, z].imshow(
-                (labels[..., z, t, None] == 1)
-                * np.array([108 / 255, 142 / 255, 191 / 255, 0.6])
-            )
-            axs[t, z].imshow(
-                (labels[..., z, t, None] == 2)
-                * np.array([214 / 255, 182 / 255, 86 / 255, 0.6])
-            )
-            axs[t, z].imshow(
-                (labels[..., z, t, None] == 3)
-                * np.array([130 / 255, 179 / 255, 102 / 255, 0.6])
-            )
-            axs[t, z].set_xticks([])
-            axs[t, z].set_yticks([])
-            if z == 0:
-                axs[t, z].set_ylabel(f"t = {t * t_step}")
-    fig1.suptitle(f"Subject {image_id} in {split} split")
-    axs[0, n_slices // 2].set_title("SAX Slices")
-    fig1.tight_layout()
-    plt.subplots_adjust(wspace=0, hspace=0)
-
-    # Create volume plot
-    xs = np.arange(n_frames) * t_step
-    rv_volumes = np.sum(labels == 1, axis=(0, 1, 2)) * 10 / 1000
-    myo_volumes = np.sum(labels == 2, axis=(0, 1, 2)) * 10 / 1000
-    lv_volumes = np.sum(labels == 3, axis=(0, 1, 2)) * 10 / 1000
-    lvef = (max(lv_volumes) - min(lv_volumes)) / max(lv_volumes) * 100
-    rvef = (max(rv_volumes) - min(rv_volumes)) / max(rv_volumes) * 100
-
-    fig2, ax = plt.subplots(figsize=(4, 4), dpi=120)
-    ax.plot(xs, rv_volumes, color="#6C8EBF", label="RV")
-    ax.plot(xs, myo_volumes, color="#D6B656", label="MYO")
-    ax.plot(xs, lv_volumes, color="#82B366", label="LV")
-    ax.set_xlabel("Frame")
-    ax.set_ylabel("Volume (ml)")
-    ax.set_title(f"LVEF = {lvef:.2f}%, RVEF = {rvef:.2f}%")
-    ax.legend(loc="lower right")
-    fig2.tight_layout()
+    fig1 = plot_segmentations(images, labels, t_step)
+    fig2 = plot_volume_changes(labels, t_step)
 
     return fig1, fig2
 
@@ -187,7 +138,7 @@ with gr.Blocks(
                 maximum=150,
                 step=1,
                 label="Choose an ACDC image, ID is between 1 and 150",
-                value=1,
+                value=150,
             )
             t_step = gr.Slider(
                 minimum=1,