148 lines
4.7 KiB
Python
148 lines
4.7 KiB
Python
from dataclasses import dataclass
|
||
from typing import Literal, Tuple
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
import requests
|
||
|
||
# 确保本地克隆的 Depth Anything V2 仓库在 sys.path 中，
|
||
# 这样其内部的 `from depth_anything_v2...` 导入才能正常工作。
|
||
_THIS_DIR = Path(__file__).resolve().parent
|
||
_DA_REPO_ROOT = _THIS_DIR / "Depth-Anything-V2"
|
||
if _DA_REPO_ROOT.is_dir():
|
||
da_path = str(_DA_REPO_ROOT)
|
||
if da_path not in sys.path:
|
||
sys.path.insert(0, da_path)
|
||
|
||
from depth_anything_v2.dpt import DepthAnythingV2 # type: ignore[import]
|
||
|
||
|
||
EncoderName = Literal["vits", "vitb", "vitl", "vitg"]
|
||
|
||
|
||
@dataclass
|
||
class DepthAnythingV2Config:
|
||
"""
|
||
Depth Anything V2 模型选择配置。
|
||
|
||
encoder: "vits" | "vitb" | "vitl" | "vitg"
|
||
device: "cuda" | "cpu"
|
||
input_size: 推理时的输入分辨率（短边），参考官方 demo，默认 518。
|
||
"""
|
||
|
||
encoder: EncoderName = "vitl"
|
||
device: str = "cuda"
|
||
input_size: int = 518
|
||
|
||
|
||
_MODEL_CONFIGS = {
|
||
"vits": {"encoder": "vits", "features": 64, "out_channels": [48, 96, 192, 384]},
|
||
"vitb": {"encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768]},
|
||
"vitl": {"encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024]},
|
||
"vitg": {"encoder": "vitg", "features": 384, "out_channels": [1536, 1536, 1536, 1536]},
|
||
}
|
||
|
||
|
||
_DA_V2_WEIGHTS_URLS = {
|
||
# 官方权重托管在 HuggingFace:
|
||
# - Small -> vits
|
||
# - Base -> vitb
|
||
# - Large -> vitl
|
||
# - Giant -> vitg
|
||
# 如需替换为国内镜像，可直接修改这些 URL。
|
||
"vits": "https://huggingface.co/depth-anything/Depth-Anything-V2-Small/resolve/main/depth_anything_v2_vits.pth",
|
||
"vitb": "https://huggingface.co/depth-anything/Depth-Anything-V2-Base/resolve/main/depth_anything_v2_vitb.pth",
|
||
"vitl": "https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth",
|
||
"vitg": "https://huggingface.co/depth-anything/Depth-Anything-V2-Giant/resolve/main/depth_anything_v2_vitg.pth",
|
||
}
|
||
|
||
|
||
def _download_if_missing(encoder: str, ckpt_path: Path) -> None:
|
||
if ckpt_path.is_file():
|
||
return
|
||
|
||
url = _DA_V2_WEIGHTS_URLS.get(encoder)
|
||
if not url:
|
||
raise FileNotFoundError(
|
||
f"找不到权重文件: {ckpt_path}\n"
|
||
f"且当前未为 encoder='{encoder}' 配置自动下载 URL，请手动下载到该路径。"
|
||
)
|
||
|
||
ckpt_path.parent.mkdir(parents=True, exist_ok=True)
|
||
print(f"自动下载 Depth Anything V2 权重 ({encoder}):\n {url}\n -> {ckpt_path}")
|
||
|
||
resp = requests.get(url, stream=True)
|
||
resp.raise_for_status()
|
||
|
||
total = int(resp.headers.get("content-length", "0") or "0")
|
||
downloaded = 0
|
||
chunk_size = 1024 * 1024
|
||
|
||
with ckpt_path.open("wb") as f:
|
||
for chunk in resp.iter_content(chunk_size=chunk_size):
|
||
if not chunk:
|
||
continue
|
||
f.write(chunk)
|
||
downloaded += len(chunk)
|
||
if total > 0:
|
||
done = int(50 * downloaded / total)
|
||
print("\r[{}{}] {:.1f}%".format("#" * done, "." * (50 - done), downloaded * 100 / total), end="")
|
||
|
||
print("\n权重下载完成。")
|
||
|
||
|
||
def load_depth_anything_v2_from_config(
|
||
cfg: DepthAnythingV2Config,
|
||
) -> Tuple[DepthAnythingV2, DepthAnythingV2Config]:
|
||
"""
|
||
根据配置加载 Depth Anything V2 模型与对应配置。
|
||
|
||
说明：
|
||
- 权重文件路径遵循官方命名约定：
|
||
checkpoints/depth_anything_v2_{encoder}.pth
|
||
例如：depth_anything_v2_vitl.pth
|
||
- 请确保上述权重文件已下载到
|
||
python_server/model/Depth/Depth-Anything-V2/checkpoints 下。
|
||
"""
|
||
if cfg.encoder not in _MODEL_CONFIGS:
|
||
raise ValueError(f"不支持的 encoder: {cfg.encoder}")
|
||
|
||
ckpt_path = _DA_REPO_ROOT / "checkpoints" / f"depth_anything_v2_{cfg.encoder}.pth"
|
||
_download_if_missing(cfg.encoder, ckpt_path)
|
||
|
||
model = DepthAnythingV2(**_MODEL_CONFIGS[cfg.encoder])
|
||
|
||
state_dict = torch.load(str(ckpt_path), map_location="cpu")
|
||
model.load_state_dict(state_dict)
|
||
|
||
if cfg.device.startswith("cuda") and torch.cuda.is_available():
|
||
device = "cuda"
|
||
else:
|
||
device = "cpu"
|
||
|
||
model = model.to(device).eval()
|
||
cfg = DepthAnythingV2Config(
|
||
encoder=cfg.encoder,
|
||
device=device,
|
||
input_size=cfg.input_size,
|
||
)
|
||
return model, cfg
|
||
|
||
|
||
def infer_depth_anything_v2(
|
||
model: DepthAnythingV2,
|
||
image_bgr: np.ndarray,
|
||
input_size: int,
|
||
) -> np.ndarray:
|
||
"""
|
||
对单张 BGR 图像做深度推理，返回 float32 深度图（未归一化）。
|
||
image_bgr: OpenCV 读取的 BGR 图像 (H, W, 3), uint8
|
||
"""
|
||
depth = model.infer_image(image_bgr, input_size)
|
||
depth = np.asarray(depth, dtype="float32")
|
||
return depth
|
||
|