Add results cache

app.py

@@ -176,6 +176,14 @@ def mae_inference(
 
 
 def mae(image_id, mask_ratio, progress=gr.Progress()):
+    # Create file path for saving MAE reconstruction plot
+    mae_path = cache_dir / f"mae_image{image_id}_mask{mask_ratio * 100:.0f}.png"
+
+    # Check if result already exists
+    if mae_path.exists():
+        progress(1, desc="Loading cached result...")
+        return str(mae_path)
+
     t = 4  # which time frame to use
     progress(0, desc="Downloading model...")
     model = CineMA.from_pretrained()
@@ -221,7 +229,6 @@ def mae(image_id, mask_ratio, progress=gr.Progress()):
     masks_dict["sax"] = np.transpose(masks_dict["sax"], (1, 0, 2))
 
     # Plot MAE reconstruction and save to file
-    mae_path = cache_dir / f"mae_image{image_id}_mask{mask_ratio:.2f}.png"
     plot_mae_reconstruction(batch, reconstructed_dict, masks_dict, mae_path)
 
     return str(mae_path)
@@ -298,6 +305,21 @@ def segmentation_sax_inference(
 
 
 def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progress()):
+    # Create file paths for saving plots
+    seg_path = (
+        cache_dir
+        / f"sax_segmentation_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.gif"
+    )
+    vol_path = (
+        cache_dir
+        / f"sax_volume_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.png"
+    )
+
+    # Check if results already exist
+    if seg_path.exists() and vol_path.exists():
+        progress(1, desc="Loading cached results...")
+        return (str(seg_path), str(vol_path))
+
     # Fixed parameters
     view = "sax"
     split = "train" if image_id <= 100 else "test"
@@ -343,18 +365,8 @@ def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progre
 
     progress(1, desc="Inference finished. Plotting ...")
 
-    # Create file paths for saving plots
-    seg_path = (
-        cache_dir
-        / f"sax_segmentation_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.gif"
-    )
-    vol_path = (
-        cache_dir
-        / f"sax_volume_{trained_dataset}_image{image_id}_seed{seed}_tstep{t_step}.png"
-    )
-
     # Plot segmentations and volume changes with file paths
-    plot_segmentations_sax(images, labels, seg_path)
+    plot_segmentations_sax(images, labels, t_step, seg_path)
     plot_volume_changes_sax(labels, t_step, vol_path)
 
     return (str(seg_path), str(vol_path))
@@ -474,6 +486,15 @@ def segmentation_lax_inference(
 
 
 def segmentation_lax(seed, image_id, progress=gr.Progress()):
+    # Create file paths for saving plots
+    seg_path = cache_dir / f"lax_segmentation_image{image_id}_seed{seed}.gif"
+    vol_path = cache_dir / f"lax_volume_image{image_id}_seed{seed}.png"
+
+    # Check if results already exist
+    if seg_path.exists() and vol_path.exists():
+        progress(1, desc="Loading cached results...")
+        return (str(seg_path), str(vol_path))
+
     # Fixed parameters
     trained_dataset = "mnms2"
     view = "lax_4c"
@@ -502,11 +523,9 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
     progress(1, desc="Inference finished. Plotting ...")
 
     # Plot segmentations and save as GIF
-    seg_path = cache_dir / f"lax_segmentation_image{image_id}_seed{seed}.gif"
     plot_segmentations_lax(images, labels, seg_path)
 
     # Plot volume changes and save as figure
-    vol_path = cache_dir / f"lax_volume_image{image_id}_seed{seed}.png"
     plot_volume_changes_lax(labels, vol_path)
 
     return (str(seg_path), str(vol_path))
@@ -647,6 +666,17 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
     view = "lax_2c" if view == "LAX 2C" else "lax_4c"
     method = method.lower()
 
+    # Create file paths for saving plots
+    landmark_path = (
+        cache_dir / f"landmark_{view}_image{image_id}_{method}_seed{seed}.gif"
+    )
+    lv_path = cache_dir / f"lv_{view}_image{image_id}_{method}_seed{seed}.png"
+
+    # Check if results already exist
+    if landmark_path.exists() and lv_path.exists():
+        progress(1, desc="Loading cached results...")
+        return (str(landmark_path), str(lv_path))
+
     # Download and load model
     progress(0, desc="Downloading model...")
     if method == "heatmap":
@@ -686,13 +716,9 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
     progress(1, desc="Inference finished. Plotting ...")
 
     # Plot landmarks in GIF
-    landmark_path = (
-        cache_dir / f"landmark_{view}_image{image_id}_{method}_seed{seed}.gif"
-    )
     plot_landmarks(images, coords, landmark_path)
 
     # Plot LV change in PNG
-    lv_path = cache_dir / f"lv_{view}_image{image_id}_{method}_seed{seed}.png"
     plot_lv(coords, lv_path)
 
     return (str(landmark_path), str(lv_path))
@@ -791,7 +817,7 @@ with gr.Blocks(
         # CineMA: A Foundation Model for Cine Cardiac MRI 🎥🫀
 
         🚀 The following demos showcase the capabilities of CineMA in multiple tasks.<br>
-        ⏱️ The examples may take 10-60 seconds to download data and model, perform inference, and render plots.<br>
+        ⏱️ The examples may take 10-60 seconds, if not cached, to download data and model, perform inference, and render plots.<br>
         🔗 For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
         """
     )

requirements.txt

@@ -17,6 +17,6 @@ scikit-learn==1.6.1
 scipy==1.15.2
 spaces==0.36.0
 timm==1.0.15
-git+https://github.com/mathpluscode/CineMA@7e86ffc7ddf06ad7283915ee143ed808c0f59576#egg=cinema
+git+https://github.com/mathpluscode/CineMA@0d1afe864d4b4c348a993b8b2b790adf8581bc03#egg=cinema
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.5.1