Update app.py

app.py

@@ -12,7 +12,6 @@ try:
     import spaces
     gpu_decorator = spaces.GPU
 except Exception:
-    # 本地环境没有 spaces 时兜底
     gpu_decorator = lambda f: f
 
 # ================== 工程路径 ==================
@@ -35,10 +34,10 @@ import matplotlib
 def colorize_depth_fixed(depth_u8: np.ndarray, cmap: str = "Spectral") -> np.ndarray:
     """
     depth_u8: uint8, 0~255
-    return: RGB uint8, 颜色全局一致
+    return: RGB uint8
     """
-    disp = depth_u8.astype(np.float32) / 255.0  # 固定映射
-    colored = matplotlib.colormaps[cmap](disp)[..., :3]  # RGB float
+    disp = depth_u8.astype(np.float32) / 255.0
+    colored = matplotlib.colormaps[cmap](disp)[..., :3]
     colored = (colored * 255).astype(np.uint8)
     return np.ascontiguousarray(colored)
 
@@ -79,15 +78,15 @@ def load_model(config_path: str):
 # ================== 启动时加载一次模型 ==================
 model = load_model(CONFIG_PATH)
 
-# ================== 推理函数（ZeroGPU 必须） ==================
+# ================== 推理函数 ==================
 @gpu_decorator
 def predict(img_rgb: np.ndarray):
     if img_rgb is None:
-        return None, None
+        return None, None, None
 
     import torch
 
-    # 输入处理
+    # -------- 输入处理 --------
     img = img_rgb.astype(np.float32) / 255.0
     tensor = torch.from_numpy(
         img.transpose(2, 0, 1)
@@ -105,14 +104,18 @@ def predict(img_rgb: np.ndarray):
         else:
             pred = outputs[0].cpu().squeeze().numpy()
 
-    # ================== 固定尺度（假设模型输出 0~1） ==================
+    # -------- 深度后处理 --------
     pred = pred.astype(np.float32)
     pred_clip = np.clip(pred, 0.0, 1.0)
 
     depth_gray = (pred_clip * 255).astype(np.uint8)
     depth_color_rgb = colorize_depth_fixed(depth_gray, cmap="Spectral")
 
-    return depth_color_rgb, depth_gray
+    # -------- 保存 npy --------
+    npy_path = "/tmp/depth.npy"
+    np.save(npy_path, pred)  # 保存原始 float32 深度
+
+    return depth_color_rgb, depth_gray, npy_path
 
 # ================== Gradio UI ==================
 demo = gr.Interface(
@@ -121,9 +124,10 @@ demo = gr.Interface(
     outputs=[
         gr.Image(type="numpy", label="Depth (Color)"),
         gr.Image(type="numpy", label="Depth (Gray)"),
+        gr.File(label="Depth (.npy)"),
     ],
     title="DAP Depth Prediction Demo",
-    description="Upload an image and get depth prediction.",
+    description="Upload an image and get depth prediction (color / gray / npy).",
 )
 
 demo.launch(