Display gifs

app.py

@@ -213,15 +213,18 @@ def mae(image_id, mask_ratio, progress=gr.Progress()):
     batch, reconstructed_dict, masks_dict = mae_inference(
         batch, transform, model, mask_ratio
     )
-    progress(1, desc="Plotting results...")
+    progress(1, desc="Inference finished. Plotting ...")
 
-    fig = plot_mae_reconstruction(
-        batch,
-        reconstructed_dict,
-        masks_dict,
-    )
-    plt.close(fig)
-    return fig
+    # (y, x, z) -> (x, y, z)
+    batch["sax"] = np.transpose(batch["sax"], (1, 0, 2))
+    reconstructed_dict["sax"] = np.transpose(reconstructed_dict["sax"], (1, 0, 2))
+    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)
 
 
 def mae_tab():
@@ -229,14 +232,16 @@ def mae_tab():
         gr.Markdown(
             """
             This page demonstrates the masking and reconstruction process of the masked autoencoder. The model was trained with a mask ratio of 0.75 over 74,000 studies.
-
-            Visualisation may take a few seconds as we download the model weights, process the data, and render the plots.
             """
         )
         with gr.Row():
             with gr.Column(scale=5):
                 gr.Markdown("## Reconstruction")
-                plot = gr.Plot(show_label=False)
+                plot = gr.Image(
+                    show_label=False,
+                    type="filepath",
+                    label="Masked Autoencoder Reconstruction",
+                )
             with gr.Column(scale=3):
                 gr.Markdown("## Data Settings")
                 image_id = gr.Slider(
@@ -254,6 +259,7 @@ def mae_tab():
                     value=0.75,
                 )
                 run_button = gr.Button("Run masked autoencoder", variant="primary")
+
         run_button.click(
             fn=mae,
             inputs=[image_id, mask_ratio],
@@ -331,13 +337,27 @@ def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progre
     images = images[..., ::t_step]
     labels = segmentation_sax_inference(images, view, transform, model, progress)
 
-    progress(1, desc="Plotting results...")
-    fig1 = plot_segmentations_sax(images, labels, t_step)
-    fig2 = plot_volume_changes_sax(labels, t_step)
-    result = (fig1, fig2)
-    plt.close(fig1)
-    plt.close(fig2)
-    return result
+    # (y, x, z, t) -> (x, y, z, t)
+    images = np.transpose(images, (1, 0, 2, 3))
+    labels = np.transpose(labels, (1, 0, 2, 3))
+
+    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_volume_changes_sax(labels, t_step, vol_path)
+
+    return (str(seg_path), str(vol_path))
 
 
 def segmentation_sax_tab():
@@ -345,8 +365,6 @@ def segmentation_sax_tab():
         gr.Markdown(
             """
             This page demonstrates the segmentation of cardiac structures in the short-axis (SAX) view.
-
-            Visualisation may take dozens of seconds to update as we download model checkpoints, process multiple time frames sequentially, and generate the final plots.
             """
         )
 
@@ -381,7 +399,7 @@ def segmentation_sax_tab():
                     maximum=10,
                     step=1,
                     label="Choose the gap between time frames",
-                    value=2,
+                    value=3,
                 )
             with gr.Column(scale=3):
                 gr.Markdown("## Model Settings")
@@ -401,16 +419,22 @@ def segmentation_sax_tab():
 
         with gr.Row():
             with gr.Column():
-                gr.Markdown("## Ventricle and Myocardium Segmentation")
-                segmentation_plot = gr.Plot(show_label=False)
+                segmentation_gif = gr.Image(
+                    show_label=True,
+                    type="filepath",
+                    label="Ventricle and Myocardium Segmentation",
+                )
             with gr.Column():
-                gr.Markdown("## Volume Estimation")
-                volume_plot = gr.Plot(show_label=False)
+                volume_plot = gr.Image(
+                    show_label=True,
+                    type="filepath",
+                    label="Ejection Fraction Estimation",
+                )
 
         run_button.click(
             fn=segmentation_sax,
             inputs=[trained_dataset, seed, image_id, t_step],
-            outputs=[segmentation_plot, volume_plot],
+            outputs=[segmentation_gif, volume_plot],
         )
     return sax_interface
 
@@ -475,13 +499,17 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
     )
     labels = segmentation_lax_inference(images, view, transform, model, progress)
 
-    progress(1, desc="Plotting results...")
-    fig1 = plot_segmentations_lax(images, labels)
-    fig2 = plot_volume_changes_lax(labels)
-    result = (fig1, fig2)
-    plt.close(fig1)
-    plt.close(fig2)
-    return result
+    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))
 
 
 def segmentation_lax_tab():
@@ -489,8 +517,6 @@ def segmentation_lax_tab():
         gr.Markdown(
             """
             This page demonstrates the segmentation of cardiac structures in the long-axis (LAX) four-chamber (4C) view.
-
-            Visualisation may take a few seconds to update as we download model checkpoints, process multiple time frames, and generate the final plots.
             """
         )
 
@@ -533,16 +559,22 @@ def segmentation_lax_tab():
 
         with gr.Row():
             with gr.Column():
-                gr.Markdown("## Ventricle and Myocardium Segmentation")
-                segmentation_plot = gr.Plot(show_label=False)
+                segmentation_gif = gr.Image(
+                    show_label=True,
+                    type="filepath",
+                    label="Ventricle and Myocardium Segmentation",
+                )
             with gr.Column():
-                gr.Markdown("## Ejection Fraction Prediction")
-                volume_plot = gr.Plot(show_label=False)
+                volume_plot = gr.Image(
+                    show_label=True,
+                    type="filepath",
+                    label="Ejection Fraction Prediction",
+                )
 
         run_button.click(
             fn=segmentation_lax,
             inputs=[seed, image_id],
-            outputs=[segmentation_plot, volume_plot],
+            outputs=[segmentation_gif, volume_plot],
         )
     return lax_interface
 
@@ -651,12 +683,19 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
     else:
         raise ValueError(f"Invalid method: {method}")
 
-    landmark_fig = plot_landmarks(images, coords)
-    lv_fig = plot_lv(coords)
-    result = (landmark_fig, lv_fig)
-    plt.close(landmark_fig)
-    plt.close(lv_fig)
-    return result
+    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))
 
 
 def landmark_tab():
@@ -664,8 +703,6 @@ def landmark_tab():
         gr.Markdown(
             """
             This page demonstrates landmark localisation in the long-axis (LAX) two-chamber (2C) and four-chamber (4C) views.
-
-            Visualisation may take a few seconds to update as we download model checkpoints, process multiple time frames, and generate the final plots.
             """
         )
 
@@ -679,7 +716,7 @@ def landmark_tab():
 
                 ### Model
 
-                The available models are finetuned on data from [Xue et al.](https://pubs.rsna.org/doi/10.1148/ryai.2021200197).
+                The available models are finetuned on data from [Xue et al.](https://pubs.rsna.org/doi/10.1148/ryai.2021200197)
                 There are two types of landmark localisation models:
 
                 - **Heatmap**: predicts dense probability maps of landmarks
@@ -726,16 +763,22 @@ def landmark_tab():
 
         with gr.Row():
             with gr.Column():
-                gr.Markdown("## Landmark Localisation")
-                landmark_plot = gr.Plot(show_label=False)
+                landmark_gif = gr.Image(
+                    show_label=True,
+                    type="filepath",
+                    label="Landmark Localisation",
+                )
             with gr.Column():
-                gr.Markdown("## Left Ventricle Length Estimation")
-                lv_plot = gr.Plot(show_label=False)
+                lv_plot = gr.Image(
+                    show_label=True,
+                    type="filepath",
+                    label="Left Ventricle Length Estimation",
+                )
 
         run_button.click(
             fn=landmark,
             inputs=[image_id, view, method, seed],
-            outputs=[landmark_plot, lv_plot],
+            outputs=[landmark_gif, lv_plot],
         )
     return landmark_interface
 
@@ -747,8 +790,9 @@ with gr.Blocks(
         """
         # CineMA: A Foundation Model for Cine Cardiac MRI 🎥🫀
 
-        The following demos showcase the capabilities of CineMA in multiple tasks.
-        For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
+        🚀 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>
+        🔗 For more details, check out our [GitHub](https://github.com/mathpluscode/CineMA).
         """
     )
 
@@ -763,4 +807,4 @@ with gr.Blocks(
             segmentation_lax_tab()
         with gr.TabItem("📍 Landmark Localisation in LAX 2C/4C View", id="landmark"):
             landmark_tab()
-demo.launch()
+demo.launch(allowed_paths=[cache_dir])

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@af1958f51e475d3d6658132c6a680d4fc4a10cac#egg=cinema
+git+https://github.com/mathpluscode/CineMA@7e86ffc7ddf06ad7283915ee143ed808c0f59576#egg=cinema
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.5.1