Update app.py

app.py

@@ -80,6 +80,23 @@ def load_model(config_path: str):
 # ================== 启动时加载一次模型 ==================
 model = load_model(CONFIG_PATH)
 
+# ================== 加载标度尺图片 ==================
+COLORBAR_DIR = os.path.join(PROJECT_ROOT, "colorbars")
+colorbar_100m_color = cv2.imread(os.path.join(COLORBAR_DIR, "colorbar_100m_color.png"))
+colorbar_100m_gray = cv2.imread(os.path.join(COLORBAR_DIR, "colorbar_100m_gray.png"))
+colorbar_10m_color = cv2.imread(os.path.join(COLORBAR_DIR, "colorbar_10m_color.png"))
+colorbar_10m_gray = cv2.imread(os.path.join(COLORBAR_DIR, "colorbar_10m_gray.png"))
+
+# 转换为RGB（Gradio需要RGB格式）
+if colorbar_100m_color is not None:
+    colorbar_100m_color = cv2.cvtColor(colorbar_100m_color, cv2.COLOR_BGR2RGB)
+if colorbar_100m_gray is not None:
+    colorbar_100m_gray = cv2.cvtColor(colorbar_100m_gray, cv2.COLOR_BGR2RGB)
+if colorbar_10m_color is not None:
+    colorbar_10m_color = cv2.cvtColor(colorbar_10m_color, cv2.COLOR_BGR2RGB)
+if colorbar_10m_gray is not None:
+    colorbar_10m_gray = cv2.cvtColor(colorbar_10m_gray, cv2.COLOR_BGR2RGB)
+
 # ================== 推理函数 ==================
 @gpu_decorator
 def infer_raw(img_rgb: np.ndarray):
@@ -109,7 +126,7 @@ def infer_raw(img_rgb: np.ndarray):
 
 def visualize_100m(pred: np.ndarray):
     if pred is None:
-        return None, None, None
+        return None, None, None, None, None
 
     pred_clip = np.clip(pred, 0.0, 1.0)
     depth_gray = (pred_clip * 255).astype(np.uint8)
@@ -118,11 +135,11 @@ def visualize_100m(pred: np.ndarray):
     npy_path = "/tmp/depth_100m.npy"
     np.save(npy_path, pred)
 
-    return depth_color, depth_gray, npy_path
+    return depth_color, depth_gray, npy_path, colorbar_100m_color, colorbar_100m_gray
 
 def visualize_10m(pred: np.ndarray):
     if pred is None:
-        return None, None, None
+        return None, None, None, None, None
 
     pred_clip = np.clip(pred, 0.0, 0.1)
     depth_gray = (pred_clip * 10 * 255).astype(np.uint8)
@@ -131,13 +148,13 @@ def visualize_10m(pred: np.ndarray):
     npy_path = "/tmp/depth_10m.npy"
     np.save(npy_path, pred)
 
-    return depth_color, depth_gray, npy_path
+    return depth_color, depth_gray, npy_path, colorbar_10m_color, colorbar_10m_gray
 
 @gpu_decorator
 def infer_and_vis_100m(img_rgb: np.ndarray):
     pred = infer_raw(img_rgb)                 # 跑模型一次（GPU）
-    color, gray, npy = visualize_100m(pred)   # 默认100m显示（CPU）
-    return pred, color, gray, npy
+    color, gray, npy, cbar_color, cbar_gray = visualize_100m(pred)   # 默认100m显示（CPU）
+    return pred, color, gray, npy, cbar_color, cbar_gray
 
 # ================== Gradio UI ==================
 example_paths = [
@@ -210,29 +227,33 @@ with gr.Blocks() as demo:
 
         # ========== Right ==========
         with gr.Column(scale=2):
-            out_color = gr.Image(label="Depth (Color)", height=260)
-            out_gray = gr.Image(label="Depth (Gray)", height=260)
+            with gr.Row():
+                out_color = gr.Image(label="Depth (Color)", height=260, scale=5)
+                colorbar_color = gr.Image(label="Scale", height=260, scale=1, show_label=False)
+            with gr.Row():
+                out_gray = gr.Image(label="Depth (Gray)", height=260, scale=5)
+                colorbar_gray = gr.Image(label="Scale", height=260, scale=1, show_label=False)
             out_npy = gr.File(label="Depth (.npy)")
 
     # 1️⃣ 跑模型
     btn_infer.click(
         fn=infer_and_vis_100m,
         inputs=inp,
-        outputs=[raw_depth, out_color, out_gray, out_npy],
+        outputs=[raw_depth, out_color, out_gray, out_npy, colorbar_color, colorbar_gray],
     )
 
     # 2️⃣ 100m
     btn_100m.click(
         fn=visualize_100m,
         inputs=raw_depth,
-        outputs=[out_color, out_gray, out_npy],
+        outputs=[out_color, out_gray, out_npy, colorbar_color, colorbar_gray],
     )
 
     # 3️⃣ 10m
     btn_10m.click(
         fn=visualize_10m,
         inputs=raw_depth,
-        outputs=[out_color, out_gray, out_npy],
+        outputs=[out_color, out_gray, out_npy, colorbar_color, colorbar_gray],
     )
 
 
@@ -248,5 +269,4 @@ if __name__ == "__main__":
         server_port=port,
         ssr_mode=False,
         show_error=True,
-    )
-
+    )