Update app.py

app.py

@@ -8,36 +8,38 @@ import gradio as gr
 from huggingface_hub import hf_hub_download
 
 # ✅ 必须最早 import spaces（在 torch / 任何 CUDA 初始化之前）
-try:
-    import spaces  # noqa: F401
-except Exception:
-    spaces = None  # 不影响本地跑
+import spaces  # 在 HF Spaces 一定存在
+
+import matplotlib  # 用你的 colormap
 
-# ========== 让 Space 能 import 你的工程 ==========
 PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(PROJECT_ROOT)
 
 from networks.models import make  # noqa: E402
 
-# ====== HF 权重仓库配置 ======
 WEIGHTS_REPO = "Insta360-Research/DAP-weights"
 WEIGHTS_FILE = "model.pth"
 CONFIG_PATH = "config/infer.yaml"
 
-# 先定义全局占位
 model = None
 device = "cpu"
 
 
-def colorize_depth(depth, colormap=cv2.COLORMAP_JET):
-    depth = depth.astype(np.float32)
-    depth_norm = (depth - depth.min()) / (depth.max() - depth.min() + 1e-6)
-    depth_u8 = (depth_norm * 255).astype(np.uint8)
-    return cv2.applyColorMap(depth_u8, colormap)  # BGR
+
+def colorize_depth_matplotlib(depth: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray:
+    if mask is None:
+        depth = np.where(depth > 0, depth, np.nan)
+    else:
+        depth = np.where((depth > 0) & mask, depth, np.nan)
+
+    disp = depth / 255.0
+
+    colored = np.nan_to_num(matplotlib.colormaps[cmap](disp)[..., :3], 0)
+    colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8))
+    return colored
 
 
 def load_model(config_path: str):
-    # ✅ torch 放到这里 import，避免在 spaces import 之前触发 CUDA
     import torch
     import torch.nn as nn
 
@@ -65,36 +67,45 @@ def load_model(config_path: str):
     return m
 
 
-# ✅ 启动时加载一次模型
 model = load_model(CONFIG_PATH)
 
+
 @spaces.GPU
 def predict(img_rgb: np.ndarray):
     if img_rgb is None:
         return None, None
 
-    import torch  # 这里用到 torch，再 import 一次没关系
+    import torch
 
     img = img_rgb.astype(np.float32) / 255.0
     tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(device)
 
-    with torch.no_grad():
+    with torch.inference_mode():
         outputs = model(tensor)
 
         if isinstance(outputs, dict) and "pred_depth" in outputs:
             if "pred_mask" in outputs:
-                outputs["pred_mask"] = 1 - outputs["pred_mask"]
-                outputs["pred_mask"] = (outputs["pred_mask"] > 0.5)
-                outputs["pred_depth"][~outputs["pred_mask"]] = 1
+                pm = 1 - outputs["pred_mask"]
+                pm = (pm > 0.5)
+                outputs["pred_depth"][~pm] = 1
             pred = outputs["pred_depth"][0].detach().cpu().squeeze().numpy()
         else:
             pred = outputs[0].detach().cpu().squeeze().numpy()
 
-    pred_clip = np.clip(pred, 0.001, 1.0)
-    depth_gray = (pred_clip * 255).astype(np.uint8)
+    # 灰度图：如果你 pred 本来就在 0~1，就直接 *255；否则先归一化
+    pred = pred.astype(np.float32)
+    pred_clip = np.clip(pred, 1e-6, np.nanmax(pred) if np.isfinite(pred).any() else 1.0)
+
+    # 让灰度输出稳定：用分位数做一次归一化
+    lo = np.nanquantile(pred_clip, 0.001)
+    hi = np.nanquantile(pred_clip, 0.99)
+    pred_norm = (pred_clip - lo) / (hi - lo + 1e-6)
+    pred_norm = np.clip(pred_norm, 0.0, 1.0)
+
+    depth_gray = (pred_norm * 255).astype(np.uint8)
 
-    depth_color_bgr = colorize_depth(pred_clip, cv2.COLORMAP_JET)
-    depth_color_rgb = cv2.cvtColor(depth_color_bgr, cv2.COLOR_BGR2RGB)
+    # 彩色图：用你改进的可视化
+    depth_color_rgb = colorize_depth_matplotlib(pred_norm, normalize=False, cmap="Spectral")
 
     return depth_color_rgb, depth_gray