init

app.py

@@ -114,23 +114,36 @@ def camera_understanding(image_src, question, seed, progress=gr.Progress(track_t
     cam = gen.get_cam(text)
     
     bgr = np.array(image)[:, :, ::-1].astype(np.float32) / 255.0
-    rgb = bgr[:, :, ::-1].copy()
+    rgb = bgr[:, :, ::-1].copy()        
     image_tensor = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
-    single_batch = {}
-    single_batch["image"] = image_tensor
-    single_batch["up_field"] = cam[:2].unsqueeze(0)
-    single_batch["latitude_field"] = cam[2:].unsqueeze(0)
+    single_batch = {
+        "image": image_tensor,
+        "up_field": cam[:2].unsqueeze(0),
+        "latitude_field": cam[2:].unsqueeze(0),
+    }
 
     figs = make_perspective_figures(single_batch, single_batch, n_pairs=1)
-    up_img = lat_img = None
+    saved_paths = []
+    save_dir = "temp/"
+    os.makedirs(save_dir, exist_ok=True)
+    
     for k, fig in figs.items():
         if "up_field" in k:
-            up_img = fig_to_image(fig)
+            suffix = "_up"
         elif "latitude_field" in k:
-            lat_img = fig_to_image(fig)
+            suffix = "_lat"
+        else:
+            suffix = f"_{k}"
+        out_path = os.path.join(save_dir, f"camera_map_vis{suffix}.png")
+        plt.tight_layout()
+        fig.savefig(out_path, dpi=200, bbox_inches='tight', pad_inches=0)
         plt.close(fig)
+        saved_paths.append(out_path)
+
+    img_up = Image.open(saved_paths[0]).convert("RGB")
+    img_lat = Image.open(saved_paths[1]).convert("RGB")
 
-    return text#, up_img, lat_img
+    return text, img_up, img_lat
 
 
 @torch.inference_mode()
@@ -227,8 +240,8 @@ with gr.Blocks(css=css) as demo:
         understanding_button = gr.Button("Chat")
         understanding_output = gr.Textbox(label="Response")
         
-        #camera1 = gr.Gallery(label="Camera Maps", columns=1, rows=1)
-        #camera2 = gr.Gallery(label="Camera Maps", columns=1, rows=1)
+        camera1 = gr.Gallery(label="Camera Maps", columns=1, rows=1)
+        camera2 = gr.Gallery(label="Camera Maps", columns=1, rows=1)
 
         with gr.Accordion("Advanced options", open=False):
             und_seed_input = gr.Number(label="Seed", precision=0, value=42)
@@ -255,7 +268,7 @@ with gr.Blocks(css=css) as demo:
     understanding_button.click(
         camera_understanding,
         inputs=[image_input, und_seed_input],
-        outputs=[understanding_output]#, camera1, camera2]
+        outputs=[understanding_output, camera1, camera2]
     )
 
 demo.launch(share=True)