app.py

from __future__ import absolute_import, division, print_function

import os
import sys
import cv2
import yaml
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download

# ================== 必须最早 import spaces ==================
try:
    import spaces  # type: ignore
    gpu_decorator = spaces.GPU
except Exception:
    gpu_decorator = lambda f: f

# ================== 工程路径：确保能 import networks ==================
# 适配：无论你从哪里启动 python app.py，都能找到项目根目录
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 = os.path.join(PROJECT_ROOT, "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]
    colored = (colored * 255).astype(np.uint8)
    return np.ascontiguousarray(colored)

# ================== 模型加载 ==================
def load_model(config_path: str):
    import torch
    import torch.nn as nn

    global device
    device = "cuda" if torch.cuda.is_available() else "cpu"

    with open(config_path, "r") as f:
        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
    )
    print(f"✅ Weights downloaded to: {model_path}")

    state = torch.load(model_path, map_location=device)

    m = make(config["model"])
    if any(k.startswith("module") for k in state.keys()):
        m = nn.DataParallel(m)

    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.eval()
    print("✅ Model loaded.")
    return m

# ================== 启动时加载一次模型 ==================
model = load_model(CONFIG_PATH)

# ================== 推理函数 ==================
@gpu_decorator
def infer_raw(img_rgb: np.ndarray):
    if img_rgb is None:
        return None

    import torch

    # 保持你原逻辑：不 resize，直接喂入
    img = img_rgb.astype(np.float32) / 255.0
    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:
                mask = 1 - outputs["pred_mask"]
                mask = mask > 0.5
                outputs["pred_depth"][~mask] = 1
            pred = outputs["pred_depth"][0].cpu().squeeze().numpy()
        else:
            # 保持你原逻辑的 fallback
            pred = outputs[0].cpu().squeeze().numpy()

    return pred.astype(np.float32)

def visualize_100m(pred: np.ndarray):
    if pred is None:
        return None, None, None

    pred_clip = np.clip(pred, 0.0, 1.0)
    depth_gray = (pred_clip * 255).astype(np.uint8)
    depth_color = colorize_depth_fixed(depth_gray, cmap="Spectral")

    npy_path = "/tmp/depth_100m.npy"
    np.save(npy_path, pred)

    return depth_color, depth_gray, npy_path

def visualize_10m(pred: np.ndarray):
    if pred is None:
        return None, None, None

    pred_clip = np.clip(pred, 0.0, 0.1)
    depth_gray = (pred_clip * 10 * 255).astype(np.uint8)
    depth_color = colorize_depth_fixed(depth_gray, cmap="Spectral")

    npy_path = "/tmp/depth_10m.npy"
    np.save(npy_path, pred)

    return depth_color, depth_gray, npy_path

@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

# ================== Gradio UI ==================
example_paths = [
    "hfdemo/01.jpg",
    "hfdemo/02.jpg",
    "hfdemo/03.jpg",
    "hfdemo/04.jpg",
    "hfdemo/05.jpg",
    "hfdemo/06.jpg",
    "hfdemo/07.jpg",
]

example_gen_paths = [
    "hfdemo/generated_00.jpg",
    "hfdemo/generated_01.jpg",
    "hfdemo/generated_02.jpg",
    "hfdemo/generated_03.jpg",
    "hfdemo/generated_04.jpg",
    "hfdemo/generated_05.jpg",
    "hfdemo/generated_06.jpg",
]

with gr.Blocks() as demo:
    gr.Markdown("# DAP Depth Prediction Demo")

    raw_depth = gr.State()  # 🔑 保存模型输出

    with gr.Row():
        # ========== Left ==========
        with gr.Column(scale=1):
            inp = gr.Image(type="numpy", label="Input Image", height=360)

            gr.Markdown("### Examples (click to load)")
            gr.Examples(examples=example_paths, inputs=inp)

            gr.Markdown("### Examples from Gemini (click to load)")
            gr.Examples(examples=example_gen_paths, inputs=inp)

            btn_infer = gr.Button("Run Inference", variant="primary")
            btn_100m = gr.Button("Visualize (100m)")
            btn_10m = gr.Button("Visualize (10m)")

            gr.Markdown(
                """
                <small>
                <b>Visualization range:</b><br>
                • <b>100m</b>: recommended for <b>outdoor</b> scenes<br>
                • <b>10m</b>: recommended for <b>indoor</b> scenes<br>
                (Only affects visualization, not the raw depth output)
                </small>
                """,
                elem_id="vis_hint",
            )

        # ========== Right ==========
        with gr.Column(scale=2):
            out_color = gr.Image(label="Depth (Color)", height=260)
            out_gray = gr.Image(label="Depth (Gray)", height=260)
            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],
    )

    # 2️⃣ 100m
    btn_100m.click(
        fn=visualize_100m,
        inputs=raw_depth,
        outputs=[out_color, out_gray, out_npy],
    )

    # 3️⃣ 10m
    btn_10m.click(
        fn=visualize_10m,
        inputs=raw_depth,
        outputs=[out_color, out_gray, out_npy],
    )

if __name__ == "__main__":
    # 适配“放到网页里”：建议用环境变量控制 host/port
    host = os.environ.get("HOST", "0.0.0.0")
    port = int(os.environ.get("PORT", "7860"))

    demo.launch(
        server_name=host,
        server_port=port,
        ssr_mode=False,
        show_error=True,
    )