Update app.py

app.py

@@ -7,16 +7,21 @@ import numpy as np
 import gradio as gr
 from huggingface_hub import hf_hub_download
 
-# ✅ 必须最早 import spaces（在 torch / 任何 CUDA 初始化之前）
-import spaces  # 在 HF Spaces 一定存在
-
-import matplotlib  # 用你的 colormap
-
+# ================== 必须最早 import spaces ==================
+try:
+    import spaces
+    gpu_decorator = spaces.GPU
+except Exception:
+    # 本地环境没有 spaces 时兜底
+    gpu_decorator = lambda f: f
+
+# ================== 工程路径 ==================
 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"
@@ -24,21 +29,20 @@ CONFIG_PATH = "config/infer.yaml"
 model = None
 device = "cpu"
 
+# ================== 固定颜色映射（颜色一致） ==================
+import matplotlib
 
+def colorize_depth_fixed(depth_u8: np.ndarray, cmap: str = "Spectral") -> np.ndarray:
+    """
+    depth_u8: uint8, 0~255
+    return: RGB uint8, 颜色全局一致
+    """
+    disp = depth_u8.astype(np.float32) / 255.0  # 固定映射
+    colored = matplotlib.colormaps[cmap](disp)[..., :3]  # RGB float
+    colored = (colored * 255).astype(np.uint8)
+    return np.ascontiguousarray(colored)
 
-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):
     import torch
     import torch.nn as nn
@@ -50,7 +54,10 @@ def load_model(config_path: str):
         config = yaml.load(f, Loader=yaml.FullLoader)
 
     print(f"Downloading weights from HF: {WEIGHTS_REPO}/{WEIGHTS_FILE}")
-    model_path = hf_hub_download(repo_id=WEIGHTS_REPO, filename=WEIGHTS_FILE)
+    model_path = hf_hub_download(
+        repo_id=WEIGHTS_REPO,
+        filename=WEIGHTS_FILE
+    )
     print(f"✅ Weights downloaded to: {model_path}")
 
     state = torch.load(model_path, map_location=device)
@@ -61,55 +68,53 @@ def load_model(config_path: str):
 
     m = m.to(device)
     m_state = m.state_dict()
-    m.load_state_dict({k: v for k, v in state.items() if k in m_state}, strict=False)
+    m.load_state_dict(
+        {k: v for k, v in state.items() if k in m_state},
+        strict=False
+    )
     m.eval()
     print("✅ Model loaded.")
     return m
 
-
+# ================== 启动时加载一次模型 ==================
 model = load_model(CONFIG_PATH)
 
-
-@spaces.GPU
+# ================== 推理函数（ZeroGPU 必须） ==================
+@gpu_decorator
 def predict(img_rgb: np.ndarray):
     if img_rgb is None:
         return None, None
 
     import torch
 
+    # 输入处理
     img = img_rgb.astype(np.float32) / 255.0
-    tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(device)
+    tensor = torch.from_numpy(
+        img.transpose(2, 0, 1)
+    ).unsqueeze(0).to(device)
 
     with torch.inference_mode():
         outputs = model(tensor)
 
         if isinstance(outputs, dict) and "pred_depth" in outputs:
             if "pred_mask" in outputs:
-                pm = 1 - outputs["pred_mask"]
-                pm = (pm > 0.5)
-                outputs["pred_depth"][~pm] = 1
-            pred = outputs["pred_depth"][0].detach().cpu().squeeze().numpy()
+                mask = 1 - outputs["pred_mask"]
+                mask = mask > 0.5
+                outputs["pred_depth"][~mask] = 1
+            pred = outputs["pred_depth"][0].cpu().squeeze().numpy()
         else:
-            pred = outputs[0].detach().cpu().squeeze().numpy()
+            pred = outputs[0].cpu().squeeze().numpy()
 
-    # 灰度图：如果你 pred 本来就在 0~1，就直接 *255；否则先归一化
+    # ================== 固定尺度（假设模型输出 0~1） ==================
     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)
+    pred_clip = np.clip(pred, 0.0, 1.0)
 
-    depth_gray = (pred_norm * 255).astype(np.uint8)
-
-    # 彩色图：用你改进的可视化
-    depth_color_rgb = colorize_depth_matplotlib(pred_norm, normalize=False, cmap="Spectral")
+    depth_gray = (pred_clip * 255).astype(np.uint8)
+    depth_color_rgb = colorize_depth_fixed(depth_gray, cmap="Spectral")
 
     return depth_color_rgb, depth_gray
 
-
+# ================== Gradio UI ==================
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy", label="Input Image"),