init

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import io
 from PIL import Image
 import numpy as np
 import spaces  # Import spaces for ZeroGPU compatibility
@@ -47,6 +48,14 @@ checkpoint_path = "checkpoints/Puffin-Base.pth"
 checkpoint = torch.load(checkpoint_path)
 info = model.load_state_dict(checkpoint, strict=False)
 
+def fig_to_image(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    buf.seek(0)
+    img = Image.open(buf).convert('RGB')
+    buf.close()
+    return img
+
 def extract_up_lat_figs(fig_dict):
     fig_up, fig_lat = None, None
     others = {}
@@ -100,11 +109,15 @@ def camera_understanding(image_src, question, seed, progress=gr.Progress(track_t
     single_batch["latitude_field"] = cam[2:].unsqueeze(0)
 
     figs = make_perspective_figures(single_batch, single_batch, n_pairs=1)
-    
-
-    fig_up, fig_lat, _ = extract_up_lat_figs(figs)
-
-    return text, fig_up, fig_lat
+    up_img = lat_img = None
+    for k, fig in figs.items():
+        if "up_field" in k:
+            up_img = fig_to_image(fig)
+        elif "latitude_field" in k:
+            lat_img = fig_to_image(fig)
+        plt.close(fig)
+
+    return text, up_img, lat_img
 
 
 @torch.inference_mode()