depth estimatro

app.py

@@ -1,28 +1,31 @@
+# app.py
 import os
 import pathlib
 import subprocess
 import gradio as gr
 import spaces
 import torch
+from PIL import Image
 
-# ---------- 权重下载：强制在 code_depth 下执行你的脚本 ----------
 BASE_DIR = pathlib.Path(__file__).resolve().parent
 SCRIPT_DIR = BASE_DIR / "code_depth"
 GET_WEIGHTS_SH = SCRIPT_DIR / "get_weights.sh"
 
-def ensure_executable(path: pathlib.Path):
-    if not path.exists():
-        raise FileNotFoundError(f"Download script not found: {path}")
-    os.chmod(path, os.stat(path).st_mode | 0o111)
-
-def ensure_weights() -> str:
-    """
-    在 code_depth 目录下运行 get_weights.sh。
-    该脚本会在 code_depth/ 下创建 checkpoints/ 并下载权重。
-    返回绝对路径：<repo_root>/code_depth/checkpoints
-    """
-    ensure_executable(GET_WEIGHTS_SH)
-    # 你脚本的工作目录需要是 code_depth
+# 让我们能 import 到 code_depth/depth_infer.py
+import sys
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.append(str(SCRIPT_DIR))
+
+from depth_infer import DepthModel  # noqa
+
+def _ensure_executable(p: pathlib.Path):
+    if not p.exists():
+        raise FileNotFoundError(f"Not found: {p}")
+    os.chmod(p, os.stat(p).st_mode | 0o111)
+
+def ensure_weights():
+    """在 code_depth 目录下运行你的 get_weights.sh。"""
+    _ensure_executable(GET_WEIGHTS_SH)
     subprocess.run(
         ["bash", str(GET_WEIGHTS_SH)],
         check=True,
@@ -30,27 +33,55 @@ def ensure_weights() -> str:
         env={**os.environ, "HF_HUB_DISABLE_TELEMETRY": "1"},
     )
     ckpt_dir = SCRIPT_DIR / "checkpoints"
+    if not ckpt_dir.exists():
+        raise RuntimeError("weights download script ran but checkpoints/ not found")
     return str(ckpt_dir)
 
-# 启动时先拉权重（不开 Persistent Storage 时，重建环境会清空；重启后会自动再拉一次）
+# 启动时下载权重（不开持久化时，若环境重建会再次下载）
 try:
     CKPT_DIR = ensure_weights()
     print(f"✅ Weights ready in: {CKPT_DIR}")
 except Exception as e:
     print(f"⚠️ Failed to prepare weights: {e}")
-    CKPT_DIR = str(SCRIPT_DIR / "checkpoints")  # 仍然给个路径，后续可检查是否存在
 
-# ---------- Gradio 推理函数 ----------
+# 模型缓存（按 encoder 复用）
+_MODELS: dict[str, DepthModel] = {}
+
+def get_model(encoder: str) -> DepthModel:
+    if encoder not in _MODELS:
+        _MODELS[encoder] = DepthModel(BASE_DIR, encoder=encoder)
+    return _MODELS[encoder]
+
 @spaces.GPU
-def greet(n: float):
-    # 在 GPU worker 里拿 device
+def infer_depth(
+    image: Image.Image,
+    encoder: str = "vitl",
+    max_res: int = 1280,
+    input_size: int = 518,
+    fp32: bool = False,
+    grayscale: bool = False,
+) -> Image.Image:
+    # 这里才真正触发 CUDA 设备占用
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    zero = torch.tensor([0.0], device=device)
-    # 仅示例输出，你可以在这里用 CKPT_DIR 加载你的模型
-    print(f"Device in greet(): {device}")
-    print(f"Using checkpoints from: {CKPT_DIR}")
-    return f"Hello {(zero + n).item()} Tensor (device={device})"
+    print(f"[infer] device={device}, encoder={encoder}, max_res={max_res}, input_size={input_size}, fp32={fp32}, gray={grayscale}")
+    model = get_model(encoder)
+    return model.infer(image, max_res=max_res, input_size=input_size, fp32=fp32, grayscale=grayscale)
+
+with gr.Blocks() as demo:
+    gr.Markdown("## GeoRemover · Depth Preview (Video-Depth-Anything)")
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(label="Upload image", type="pil")
+            encoder = gr.Dropdown(["vits", "vitl"], value="vitl", label="Encoder")
+            max_res = gr.Slider(512, 2048, value=1280, step=64, label="Max resolution")
+            input_size = gr.Slider(256, 1024, value=518, step=2, label="Model input_size")
+            fp32 = gr.Checkbox(False, label="Use FP32 (default FP16)")
+            gray = gr.Checkbox(False, label="Grayscale depth")
+            btn = gr.Button("Run")
+        with gr.Column():
+            out = gr.Image(label="Depth visualization")
+
+    btn.click(fn=infer_depth, inputs=[inp, encoder, max_res, input_size, fp32, gray], outputs=[out])
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(label="n"), outputs=gr.Text())
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)

code_depth/depth_infer.py

@@ -0,0 +1,87 @@
+# app.py
+import os
+import pathlib
+import subprocess
+import gradio as gr
+import spaces
+import torch
+from PIL import Image
+
+BASE_DIR = pathlib.Path(__file__).resolve().parent
+SCRIPT_DIR = BASE_DIR / "code_depth"
+GET_WEIGHTS_SH = SCRIPT_DIR / "get_weights.sh"
+
+# 让我们能 import 到 code_depth/depth_infer.py
+import sys
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.append(str(SCRIPT_DIR))
+
+from depth_infer import DepthModel  # noqa
+
+def _ensure_executable(p: pathlib.Path):
+    if not p.exists():
+        raise FileNotFoundError(f"Not found: {p}")
+    os.chmod(p, os.stat(p).st_mode | 0o111)
+
+def ensure_weights():
+    """在 code_depth 目录下运行你的 get_weights.sh。"""
+    _ensure_executable(GET_WEIGHTS_SH)
+    subprocess.run(
+        ["bash", str(GET_WEIGHTS_SH)],
+        check=True,
+        cwd=str(SCRIPT_DIR),
+        env={**os.environ, "HF_HUB_DISABLE_TELEMETRY": "1"},
+    )
+    ckpt_dir = SCRIPT_DIR / "checkpoints"
+    if not ckpt_dir.exists():
+        raise RuntimeError("weights download script ran but checkpoints/ not found")
+    return str(ckpt_dir)
+
+# 启动时下载权重（不开持久化时，若环境重建会再次下载）
+try:
+    CKPT_DIR = ensure_weights()
+    print(f"✅ Weights ready in: {CKPT_DIR}")
+except Exception as e:
+    print(f"⚠️ Failed to prepare weights: {e}")
+
+# 模型缓存（按 encoder 复用）
+_MODELS: dict[str, DepthModel] = {}
+
+def get_model(encoder: str) -> DepthModel:
+    if encoder not in _MODELS:
+        _MODELS[encoder] = DepthModel(BASE_DIR, encoder=encoder)
+    return _MODELS[encoder]
+
+@spaces.GPU
+def infer_depth(
+    image: Image.Image,
+    encoder: str = "vitl",
+    max_res: int = 1280,
+    input_size: int = 518,
+    fp32: bool = False,
+    grayscale: bool = False,
+) -> Image.Image:
+    # 这里才真正触发 CUDA 设备占用
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"[infer] device={device}, encoder={encoder}, max_res={max_res}, input_size={input_size}, fp32={fp32}, gray={grayscale}")
+    model = get_model(encoder)
+    return model.infer(image, max_res=max_res, input_size=input_size, fp32=fp32, grayscale=grayscale)
+
+with gr.Blocks() as demo:
+    gr.Markdown("## GeoRemover · Depth Preview (Video-Depth-Anything)")
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(label="Upload image", type="pil")
+            encoder = gr.Dropdown(["vits", "vitl"], value="vitl", label="Encoder")
+            max_res = gr.Slider(512, 2048, value=1280, step=64, label="Max resolution")
+            input_size = gr.Slider(256, 1024, value=518, step=2, label="Model input_size")
+            fp32 = gr.Checkbox(False, label="Use FP32 (default FP16)")
+            gray = gr.Checkbox(False, label="Grayscale depth")
+            btn = gr.Button("Run")
+        with gr.Column():
+            out = gr.Image(label="Depth visualization")
+
+    btn.click(fn=infer_depth, inputs=[inp, encoder, max_res, input_size, fp32, gray], outputs=[out])
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)