Close matplotlib figures

app.py

@@ -1,11 +1,16 @@
 from pathlib import Path
 
 import gradio as gr
+import matplotlib.pyplot as plt
 import numpy as np
 import requests
 import SimpleITK as sitk  # noqa: N813
 import spaces
 import torch
+from huggingface_hub import hf_hub_download
+from monai.transforms import Compose, ScaleIntensityd, SpatialPadd
+from tqdm import tqdm
+
 from cinema import CineMA, ConvUNetR, ConvViT, heatmap_soft_argmax
 from cinema.examples.cine_cmr import plot_cmr_views
 from cinema.examples.inference.landmark_heatmap import (
@@ -29,9 +34,6 @@ from cinema.examples.inference.segmentation_sax import (
 from cinema.examples.inference.segmentation_sax import (
     plot_volume_changes as plot_volume_changes_sax,
 )
-from huggingface_hub import hf_hub_download
-from monai.transforms import Compose, ScaleIntensityd, SpatialPadd
-from tqdm import tqdm
 
 # cache directories
 cache_dir = Path("/tmp/.cinema")
@@ -220,6 +222,7 @@ def mae(image_id, mask_ratio, progress=gr.Progress()):
         reconstructed_dict,
         masks_dict,
     )
+    plt.close(fig)
     return fig
 
 
@@ -331,8 +334,10 @@ def segmentation_sax(trained_dataset, seed, image_id, t_step, progress=gr.Progre
     progress(1, desc="Plotting results...")
     fig1 = plot_segmentations_sax(images, labels, t_step)
     fig2 = plot_volume_changes_sax(labels, t_step)
-
-    return fig1, fig2
+    result = (fig1, fig2)
+    plt.close(fig1)
+    plt.close(fig2)
+    return result
 
 
 def segmentation_sax_tab():
@@ -471,8 +476,10 @@ def segmentation_lax(seed, image_id, progress=gr.Progress()):
     progress(1, desc="Plotting results...")
     fig1 = plot_segmentations_lax(images, labels)
     fig2 = plot_volume_changes_lax(labels)
-
-    return fig1, fig2
+    result = (fig1, fig2)
+    plt.close(fig1)
+    plt.close(fig2)
+    return result
 
 
 def segmentation_lax_tab():
@@ -642,7 +649,10 @@ def landmark(image_id, view, method, seed, progress=gr.Progress()):
 
     landmark_fig = plot_landmarks(images, coords)
     lv_fig = plot_lv(coords)
-    return landmark_fig, lv_fig
+    result = (landmark_fig, lv_fig)
+    plt.close(landmark_fig)
+    plt.close(lv_fig)
+    return result
 
 
 def landmark_tab():

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