Spaces:
Running on Zero
Running on Zero
Initial MetaView novel view synthesis demo
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +4 -0
- README.md +26 -7
- app.py +270 -0
- depth_anything_3/api.py +451 -0
- depth_anything_3/app/css_and_html.py +594 -0
- depth_anything_3/app/gradio_app.py +724 -0
- depth_anything_3/app/modules/__init__.py +43 -0
- depth_anything_3/app/modules/event_handlers.py +622 -0
- depth_anything_3/app/modules/file_handlers.py +355 -0
- depth_anything_3/app/modules/model_inference.py +260 -0
- depth_anything_3/app/modules/ui_components.py +477 -0
- depth_anything_3/app/modules/utils.py +207 -0
- depth_anything_3/app/modules/visualization.py +434 -0
- depth_anything_3/bench/__init__.py +45 -0
- depth_anything_3/bench/configs/eval_bench.yaml +98 -0
- depth_anything_3/bench/dataset.py +136 -0
- depth_anything_3/bench/datasets/__init__.py +21 -0
- depth_anything_3/bench/datasets/dtu.py +681 -0
- depth_anything_3/bench/datasets/dtu64.py +182 -0
- depth_anything_3/bench/datasets/eth3d.py +594 -0
- depth_anything_3/bench/datasets/hiroom.py +440 -0
- depth_anything_3/bench/datasets/scannetpp.py +591 -0
- depth_anything_3/bench/datasets/sevenscenes.py +449 -0
- depth_anything_3/bench/evaluator.py +752 -0
- depth_anything_3/bench/print_metrics.py +618 -0
- depth_anything_3/bench/registries.py +85 -0
- depth_anything_3/bench/utils.py +525 -0
- depth_anything_3/cfg.py +144 -0
- depth_anything_3/cli.py +824 -0
- depth_anything_3/configs/da3-base.yaml +45 -0
- depth_anything_3/configs/da3-giant.yaml +71 -0
- depth_anything_3/configs/da3-large.yaml +45 -0
- depth_anything_3/configs/da3-small.yaml +45 -0
- depth_anything_3/configs/da3metric-large.yaml +28 -0
- depth_anything_3/configs/da3mono-large.yaml +28 -0
- depth_anything_3/configs/da3nested-giant-large.yaml +10 -0
- depth_anything_3/model/__init__.py +20 -0
- depth_anything_3/model/cam_dec.py +45 -0
- depth_anything_3/model/cam_enc.py +80 -0
- depth_anything_3/model/da3.py +442 -0
- depth_anything_3/model/dinov2/dinov2.py +64 -0
- depth_anything_3/model/dinov2/layers/__init__.py +25 -0
- depth_anything_3/model/dinov2/layers/attention.py +100 -0
- depth_anything_3/model/dinov2/layers/block.py +143 -0
- depth_anything_3/model/dinov2/layers/drop_path.py +35 -0
- depth_anything_3/model/dinov2/layers/layer_scale.py +31 -0
- depth_anything_3/model/dinov2/layers/mlp.py +40 -0
- depth_anything_3/model/dinov2/layers/patch_embed.py +94 -0
- depth_anything_3/model/dinov2/layers/rope.py +200 -0
- depth_anything_3/model/dinov2/layers/swiglu_ffn.py +62 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
examples/1.png filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
examples/12.png filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
examples/5.png filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
examples/9.png filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -1,13 +1,32 @@
|
|
| 1 |
---
|
| 2 |
-
title:
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: gradio
|
| 7 |
-
sdk_version:
|
| 8 |
-
python_version: '3.12'
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
|
|
|
|
|
|
|
|
|
|
|
|
| 11 |
---
|
| 12 |
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: MetaView Novel View Synthesis
|
| 3 |
+
emoji: 🎥
|
| 4 |
+
colorFrom: pink
|
| 5 |
+
colorTo: purple
|
| 6 |
sdk: gradio
|
| 7 |
+
sdk_version: 5.49.1
|
|
|
|
| 8 |
app_file: app.py
|
| 9 |
pinned: false
|
| 10 |
+
license: apache-2.0
|
| 11 |
+
short_description: Monocular novel view synthesis from a single image
|
| 12 |
+
python_version: "3.10"
|
| 13 |
+
startup_duration_timeout: 1h
|
| 14 |
---
|
| 15 |
|
| 16 |
+
# MetaView — Monocular Novel View Synthesis
|
| 17 |
+
|
| 18 |
+
Interactive demo for [**MetaView**](https://huggingface.co/Kwai-Kolors/MetaView)
|
| 19 |
+
(ECCV 2026): *Monocular Novel View Synthesis with Scale-Aware Implicit Geometry
|
| 20 |
+
Priors*.
|
| 21 |
+
|
| 22 |
+
Upload a single image, choose a target camera **yaw** (left/right) and **pitch**
|
| 23 |
+
(up/down), and MetaView renders the scene from the new viewpoint.
|
| 24 |
+
|
| 25 |
+
The pipeline combines:
|
| 26 |
+
- **Qwen-Image-Edit** MM-DiT backbone (novel-view generation),
|
| 27 |
+
- **Depth-Anything-3** GIANT (implicit 3D feature priors) + NESTED (dense depth),
|
| 28 |
+
- **Modified RoPE (PRoPE)** for metric scale anchoring of the camera pose.
|
| 29 |
+
|
| 30 |
+
Links: [Model](https://huggingface.co/Kwai-Kolors/MetaView) ·
|
| 31 |
+
[Code](https://github.com/KlingAIResearch/MetaView) ·
|
| 32 |
+
[Paper](https://arxiv.org/abs/2607.12000)
|
app.py
ADDED
|
@@ -0,0 +1,270 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
|
| 3 |
+
# Allocator config to survive transient memory spikes from the large DiT.
|
| 4 |
+
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
| 5 |
+
# DiffSynth defaults to ModelScope; force Hugging Face so all weights come from the Hub.
|
| 6 |
+
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
|
| 7 |
+
|
| 8 |
+
import glob
|
| 9 |
+
import math
|
| 10 |
+
import sys
|
| 11 |
+
import tempfile
|
| 12 |
+
|
| 13 |
+
import spaces # must come before torch / any CUDA-touching import
|
| 14 |
+
import numpy as np
|
| 15 |
+
import torch
|
| 16 |
+
import torch.nn.functional as F
|
| 17 |
+
import gradio as gr
|
| 18 |
+
from PIL import Image
|
| 19 |
+
from huggingface_hub import snapshot_download
|
| 20 |
+
|
| 21 |
+
# Local vendored packages (MetaView `src/` + `diffsynth/` and Depth-Anything-3).
|
| 22 |
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
| 23 |
+
|
| 24 |
+
from depth_anything_3.api import DepthAnything3
|
| 25 |
+
from diffsynth.core import ModelConfig
|
| 26 |
+
from diffsynth import load_state_dict
|
| 27 |
+
from src.MetaView_pipeline import MetaViewPipeline
|
| 28 |
+
|
| 29 |
+
# ---------------------------------------------------------------------------
|
| 30 |
+
# Model ids
|
| 31 |
+
# ---------------------------------------------------------------------------
|
| 32 |
+
METAVIEW_REPO = "Kwai-Kolors/MetaView"
|
| 33 |
+
DA3_GIANT_REPO = "depth-anything/DA3-GIANT-1.1"
|
| 34 |
+
DA3_NESTED_REPO = "depth-anything/DA3NESTED-GIANT-LARGE-1.1"
|
| 35 |
+
QWEN_EDIT_REPO = "Qwen/Qwen-Image-Edit"
|
| 36 |
+
|
| 37 |
+
# MetaView global config (matches the reference inference script).
|
| 38 |
+
EXPORT_3D_FEAT_LAYERS = [19, 27, 33, 39]
|
| 39 |
+
PROPE_DIM_ARRANGE = [64, 20, 20, 24]
|
| 40 |
+
ADD_DEPTH = len(PROPE_DIM_ARRANGE) == 4
|
| 41 |
+
MERGE_3D = True
|
| 42 |
+
PROMPT = ["镜头视角转到指定位置"] # "move the camera view to the target position"
|
| 43 |
+
GEN_W, GEN_H = 960, 528
|
| 44 |
+
NUM_STEPS = 40
|
| 45 |
+
|
| 46 |
+
device = "cuda"
|
| 47 |
+
dtype = torch.bfloat16
|
| 48 |
+
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
# Load everything at module scope. `import spaces` above has hijacked
|
| 51 |
+
# torch.cuda.*, so `.to("cuda")` here packs weights to disk and streams them
|
| 52 |
+
# into VRAM on the first @spaces.GPU call.
|
| 53 |
+
# ---------------------------------------------------------------------------
|
| 54 |
+
print("[*] Downloading MetaView checkpoint...")
|
| 55 |
+
metaview_dir = snapshot_download(METAVIEW_REPO)
|
| 56 |
+
metaview_ckpt = glob.glob(os.path.join(metaview_dir, "*.safetensors"))[0]
|
| 57 |
+
|
| 58 |
+
print("[*] Downloading Depth-Anything-3 models...")
|
| 59 |
+
da3_giant_dir = snapshot_download(DA3_GIANT_REPO)
|
| 60 |
+
da3_nested_dir = snapshot_download(DA3_NESTED_REPO)
|
| 61 |
+
|
| 62 |
+
# Download the Qwen-Image-Edit backbone components from the Hub explicitly so we
|
| 63 |
+
# can hand DiffSynth concrete local file paths (avoids the ModelScope default
|
| 64 |
+
# download path and the processor-dir globbing quirk).
|
| 65 |
+
print("[*] Downloading Qwen-Image-Edit backbone...")
|
| 66 |
+
qwen_edit_dir = snapshot_download(
|
| 67 |
+
QWEN_EDIT_REPO,
|
| 68 |
+
allow_patterns=["transformer/diffusion_pytorch_model*.safetensors",
|
| 69 |
+
"transformer/*.json", "processor/*"],
|
| 70 |
+
)
|
| 71 |
+
qwen_image_dir = snapshot_download(
|
| 72 |
+
"Qwen/Qwen-Image",
|
| 73 |
+
allow_patterns=["text_encoder/model*.safetensors", "text_encoder/*.json",
|
| 74 |
+
"vae/diffusion_pytorch_model.safetensors", "vae/*.json",
|
| 75 |
+
"tokenizer/*"],
|
| 76 |
+
)
|
| 77 |
+
|
| 78 |
+
transformer_files = sorted(glob.glob(
|
| 79 |
+
os.path.join(qwen_edit_dir, "transformer", "diffusion_pytorch_model*.safetensors")))
|
| 80 |
+
text_encoder_files = sorted(glob.glob(
|
| 81 |
+
os.path.join(qwen_image_dir, "text_encoder", "model*.safetensors")))
|
| 82 |
+
vae_file = os.path.join(qwen_image_dir, "vae", "diffusion_pytorch_model.safetensors")
|
| 83 |
+
processor_path = os.path.join(qwen_edit_dir, "processor")
|
| 84 |
+
tokenizer_path = os.path.join(qwen_image_dir, "tokenizer")
|
| 85 |
+
|
| 86 |
+
print("[*] Loading Depth-Anything-3 GIANT (3D feature extractor)...")
|
| 87 |
+
da3_giant = DepthAnything3.from_pretrained(da3_giant_dir).to(device=device).eval()
|
| 88 |
+
|
| 89 |
+
print("[*] Loading Depth-Anything-3 NESTED (dense depth)...")
|
| 90 |
+
da3_nested = DepthAnything3.from_pretrained(da3_nested_dir).to(device=device).eval()
|
| 91 |
+
|
| 92 |
+
print("[*] Loading MetaView / Qwen-Image-Edit pipeline...")
|
| 93 |
+
pipe = MetaViewPipeline.from_pretrained(
|
| 94 |
+
torch_dtype=dtype,
|
| 95 |
+
device=device,
|
| 96 |
+
model_configs=[
|
| 97 |
+
ModelConfig(path=transformer_files),
|
| 98 |
+
ModelConfig(path=text_encoder_files),
|
| 99 |
+
ModelConfig(path=vae_file),
|
| 100 |
+
],
|
| 101 |
+
tokenizer_config=ModelConfig(path=tokenizer_path),
|
| 102 |
+
processor_config=ModelConfig(path=processor_path),
|
| 103 |
+
)
|
| 104 |
+
|
| 105 |
+
print(f"[*] Applying MetaView weights from {metaview_ckpt}...")
|
| 106 |
+
state_dict = load_state_dict(metaview_ckpt)
|
| 107 |
+
pipe.dit.load_state_dict(state_dict, strict=False)
|
| 108 |
+
del state_dict
|
| 109 |
+
print("[*] All models loaded.")
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def compute_target_extrinsic(yaw_deg, pitch_deg, radius):
|
| 113 |
+
"""Camera World-to-Camera extrinsic for a rotation around a sphere center
|
| 114 |
+
in front of the camera (yaw = left/right, pitch = up/down)."""
|
| 115 |
+
yaw = math.radians(yaw_deg)
|
| 116 |
+
pitch = math.radians(pitch_deg)
|
| 117 |
+
R_y = np.array([[np.cos(yaw), 0, np.sin(yaw)],
|
| 118 |
+
[0, 1, 0],
|
| 119 |
+
[-np.sin(yaw), 0, np.cos(yaw)]])
|
| 120 |
+
R_x = np.array([[1, 0, 0],
|
| 121 |
+
[0, np.cos(pitch), -np.sin(pitch)],
|
| 122 |
+
[0, np.sin(pitch), np.cos(pitch)]])
|
| 123 |
+
R = R_y @ R_x
|
| 124 |
+
C = np.array([0.0, 0.0, radius])
|
| 125 |
+
t = C - R @ C
|
| 126 |
+
T = np.eye(4)
|
| 127 |
+
T[:3, :3] = R
|
| 128 |
+
T[:3, 3] = t
|
| 129 |
+
return T
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
@spaces.GPU(duration=150, size="xlarge")
|
| 133 |
+
def synthesize(image, yaw, pitch, radius, progress=gr.Progress(track_tqdm=True)):
|
| 134 |
+
"""Synthesize a novel view of a single input image at a target camera pose.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
image: Source image (a single monocular view).
|
| 138 |
+
yaw: Horizontal camera rotation in degrees (positive = right, negative = left).
|
| 139 |
+
pitch: Vertical camera rotation in degrees (positive = up, negative = down).
|
| 140 |
+
radius: Rotation radius. If 0, it is auto-derived from the scene center depth.
|
| 141 |
+
|
| 142 |
+
Returns:
|
| 143 |
+
The synthesized novel-view image at the requested camera pose.
|
| 144 |
+
"""
|
| 145 |
+
if image is None:
|
| 146 |
+
raise gr.Error("Please provide an input image.")
|
| 147 |
+
|
| 148 |
+
original_image = image.convert("RGB")
|
| 149 |
+
edit_image = original_image.resize((GEN_W, GEN_H))
|
| 150 |
+
|
| 151 |
+
with torch.inference_mode():
|
| 152 |
+
# --- 1. 3D feature extraction (DA3 GIANT) + intrinsics ---
|
| 153 |
+
feat_out = da3_giant.inference(
|
| 154 |
+
[edit_image], export_feat_layers=EXPORT_3D_FEAT_LAYERS, process_res=840
|
| 155 |
+
)
|
| 156 |
+
intri = feat_out.intrinsics[0]
|
| 157 |
+
width = intri[0, 2] * 2
|
| 158 |
+
height = intri[1, 2] * 2
|
| 159 |
+
Ks_matrix = [
|
| 160 |
+
[intri[0, 0] / width, 0.0, 0.0],
|
| 161 |
+
[0.0, intri[1, 1] / height, 0.0],
|
| 162 |
+
[0.0, 0.0, 1.0],
|
| 163 |
+
]
|
| 164 |
+
Ks = torch.Tensor(Ks_matrix)
|
| 165 |
+
Ks = torch.stack([Ks, Ks], dim=0).unsqueeze(0) # (1, 2, 3, 3)
|
| 166 |
+
|
| 167 |
+
feats = [torch.from_numpy(feat_out.aux[f"feat_layer_{layer}"])
|
| 168 |
+
for layer in EXPORT_3D_FEAT_LAYERS]
|
| 169 |
+
feat_3D = torch.cat(feats, dim=-1).to(dtype=dtype, device=device)
|
| 170 |
+
|
| 171 |
+
# --- 2. Dense depth estimation (DA3 NESTED) ---
|
| 172 |
+
prediction = da3_nested.inference([edit_image], process_res=840)
|
| 173 |
+
depth_edit = torch.Tensor(prediction.depth).unsqueeze(0)
|
| 174 |
+
depth_edit = F.interpolate(depth_edit, size=(GEN_H, GEN_W),
|
| 175 |
+
mode="bilinear", align_corners=False)[0]
|
| 176 |
+
depth_latent = torch.zeros_like(depth_edit)
|
| 177 |
+
depth = torch.cat([depth_latent, depth_edit], dim=0).unsqueeze(0) # (1, 2, H, W)
|
| 178 |
+
|
| 179 |
+
# --- 3. Target pose ---
|
| 180 |
+
r = float(radius)
|
| 181 |
+
if r <= 0:
|
| 182 |
+
depth_squeeze = depth[0, 1]
|
| 183 |
+
r = depth_squeeze[depth_squeeze.shape[0] // 2,
|
| 184 |
+
depth_squeeze.shape[1] // 2].item()
|
| 185 |
+
extrinsic_target = compute_target_extrinsic(float(yaw), float(pitch), r)
|
| 186 |
+
extrinsic_source = np.eye(4)
|
| 187 |
+
viewmats = torch.Tensor(
|
| 188 |
+
np.stack((extrinsic_target, extrinsic_source), axis=0)
|
| 189 |
+
).unsqueeze(0) # (1, 2, 4, 4) -> [target, source]
|
| 190 |
+
|
| 191 |
+
# --- 4. Novel view generation (MetaView DiT) ---
|
| 192 |
+
generated_image = pipe(
|
| 193 |
+
PROMPT, edit_image=edit_image, edit_image_auto_resize=False,
|
| 194 |
+
seed=0,
|
| 195 |
+
viewmats=viewmats.to(device=device, dtype=dtype),
|
| 196 |
+
Ks=Ks.to(device=device, dtype=dtype),
|
| 197 |
+
prope_dim_arrange=PROPE_DIM_ARRANGE,
|
| 198 |
+
add_attn=True,
|
| 199 |
+
add_3D=True,
|
| 200 |
+
feat_3D=feat_3D,
|
| 201 |
+
depth=depth.to(device=device, dtype=dtype) if ADD_DEPTH else None,
|
| 202 |
+
merge_3D=MERGE_3D,
|
| 203 |
+
val=True,
|
| 204 |
+
num_inference_steps=NUM_STEPS,
|
| 205 |
+
height=GEN_H, width=GEN_W,
|
| 206 |
+
)
|
| 207 |
+
|
| 208 |
+
return generated_image
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
CSS = """
|
| 212 |
+
#col-container { max-width: 1100px; margin: 0 auto; }
|
| 213 |
+
.dark .gradio-container { color: var(--body-text-color); }
|
| 214 |
+
"""
|
| 215 |
+
|
| 216 |
+
with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
|
| 217 |
+
with gr.Column(elem_id="col-container"):
|
| 218 |
+
gr.Markdown(
|
| 219 |
+
"""
|
| 220 |
+
# MetaView — Monocular Novel View Synthesis
|
| 221 |
+
Synthesize a **novel camera view** from a single image. Upload an image,
|
| 222 |
+
pick a target camera **yaw** (left/right) and **pitch** (up/down), and
|
| 223 |
+
MetaView renders the scene from that new viewpoint.
|
| 224 |
+
|
| 225 |
+
Built on Qwen-Image-Edit + Depth-Anything-3 geometry priors.
|
| 226 |
+
[Model](https://huggingface.co/Kwai-Kolors/MetaView) ·
|
| 227 |
+
[Code](https://github.com/KlingAIResearch/MetaView) ·
|
| 228 |
+
[Paper](https://arxiv.org/abs/2607.12000)
|
| 229 |
+
"""
|
| 230 |
+
)
|
| 231 |
+
with gr.Row():
|
| 232 |
+
with gr.Column():
|
| 233 |
+
image = gr.Image(label="Input image", type="pil", height=340)
|
| 234 |
+
with gr.Row():
|
| 235 |
+
yaw = gr.Slider(-60, 60, value=-30, step=1,
|
| 236 |
+
label="Yaw (°) ← left | right →")
|
| 237 |
+
pitch = gr.Slider(-45, 45, value=10, step=1,
|
| 238 |
+
label="Pitch (°) ↓ down | up ↑")
|
| 239 |
+
run = gr.Button("Synthesize novel view", variant="primary")
|
| 240 |
+
with gr.Accordion("Advanced settings", open=False):
|
| 241 |
+
radius = gr.Slider(
|
| 242 |
+
0.0, 10.0, value=0.0, step=0.1,
|
| 243 |
+
label="Rotation radius (0 = auto from center depth)",
|
| 244 |
+
)
|
| 245 |
+
with gr.Column():
|
| 246 |
+
output = gr.Image(label="Novel view", height=340)
|
| 247 |
+
|
| 248 |
+
gr.Examples(
|
| 249 |
+
examples=[
|
| 250 |
+
["examples/1.png", -30, 10],
|
| 251 |
+
["examples/5.png", 30, 0],
|
| 252 |
+
["examples/9.png", -25, 15],
|
| 253 |
+
["examples/12.png", 40, -10],
|
| 254 |
+
],
|
| 255 |
+
inputs=[image, yaw, pitch],
|
| 256 |
+
outputs=output,
|
| 257 |
+
fn=synthesize,
|
| 258 |
+
cache_examples=True,
|
| 259 |
+
cache_mode="lazy",
|
| 260 |
+
)
|
| 261 |
+
|
| 262 |
+
run.click(
|
| 263 |
+
synthesize,
|
| 264 |
+
inputs=[image, yaw, pitch, radius],
|
| 265 |
+
outputs=output,
|
| 266 |
+
api_name="synthesize",
|
| 267 |
+
)
|
| 268 |
+
|
| 269 |
+
if __name__ == "__main__":
|
| 270 |
+
demo.launch(mcp_server=True)
|
depth_anything_3/api.py
ADDED
|
@@ -0,0 +1,451 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""
|
| 15 |
+
Depth Anything 3 API module.
|
| 16 |
+
|
| 17 |
+
This module provides the main API for Depth Anything 3, including model loading,
|
| 18 |
+
inference, and export capabilities. It supports both single and nested model architectures.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
from __future__ import annotations
|
| 22 |
+
|
| 23 |
+
import time
|
| 24 |
+
from typing import Optional, Sequence
|
| 25 |
+
import numpy as np
|
| 26 |
+
import torch
|
| 27 |
+
import torch.nn as nn
|
| 28 |
+
from huggingface_hub import PyTorchModelHubMixin
|
| 29 |
+
from PIL import Image
|
| 30 |
+
|
| 31 |
+
from depth_anything_3.cfg import create_object, load_config
|
| 32 |
+
from depth_anything_3.registry import MODEL_REGISTRY
|
| 33 |
+
from depth_anything_3.specs import Prediction
|
| 34 |
+
from depth_anything_3.utils.geometry import affine_inverse
|
| 35 |
+
from depth_anything_3.utils.io.input_processor import InputProcessor
|
| 36 |
+
from depth_anything_3.utils.io.output_processor import OutputProcessor
|
| 37 |
+
from depth_anything_3.utils.logger import logger
|
| 38 |
+
|
| 39 |
+
# NOTE: `export` (gsplat / moviepy / open3d / pycolmap) and `align_poses_umeyama`
|
| 40 |
+
# (evo) are only used in the export / pose-alignment code paths that this demo
|
| 41 |
+
# never exercises. Import them lazily inside the functions that use them so the
|
| 42 |
+
# heavy, hard-to-build deps aren't required at module import time.
|
| 43 |
+
|
| 44 |
+
torch.backends.cudnn.benchmark = False
|
| 45 |
+
# logger.info("CUDNN Benchmark Disabled")
|
| 46 |
+
|
| 47 |
+
SAFETENSORS_NAME = "model.safetensors"
|
| 48 |
+
CONFIG_NAME = "config.json"
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
class DepthAnything3(nn.Module, PyTorchModelHubMixin):
|
| 52 |
+
"""
|
| 53 |
+
Depth Anything 3 main API class.
|
| 54 |
+
|
| 55 |
+
This class provides a high-level interface for depth estimation using Depth Anything 3.
|
| 56 |
+
It supports both single and nested model architectures with metric scaling capabilities.
|
| 57 |
+
|
| 58 |
+
Features:
|
| 59 |
+
- Hugging Face Hub integration via PyTorchModelHubMixin
|
| 60 |
+
- Support for multiple model presets (vitb, vitg, nested variants)
|
| 61 |
+
- Automatic mixed precision inference
|
| 62 |
+
- Export capabilities for various formats (GLB, PLY, NPZ, etc.)
|
| 63 |
+
- Camera pose estimation and metric depth scaling
|
| 64 |
+
|
| 65 |
+
Usage:
|
| 66 |
+
# Load from Hugging Face Hub
|
| 67 |
+
model = DepthAnything3.from_pretrained("huggingface/model-name")
|
| 68 |
+
|
| 69 |
+
# Or create with specific preset
|
| 70 |
+
model = DepthAnything3(preset="vitg")
|
| 71 |
+
|
| 72 |
+
# Run inference
|
| 73 |
+
prediction = model.inference(images, export_dir="output", export_format="glb")
|
| 74 |
+
"""
|
| 75 |
+
|
| 76 |
+
_commit_hash: str | None = None # Set by mixin when loading from Hub
|
| 77 |
+
|
| 78 |
+
def __init__(self, model_name: str = "da3-large", **kwargs):
|
| 79 |
+
"""
|
| 80 |
+
Initialize DepthAnything3 with specified preset.
|
| 81 |
+
|
| 82 |
+
Args:
|
| 83 |
+
model_name: The name of the model preset to use.
|
| 84 |
+
Examples: 'da3-giant', 'da3-large', 'da3metric-large', 'da3nested-giant-large'.
|
| 85 |
+
**kwargs: Additional keyword arguments (currently unused).
|
| 86 |
+
"""
|
| 87 |
+
super().__init__()
|
| 88 |
+
self.model_name = model_name
|
| 89 |
+
|
| 90 |
+
# Build the underlying network
|
| 91 |
+
self.config = load_config(MODEL_REGISTRY[self.model_name])
|
| 92 |
+
self.model = create_object(self.config)
|
| 93 |
+
self.model.eval()
|
| 94 |
+
|
| 95 |
+
# Initialize processors
|
| 96 |
+
self.input_processor = InputProcessor()
|
| 97 |
+
self.output_processor = OutputProcessor()
|
| 98 |
+
|
| 99 |
+
# Device management (set by user)
|
| 100 |
+
self.device = None
|
| 101 |
+
|
| 102 |
+
@torch.inference_mode()
|
| 103 |
+
def forward(
|
| 104 |
+
self,
|
| 105 |
+
image: torch.Tensor,
|
| 106 |
+
extrinsics: torch.Tensor | None = None,
|
| 107 |
+
intrinsics: torch.Tensor | None = None,
|
| 108 |
+
export_feat_layers: list[int] | None = None,
|
| 109 |
+
infer_gs: bool = False,
|
| 110 |
+
use_ray_pose: bool = False,
|
| 111 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 112 |
+
) -> dict[str, torch.Tensor]:
|
| 113 |
+
"""
|
| 114 |
+
Forward pass through the model.
|
| 115 |
+
|
| 116 |
+
Args:
|
| 117 |
+
image: Input batch with shape ``(B, N, 3, H, W)`` on the model device.
|
| 118 |
+
extrinsics: Optional camera extrinsics with shape ``(B, N, 4, 4)``.
|
| 119 |
+
intrinsics: Optional camera intrinsics with shape ``(B, N, 3, 3)``.
|
| 120 |
+
export_feat_layers: Layer indices to return intermediate features for.
|
| 121 |
+
infer_gs: Enable Gaussian Splatting branch.
|
| 122 |
+
use_ray_pose: Use ray-based pose estimation instead of camera decoder.
|
| 123 |
+
ref_view_strategy: Strategy for selecting reference view from multiple views.
|
| 124 |
+
|
| 125 |
+
Returns:
|
| 126 |
+
Dictionary containing model predictions
|
| 127 |
+
"""
|
| 128 |
+
# Determine optimal autocast dtype
|
| 129 |
+
autocast_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
| 130 |
+
with torch.no_grad():
|
| 131 |
+
with torch.autocast(device_type=image.device.type, dtype=autocast_dtype):
|
| 132 |
+
return self.model(
|
| 133 |
+
image, extrinsics, intrinsics, export_feat_layers, infer_gs, use_ray_pose, ref_view_strategy
|
| 134 |
+
)
|
| 135 |
+
|
| 136 |
+
def inference(
|
| 137 |
+
self,
|
| 138 |
+
image: list[np.ndarray | Image.Image | str],
|
| 139 |
+
extrinsics: np.ndarray | None = None,
|
| 140 |
+
intrinsics: np.ndarray | None = None,
|
| 141 |
+
align_to_input_ext_scale: bool = True,
|
| 142 |
+
infer_gs: bool = False,
|
| 143 |
+
use_ray_pose: bool = False,
|
| 144 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 145 |
+
render_exts: np.ndarray | None = None,
|
| 146 |
+
render_ixts: np.ndarray | None = None,
|
| 147 |
+
render_hw: tuple[int, int] | None = None,
|
| 148 |
+
process_res: int = 504,
|
| 149 |
+
process_res_method: str = "upper_bound_resize",
|
| 150 |
+
export_dir: str | None = None,
|
| 151 |
+
export_format: str = "mini_npz",
|
| 152 |
+
export_feat_layers: Sequence[int] | None = None,
|
| 153 |
+
# GLB export parameters
|
| 154 |
+
conf_thresh_percentile: float = 40.0,
|
| 155 |
+
num_max_points: int = 1_000_000,
|
| 156 |
+
show_cameras: bool = True,
|
| 157 |
+
# Feat_vis export parameters
|
| 158 |
+
feat_vis_fps: int = 15,
|
| 159 |
+
# Other export parameters, e.g., gs_ply, gs_video
|
| 160 |
+
export_kwargs: Optional[dict] = {},
|
| 161 |
+
) -> Prediction:
|
| 162 |
+
"""
|
| 163 |
+
Run inference on input images.
|
| 164 |
+
|
| 165 |
+
Args:
|
| 166 |
+
image: List of input images (numpy arrays, PIL Images, or file paths)
|
| 167 |
+
extrinsics: Camera extrinsics (N, 4, 4)
|
| 168 |
+
intrinsics: Camera intrinsics (N, 3, 3)
|
| 169 |
+
align_to_input_ext_scale: whether to align the input pose scale to the prediction
|
| 170 |
+
infer_gs: Enable the 3D Gaussian branch (needed for `gs_ply`/`gs_video` exports)
|
| 171 |
+
use_ray_pose: Use ray-based pose estimation instead of camera decoder (default: False)
|
| 172 |
+
ref_view_strategy: Strategy for selecting reference view from multiple views.
|
| 173 |
+
Options: "first", "middle", "saddle_balanced", "saddle_sim_range".
|
| 174 |
+
Default: "saddle_balanced". For single view input (S ≤ 2), no reordering is performed.
|
| 175 |
+
render_exts: Optional render extrinsics for Gaussian video export
|
| 176 |
+
render_ixts: Optional render intrinsics for Gaussian video export
|
| 177 |
+
render_hw: Optional render resolution for Gaussian video export
|
| 178 |
+
process_res: Processing resolution
|
| 179 |
+
process_res_method: Resize method for processing
|
| 180 |
+
export_dir: Directory to export results
|
| 181 |
+
export_format: Export format (mini_npz, npz, glb, ply, gs, gs_video)
|
| 182 |
+
export_feat_layers: Layer indices to export intermediate features from
|
| 183 |
+
conf_thresh_percentile: [GLB] Lower percentile for adaptive confidence threshold (default: 40.0) # noqa: E501
|
| 184 |
+
num_max_points: [GLB] Maximum number of points in the point cloud (default: 1,000,000)
|
| 185 |
+
show_cameras: [GLB] Show camera wireframes in the exported scene (default: True)
|
| 186 |
+
feat_vis_fps: [FEAT_VIS] Frame rate for output video (default: 15)
|
| 187 |
+
export_kwargs: additional arguments to export functions.
|
| 188 |
+
|
| 189 |
+
Returns:
|
| 190 |
+
Prediction object containing depth maps and camera parameters
|
| 191 |
+
"""
|
| 192 |
+
if "gs" in export_format:
|
| 193 |
+
assert infer_gs, "must set `infer_gs=True` to perform gs-related export."
|
| 194 |
+
|
| 195 |
+
if "colmap" in export_format:
|
| 196 |
+
assert isinstance(image[0], str), "`image` must be image paths for COLMAP export."
|
| 197 |
+
|
| 198 |
+
# Preprocess images
|
| 199 |
+
imgs_cpu, extrinsics, intrinsics = self._preprocess_inputs(
|
| 200 |
+
image, extrinsics, intrinsics, process_res, process_res_method
|
| 201 |
+
)
|
| 202 |
+
|
| 203 |
+
# Prepare tensors for model
|
| 204 |
+
imgs, ex_t, in_t = self._prepare_model_inputs(imgs_cpu, extrinsics, intrinsics)
|
| 205 |
+
|
| 206 |
+
# Normalize extrinsics
|
| 207 |
+
ex_t_norm = self._normalize_extrinsics(ex_t.clone() if ex_t is not None else None)
|
| 208 |
+
|
| 209 |
+
# Run model forward pass
|
| 210 |
+
export_feat_layers = list(export_feat_layers) if export_feat_layers is not None else []
|
| 211 |
+
|
| 212 |
+
raw_output = self._run_model_forward(
|
| 213 |
+
imgs, ex_t_norm, in_t, export_feat_layers, infer_gs, use_ray_pose, ref_view_strategy
|
| 214 |
+
)
|
| 215 |
+
|
| 216 |
+
# Convert raw output to prediction
|
| 217 |
+
prediction = self._convert_to_prediction(raw_output)
|
| 218 |
+
|
| 219 |
+
# Align prediction to extrinsincs
|
| 220 |
+
prediction = self._align_to_input_extrinsics_intrinsics(
|
| 221 |
+
extrinsics, intrinsics, prediction, align_to_input_ext_scale
|
| 222 |
+
)
|
| 223 |
+
|
| 224 |
+
# Add processed images for visualization
|
| 225 |
+
prediction = self._add_processed_images(prediction, imgs_cpu)
|
| 226 |
+
|
| 227 |
+
# Export if requested
|
| 228 |
+
if export_dir is not None:
|
| 229 |
+
|
| 230 |
+
if "gs" in export_format:
|
| 231 |
+
if infer_gs and "gs_video" not in export_format:
|
| 232 |
+
export_format = f"{export_format}-gs_video"
|
| 233 |
+
if "gs_video" in export_format:
|
| 234 |
+
if "gs_video" not in export_kwargs:
|
| 235 |
+
export_kwargs["gs_video"] = {}
|
| 236 |
+
export_kwargs["gs_video"].update(
|
| 237 |
+
{
|
| 238 |
+
"extrinsics": render_exts,
|
| 239 |
+
"intrinsics": render_ixts,
|
| 240 |
+
"out_image_hw": render_hw,
|
| 241 |
+
}
|
| 242 |
+
)
|
| 243 |
+
# Add GLB export parameters
|
| 244 |
+
if "glb" in export_format:
|
| 245 |
+
if "glb" not in export_kwargs:
|
| 246 |
+
export_kwargs["glb"] = {}
|
| 247 |
+
export_kwargs["glb"].update(
|
| 248 |
+
{
|
| 249 |
+
"conf_thresh_percentile": conf_thresh_percentile,
|
| 250 |
+
"num_max_points": num_max_points,
|
| 251 |
+
"show_cameras": show_cameras,
|
| 252 |
+
}
|
| 253 |
+
)
|
| 254 |
+
# Add Feat_vis export parameters
|
| 255 |
+
if "feat_vis" in export_format:
|
| 256 |
+
if "feat_vis" not in export_kwargs:
|
| 257 |
+
export_kwargs["feat_vis"] = {}
|
| 258 |
+
export_kwargs["feat_vis"].update(
|
| 259 |
+
{
|
| 260 |
+
"fps": feat_vis_fps,
|
| 261 |
+
}
|
| 262 |
+
)
|
| 263 |
+
# Add COLMAP export parameters
|
| 264 |
+
if "colmap" in export_format:
|
| 265 |
+
if "colmap" not in export_kwargs:
|
| 266 |
+
export_kwargs["colmap"] = {}
|
| 267 |
+
export_kwargs["colmap"].update(
|
| 268 |
+
{
|
| 269 |
+
"image_paths": image,
|
| 270 |
+
"conf_thresh_percentile": conf_thresh_percentile,
|
| 271 |
+
"process_res_method": process_res_method,
|
| 272 |
+
}
|
| 273 |
+
)
|
| 274 |
+
self._export_results(prediction, export_format, export_dir, **export_kwargs)
|
| 275 |
+
|
| 276 |
+
return prediction
|
| 277 |
+
|
| 278 |
+
def _preprocess_inputs(
|
| 279 |
+
self,
|
| 280 |
+
image: list[np.ndarray | Image.Image | str],
|
| 281 |
+
extrinsics: np.ndarray | None = None,
|
| 282 |
+
intrinsics: np.ndarray | None = None,
|
| 283 |
+
process_res: int = 504,
|
| 284 |
+
process_res_method: str = "upper_bound_resize",
|
| 285 |
+
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
|
| 286 |
+
"""Preprocess input images using input processor."""
|
| 287 |
+
start_time = time.time()
|
| 288 |
+
imgs_cpu, extrinsics, intrinsics = self.input_processor(
|
| 289 |
+
image,
|
| 290 |
+
extrinsics.copy() if extrinsics is not None else None,
|
| 291 |
+
intrinsics.copy() if intrinsics is not None else None,
|
| 292 |
+
process_res,
|
| 293 |
+
process_res_method,
|
| 294 |
+
)
|
| 295 |
+
end_time = time.time()
|
| 296 |
+
logger.info(
|
| 297 |
+
"Processed Images Done taking",
|
| 298 |
+
end_time - start_time,
|
| 299 |
+
"seconds. Shape: ",
|
| 300 |
+
imgs_cpu.shape,
|
| 301 |
+
)
|
| 302 |
+
return imgs_cpu, extrinsics, intrinsics
|
| 303 |
+
|
| 304 |
+
def _prepare_model_inputs(
|
| 305 |
+
self,
|
| 306 |
+
imgs_cpu: torch.Tensor,
|
| 307 |
+
extrinsics: torch.Tensor | None,
|
| 308 |
+
intrinsics: torch.Tensor | None,
|
| 309 |
+
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
|
| 310 |
+
"""Prepare tensors for model input."""
|
| 311 |
+
device = self._get_model_device()
|
| 312 |
+
|
| 313 |
+
# Move images to model device
|
| 314 |
+
imgs = imgs_cpu.to(device, non_blocking=True)[None].float()
|
| 315 |
+
|
| 316 |
+
# Convert camera parameters to tensors
|
| 317 |
+
ex_t = (
|
| 318 |
+
extrinsics.to(device, non_blocking=True)[None].float()
|
| 319 |
+
if extrinsics is not None
|
| 320 |
+
else None
|
| 321 |
+
)
|
| 322 |
+
in_t = (
|
| 323 |
+
intrinsics.to(device, non_blocking=True)[None].float()
|
| 324 |
+
if intrinsics is not None
|
| 325 |
+
else None
|
| 326 |
+
)
|
| 327 |
+
|
| 328 |
+
return imgs, ex_t, in_t
|
| 329 |
+
|
| 330 |
+
def _normalize_extrinsics(self, ex_t: torch.Tensor | None) -> torch.Tensor | None:
|
| 331 |
+
"""Normalize extrinsics"""
|
| 332 |
+
if ex_t is None:
|
| 333 |
+
return None
|
| 334 |
+
transform = affine_inverse(ex_t[:, :1])
|
| 335 |
+
ex_t_norm = ex_t @ transform
|
| 336 |
+
c2ws = affine_inverse(ex_t_norm)
|
| 337 |
+
translations = c2ws[..., :3, 3]
|
| 338 |
+
dists = translations.norm(dim=-1)
|
| 339 |
+
median_dist = torch.median(dists)
|
| 340 |
+
median_dist = torch.clamp(median_dist, min=1e-1)
|
| 341 |
+
ex_t_norm[..., :3, 3] = ex_t_norm[..., :3, 3] / median_dist
|
| 342 |
+
return ex_t_norm
|
| 343 |
+
|
| 344 |
+
def _align_to_input_extrinsics_intrinsics(
|
| 345 |
+
self,
|
| 346 |
+
extrinsics: torch.Tensor | None,
|
| 347 |
+
intrinsics: torch.Tensor | None,
|
| 348 |
+
prediction: Prediction,
|
| 349 |
+
align_to_input_ext_scale: bool = True,
|
| 350 |
+
ransac_view_thresh: int = 10,
|
| 351 |
+
) -> Prediction:
|
| 352 |
+
"""Align depth map to input extrinsics"""
|
| 353 |
+
if extrinsics is None:
|
| 354 |
+
return prediction
|
| 355 |
+
from depth_anything_3.utils.pose_align import align_poses_umeyama
|
| 356 |
+
prediction.intrinsics = intrinsics.numpy()
|
| 357 |
+
_, _, scale, aligned_extrinsics = align_poses_umeyama(
|
| 358 |
+
prediction.extrinsics,
|
| 359 |
+
extrinsics.numpy(),
|
| 360 |
+
ransac=len(extrinsics) >= ransac_view_thresh,
|
| 361 |
+
return_aligned=True,
|
| 362 |
+
random_state=42,
|
| 363 |
+
)
|
| 364 |
+
if align_to_input_ext_scale:
|
| 365 |
+
prediction.extrinsics = extrinsics[..., :3, :].numpy()
|
| 366 |
+
prediction.depth /= scale
|
| 367 |
+
else:
|
| 368 |
+
prediction.extrinsics = aligned_extrinsics
|
| 369 |
+
return prediction
|
| 370 |
+
|
| 371 |
+
def _run_model_forward(
|
| 372 |
+
self,
|
| 373 |
+
imgs: torch.Tensor,
|
| 374 |
+
ex_t: torch.Tensor | None,
|
| 375 |
+
in_t: torch.Tensor | None,
|
| 376 |
+
export_feat_layers: Sequence[int] | None = None,
|
| 377 |
+
infer_gs: bool = False,
|
| 378 |
+
use_ray_pose: bool = False,
|
| 379 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 380 |
+
) -> dict[str, torch.Tensor]:
|
| 381 |
+
"""Run model forward pass."""
|
| 382 |
+
device = imgs.device
|
| 383 |
+
need_sync = device.type == "cuda"
|
| 384 |
+
if need_sync:
|
| 385 |
+
torch.cuda.synchronize(device)
|
| 386 |
+
start_time = time.time()
|
| 387 |
+
feat_layers = list(export_feat_layers) if export_feat_layers is not None else None
|
| 388 |
+
output = self.forward(imgs, ex_t, in_t, feat_layers, infer_gs, use_ray_pose, ref_view_strategy)
|
| 389 |
+
if need_sync:
|
| 390 |
+
torch.cuda.synchronize(device)
|
| 391 |
+
end_time = time.time()
|
| 392 |
+
logger.info(f"Model Forward Pass Done. Time: {end_time - start_time} seconds")
|
| 393 |
+
return output
|
| 394 |
+
|
| 395 |
+
def _convert_to_prediction(self, raw_output: dict[str, torch.Tensor]) -> Prediction:
|
| 396 |
+
"""Convert raw model output to Prediction object."""
|
| 397 |
+
start_time = time.time()
|
| 398 |
+
output = self.output_processor(raw_output)
|
| 399 |
+
end_time = time.time()
|
| 400 |
+
logger.info(f"Conversion to Prediction Done. Time: {end_time - start_time} seconds")
|
| 401 |
+
return output
|
| 402 |
+
|
| 403 |
+
def _add_processed_images(self, prediction: Prediction, imgs_cpu: torch.Tensor) -> Prediction:
|
| 404 |
+
"""Add processed images to prediction for visualization."""
|
| 405 |
+
# Convert from (N, 3, H, W) to (N, H, W, 3) and denormalize
|
| 406 |
+
processed_imgs = imgs_cpu.permute(0, 2, 3, 1).cpu().numpy() # (N, H, W, 3)
|
| 407 |
+
|
| 408 |
+
# Denormalize from ImageNet normalization
|
| 409 |
+
mean = np.array([0.485, 0.456, 0.406])
|
| 410 |
+
std = np.array([0.229, 0.224, 0.225])
|
| 411 |
+
processed_imgs = processed_imgs * std + mean
|
| 412 |
+
processed_imgs = np.clip(processed_imgs, 0, 1)
|
| 413 |
+
processed_imgs = (processed_imgs * 255).astype(np.uint8)
|
| 414 |
+
|
| 415 |
+
prediction.processed_images = processed_imgs
|
| 416 |
+
return prediction
|
| 417 |
+
|
| 418 |
+
def _export_results(
|
| 419 |
+
self, prediction: Prediction, export_format: str, export_dir: str, **kwargs
|
| 420 |
+
) -> None:
|
| 421 |
+
"""Export results to specified format and directory."""
|
| 422 |
+
from depth_anything_3.utils.export import export
|
| 423 |
+
start_time = time.time()
|
| 424 |
+
export(prediction, export_format, export_dir, **kwargs)
|
| 425 |
+
end_time = time.time()
|
| 426 |
+
logger.info(f"Export Results Done. Time: {end_time - start_time} seconds")
|
| 427 |
+
|
| 428 |
+
def _get_model_device(self) -> torch.device:
|
| 429 |
+
"""
|
| 430 |
+
Get the device where the model is located.
|
| 431 |
+
|
| 432 |
+
Returns:
|
| 433 |
+
Device where the model parameters are located
|
| 434 |
+
|
| 435 |
+
Raises:
|
| 436 |
+
ValueError: If no tensors are found in the model
|
| 437 |
+
"""
|
| 438 |
+
if self.device is not None:
|
| 439 |
+
return self.device
|
| 440 |
+
|
| 441 |
+
# Find device from parameters
|
| 442 |
+
for param in self.parameters():
|
| 443 |
+
self.device = param.device
|
| 444 |
+
return param.device
|
| 445 |
+
|
| 446 |
+
# Find device from buffers
|
| 447 |
+
for buffer in self.buffers():
|
| 448 |
+
self.device = buffer.device
|
| 449 |
+
return buffer.device
|
| 450 |
+
|
| 451 |
+
raise ValueError("No tensor found in model")
|
depth_anything_3/app/css_and_html.py
ADDED
|
@@ -0,0 +1,594 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# flake8: noqa: E501
|
| 2 |
+
|
| 3 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 4 |
+
#
|
| 5 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 6 |
+
# you may not use this file except in compliance with the License.
|
| 7 |
+
# You may obtain a copy of the License at
|
| 8 |
+
#
|
| 9 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 10 |
+
#
|
| 11 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 12 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 13 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 14 |
+
# See the License for the specific language governing permissions and
|
| 15 |
+
# limitations under the License.
|
| 16 |
+
|
| 17 |
+
"""
|
| 18 |
+
CSS and HTML content for the Depth Anything 3 Gradio application.
|
| 19 |
+
This module contains all the CSS styles and HTML content blocks
|
| 20 |
+
used in the Gradio interface.
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
# CSS Styles for the Gradio interface
|
| 24 |
+
GRADIO_CSS = """
|
| 25 |
+
/* Add Font Awesome CDN with all styles including brands and colors */
|
| 26 |
+
@import url('https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css');
|
| 27 |
+
|
| 28 |
+
/* Add custom styles for colored icons */
|
| 29 |
+
.fa-color-blue {
|
| 30 |
+
color: #3b82f6;
|
| 31 |
+
}
|
| 32 |
+
|
| 33 |
+
.fa-color-purple {
|
| 34 |
+
color: #8b5cf6;
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
.fa-color-cyan {
|
| 38 |
+
color: #06b6d4;
|
| 39 |
+
}
|
| 40 |
+
|
| 41 |
+
.fa-color-green {
|
| 42 |
+
color: #10b981;
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
.fa-color-yellow {
|
| 46 |
+
color: #f59e0b;
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
.fa-color-red {
|
| 50 |
+
color: #ef4444;
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
.link-btn {
|
| 54 |
+
display: inline-flex;
|
| 55 |
+
align-items: center;
|
| 56 |
+
gap: 8px;
|
| 57 |
+
text-decoration: none;
|
| 58 |
+
padding: 12px 24px;
|
| 59 |
+
border-radius: 50px;
|
| 60 |
+
font-weight: 500;
|
| 61 |
+
transition: all 0.3s ease;
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
/* Dark mode tech theme */
|
| 65 |
+
@media (prefers-color-scheme: dark) {
|
| 66 |
+
html, body {
|
| 67 |
+
background: #1e293b;
|
| 68 |
+
color: #ffffff;
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
.gradio-container {
|
| 72 |
+
background: #1e293b;
|
| 73 |
+
color: #ffffff;
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
.link-btn {
|
| 77 |
+
background: rgba(255, 255, 255, 0.2);
|
| 78 |
+
color: white;
|
| 79 |
+
backdrop-filter: blur(10px);
|
| 80 |
+
border: 1px solid rgba(255, 255, 255, 0.3);
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
.link-btn:hover {
|
| 84 |
+
background: rgba(255, 255, 255, 0.3);
|
| 85 |
+
transform: translateY(-2px);
|
| 86 |
+
box-shadow: 0 8px 25px rgba(0, 0, 0, 0.2);
|
| 87 |
+
}
|
| 88 |
+
|
| 89 |
+
.tech-bg {
|
| 90 |
+
background: linear-gradient(135deg, #0f172a, #1e293b); /* Darker colors */
|
| 91 |
+
position: relative;
|
| 92 |
+
overflow: hidden;
|
| 93 |
+
}
|
| 94 |
+
|
| 95 |
+
.tech-bg::before {
|
| 96 |
+
content: '';
|
| 97 |
+
position: absolute;
|
| 98 |
+
top: 0;
|
| 99 |
+
left: 0;
|
| 100 |
+
right: 0;
|
| 101 |
+
bottom: 0;
|
| 102 |
+
background:
|
| 103 |
+
radial-gradient(circle at 20% 80%, rgba(59, 130, 246, 0.15) 0%, transparent 50%), /* Reduced opacity */
|
| 104 |
+
radial-gradient(circle at 80% 20%, rgba(139, 92, 246, 0.15) 0%, transparent 50%), /* Reduced opacity */
|
| 105 |
+
radial-gradient(circle at 40% 40%, rgba(18, 194, 233, 0.1) 0%, transparent 50%); /* Reduced opacity */
|
| 106 |
+
animation: techPulse 8s ease-in-out infinite;
|
| 107 |
+
}
|
| 108 |
+
|
| 109 |
+
.gradio-container .panel,
|
| 110 |
+
.gradio-container .block,
|
| 111 |
+
.gradio-container .form {
|
| 112 |
+
background: rgba(0, 0, 0, 0.3);
|
| 113 |
+
border: 1px solid rgba(59, 130, 246, 0.2);
|
| 114 |
+
border-radius: 10px;
|
| 115 |
+
}
|
| 116 |
+
|
| 117 |
+
.gradio-container * {
|
| 118 |
+
color: #ffffff;
|
| 119 |
+
}
|
| 120 |
+
|
| 121 |
+
.gradio-container label {
|
| 122 |
+
color: #e0e0e0;
|
| 123 |
+
}
|
| 124 |
+
|
| 125 |
+
.gradio-container .markdown {
|
| 126 |
+
color: #e0e0e0;
|
| 127 |
+
}
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
/* Light mode tech theme */
|
| 131 |
+
@media (prefers-color-scheme: light) {
|
| 132 |
+
html, body {
|
| 133 |
+
background: #ffffff;
|
| 134 |
+
color: #1e293b;
|
| 135 |
+
}
|
| 136 |
+
|
| 137 |
+
.gradio-container {
|
| 138 |
+
background: #ffffff;
|
| 139 |
+
color: #1e293b;
|
| 140 |
+
}
|
| 141 |
+
|
| 142 |
+
.tech-bg {
|
| 143 |
+
background: linear-gradient(135deg, #ffffff, #f1f5f9);
|
| 144 |
+
position: relative;
|
| 145 |
+
overflow: hidden;
|
| 146 |
+
}
|
| 147 |
+
|
| 148 |
+
.link-btn {
|
| 149 |
+
background: rgba(59, 130, 246, 0.15);
|
| 150 |
+
color: var(--body-text-color);
|
| 151 |
+
border: 1px solid rgba(59, 130, 246, 0.3);
|
| 152 |
+
}
|
| 153 |
+
|
| 154 |
+
.link-btn:hover {
|
| 155 |
+
background: rgba(59, 130, 246, 0.25);
|
| 156 |
+
transform: translateY(-2px);
|
| 157 |
+
box-shadow: 0 8px 25px rgba(59, 130, 246, 0.2);
|
| 158 |
+
}
|
| 159 |
+
|
| 160 |
+
.tech-bg::before {
|
| 161 |
+
content: '';
|
| 162 |
+
position: absolute;
|
| 163 |
+
top: 0;
|
| 164 |
+
left: 0;
|
| 165 |
+
right: 0;
|
| 166 |
+
bottom: 0;
|
| 167 |
+
background:
|
| 168 |
+
radial-gradient(circle at 20% 80%, rgba(59, 130, 246, 0.1) 0%, transparent 50%),
|
| 169 |
+
radial-gradient(circle at 80% 20%, rgba(139, 92, 246, 0.1) 0%, transparent 50%),
|
| 170 |
+
radial-gradient(circle at 40% 40%, rgba(18, 194, 233, 0.08) 0%, transparent 50%);
|
| 171 |
+
animation: techPulse 8s ease-in-out infinite;
|
| 172 |
+
}
|
| 173 |
+
|
| 174 |
+
.gradio-container .panel,
|
| 175 |
+
.gradio-container .block,
|
| 176 |
+
.gradio-container .form {
|
| 177 |
+
background: rgba(255, 255, 255, 0.8);
|
| 178 |
+
border: 1px solid rgba(59, 130, 246, 0.3);
|
| 179 |
+
border-radius: 10px;
|
| 180 |
+
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
|
| 181 |
+
}
|
| 182 |
+
|
| 183 |
+
.gradio-container * {
|
| 184 |
+
color: #1e293b;
|
| 185 |
+
}
|
| 186 |
+
|
| 187 |
+
.gradio-container label {
|
| 188 |
+
color: #334155;
|
| 189 |
+
}
|
| 190 |
+
|
| 191 |
+
.gradio-container .markdown {
|
| 192 |
+
color: #334155;
|
| 193 |
+
}
|
| 194 |
+
}
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
@keyframes techPulse {
|
| 200 |
+
0%, 100% { opacity: 0.5; }
|
| 201 |
+
50% { opacity: 0.8; }
|
| 202 |
+
}
|
| 203 |
+
|
| 204 |
+
/* Custom log with tech gradient */
|
| 205 |
+
.custom-log * {
|
| 206 |
+
font-style: italic;
|
| 207 |
+
font-size: 22px !important;
|
| 208 |
+
background: linear-gradient(135deg, #3b82f6, #8b5cf6);
|
| 209 |
+
background-size: 400% 400%;
|
| 210 |
+
-webkit-background-clip: text;
|
| 211 |
+
background-clip: text;
|
| 212 |
+
font-weight: bold !important;
|
| 213 |
+
color: transparent !important;
|
| 214 |
+
text-align: center !important;
|
| 215 |
+
animation: techGradient 3s ease infinite;
|
| 216 |
+
}
|
| 217 |
+
|
| 218 |
+
@keyframes techGradient {
|
| 219 |
+
0% { background-position: 0% 50%; }
|
| 220 |
+
50% { background-position: 100% 50%; }
|
| 221 |
+
100% { background-position: 0% 50%; }
|
| 222 |
+
}
|
| 223 |
+
|
| 224 |
+
@keyframes metricPulse {
|
| 225 |
+
0%, 100% { background-position: 0% 50%; }
|
| 226 |
+
50% { background-position: 100% 50%; }
|
| 227 |
+
}
|
| 228 |
+
|
| 229 |
+
@keyframes pointcloudPulse {
|
| 230 |
+
0%, 100% { background-position: 0% 50%; }
|
| 231 |
+
50% { background-position: 100% 50%; }
|
| 232 |
+
}
|
| 233 |
+
|
| 234 |
+
@keyframes camerasPulse {
|
| 235 |
+
0%, 100% { background-position: 0% 50%; }
|
| 236 |
+
50% { background-position: 100% 50%; }
|
| 237 |
+
}
|
| 238 |
+
|
| 239 |
+
@keyframes gaussiansPulse {
|
| 240 |
+
0%, 100% { background-position: 0% 50%; }
|
| 241 |
+
50% { background-position: 100% 50%; }
|
| 242 |
+
}
|
| 243 |
+
|
| 244 |
+
/* Special colors for key terms - Global styles */
|
| 245 |
+
.metric-text {
|
| 246 |
+
background: linear-gradient(45deg, #ff6b6b, #ff8e53, #ff6b6b);
|
| 247 |
+
background-size: 200% 200%;
|
| 248 |
+
-webkit-background-clip: text;
|
| 249 |
+
background-clip: text;
|
| 250 |
+
color: transparent !important;
|
| 251 |
+
animation: metricPulse 2s ease-in-out infinite;
|
| 252 |
+
font-weight: 700;
|
| 253 |
+
text-shadow: 0 0 10px rgba(255, 107, 107, 0.5);
|
| 254 |
+
}
|
| 255 |
+
|
| 256 |
+
.pointcloud-text {
|
| 257 |
+
background: linear-gradient(45deg, #4ecdc4, #44a08d, #4ecdc4);
|
| 258 |
+
background-size: 200% 200%;
|
| 259 |
+
-webkit-background-clip: text;
|
| 260 |
+
background-clip: text;
|
| 261 |
+
color: transparent !important;
|
| 262 |
+
animation: pointcloudPulse 2.5s ease-in-out infinite;
|
| 263 |
+
font-weight: 700;
|
| 264 |
+
text-shadow: 0 0 10px rgba(78, 205, 196, 0.5);
|
| 265 |
+
}
|
| 266 |
+
|
| 267 |
+
.cameras-text {
|
| 268 |
+
background: linear-gradient(45deg, #667eea, #764ba2, #667eea);
|
| 269 |
+
background-size: 200% 200%;
|
| 270 |
+
-webkit-background-clip: text;
|
| 271 |
+
background-clip: text;
|
| 272 |
+
color: transparent !important;
|
| 273 |
+
animation: camerasPulse 3s ease-in-out infinite;
|
| 274 |
+
font-weight: 700;
|
| 275 |
+
text-shadow: 0 0 10px rgba(102, 126, 234, 0.5);
|
| 276 |
+
}
|
| 277 |
+
|
| 278 |
+
.gaussians-text {
|
| 279 |
+
background: linear-gradient(45deg, #f093fb, #f5576c, #f093fb);
|
| 280 |
+
background-size: 200% 200%;
|
| 281 |
+
-webkit-background-clip: text;
|
| 282 |
+
background-clip: text;
|
| 283 |
+
color: transparent !important;
|
| 284 |
+
animation: gaussiansPulse 2.2s ease-in-out infinite;
|
| 285 |
+
font-weight: 700;
|
| 286 |
+
text-shadow: 0 0 10px rgba(240, 147, 251, 0.5);
|
| 287 |
+
}
|
| 288 |
+
|
| 289 |
+
.example-log * {
|
| 290 |
+
font-style: italic;
|
| 291 |
+
font-size: 16px !important;
|
| 292 |
+
background: linear-gradient(135deg, #3b82f6, #8b5cf6);
|
| 293 |
+
-webkit-background-clip: text;
|
| 294 |
+
background-clip: text;
|
| 295 |
+
color: transparent !important;
|
| 296 |
+
}
|
| 297 |
+
|
| 298 |
+
#my_radio .wrap {
|
| 299 |
+
display: flex;
|
| 300 |
+
flex-wrap: nowrap;
|
| 301 |
+
justify-content: center;
|
| 302 |
+
align-items: center;
|
| 303 |
+
}
|
| 304 |
+
|
| 305 |
+
#my_radio .wrap label {
|
| 306 |
+
display: flex;
|
| 307 |
+
width: 50%;
|
| 308 |
+
justify-content: center;
|
| 309 |
+
align-items: center;
|
| 310 |
+
margin: 0;
|
| 311 |
+
padding: 10px 0;
|
| 312 |
+
box-sizing: border-box;
|
| 313 |
+
}
|
| 314 |
+
|
| 315 |
+
/* Align navigation buttons with dropdown bottom */
|
| 316 |
+
.navigation-row {
|
| 317 |
+
display: flex !important;
|
| 318 |
+
align-items: flex-end !important;
|
| 319 |
+
gap: 8px !important;
|
| 320 |
+
}
|
| 321 |
+
|
| 322 |
+
.navigation-row > div:nth-child(1),
|
| 323 |
+
.navigation-row > div:nth-child(3) {
|
| 324 |
+
align-self: flex-end !important;
|
| 325 |
+
}
|
| 326 |
+
|
| 327 |
+
.navigation-row > div:nth-child(2) {
|
| 328 |
+
flex: 1 !important;
|
| 329 |
+
}
|
| 330 |
+
|
| 331 |
+
/* Make thumbnails clickable with pointer cursor */
|
| 332 |
+
.clickable-thumbnail img {
|
| 333 |
+
cursor: pointer !important;
|
| 334 |
+
}
|
| 335 |
+
|
| 336 |
+
.clickable-thumbnail:hover img {
|
| 337 |
+
cursor: pointer !important;
|
| 338 |
+
opacity: 0.8;
|
| 339 |
+
transition: opacity 0.3s ease;
|
| 340 |
+
}
|
| 341 |
+
|
| 342 |
+
/* Make thumbnail containers narrower horizontally */
|
| 343 |
+
.clickable-thumbnail {
|
| 344 |
+
padding: 5px 2px !important;
|
| 345 |
+
margin: 0 2px !important;
|
| 346 |
+
}
|
| 347 |
+
|
| 348 |
+
.clickable-thumbnail .image-container {
|
| 349 |
+
margin: 0 !important;
|
| 350 |
+
padding: 0 !important;
|
| 351 |
+
}
|
| 352 |
+
|
| 353 |
+
.scene-info {
|
| 354 |
+
text-align: center !important;
|
| 355 |
+
padding: 5px 2px !important;
|
| 356 |
+
margin: 0 !important;
|
| 357 |
+
}
|
| 358 |
+
"""
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
def get_header_html(logo_base64=None):
|
| 362 |
+
"""
|
| 363 |
+
Generate the main header HTML with logo and title.
|
| 364 |
+
|
| 365 |
+
Args:
|
| 366 |
+
logo_base64 (str, optional): Base64 encoded logo image
|
| 367 |
+
|
| 368 |
+
Returns:
|
| 369 |
+
str: HTML string for the header
|
| 370 |
+
"""
|
| 371 |
+
return """
|
| 372 |
+
<div class="tech-bg" style="text-align: center; margin-bottom: 5px; padding: 40px 20px; border-radius: 15px; position: relative; overflow: hidden;">
|
| 373 |
+
<div style="position: relative; z-index: 2;">
|
| 374 |
+
<h1 style="margin: 0; font-size: 3.5em; font-weight: 700;
|
| 375 |
+
background: linear-gradient(135deg, #3b82f6, #8b5cf6);
|
| 376 |
+
background-size: 400% 400%;
|
| 377 |
+
-webkit-background-clip: text;
|
| 378 |
+
background-clip: text;
|
| 379 |
+
color: transparent;
|
| 380 |
+
animation: techGradient 3s ease infinite;
|
| 381 |
+
text-shadow: 0 0 30px rgba(59, 130, 246, 0.5);
|
| 382 |
+
letter-spacing: 2px;">
|
| 383 |
+
Depth Anything 3
|
| 384 |
+
</h1>
|
| 385 |
+
<p style="margin: 15px 0 0 0; font-size: 2.16em; font-weight: 300;" class="header-subtitle">
|
| 386 |
+
Recovering the Visual Space from Any Views
|
| 387 |
+
</p>
|
| 388 |
+
<div style="margin-top: 20px;">
|
| 389 |
+
<!-- Revert buttons to original inline styles -->
|
| 390 |
+
<a href="https://depth-anything-3.github.io" target="_blank" class="link-btn">
|
| 391 |
+
<i class="fas fa-globe" style="margin-right: 8px;"></i> Project Page
|
| 392 |
+
</a>
|
| 393 |
+
<a href="https://arxiv.org/abs/2406.09414" target="_blank" class="link-btn">
|
| 394 |
+
<i class="fas fa-file-pdf" style="margin-right: 8px;"></i> Paper
|
| 395 |
+
</a>
|
| 396 |
+
<a href="https://github.com/ByteDance-Seed/Depth-Anything-3" target="_blank" class="link-btn">
|
| 397 |
+
<i class="fab fa-github" style="margin-right: 8px;"></i> Code
|
| 398 |
+
</a>
|
| 399 |
+
</div>
|
| 400 |
+
</div>
|
| 401 |
+
</div>
|
| 402 |
+
|
| 403 |
+
<style>
|
| 404 |
+
/* Ensure tech-bg class is properly applied in dark mode */
|
| 405 |
+
@media (prefers-color-scheme: dark) {
|
| 406 |
+
.header-subtitle {
|
| 407 |
+
color: #cbd5e1;
|
| 408 |
+
}
|
| 409 |
+
/* Increase priority to ensure background color is properly applied */
|
| 410 |
+
.tech-bg {
|
| 411 |
+
background: linear-gradient(135deg, #0f172a, #1e293b) !important;
|
| 412 |
+
}
|
| 413 |
+
}
|
| 414 |
+
|
| 415 |
+
@media (prefers-color-scheme: light) {
|
| 416 |
+
.header-subtitle {
|
| 417 |
+
color: #475569;
|
| 418 |
+
}
|
| 419 |
+
/* Also add explicit background color for light mode */
|
| 420 |
+
.tech-bg {
|
| 421 |
+
background: linear-gradient(135deg, rgba(59, 130, 246, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%) !important;
|
| 422 |
+
}
|
| 423 |
+
}
|
| 424 |
+
</style>
|
| 425 |
+
"""
|
| 426 |
+
|
| 427 |
+
|
| 428 |
+
def get_description_html():
|
| 429 |
+
"""
|
| 430 |
+
Generate the main description and getting started HTML.
|
| 431 |
+
|
| 432 |
+
Returns:
|
| 433 |
+
str: HTML string for the description
|
| 434 |
+
"""
|
| 435 |
+
return """
|
| 436 |
+
<div class="description-container" style="padding: 25px; border-radius: 15px; margin: 0 0 20px 0;">
|
| 437 |
+
<h2 class="description-title" style="margin-top: 0; font-size: 1.6em; text-align: center;">
|
| 438 |
+
<i class="fas fa-bullseye fa-color-red" style="margin-right: 8px;"></i> What This Demo Does
|
| 439 |
+
</h2>
|
| 440 |
+
<div class="description-content" style="padding: 20px; border-radius: 10px; margin: 15px 0; text-align: center;">
|
| 441 |
+
<p class="description-main" style="line-height: 1.6; margin: 0; font-size: 1.45em;">
|
| 442 |
+
<strong>Upload images or videos</strong> → <strong>Get <span class="metric-text">Metric</span> <span class="pointcloud-text">Point Clouds</span>, <span class="cameras-text">Cameras</span> and <span class="gaussians-text">Novel Views</span></strong> → <strong>Explore in 3D</strong>
|
| 443 |
+
</p>
|
| 444 |
+
</div>
|
| 445 |
+
|
| 446 |
+
<div style="text-align: center; margin-top: 15px;">
|
| 447 |
+
<p class="description-tip" style="font-style: italic; margin: 0;">
|
| 448 |
+
<i class="fas fa-lightbulb fa-color-yellow" style="margin-right: 8px;"></i> <strong>Tip:</strong> Landscape-oriented images or videos are preferred for best 3D recovering.
|
| 449 |
+
</p>
|
| 450 |
+
</div>
|
| 451 |
+
</div>
|
| 452 |
+
|
| 453 |
+
<style>
|
| 454 |
+
@media (prefers-color-scheme: dark) {
|
| 455 |
+
.description-container {
|
| 456 |
+
background: linear-gradient(135deg, rgba(59, 130, 246, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%);
|
| 457 |
+
border: 1px solid rgba(59, 130, 246, 0.2);
|
| 458 |
+
}
|
| 459 |
+
.description-title { color: #3b82f6; }
|
| 460 |
+
.description-content { background: rgba(0, 0, 0, 0.3); }
|
| 461 |
+
.description-main { color: #e0e0e0; }
|
| 462 |
+
.description-text { color: #cbd5e1; }
|
| 463 |
+
.description-tip { color: #cbd5e1; }
|
| 464 |
+
}
|
| 465 |
+
|
| 466 |
+
@media (prefers-color-scheme: light) {
|
| 467 |
+
.description-container {
|
| 468 |
+
background: linear-gradient(135deg, rgba(59, 130, 246, 0.05) 0%, rgba(139, 92, 246, 0.05) 100%);
|
| 469 |
+
border: 1px solid rgba(59, 130, 246, 0.3);
|
| 470 |
+
}
|
| 471 |
+
.description-title { color: #3b82f6; }
|
| 472 |
+
.description-content { background: transparent; }
|
| 473 |
+
.description-main { color: #1e293b; }
|
| 474 |
+
.description-text { color: #475569; }
|
| 475 |
+
.description-tip { color: #475569; }
|
| 476 |
+
}
|
| 477 |
+
</style>
|
| 478 |
+
"""
|
| 479 |
+
|
| 480 |
+
|
| 481 |
+
def get_acknowledgements_html():
|
| 482 |
+
"""
|
| 483 |
+
Generate the acknowledgements section HTML.
|
| 484 |
+
|
| 485 |
+
Returns:
|
| 486 |
+
str: HTML string for the acknowledgements
|
| 487 |
+
"""
|
| 488 |
+
return """
|
| 489 |
+
<div style="background: linear-gradient(135deg, rgba(59, 130, 246, 0.1) 0%, rgba(139, 92, 246, 0.1) 100%);
|
| 490 |
+
padding: 25px; border-radius: 15px; margin: 20px 0; border: 1px solid rgba(59, 130, 246, 0.2);">
|
| 491 |
+
<h3 style="color: #3b82f6; margin-top: 0; text-align: center; font-size: 1.4em;">
|
| 492 |
+
<i class="fas fa-trophy fa-color-yellow" style="margin-right: 8px;"></i> Research Credits & Acknowledgments
|
| 493 |
+
</h3>
|
| 494 |
+
|
| 495 |
+
<div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin: 15px 0;">
|
| 496 |
+
<!-- Original Research Section (Left) -->
|
| 497 |
+
<div style="text-align: center;">
|
| 498 |
+
<h4 style="color: #8b5cf6; margin: 10px 0;"><i class="fas fa-flask fa-color-green" style="margin-right: 8px;"></i> Original Research</h4>
|
| 499 |
+
<p style="color: #e0e0e0; margin: 5px 0;">
|
| 500 |
+
<a href="https://depth-anything-3.github.io" target="_blank"
|
| 501 |
+
style="color: #3b82f6; text-decoration: none; font-weight: 600;">
|
| 502 |
+
Depth Anything 3
|
| 503 |
+
</a>
|
| 504 |
+
</p>
|
| 505 |
+
</div>
|
| 506 |
+
|
| 507 |
+
<!-- Previous Versions Section (Right) -->
|
| 508 |
+
<div style="text-align: center;">
|
| 509 |
+
<h4 style="color: #8b5cf6; margin: 10px 0;"><i class="fas fa-history fa-color-blue" style="margin-right: 8px;"></i> Previous Versions</h4>
|
| 510 |
+
<div style="display: flex; flex-direction: row; gap: 15px; justify-content: center; align-items: center;">
|
| 511 |
+
<p style="color: #e0e0e0; margin: 0;">
|
| 512 |
+
<a href="https://huggingface.co/spaces/LiheYoung/Depth-Anything" target="_blank"
|
| 513 |
+
style="color: #3b82f6; text-decoration: none; font-weight: 600;">
|
| 514 |
+
Depth-Anything
|
| 515 |
+
</a>
|
| 516 |
+
</p>
|
| 517 |
+
<span style="color: #e0e0e0;">•</span>
|
| 518 |
+
<p style="color: #e0e0e0; margin: 0;">
|
| 519 |
+
<a href="https://huggingface.co/spaces/depth-anything/Depth-Anything-V2" target="_blank"
|
| 520 |
+
style="color: #3b82f6; text-decoration: none; font-weight: 600;">
|
| 521 |
+
Depth-Anything-V2
|
| 522 |
+
</a>
|
| 523 |
+
</p>
|
| 524 |
+
</div>
|
| 525 |
+
</div>
|
| 526 |
+
</div>
|
| 527 |
+
|
| 528 |
+
<!-- HF Demo Adapted from - Centered at the bottom of the whole block -->
|
| 529 |
+
<div style="margin-top: 20px; padding-top: 15px; border-top: 1px solid rgba(59, 130, 246, 0.3); text-align: center;">
|
| 530 |
+
<p style="color: #a0a0a0; font-size: 0.9em; margin: 0;">
|
| 531 |
+
<i class="fas fa-code-branch fa-color-gray" style="margin-right: 5px;"></i> HF demo adapted from <a href="https://huggingface.co/spaces/facebook/map-anything" target="_blank" style="color: inherit; text-decoration: none;">Map Anything</a>
|
| 532 |
+
</p>
|
| 533 |
+
</div>
|
| 534 |
+
</div>
|
| 535 |
+
"""
|
| 536 |
+
|
| 537 |
+
|
| 538 |
+
def get_gradio_theme():
|
| 539 |
+
"""
|
| 540 |
+
Get the configured Gradio theme with adaptive tech colors.
|
| 541 |
+
|
| 542 |
+
Returns:
|
| 543 |
+
gr.themes.Base: Configured Gradio theme
|
| 544 |
+
"""
|
| 545 |
+
import gradio as gr
|
| 546 |
+
|
| 547 |
+
return gr.themes.Base(
|
| 548 |
+
primary_hue=gr.themes.Color(
|
| 549 |
+
c50="#eff6ff",
|
| 550 |
+
c100="#dbeafe",
|
| 551 |
+
c200="#bfdbfe",
|
| 552 |
+
c300="#93c5fd",
|
| 553 |
+
c400="#60a5fa",
|
| 554 |
+
c500="#3b82f6",
|
| 555 |
+
c600="#2563eb",
|
| 556 |
+
c700="#1d4ed8",
|
| 557 |
+
c800="#1e40af",
|
| 558 |
+
c900="#1e3a8a",
|
| 559 |
+
c950="#172554",
|
| 560 |
+
),
|
| 561 |
+
secondary_hue=gr.themes.Color(
|
| 562 |
+
c50="#f5f3ff",
|
| 563 |
+
c100="#ede9fe",
|
| 564 |
+
c200="#ddd6fe",
|
| 565 |
+
c300="#c4b5fd",
|
| 566 |
+
c400="#a78bfa",
|
| 567 |
+
c500="#8b5cf6",
|
| 568 |
+
c600="#7c3aed",
|
| 569 |
+
c700="#6d28d9",
|
| 570 |
+
c800="#5b21b6",
|
| 571 |
+
c900="#4c1d95",
|
| 572 |
+
c950="#2e1065",
|
| 573 |
+
),
|
| 574 |
+
neutral_hue=gr.themes.Color(
|
| 575 |
+
c50="#f8fafc",
|
| 576 |
+
c100="#f1f5f9",
|
| 577 |
+
c200="#e2e8f0",
|
| 578 |
+
c300="#cbd5e1",
|
| 579 |
+
c400="#94a3b8",
|
| 580 |
+
c500="#64748b",
|
| 581 |
+
c600="#475569",
|
| 582 |
+
c700="#334155",
|
| 583 |
+
c800="#1e293b",
|
| 584 |
+
c900="#0f172a",
|
| 585 |
+
c950="#020617",
|
| 586 |
+
),
|
| 587 |
+
)
|
| 588 |
+
|
| 589 |
+
|
| 590 |
+
# Measure tab instructions HTML
|
| 591 |
+
MEASURE_INSTRUCTIONS_HTML = """
|
| 592 |
+
### Click points on the image to compute distance.
|
| 593 |
+
> <i class="fas fa-triangle-exclamation fa-color-red" style="margin-right: 5px;"></i> Metric scale estimation is difficult on aerial/drone images.
|
| 594 |
+
"""
|
depth_anything_3/app/gradio_app.py
ADDED
|
@@ -0,0 +1,724 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Refactored Gradio App for Depth Anything 3.
|
| 17 |
+
|
| 18 |
+
This is the main application file that orchestrates all components.
|
| 19 |
+
The original functionality has been split into modular components for better maintainability.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import argparse
|
| 23 |
+
import os
|
| 24 |
+
from typing import Any, Dict, List
|
| 25 |
+
import gradio as gr
|
| 26 |
+
|
| 27 |
+
from depth_anything_3.app.css_and_html import GRADIO_CSS, get_gradio_theme
|
| 28 |
+
from depth_anything_3.app.modules.event_handlers import EventHandlers
|
| 29 |
+
from depth_anything_3.app.modules.ui_components import UIComponents
|
| 30 |
+
|
| 31 |
+
# Set environment variables
|
| 32 |
+
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class DepthAnything3App:
|
| 36 |
+
"""
|
| 37 |
+
Main application class for Depth Anything 3 Gradio app.
|
| 38 |
+
"""
|
| 39 |
+
|
| 40 |
+
def __init__(self, model_dir: str = None, workspace_dir: str = None, gallery_dir: str = None):
|
| 41 |
+
"""
|
| 42 |
+
Initialize the application.
|
| 43 |
+
|
| 44 |
+
Args:
|
| 45 |
+
model_dir: Path to the model directory
|
| 46 |
+
workspace_dir: Path to the workspace directory
|
| 47 |
+
gallery_dir: Path to the gallery directory
|
| 48 |
+
"""
|
| 49 |
+
self.model_dir = model_dir
|
| 50 |
+
self.workspace_dir = workspace_dir
|
| 51 |
+
self.gallery_dir = gallery_dir
|
| 52 |
+
|
| 53 |
+
# Set environment variables for directories
|
| 54 |
+
if self.model_dir:
|
| 55 |
+
os.environ["DA3_MODEL_DIR"] = self.model_dir
|
| 56 |
+
if self.workspace_dir:
|
| 57 |
+
os.environ["DA3_WORKSPACE_DIR"] = self.workspace_dir
|
| 58 |
+
if self.gallery_dir:
|
| 59 |
+
os.environ["DA3_GALLERY_DIR"] = self.gallery_dir
|
| 60 |
+
|
| 61 |
+
self.event_handlers = EventHandlers()
|
| 62 |
+
self.ui_components = UIComponents()
|
| 63 |
+
|
| 64 |
+
def cache_examples(
|
| 65 |
+
self,
|
| 66 |
+
show_cam: bool = True,
|
| 67 |
+
filter_black_bg: bool = False,
|
| 68 |
+
filter_white_bg: bool = False,
|
| 69 |
+
save_percentage: float = 20.0,
|
| 70 |
+
num_max_points: int = 1000,
|
| 71 |
+
cache_gs_tag: str = "",
|
| 72 |
+
gs_trj_mode: str = "smooth",
|
| 73 |
+
gs_video_quality: str = "low",
|
| 74 |
+
) -> None:
|
| 75 |
+
"""
|
| 76 |
+
Pre-cache all example scenes at startup.
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
show_cam: Whether to show camera in visualization
|
| 80 |
+
filter_black_bg: Whether to filter black background
|
| 81 |
+
filter_white_bg: Whether to filter white background
|
| 82 |
+
save_percentage: Filter percentage for point cloud
|
| 83 |
+
num_max_points: Maximum number of points
|
| 84 |
+
cache_gs_tag: Tag to match scene names for high-res+3DGS caching (e.g., "dl3dv")
|
| 85 |
+
gs_trj_mode: Trajectory mode for 3DGS
|
| 86 |
+
gs_video_quality: Video quality for 3DGS
|
| 87 |
+
"""
|
| 88 |
+
from depth_anything_3.app.modules.utils import get_scene_info
|
| 89 |
+
|
| 90 |
+
examples_dir = os.path.join(self.workspace_dir, "examples")
|
| 91 |
+
if not os.path.exists(examples_dir):
|
| 92 |
+
print(f"Examples directory not found: {examples_dir}")
|
| 93 |
+
return
|
| 94 |
+
|
| 95 |
+
scenes = get_scene_info(examples_dir)
|
| 96 |
+
if not scenes:
|
| 97 |
+
print("No example scenes found to cache.")
|
| 98 |
+
return
|
| 99 |
+
|
| 100 |
+
print(f"\n{'='*60}")
|
| 101 |
+
print(f"Caching {len(scenes)} example scenes...")
|
| 102 |
+
print(f"{'='*60}\n")
|
| 103 |
+
|
| 104 |
+
for i, scene in enumerate(scenes, 1):
|
| 105 |
+
scene_name = scene["name"]
|
| 106 |
+
|
| 107 |
+
# Check if scene name matches the gs tag for high-res+3DGS caching
|
| 108 |
+
use_high_res_gs = cache_gs_tag and cache_gs_tag.lower() in scene_name.lower()
|
| 109 |
+
|
| 110 |
+
if use_high_res_gs:
|
| 111 |
+
print(f"[{i}/{len(scenes)}] Caching scene: {scene_name} (HIGH-RES + 3DGS)")
|
| 112 |
+
print(f" - Number of images: {scene['num_images']}")
|
| 113 |
+
print(f" - Matched tag: '{cache_gs_tag}' - using high_res + 3DGS")
|
| 114 |
+
else:
|
| 115 |
+
print(f"[{i}/{len(scenes)}] Caching scene: {scene_name} (LOW-RES)")
|
| 116 |
+
print(f" - Number of images: {scene['num_images']}")
|
| 117 |
+
|
| 118 |
+
try:
|
| 119 |
+
# Load example scene
|
| 120 |
+
_, target_dir, _, _, _, _, _, _, _ = self.event_handlers.load_example_scene(
|
| 121 |
+
scene_name
|
| 122 |
+
)
|
| 123 |
+
|
| 124 |
+
if target_dir and target_dir != "None":
|
| 125 |
+
# Run reconstruction with appropriate settings
|
| 126 |
+
print(" - Running reconstruction...")
|
| 127 |
+
result = self.event_handlers.gradio_demo(
|
| 128 |
+
target_dir=target_dir,
|
| 129 |
+
show_cam=show_cam,
|
| 130 |
+
filter_black_bg=filter_black_bg,
|
| 131 |
+
filter_white_bg=filter_white_bg,
|
| 132 |
+
process_res_method="high_res" if use_high_res_gs else "low_res",
|
| 133 |
+
save_percentage=save_percentage,
|
| 134 |
+
num_max_points=num_max_points,
|
| 135 |
+
infer_gs=use_high_res_gs,
|
| 136 |
+
ref_view_strategy="saddle_balanced",
|
| 137 |
+
gs_trj_mode=gs_trj_mode,
|
| 138 |
+
gs_video_quality=gs_video_quality,
|
| 139 |
+
)
|
| 140 |
+
|
| 141 |
+
# Check if successful
|
| 142 |
+
if result[0] is not None: # reconstruction_output
|
| 143 |
+
print(f" ✓ Scene '{scene_name}' cached successfully")
|
| 144 |
+
else:
|
| 145 |
+
print(f" ✗ Scene '{scene_name}' caching failed: {result[1]}")
|
| 146 |
+
else:
|
| 147 |
+
print(f" ✗ Scene '{scene_name}' loading failed")
|
| 148 |
+
|
| 149 |
+
except Exception as e:
|
| 150 |
+
print(f" ✗ Error caching scene '{scene_name}': {str(e)}")
|
| 151 |
+
|
| 152 |
+
print()
|
| 153 |
+
|
| 154 |
+
print("=" * 60)
|
| 155 |
+
print("Example scene caching completed!")
|
| 156 |
+
print("=" * 60 + "\n")
|
| 157 |
+
|
| 158 |
+
def create_app(self) -> gr.Blocks:
|
| 159 |
+
"""
|
| 160 |
+
Create and configure the Gradio application.
|
| 161 |
+
|
| 162 |
+
Returns:
|
| 163 |
+
Configured Gradio Blocks interface
|
| 164 |
+
"""
|
| 165 |
+
|
| 166 |
+
# Initialize theme
|
| 167 |
+
def get_theme():
|
| 168 |
+
return get_gradio_theme()
|
| 169 |
+
|
| 170 |
+
with gr.Blocks(theme=get_theme(), css=GRADIO_CSS) as demo:
|
| 171 |
+
# State variables for the tabbed interface
|
| 172 |
+
is_example = gr.Textbox(label="is_example", visible=False, value="None")
|
| 173 |
+
processed_data_state = gr.State(value=None)
|
| 174 |
+
measure_points_state = gr.State(value=[])
|
| 175 |
+
selected_image_index_state = gr.State(value=0) # Track selected image index
|
| 176 |
+
# current_view_index = gr.State(value=0) # noqa: F841 Track current view index
|
| 177 |
+
|
| 178 |
+
# Header and description
|
| 179 |
+
self.ui_components.create_header_section()
|
| 180 |
+
self.ui_components.create_description_section()
|
| 181 |
+
|
| 182 |
+
target_dir_output = gr.Textbox(label="Target Dir", visible=False, value="None")
|
| 183 |
+
|
| 184 |
+
# Main content area
|
| 185 |
+
with gr.Row():
|
| 186 |
+
with gr.Column(scale=2):
|
| 187 |
+
# Upload section
|
| 188 |
+
(
|
| 189 |
+
input_video,
|
| 190 |
+
s_time_interval,
|
| 191 |
+
input_images,
|
| 192 |
+
image_gallery,
|
| 193 |
+
) = self.ui_components.create_upload_section()
|
| 194 |
+
|
| 195 |
+
with gr.Column(scale=4):
|
| 196 |
+
with gr.Column():
|
| 197 |
+
# gr.Markdown("**Metric 3D Reconstruction (Point Cloud and Camera Poses)**")
|
| 198 |
+
# Reconstruction control section (buttons) - moved below tabs
|
| 199 |
+
|
| 200 |
+
log_output = gr.Markdown(
|
| 201 |
+
"Please upload a video or images, then click Reconstruct.",
|
| 202 |
+
elem_classes=["custom-log"],
|
| 203 |
+
)
|
| 204 |
+
|
| 205 |
+
# Tabbed interface
|
| 206 |
+
with gr.Tabs():
|
| 207 |
+
with gr.Tab("Point Cloud & Cameras"):
|
| 208 |
+
reconstruction_output = (
|
| 209 |
+
self.ui_components.create_3d_viewer_section()
|
| 210 |
+
)
|
| 211 |
+
|
| 212 |
+
with gr.Tab("Metric Depth"):
|
| 213 |
+
(
|
| 214 |
+
prev_measure_btn,
|
| 215 |
+
measure_view_selector,
|
| 216 |
+
next_measure_btn,
|
| 217 |
+
measure_image,
|
| 218 |
+
measure_depth_image,
|
| 219 |
+
measure_text,
|
| 220 |
+
) = self.ui_components.create_measure_section()
|
| 221 |
+
|
| 222 |
+
with gr.Tab("3DGS Rendered Novel Views"):
|
| 223 |
+
gs_video, gs_info = self.ui_components.create_nvs_video()
|
| 224 |
+
|
| 225 |
+
# Inference control section (before inference)
|
| 226 |
+
(process_res_method_dropdown, infer_gs, ref_view_strategy_dropdown) = (
|
| 227 |
+
self.ui_components.create_inference_control_section()
|
| 228 |
+
)
|
| 229 |
+
|
| 230 |
+
# Display control section - includes 3DGS options, buttons, and Visualization Options # noqa: E501
|
| 231 |
+
(
|
| 232 |
+
show_cam,
|
| 233 |
+
filter_black_bg,
|
| 234 |
+
filter_white_bg,
|
| 235 |
+
save_percentage,
|
| 236 |
+
num_max_points,
|
| 237 |
+
gs_trj_mode,
|
| 238 |
+
gs_video_quality,
|
| 239 |
+
submit_btn,
|
| 240 |
+
clear_btn,
|
| 241 |
+
) = self.ui_components.create_display_control_section()
|
| 242 |
+
|
| 243 |
+
# bind visibility of gs_trj_mode to infer_gs
|
| 244 |
+
infer_gs.change(
|
| 245 |
+
fn=lambda checked: (
|
| 246 |
+
gr.update(visible=checked),
|
| 247 |
+
gr.update(visible=checked),
|
| 248 |
+
gr.update(visible=checked),
|
| 249 |
+
gr.update(visible=(not checked)),
|
| 250 |
+
),
|
| 251 |
+
inputs=infer_gs,
|
| 252 |
+
outputs=[gs_trj_mode, gs_video_quality, gs_video, gs_info],
|
| 253 |
+
)
|
| 254 |
+
|
| 255 |
+
# Example scenes section
|
| 256 |
+
gr.Markdown("## Example Scenes")
|
| 257 |
+
|
| 258 |
+
scenes = self.ui_components.create_example_scenes_section()
|
| 259 |
+
scene_components = self.ui_components.create_example_scene_grid(scenes)
|
| 260 |
+
|
| 261 |
+
# Set up event handlers
|
| 262 |
+
self._setup_event_handlers(
|
| 263 |
+
demo,
|
| 264 |
+
is_example,
|
| 265 |
+
processed_data_state,
|
| 266 |
+
measure_points_state,
|
| 267 |
+
target_dir_output,
|
| 268 |
+
input_video,
|
| 269 |
+
input_images,
|
| 270 |
+
s_time_interval,
|
| 271 |
+
image_gallery,
|
| 272 |
+
reconstruction_output,
|
| 273 |
+
log_output,
|
| 274 |
+
show_cam,
|
| 275 |
+
filter_black_bg,
|
| 276 |
+
filter_white_bg,
|
| 277 |
+
process_res_method_dropdown,
|
| 278 |
+
save_percentage,
|
| 279 |
+
submit_btn,
|
| 280 |
+
clear_btn,
|
| 281 |
+
num_max_points,
|
| 282 |
+
infer_gs,
|
| 283 |
+
ref_view_strategy_dropdown,
|
| 284 |
+
selected_image_index_state,
|
| 285 |
+
measure_view_selector,
|
| 286 |
+
measure_image,
|
| 287 |
+
measure_depth_image,
|
| 288 |
+
measure_text,
|
| 289 |
+
prev_measure_btn,
|
| 290 |
+
next_measure_btn,
|
| 291 |
+
scenes,
|
| 292 |
+
scene_components,
|
| 293 |
+
gs_video,
|
| 294 |
+
gs_info,
|
| 295 |
+
gs_trj_mode,
|
| 296 |
+
gs_video_quality,
|
| 297 |
+
)
|
| 298 |
+
|
| 299 |
+
# Acknowledgements
|
| 300 |
+
self.ui_components.create_acknowledgements_section()
|
| 301 |
+
|
| 302 |
+
return demo
|
| 303 |
+
|
| 304 |
+
def _setup_event_handlers(
|
| 305 |
+
self,
|
| 306 |
+
demo: gr.Blocks,
|
| 307 |
+
is_example: gr.Textbox,
|
| 308 |
+
processed_data_state: gr.State,
|
| 309 |
+
measure_points_state: gr.State,
|
| 310 |
+
target_dir_output: gr.Textbox,
|
| 311 |
+
input_video: gr.Video,
|
| 312 |
+
input_images: gr.File,
|
| 313 |
+
s_time_interval: gr.Slider,
|
| 314 |
+
image_gallery: gr.Gallery,
|
| 315 |
+
reconstruction_output: gr.Model3D,
|
| 316 |
+
log_output: gr.Markdown,
|
| 317 |
+
show_cam: gr.Checkbox,
|
| 318 |
+
filter_black_bg: gr.Checkbox,
|
| 319 |
+
filter_white_bg: gr.Checkbox,
|
| 320 |
+
process_res_method_dropdown: gr.Dropdown,
|
| 321 |
+
save_percentage: gr.Slider,
|
| 322 |
+
submit_btn: gr.Button,
|
| 323 |
+
clear_btn: gr.ClearButton,
|
| 324 |
+
num_max_points: gr.Slider,
|
| 325 |
+
infer_gs: gr.Checkbox,
|
| 326 |
+
ref_view_strategy_dropdown: gr.Dropdown,
|
| 327 |
+
selected_image_index_state: gr.State,
|
| 328 |
+
measure_view_selector: gr.Dropdown,
|
| 329 |
+
measure_image: gr.Image,
|
| 330 |
+
measure_depth_image: gr.Image,
|
| 331 |
+
measure_text: gr.Markdown,
|
| 332 |
+
prev_measure_btn: gr.Button,
|
| 333 |
+
next_measure_btn: gr.Button,
|
| 334 |
+
scenes: List[Dict[str, Any]],
|
| 335 |
+
scene_components: List[gr.Image],
|
| 336 |
+
gs_video: gr.Video,
|
| 337 |
+
gs_info: gr.Markdown,
|
| 338 |
+
gs_trj_mode: gr.Dropdown,
|
| 339 |
+
gs_video_quality: gr.Dropdown,
|
| 340 |
+
) -> None:
|
| 341 |
+
"""
|
| 342 |
+
Set up all event handlers for the application.
|
| 343 |
+
|
| 344 |
+
Args:
|
| 345 |
+
demo: Gradio Blocks interface
|
| 346 |
+
All other arguments: Gradio components to connect
|
| 347 |
+
"""
|
| 348 |
+
# Configure clear button
|
| 349 |
+
clear_btn.add(
|
| 350 |
+
[
|
| 351 |
+
input_video,
|
| 352 |
+
input_images,
|
| 353 |
+
reconstruction_output,
|
| 354 |
+
log_output,
|
| 355 |
+
target_dir_output,
|
| 356 |
+
image_gallery,
|
| 357 |
+
gs_video,
|
| 358 |
+
]
|
| 359 |
+
)
|
| 360 |
+
|
| 361 |
+
# Main reconstruction button
|
| 362 |
+
submit_btn.click(
|
| 363 |
+
fn=self.event_handlers.clear_fields, inputs=[], outputs=[reconstruction_output]
|
| 364 |
+
).then(fn=self.event_handlers.update_log, inputs=[], outputs=[log_output]).then(
|
| 365 |
+
fn=self.event_handlers.gradio_demo,
|
| 366 |
+
inputs=[
|
| 367 |
+
target_dir_output,
|
| 368 |
+
show_cam,
|
| 369 |
+
filter_black_bg,
|
| 370 |
+
filter_white_bg,
|
| 371 |
+
process_res_method_dropdown,
|
| 372 |
+
save_percentage,
|
| 373 |
+
# pass num_max_points
|
| 374 |
+
num_max_points,
|
| 375 |
+
infer_gs,
|
| 376 |
+
ref_view_strategy_dropdown,
|
| 377 |
+
gs_trj_mode,
|
| 378 |
+
gs_video_quality,
|
| 379 |
+
],
|
| 380 |
+
outputs=[
|
| 381 |
+
reconstruction_output,
|
| 382 |
+
log_output,
|
| 383 |
+
processed_data_state,
|
| 384 |
+
measure_image,
|
| 385 |
+
measure_depth_image,
|
| 386 |
+
measure_text,
|
| 387 |
+
measure_view_selector,
|
| 388 |
+
gs_video,
|
| 389 |
+
gs_video, # gs_video visibility
|
| 390 |
+
gs_info, # gs_info visibility
|
| 391 |
+
],
|
| 392 |
+
).then(
|
| 393 |
+
fn=lambda: "False",
|
| 394 |
+
inputs=[],
|
| 395 |
+
outputs=[is_example], # set is_example to "False"
|
| 396 |
+
)
|
| 397 |
+
|
| 398 |
+
# Real-time visualization updates
|
| 399 |
+
self._setup_visualization_handlers(
|
| 400 |
+
show_cam,
|
| 401 |
+
filter_black_bg,
|
| 402 |
+
filter_white_bg,
|
| 403 |
+
process_res_method_dropdown,
|
| 404 |
+
target_dir_output,
|
| 405 |
+
is_example,
|
| 406 |
+
reconstruction_output,
|
| 407 |
+
log_output,
|
| 408 |
+
)
|
| 409 |
+
|
| 410 |
+
# File upload handlers
|
| 411 |
+
input_video.change(
|
| 412 |
+
fn=self.event_handlers.handle_uploads,
|
| 413 |
+
inputs=[input_video, input_images, s_time_interval],
|
| 414 |
+
outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
|
| 415 |
+
)
|
| 416 |
+
input_images.change(
|
| 417 |
+
fn=self.event_handlers.handle_uploads,
|
| 418 |
+
inputs=[input_video, input_images, s_time_interval],
|
| 419 |
+
outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
|
| 420 |
+
)
|
| 421 |
+
|
| 422 |
+
# Navigation handlers
|
| 423 |
+
self._setup_navigation_handlers(
|
| 424 |
+
prev_measure_btn,
|
| 425 |
+
next_measure_btn,
|
| 426 |
+
measure_view_selector,
|
| 427 |
+
measure_image,
|
| 428 |
+
measure_depth_image,
|
| 429 |
+
measure_points_state,
|
| 430 |
+
processed_data_state,
|
| 431 |
+
)
|
| 432 |
+
|
| 433 |
+
# Measurement handler
|
| 434 |
+
measure_image.select(
|
| 435 |
+
fn=self.event_handlers.measure,
|
| 436 |
+
inputs=[processed_data_state, measure_points_state, measure_view_selector],
|
| 437 |
+
outputs=[measure_image, measure_depth_image, measure_points_state, measure_text],
|
| 438 |
+
)
|
| 439 |
+
|
| 440 |
+
# Example scene handlers
|
| 441 |
+
self._setup_example_scene_handlers(
|
| 442 |
+
scenes,
|
| 443 |
+
scene_components,
|
| 444 |
+
reconstruction_output,
|
| 445 |
+
target_dir_output,
|
| 446 |
+
image_gallery,
|
| 447 |
+
log_output,
|
| 448 |
+
is_example,
|
| 449 |
+
processed_data_state,
|
| 450 |
+
measure_view_selector,
|
| 451 |
+
measure_image,
|
| 452 |
+
measure_depth_image,
|
| 453 |
+
gs_video,
|
| 454 |
+
gs_info,
|
| 455 |
+
)
|
| 456 |
+
|
| 457 |
+
def _setup_visualization_handlers(
|
| 458 |
+
self,
|
| 459 |
+
show_cam: gr.Checkbox,
|
| 460 |
+
filter_black_bg: gr.Checkbox,
|
| 461 |
+
filter_white_bg: gr.Checkbox,
|
| 462 |
+
process_res_method_dropdown: gr.Dropdown,
|
| 463 |
+
target_dir_output: gr.Textbox,
|
| 464 |
+
is_example: gr.Textbox,
|
| 465 |
+
reconstruction_output: gr.Model3D,
|
| 466 |
+
log_output: gr.Markdown,
|
| 467 |
+
) -> None:
|
| 468 |
+
"""Set up visualization update handlers."""
|
| 469 |
+
# Common inputs for visualization updates
|
| 470 |
+
viz_inputs = [
|
| 471 |
+
target_dir_output,
|
| 472 |
+
show_cam,
|
| 473 |
+
is_example,
|
| 474 |
+
filter_black_bg,
|
| 475 |
+
filter_white_bg,
|
| 476 |
+
process_res_method_dropdown,
|
| 477 |
+
]
|
| 478 |
+
|
| 479 |
+
# Set up change handlers for all visualization controls
|
| 480 |
+
for component in [show_cam, filter_black_bg, filter_white_bg]:
|
| 481 |
+
component.change(
|
| 482 |
+
fn=self.event_handlers.update_visualization,
|
| 483 |
+
inputs=viz_inputs,
|
| 484 |
+
outputs=[reconstruction_output, log_output],
|
| 485 |
+
)
|
| 486 |
+
|
| 487 |
+
def _setup_navigation_handlers(
|
| 488 |
+
self,
|
| 489 |
+
prev_measure_btn: gr.Button,
|
| 490 |
+
next_measure_btn: gr.Button,
|
| 491 |
+
measure_view_selector: gr.Dropdown,
|
| 492 |
+
measure_image: gr.Image,
|
| 493 |
+
measure_depth_image: gr.Image,
|
| 494 |
+
measure_points_state: gr.State,
|
| 495 |
+
processed_data_state: gr.State,
|
| 496 |
+
) -> None:
|
| 497 |
+
"""Set up navigation handlers for measure tab."""
|
| 498 |
+
# Measure tab navigation
|
| 499 |
+
prev_measure_btn.click(
|
| 500 |
+
fn=lambda processed_data, current_selector: self.event_handlers.navigate_measure_view(
|
| 501 |
+
processed_data, current_selector, -1
|
| 502 |
+
),
|
| 503 |
+
inputs=[processed_data_state, measure_view_selector],
|
| 504 |
+
outputs=[
|
| 505 |
+
measure_view_selector,
|
| 506 |
+
measure_image,
|
| 507 |
+
measure_depth_image,
|
| 508 |
+
measure_points_state,
|
| 509 |
+
],
|
| 510 |
+
)
|
| 511 |
+
|
| 512 |
+
next_measure_btn.click(
|
| 513 |
+
fn=lambda processed_data, current_selector: self.event_handlers.navigate_measure_view(
|
| 514 |
+
processed_data, current_selector, 1
|
| 515 |
+
),
|
| 516 |
+
inputs=[processed_data_state, measure_view_selector],
|
| 517 |
+
outputs=[
|
| 518 |
+
measure_view_selector,
|
| 519 |
+
measure_image,
|
| 520 |
+
measure_depth_image,
|
| 521 |
+
measure_points_state,
|
| 522 |
+
],
|
| 523 |
+
)
|
| 524 |
+
|
| 525 |
+
measure_view_selector.change(
|
| 526 |
+
fn=lambda processed_data, selector_value: (
|
| 527 |
+
self.event_handlers.update_measure_view(
|
| 528 |
+
processed_data, int(selector_value.split()[1]) - 1
|
| 529 |
+
)
|
| 530 |
+
if selector_value
|
| 531 |
+
else (None, None, [])
|
| 532 |
+
),
|
| 533 |
+
inputs=[processed_data_state, measure_view_selector],
|
| 534 |
+
outputs=[measure_image, measure_depth_image, measure_points_state],
|
| 535 |
+
)
|
| 536 |
+
|
| 537 |
+
def _setup_example_scene_handlers(
|
| 538 |
+
self,
|
| 539 |
+
scenes: List[Dict[str, Any]],
|
| 540 |
+
scene_components: List[gr.Image],
|
| 541 |
+
reconstruction_output: gr.Model3D,
|
| 542 |
+
target_dir_output: gr.Textbox,
|
| 543 |
+
image_gallery: gr.Gallery,
|
| 544 |
+
log_output: gr.Markdown,
|
| 545 |
+
is_example: gr.Textbox,
|
| 546 |
+
processed_data_state: gr.State,
|
| 547 |
+
measure_view_selector: gr.Dropdown,
|
| 548 |
+
measure_image: gr.Image,
|
| 549 |
+
measure_depth_image: gr.Image,
|
| 550 |
+
gs_video: gr.Video,
|
| 551 |
+
gs_info: gr.Markdown,
|
| 552 |
+
) -> None:
|
| 553 |
+
"""Set up example scene handlers."""
|
| 554 |
+
|
| 555 |
+
def load_and_update_measure(name):
|
| 556 |
+
result = self.event_handlers.load_example_scene(name)
|
| 557 |
+
# result = (reconstruction_output, target_dir, image_paths, log_message, processed_data, measure_view_selector, gs_video, gs_video_vis, gs_info_vis) # noqa: E501
|
| 558 |
+
|
| 559 |
+
# Update measure view if processed_data is available
|
| 560 |
+
measure_img = None
|
| 561 |
+
measure_depth = None
|
| 562 |
+
if result[4] is not None: # processed_data exists
|
| 563 |
+
measure_img, measure_depth, _ = (
|
| 564 |
+
self.event_handlers.visualization_handler.update_measure_view(result[4], 0)
|
| 565 |
+
)
|
| 566 |
+
|
| 567 |
+
return result + ("True", measure_img, measure_depth)
|
| 568 |
+
|
| 569 |
+
for i, scene in enumerate(scenes):
|
| 570 |
+
if i < len(scene_components):
|
| 571 |
+
scene_components[i].select(
|
| 572 |
+
fn=lambda name=scene["name"]: load_and_update_measure(name),
|
| 573 |
+
outputs=[
|
| 574 |
+
reconstruction_output,
|
| 575 |
+
target_dir_output,
|
| 576 |
+
image_gallery,
|
| 577 |
+
log_output,
|
| 578 |
+
processed_data_state,
|
| 579 |
+
measure_view_selector,
|
| 580 |
+
gs_video,
|
| 581 |
+
gs_video, # gs_video_visibility
|
| 582 |
+
gs_info, # gs_info_visibility
|
| 583 |
+
is_example,
|
| 584 |
+
measure_image,
|
| 585 |
+
measure_depth_image,
|
| 586 |
+
],
|
| 587 |
+
)
|
| 588 |
+
|
| 589 |
+
def launch(self, host: str = "127.0.0.1", port: int = 7860, **kwargs) -> None:
|
| 590 |
+
"""
|
| 591 |
+
Launch the application.
|
| 592 |
+
|
| 593 |
+
Args:
|
| 594 |
+
host: Host address to bind to
|
| 595 |
+
port: Port number to bind to
|
| 596 |
+
**kwargs: Additional arguments for demo.launch()
|
| 597 |
+
"""
|
| 598 |
+
demo = self.create_app()
|
| 599 |
+
demo.queue(max_size=20).launch(
|
| 600 |
+
show_error=True, ssr_mode=False, server_name=host, server_port=port, **kwargs
|
| 601 |
+
)
|
| 602 |
+
|
| 603 |
+
|
| 604 |
+
def main():
|
| 605 |
+
"""Main function to run the application."""
|
| 606 |
+
parser = argparse.ArgumentParser(
|
| 607 |
+
description="Depth Anything 3 Gradio Application",
|
| 608 |
+
formatter_class=argparse.RawDescriptionHelpFormatter,
|
| 609 |
+
epilog="""
|
| 610 |
+
Examples:
|
| 611 |
+
# Basic usage
|
| 612 |
+
python gradio_app.py --help
|
| 613 |
+
python gradio_app.py --host 0.0.0.0 --port 8080
|
| 614 |
+
python gradio_app.py --model-dir /path/to/model --workspace-dir /path/to/workspace
|
| 615 |
+
|
| 616 |
+
# Cache examples at startup (all low-res)
|
| 617 |
+
python gradio_app.py --cache-examples
|
| 618 |
+
|
| 619 |
+
# Cache with selective high-res+3DGS for scenes matching tag
|
| 620 |
+
python gradio_app.py --cache-examples --cache-gs-tag dl3dv
|
| 621 |
+
# This will use high-res + 3DGS for scenes containing "dl3dv" in their name,
|
| 622 |
+
# and low-res only for other scenes
|
| 623 |
+
""",
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Server configuration
|
| 627 |
+
parser.add_argument(
|
| 628 |
+
"--host", default="127.0.0.1", help="Host address to bind to (default: 127.0.0.1)"
|
| 629 |
+
)
|
| 630 |
+
parser.add_argument(
|
| 631 |
+
"--port", type=int, default=7860, help="Port number to bind to (default: 7860)"
|
| 632 |
+
)
|
| 633 |
+
|
| 634 |
+
# Directory configuration
|
| 635 |
+
parser.add_argument(
|
| 636 |
+
"--model-dir",
|
| 637 |
+
default="depth-anything/DA3NESTED-GIANT-LARGE",
|
| 638 |
+
help="Path to the model directory (default: depth-anything/DA3NESTED-GIANT-LARGE)",
|
| 639 |
+
)
|
| 640 |
+
parser.add_argument(
|
| 641 |
+
"--workspace-dir",
|
| 642 |
+
default="workspace/gradio", # noqa: E501
|
| 643 |
+
help="Path to the workspace directory (default: workspace/gradio)", # noqa: E501
|
| 644 |
+
)
|
| 645 |
+
parser.add_argument(
|
| 646 |
+
"--gallery-dir",
|
| 647 |
+
default="workspace/gallery",
|
| 648 |
+
help="Path to the gallery directory (default: workspace/gallery)", # noqa: E501
|
| 649 |
+
)
|
| 650 |
+
|
| 651 |
+
# Additional Gradio options
|
| 652 |
+
parser.add_argument("--share", action="store_true", help="Create a public link for the app")
|
| 653 |
+
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
|
| 654 |
+
|
| 655 |
+
# Example caching options
|
| 656 |
+
parser.add_argument(
|
| 657 |
+
"--cache-examples",
|
| 658 |
+
action="store_true",
|
| 659 |
+
help="Pre-cache all example scenes at startup for faster loading",
|
| 660 |
+
)
|
| 661 |
+
parser.add_argument(
|
| 662 |
+
"--cache-gs-tag",
|
| 663 |
+
type=str,
|
| 664 |
+
default="",
|
| 665 |
+
help="Tag to match scene names for high-res+3DGS caching (e.g., 'dl3dv'). Scenes containing this tag will use high_res and infer_gs=True; others will use low_res only.", # noqa: E501
|
| 666 |
+
)
|
| 667 |
+
|
| 668 |
+
args = parser.parse_args()
|
| 669 |
+
|
| 670 |
+
# Create directories if they don't exist
|
| 671 |
+
os.makedirs(args.workspace_dir, exist_ok=True)
|
| 672 |
+
os.makedirs(args.gallery_dir, exist_ok=True)
|
| 673 |
+
|
| 674 |
+
# Initialize and launch the application
|
| 675 |
+
app = DepthAnything3App(
|
| 676 |
+
model_dir=args.model_dir, workspace_dir=args.workspace_dir, gallery_dir=args.gallery_dir
|
| 677 |
+
)
|
| 678 |
+
|
| 679 |
+
# Prepare launch arguments
|
| 680 |
+
launch_kwargs = {"share": args.share, "debug": args.debug}
|
| 681 |
+
|
| 682 |
+
print("Starting Depth Anything 3 Gradio App...")
|
| 683 |
+
print(f"Host: {args.host}")
|
| 684 |
+
print(f"Port: {args.port}")
|
| 685 |
+
print(f"Model Directory: {args.model_dir}")
|
| 686 |
+
print(f"Workspace Directory: {args.workspace_dir}")
|
| 687 |
+
print(f"Gallery Directory: {args.gallery_dir}")
|
| 688 |
+
print(f"Share: {args.share}")
|
| 689 |
+
print(f"Debug: {args.debug}")
|
| 690 |
+
print(f"Cache Examples: {args.cache_examples}")
|
| 691 |
+
if args.cache_examples:
|
| 692 |
+
if args.cache_gs_tag:
|
| 693 |
+
print(
|
| 694 |
+
f"Cache GS Tag: '{args.cache_gs_tag}' (scenes matching this tag will use high-res + 3DGS)" # noqa: E501
|
| 695 |
+
) # noqa: E501
|
| 696 |
+
else:
|
| 697 |
+
print("Cache GS Tag: None (all scenes will use low-res only)")
|
| 698 |
+
|
| 699 |
+
# Pre-cache examples if requested
|
| 700 |
+
if args.cache_examples:
|
| 701 |
+
print("\n" + "=" * 60)
|
| 702 |
+
print("Pre-caching mode enabled")
|
| 703 |
+
if args.cache_gs_tag:
|
| 704 |
+
print(f"Scenes containing '{args.cache_gs_tag}' will use HIGH-RES + 3DGS")
|
| 705 |
+
print("Other scenes will use LOW-RES only")
|
| 706 |
+
else:
|
| 707 |
+
print("All scenes will use LOW-RES only")
|
| 708 |
+
print("=" * 60)
|
| 709 |
+
app.cache_examples(
|
| 710 |
+
show_cam=True,
|
| 711 |
+
filter_black_bg=False,
|
| 712 |
+
filter_white_bg=False,
|
| 713 |
+
save_percentage=5.0,
|
| 714 |
+
num_max_points=1000,
|
| 715 |
+
cache_gs_tag=args.cache_gs_tag,
|
| 716 |
+
gs_trj_mode="smooth",
|
| 717 |
+
gs_video_quality="low",
|
| 718 |
+
)
|
| 719 |
+
|
| 720 |
+
app.launch(host=args.host, port=args.port, **launch_kwargs)
|
| 721 |
+
|
| 722 |
+
|
| 723 |
+
if __name__ == "__main__":
|
| 724 |
+
main()
|
depth_anything_3/app/modules/__init__.py
ADDED
|
@@ -0,0 +1,43 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Modules package for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This package contains all the modular components for the Gradio application.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
from depth_anything_3.app.modules.event_handlers import EventHandlers
|
| 22 |
+
from depth_anything_3.app.modules.file_handlers import FileHandler
|
| 23 |
+
from depth_anything_3.app.modules.model_inference import ModelInference
|
| 24 |
+
from depth_anything_3.app.modules.ui_components import UIComponents
|
| 25 |
+
from depth_anything_3.app.modules.utils import (
|
| 26 |
+
create_depth_visualization,
|
| 27 |
+
get_logo_base64,
|
| 28 |
+
get_scene_info,
|
| 29 |
+
save_to_gallery_func,
|
| 30 |
+
)
|
| 31 |
+
from depth_anything_3.app.modules.visualization import VisualizationHandler
|
| 32 |
+
|
| 33 |
+
__all__ = [
|
| 34 |
+
"ModelInference",
|
| 35 |
+
"FileHandler",
|
| 36 |
+
"VisualizationHandler",
|
| 37 |
+
"EventHandlers",
|
| 38 |
+
"UIComponents",
|
| 39 |
+
"create_depth_visualization",
|
| 40 |
+
"save_to_gallery_func",
|
| 41 |
+
"get_scene_info",
|
| 42 |
+
"get_logo_base64",
|
| 43 |
+
]
|
depth_anything_3/app/modules/event_handlers.py
ADDED
|
@@ -0,0 +1,622 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Event handling module for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This module handles all event callbacks and user interactions.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import os
|
| 22 |
+
import time
|
| 23 |
+
from glob import glob
|
| 24 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 25 |
+
import gradio as gr
|
| 26 |
+
import numpy as np
|
| 27 |
+
import torch
|
| 28 |
+
|
| 29 |
+
from depth_anything_3.app.modules.file_handlers import FileHandler
|
| 30 |
+
from depth_anything_3.app.modules.model_inference import ModelInference
|
| 31 |
+
from depth_anything_3.utils.memory import cleanup_cuda_memory
|
| 32 |
+
from depth_anything_3.app.modules.visualization import VisualizationHandler
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class EventHandlers:
|
| 36 |
+
"""
|
| 37 |
+
Handles all event callbacks and user interactions for the Gradio app.
|
| 38 |
+
"""
|
| 39 |
+
|
| 40 |
+
def __init__(self):
|
| 41 |
+
"""Initialize the event handlers."""
|
| 42 |
+
self.model_inference = ModelInference()
|
| 43 |
+
self.file_handler = FileHandler()
|
| 44 |
+
self.visualization_handler = VisualizationHandler()
|
| 45 |
+
|
| 46 |
+
def clear_fields(self) -> None:
|
| 47 |
+
"""
|
| 48 |
+
Clears the 3D viewer, the stored target_dir, and empties the gallery.
|
| 49 |
+
"""
|
| 50 |
+
return None
|
| 51 |
+
|
| 52 |
+
def update_log(self) -> str:
|
| 53 |
+
"""
|
| 54 |
+
Display a quick log message while waiting.
|
| 55 |
+
"""
|
| 56 |
+
return "Loading and Reconstructing..."
|
| 57 |
+
|
| 58 |
+
def save_current_visualization(
|
| 59 |
+
self,
|
| 60 |
+
target_dir: str,
|
| 61 |
+
save_percentage: float,
|
| 62 |
+
show_cam: bool,
|
| 63 |
+
filter_black_bg: bool,
|
| 64 |
+
filter_white_bg: bool,
|
| 65 |
+
processed_data: Optional[Dict],
|
| 66 |
+
scene_name: str = "",
|
| 67 |
+
) -> str:
|
| 68 |
+
"""
|
| 69 |
+
Save current visualization results to gallery with specified save percentage.
|
| 70 |
+
|
| 71 |
+
Args:
|
| 72 |
+
target_dir: Directory containing results
|
| 73 |
+
save_percentage: Percentage of points to save (0-100)
|
| 74 |
+
show_cam: Whether to show cameras
|
| 75 |
+
filter_black_bg: Whether to filter black background
|
| 76 |
+
filter_white_bg: Whether to filter white background
|
| 77 |
+
processed_data: Processed data from reconstruction
|
| 78 |
+
|
| 79 |
+
Returns:
|
| 80 |
+
Status message
|
| 81 |
+
"""
|
| 82 |
+
if not self.file_handler.is_valid_session_dir(target_dir):
|
| 83 |
+
return "No reconstruction available. Please run 'Reconstruct' first."
|
| 84 |
+
|
| 85 |
+
if processed_data is None:
|
| 86 |
+
return "No processed data available. Please run 'Reconstruct' first."
|
| 87 |
+
|
| 88 |
+
try:
|
| 89 |
+
# Add debug information
|
| 90 |
+
print("[DEBUG] save_current_visualization called with:")
|
| 91 |
+
print(f" target_dir: {target_dir}")
|
| 92 |
+
print(f" save_percentage: {save_percentage}")
|
| 93 |
+
print(f" show_cam: {show_cam}")
|
| 94 |
+
print(f" filter_black_bg: {filter_black_bg}")
|
| 95 |
+
print(f" filter_white_bg: {filter_white_bg}")
|
| 96 |
+
print(f" processed_data: {processed_data is not None}")
|
| 97 |
+
|
| 98 |
+
# Import the gallery save function
|
| 99 |
+
# Create gallery name with user input or auto-generated
|
| 100 |
+
import datetime
|
| 101 |
+
|
| 102 |
+
from .utils import save_to_gallery_func
|
| 103 |
+
|
| 104 |
+
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
| 105 |
+
if scene_name and scene_name.strip():
|
| 106 |
+
gallery_name = f"{scene_name.strip()}_{timestamp}_pct{save_percentage:.0f}"
|
| 107 |
+
else:
|
| 108 |
+
gallery_name = f"save_{timestamp}_pct{save_percentage:.0f}"
|
| 109 |
+
|
| 110 |
+
print(f"[DEBUG] Saving to gallery with name: {gallery_name}")
|
| 111 |
+
|
| 112 |
+
# Save entire process folder to gallery
|
| 113 |
+
success, message = save_to_gallery_func(
|
| 114 |
+
target_dir=target_dir, processed_data=processed_data, gallery_name=gallery_name
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
if success:
|
| 118 |
+
print(f"[DEBUG] Gallery save completed successfully: {message}")
|
| 119 |
+
return (
|
| 120 |
+
"Successfully saved to gallery!\n"
|
| 121 |
+
f"Gallery name: {gallery_name}\n"
|
| 122 |
+
f"Save percentage: {save_percentage}%\n"
|
| 123 |
+
f"Show cameras: {show_cam}\n"
|
| 124 |
+
f"Filter black bg: {filter_black_bg}\n"
|
| 125 |
+
f"Filter white bg: {filter_white_bg}\n\n"
|
| 126 |
+
f"{message}"
|
| 127 |
+
)
|
| 128 |
+
else:
|
| 129 |
+
print(f"[DEBUG] Gallery save failed: {message}")
|
| 130 |
+
return f"Failed to save to gallery: {message}"
|
| 131 |
+
|
| 132 |
+
except Exception as e:
|
| 133 |
+
return f"Error saving visualization: {str(e)}"
|
| 134 |
+
|
| 135 |
+
def gradio_demo(
|
| 136 |
+
self,
|
| 137 |
+
target_dir: str,
|
| 138 |
+
show_cam: bool = True,
|
| 139 |
+
filter_black_bg: bool = False,
|
| 140 |
+
filter_white_bg: bool = False,
|
| 141 |
+
process_res_method: str = "upper_bound_resize",
|
| 142 |
+
save_percentage: float = 30.0,
|
| 143 |
+
num_max_points: int = 1_000_000,
|
| 144 |
+
infer_gs: bool = False,
|
| 145 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 146 |
+
gs_trj_mode: str = "extend",
|
| 147 |
+
gs_video_quality: str = "high",
|
| 148 |
+
) -> Tuple[
|
| 149 |
+
Optional[str],
|
| 150 |
+
str,
|
| 151 |
+
Optional[Dict],
|
| 152 |
+
Optional[np.ndarray],
|
| 153 |
+
Optional[np.ndarray],
|
| 154 |
+
str,
|
| 155 |
+
gr.Dropdown,
|
| 156 |
+
Optional[str], # gs video path
|
| 157 |
+
gr.update, # gs video visibility update
|
| 158 |
+
gr.update, # gs info visibility update
|
| 159 |
+
]:
|
| 160 |
+
"""
|
| 161 |
+
Perform reconstruction using the already-created target_dir/images.
|
| 162 |
+
|
| 163 |
+
Args:
|
| 164 |
+
target_dir: Directory containing images
|
| 165 |
+
show_cam: Whether to show camera
|
| 166 |
+
filter_black_bg: Whether to filter black background
|
| 167 |
+
filter_white_bg: Whether to filter white background
|
| 168 |
+
process_res_method: Method for resizing input images
|
| 169 |
+
save_percentage: Filter percentage for point cloud
|
| 170 |
+
num_max_points: Maximum number of points
|
| 171 |
+
infer_gs: Whether to infer 3D Gaussian Splatting
|
| 172 |
+
ref_view_strategy: Reference view selection strategy
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Tuple of reconstruction results
|
| 176 |
+
"""
|
| 177 |
+
if not self.file_handler.is_valid_session_dir(target_dir):
|
| 178 |
+
return (
|
| 179 |
+
None,
|
| 180 |
+
"No valid target directory found. Please upload first.",
|
| 181 |
+
None,
|
| 182 |
+
None,
|
| 183 |
+
None,
|
| 184 |
+
"",
|
| 185 |
+
None,
|
| 186 |
+
None,
|
| 187 |
+
gr.update(visible=False), # gs_video
|
| 188 |
+
gr.update(visible=True), # gs_info
|
| 189 |
+
)
|
| 190 |
+
|
| 191 |
+
start_time = time.time()
|
| 192 |
+
cleanup_cuda_memory()
|
| 193 |
+
|
| 194 |
+
# Get image files for logging
|
| 195 |
+
target_dir_images = os.path.join(target_dir, "images")
|
| 196 |
+
all_files = (
|
| 197 |
+
sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else []
|
| 198 |
+
)
|
| 199 |
+
|
| 200 |
+
print("Running DepthAnything3 model...")
|
| 201 |
+
print(f"Reference view strategy: {ref_view_strategy}")
|
| 202 |
+
|
| 203 |
+
with torch.no_grad():
|
| 204 |
+
prediction, processed_data = self.model_inference.run_inference(
|
| 205 |
+
target_dir,
|
| 206 |
+
process_res_method=process_res_method,
|
| 207 |
+
show_camera=show_cam,
|
| 208 |
+
save_percentage=save_percentage,
|
| 209 |
+
num_max_points=int(num_max_points * 1000), # Convert K to actual count
|
| 210 |
+
infer_gs=infer_gs,
|
| 211 |
+
ref_view_strategy=ref_view_strategy,
|
| 212 |
+
gs_trj_mode=gs_trj_mode,
|
| 213 |
+
gs_video_quality=gs_video_quality,
|
| 214 |
+
)
|
| 215 |
+
|
| 216 |
+
# The GLB file is already generated by the API
|
| 217 |
+
glbfile = os.path.join(target_dir, "scene.glb")
|
| 218 |
+
|
| 219 |
+
# Handle 3DGS video based on infer_gs flag
|
| 220 |
+
gsvideo_path = None
|
| 221 |
+
gs_video_visible = False
|
| 222 |
+
gs_info_visible = True
|
| 223 |
+
|
| 224 |
+
if infer_gs:
|
| 225 |
+
try:
|
| 226 |
+
gsvideo_path = sorted(glob(os.path.join(target_dir, "gs_video", "*.mp4")))[-1]
|
| 227 |
+
gs_video_visible = True
|
| 228 |
+
gs_info_visible = False
|
| 229 |
+
except IndexError:
|
| 230 |
+
gsvideo_path = None
|
| 231 |
+
print("3DGS video not found, but infer_gs was enabled")
|
| 232 |
+
|
| 233 |
+
# Cleanup
|
| 234 |
+
cleanup_cuda_memory()
|
| 235 |
+
|
| 236 |
+
end_time = time.time()
|
| 237 |
+
print(f"Total time: {end_time - start_time:.2f} seconds")
|
| 238 |
+
log_msg = f"Reconstruction Success ({len(all_files)} frames). Waiting for visualization."
|
| 239 |
+
|
| 240 |
+
# Populate visualization tabs with processed data
|
| 241 |
+
depth_vis, measure_img, measure_depth_vis, measure_pts = (
|
| 242 |
+
self.visualization_handler.populate_visualization_tabs(processed_data)
|
| 243 |
+
)
|
| 244 |
+
|
| 245 |
+
# Update view selectors based on available views
|
| 246 |
+
depth_selector, measure_selector = self.visualization_handler.update_view_selectors(
|
| 247 |
+
processed_data
|
| 248 |
+
)
|
| 249 |
+
|
| 250 |
+
return (
|
| 251 |
+
glbfile,
|
| 252 |
+
log_msg,
|
| 253 |
+
processed_data,
|
| 254 |
+
measure_img, # measure_image
|
| 255 |
+
measure_depth_vis, # measure_depth_image
|
| 256 |
+
"", # measure_text (empty initially)
|
| 257 |
+
measure_selector, # measure_view_selector
|
| 258 |
+
gsvideo_path,
|
| 259 |
+
gr.update(visible=gs_video_visible), # gs_video visibility
|
| 260 |
+
gr.update(visible=gs_info_visible), # gs_info visibility
|
| 261 |
+
)
|
| 262 |
+
|
| 263 |
+
def update_visualization(
|
| 264 |
+
self,
|
| 265 |
+
target_dir: str,
|
| 266 |
+
show_cam: bool,
|
| 267 |
+
is_example: str,
|
| 268 |
+
filter_black_bg: bool = False,
|
| 269 |
+
filter_white_bg: bool = False,
|
| 270 |
+
process_res_method: str = "upper_bound_resize",
|
| 271 |
+
) -> Tuple[gr.update, str]:
|
| 272 |
+
"""
|
| 273 |
+
Reload saved predictions from npz, create (or reuse) the GLB for new parameters,
|
| 274 |
+
and return it for the 3D viewer.
|
| 275 |
+
|
| 276 |
+
Args:
|
| 277 |
+
target_dir: Directory containing results
|
| 278 |
+
show_cam: Whether to show camera
|
| 279 |
+
is_example: Whether this is an example scene
|
| 280 |
+
filter_black_bg: Whether to filter black background
|
| 281 |
+
filter_white_bg: Whether to filter white background
|
| 282 |
+
process_res_method: Method for resizing input images
|
| 283 |
+
|
| 284 |
+
Returns:
|
| 285 |
+
Tuple of (glb_file, log_message)
|
| 286 |
+
"""
|
| 287 |
+
if not self.file_handler.is_valid_session_dir(target_dir):
|
| 288 |
+
return (
|
| 289 |
+
gr.update(),
|
| 290 |
+
"No reconstruction available. Please click the Reconstruct button first.",
|
| 291 |
+
)
|
| 292 |
+
|
| 293 |
+
# Check if GLB exists (could be cached example or reconstructed scene)
|
| 294 |
+
glbfile = os.path.join(target_dir, "scene.glb")
|
| 295 |
+
if os.path.exists(glbfile):
|
| 296 |
+
return (
|
| 297 |
+
glbfile,
|
| 298 |
+
(
|
| 299 |
+
"Visualization loaded from cache."
|
| 300 |
+
if is_example == "True"
|
| 301 |
+
else "Visualization updated."
|
| 302 |
+
),
|
| 303 |
+
)
|
| 304 |
+
|
| 305 |
+
# If no GLB but it's an example that hasn't been reconstructed yet
|
| 306 |
+
if is_example == "True":
|
| 307 |
+
return (
|
| 308 |
+
gr.update(),
|
| 309 |
+
"No reconstruction available. Please click the Reconstruct button first.",
|
| 310 |
+
)
|
| 311 |
+
|
| 312 |
+
# For non-examples, check predictions.npz
|
| 313 |
+
predictions_path = os.path.join(target_dir, "predictions.npz")
|
| 314 |
+
if not os.path.exists(predictions_path):
|
| 315 |
+
error_message = (
|
| 316 |
+
f"No reconstruction available at {predictions_path}. "
|
| 317 |
+
"Please run 'Reconstruct' first."
|
| 318 |
+
)
|
| 319 |
+
return gr.update(), error_message
|
| 320 |
+
|
| 321 |
+
try:
|
| 322 |
+
loaded = np.load(predictions_path, allow_pickle=False)
|
| 323 |
+
predictions = {key: loaded[key] for key in loaded.keys()} # noqa: F841
|
| 324 |
+
except Exception as e:
|
| 325 |
+
return gr.update(), f"Cached results could not be read: {e}"
|
| 326 |
+
|
| 327 |
+
return (
|
| 328 |
+
glbfile,
|
| 329 |
+
"Visualization updated.",
|
| 330 |
+
)
|
| 331 |
+
|
| 332 |
+
def handle_uploads(
|
| 333 |
+
self,
|
| 334 |
+
input_video: Optional[str],
|
| 335 |
+
input_images: Optional[List],
|
| 336 |
+
s_time_interval: float = 10.0,
|
| 337 |
+
) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[str]]:
|
| 338 |
+
"""
|
| 339 |
+
Handle file uploads and update gallery.
|
| 340 |
+
|
| 341 |
+
Args:
|
| 342 |
+
input_video: Path to input video file
|
| 343 |
+
input_images: List of input image files
|
| 344 |
+
s_time_interval: Sampling FPS (frames per second) for frame extraction
|
| 345 |
+
|
| 346 |
+
Returns:
|
| 347 |
+
Tuple of (reconstruction_output, target_dir, image_paths, log_message)
|
| 348 |
+
"""
|
| 349 |
+
return self.file_handler.update_gallery_on_upload(
|
| 350 |
+
input_video, input_images, s_time_interval
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
def load_example_scene(self, scene_name: str, examples_dir: str = None) -> Tuple[
|
| 354 |
+
Optional[str],
|
| 355 |
+
Optional[str],
|
| 356 |
+
Optional[List],
|
| 357 |
+
str,
|
| 358 |
+
Optional[Dict],
|
| 359 |
+
gr.Dropdown,
|
| 360 |
+
Optional[str],
|
| 361 |
+
gr.update,
|
| 362 |
+
gr.update,
|
| 363 |
+
]:
|
| 364 |
+
"""
|
| 365 |
+
Load a scene from examples directory.
|
| 366 |
+
|
| 367 |
+
Args:
|
| 368 |
+
scene_name: Name of the scene to load
|
| 369 |
+
examples_dir: Path to examples directory (if None, uses workspace_dir/examples)
|
| 370 |
+
|
| 371 |
+
Returns:
|
| 372 |
+
Tuple of (reconstruction_output, target_dir, image_paths, log_message, processed_data, measure_view_selector, gs_video, gs_video_vis, gs_info_vis) # noqa: E501
|
| 373 |
+
"""
|
| 374 |
+
if examples_dir is None:
|
| 375 |
+
# Get workspace directory from environment variable
|
| 376 |
+
workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace")
|
| 377 |
+
examples_dir = os.path.join(workspace_dir, "examples")
|
| 378 |
+
|
| 379 |
+
reconstruction_output, target_dir, image_paths, log_message = (
|
| 380 |
+
self.file_handler.load_example_scene(scene_name, examples_dir)
|
| 381 |
+
)
|
| 382 |
+
|
| 383 |
+
# Try to load cached processed data if available
|
| 384 |
+
processed_data = None
|
| 385 |
+
measure_view_selector = gr.Dropdown(choices=["View 1"], value="View 1")
|
| 386 |
+
gs_video_path = None
|
| 387 |
+
gs_video_visible = False
|
| 388 |
+
gs_info_visible = True
|
| 389 |
+
|
| 390 |
+
if self.file_handler.is_valid_session_dir(target_dir):
|
| 391 |
+
predictions_path = os.path.join(target_dir, "predictions.npz")
|
| 392 |
+
if os.path.exists(predictions_path):
|
| 393 |
+
try:
|
| 394 |
+
# Load predictions from cache
|
| 395 |
+
loaded = np.load(predictions_path, allow_pickle=False)
|
| 396 |
+
predictions = {key: loaded[key] for key in loaded.keys()}
|
| 397 |
+
|
| 398 |
+
# Reconstruct processed_data structure
|
| 399 |
+
num_images = len(predictions.get("images", []))
|
| 400 |
+
processed_data = {}
|
| 401 |
+
|
| 402 |
+
for i in range(num_images):
|
| 403 |
+
processed_data[i] = {
|
| 404 |
+
"image": predictions["images"][i] if "images" in predictions else None,
|
| 405 |
+
"depth": predictions["depths"][i] if "depths" in predictions else None,
|
| 406 |
+
"depth_image": os.path.join(
|
| 407 |
+
target_dir, "depth_vis", f"{i:04d}.jpg" # Fixed: use .jpg not .png
|
| 408 |
+
),
|
| 409 |
+
"intrinsics": (
|
| 410 |
+
predictions["intrinsics"][i]
|
| 411 |
+
if "intrinsics" in predictions
|
| 412 |
+
and i < len(predictions["intrinsics"])
|
| 413 |
+
else None
|
| 414 |
+
),
|
| 415 |
+
"mask": None,
|
| 416 |
+
}
|
| 417 |
+
|
| 418 |
+
# Update measure view selector
|
| 419 |
+
choices = [f"View {i + 1}" for i in range(num_images)]
|
| 420 |
+
measure_view_selector = gr.Dropdown(choices=choices, value=choices[0])
|
| 421 |
+
|
| 422 |
+
except Exception as e:
|
| 423 |
+
print(f"Error loading cached data: {e}")
|
| 424 |
+
|
| 425 |
+
# Check for cached 3DGS video
|
| 426 |
+
gs_video_dir = os.path.join(target_dir, "gs_video")
|
| 427 |
+
if os.path.exists(gs_video_dir):
|
| 428 |
+
try:
|
| 429 |
+
from glob import glob
|
| 430 |
+
|
| 431 |
+
gs_videos = sorted(glob(os.path.join(gs_video_dir, "*.mp4")))
|
| 432 |
+
if gs_videos:
|
| 433 |
+
gs_video_path = gs_videos[-1]
|
| 434 |
+
gs_video_visible = True
|
| 435 |
+
gs_info_visible = False
|
| 436 |
+
print(f"Loaded cached 3DGS video: {gs_video_path}")
|
| 437 |
+
except Exception as e:
|
| 438 |
+
print(f"Error loading cached 3DGS video: {e}")
|
| 439 |
+
|
| 440 |
+
return (
|
| 441 |
+
reconstruction_output,
|
| 442 |
+
target_dir,
|
| 443 |
+
image_paths,
|
| 444 |
+
log_message,
|
| 445 |
+
processed_data,
|
| 446 |
+
measure_view_selector,
|
| 447 |
+
gs_video_path,
|
| 448 |
+
gr.update(visible=gs_video_visible),
|
| 449 |
+
gr.update(visible=gs_info_visible),
|
| 450 |
+
)
|
| 451 |
+
|
| 452 |
+
def navigate_depth_view(
|
| 453 |
+
self,
|
| 454 |
+
processed_data: Optional[Dict[int, Dict[str, Any]]],
|
| 455 |
+
current_selector: str,
|
| 456 |
+
direction: int,
|
| 457 |
+
) -> Tuple[str, Optional[str]]:
|
| 458 |
+
"""
|
| 459 |
+
Navigate depth view.
|
| 460 |
+
|
| 461 |
+
Args:
|
| 462 |
+
processed_data: Processed data dictionary
|
| 463 |
+
current_selector: Current selector value
|
| 464 |
+
direction: Direction to navigate
|
| 465 |
+
|
| 466 |
+
Returns:
|
| 467 |
+
Tuple of (new_selector_value, depth_vis)
|
| 468 |
+
"""
|
| 469 |
+
return self.visualization_handler.navigate_depth_view(
|
| 470 |
+
processed_data, current_selector, direction
|
| 471 |
+
)
|
| 472 |
+
|
| 473 |
+
def update_depth_view(
|
| 474 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int
|
| 475 |
+
) -> Optional[str]:
|
| 476 |
+
"""
|
| 477 |
+
Update depth view for a specific view index.
|
| 478 |
+
|
| 479 |
+
Args:
|
| 480 |
+
processed_data: Processed data dictionary
|
| 481 |
+
view_index: Index of the view to update
|
| 482 |
+
|
| 483 |
+
Returns:
|
| 484 |
+
Path to depth visualization image or None
|
| 485 |
+
"""
|
| 486 |
+
return self.visualization_handler.update_depth_view(processed_data, view_index)
|
| 487 |
+
|
| 488 |
+
def navigate_measure_view(
|
| 489 |
+
self,
|
| 490 |
+
processed_data: Optional[Dict[int, Dict[str, Any]]],
|
| 491 |
+
current_selector: str,
|
| 492 |
+
direction: int,
|
| 493 |
+
) -> Tuple[str, Optional[np.ndarray], Optional[np.ndarray], List]:
|
| 494 |
+
"""
|
| 495 |
+
Navigate measure view.
|
| 496 |
+
|
| 497 |
+
Args:
|
| 498 |
+
processed_data: Processed data dictionary
|
| 499 |
+
current_selector: Current selector value
|
| 500 |
+
direction: Direction to navigate
|
| 501 |
+
|
| 502 |
+
Returns:
|
| 503 |
+
Tuple of (new_selector_value, measure_image, depth_right_half, measure_points)
|
| 504 |
+
"""
|
| 505 |
+
return self.visualization_handler.navigate_measure_view(
|
| 506 |
+
processed_data, current_selector, direction
|
| 507 |
+
)
|
| 508 |
+
|
| 509 |
+
def update_measure_view(
|
| 510 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int
|
| 511 |
+
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List]:
|
| 512 |
+
"""
|
| 513 |
+
Update measure view for a specific view index.
|
| 514 |
+
|
| 515 |
+
Args:
|
| 516 |
+
processed_data: Processed data dictionary
|
| 517 |
+
view_index: Index of the view to update
|
| 518 |
+
|
| 519 |
+
Returns:
|
| 520 |
+
Tuple of (measure_image, depth_right_half, measure_points)
|
| 521 |
+
"""
|
| 522 |
+
return self.visualization_handler.update_measure_view(processed_data, view_index)
|
| 523 |
+
|
| 524 |
+
def measure(
|
| 525 |
+
self,
|
| 526 |
+
processed_data: Optional[Dict[int, Dict[str, Any]]],
|
| 527 |
+
measure_points: List,
|
| 528 |
+
current_view_selector: str,
|
| 529 |
+
event: gr.SelectData,
|
| 530 |
+
) -> List:
|
| 531 |
+
"""
|
| 532 |
+
Handle measurement on images.
|
| 533 |
+
|
| 534 |
+
Args:
|
| 535 |
+
processed_data: Processed data dictionary
|
| 536 |
+
measure_points: List of current measure points
|
| 537 |
+
current_view_selector: Current view selector value
|
| 538 |
+
event: Gradio select event
|
| 539 |
+
|
| 540 |
+
Returns:
|
| 541 |
+
List of [image, depth_right_half, measure_points, text]
|
| 542 |
+
"""
|
| 543 |
+
return self.visualization_handler.measure(
|
| 544 |
+
processed_data, measure_points, current_view_selector, event
|
| 545 |
+
)
|
| 546 |
+
|
| 547 |
+
def select_first_frame(
|
| 548 |
+
self, image_gallery: List, selected_index: int = 0
|
| 549 |
+
) -> Tuple[List, str, str]:
|
| 550 |
+
"""
|
| 551 |
+
Select the first frame from the image gallery.
|
| 552 |
+
|
| 553 |
+
Args:
|
| 554 |
+
image_gallery: List of images in the gallery
|
| 555 |
+
selected_index: Index of the selected image (default: 0)
|
| 556 |
+
|
| 557 |
+
Returns:
|
| 558 |
+
Tuple of (updated_image_gallery, log_message, selected_frame_path)
|
| 559 |
+
"""
|
| 560 |
+
try:
|
| 561 |
+
if not image_gallery or len(image_gallery) == 0:
|
| 562 |
+
return image_gallery, "No images available to select as first frame.", ""
|
| 563 |
+
|
| 564 |
+
# Handle None or invalid selected_index
|
| 565 |
+
if (
|
| 566 |
+
selected_index is None
|
| 567 |
+
or selected_index < 0
|
| 568 |
+
or selected_index >= len(image_gallery)
|
| 569 |
+
):
|
| 570 |
+
selected_index = 0
|
| 571 |
+
print(f"Invalid selected_index: {selected_index}, using default: 0")
|
| 572 |
+
|
| 573 |
+
# Get the selected image based on index
|
| 574 |
+
selected_image = image_gallery[selected_index]
|
| 575 |
+
print(f"Selected image index: {selected_index}")
|
| 576 |
+
print(f"Total images: {len(image_gallery)}")
|
| 577 |
+
|
| 578 |
+
# Extract the file path from the selected image
|
| 579 |
+
selected_frame_path = ""
|
| 580 |
+
print(f"Selected image type: {type(selected_image)}")
|
| 581 |
+
print(f"Selected image: {selected_image}")
|
| 582 |
+
|
| 583 |
+
if isinstance(selected_image, tuple):
|
| 584 |
+
# Gradio Gallery returns tuple (path, None)
|
| 585 |
+
selected_frame_path = selected_image[0]
|
| 586 |
+
elif isinstance(selected_image, str):
|
| 587 |
+
selected_frame_path = selected_image
|
| 588 |
+
elif hasattr(selected_image, "name"):
|
| 589 |
+
selected_frame_path = selected_image.name
|
| 590 |
+
elif isinstance(selected_image, dict):
|
| 591 |
+
if "name" in selected_image:
|
| 592 |
+
selected_frame_path = selected_image["name"]
|
| 593 |
+
elif "path" in selected_image:
|
| 594 |
+
selected_frame_path = selected_image["path"]
|
| 595 |
+
elif "src" in selected_image:
|
| 596 |
+
selected_frame_path = selected_image["src"]
|
| 597 |
+
else:
|
| 598 |
+
# Try to convert to string
|
| 599 |
+
selected_frame_path = str(selected_image)
|
| 600 |
+
|
| 601 |
+
print(f"Extracted path: {selected_frame_path}")
|
| 602 |
+
|
| 603 |
+
# Extract filename from the path for matching
|
| 604 |
+
import os
|
| 605 |
+
|
| 606 |
+
selected_filename = os.path.basename(selected_frame_path)
|
| 607 |
+
print(f"Selected filename: {selected_filename}")
|
| 608 |
+
|
| 609 |
+
# Move the selected image to the front
|
| 610 |
+
updated_gallery = [selected_image] + [
|
| 611 |
+
img for img in image_gallery if img != selected_image
|
| 612 |
+
]
|
| 613 |
+
|
| 614 |
+
log_message = (
|
| 615 |
+
f"Selected frame: {selected_filename}. "
|
| 616 |
+
f"Moved to first position. Total frames: {len(updated_gallery)}"
|
| 617 |
+
)
|
| 618 |
+
return updated_gallery, log_message, selected_filename
|
| 619 |
+
|
| 620 |
+
except Exception as e:
|
| 621 |
+
print(f"Error selecting first frame: {e}")
|
| 622 |
+
return image_gallery, f"Error selecting first frame: {e}", ""
|
depth_anything_3/app/modules/file_handlers.py
ADDED
|
@@ -0,0 +1,355 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
File handling module for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This module handles file uploads, video processing, and file operations.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import os
|
| 22 |
+
import shutil
|
| 23 |
+
import time
|
| 24 |
+
from datetime import datetime
|
| 25 |
+
from typing import List, Optional, Tuple
|
| 26 |
+
import cv2
|
| 27 |
+
from PIL import Image
|
| 28 |
+
from pillow_heif import register_heif_opener
|
| 29 |
+
|
| 30 |
+
register_heif_opener()
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
ALLOWED_IMAGE_EXTENSIONS = {
|
| 34 |
+
".png",
|
| 35 |
+
".jpg",
|
| 36 |
+
".jpeg",
|
| 37 |
+
".bmp",
|
| 38 |
+
".tiff",
|
| 39 |
+
".tif",
|
| 40 |
+
".webp",
|
| 41 |
+
".heic",
|
| 42 |
+
".heif",
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def _looks_like_image(file_path: str) -> bool:
|
| 47 |
+
"""Best-effort check that a file is actually a decodable image, not just named like one."""
|
| 48 |
+
try:
|
| 49 |
+
with Image.open(file_path) as img:
|
| 50 |
+
img.verify()
|
| 51 |
+
return True
|
| 52 |
+
except Exception:
|
| 53 |
+
return False
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
class FileHandler:
|
| 57 |
+
"""
|
| 58 |
+
Handles file uploads and processing for the Gradio app.
|
| 59 |
+
"""
|
| 60 |
+
|
| 61 |
+
def __init__(self):
|
| 62 |
+
"""Initialize the file handler."""
|
| 63 |
+
|
| 64 |
+
def is_valid_session_dir(self, target_dir: Optional[str]) -> bool:
|
| 65 |
+
"""Check that target_dir is a session/example directory this handler itself
|
| 66 |
+
creates under the configured workspace, rather than an arbitrary path."""
|
| 67 |
+
if not target_dir or target_dir == "None":
|
| 68 |
+
return False
|
| 69 |
+
|
| 70 |
+
try:
|
| 71 |
+
workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace")
|
| 72 |
+
input_images_dir = os.path.realpath(os.path.join(workspace_dir, "input_images"))
|
| 73 |
+
real_target = os.path.realpath(target_dir)
|
| 74 |
+
|
| 75 |
+
if real_target != input_images_dir and not real_target.startswith(
|
| 76 |
+
input_images_dir + os.sep
|
| 77 |
+
):
|
| 78 |
+
return False
|
| 79 |
+
|
| 80 |
+
rel = os.path.relpath(real_target, input_images_dir)
|
| 81 |
+
if os.sep in rel or not (rel.startswith("session_") or rel.startswith("example_")):
|
| 82 |
+
return False
|
| 83 |
+
|
| 84 |
+
return os.path.isdir(real_target)
|
| 85 |
+
except (OSError, ValueError):
|
| 86 |
+
return False
|
| 87 |
+
|
| 88 |
+
def handle_uploads(
|
| 89 |
+
self,
|
| 90 |
+
input_video: Optional[str],
|
| 91 |
+
input_images: Optional[List],
|
| 92 |
+
s_time_interval: float = 10.0,
|
| 93 |
+
) -> Tuple[str, List[str]]:
|
| 94 |
+
"""
|
| 95 |
+
Create a new 'target_dir' + 'images' subfolder, and place user-uploaded
|
| 96 |
+
images or extracted frames from video into it.
|
| 97 |
+
|
| 98 |
+
Args:
|
| 99 |
+
input_video: Path to input video file
|
| 100 |
+
input_images: List of input image files
|
| 101 |
+
s_time_interval: Sampling FPS (frames per second) for frame extraction
|
| 102 |
+
|
| 103 |
+
Returns:
|
| 104 |
+
Tuple of (target_dir, image_paths)
|
| 105 |
+
"""
|
| 106 |
+
start_time = time.time()
|
| 107 |
+
|
| 108 |
+
# Get workspace directory from environment variable or use default
|
| 109 |
+
workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace")
|
| 110 |
+
if not os.path.exists(workspace_dir):
|
| 111 |
+
os.makedirs(workspace_dir)
|
| 112 |
+
|
| 113 |
+
# Create input_images subdirectory
|
| 114 |
+
input_images_dir = os.path.join(workspace_dir, "input_images")
|
| 115 |
+
if not os.path.exists(input_images_dir):
|
| 116 |
+
os.makedirs(input_images_dir)
|
| 117 |
+
|
| 118 |
+
# Create a unique folder name within input_images
|
| 119 |
+
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
| 120 |
+
target_dir = os.path.join(input_images_dir, f"session_{timestamp}")
|
| 121 |
+
target_dir_images = os.path.join(target_dir, "images")
|
| 122 |
+
|
| 123 |
+
# Clean up if somehow that folder already exists
|
| 124 |
+
if os.path.exists(target_dir):
|
| 125 |
+
shutil.rmtree(target_dir)
|
| 126 |
+
os.makedirs(target_dir)
|
| 127 |
+
os.makedirs(target_dir_images)
|
| 128 |
+
|
| 129 |
+
image_paths = []
|
| 130 |
+
|
| 131 |
+
# Handle images
|
| 132 |
+
if input_images is not None:
|
| 133 |
+
image_paths.extend(self._process_images(input_images, target_dir_images))
|
| 134 |
+
|
| 135 |
+
# Handle video
|
| 136 |
+
if input_video is not None:
|
| 137 |
+
image_paths.extend(
|
| 138 |
+
self._process_video(input_video, target_dir_images, s_time_interval)
|
| 139 |
+
)
|
| 140 |
+
|
| 141 |
+
# Sort final images for gallery
|
| 142 |
+
image_paths = sorted(image_paths)
|
| 143 |
+
|
| 144 |
+
end_time = time.time()
|
| 145 |
+
print(f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds")
|
| 146 |
+
return target_dir, image_paths
|
| 147 |
+
|
| 148 |
+
def _process_images(self, input_images: List, target_dir_images: str) -> List[str]:
|
| 149 |
+
"""
|
| 150 |
+
Process uploaded images.
|
| 151 |
+
|
| 152 |
+
Args:
|
| 153 |
+
input_images: List of input image files
|
| 154 |
+
target_dir_images: Target directory for images
|
| 155 |
+
|
| 156 |
+
Returns:
|
| 157 |
+
List of processed image paths
|
| 158 |
+
"""
|
| 159 |
+
image_paths = []
|
| 160 |
+
|
| 161 |
+
for file_data in input_images:
|
| 162 |
+
if isinstance(file_data, dict) and "name" in file_data:
|
| 163 |
+
file_path = file_data["name"]
|
| 164 |
+
else:
|
| 165 |
+
file_path = file_data
|
| 166 |
+
|
| 167 |
+
file_ext = os.path.splitext(file_path)[1].lower()
|
| 168 |
+
if file_ext not in ALLOWED_IMAGE_EXTENSIONS or not _looks_like_image(file_path):
|
| 169 |
+
print(f"Skipping non-image upload: {os.path.basename(file_path)}")
|
| 170 |
+
continue
|
| 171 |
+
|
| 172 |
+
# Check if the file is a HEIC image
|
| 173 |
+
if file_ext in [".heic", ".heif"]:
|
| 174 |
+
# Convert HEIC to JPEG for better gallery compatibility
|
| 175 |
+
try:
|
| 176 |
+
with Image.open(file_path) as img:
|
| 177 |
+
# Convert to RGB if necessary (HEIC can have different color modes)
|
| 178 |
+
if img.mode not in ("RGB", "L"):
|
| 179 |
+
img = img.convert("RGB")
|
| 180 |
+
|
| 181 |
+
# Create JPEG filename
|
| 182 |
+
base_name = os.path.splitext(os.path.basename(file_path))[0]
|
| 183 |
+
dst_path = os.path.join(target_dir_images, f"{base_name}.jpg")
|
| 184 |
+
|
| 185 |
+
# Save as JPEG with high quality
|
| 186 |
+
img.save(dst_path, "JPEG", quality=95)
|
| 187 |
+
image_paths.append(dst_path)
|
| 188 |
+
print(
|
| 189 |
+
f"Converted HEIC to JPEG: {os.path.basename(file_path)} -> "
|
| 190 |
+
f"{os.path.basename(dst_path)}"
|
| 191 |
+
)
|
| 192 |
+
except Exception as e:
|
| 193 |
+
print(f"Error converting HEIC file {file_path}: {e}")
|
| 194 |
+
# Fall back to copying as is
|
| 195 |
+
dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
|
| 196 |
+
shutil.copy(file_path, dst_path)
|
| 197 |
+
image_paths.append(dst_path)
|
| 198 |
+
else:
|
| 199 |
+
# Regular image files - copy as is
|
| 200 |
+
dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
|
| 201 |
+
shutil.copy(file_path, dst_path)
|
| 202 |
+
image_paths.append(dst_path)
|
| 203 |
+
|
| 204 |
+
return image_paths
|
| 205 |
+
|
| 206 |
+
def _process_video(
|
| 207 |
+
self, input_video: str, target_dir_images: str, s_time_interval: float
|
| 208 |
+
) -> List[str]:
|
| 209 |
+
"""
|
| 210 |
+
Process video file and extract frames.
|
| 211 |
+
|
| 212 |
+
Args:
|
| 213 |
+
input_video: Path to input video file
|
| 214 |
+
target_dir_images: Target directory for extracted frames
|
| 215 |
+
s_time_interval: Sampling FPS (frames per second) for frame extraction
|
| 216 |
+
|
| 217 |
+
Returns:
|
| 218 |
+
List of extracted frame paths
|
| 219 |
+
"""
|
| 220 |
+
image_paths = []
|
| 221 |
+
|
| 222 |
+
if isinstance(input_video, dict) and "name" in input_video:
|
| 223 |
+
video_path = input_video["name"]
|
| 224 |
+
else:
|
| 225 |
+
video_path = input_video
|
| 226 |
+
|
| 227 |
+
vs = cv2.VideoCapture(video_path)
|
| 228 |
+
fps = vs.get(cv2.CAP_PROP_FPS)
|
| 229 |
+
frame_interval = max(1, int(fps / s_time_interval)) # Convert FPS to frame interval
|
| 230 |
+
|
| 231 |
+
count = 0
|
| 232 |
+
video_frame_num = 0
|
| 233 |
+
while True:
|
| 234 |
+
gotit, frame = vs.read()
|
| 235 |
+
if not gotit:
|
| 236 |
+
break
|
| 237 |
+
count += 1
|
| 238 |
+
if count % frame_interval == 0:
|
| 239 |
+
image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png")
|
| 240 |
+
cv2.imwrite(image_path, frame)
|
| 241 |
+
image_paths.append(image_path)
|
| 242 |
+
video_frame_num += 1
|
| 243 |
+
|
| 244 |
+
return image_paths
|
| 245 |
+
|
| 246 |
+
def update_gallery_on_upload(
|
| 247 |
+
self,
|
| 248 |
+
input_video: Optional[str],
|
| 249 |
+
input_images: Optional[List],
|
| 250 |
+
s_time_interval: float = 10.0,
|
| 251 |
+
) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[str]]:
|
| 252 |
+
"""
|
| 253 |
+
Handle file uploads and update gallery.
|
| 254 |
+
|
| 255 |
+
Args:
|
| 256 |
+
input_video: Path to input video file
|
| 257 |
+
input_images: List of input image files
|
| 258 |
+
s_time_interval: Sampling FPS (frames per second) for frame extraction
|
| 259 |
+
|
| 260 |
+
Returns:
|
| 261 |
+
Tuple of (reconstruction_output, target_dir, image_paths, log_message)
|
| 262 |
+
"""
|
| 263 |
+
if not input_video and not input_images:
|
| 264 |
+
return None, None, None, None
|
| 265 |
+
|
| 266 |
+
target_dir, image_paths = self.handle_uploads(input_video, input_images, s_time_interval)
|
| 267 |
+
return (
|
| 268 |
+
None,
|
| 269 |
+
target_dir,
|
| 270 |
+
image_paths,
|
| 271 |
+
"Upload complete. Click 'Reconstruct' to begin 3D processing.",
|
| 272 |
+
)
|
| 273 |
+
|
| 274 |
+
def load_example_scene(
|
| 275 |
+
self, scene_name: str, examples_dir: str = "examples"
|
| 276 |
+
) -> Tuple[Optional[str], Optional[str], Optional[List], str]:
|
| 277 |
+
"""
|
| 278 |
+
Load a scene from examples directory.
|
| 279 |
+
|
| 280 |
+
Args:
|
| 281 |
+
scene_name: Name of the scene to load
|
| 282 |
+
examples_dir: Path to examples directory
|
| 283 |
+
|
| 284 |
+
Returns:
|
| 285 |
+
Tuple of (reconstruction_output, target_dir, image_paths, log_message)
|
| 286 |
+
"""
|
| 287 |
+
from depth_anything_3.app.modules.utils import get_scene_info
|
| 288 |
+
|
| 289 |
+
scenes = get_scene_info(examples_dir)
|
| 290 |
+
|
| 291 |
+
# Find the selected scene
|
| 292 |
+
selected_scene = None
|
| 293 |
+
for scene in scenes:
|
| 294 |
+
if scene["name"] == scene_name:
|
| 295 |
+
selected_scene = scene
|
| 296 |
+
break
|
| 297 |
+
|
| 298 |
+
if selected_scene is None:
|
| 299 |
+
return None, None, None, "Scene not found"
|
| 300 |
+
|
| 301 |
+
# Use fixed directory name for examples (not timestamp-based)
|
| 302 |
+
workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace")
|
| 303 |
+
input_images_dir = os.path.join(workspace_dir, "input_images")
|
| 304 |
+
if not os.path.exists(input_images_dir):
|
| 305 |
+
os.makedirs(input_images_dir)
|
| 306 |
+
|
| 307 |
+
# Create a fixed folder name based on scene name
|
| 308 |
+
target_dir = os.path.join(input_images_dir, f"example_{scene_name}")
|
| 309 |
+
target_dir_images = os.path.join(target_dir, "images")
|
| 310 |
+
|
| 311 |
+
# Check if already cached (GLB file exists)
|
| 312 |
+
glb_path = os.path.join(target_dir, "scene.glb")
|
| 313 |
+
is_cached = os.path.exists(glb_path)
|
| 314 |
+
|
| 315 |
+
# Create directory if it doesn't exist
|
| 316 |
+
if not os.path.exists(target_dir):
|
| 317 |
+
os.makedirs(target_dir)
|
| 318 |
+
os.makedirs(target_dir_images)
|
| 319 |
+
|
| 320 |
+
# Copy images if directory is new or empty
|
| 321 |
+
if not os.path.exists(target_dir_images) or len(os.listdir(target_dir_images)) == 0:
|
| 322 |
+
os.makedirs(target_dir_images, exist_ok=True)
|
| 323 |
+
image_paths = []
|
| 324 |
+
for file_path in selected_scene["image_files"]:
|
| 325 |
+
dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
|
| 326 |
+
shutil.copy(file_path, dst_path)
|
| 327 |
+
image_paths.append(dst_path)
|
| 328 |
+
else:
|
| 329 |
+
# Use existing images
|
| 330 |
+
image_paths = sorted(
|
| 331 |
+
[
|
| 332 |
+
os.path.join(target_dir_images, f)
|
| 333 |
+
for f in os.listdir(target_dir_images)
|
| 334 |
+
if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif"))
|
| 335 |
+
]
|
| 336 |
+
)
|
| 337 |
+
|
| 338 |
+
# Return cached GLB if available
|
| 339 |
+
if is_cached:
|
| 340 |
+
return (
|
| 341 |
+
glb_path, # Return cached reconstruction
|
| 342 |
+
target_dir, # Set target directory
|
| 343 |
+
image_paths, # Set gallery
|
| 344 |
+
f"Loaded cached scene '{scene_name}' with {selected_scene['num_images']} images.",
|
| 345 |
+
)
|
| 346 |
+
else:
|
| 347 |
+
return (
|
| 348 |
+
None, # No cached reconstruction
|
| 349 |
+
target_dir, # Set target directory
|
| 350 |
+
image_paths, # Set gallery
|
| 351 |
+
(
|
| 352 |
+
f"Loaded scene '{scene_name}' with {selected_scene['num_images']} images. "
|
| 353 |
+
"Click 'Reconstruct' to begin 3D processing."
|
| 354 |
+
),
|
| 355 |
+
)
|
depth_anything_3/app/modules/model_inference.py
ADDED
|
@@ -0,0 +1,260 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Model inference module for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This module handles all model-related operations including inference,
|
| 19 |
+
data processing, and result preparation.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import glob
|
| 23 |
+
import os
|
| 24 |
+
from typing import Any, Dict, Optional, Tuple
|
| 25 |
+
import numpy as np
|
| 26 |
+
import torch
|
| 27 |
+
|
| 28 |
+
from depth_anything_3.api import DepthAnything3
|
| 29 |
+
from depth_anything_3.utils.memory import cleanup_cuda_memory
|
| 30 |
+
from depth_anything_3.utils.export.glb import export_to_glb
|
| 31 |
+
from depth_anything_3.utils.export.gs import export_to_gs_video
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class ModelInference:
|
| 35 |
+
"""
|
| 36 |
+
Handles model inference and data processing for Depth Anything 3.
|
| 37 |
+
"""
|
| 38 |
+
|
| 39 |
+
def __init__(self):
|
| 40 |
+
"""Initialize the model inference handler."""
|
| 41 |
+
self.model = None
|
| 42 |
+
|
| 43 |
+
def initialize_model(self, device: str = "cuda") -> None:
|
| 44 |
+
"""
|
| 45 |
+
Initialize the DepthAnything3 model.
|
| 46 |
+
|
| 47 |
+
Args:
|
| 48 |
+
device: Device to load the model on
|
| 49 |
+
"""
|
| 50 |
+
if self.model is None:
|
| 51 |
+
# Get model directory from environment variable or use default
|
| 52 |
+
model_dir = os.environ.get(
|
| 53 |
+
"DA3_MODEL_DIR", "/dev/shm/da3_models/DA3HF-VITG-METRIC_VITL"
|
| 54 |
+
)
|
| 55 |
+
self.model = DepthAnything3.from_pretrained(model_dir)
|
| 56 |
+
self.model = self.model.to(device)
|
| 57 |
+
else:
|
| 58 |
+
self.model = self.model.to(device)
|
| 59 |
+
|
| 60 |
+
self.model.eval()
|
| 61 |
+
|
| 62 |
+
def run_inference(
|
| 63 |
+
self,
|
| 64 |
+
target_dir: str,
|
| 65 |
+
filter_black_bg: bool = False,
|
| 66 |
+
filter_white_bg: bool = False,
|
| 67 |
+
process_res_method: str = "upper_bound_resize",
|
| 68 |
+
show_camera: bool = True,
|
| 69 |
+
save_percentage: float = 30.0,
|
| 70 |
+
num_max_points: int = 1_000_000,
|
| 71 |
+
infer_gs: bool = False,
|
| 72 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 73 |
+
gs_trj_mode: str = "extend",
|
| 74 |
+
gs_video_quality: str = "high",
|
| 75 |
+
) -> Tuple[Any, Dict[int, Dict[str, Any]]]:
|
| 76 |
+
"""
|
| 77 |
+
Run DepthAnything3 model inference on images.
|
| 78 |
+
|
| 79 |
+
Args:
|
| 80 |
+
target_dir: Directory containing images
|
| 81 |
+
filter_black_bg: Whether to filter black background
|
| 82 |
+
filter_white_bg: Whether to filter white background
|
| 83 |
+
process_res_method: Method for resizing input images
|
| 84 |
+
show_camera: Whether to show camera in 3D view
|
| 85 |
+
save_percentage: Percentage of points to save (0-100)
|
| 86 |
+
num_max_points: Maximum number of points in point cloud
|
| 87 |
+
infer_gs: Whether to infer 3D Gaussian Splatting
|
| 88 |
+
ref_view_strategy: Reference view selection strategy
|
| 89 |
+
gs_trj_mode: Trajectory mode for 3DGS
|
| 90 |
+
gs_video_quality: Video quality for 3DGS
|
| 91 |
+
|
| 92 |
+
Returns:
|
| 93 |
+
Tuple of (prediction, processed_data)
|
| 94 |
+
"""
|
| 95 |
+
print(f"Processing images from {target_dir}")
|
| 96 |
+
|
| 97 |
+
# Device check
|
| 98 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 99 |
+
device = torch.device(device)
|
| 100 |
+
|
| 101 |
+
# Initialize model if needed
|
| 102 |
+
self.initialize_model(device)
|
| 103 |
+
|
| 104 |
+
# Get image paths
|
| 105 |
+
print("Loading images...")
|
| 106 |
+
image_folder_path = os.path.join(target_dir, "images")
|
| 107 |
+
all_image_paths = sorted(glob.glob(os.path.join(image_folder_path, "*")))
|
| 108 |
+
|
| 109 |
+
# Filter for image files
|
| 110 |
+
image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]
|
| 111 |
+
all_image_paths = [
|
| 112 |
+
path
|
| 113 |
+
for path in all_image_paths
|
| 114 |
+
if any(path.lower().endswith(ext) for ext in image_extensions)
|
| 115 |
+
]
|
| 116 |
+
|
| 117 |
+
print(f"Found {len(all_image_paths)} images")
|
| 118 |
+
print(f"All image paths: {all_image_paths}")
|
| 119 |
+
|
| 120 |
+
# Use sorted image order (reference view will be selected automatically)
|
| 121 |
+
image_paths = all_image_paths
|
| 122 |
+
print(f"Reference view selection strategy: {ref_view_strategy}")
|
| 123 |
+
|
| 124 |
+
if len(image_paths) == 0:
|
| 125 |
+
raise ValueError("No images found. Check your upload.")
|
| 126 |
+
|
| 127 |
+
# Map UI options to actual method names
|
| 128 |
+
method_mapping = {"high_res": "lower_bound_resize", "low_res": "upper_bound_resize"}
|
| 129 |
+
actual_method = method_mapping.get(process_res_method, "upper_bound_crop")
|
| 130 |
+
|
| 131 |
+
# Run model inference
|
| 132 |
+
print(f"Running inference with method: {actual_method}")
|
| 133 |
+
with torch.no_grad():
|
| 134 |
+
prediction = self.model.inference(
|
| 135 |
+
image_paths,
|
| 136 |
+
export_dir=None,
|
| 137 |
+
process_res_method=actual_method,
|
| 138 |
+
infer_gs=infer_gs,
|
| 139 |
+
ref_view_strategy=ref_view_strategy,
|
| 140 |
+
)
|
| 141 |
+
# num_max_points: int = 1_000_000,
|
| 142 |
+
export_to_glb(
|
| 143 |
+
prediction,
|
| 144 |
+
filter_black_bg=filter_black_bg,
|
| 145 |
+
filter_white_bg=filter_white_bg,
|
| 146 |
+
export_dir=target_dir,
|
| 147 |
+
show_cameras=show_camera,
|
| 148 |
+
conf_thresh_percentile=save_percentage,
|
| 149 |
+
num_max_points=int(num_max_points),
|
| 150 |
+
)
|
| 151 |
+
|
| 152 |
+
# export to gs video if needed
|
| 153 |
+
if infer_gs:
|
| 154 |
+
mode_mapping = {"extend": "extend", "smooth": "interpolate_smooth"}
|
| 155 |
+
print(f"GS mode: {gs_trj_mode}; Backend mode: {mode_mapping[gs_trj_mode]}")
|
| 156 |
+
export_to_gs_video(
|
| 157 |
+
prediction,
|
| 158 |
+
export_dir=target_dir,
|
| 159 |
+
chunk_size=4,
|
| 160 |
+
trj_mode=mode_mapping.get(gs_trj_mode, "extend"),
|
| 161 |
+
enable_tqdm=True,
|
| 162 |
+
vis_depth="hcat",
|
| 163 |
+
video_quality=gs_video_quality,
|
| 164 |
+
)
|
| 165 |
+
|
| 166 |
+
# Save predictions.npz for caching metric depth data
|
| 167 |
+
self._save_predictions_cache(target_dir, prediction)
|
| 168 |
+
|
| 169 |
+
# Process results
|
| 170 |
+
processed_data = self._process_results(target_dir, prediction, image_paths)
|
| 171 |
+
|
| 172 |
+
# Clean up using centralized memory utilities for consistency with backend
|
| 173 |
+
cleanup_cuda_memory()
|
| 174 |
+
|
| 175 |
+
return prediction, processed_data
|
| 176 |
+
|
| 177 |
+
def _save_predictions_cache(self, target_dir: str, prediction: Any) -> None:
|
| 178 |
+
"""
|
| 179 |
+
Save predictions data to predictions.npz for caching.
|
| 180 |
+
|
| 181 |
+
Args:
|
| 182 |
+
target_dir: Directory to save the cache
|
| 183 |
+
prediction: Model prediction object
|
| 184 |
+
"""
|
| 185 |
+
try:
|
| 186 |
+
output_file = os.path.join(target_dir, "predictions.npz")
|
| 187 |
+
|
| 188 |
+
# Build save dict with prediction data
|
| 189 |
+
save_dict = {}
|
| 190 |
+
|
| 191 |
+
# Save processed images if available
|
| 192 |
+
if prediction.processed_images is not None:
|
| 193 |
+
save_dict["images"] = prediction.processed_images
|
| 194 |
+
|
| 195 |
+
# Save depth data
|
| 196 |
+
if prediction.depth is not None:
|
| 197 |
+
save_dict["depths"] = np.round(prediction.depth, 6)
|
| 198 |
+
|
| 199 |
+
# Save confidence if available
|
| 200 |
+
if prediction.conf is not None:
|
| 201 |
+
save_dict["conf"] = np.round(prediction.conf, 2)
|
| 202 |
+
|
| 203 |
+
# Save camera parameters
|
| 204 |
+
if prediction.extrinsics is not None:
|
| 205 |
+
save_dict["extrinsics"] = prediction.extrinsics
|
| 206 |
+
if prediction.intrinsics is not None:
|
| 207 |
+
save_dict["intrinsics"] = prediction.intrinsics
|
| 208 |
+
|
| 209 |
+
# Save to file
|
| 210 |
+
np.savez_compressed(output_file, **save_dict)
|
| 211 |
+
print(f"Saved predictions cache to: {output_file}")
|
| 212 |
+
|
| 213 |
+
except Exception as e:
|
| 214 |
+
print(f"Warning: Failed to save predictions cache: {e}")
|
| 215 |
+
|
| 216 |
+
def _process_results(
|
| 217 |
+
self, target_dir: str, prediction: Any, image_paths: list
|
| 218 |
+
) -> Dict[int, Dict[str, Any]]:
|
| 219 |
+
"""
|
| 220 |
+
Process model results into structured data.
|
| 221 |
+
|
| 222 |
+
Args:
|
| 223 |
+
target_dir: Directory containing results
|
| 224 |
+
prediction: Model prediction object
|
| 225 |
+
image_paths: List of input image paths
|
| 226 |
+
|
| 227 |
+
Returns:
|
| 228 |
+
Dictionary containing processed data for each view
|
| 229 |
+
"""
|
| 230 |
+
processed_data = {}
|
| 231 |
+
|
| 232 |
+
# Read generated depth visualization files
|
| 233 |
+
depth_vis_dir = os.path.join(target_dir, "depth_vis")
|
| 234 |
+
|
| 235 |
+
if os.path.exists(depth_vis_dir):
|
| 236 |
+
depth_files = sorted(glob.glob(os.path.join(depth_vis_dir, "*.jpg")))
|
| 237 |
+
for i, depth_file in enumerate(depth_files):
|
| 238 |
+
# Use processed images directly from API
|
| 239 |
+
processed_image = None
|
| 240 |
+
if prediction.processed_images is not None and i < len(
|
| 241 |
+
prediction.processed_images
|
| 242 |
+
):
|
| 243 |
+
processed_image = prediction.processed_images[i]
|
| 244 |
+
|
| 245 |
+
processed_data[i] = {
|
| 246 |
+
"depth_image": depth_file,
|
| 247 |
+
"image": processed_image,
|
| 248 |
+
"original_image_path": image_paths[i] if i < len(image_paths) else None,
|
| 249 |
+
"depth": prediction.depth[i] if i < len(prediction.depth) else None,
|
| 250 |
+
"intrinsics": (
|
| 251 |
+
prediction.intrinsics[i]
|
| 252 |
+
if prediction.intrinsics is not None and i < len(prediction.intrinsics)
|
| 253 |
+
else None
|
| 254 |
+
),
|
| 255 |
+
"mask": None, # No mask information available
|
| 256 |
+
}
|
| 257 |
+
|
| 258 |
+
return processed_data
|
| 259 |
+
|
| 260 |
+
# cleanup() removed: call cleanup_cuda_memory() directly where needed.
|
depth_anything_3/app/modules/ui_components.py
ADDED
|
@@ -0,0 +1,477 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
UI components module for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This module contains UI component definitions and layout functions.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import os
|
| 22 |
+
from typing import Any, Dict, List, Tuple
|
| 23 |
+
import gradio as gr
|
| 24 |
+
|
| 25 |
+
from depth_anything_3.app.modules.utils import get_logo_base64, get_scene_info
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class UIComponents:
|
| 29 |
+
"""
|
| 30 |
+
Handles UI component creation and layout for the Gradio app.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
def __init__(self):
|
| 34 |
+
"""Initialize the UI components handler."""
|
| 35 |
+
|
| 36 |
+
def create_upload_section(self) -> Tuple[gr.Video, gr.Slider, gr.File, gr.Gallery]:
|
| 37 |
+
"""
|
| 38 |
+
Create the upload section with video, images, and gallery components.
|
| 39 |
+
|
| 40 |
+
Returns:
|
| 41 |
+
A tuple of Gradio components: (input_video, s_time_interval, input_images, image_gallery).
|
| 42 |
+
"""
|
| 43 |
+
input_video = gr.Video(label="Upload Video", interactive=True)
|
| 44 |
+
s_time_interval = gr.Slider(
|
| 45 |
+
minimum=0.1,
|
| 46 |
+
maximum=60,
|
| 47 |
+
value=10,
|
| 48 |
+
step=0.1,
|
| 49 |
+
label="Sampling FPS (Frames Per Second)",
|
| 50 |
+
interactive=True,
|
| 51 |
+
visible=True,
|
| 52 |
+
)
|
| 53 |
+
input_images = gr.File(file_count="multiple", label="Upload Images", interactive=True)
|
| 54 |
+
image_gallery = gr.Gallery(
|
| 55 |
+
label="Preview",
|
| 56 |
+
columns=4,
|
| 57 |
+
height="300px",
|
| 58 |
+
show_download_button=True,
|
| 59 |
+
object_fit="contain",
|
| 60 |
+
preview=True,
|
| 61 |
+
interactive=False,
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
return input_video, s_time_interval, input_images, image_gallery
|
| 65 |
+
|
| 66 |
+
def create_3d_viewer_section(self) -> gr.Model3D:
|
| 67 |
+
"""
|
| 68 |
+
Create the 3D viewer component.
|
| 69 |
+
|
| 70 |
+
Returns:
|
| 71 |
+
3D model viewer component
|
| 72 |
+
"""
|
| 73 |
+
return gr.Model3D(
|
| 74 |
+
height=520,
|
| 75 |
+
zoom_speed=0.5,
|
| 76 |
+
pan_speed=0.5,
|
| 77 |
+
clear_color=[0.0, 0.0, 0.0, 0.0],
|
| 78 |
+
key="persistent_3d_viewer",
|
| 79 |
+
elem_id="reconstruction_3d_viewer",
|
| 80 |
+
)
|
| 81 |
+
|
| 82 |
+
def create_nvs_video(self) -> Tuple[gr.Video, gr.Markdown]:
|
| 83 |
+
"""
|
| 84 |
+
Create the 3DGS rendered video display component and info message.
|
| 85 |
+
|
| 86 |
+
Returns:
|
| 87 |
+
Tuple of (video component, info message component)
|
| 88 |
+
"""
|
| 89 |
+
with gr.Column():
|
| 90 |
+
gs_info = gr.Markdown(
|
| 91 |
+
(
|
| 92 |
+
"‼️ **3D Gaussian Splatting rendering is currently DISABLED.** <br><br><br>"
|
| 93 |
+
"To render novel views from 3DGS, "
|
| 94 |
+
"enable **Infer 3D Gaussian Splatting** below. <br>"
|
| 95 |
+
"Next, in **Visualization Options**, "
|
| 96 |
+
"*optionally* configure the **rendering trajectory** (default: smooth) "
|
| 97 |
+
"and **video quality** (default: low), "
|
| 98 |
+
"then click **Reconstruct**."
|
| 99 |
+
),
|
| 100 |
+
visible=True,
|
| 101 |
+
height=520,
|
| 102 |
+
)
|
| 103 |
+
gs_video = gr.Video(
|
| 104 |
+
height=520,
|
| 105 |
+
label="3DGS Rendered NVS Video (depth shown for reference only)",
|
| 106 |
+
interactive=False,
|
| 107 |
+
visible=False,
|
| 108 |
+
)
|
| 109 |
+
return gs_video, gs_info
|
| 110 |
+
|
| 111 |
+
def create_depth_section(self) -> Tuple[gr.Button, gr.Dropdown, gr.Button, gr.Image]:
|
| 112 |
+
"""
|
| 113 |
+
Create the depth visualization section.
|
| 114 |
+
|
| 115 |
+
Returns:
|
| 116 |
+
A tuple of (prev_depth_btn, depth_view_selector, next_depth_btn, depth_map)
|
| 117 |
+
"""
|
| 118 |
+
with gr.Row(elem_classes=["navigation-row"]):
|
| 119 |
+
prev_depth_btn = gr.Button("◀ Previous", size="sm", scale=1)
|
| 120 |
+
depth_view_selector = gr.Dropdown(
|
| 121 |
+
choices=["View 1"],
|
| 122 |
+
value="View 1",
|
| 123 |
+
label="Select View",
|
| 124 |
+
scale=2,
|
| 125 |
+
interactive=True,
|
| 126 |
+
allow_custom_value=True,
|
| 127 |
+
)
|
| 128 |
+
next_depth_btn = gr.Button("Next ▶", size="sm", scale=1)
|
| 129 |
+
depth_map = gr.Image(
|
| 130 |
+
type="numpy",
|
| 131 |
+
label="Colorized Depth Map",
|
| 132 |
+
format="png",
|
| 133 |
+
interactive=False,
|
| 134 |
+
)
|
| 135 |
+
|
| 136 |
+
return prev_depth_btn, depth_view_selector, next_depth_btn, depth_map
|
| 137 |
+
|
| 138 |
+
def create_measure_section(
|
| 139 |
+
self,
|
| 140 |
+
) -> Tuple[gr.Button, gr.Dropdown, gr.Button, gr.Image, gr.Image, gr.Markdown]:
|
| 141 |
+
"""
|
| 142 |
+
Create the measurement section.
|
| 143 |
+
|
| 144 |
+
Returns:
|
| 145 |
+
A tuple of (prev_measure_btn, measure_view_selector, next_measure_btn, measure_image,
|
| 146 |
+
measure_depth_image, measure_text)
|
| 147 |
+
"""
|
| 148 |
+
from depth_anything_3.app.css_and_html import MEASURE_INSTRUCTIONS_HTML
|
| 149 |
+
|
| 150 |
+
gr.Markdown(MEASURE_INSTRUCTIONS_HTML)
|
| 151 |
+
with gr.Row(elem_classes=["navigation-row"]):
|
| 152 |
+
prev_measure_btn = gr.Button("◀ Previous", size="sm", scale=1)
|
| 153 |
+
measure_view_selector = gr.Dropdown(
|
| 154 |
+
choices=["View 1"],
|
| 155 |
+
value="View 1",
|
| 156 |
+
label="Select View",
|
| 157 |
+
scale=2,
|
| 158 |
+
interactive=True,
|
| 159 |
+
allow_custom_value=True,
|
| 160 |
+
)
|
| 161 |
+
next_measure_btn = gr.Button("Next ▶", size="sm", scale=1)
|
| 162 |
+
with gr.Row():
|
| 163 |
+
measure_image = gr.Image(
|
| 164 |
+
type="numpy",
|
| 165 |
+
show_label=False,
|
| 166 |
+
format="webp",
|
| 167 |
+
interactive=False,
|
| 168 |
+
sources=[],
|
| 169 |
+
label="RGB Image",
|
| 170 |
+
scale=1,
|
| 171 |
+
height=275,
|
| 172 |
+
)
|
| 173 |
+
measure_depth_image = gr.Image(
|
| 174 |
+
type="numpy",
|
| 175 |
+
show_label=False,
|
| 176 |
+
format="webp",
|
| 177 |
+
interactive=False,
|
| 178 |
+
sources=[],
|
| 179 |
+
label="Depth Visualization (Right Half)",
|
| 180 |
+
scale=1,
|
| 181 |
+
height=275,
|
| 182 |
+
)
|
| 183 |
+
gr.Markdown(
|
| 184 |
+
"**Note:** Images have been adjusted to model processing size. "
|
| 185 |
+
"Click two points on the RGB image to measure distance."
|
| 186 |
+
)
|
| 187 |
+
measure_text = gr.Markdown("")
|
| 188 |
+
|
| 189 |
+
return (
|
| 190 |
+
prev_measure_btn,
|
| 191 |
+
measure_view_selector,
|
| 192 |
+
next_measure_btn,
|
| 193 |
+
measure_image,
|
| 194 |
+
measure_depth_image,
|
| 195 |
+
measure_text,
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
def create_inference_control_section(self) -> Tuple[gr.Dropdown, gr.Checkbox, gr.Dropdown]:
|
| 199 |
+
"""
|
| 200 |
+
Create the inference control section (before inference).
|
| 201 |
+
|
| 202 |
+
Returns:
|
| 203 |
+
Tuple of (process_res_method_dropdown, infer_gs, ref_view_strategy)
|
| 204 |
+
"""
|
| 205 |
+
with gr.Row():
|
| 206 |
+
process_res_method_dropdown = gr.Dropdown(
|
| 207 |
+
choices=["high_res", "low_res"],
|
| 208 |
+
value="low_res",
|
| 209 |
+
label="Image Processing Method",
|
| 210 |
+
info="low_res for much more images",
|
| 211 |
+
scale=1,
|
| 212 |
+
)
|
| 213 |
+
# Modify line 220, add color class
|
| 214 |
+
infer_gs = gr.Checkbox(
|
| 215 |
+
label="Infer 3D Gaussian Splatting",
|
| 216 |
+
value=False,
|
| 217 |
+
info=(
|
| 218 |
+
'Enable novel view rendering from 3DGS (<i class="fas fa-triangle-exclamation '
|
| 219 |
+
'fa-color-red"></i> requires extra processing time)'
|
| 220 |
+
),
|
| 221 |
+
scale=1,
|
| 222 |
+
)
|
| 223 |
+
ref_view_strategy = gr.Dropdown(
|
| 224 |
+
choices=["saddle_balanced", "saddle_sim_range", "first", "middle"],
|
| 225 |
+
value="saddle_balanced",
|
| 226 |
+
label="Reference View Strategy",
|
| 227 |
+
info="Strategy for selecting reference view from multiple inputs",
|
| 228 |
+
scale=1,
|
| 229 |
+
)
|
| 230 |
+
|
| 231 |
+
return (process_res_method_dropdown, infer_gs, ref_view_strategy)
|
| 232 |
+
|
| 233 |
+
def create_display_control_section(
|
| 234 |
+
self,
|
| 235 |
+
) -> Tuple[
|
| 236 |
+
gr.Checkbox,
|
| 237 |
+
gr.Checkbox,
|
| 238 |
+
gr.Checkbox,
|
| 239 |
+
gr.Slider,
|
| 240 |
+
gr.Slider,
|
| 241 |
+
gr.Dropdown,
|
| 242 |
+
gr.Dropdown,
|
| 243 |
+
gr.Button,
|
| 244 |
+
gr.ClearButton,
|
| 245 |
+
]:
|
| 246 |
+
"""
|
| 247 |
+
Create the display control section (options for visualization).
|
| 248 |
+
|
| 249 |
+
Returns:
|
| 250 |
+
Tuple of display control components including buttons
|
| 251 |
+
"""
|
| 252 |
+
with gr.Column():
|
| 253 |
+
# 3DGS options at the top
|
| 254 |
+
with gr.Row():
|
| 255 |
+
gs_trj_mode = gr.Dropdown(
|
| 256 |
+
choices=["smooth", "extend"],
|
| 257 |
+
value="smooth",
|
| 258 |
+
label=("Rendering trajectory for 3DGS viewpoints (requires n_views ≥ 2)"),
|
| 259 |
+
info=("'smooth' for view interpolation; 'extend' for longer trajectory"),
|
| 260 |
+
visible=False, # initially hidden
|
| 261 |
+
)
|
| 262 |
+
gs_video_quality = gr.Dropdown(
|
| 263 |
+
choices=["low", "medium", "high"],
|
| 264 |
+
value="low",
|
| 265 |
+
label=("Video quality for 3DGS rendered outputs"),
|
| 266 |
+
info=("'low' for faster loading speed; 'high' for better visual quality"),
|
| 267 |
+
visible=False, # initially hidden
|
| 268 |
+
)
|
| 269 |
+
|
| 270 |
+
# Reconstruct and Clear buttons (before Visualization Options)
|
| 271 |
+
with gr.Row():
|
| 272 |
+
submit_btn = gr.Button("Reconstruct", scale=1, variant="primary")
|
| 273 |
+
clear_btn = gr.ClearButton(scale=1)
|
| 274 |
+
|
| 275 |
+
gr.Markdown("### Visualization Options: (Click Reconstruct to update)")
|
| 276 |
+
show_cam = gr.Checkbox(label="Show Camera", value=True)
|
| 277 |
+
filter_black_bg = gr.Checkbox(label="Filter Black Background", value=False)
|
| 278 |
+
filter_white_bg = gr.Checkbox(label="Filter White Background", value=False)
|
| 279 |
+
save_percentage = gr.Slider(
|
| 280 |
+
minimum=0,
|
| 281 |
+
maximum=100,
|
| 282 |
+
value=10,
|
| 283 |
+
step=1,
|
| 284 |
+
label="Filter Percentage",
|
| 285 |
+
info="Confidence Threshold (%): Higher values filter more points.",
|
| 286 |
+
)
|
| 287 |
+
num_max_points = gr.Slider(
|
| 288 |
+
minimum=1000,
|
| 289 |
+
maximum=100000,
|
| 290 |
+
value=1000,
|
| 291 |
+
step=1000,
|
| 292 |
+
label="Max Points (K points)",
|
| 293 |
+
info="Maximum number of points to export to GLB (in thousands)",
|
| 294 |
+
)
|
| 295 |
+
|
| 296 |
+
return (
|
| 297 |
+
show_cam,
|
| 298 |
+
filter_black_bg,
|
| 299 |
+
filter_white_bg,
|
| 300 |
+
save_percentage,
|
| 301 |
+
num_max_points,
|
| 302 |
+
gs_trj_mode,
|
| 303 |
+
gs_video_quality,
|
| 304 |
+
submit_btn,
|
| 305 |
+
clear_btn,
|
| 306 |
+
)
|
| 307 |
+
|
| 308 |
+
def create_control_section(
|
| 309 |
+
self,
|
| 310 |
+
) -> Tuple[
|
| 311 |
+
gr.Button,
|
| 312 |
+
gr.ClearButton,
|
| 313 |
+
gr.Dropdown,
|
| 314 |
+
gr.Checkbox,
|
| 315 |
+
gr.Checkbox,
|
| 316 |
+
gr.Checkbox,
|
| 317 |
+
gr.Checkbox,
|
| 318 |
+
gr.Checkbox,
|
| 319 |
+
gr.Dropdown,
|
| 320 |
+
gr.Checkbox,
|
| 321 |
+
gr.Textbox,
|
| 322 |
+
]:
|
| 323 |
+
"""
|
| 324 |
+
Create the control section with buttons and options.
|
| 325 |
+
|
| 326 |
+
Returns:
|
| 327 |
+
Tuple of control components
|
| 328 |
+
"""
|
| 329 |
+
with gr.Row():
|
| 330 |
+
submit_btn = gr.Button("Reconstruct", scale=1, variant="primary")
|
| 331 |
+
clear_btn = gr.ClearButton(
|
| 332 |
+
scale=1,
|
| 333 |
+
)
|
| 334 |
+
|
| 335 |
+
with gr.Row():
|
| 336 |
+
frame_filter = gr.Dropdown(
|
| 337 |
+
choices=["All"], value="All", label="Show Points from Frame"
|
| 338 |
+
)
|
| 339 |
+
with gr.Column():
|
| 340 |
+
gr.Markdown("### Visualization Option: (Click Reconstruct to update)")
|
| 341 |
+
show_cam = gr.Checkbox(label="Show Camera", value=True)
|
| 342 |
+
show_mesh = gr.Checkbox(label="Show Mesh", value=True)
|
| 343 |
+
filter_black_bg = gr.Checkbox(label="Filter Black Background", value=False)
|
| 344 |
+
filter_white_bg = gr.Checkbox(label="Filter White Background", value=False)
|
| 345 |
+
gr.Markdown("### Reconstruction Options: (updated on next run)")
|
| 346 |
+
apply_mask_checkbox = gr.Checkbox(
|
| 347 |
+
label="Apply mask for predicted ambiguous depth classes & edges",
|
| 348 |
+
value=True,
|
| 349 |
+
)
|
| 350 |
+
process_res_method_dropdown = gr.Dropdown(
|
| 351 |
+
choices=[
|
| 352 |
+
"upper_bound_resize",
|
| 353 |
+
"upper_bound_crop",
|
| 354 |
+
"lower_bound_resize",
|
| 355 |
+
"lower_bound_crop",
|
| 356 |
+
],
|
| 357 |
+
value="upper_bound_resize",
|
| 358 |
+
label="Image Processing Method",
|
| 359 |
+
info="Method for resizing input images",
|
| 360 |
+
)
|
| 361 |
+
save_to_gallery_checkbox = gr.Checkbox(
|
| 362 |
+
label="Save to Gallery",
|
| 363 |
+
value=False,
|
| 364 |
+
info="Save current reconstruction results to gallery directory",
|
| 365 |
+
)
|
| 366 |
+
gallery_name_input = gr.Textbox(
|
| 367 |
+
label="Gallery Name",
|
| 368 |
+
placeholder="Enter a name for the gallery folder",
|
| 369 |
+
value="",
|
| 370 |
+
info="Leave empty for auto-generated name with timestamp",
|
| 371 |
+
)
|
| 372 |
+
|
| 373 |
+
return (
|
| 374 |
+
submit_btn,
|
| 375 |
+
clear_btn,
|
| 376 |
+
frame_filter,
|
| 377 |
+
show_cam,
|
| 378 |
+
show_mesh,
|
| 379 |
+
filter_black_bg,
|
| 380 |
+
filter_white_bg,
|
| 381 |
+
apply_mask_checkbox,
|
| 382 |
+
process_res_method_dropdown,
|
| 383 |
+
save_to_gallery_checkbox,
|
| 384 |
+
gallery_name_input,
|
| 385 |
+
)
|
| 386 |
+
|
| 387 |
+
def create_example_scenes_section(self) -> List[Dict[str, Any]]:
|
| 388 |
+
"""
|
| 389 |
+
Create the example scenes section.
|
| 390 |
+
|
| 391 |
+
Returns:
|
| 392 |
+
List of scene information dictionaries
|
| 393 |
+
"""
|
| 394 |
+
# Get workspace directory from environment variable
|
| 395 |
+
workspace_dir = os.environ.get("DA3_WORKSPACE_DIR", "gradio_workspace")
|
| 396 |
+
examples_dir = os.path.join(workspace_dir, "examples")
|
| 397 |
+
|
| 398 |
+
# Get scene information
|
| 399 |
+
scenes = get_scene_info(examples_dir)
|
| 400 |
+
|
| 401 |
+
return scenes
|
| 402 |
+
|
| 403 |
+
def create_example_scene_grid(self, scenes: List[Dict[str, Any]]) -> List[gr.Image]:
|
| 404 |
+
"""
|
| 405 |
+
Create the example scene grid.
|
| 406 |
+
|
| 407 |
+
Args:
|
| 408 |
+
scenes: List of scene information dictionaries
|
| 409 |
+
|
| 410 |
+
Returns:
|
| 411 |
+
List of scene image components
|
| 412 |
+
"""
|
| 413 |
+
scene_components = []
|
| 414 |
+
|
| 415 |
+
if scenes:
|
| 416 |
+
for i in range(0, len(scenes), 4): # Process 4 scenes per row
|
| 417 |
+
with gr.Row():
|
| 418 |
+
for j in range(4):
|
| 419 |
+
scene_idx = i + j
|
| 420 |
+
if scene_idx < len(scenes):
|
| 421 |
+
scene = scenes[scene_idx]
|
| 422 |
+
with gr.Column(scale=1, elem_classes=["clickable-thumbnail"]):
|
| 423 |
+
# Clickable thumbnail
|
| 424 |
+
scene_img = gr.Image(
|
| 425 |
+
value=scene["thumbnail"],
|
| 426 |
+
height=150,
|
| 427 |
+
interactive=False,
|
| 428 |
+
show_label=False,
|
| 429 |
+
elem_id=f"scene_thumb_{scene['name']}",
|
| 430 |
+
sources=[],
|
| 431 |
+
)
|
| 432 |
+
scene_components.append(scene_img)
|
| 433 |
+
|
| 434 |
+
# Scene name and image count as text below thumbnail
|
| 435 |
+
gr.Markdown(
|
| 436 |
+
f"**{scene['name']}** \n {scene['num_images']} images",
|
| 437 |
+
elem_classes=["scene-info"],
|
| 438 |
+
)
|
| 439 |
+
else:
|
| 440 |
+
# Empty column to maintain grid structure
|
| 441 |
+
with gr.Column(scale=1):
|
| 442 |
+
pass
|
| 443 |
+
|
| 444 |
+
return scene_components
|
| 445 |
+
|
| 446 |
+
def create_header_section(self) -> gr.HTML:
|
| 447 |
+
"""
|
| 448 |
+
Create the header section with logo and title.
|
| 449 |
+
|
| 450 |
+
Returns:
|
| 451 |
+
Header HTML component
|
| 452 |
+
"""
|
| 453 |
+
from depth_anything_3.app.css_and_html import get_header_html
|
| 454 |
+
|
| 455 |
+
return gr.HTML(get_header_html(get_logo_base64()))
|
| 456 |
+
|
| 457 |
+
def create_description_section(self) -> gr.HTML:
|
| 458 |
+
"""
|
| 459 |
+
Create the description section.
|
| 460 |
+
|
| 461 |
+
Returns:
|
| 462 |
+
Description HTML component
|
| 463 |
+
"""
|
| 464 |
+
from depth_anything_3.app.css_and_html import get_description_html
|
| 465 |
+
|
| 466 |
+
return gr.HTML(get_description_html())
|
| 467 |
+
|
| 468 |
+
def create_acknowledgements_section(self) -> gr.HTML:
|
| 469 |
+
"""
|
| 470 |
+
Create the acknowledgements section.
|
| 471 |
+
|
| 472 |
+
Returns:
|
| 473 |
+
Acknowledgements HTML component
|
| 474 |
+
"""
|
| 475 |
+
from depth_anything_3.app.css_and_html import get_acknowledgements_html
|
| 476 |
+
|
| 477 |
+
return gr.HTML(get_acknowledgements_html())
|
depth_anything_3/app/modules/utils.py
ADDED
|
@@ -0,0 +1,207 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Utility functions for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This module contains helper functions for data processing, visualization,
|
| 19 |
+
and file operations.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
import json
|
| 24 |
+
import os
|
| 25 |
+
import shutil
|
| 26 |
+
from datetime import datetime
|
| 27 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 28 |
+
import numpy as np
|
| 29 |
+
|
| 30 |
+
def create_depth_visualization(depth: np.ndarray) -> Optional[np.ndarray]:
|
| 31 |
+
"""
|
| 32 |
+
Create a colored depth visualization.
|
| 33 |
+
|
| 34 |
+
Args:
|
| 35 |
+
depth: Depth array
|
| 36 |
+
|
| 37 |
+
Returns:
|
| 38 |
+
Colored depth visualization or None
|
| 39 |
+
"""
|
| 40 |
+
if depth is None:
|
| 41 |
+
return None
|
| 42 |
+
|
| 43 |
+
# Normalize depth to 0-1 range
|
| 44 |
+
depth_min = depth[depth > 0].min() if (depth > 0).any() else 0
|
| 45 |
+
depth_max = depth.max()
|
| 46 |
+
|
| 47 |
+
if depth_max <= depth_min:
|
| 48 |
+
return None
|
| 49 |
+
|
| 50 |
+
# Normalize depth
|
| 51 |
+
depth_norm = (depth - depth_min) / (depth_max - depth_min)
|
| 52 |
+
depth_norm = np.clip(depth_norm, 0, 1)
|
| 53 |
+
|
| 54 |
+
# Apply colormap (using matplotlib's viridis colormap)
|
| 55 |
+
import matplotlib.cm as cm
|
| 56 |
+
|
| 57 |
+
# Convert to colored image
|
| 58 |
+
depth_colored = cm.viridis(depth_norm)[:, :, :3] # Remove alpha channel
|
| 59 |
+
depth_colored = (depth_colored * 255).astype(np.uint8)
|
| 60 |
+
|
| 61 |
+
return depth_colored
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def save_to_gallery_func(
|
| 65 |
+
target_dir: str, processed_data: Dict[int, Dict[str, Any]], gallery_name: Optional[str] = None
|
| 66 |
+
) -> Tuple[bool, str]:
|
| 67 |
+
"""
|
| 68 |
+
Save the current reconstruction results to the gallery directory.
|
| 69 |
+
|
| 70 |
+
Args:
|
| 71 |
+
target_dir: Source directory containing reconstruction results
|
| 72 |
+
processed_data: Processed data dictionary
|
| 73 |
+
gallery_name: Name for the gallery folder
|
| 74 |
+
|
| 75 |
+
Returns:
|
| 76 |
+
Tuple of (success, message)
|
| 77 |
+
"""
|
| 78 |
+
try:
|
| 79 |
+
# Get gallery directory from environment variable or use default
|
| 80 |
+
gallery_dir = os.environ.get(
|
| 81 |
+
"DA3_GALLERY_DIR",
|
| 82 |
+
"workspace/gallery",
|
| 83 |
+
)
|
| 84 |
+
if not os.path.exists(gallery_dir):
|
| 85 |
+
os.makedirs(gallery_dir)
|
| 86 |
+
|
| 87 |
+
# Use provided name or create a unique name
|
| 88 |
+
if gallery_name is None or gallery_name.strip() == "":
|
| 89 |
+
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
| 90 |
+
gallery_name = f"reconstruction_{timestamp}"
|
| 91 |
+
|
| 92 |
+
gallery_path = os.path.join(gallery_dir, gallery_name)
|
| 93 |
+
|
| 94 |
+
# Check if directory already exists
|
| 95 |
+
if os.path.exists(gallery_path):
|
| 96 |
+
return False, f"Save failed: folder '{gallery_name}' already exists"
|
| 97 |
+
|
| 98 |
+
# Create the gallery directory
|
| 99 |
+
os.makedirs(gallery_path, exist_ok=True)
|
| 100 |
+
|
| 101 |
+
# Copy GLB file
|
| 102 |
+
glb_source = os.path.join(target_dir, "scene.glb")
|
| 103 |
+
glb_dest = os.path.join(gallery_path, "scene.glb")
|
| 104 |
+
if os.path.exists(glb_source):
|
| 105 |
+
shutil.copy2(glb_source, glb_dest)
|
| 106 |
+
|
| 107 |
+
# Copy depth visualization images
|
| 108 |
+
depth_vis_dir = os.path.join(target_dir, "depth_vis")
|
| 109 |
+
if os.path.exists(depth_vis_dir):
|
| 110 |
+
gallery_depth_vis = os.path.join(gallery_path, "depth_vis")
|
| 111 |
+
shutil.copytree(depth_vis_dir, gallery_depth_vis)
|
| 112 |
+
|
| 113 |
+
# Copy original images
|
| 114 |
+
images_source = os.path.join(target_dir, "images")
|
| 115 |
+
if os.path.exists(images_source):
|
| 116 |
+
gallery_images = os.path.join(gallery_path, "images")
|
| 117 |
+
shutil.copytree(images_source, gallery_images)
|
| 118 |
+
|
| 119 |
+
scene_preview_source = os.path.join(target_dir, "scene.jpg")
|
| 120 |
+
scene_preview_dest = os.path.join(gallery_path, "scene.jpg")
|
| 121 |
+
shutil.copy2(scene_preview_source, scene_preview_dest)
|
| 122 |
+
|
| 123 |
+
# Save metadata
|
| 124 |
+
metadata = {
|
| 125 |
+
"timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
|
| 126 |
+
"num_images": len(processed_data) if processed_data else 0,
|
| 127 |
+
"gallery_name": gallery_name,
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
with open(os.path.join(gallery_path, "metadata.json"), "w") as f:
|
| 131 |
+
json.dump(metadata, f, indent=2)
|
| 132 |
+
|
| 133 |
+
print(f"Saved reconstruction to gallery: {gallery_path}")
|
| 134 |
+
return True, f"Save successful: saved to {gallery_path}"
|
| 135 |
+
|
| 136 |
+
except Exception as e:
|
| 137 |
+
print(f"Error saving to gallery: {e}")
|
| 138 |
+
return False, f"Save failed: {str(e)}"
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def get_scene_info(examples_dir: str) -> List[Dict[str, Any]]:
|
| 142 |
+
"""
|
| 143 |
+
Get information about scenes in the examples directory.
|
| 144 |
+
|
| 145 |
+
Args:
|
| 146 |
+
examples_dir: Path to examples directory
|
| 147 |
+
|
| 148 |
+
Returns:
|
| 149 |
+
List of scene information dictionaries
|
| 150 |
+
"""
|
| 151 |
+
import glob
|
| 152 |
+
|
| 153 |
+
scenes = []
|
| 154 |
+
if not os.path.exists(examples_dir):
|
| 155 |
+
return scenes
|
| 156 |
+
|
| 157 |
+
for scene_folder in sorted(os.listdir(examples_dir)):
|
| 158 |
+
scene_path = os.path.join(examples_dir, scene_folder)
|
| 159 |
+
if os.path.isdir(scene_path):
|
| 160 |
+
# Find all image files in the scene folder
|
| 161 |
+
image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.tiff", "*.tif"]
|
| 162 |
+
image_files = []
|
| 163 |
+
for ext in image_extensions:
|
| 164 |
+
image_files.extend(glob.glob(os.path.join(scene_path, ext)))
|
| 165 |
+
image_files.extend(glob.glob(os.path.join(scene_path, ext.upper())))
|
| 166 |
+
|
| 167 |
+
if image_files:
|
| 168 |
+
# Sort images and get the first one for thumbnail
|
| 169 |
+
image_files = sorted(image_files)
|
| 170 |
+
first_image = image_files[0]
|
| 171 |
+
num_images = len(image_files)
|
| 172 |
+
|
| 173 |
+
scenes.append(
|
| 174 |
+
{
|
| 175 |
+
"name": scene_folder,
|
| 176 |
+
"path": scene_path,
|
| 177 |
+
"thumbnail": first_image,
|
| 178 |
+
"num_images": num_images,
|
| 179 |
+
"image_files": image_files,
|
| 180 |
+
}
|
| 181 |
+
)
|
| 182 |
+
|
| 183 |
+
return scenes
|
| 184 |
+
|
| 185 |
+
|
| 186 |
+
# NOTE: cleanup was moved to a single canonical helper in
|
| 187 |
+
# `depth_anything_3.utils.memory.cleanup_cuda_memory`.
|
| 188 |
+
# Callers should import and call that directly instead of using this module.
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
def get_logo_base64() -> Optional[str]:
|
| 192 |
+
"""
|
| 193 |
+
Convert WAI logo to base64 for embedding in HTML.
|
| 194 |
+
|
| 195 |
+
Returns:
|
| 196 |
+
Base64 encoded logo string or None
|
| 197 |
+
"""
|
| 198 |
+
import base64
|
| 199 |
+
|
| 200 |
+
logo_path = "examples/WAI-Logo/wai_logo.png"
|
| 201 |
+
try:
|
| 202 |
+
with open(logo_path, "rb") as img_file:
|
| 203 |
+
img_data = img_file.read()
|
| 204 |
+
base64_str = base64.b64encode(img_data).decode()
|
| 205 |
+
return f"data:image/png;base64,{base64_str}"
|
| 206 |
+
except FileNotFoundError:
|
| 207 |
+
return None
|
depth_anything_3/app/modules/visualization.py
ADDED
|
@@ -0,0 +1,434 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Visualization module for Depth Anything 3 Gradio app.
|
| 17 |
+
|
| 18 |
+
This module handles visualization updates, navigation, and measurement functionality.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import os
|
| 22 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 23 |
+
import cv2
|
| 24 |
+
import gradio as gr
|
| 25 |
+
import numpy as np
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class VisualizationHandler:
|
| 29 |
+
"""
|
| 30 |
+
Handles visualization updates and navigation for the Gradio app.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
def __init__(self):
|
| 34 |
+
"""Initialize the visualization handler."""
|
| 35 |
+
|
| 36 |
+
def update_view_selectors(
|
| 37 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]]
|
| 38 |
+
) -> Tuple[gr.Dropdown, gr.Dropdown]:
|
| 39 |
+
"""
|
| 40 |
+
Update view selector dropdowns based on available views.
|
| 41 |
+
|
| 42 |
+
Args:
|
| 43 |
+
processed_data: Processed data dictionary
|
| 44 |
+
|
| 45 |
+
Returns:
|
| 46 |
+
Tuple of (depth_view_selector, measure_view_selector)
|
| 47 |
+
"""
|
| 48 |
+
if processed_data is None or len(processed_data) == 0:
|
| 49 |
+
choices = ["View 1"]
|
| 50 |
+
else:
|
| 51 |
+
num_views = len(processed_data)
|
| 52 |
+
choices = [f"View {i + 1}" for i in range(num_views)]
|
| 53 |
+
|
| 54 |
+
return (
|
| 55 |
+
gr.Dropdown(choices=choices, value=choices[0]), # depth_view_selector
|
| 56 |
+
gr.Dropdown(choices=choices, value=choices[0]), # measure_view_selector
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
def get_view_data_by_index(
|
| 60 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int
|
| 61 |
+
) -> Optional[Dict[str, Any]]:
|
| 62 |
+
"""
|
| 63 |
+
Get view data by index, handling bounds.
|
| 64 |
+
|
| 65 |
+
Args:
|
| 66 |
+
processed_data: Processed data dictionary
|
| 67 |
+
view_index: Index of the view to get
|
| 68 |
+
|
| 69 |
+
Returns:
|
| 70 |
+
View data dictionary or None
|
| 71 |
+
"""
|
| 72 |
+
if processed_data is None or len(processed_data) == 0:
|
| 73 |
+
return None
|
| 74 |
+
|
| 75 |
+
view_keys = list(processed_data.keys())
|
| 76 |
+
if view_index < 0 or view_index >= len(view_keys):
|
| 77 |
+
view_index = 0
|
| 78 |
+
|
| 79 |
+
return processed_data[view_keys[view_index]]
|
| 80 |
+
|
| 81 |
+
def update_depth_view(
|
| 82 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int
|
| 83 |
+
) -> Optional[str]:
|
| 84 |
+
"""
|
| 85 |
+
Update depth view for a specific view index.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
processed_data: Processed data dictionary
|
| 89 |
+
view_index: Index of the view to update
|
| 90 |
+
|
| 91 |
+
Returns:
|
| 92 |
+
Path to depth visualization image or None
|
| 93 |
+
"""
|
| 94 |
+
view_data = self.get_view_data_by_index(processed_data, view_index)
|
| 95 |
+
if view_data is None or view_data.get("depth_image") is None:
|
| 96 |
+
return None
|
| 97 |
+
|
| 98 |
+
# Return the depth visualization image directly
|
| 99 |
+
return view_data["depth_image"]
|
| 100 |
+
|
| 101 |
+
def navigate_depth_view(
|
| 102 |
+
self,
|
| 103 |
+
processed_data: Optional[Dict[int, Dict[str, Any]]],
|
| 104 |
+
current_selector_value: str,
|
| 105 |
+
direction: int,
|
| 106 |
+
) -> Tuple[str, Optional[str]]:
|
| 107 |
+
"""
|
| 108 |
+
Navigate depth view (direction: -1 for previous, +1 for next).
|
| 109 |
+
|
| 110 |
+
Args:
|
| 111 |
+
processed_data: Processed data dictionary
|
| 112 |
+
current_selector_value: Current selector value
|
| 113 |
+
direction: Direction to navigate (-1 for previous, +1 for next)
|
| 114 |
+
|
| 115 |
+
Returns:
|
| 116 |
+
Tuple of (new_selector_value, depth_vis)
|
| 117 |
+
"""
|
| 118 |
+
if processed_data is None or len(processed_data) == 0:
|
| 119 |
+
return "View 1", None
|
| 120 |
+
|
| 121 |
+
# Parse current view number
|
| 122 |
+
try:
|
| 123 |
+
current_view = int(current_selector_value.split()[1]) - 1
|
| 124 |
+
except: # noqa
|
| 125 |
+
current_view = 0
|
| 126 |
+
|
| 127 |
+
num_views = len(processed_data)
|
| 128 |
+
new_view = (current_view + direction) % num_views
|
| 129 |
+
|
| 130 |
+
new_selector_value = f"View {new_view + 1}"
|
| 131 |
+
depth_vis = self.update_depth_view(processed_data, new_view)
|
| 132 |
+
|
| 133 |
+
return new_selector_value, depth_vis
|
| 134 |
+
|
| 135 |
+
def update_measure_view(
|
| 136 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]], view_index: int
|
| 137 |
+
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List]:
|
| 138 |
+
"""
|
| 139 |
+
Update measure view for a specific view index.
|
| 140 |
+
|
| 141 |
+
Args:
|
| 142 |
+
processed_data: Processed data dictionary
|
| 143 |
+
view_index: Index of the view to update
|
| 144 |
+
|
| 145 |
+
Returns:
|
| 146 |
+
Tuple of (measure_image, depth_right_half, measure_points)
|
| 147 |
+
"""
|
| 148 |
+
view_data = self.get_view_data_by_index(processed_data, view_index)
|
| 149 |
+
if view_data is None:
|
| 150 |
+
return None, None, [] # image, depth_right_half, measure_points
|
| 151 |
+
|
| 152 |
+
# Get the processed (resized) image
|
| 153 |
+
if "image" in view_data and view_data["image"] is not None:
|
| 154 |
+
image = view_data["image"].copy()
|
| 155 |
+
else:
|
| 156 |
+
return None, None, []
|
| 157 |
+
|
| 158 |
+
# Ensure image is in uint8 format
|
| 159 |
+
if image.dtype != np.uint8:
|
| 160 |
+
if image.max() <= 1.0:
|
| 161 |
+
image = (image * 255).astype(np.uint8)
|
| 162 |
+
else:
|
| 163 |
+
image = image.astype(np.uint8)
|
| 164 |
+
|
| 165 |
+
# Extract right half of the depth visualization (pure depth part)
|
| 166 |
+
depth_image_path = view_data.get("depth_image", None)
|
| 167 |
+
depth_right_half = None
|
| 168 |
+
|
| 169 |
+
if depth_image_path and os.path.exists(depth_image_path):
|
| 170 |
+
try:
|
| 171 |
+
# Load the combined depth visualization image
|
| 172 |
+
depth_combined = cv2.imread(depth_image_path)
|
| 173 |
+
depth_combined = cv2.cvtColor(depth_combined, cv2.COLOR_BGR2RGB)
|
| 174 |
+
if depth_combined is not None:
|
| 175 |
+
height, width = depth_combined.shape[:2]
|
| 176 |
+
# Extract right half (depth visualization part)
|
| 177 |
+
depth_right_half = depth_combined[:, width // 2 :]
|
| 178 |
+
except Exception as e:
|
| 179 |
+
print(f"Error extracting depth right half: {e}")
|
| 180 |
+
|
| 181 |
+
return image, depth_right_half, []
|
| 182 |
+
|
| 183 |
+
def navigate_measure_view(
|
| 184 |
+
self,
|
| 185 |
+
processed_data: Optional[Dict[int, Dict[str, Any]]],
|
| 186 |
+
current_selector_value: str,
|
| 187 |
+
direction: int,
|
| 188 |
+
) -> Tuple[str, Optional[np.ndarray], Optional[str], List]:
|
| 189 |
+
"""
|
| 190 |
+
Navigate measure view (direction: -1 for previous, +1 for next).
|
| 191 |
+
|
| 192 |
+
Args:
|
| 193 |
+
processed_data: Processed data dictionary
|
| 194 |
+
current_selector_value: Current selector value
|
| 195 |
+
direction: Direction to navigate (-1 for previous, +1 for next)
|
| 196 |
+
|
| 197 |
+
Returns:
|
| 198 |
+
Tuple of (new_selector_value, measure_image, depth_image_path, measure_points)
|
| 199 |
+
"""
|
| 200 |
+
if processed_data is None or len(processed_data) == 0:
|
| 201 |
+
return "View 1", None, None, []
|
| 202 |
+
|
| 203 |
+
# Parse current view number
|
| 204 |
+
try:
|
| 205 |
+
current_view = int(current_selector_value.split()[1]) - 1
|
| 206 |
+
except: # noqa
|
| 207 |
+
current_view = 0
|
| 208 |
+
|
| 209 |
+
num_views = len(processed_data)
|
| 210 |
+
new_view = (current_view + direction) % num_views
|
| 211 |
+
|
| 212 |
+
new_selector_value = f"View {new_view + 1}"
|
| 213 |
+
measure_image, depth_right_half, measure_points = self.update_measure_view(
|
| 214 |
+
processed_data, new_view
|
| 215 |
+
)
|
| 216 |
+
|
| 217 |
+
return new_selector_value, measure_image, depth_right_half, measure_points
|
| 218 |
+
|
| 219 |
+
def populate_visualization_tabs(
|
| 220 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]]
|
| 221 |
+
) -> Tuple[Optional[str], Optional[np.ndarray], Optional[str], List]:
|
| 222 |
+
"""
|
| 223 |
+
Populate the depth and measure tabs with processed data.
|
| 224 |
+
|
| 225 |
+
Args:
|
| 226 |
+
processed_data: Processed data dictionary
|
| 227 |
+
|
| 228 |
+
Returns:
|
| 229 |
+
Tuple of (depth_vis, measure_img, depth_image_path, measure_points)
|
| 230 |
+
"""
|
| 231 |
+
if processed_data is None or len(processed_data) == 0:
|
| 232 |
+
return None, None, None, []
|
| 233 |
+
|
| 234 |
+
# Use update function to get depth visualization
|
| 235 |
+
depth_vis = self.update_depth_view(processed_data, 0)
|
| 236 |
+
measure_img, depth_right_half, _ = self.update_measure_view(processed_data, 0)
|
| 237 |
+
|
| 238 |
+
return depth_vis, measure_img, depth_right_half, []
|
| 239 |
+
|
| 240 |
+
def reset_measure(
|
| 241 |
+
self, processed_data: Optional[Dict[int, Dict[str, Any]]]
|
| 242 |
+
) -> Tuple[Optional[np.ndarray], List, str]:
|
| 243 |
+
"""
|
| 244 |
+
Reset measure points.
|
| 245 |
+
|
| 246 |
+
Args:
|
| 247 |
+
processed_data: Processed data dictionary
|
| 248 |
+
|
| 249 |
+
Returns:
|
| 250 |
+
Tuple of (image, measure_points, text)
|
| 251 |
+
"""
|
| 252 |
+
if processed_data is None or len(processed_data) == 0:
|
| 253 |
+
return None, [], ""
|
| 254 |
+
|
| 255 |
+
# Return the first view image
|
| 256 |
+
first_view = list(processed_data.values())[0]
|
| 257 |
+
return first_view["image"], [], ""
|
| 258 |
+
|
| 259 |
+
def measure(
|
| 260 |
+
self,
|
| 261 |
+
processed_data: Optional[Dict[int, Dict[str, Any]]],
|
| 262 |
+
measure_points: List,
|
| 263 |
+
current_view_selector: str,
|
| 264 |
+
event: gr.SelectData,
|
| 265 |
+
) -> List:
|
| 266 |
+
"""
|
| 267 |
+
Handle measurement on images.
|
| 268 |
+
|
| 269 |
+
Args:
|
| 270 |
+
processed_data: Processed data dictionary
|
| 271 |
+
measure_points: List of current measure points
|
| 272 |
+
current_view_selector: Current view selector value
|
| 273 |
+
event: Gradio select event
|
| 274 |
+
|
| 275 |
+
Returns:
|
| 276 |
+
List of [image, depth_right_half, measure_points, text]
|
| 277 |
+
"""
|
| 278 |
+
try:
|
| 279 |
+
print(f"Measure function called with selector: {current_view_selector}")
|
| 280 |
+
|
| 281 |
+
if processed_data is None or len(processed_data) == 0:
|
| 282 |
+
return [None, [], "No data available"]
|
| 283 |
+
|
| 284 |
+
# Use the currently selected view instead of always using the first view
|
| 285 |
+
try:
|
| 286 |
+
current_view_index = int(current_view_selector.split()[1]) - 1
|
| 287 |
+
except: # noqa
|
| 288 |
+
current_view_index = 0
|
| 289 |
+
|
| 290 |
+
print(f"Using view index: {current_view_index}")
|
| 291 |
+
|
| 292 |
+
# Get view data safely
|
| 293 |
+
if current_view_index < 0 or current_view_index >= len(processed_data):
|
| 294 |
+
current_view_index = 0
|
| 295 |
+
|
| 296 |
+
view_keys = list(processed_data.keys())
|
| 297 |
+
current_view = processed_data[view_keys[current_view_index]]
|
| 298 |
+
|
| 299 |
+
if current_view is None:
|
| 300 |
+
return [None, [], "No view data available"]
|
| 301 |
+
|
| 302 |
+
point2d = event.index[0], event.index[1]
|
| 303 |
+
print(f"Clicked point: {point2d}")
|
| 304 |
+
|
| 305 |
+
measure_points.append(point2d)
|
| 306 |
+
|
| 307 |
+
# Get image and depth visualization
|
| 308 |
+
image, depth_right_half, _ = self.update_measure_view(
|
| 309 |
+
processed_data, current_view_index
|
| 310 |
+
)
|
| 311 |
+
if image is None:
|
| 312 |
+
return [None, [], "No image available"]
|
| 313 |
+
|
| 314 |
+
image = image.copy()
|
| 315 |
+
|
| 316 |
+
# Ensure image is in uint8 format for proper cv2 operations
|
| 317 |
+
try:
|
| 318 |
+
if image.dtype != np.uint8:
|
| 319 |
+
if image.max() <= 1.0:
|
| 320 |
+
# Image is in [0, 1] range, convert to [0, 255]
|
| 321 |
+
image = (image * 255).astype(np.uint8)
|
| 322 |
+
else:
|
| 323 |
+
# Image is already in [0, 255] range
|
| 324 |
+
image = image.astype(np.uint8)
|
| 325 |
+
except Exception as e:
|
| 326 |
+
print(f"Image conversion error: {e}")
|
| 327 |
+
return [None, [], f"Image conversion error: {e}"]
|
| 328 |
+
|
| 329 |
+
# Draw circles for points
|
| 330 |
+
try:
|
| 331 |
+
for p in measure_points:
|
| 332 |
+
if 0 <= p[0] < image.shape[1] and 0 <= p[1] < image.shape[0]:
|
| 333 |
+
image = cv2.circle(image, p, radius=5, color=(255, 0, 0), thickness=2)
|
| 334 |
+
except Exception as e:
|
| 335 |
+
print(f"Drawing error: {e}")
|
| 336 |
+
return [None, [], f"Drawing error: {e}"]
|
| 337 |
+
|
| 338 |
+
# Get depth information from processed_data
|
| 339 |
+
depth_text = ""
|
| 340 |
+
try:
|
| 341 |
+
for i, p in enumerate(measure_points):
|
| 342 |
+
if (
|
| 343 |
+
current_view["depth"] is not None
|
| 344 |
+
and 0 <= p[1] < current_view["depth"].shape[0]
|
| 345 |
+
and 0 <= p[0] < current_view["depth"].shape[1]
|
| 346 |
+
):
|
| 347 |
+
d = current_view["depth"][p[1], p[0]]
|
| 348 |
+
depth_text += f"- **P{i + 1} depth: {d:.2f}m**\n"
|
| 349 |
+
else:
|
| 350 |
+
depth_text += f"- **P{i + 1}: Click position ({p[0]}, {p[1]}) - No depth information**\n" # noqa: E501
|
| 351 |
+
except Exception as e:
|
| 352 |
+
print(f"Depth text error: {e}")
|
| 353 |
+
depth_text = f"Error computing depth: {e}\n"
|
| 354 |
+
|
| 355 |
+
if len(measure_points) == 2:
|
| 356 |
+
try:
|
| 357 |
+
point1, point2 = measure_points
|
| 358 |
+
# Draw line
|
| 359 |
+
if (
|
| 360 |
+
0 <= point1[0] < image.shape[1]
|
| 361 |
+
and 0 <= point1[1] < image.shape[0]
|
| 362 |
+
and 0 <= point2[0] < image.shape[1]
|
| 363 |
+
and 0 <= point2[1] < image.shape[0]
|
| 364 |
+
):
|
| 365 |
+
image = cv2.line(image, point1, point2, color=(255, 0, 0), thickness=2)
|
| 366 |
+
|
| 367 |
+
# Compute 3D distance using depth information and camera intrinsics
|
| 368 |
+
distance_text = "- **Distance: Unable to calculate 3D distance**"
|
| 369 |
+
if (
|
| 370 |
+
current_view["depth"] is not None
|
| 371 |
+
and 0 <= point1[1] < current_view["depth"].shape[0]
|
| 372 |
+
and 0 <= point1[0] < current_view["depth"].shape[1]
|
| 373 |
+
and 0 <= point2[1] < current_view["depth"].shape[0]
|
| 374 |
+
and 0 <= point2[0] < current_view["depth"].shape[1]
|
| 375 |
+
):
|
| 376 |
+
try:
|
| 377 |
+
# Get depth values at the two points
|
| 378 |
+
d1 = current_view["depth"][point1[1], point1[0]]
|
| 379 |
+
d2 = current_view["depth"][point2[1], point2[0]]
|
| 380 |
+
|
| 381 |
+
# Convert 2D pixel coordinates to 3D world coordinates
|
| 382 |
+
if current_view["intrinsics"] is not None:
|
| 383 |
+
# Get camera intrinsics
|
| 384 |
+
K = current_view["intrinsics"] # 3x3 intrinsic matrix
|
| 385 |
+
fx, fy = K[0, 0], K[1, 1] # focal lengths
|
| 386 |
+
cx, cy = K[0, 2], K[1, 2] # principal point
|
| 387 |
+
|
| 388 |
+
# Convert pixel coordinates to normalized camera coordinates
|
| 389 |
+
# Point 1: (u1, v1) -> (x1, y1, z1)
|
| 390 |
+
u1, v1 = point1[0], point1[1]
|
| 391 |
+
x1 = (u1 - cx) * d1 / fx
|
| 392 |
+
y1 = (v1 - cy) * d1 / fy
|
| 393 |
+
z1 = d1
|
| 394 |
+
|
| 395 |
+
# Point 2: (u2, v2) -> (x2, y2, z2)
|
| 396 |
+
u2, v2 = point2[0], point2[1]
|
| 397 |
+
x2 = (u2 - cx) * d2 / fx
|
| 398 |
+
y2 = (v2 - cy) * d2 / fy
|
| 399 |
+
z2 = d2
|
| 400 |
+
|
| 401 |
+
# Calculate 3D Euclidean distance
|
| 402 |
+
p1_3d = np.array([x1, y1, z1])
|
| 403 |
+
p2_3d = np.array([x2, y2, z2])
|
| 404 |
+
distance_3d = np.linalg.norm(p1_3d - p2_3d)
|
| 405 |
+
|
| 406 |
+
distance_text = f"- **Distance: {distance_3d:.2f}m**"
|
| 407 |
+
else:
|
| 408 |
+
# Fallback to simplified calculation if no intrinsics
|
| 409 |
+
pixel_distance = np.sqrt(
|
| 410 |
+
(point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2
|
| 411 |
+
)
|
| 412 |
+
avg_depth = (d1 + d2) / 2
|
| 413 |
+
scale_factor = avg_depth / 1000 # Rough scaling factor
|
| 414 |
+
estimated_3d_distance = pixel_distance * scale_factor
|
| 415 |
+
distance_text = f"- **Distance: {estimated_3d_distance:.2f}m (estimated, no intrinsics)**" # noqa: E501
|
| 416 |
+
|
| 417 |
+
except Exception as e:
|
| 418 |
+
print(f"Distance computation error: {e}")
|
| 419 |
+
distance_text = f"- **Distance computation error: {e}**"
|
| 420 |
+
|
| 421 |
+
measure_points = []
|
| 422 |
+
text = depth_text + distance_text
|
| 423 |
+
print(f"Measurement complete: {text}")
|
| 424 |
+
return [image, depth_right_half, measure_points, text]
|
| 425 |
+
except Exception as e:
|
| 426 |
+
print(f"Final measurement error: {e}")
|
| 427 |
+
return [None, [], f"Measurement error: {e}"]
|
| 428 |
+
else:
|
| 429 |
+
print(f"Single point measurement: {depth_text}")
|
| 430 |
+
return [image, depth_right_half, measure_points, depth_text]
|
| 431 |
+
|
| 432 |
+
except Exception as e:
|
| 433 |
+
print(f"Overall measure function error: {e}")
|
| 434 |
+
return [None, [], f"Measure function error: {e}"]
|
depth_anything_3/bench/__init__.py
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Depth Anything 3 Benchmark Evaluation Module.
|
| 17 |
+
|
| 18 |
+
This module provides tools for evaluating DepthAnything3 model on various benchmark datasets.
|
| 19 |
+
Currently supported datasets:
|
| 20 |
+
- DTU (3D Reconstruction)
|
| 21 |
+
- DTU-64 (Pose Evaluation Only)
|
| 22 |
+
- ETH3D (3D Reconstruction)
|
| 23 |
+
- 7Scenes (3D Reconstruction)
|
| 24 |
+
- ScanNet++ (3D Reconstruction)
|
| 25 |
+
- HiRoom (3D Reconstruction)
|
| 26 |
+
|
| 27 |
+
Supported evaluation modes:
|
| 28 |
+
- pose: Camera pose estimation evaluation
|
| 29 |
+
- recon_unposed: 3D reconstruction with predicted poses
|
| 30 |
+
- recon_posed: 3D reconstruction with ground truth poses
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
from depth_anything_3.bench.registries import MV_REGISTRY, MONO_REGISTRY
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def __getattr__(name):
|
| 37 |
+
"""Lazy import to avoid circular import when running as __main__."""
|
| 38 |
+
if name == "Evaluator":
|
| 39 |
+
from depth_anything_3.bench.evaluator import Evaluator
|
| 40 |
+
return Evaluator
|
| 41 |
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
__all__ = ["Evaluator", "MV_REGISTRY", "MONO_REGISTRY"]
|
| 45 |
+
|
depth_anything_3/bench/configs/eval_bench.yaml
ADDED
|
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# DepthAnything3 Benchmark Evaluation Configuration
|
| 2 |
+
#
|
| 3 |
+
# This config can be loaded and overridden via command line.
|
| 4 |
+
# Example: python -m depth_anything_3.bench.evaluator --model /path/to/model --work_dir /path/to/workspace
|
| 5 |
+
#
|
| 6 |
+
# See depth_anything_3.cfg for config utility functions.
|
| 7 |
+
|
| 8 |
+
# ==============================================================================
|
| 9 |
+
# Model Configuration
|
| 10 |
+
# ==============================================================================
|
| 11 |
+
model:
|
| 12 |
+
# Path to model checkpoint or HuggingFace model ID
|
| 13 |
+
path: depth-anything/DA3-GIANT
|
| 14 |
+
|
| 15 |
+
# ==============================================================================
|
| 16 |
+
# Workspace Configuration
|
| 17 |
+
# ==============================================================================
|
| 18 |
+
workspace:
|
| 19 |
+
# Working directory for outputs (model results, metrics, etc.)
|
| 20 |
+
work_dir: ./workspace/evaluation
|
| 21 |
+
|
| 22 |
+
# ==============================================================================
|
| 23 |
+
# Evaluation Configuration
|
| 24 |
+
# ==============================================================================
|
| 25 |
+
eval:
|
| 26 |
+
# Datasets to evaluate
|
| 27 |
+
# Options: dtu, dtu64, eth3d, 7scenes (sevenscenes), scannetpp, hiroom
|
| 28 |
+
datasets:
|
| 29 |
+
- eth3d
|
| 30 |
+
- 7scenes
|
| 31 |
+
- scannetpp
|
| 32 |
+
- hiroom
|
| 33 |
+
- dtu
|
| 34 |
+
- dtu64
|
| 35 |
+
|
| 36 |
+
# Evaluation modes
|
| 37 |
+
# Options: pose, recon_unposed, recon_posed, view_syn
|
| 38 |
+
modes:
|
| 39 |
+
- pose
|
| 40 |
+
- recon_unposed
|
| 41 |
+
- recon_posed
|
| 42 |
+
|
| 43 |
+
# Reference view selection strategy for inference
|
| 44 |
+
# Options: first, saddle_balanced, auto, mid
|
| 45 |
+
ref_view_strategy: "first"
|
| 46 |
+
|
| 47 |
+
# Specific scenes to evaluate (null = all scenes)
|
| 48 |
+
# Example: [courtyard, relief] for eth3d
|
| 49 |
+
scenes: null
|
| 50 |
+
|
| 51 |
+
# Maximum number of frames per scene (for sampling)
|
| 52 |
+
# If a scene has more frames, randomly sample to this limit.
|
| 53 |
+
# Set to -1 to disable sampling.
|
| 54 |
+
max_frames: 100
|
| 55 |
+
|
| 56 |
+
# Only run evaluation (skip inference)
|
| 57 |
+
eval_only: false
|
| 58 |
+
|
| 59 |
+
# Only print saved metrics (skip inference and evaluation)
|
| 60 |
+
print_only: false
|
| 61 |
+
|
| 62 |
+
# ==============================================================================
|
| 63 |
+
# Inference Configuration
|
| 64 |
+
# ==============================================================================
|
| 65 |
+
inference:
|
| 66 |
+
# Number of parallel workers for TSDF fusion
|
| 67 |
+
num_fusion_workers: 4
|
| 68 |
+
|
| 69 |
+
# Enable debug mode with verbose output
|
| 70 |
+
debug: false
|
| 71 |
+
|
| 72 |
+
# ==============================================================================
|
| 73 |
+
# Preset Configurations
|
| 74 |
+
# ==============================================================================
|
| 75 |
+
# These can be activated via command line: --preset full_eval
|
| 76 |
+
|
| 77 |
+
presets:
|
| 78 |
+
# Full evaluation on all 6 datasets
|
| 79 |
+
full_eval:
|
| 80 |
+
datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu, dtu64]
|
| 81 |
+
modes: [pose, recon_unposed, recon_posed]
|
| 82 |
+
|
| 83 |
+
# Pose-only evaluation
|
| 84 |
+
pose_only:
|
| 85 |
+
datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu64]
|
| 86 |
+
modes: [pose]
|
| 87 |
+
|
| 88 |
+
# Reconstruction-only evaluation (5 datasets, excluding dtu64)
|
| 89 |
+
recon_only:
|
| 90 |
+
datasets: [eth3d, 7scenes, scannetpp, hiroom, dtu]
|
| 91 |
+
modes: [recon_unposed, recon_posed]
|
| 92 |
+
|
| 93 |
+
# Quick test (single scene per dataset)
|
| 94 |
+
quick_test:
|
| 95 |
+
datasets: [eth3d]
|
| 96 |
+
modes: [pose, recon_unposed]
|
| 97 |
+
scenes: [courtyard]
|
| 98 |
+
|
depth_anything_3/bench/dataset.py
ADDED
|
@@ -0,0 +1,136 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Base dataset class for benchmark evaluation.
|
| 17 |
+
|
| 18 |
+
All dataset implementations should inherit from this class and implement
|
| 19 |
+
the required abstract methods.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import os
|
| 23 |
+
import time
|
| 24 |
+
from abc import abstractmethod
|
| 25 |
+
from typing import Dict as TDict
|
| 26 |
+
|
| 27 |
+
import numpy as np
|
| 28 |
+
import torch
|
| 29 |
+
from addict import Dict
|
| 30 |
+
|
| 31 |
+
from depth_anything_3.bench.utils import compute_pose
|
| 32 |
+
from depth_anything_3.utils.geometry import as_homogeneous
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def _wait_for_file_ready(path: str, timeout: float = 3.0, interval: float = 0.2) -> None:
|
| 36 |
+
"""Wait until file size stabilizes for 2 consecutive checks."""
|
| 37 |
+
last_size = -1
|
| 38 |
+
stable_count = 0
|
| 39 |
+
start = time.time()
|
| 40 |
+
while time.time() - start < timeout:
|
| 41 |
+
time.sleep(interval)
|
| 42 |
+
size = os.path.getsize(path)
|
| 43 |
+
if size == last_size and size > 0:
|
| 44 |
+
stable_count += 1
|
| 45 |
+
if stable_count >= 2: # Need 2 consecutive stable checks
|
| 46 |
+
return
|
| 47 |
+
else:
|
| 48 |
+
stable_count = 0
|
| 49 |
+
last_size = size
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
class Dataset:
|
| 53 |
+
"""
|
| 54 |
+
Base class for all benchmark datasets.
|
| 55 |
+
|
| 56 |
+
Subclasses must implement:
|
| 57 |
+
- SCENES: List of scene identifiers
|
| 58 |
+
- data_root: Path to dataset root
|
| 59 |
+
- get_data(scene): Return scene data (images, intrinsics, extrinsics, etc.)
|
| 60 |
+
- eval3d(scene, fuse_path): Evaluate 3D reconstruction
|
| 61 |
+
- fuse3d(scene, result_path, fuse_path, mode): Fuse depth maps into point cloud
|
| 62 |
+
|
| 63 |
+
Optional overrides:
|
| 64 |
+
- eval_pose(scene, result_path): Evaluate pose estimation (default provided)
|
| 65 |
+
"""
|
| 66 |
+
|
| 67 |
+
# Subclasses should define these
|
| 68 |
+
SCENES: list = []
|
| 69 |
+
data_root: str = ""
|
| 70 |
+
|
| 71 |
+
def __init__(self):
|
| 72 |
+
pass
|
| 73 |
+
|
| 74 |
+
def eval_pose(self, scene: str, result_path: str) -> TDict[str, float]:
|
| 75 |
+
"""
|
| 76 |
+
Evaluate camera pose estimation accuracy.
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
scene: Scene identifier
|
| 80 |
+
result_path: Path to .npz file containing predicted extrinsics
|
| 81 |
+
|
| 82 |
+
Returns:
|
| 83 |
+
Dict with pose metrics (auc30, auc15, auc05, auc03)
|
| 84 |
+
"""
|
| 85 |
+
_wait_for_file_ready(result_path)
|
| 86 |
+
pred = np.load(result_path)
|
| 87 |
+
gt = self.get_data(scene)
|
| 88 |
+
return compute_pose(
|
| 89 |
+
torch.from_numpy(as_homogeneous(pred["extrinsics"])),
|
| 90 |
+
torch.from_numpy(as_homogeneous(gt["extrinsics"])),
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
@abstractmethod
|
| 94 |
+
def get_data(self, scene: str) -> Dict:
|
| 95 |
+
"""
|
| 96 |
+
Get scene data including images, camera parameters, and auxiliary info.
|
| 97 |
+
|
| 98 |
+
Args:
|
| 99 |
+
scene: Scene identifier
|
| 100 |
+
|
| 101 |
+
Returns:
|
| 102 |
+
Dict with:
|
| 103 |
+
- image_files: List[str] - paths to images
|
| 104 |
+
- extrinsics: np.ndarray [N, 4, 4] - camera extrinsics (world-to-camera)
|
| 105 |
+
- intrinsics: np.ndarray [N, 3, 3] - camera intrinsics
|
| 106 |
+
- aux: Dict - auxiliary data (masks, GT paths, etc.)
|
| 107 |
+
"""
|
| 108 |
+
raise NotImplementedError
|
| 109 |
+
|
| 110 |
+
@abstractmethod
|
| 111 |
+
def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]:
|
| 112 |
+
"""
|
| 113 |
+
Evaluate 3D reconstruction quality against ground truth.
|
| 114 |
+
|
| 115 |
+
Args:
|
| 116 |
+
scene: Scene identifier
|
| 117 |
+
fuse_path: Path to fused point cloud (.ply)
|
| 118 |
+
|
| 119 |
+
Returns:
|
| 120 |
+
Dict with reconstruction metrics (e.g., acc, comp, overall)
|
| 121 |
+
"""
|
| 122 |
+
raise NotImplementedError
|
| 123 |
+
|
| 124 |
+
@abstractmethod
|
| 125 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 126 |
+
"""
|
| 127 |
+
Fuse per-view depth maps into a single point cloud.
|
| 128 |
+
|
| 129 |
+
Args:
|
| 130 |
+
scene: Scene identifier
|
| 131 |
+
result_path: Path to .npz file with predicted depths and poses
|
| 132 |
+
fuse_path: Output path for fused point cloud (.ply)
|
| 133 |
+
mode: Fusion mode ("recon_unposed" or "recon_posed")
|
| 134 |
+
"""
|
| 135 |
+
raise NotImplementedError
|
| 136 |
+
|
depth_anything_3/bench/datasets/__init__.py
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Benchmark dataset implementations.
|
| 17 |
+
|
| 18 |
+
Datasets are auto-registered via decorators when imported.
|
| 19 |
+
Add new dataset files here and they will be automatically discovered.
|
| 20 |
+
"""
|
| 21 |
+
|
depth_anything_3/bench/datasets/dtu.py
ADDED
|
@@ -0,0 +1,681 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
DTU Benchmark dataset implementation.
|
| 17 |
+
|
| 18 |
+
DTU is a multi-view stereo benchmark for 3D reconstruction evaluation.
|
| 19 |
+
Reference: https://roboimagedata.compute.dtu.dk/
|
| 20 |
+
|
| 21 |
+
Note: DepthAnything3 was never trained on any images from DTU.
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
import glob
|
| 25 |
+
import os
|
| 26 |
+
from typing import Dict as TDict, List
|
| 27 |
+
|
| 28 |
+
import numpy as np
|
| 29 |
+
import open3d as o3d
|
| 30 |
+
import torch
|
| 31 |
+
import torch.nn.functional as F
|
| 32 |
+
from addict import Dict
|
| 33 |
+
from PIL import Image
|
| 34 |
+
from plyfile import PlyData
|
| 35 |
+
from scipy.io import loadmat
|
| 36 |
+
from sklearn import neighbors as skln
|
| 37 |
+
from tqdm import tqdm
|
| 38 |
+
|
| 39 |
+
from depth_anything_3.bench.dataset import Dataset
|
| 40 |
+
from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY
|
| 41 |
+
from depth_anything_3.utils.constants import (
|
| 42 |
+
DTU_DIST_THRESH,
|
| 43 |
+
DTU_EVAL_DATA_ROOT,
|
| 44 |
+
DTU_MAX_POINTS,
|
| 45 |
+
DTU_NUM_CONSIST,
|
| 46 |
+
DTU_SCENES,
|
| 47 |
+
)
|
| 48 |
+
from depth_anything_3.utils.pose_align import align_poses_umeyama
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
@MV_REGISTRY.register(name="dtu")
|
| 52 |
+
@MONO_REGISTRY.register(name="dtu")
|
| 53 |
+
class DTU(Dataset):
|
| 54 |
+
"""
|
| 55 |
+
DTU Benchmark dataset wrapper for DepthAnything3 evaluation.
|
| 56 |
+
|
| 57 |
+
Supports:
|
| 58 |
+
- Camera pose estimation evaluation (AUC metrics)
|
| 59 |
+
- 3D reconstruction evaluation (accuracy, completeness, overall)
|
| 60 |
+
- Point cloud fusion from depth maps
|
| 61 |
+
|
| 62 |
+
The dataset uses MVSNet evaluation protocol:
|
| 63 |
+
https://drive.google.com/file/d/1rX0EXlUL4prRxrRu2DgLJv2j7-tpUD4D/view
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
data_root = DTU_EVAL_DATA_ROOT
|
| 67 |
+
SCENES = DTU_SCENES
|
| 68 |
+
|
| 69 |
+
# Evaluation/triangulation hyperparameters from constants
|
| 70 |
+
dist_thresh = DTU_DIST_THRESH
|
| 71 |
+
num_consist = DTU_NUM_CONSIST
|
| 72 |
+
|
| 73 |
+
# ------------------------------
|
| 74 |
+
# Public API
|
| 75 |
+
# ------------------------------
|
| 76 |
+
|
| 77 |
+
def read_cam_file(self, filename: str) -> tuple:
|
| 78 |
+
"""
|
| 79 |
+
Read DTU camera file containing extrinsics and intrinsics.
|
| 80 |
+
|
| 81 |
+
Args:
|
| 82 |
+
filename: Path to camera text file
|
| 83 |
+
|
| 84 |
+
Returns:
|
| 85 |
+
Tuple of (intrinsics [3,3], extrinsics [4,4])
|
| 86 |
+
"""
|
| 87 |
+
with open(filename) as f:
|
| 88 |
+
lines = [line.rstrip() for line in f.readlines()]
|
| 89 |
+
extrinsics = np.fromstring(" ".join(lines[1:5]), dtype=np.float32, sep=" ").reshape((4, 4))
|
| 90 |
+
intrinsics = np.fromstring(" ".join(lines[7:10]), dtype=np.float32, sep=" ").reshape((3, 3))
|
| 91 |
+
return intrinsics, extrinsics
|
| 92 |
+
|
| 93 |
+
def get_data(self, scene: str) -> Dict:
|
| 94 |
+
"""
|
| 95 |
+
Collect per-view image paths, intrinsics/extrinsics, and GT masks.
|
| 96 |
+
|
| 97 |
+
Args:
|
| 98 |
+
scene: Scene identifier (e.g., "scan1")
|
| 99 |
+
|
| 100 |
+
Returns:
|
| 101 |
+
Dict with:
|
| 102 |
+
- image_files: List[str] - paths to images
|
| 103 |
+
- extrinsics: np.ndarray [N, 4, 4]
|
| 104 |
+
- intrinsics: np.ndarray [N, 3, 3]
|
| 105 |
+
- aux.mask_files: List[str] - paths to depth masks
|
| 106 |
+
"""
|
| 107 |
+
rgb_folder = os.path.join(self.data_root, "Rectified", scene)
|
| 108 |
+
camera_folder = os.path.join(self.data_root, "Cameras")
|
| 109 |
+
|
| 110 |
+
files = sorted(glob.glob(os.path.join(rgb_folder, "*.png")))
|
| 111 |
+
# Reorder: place index 33 first (reference view convention)
|
| 112 |
+
files = [files[33]] + files[:33] + files[34:]
|
| 113 |
+
|
| 114 |
+
out = Dict(
|
| 115 |
+
{
|
| 116 |
+
"image_files": files,
|
| 117 |
+
"extrinsics": [],
|
| 118 |
+
"intrinsics": [],
|
| 119 |
+
"aux": Dict({"mask_files": []}),
|
| 120 |
+
}
|
| 121 |
+
)
|
| 122 |
+
|
| 123 |
+
for rgb_file in files:
|
| 124 |
+
basename = os.path.basename(rgb_file)
|
| 125 |
+
file_idx = basename.split("_")[1]
|
| 126 |
+
cam_idx = depth_idx = int(file_idx) - 1
|
| 127 |
+
|
| 128 |
+
mask_file = self._depth_mask_path(scene, depth_idx)
|
| 129 |
+
proj_mat_filename = os.path.join(camera_folder, f"{cam_idx:0>8}_cam.txt")
|
| 130 |
+
|
| 131 |
+
ixt, ext = self.read_cam_file(proj_mat_filename)
|
| 132 |
+
out.extrinsics.append(ext)
|
| 133 |
+
out.intrinsics.append(ixt)
|
| 134 |
+
out.aux.mask_files.append(mask_file)
|
| 135 |
+
|
| 136 |
+
out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32)
|
| 137 |
+
out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32)
|
| 138 |
+
return out
|
| 139 |
+
|
| 140 |
+
def get_3dgtpath(self, scene: str) -> str:
|
| 141 |
+
"""Get path to ground truth point cloud for a scene."""
|
| 142 |
+
scene_id = int(scene[4:])
|
| 143 |
+
return os.path.join(self.data_root, f"Points/stl/stl{scene_id:03}_total.ply")
|
| 144 |
+
|
| 145 |
+
def eval3d(self, scene: str, fuse_path: str, use_gpu: bool = False) -> TDict[str, float]:
|
| 146 |
+
"""
|
| 147 |
+
Evaluate fused point cloud against DTU GT with ObsMask/Plane.
|
| 148 |
+
|
| 149 |
+
Args:
|
| 150 |
+
scene: Scene identifier
|
| 151 |
+
fuse_path: Path to fused point cloud
|
| 152 |
+
use_gpu: If True, use GPU-accelerated distance computation (faster but may have minor numerical differences)
|
| 153 |
+
|
| 154 |
+
Returns:
|
| 155 |
+
Dict with metrics: {"comp": float, "acc": float, "overall": float}
|
| 156 |
+
"""
|
| 157 |
+
scene_id = int(scene[4:])
|
| 158 |
+
gt_ply = os.path.join(self.data_root, f"Points/stl/stl{scene_id:03}_total.ply")
|
| 159 |
+
mask_file = os.path.join(
|
| 160 |
+
self.data_root, f"SampleSet/mvs_data/ObsMask/ObsMask{scene_id}_10.mat"
|
| 161 |
+
)
|
| 162 |
+
plane_file = os.path.join(
|
| 163 |
+
self.data_root, f"SampleSet/mvs_data/ObsMask/Plane{scene_id}.mat"
|
| 164 |
+
)
|
| 165 |
+
result = self._evaluate_reconstruction(
|
| 166 |
+
scene, fuse_path, gt_ply, mask_file, plane_file, use_gpu=use_gpu
|
| 167 |
+
)
|
| 168 |
+
return {"comp": result[0], "acc": result[1], "overall": result[2]}
|
| 169 |
+
|
| 170 |
+
def load_masks(self, mask_files: List[str]) -> np.ndarray:
|
| 171 |
+
"""
|
| 172 |
+
Load DTU depth validity masks.
|
| 173 |
+
|
| 174 |
+
Args:
|
| 175 |
+
mask_files: List of paths to mask images
|
| 176 |
+
|
| 177 |
+
Returns:
|
| 178 |
+
Boolean array [N, H, W] indicating valid depth regions
|
| 179 |
+
"""
|
| 180 |
+
masks = []
|
| 181 |
+
for mask_file in mask_files:
|
| 182 |
+
mask = Image.open(mask_file)
|
| 183 |
+
mask = np.array(mask, dtype=np.float32)
|
| 184 |
+
masks.append(mask > 10)
|
| 185 |
+
return np.asarray(masks)
|
| 186 |
+
|
| 187 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 188 |
+
"""
|
| 189 |
+
Fuse per-view depths into a point cloud and save to PLY.
|
| 190 |
+
|
| 191 |
+
Args:
|
| 192 |
+
scene: Scene identifier (e.g., "scan114")
|
| 193 |
+
result_path: Path to npz file containing predicted depths/poses
|
| 194 |
+
fuse_path: Output path for fused point cloud (.ply)
|
| 195 |
+
mode: "recon_unposed" or "recon_posed"
|
| 196 |
+
"""
|
| 197 |
+
gt_data = self.get_data(scene)
|
| 198 |
+
pred_data = Dict({k: v for k, v in np.load(result_path).items()})
|
| 199 |
+
masks = self.load_masks(gt_data.aux.mask_files)
|
| 200 |
+
|
| 201 |
+
if mode == "recon_unposed":
|
| 202 |
+
depths, intrinsics, extrinsics = self._prep_unposed(pred_data, gt_data, masks)
|
| 203 |
+
elif mode == "recon_posed":
|
| 204 |
+
depths, intrinsics, extrinsics = self._prep_posed(pred_data, gt_data, masks)
|
| 205 |
+
else:
|
| 206 |
+
raise ValueError(f"Invalid mode: {mode}")
|
| 207 |
+
|
| 208 |
+
proj_mat = self._build_proj_mats(intrinsics, extrinsics)
|
| 209 |
+
|
| 210 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 211 |
+
dtype = torch.float32
|
| 212 |
+
depths_t = torch.from_numpy(depths).to(device=device, dtype=dtype).unsqueeze(1)
|
| 213 |
+
proj_t = torch.from_numpy(proj_mat).to(device=device, dtype=dtype)
|
| 214 |
+
height, width = depths_t.shape[-2:]
|
| 215 |
+
|
| 216 |
+
points: List[np.ndarray] = []
|
| 217 |
+
for idx in range(len(gt_data.image_files)):
|
| 218 |
+
if mode == "recon_unposed":
|
| 219 |
+
# Simple unfiltered back-projection per frame
|
| 220 |
+
cur_p_pcd = self._generate_points_from_depth(
|
| 221 |
+
depths_t[idx : idx + 1], proj_t[idx : idx + 1]
|
| 222 |
+
)
|
| 223 |
+
mask = (depths_t[idx : idx + 1] > 0.001).squeeze()
|
| 224 |
+
cur_p_pcd = cur_p_pcd[:, :, mask]
|
| 225 |
+
no_filter_pc = cur_p_pcd.squeeze(0).permute(1, 0).cpu().numpy()
|
| 226 |
+
points.append(no_filter_pc)
|
| 227 |
+
else: # recon_posed
|
| 228 |
+
final_pc = self._fuse_consistent_points(depths_t, proj_t, idx, height, width)
|
| 229 |
+
points.append(final_pc)
|
| 230 |
+
|
| 231 |
+
# Concatenate and optionally downsample to hard cap
|
| 232 |
+
points_np = np.concatenate(points, axis=0)
|
| 233 |
+
points_np = self._cap_points(points_np, max_points=DTU_MAX_POINTS)
|
| 234 |
+
|
| 235 |
+
os.makedirs(os.path.dirname(fuse_path), exist_ok=True)
|
| 236 |
+
pcd = o3d.geometry.PointCloud()
|
| 237 |
+
pcd.points = o3d.utility.Vector3dVector(points_np)
|
| 238 |
+
o3d.io.write_point_cloud(fuse_path, pcd)
|
| 239 |
+
|
| 240 |
+
# ------------------------------
|
| 241 |
+
# Geometry helpers
|
| 242 |
+
# ------------------------------
|
| 243 |
+
|
| 244 |
+
def _generate_points_from_depth(
|
| 245 |
+
self, depth: torch.Tensor, proj: torch.Tensor
|
| 246 |
+
) -> torch.Tensor:
|
| 247 |
+
"""
|
| 248 |
+
Back-project depth map into 3D world coordinates.
|
| 249 |
+
|
| 250 |
+
Args:
|
| 251 |
+
depth: Depth tensor [B, 1, H, W]
|
| 252 |
+
proj: Projection matrix [B, 4, 4] = [[K@R, K@t], [0,0,0,1]]
|
| 253 |
+
|
| 254 |
+
Returns:
|
| 255 |
+
Point cloud tensor [B, 3, H, W]
|
| 256 |
+
"""
|
| 257 |
+
batch, height, width = depth.shape[0], depth.shape[2], depth.shape[3]
|
| 258 |
+
inv_proj = torch.inverse(proj)
|
| 259 |
+
rot = inv_proj[:, :3, :3]
|
| 260 |
+
trans = inv_proj[:, :3, 3:4]
|
| 261 |
+
|
| 262 |
+
y, x = torch.meshgrid(
|
| 263 |
+
[
|
| 264 |
+
torch.arange(0, height, dtype=torch.float32, device=depth.device),
|
| 265 |
+
torch.arange(0, width, dtype=torch.float32, device=depth.device),
|
| 266 |
+
],
|
| 267 |
+
indexing="ij",
|
| 268 |
+
)
|
| 269 |
+
y, x = y.contiguous(), x.contiguous()
|
| 270 |
+
y, x = y.view(height * width), x.view(height * width)
|
| 271 |
+
xyz = torch.stack((x, y, torch.ones_like(x)))
|
| 272 |
+
xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1)
|
| 273 |
+
rot_xyz = torch.matmul(rot, xyz)
|
| 274 |
+
rot_depth_xyz = rot_xyz * depth.view(batch, 1, -1)
|
| 275 |
+
proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1)
|
| 276 |
+
return proj_xyz.view(batch, 3, height, width)
|
| 277 |
+
|
| 278 |
+
def _homo_warping(
|
| 279 |
+
self,
|
| 280 |
+
src_fea: torch.Tensor,
|
| 281 |
+
src_proj: torch.Tensor,
|
| 282 |
+
ref_proj: torch.Tensor,
|
| 283 |
+
depth_values: torch.Tensor,
|
| 284 |
+
) -> torch.Tensor:
|
| 285 |
+
"""
|
| 286 |
+
Homography warping for multi-view consistency checking.
|
| 287 |
+
|
| 288 |
+
Args:
|
| 289 |
+
src_fea: Source features [B, C, H, W]
|
| 290 |
+
src_proj: Source projection [B, 4, 4]
|
| 291 |
+
ref_proj: Reference projection [B, 4, 4]
|
| 292 |
+
depth_values: Depth values [B, Ndepth] or [B, Ndepth, H, W]
|
| 293 |
+
|
| 294 |
+
Returns:
|
| 295 |
+
Warped features [B, C, H, W]
|
| 296 |
+
"""
|
| 297 |
+
batch, channels = src_fea.shape[0], src_fea.shape[1]
|
| 298 |
+
height, width = src_fea.shape[2], src_fea.shape[3]
|
| 299 |
+
|
| 300 |
+
with torch.no_grad():
|
| 301 |
+
proj = torch.matmul(src_proj, torch.inverse(ref_proj))
|
| 302 |
+
rot = proj[:, :3, :3]
|
| 303 |
+
trans = proj[:, :3, 3:4]
|
| 304 |
+
|
| 305 |
+
y, x = torch.meshgrid(
|
| 306 |
+
[
|
| 307 |
+
torch.arange(0, height, dtype=torch.float32, device=src_fea.device),
|
| 308 |
+
torch.arange(0, width, dtype=torch.float32, device=src_fea.device),
|
| 309 |
+
],
|
| 310 |
+
indexing="ij",
|
| 311 |
+
)
|
| 312 |
+
y, x = y.contiguous(), x.contiguous()
|
| 313 |
+
y, x = y.view(height * width), x.view(height * width)
|
| 314 |
+
xyz = torch.stack((x, y, torch.ones_like(x)))
|
| 315 |
+
xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1)
|
| 316 |
+
rot_xyz = torch.matmul(rot, xyz)
|
| 317 |
+
|
| 318 |
+
rot_depth_xyz = rot_xyz.unsqueeze(2) * depth_values.view(-1, 1, 1, height * width)
|
| 319 |
+
proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1)
|
| 320 |
+
proj_xy = proj_xyz[:, :2, :, :] / proj_xyz[:, 2:3, :, :]
|
| 321 |
+
proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1
|
| 322 |
+
proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1
|
| 323 |
+
grid = torch.stack((proj_x_normalized, proj_y_normalized), dim=3)
|
| 324 |
+
|
| 325 |
+
warped_src_fea = F.grid_sample(
|
| 326 |
+
src_fea,
|
| 327 |
+
grid.view(batch, height, width, 2),
|
| 328 |
+
mode="bilinear",
|
| 329 |
+
padding_mode="zeros",
|
| 330 |
+
align_corners=True,
|
| 331 |
+
)
|
| 332 |
+
return warped_src_fea.view(batch, channels, height, width)
|
| 333 |
+
|
| 334 |
+
def _filter_depth(
|
| 335 |
+
self,
|
| 336 |
+
ref_depth: torch.Tensor,
|
| 337 |
+
src_depths: torch.Tensor,
|
| 338 |
+
ref_proj: torch.Tensor,
|
| 339 |
+
src_projs: torch.Tensor,
|
| 340 |
+
) -> tuple:
|
| 341 |
+
"""
|
| 342 |
+
Compute geometric consistency between reference and source depths.
|
| 343 |
+
|
| 344 |
+
Args:
|
| 345 |
+
ref_depth: Reference depth [1, 1, H, W]
|
| 346 |
+
src_depths: Source depths [B, 1, H, W]
|
| 347 |
+
ref_proj: Reference projection [1, 4, 4]
|
| 348 |
+
src_projs: Source projections [B, 4, 4]
|
| 349 |
+
|
| 350 |
+
Returns:
|
| 351 |
+
Tuple of (ref_pc, aligned_pcs, dist)
|
| 352 |
+
"""
|
| 353 |
+
ref_pc = self._generate_points_from_depth(ref_depth, ref_proj)
|
| 354 |
+
src_pcs = self._generate_points_from_depth(src_depths, src_projs)
|
| 355 |
+
aligned_pcs = self._homo_warping(src_pcs, src_projs, ref_proj, ref_depth)
|
| 356 |
+
x_2 = (ref_pc[:, 0] - aligned_pcs[:, 0]) ** 2
|
| 357 |
+
y_2 = (ref_pc[:, 1] - aligned_pcs[:, 1]) ** 2
|
| 358 |
+
z_2 = (ref_pc[:, 2] - aligned_pcs[:, 2]) ** 2
|
| 359 |
+
dist = torch.sqrt(x_2 + y_2 + z_2).unsqueeze(1)
|
| 360 |
+
return ref_pc, aligned_pcs, dist
|
| 361 |
+
|
| 362 |
+
def _extract_points(
|
| 363 |
+
self, pc: torch.Tensor, mask: torch.Tensor, rgb: np.ndarray = None
|
| 364 |
+
) -> np.ndarray:
|
| 365 |
+
"""Extract masked points from a dense grid."""
|
| 366 |
+
pc = pc.cpu().numpy()
|
| 367 |
+
mask = mask.cpu().numpy().reshape(-1)
|
| 368 |
+
pc = pc.reshape(-1, 3)
|
| 369 |
+
points = pc[np.where(mask)]
|
| 370 |
+
if rgb is not None:
|
| 371 |
+
rgb = rgb.reshape(-1, 3)
|
| 372 |
+
colors = rgb[np.where(mask)]
|
| 373 |
+
return np.concatenate([points, colors], axis=1)
|
| 374 |
+
return points
|
| 375 |
+
|
| 376 |
+
# ------------------------------
|
| 377 |
+
# 3D Reconstruction Evaluation
|
| 378 |
+
# ------------------------------
|
| 379 |
+
|
| 380 |
+
def _evaluate_reconstruction(
|
| 381 |
+
self,
|
| 382 |
+
scanid: str,
|
| 383 |
+
pred_ply: str,
|
| 384 |
+
gt_ply: str,
|
| 385 |
+
mask_file: str,
|
| 386 |
+
plane_file: str,
|
| 387 |
+
down_dense: float = 0.2,
|
| 388 |
+
patch: int = 60,
|
| 389 |
+
max_dist: int = 20,
|
| 390 |
+
use_gpu: bool = False,
|
| 391 |
+
) -> tuple:
|
| 392 |
+
"""
|
| 393 |
+
Compute accuracy, completeness, and overall metrics for one scan.
|
| 394 |
+
|
| 395 |
+
Args:
|
| 396 |
+
scanid: Scan identifier
|
| 397 |
+
pred_ply: Predicted point cloud path or array
|
| 398 |
+
gt_ply: Ground truth point cloud path or array
|
| 399 |
+
mask_file: ObsMask file path
|
| 400 |
+
plane_file: Plane file path
|
| 401 |
+
down_dense: Downsample density (min distance between points)
|
| 402 |
+
patch: Patch size for boundary
|
| 403 |
+
max_dist: Outlier threshold in mm
|
| 404 |
+
use_gpu: If True, use GPU-accelerated distance computation
|
| 405 |
+
|
| 406 |
+
Returns:
|
| 407 |
+
Tuple of (mean_d2s, mean_s2d, overall)
|
| 408 |
+
"""
|
| 409 |
+
thresh = down_dense
|
| 410 |
+
|
| 411 |
+
# Load and downsample predicted point cloud
|
| 412 |
+
data_pcd = self._read_ply(pred_ply) if isinstance(pred_ply, str) else pred_ply
|
| 413 |
+
# Use fixed seed for reproducibility
|
| 414 |
+
shuffle_rng = np.random.default_rng(seed=42)
|
| 415 |
+
shuffle_rng.shuffle(data_pcd, axis=0)
|
| 416 |
+
|
| 417 |
+
# Downsample point cloud
|
| 418 |
+
nn_engine = skln.NearestNeighbors(
|
| 419 |
+
n_neighbors=1, radius=thresh, algorithm="kd_tree", n_jobs=-1
|
| 420 |
+
)
|
| 421 |
+
nn_engine.fit(data_pcd)
|
| 422 |
+
rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False)
|
| 423 |
+
mask = np.ones(data_pcd.shape[0], dtype=np.bool_)
|
| 424 |
+
for curr, idxs in enumerate(rnn_idxs):
|
| 425 |
+
if mask[curr]:
|
| 426 |
+
mask[idxs] = 0
|
| 427 |
+
mask[curr] = 1
|
| 428 |
+
data_down = data_pcd[mask]
|
| 429 |
+
|
| 430 |
+
# Restrict to observed volume (ObsMask)
|
| 431 |
+
obs_mask_file = loadmat(mask_file)
|
| 432 |
+
ObsMask, BB, Res = (obs_mask_file[attr] for attr in ["ObsMask", "BB", "Res"])
|
| 433 |
+
BB = BB.astype(np.float32)
|
| 434 |
+
|
| 435 |
+
inbound = ((data_down >= BB[:1] - patch) & (data_down < BB[1:] + patch * 2)).sum(
|
| 436 |
+
axis=-1
|
| 437 |
+
) == 3
|
| 438 |
+
data_in = data_down[inbound]
|
| 439 |
+
|
| 440 |
+
data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32)
|
| 441 |
+
grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum(
|
| 442 |
+
axis=-1
|
| 443 |
+
) == 3
|
| 444 |
+
data_grid_in = data_grid[grid_inbound]
|
| 445 |
+
in_obs = ObsMask[data_grid_in[:, 0], data_grid_in[:, 1], data_grid_in[:, 2]].astype(
|
| 446 |
+
np.bool_
|
| 447 |
+
)
|
| 448 |
+
data_in_obs = data_in[grid_inbound][in_obs]
|
| 449 |
+
|
| 450 |
+
# Compute accuracy (pred -> GT) and completeness (GT -> pred)
|
| 451 |
+
stl = self._read_ply(gt_ply) if isinstance(gt_ply, str) else gt_ply
|
| 452 |
+
|
| 453 |
+
if use_gpu and torch.cuda.is_available():
|
| 454 |
+
# GPU-accelerated distance computation
|
| 455 |
+
mean_d2s = self._knn_dist_gpu(data_in_obs, stl, max_dist)
|
| 456 |
+
else:
|
| 457 |
+
# CPU version (original, for exact reproduction)
|
| 458 |
+
nn_engine.fit(stl)
|
| 459 |
+
dist_d2s, _ = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True)
|
| 460 |
+
mean_d2s = dist_d2s[dist_d2s < max_dist].mean()
|
| 461 |
+
|
| 462 |
+
ground_plane = loadmat(plane_file)["P"]
|
| 463 |
+
stl_hom = np.concatenate([stl, np.ones_like(stl[:, :1])], -1)
|
| 464 |
+
above = (ground_plane.reshape((1, 4)) * stl_hom).sum(-1) > 0
|
| 465 |
+
stl_above = stl[above]
|
| 466 |
+
|
| 467 |
+
if use_gpu and torch.cuda.is_available():
|
| 468 |
+
# GPU-accelerated distance computation
|
| 469 |
+
mean_s2d = self._knn_dist_gpu(stl_above, data_in, max_dist)
|
| 470 |
+
else:
|
| 471 |
+
# CPU version (original, for exact reproduction)
|
| 472 |
+
nn_engine.fit(data_in)
|
| 473 |
+
dist_s2d, _ = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True)
|
| 474 |
+
mean_s2d = dist_s2d[dist_s2d < max_dist].mean()
|
| 475 |
+
|
| 476 |
+
overall = (mean_d2s + mean_s2d) / 2
|
| 477 |
+
return mean_d2s, mean_s2d, overall
|
| 478 |
+
|
| 479 |
+
def _knn_dist_gpu(
|
| 480 |
+
self,
|
| 481 |
+
query: np.ndarray,
|
| 482 |
+
target: np.ndarray,
|
| 483 |
+
max_dist: float,
|
| 484 |
+
batch_size: int = 8192,
|
| 485 |
+
target_batch_size: int = 50000,
|
| 486 |
+
) -> float:
|
| 487 |
+
"""
|
| 488 |
+
GPU-accelerated nearest neighbor distance computation.
|
| 489 |
+
|
| 490 |
+
Args:
|
| 491 |
+
query: Query points [N, 3]
|
| 492 |
+
target: Target points [M, 3]
|
| 493 |
+
max_dist: Outlier threshold
|
| 494 |
+
batch_size: Batch size for query to avoid OOM (tuned for 16GB GPU)
|
| 495 |
+
target_batch_size: Batch size for target to avoid OOM
|
| 496 |
+
|
| 497 |
+
Returns:
|
| 498 |
+
Mean distance (excluding outliers)
|
| 499 |
+
"""
|
| 500 |
+
device = torch.device("cuda")
|
| 501 |
+
|
| 502 |
+
all_min_dists = []
|
| 503 |
+
n_query_batches = (len(query) + batch_size - 1) // batch_size
|
| 504 |
+
n_target_batches = (len(target) + target_batch_size - 1) // target_batch_size
|
| 505 |
+
|
| 506 |
+
# Pre-load target batches to GPU to avoid repeated transfers
|
| 507 |
+
# Memory: ~50000 pts * 3 coords * 4 bytes * n_batches
|
| 508 |
+
target_batches = []
|
| 509 |
+
for j in range(0, len(target), target_batch_size):
|
| 510 |
+
target_batch = target[j : j + target_batch_size]
|
| 511 |
+
target_t = torch.from_numpy(target_batch).float().to(device)
|
| 512 |
+
target_batches.append(target_t)
|
| 513 |
+
|
| 514 |
+
with tqdm(total=n_query_batches, desc=" GPU KNN", leave=False, ncols=100) as pbar:
|
| 515 |
+
for i in range(0, len(query), batch_size):
|
| 516 |
+
batch = query[i : i + batch_size]
|
| 517 |
+
query_t = torch.from_numpy(batch).float().to(device)
|
| 518 |
+
|
| 519 |
+
# Compute distances to all target batches
|
| 520 |
+
# Memory peak: query_batch × target_batch_size × 4 bytes
|
| 521 |
+
# = 8192 × 50000 × 4 = ~1.6 GB per cdist call
|
| 522 |
+
batch_min_dists = []
|
| 523 |
+
for target_t in target_batches:
|
| 524 |
+
dists = torch.cdist(query_t, target_t)
|
| 525 |
+
batch_min_dists.append(dists.min(dim=1).values)
|
| 526 |
+
del dists # Free immediately
|
| 527 |
+
|
| 528 |
+
# Get minimum distance across all target batches
|
| 529 |
+
min_dists = torch.stack(batch_min_dists, dim=1).min(dim=1).values
|
| 530 |
+
all_min_dists.append(min_dists.cpu().numpy())
|
| 531 |
+
|
| 532 |
+
del query_t, min_dists, batch_min_dists
|
| 533 |
+
pbar.update(1)
|
| 534 |
+
|
| 535 |
+
# Clean up target batches
|
| 536 |
+
for target_t in target_batches:
|
| 537 |
+
del target_t
|
| 538 |
+
torch.cuda.empty_cache()
|
| 539 |
+
|
| 540 |
+
all_min_dists = np.concatenate(all_min_dists)
|
| 541 |
+
return all_min_dists[all_min_dists < max_dist].mean()
|
| 542 |
+
|
| 543 |
+
def _read_ply(self, file: str) -> np.ndarray:
|
| 544 |
+
"""Read point cloud from PLY file."""
|
| 545 |
+
data = PlyData.read(file)
|
| 546 |
+
vertex = data["vertex"]
|
| 547 |
+
return np.stack([vertex["x"], vertex["y"], vertex["z"]], axis=-1)
|
| 548 |
+
|
| 549 |
+
# ------------------------------
|
| 550 |
+
# Private helpers
|
| 551 |
+
# ------------------------------
|
| 552 |
+
|
| 553 |
+
def _depth_mask_path(self, scene: str, depth_idx: int) -> str:
|
| 554 |
+
"""Get path to depth mask for a scene and frame."""
|
| 555 |
+
return os.path.join(
|
| 556 |
+
self.data_root, "depth_raw", "Depths", scene, f"depth_visual_{depth_idx:04d}.png"
|
| 557 |
+
)
|
| 558 |
+
|
| 559 |
+
def _prep_unposed(
|
| 560 |
+
self, pred_data: Dict, gt_data: Dict, masks: np.ndarray
|
| 561 |
+
) -> tuple:
|
| 562 |
+
"""
|
| 563 |
+
Prepare depths/intrinsics/extrinsics for recon_unposed mode.
|
| 564 |
+
|
| 565 |
+
Applies Umeyama scale, rescales intrinsics if depth resolution differs,
|
| 566 |
+
and zeroes invalid-mask depths (nearest interpolation as in paper).
|
| 567 |
+
"""
|
| 568 |
+
_, _, scale, extrinsics = align_poses_umeyama(
|
| 569 |
+
gt_data.extrinsics.copy(),
|
| 570 |
+
pred_data.extrinsics.copy(),
|
| 571 |
+
ransac=True,
|
| 572 |
+
return_aligned=True,
|
| 573 |
+
random_state=42,
|
| 574 |
+
)
|
| 575 |
+
depths = pred_data.depth * scale
|
| 576 |
+
intrinsics = pred_data.intrinsics.copy()
|
| 577 |
+
|
| 578 |
+
if depths.shape[-2:] != masks.shape[-2:]:
|
| 579 |
+
# When resizing depths to mask size, adjust intrinsics accordingly
|
| 580 |
+
sx = masks.shape[-1] / depths.shape[-1]
|
| 581 |
+
sy = masks.shape[-2] / depths.shape[-2]
|
| 582 |
+
intrinsics[:, 0:1] *= sx
|
| 583 |
+
intrinsics[:, 1:2] *= sy
|
| 584 |
+
depths = F.interpolate(
|
| 585 |
+
torch.from_numpy(depths)[None].float(),
|
| 586 |
+
size=(masks.shape[-2], masks.shape[-1]),
|
| 587 |
+
mode="nearest",
|
| 588 |
+
)[0].numpy()
|
| 589 |
+
depths[masks == False] = 0.0 # noqa: E712
|
| 590 |
+
|
| 591 |
+
return depths, intrinsics, extrinsics
|
| 592 |
+
|
| 593 |
+
def _prep_posed(
|
| 594 |
+
self, pred_data: Dict, gt_data: Dict, masks: np.ndarray
|
| 595 |
+
) -> tuple:
|
| 596 |
+
"""
|
| 597 |
+
Prepare depths/intrinsics/extrinsics for recon_posed mode.
|
| 598 |
+
|
| 599 |
+
Uses GT intrinsics/extrinsics but aligns scale via Umeyama.
|
| 600 |
+
Same mask order as other datasets: mask BEFORE scale.
|
| 601 |
+
"""
|
| 602 |
+
_, _, scale, _ = align_poses_umeyama(
|
| 603 |
+
gt_data.extrinsics.copy(),
|
| 604 |
+
pred_data.extrinsics.copy(),
|
| 605 |
+
ransac=True,
|
| 606 |
+
return_aligned=True,
|
| 607 |
+
random_state=42,
|
| 608 |
+
)
|
| 609 |
+
depths = pred_data.depth.copy()
|
| 610 |
+
intrinsics = gt_data.intrinsics.copy()
|
| 611 |
+
extrinsics = gt_data.extrinsics.copy()
|
| 612 |
+
|
| 613 |
+
if depths.shape[-2:] != masks.shape[-2:]:
|
| 614 |
+
depths = F.interpolate(
|
| 615 |
+
torch.from_numpy(depths)[None].float(),
|
| 616 |
+
size=(masks.shape[-2], masks.shape[-1]),
|
| 617 |
+
mode="nearest",
|
| 618 |
+
)[0].numpy()
|
| 619 |
+
|
| 620 |
+
# Mask BEFORE scale (same as other datasets)
|
| 621 |
+
depths[masks == False] = 0.0 # noqa: E712
|
| 622 |
+
depths = depths * scale
|
| 623 |
+
|
| 624 |
+
return depths, intrinsics, extrinsics
|
| 625 |
+
|
| 626 |
+
def _build_proj_mats(
|
| 627 |
+
self, intrinsics: np.ndarray, extrinsics: np.ndarray
|
| 628 |
+
) -> np.ndarray:
|
| 629 |
+
"""Compute per-view 4x4 projection matrices from K and [R|t]."""
|
| 630 |
+
proj_mat_list = []
|
| 631 |
+
for i in range(len(intrinsics)):
|
| 632 |
+
proj_mat = np.eye(4, dtype=np.float32)
|
| 633 |
+
proj_mat[:3, :4] = np.dot(intrinsics[i], extrinsics[i][:3])
|
| 634 |
+
proj_mat_list.append(proj_mat)
|
| 635 |
+
return np.stack(proj_mat_list, axis=0)
|
| 636 |
+
|
| 637 |
+
def _fuse_consistent_points(
|
| 638 |
+
self,
|
| 639 |
+
depths_t: torch.Tensor,
|
| 640 |
+
proj_t: torch.Tensor,
|
| 641 |
+
idx: int,
|
| 642 |
+
H: int,
|
| 643 |
+
W: int,
|
| 644 |
+
) -> np.ndarray:
|
| 645 |
+
"""Fuse points consistent across multiple source views for a reference index."""
|
| 646 |
+
device, dtype = depths_t.device, depths_t.dtype
|
| 647 |
+
pc_buff = torch.zeros((3, H, W), device=device, dtype=dtype)
|
| 648 |
+
val_cnt = torch.zeros((1, H, W), device=device, dtype=dtype)
|
| 649 |
+
|
| 650 |
+
j = 0
|
| 651 |
+
batch_size = 20
|
| 652 |
+
tot_frame = depths_t.shape[0]
|
| 653 |
+
while True:
|
| 654 |
+
ref_pc, pcs, dist = self._filter_depth(
|
| 655 |
+
ref_depth=depths_t[idx : idx + 1],
|
| 656 |
+
src_depths=depths_t[j : min(j + batch_size, tot_frame)],
|
| 657 |
+
ref_proj=proj_t[idx : idx + 1],
|
| 658 |
+
src_projs=proj_t[j : min(j + batch_size, tot_frame)],
|
| 659 |
+
)
|
| 660 |
+
masks = (dist < self.dist_thresh).float()
|
| 661 |
+
masked_pc = pcs * masks
|
| 662 |
+
pc_buff += masked_pc.sum(dim=0, keepdim=False)
|
| 663 |
+
val_cnt += masks.sum(dim=0, keepdim=False)
|
| 664 |
+
j += batch_size
|
| 665 |
+
if j >= tot_frame:
|
| 666 |
+
break
|
| 667 |
+
|
| 668 |
+
final_mask = (val_cnt >= self.num_consist).squeeze(0)
|
| 669 |
+
avg_points = torch.div(pc_buff, val_cnt).permute(1, 2, 0)
|
| 670 |
+
final_pc = self._extract_points(avg_points, final_mask)
|
| 671 |
+
return final_pc
|
| 672 |
+
|
| 673 |
+
def _cap_points(self, points: np.ndarray, max_points: int) -> np.ndarray:
|
| 674 |
+
"""Downsample points if exceeding max count."""
|
| 675 |
+
if len(points) <= max_points:
|
| 676 |
+
return points
|
| 677 |
+
# Use fixed seed for reproducibility
|
| 678 |
+
rng = np.random.default_rng(seed=42)
|
| 679 |
+
random_idx = rng.choice(len(points), max_points, replace=False)
|
| 680 |
+
return points[random_idx]
|
| 681 |
+
|
depth_anything_3/bench/datasets/dtu64.py
ADDED
|
@@ -0,0 +1,182 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
DTU-64 Dataset implementation for POSE EVALUATION ONLY.
|
| 17 |
+
|
| 18 |
+
This is a subset of DTU with 64 images per scene, specifically designed for
|
| 19 |
+
camera pose estimation evaluation. It does NOT support 3D reconstruction.
|
| 20 |
+
|
| 21 |
+
Note: GT depth loading is not implemented as it's not needed for pose evaluation.
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
import glob
|
| 25 |
+
import os
|
| 26 |
+
from typing import Dict as TDict
|
| 27 |
+
|
| 28 |
+
import numpy as np
|
| 29 |
+
from addict import Dict
|
| 30 |
+
|
| 31 |
+
from depth_anything_3.bench.dataset import Dataset
|
| 32 |
+
from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY
|
| 33 |
+
from depth_anything_3.utils.constants import (
|
| 34 |
+
DTU64_CAMERA_ROOT,
|
| 35 |
+
DTU64_EVAL_DATA_ROOT,
|
| 36 |
+
DTU64_SCENES,
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
@MV_REGISTRY.register(name="dtu64")
|
| 41 |
+
@MONO_REGISTRY.register(name="dtu64")
|
| 42 |
+
class DTU64(Dataset):
|
| 43 |
+
"""
|
| 44 |
+
DTU-64 Dataset wrapper for DepthAnything3 POSE EVALUATION ONLY.
|
| 45 |
+
|
| 46 |
+
This dataset is a subset of DTU with 64 images per scene.
|
| 47 |
+
It is specifically designed for camera pose estimation evaluation
|
| 48 |
+
and does NOT support 3D reconstruction evaluation.
|
| 49 |
+
|
| 50 |
+
Dataset structure:
|
| 51 |
+
DTU/scans/
|
| 52 |
+
├── {scene}/
|
| 53 |
+
│ └── image/ # RGB images (64 per scene)
|
| 54 |
+
└── Cameras/
|
| 55 |
+
└── {idx}_cam.txt # Camera parameters
|
| 56 |
+
|
| 57 |
+
Supported modes:
|
| 58 |
+
- pose: Camera pose estimation evaluation
|
| 59 |
+
|
| 60 |
+
NOT supported:
|
| 61 |
+
- recon_unposed: 3D reconstruction (no GT depth available)
|
| 62 |
+
- recon_posed: 3D reconstruction (no GT depth available)
|
| 63 |
+
"""
|
| 64 |
+
|
| 65 |
+
data_root = DTU64_EVAL_DATA_ROOT
|
| 66 |
+
camera_root = DTU64_CAMERA_ROOT
|
| 67 |
+
SCENES = DTU64_SCENES
|
| 68 |
+
|
| 69 |
+
def __init__(self):
|
| 70 |
+
super().__init__()
|
| 71 |
+
self._scene_cache = {}
|
| 72 |
+
|
| 73 |
+
# ------------------------------
|
| 74 |
+
# Camera file parsing
|
| 75 |
+
# ------------------------------
|
| 76 |
+
|
| 77 |
+
def read_cam_file(self, filename: str) -> tuple:
|
| 78 |
+
"""
|
| 79 |
+
Read DTU camera file containing extrinsics and intrinsics.
|
| 80 |
+
|
| 81 |
+
Args:
|
| 82 |
+
filename: Path to camera text file
|
| 83 |
+
|
| 84 |
+
Returns:
|
| 85 |
+
Tuple of (intrinsics [3,3], extrinsics [4,4])
|
| 86 |
+
"""
|
| 87 |
+
with open(filename) as f:
|
| 88 |
+
lines = [line.rstrip() for line in f.readlines()]
|
| 89 |
+
# extrinsics: line [1,5), 4x4 matrix
|
| 90 |
+
extrinsics = np.fromstring(" ".join(lines[1:5]), dtype=np.float32, sep=" ").reshape((4, 4))
|
| 91 |
+
# intrinsics: line [7-10), 3x3 matrix
|
| 92 |
+
intrinsics = np.fromstring(" ".join(lines[7:10]), dtype=np.float32, sep=" ").reshape((3, 3))
|
| 93 |
+
return intrinsics, extrinsics
|
| 94 |
+
|
| 95 |
+
# ------------------------------
|
| 96 |
+
# Public API
|
| 97 |
+
# ------------------------------
|
| 98 |
+
|
| 99 |
+
def get_data(self, scene: str) -> Dict:
|
| 100 |
+
"""
|
| 101 |
+
Collect per-view image paths, intrinsics/extrinsics for a scene.
|
| 102 |
+
|
| 103 |
+
Args:
|
| 104 |
+
scene: Scene identifier (e.g., "scan105")
|
| 105 |
+
|
| 106 |
+
Returns:
|
| 107 |
+
Dict with:
|
| 108 |
+
- image_files: List[str] - paths to images (64 per scene)
|
| 109 |
+
- extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms
|
| 110 |
+
- intrinsics: np.ndarray [N, 3, 3] - camera intrinsics
|
| 111 |
+
- aux: Dict (empty for this dataset)
|
| 112 |
+
"""
|
| 113 |
+
if scene in self._scene_cache:
|
| 114 |
+
return self._scene_cache[scene]
|
| 115 |
+
|
| 116 |
+
rgb_folder = os.path.join(self.data_root, scene, "image")
|
| 117 |
+
|
| 118 |
+
# Get all PNG files sorted
|
| 119 |
+
files = sorted(glob.glob(os.path.join(rgb_folder, "*.png")))
|
| 120 |
+
|
| 121 |
+
# Reorder: place index 33 first (reference view convention)
|
| 122 |
+
if len(files) > 33:
|
| 123 |
+
files = [files[33]] + files[:33] + files[34:]
|
| 124 |
+
|
| 125 |
+
out = Dict({
|
| 126 |
+
"image_files": [],
|
| 127 |
+
"extrinsics": [],
|
| 128 |
+
"intrinsics": [],
|
| 129 |
+
"aux": Dict({}),
|
| 130 |
+
})
|
| 131 |
+
|
| 132 |
+
for rgb_file in files:
|
| 133 |
+
basename = os.path.basename(rgb_file)
|
| 134 |
+
# File naming: "00000033.png" -> cam_idx = 33
|
| 135 |
+
file_idx = basename.split(".")[0]
|
| 136 |
+
cam_idx = int(file_idx)
|
| 137 |
+
|
| 138 |
+
# Camera file path
|
| 139 |
+
cam_file = os.path.join(self.camera_root, f"{cam_idx:0>8}_cam.txt")
|
| 140 |
+
|
| 141 |
+
if not os.path.exists(cam_file):
|
| 142 |
+
print(f"[DTU-64] Warning: Camera file not found: {cam_file}")
|
| 143 |
+
continue
|
| 144 |
+
|
| 145 |
+
intrinsics, extrinsics = self.read_cam_file(cam_file)
|
| 146 |
+
|
| 147 |
+
out.image_files.append(rgb_file)
|
| 148 |
+
out.extrinsics.append(extrinsics)
|
| 149 |
+
out.intrinsics.append(intrinsics)
|
| 150 |
+
|
| 151 |
+
out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32)
|
| 152 |
+
out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32)
|
| 153 |
+
|
| 154 |
+
print(f"[DTU-64] {scene}: {len(out.image_files)} images (pose evaluation only)")
|
| 155 |
+
|
| 156 |
+
self._scene_cache[scene] = out
|
| 157 |
+
return out
|
| 158 |
+
|
| 159 |
+
def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]:
|
| 160 |
+
"""
|
| 161 |
+
NOT SUPPORTED for DTU-64.
|
| 162 |
+
|
| 163 |
+
DTU-64 is only for pose evaluation, not 3D reconstruction.
|
| 164 |
+
"""
|
| 165 |
+
raise NotImplementedError(
|
| 166 |
+
"DTU-64 dataset is for POSE EVALUATION ONLY. "
|
| 167 |
+
"3D reconstruction evaluation is not supported. "
|
| 168 |
+
"Use the standard 'dtu' dataset for 3D reconstruction evaluation."
|
| 169 |
+
)
|
| 170 |
+
|
| 171 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 172 |
+
"""
|
| 173 |
+
NOT SUPPORTED for DTU-64.
|
| 174 |
+
|
| 175 |
+
DTU-64 is only for pose evaluation, not 3D reconstruction.
|
| 176 |
+
"""
|
| 177 |
+
raise NotImplementedError(
|
| 178 |
+
"DTU-64 dataset is for POSE EVALUATION ONLY. "
|
| 179 |
+
"3D reconstruction (fuse3d) is not supported. "
|
| 180 |
+
"Use the standard 'dtu' dataset for 3D reconstruction."
|
| 181 |
+
)
|
| 182 |
+
|
depth_anything_3/bench/datasets/eth3d.py
ADDED
|
@@ -0,0 +1,594 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
ETH3D Benchmark dataset implementation.
|
| 17 |
+
|
| 18 |
+
ETH3D is a multi-view stereo benchmark with high-resolution images and
|
| 19 |
+
accurate ground truth geometry from laser scanning.
|
| 20 |
+
Reference: https://www.eth3d.net/
|
| 21 |
+
|
| 22 |
+
Evaluation metrics:
|
| 23 |
+
- 3D reconstruction: Accuracy, Completeness, F-score
|
| 24 |
+
- Camera pose estimation: AUC metrics
|
| 25 |
+
"""
|
| 26 |
+
|
| 27 |
+
import glob
|
| 28 |
+
import os
|
| 29 |
+
from typing import Dict as TDict, List, Optional
|
| 30 |
+
|
| 31 |
+
import cv2
|
| 32 |
+
import numpy as np
|
| 33 |
+
import open3d as o3d
|
| 34 |
+
import torch
|
| 35 |
+
import torch.nn.functional as F
|
| 36 |
+
from addict import Dict
|
| 37 |
+
from PIL import Image
|
| 38 |
+
|
| 39 |
+
from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready
|
| 40 |
+
from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY
|
| 41 |
+
from depth_anything_3.bench.utils import (
|
| 42 |
+
create_tsdf_volume,
|
| 43 |
+
evaluate_3d_reconstruction,
|
| 44 |
+
fuse_depth_to_tsdf,
|
| 45 |
+
quat2rotmat,
|
| 46 |
+
sample_points_from_mesh,
|
| 47 |
+
)
|
| 48 |
+
from depth_anything_3.utils.constants import (
|
| 49 |
+
ETH3D_DOWN_SAMPLE,
|
| 50 |
+
ETH3D_EVAL_DATA_ROOT,
|
| 51 |
+
ETH3D_EVAL_THRESHOLD,
|
| 52 |
+
ETH3D_FILTER_KEYS,
|
| 53 |
+
ETH3D_MAX_DEPTH,
|
| 54 |
+
ETH3D_SAMPLING_NUMBER,
|
| 55 |
+
ETH3D_SCENES,
|
| 56 |
+
ETH3D_SDF_TRUNC,
|
| 57 |
+
ETH3D_VOXEL_LENGTH,
|
| 58 |
+
)
|
| 59 |
+
from depth_anything_3.utils.pose_align import align_poses_umeyama
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
@MV_REGISTRY.register(name="eth3d")
|
| 63 |
+
@MONO_REGISTRY.register(name="eth3d")
|
| 64 |
+
class ETH3D(Dataset):
|
| 65 |
+
"""
|
| 66 |
+
ETH3D Benchmark dataset wrapper for DepthAnything3 evaluation.
|
| 67 |
+
|
| 68 |
+
Supports:
|
| 69 |
+
- Camera pose estimation evaluation (AUC metrics)
|
| 70 |
+
- 3D reconstruction evaluation (Accuracy, Completeness, F-score)
|
| 71 |
+
- TSDF-based point cloud fusion
|
| 72 |
+
|
| 73 |
+
Dataset structure:
|
| 74 |
+
eth3d/multiview/
|
| 75 |
+
├── scene_name/
|
| 76 |
+
│ ├── images/ # RGB images
|
| 77 |
+
│ ├── dslr_calibration_jpg/
|
| 78 |
+
│ │ ├── cameras.txt # Camera intrinsics
|
| 79 |
+
│ │ └── images.txt # Camera poses
|
| 80 |
+
│ ├── combined_mesh.ply # Ground truth mesh
|
| 81 |
+
│ └── ground_truth_depth/ # GT depth maps (optional)
|
| 82 |
+
"""
|
| 83 |
+
|
| 84 |
+
data_root = ETH3D_EVAL_DATA_ROOT
|
| 85 |
+
SCENES = ETH3D_SCENES
|
| 86 |
+
|
| 87 |
+
# Evaluation hyperparameters from constants
|
| 88 |
+
max_depth = ETH3D_MAX_DEPTH
|
| 89 |
+
sampling_number = ETH3D_SAMPLING_NUMBER
|
| 90 |
+
voxel_length = ETH3D_VOXEL_LENGTH
|
| 91 |
+
sdf_trunc = ETH3D_SDF_TRUNC
|
| 92 |
+
eval_threshold = ETH3D_EVAL_THRESHOLD
|
| 93 |
+
down_sample = ETH3D_DOWN_SAMPLE
|
| 94 |
+
|
| 95 |
+
def __init__(self):
|
| 96 |
+
super().__init__()
|
| 97 |
+
# Pre-load scene data for efficiency
|
| 98 |
+
self._scene_cache = {}
|
| 99 |
+
|
| 100 |
+
# ------------------------------
|
| 101 |
+
# Camera file parsing
|
| 102 |
+
# ------------------------------
|
| 103 |
+
|
| 104 |
+
def _parse_cameras_txt(self, filepath: str) -> dict:
|
| 105 |
+
"""
|
| 106 |
+
Parse COLMAP-style cameras.txt file.
|
| 107 |
+
|
| 108 |
+
Returns:
|
| 109 |
+
Dict mapping camera_id to intrinsic parameters
|
| 110 |
+
"""
|
| 111 |
+
camera_dict = {}
|
| 112 |
+
with open(filepath) as f:
|
| 113 |
+
lines = f.readlines()
|
| 114 |
+
for line in lines[3:]: # Skip header
|
| 115 |
+
line = line.strip()
|
| 116 |
+
if not line or line.startswith("#"):
|
| 117 |
+
continue
|
| 118 |
+
parts = line.split()
|
| 119 |
+
if len(parts) < 8:
|
| 120 |
+
continue
|
| 121 |
+
cam_id = parts[0]
|
| 122 |
+
# Format: ID, MODEL, WIDTH, HEIGHT, fx, fy, cx, cy, [distortion params...]
|
| 123 |
+
camera_dict[cam_id] = {
|
| 124 |
+
"width": float(parts[2]),
|
| 125 |
+
"height": float(parts[3]),
|
| 126 |
+
"fx": float(parts[4]),
|
| 127 |
+
"fy": float(parts[5]),
|
| 128 |
+
"cx": float(parts[6]),
|
| 129 |
+
"cy": float(parts[7]),
|
| 130 |
+
}
|
| 131 |
+
return camera_dict
|
| 132 |
+
|
| 133 |
+
def _parse_images_txt(self, filepath: str) -> dict:
|
| 134 |
+
"""
|
| 135 |
+
Parse COLMAP-style images.txt file.
|
| 136 |
+
|
| 137 |
+
Returns:
|
| 138 |
+
Dict mapping image path to pose parameters
|
| 139 |
+
"""
|
| 140 |
+
pose_dict = {}
|
| 141 |
+
with open(filepath) as f:
|
| 142 |
+
lines = f.readlines()
|
| 143 |
+
for idx, line in enumerate(lines[4:]): # Skip header
|
| 144 |
+
line = line.strip()
|
| 145 |
+
if not line or line.startswith("#"):
|
| 146 |
+
continue
|
| 147 |
+
# Every other line contains pose info
|
| 148 |
+
if idx % 2 == 0:
|
| 149 |
+
parts = line.split()
|
| 150 |
+
if len(parts) < 10:
|
| 151 |
+
continue
|
| 152 |
+
# Format: IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME
|
| 153 |
+
image_id = parts[0]
|
| 154 |
+
qw, qx, qy, qz = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])
|
| 155 |
+
tx, ty, tz = float(parts[5]), float(parts[6]), float(parts[7])
|
| 156 |
+
camera_id = parts[8]
|
| 157 |
+
name = parts[9]
|
| 158 |
+
pose_dict[name] = {
|
| 159 |
+
"image_id": image_id,
|
| 160 |
+
"quat": [qw, qx, qy, qz],
|
| 161 |
+
"trans": [tx, ty, tz],
|
| 162 |
+
"camera_id": camera_id,
|
| 163 |
+
}
|
| 164 |
+
return pose_dict
|
| 165 |
+
|
| 166 |
+
def _should_filter_image(self, scene: str, image_name: str) -> bool:
|
| 167 |
+
"""Check if image should be filtered out based on known problematic views."""
|
| 168 |
+
filter_keys = ETH3D_FILTER_KEYS.get(scene, [])
|
| 169 |
+
for key in filter_keys:
|
| 170 |
+
if image_name.endswith(key):
|
| 171 |
+
return True
|
| 172 |
+
return False
|
| 173 |
+
|
| 174 |
+
# ------------------------------
|
| 175 |
+
# Public API
|
| 176 |
+
# ------------------------------
|
| 177 |
+
|
| 178 |
+
def get_data(self, scene: str) -> Dict:
|
| 179 |
+
"""
|
| 180 |
+
Collect per-view image paths, intrinsics/extrinsics for a scene.
|
| 181 |
+
|
| 182 |
+
Args:
|
| 183 |
+
scene: Scene identifier (e.g., "courtyard")
|
| 184 |
+
|
| 185 |
+
Returns:
|
| 186 |
+
Dict with:
|
| 187 |
+
- image_files: List[str] - paths to images
|
| 188 |
+
- extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms
|
| 189 |
+
- intrinsics: np.ndarray [N, 3, 3] - camera intrinsics
|
| 190 |
+
- aux: Dict with gt_mesh_path
|
| 191 |
+
"""
|
| 192 |
+
# Check cache
|
| 193 |
+
if scene in self._scene_cache:
|
| 194 |
+
return self._scene_cache[scene]
|
| 195 |
+
|
| 196 |
+
scene_dir = os.path.join(self.data_root, scene)
|
| 197 |
+
|
| 198 |
+
# Parse camera files
|
| 199 |
+
cameras_file = os.path.join(scene_dir, "dslr_calibration_jpg", "cameras.txt")
|
| 200 |
+
images_file = os.path.join(scene_dir, "dslr_calibration_jpg", "images.txt")
|
| 201 |
+
camera_dict = self._parse_cameras_txt(cameras_file)
|
| 202 |
+
pose_dict = self._parse_images_txt(images_file)
|
| 203 |
+
|
| 204 |
+
# Ground truth mesh path
|
| 205 |
+
gt_mesh_path = os.path.join(scene_dir, "combined_mesh.ply")
|
| 206 |
+
|
| 207 |
+
out = Dict({
|
| 208 |
+
"image_files": [],
|
| 209 |
+
"extrinsics": [],
|
| 210 |
+
"intrinsics": [],
|
| 211 |
+
"aux": Dict({
|
| 212 |
+
"gt_mesh_path": gt_mesh_path,
|
| 213 |
+
"heights": [],
|
| 214 |
+
"widths": [],
|
| 215 |
+
}),
|
| 216 |
+
})
|
| 217 |
+
|
| 218 |
+
# Process each image (preserve original order from images.txt)
|
| 219 |
+
filtered_count = 0
|
| 220 |
+
for image_name, pose_info in pose_dict.items():
|
| 221 |
+
# Filter problematic views
|
| 222 |
+
if self._should_filter_image(scene, image_name):
|
| 223 |
+
filtered_count += 1
|
| 224 |
+
continue
|
| 225 |
+
|
| 226 |
+
image_path = os.path.join(scene_dir, "images", image_name)
|
| 227 |
+
if not os.path.exists(image_path):
|
| 228 |
+
continue
|
| 229 |
+
|
| 230 |
+
cam_info = camera_dict.get(pose_info["camera_id"])
|
| 231 |
+
if cam_info is None:
|
| 232 |
+
continue
|
| 233 |
+
|
| 234 |
+
# Build intrinsics matrix
|
| 235 |
+
ixt = np.array([
|
| 236 |
+
[cam_info["fx"], 0, cam_info["cx"]],
|
| 237 |
+
[0, cam_info["fy"], cam_info["cy"]],
|
| 238 |
+
[0, 0, 1],
|
| 239 |
+
], dtype=np.float32)
|
| 240 |
+
|
| 241 |
+
# Build extrinsics matrix (world-to-camera)
|
| 242 |
+
# COLMAP format: world point -> camera point
|
| 243 |
+
rot = quat2rotmat(pose_info["quat"])
|
| 244 |
+
ext = np.eye(4, dtype=np.float32)
|
| 245 |
+
ext[:3, :3] = rot
|
| 246 |
+
ext[:3, 3] = pose_info["trans"]
|
| 247 |
+
|
| 248 |
+
out.image_files.append(image_path)
|
| 249 |
+
out.extrinsics.append(ext)
|
| 250 |
+
out.intrinsics.append(ixt)
|
| 251 |
+
out.aux.heights.append(cam_info["height"])
|
| 252 |
+
out.aux.widths.append(cam_info["width"])
|
| 253 |
+
|
| 254 |
+
out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32)
|
| 255 |
+
out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32)
|
| 256 |
+
|
| 257 |
+
# Print scene info
|
| 258 |
+
total_images = len(pose_dict)
|
| 259 |
+
used_images = len(out.image_files)
|
| 260 |
+
print(f"[ETH3D] {scene}: {used_images}/{total_images} images "
|
| 261 |
+
f"(filtered {filtered_count}, missing {total_images - used_images - filtered_count})")
|
| 262 |
+
|
| 263 |
+
if used_images < 3:
|
| 264 |
+
print(f"[ETH3D] ⚠️ WARNING: {scene} has only {used_images} images - evaluation may fail!")
|
| 265 |
+
|
| 266 |
+
# Cache result
|
| 267 |
+
self._scene_cache[scene] = out
|
| 268 |
+
return out
|
| 269 |
+
|
| 270 |
+
def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]:
|
| 271 |
+
"""
|
| 272 |
+
Evaluate fused point cloud against ETH3D ground truth mesh.
|
| 273 |
+
|
| 274 |
+
Args:
|
| 275 |
+
scene: Scene identifier
|
| 276 |
+
fuse_path: Path to fused point cloud (.ply)
|
| 277 |
+
|
| 278 |
+
Returns:
|
| 279 |
+
Dict with metrics: acc, comp, overall, precision, recall, fscore
|
| 280 |
+
"""
|
| 281 |
+
gt_data = self.get_data(scene)
|
| 282 |
+
gt_mesh_path = gt_data.aux.gt_mesh_path
|
| 283 |
+
|
| 284 |
+
# Load and sample ground truth mesh
|
| 285 |
+
gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path)
|
| 286 |
+
gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number)
|
| 287 |
+
|
| 288 |
+
# Load predicted point cloud
|
| 289 |
+
pred_pcd = o3d.io.read_point_cloud(fuse_path)
|
| 290 |
+
|
| 291 |
+
# Evaluate using shared utility function
|
| 292 |
+
metrics = evaluate_3d_reconstruction(
|
| 293 |
+
pred_pcd,
|
| 294 |
+
gt_pcd,
|
| 295 |
+
threshold=self.eval_threshold,
|
| 296 |
+
down_sample=self.down_sample,
|
| 297 |
+
)
|
| 298 |
+
|
| 299 |
+
return metrics
|
| 300 |
+
|
| 301 |
+
def _load_gt_meta(self, result_path: str) -> Dict:
|
| 302 |
+
"""
|
| 303 |
+
Load saved GT meta (extrinsics, intrinsics, image_files) for fusion.
|
| 304 |
+
|
| 305 |
+
This is needed when frames are sampled, so fuse3d uses the correct
|
| 306 |
+
(sampled) GT instead of full dataset GT.
|
| 307 |
+
|
| 308 |
+
Args:
|
| 309 |
+
result_path: Path to npz file (used to derive gt_meta.npz path)
|
| 310 |
+
|
| 311 |
+
Returns:
|
| 312 |
+
Dict with GT data, or None if gt_meta.npz doesn't exist
|
| 313 |
+
"""
|
| 314 |
+
# gt_meta.npz is in the same exports/ directory as results.npz
|
| 315 |
+
export_dir = os.path.dirname(result_path) # exports/mini_npz/
|
| 316 |
+
gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz")
|
| 317 |
+
|
| 318 |
+
if os.path.exists(gt_meta_path):
|
| 319 |
+
data = np.load(gt_meta_path, allow_pickle=True)
|
| 320 |
+
return Dict({
|
| 321 |
+
"extrinsics": data["extrinsics"],
|
| 322 |
+
"intrinsics": data["intrinsics"],
|
| 323 |
+
"image_files": data["image_files"] if "image_files" in data else None,
|
| 324 |
+
})
|
| 325 |
+
return None
|
| 326 |
+
|
| 327 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 328 |
+
"""
|
| 329 |
+
Fuse per-view depths into a point cloud using TSDF fusion.
|
| 330 |
+
|
| 331 |
+
Pipeline:
|
| 332 |
+
1. Load original images (keep original size)
|
| 333 |
+
2. Resize depth to original image size (nearest interpolation)
|
| 334 |
+
3. Adjust intrinsics to original image size
|
| 335 |
+
4. Apply scale alignment and mask invalid depths
|
| 336 |
+
5. TSDF fusion
|
| 337 |
+
|
| 338 |
+
Args:
|
| 339 |
+
scene: Scene identifier
|
| 340 |
+
result_path: Path to npz file with predicted depths/poses
|
| 341 |
+
fuse_path: Output path for fused point cloud (.ply)
|
| 342 |
+
mode: "recon_unposed" or "recon_posed"
|
| 343 |
+
"""
|
| 344 |
+
# Try to load saved GT meta (handles frame sampling)
|
| 345 |
+
gt_meta = self._load_gt_meta(result_path)
|
| 346 |
+
if gt_meta is not None:
|
| 347 |
+
gt_data = gt_meta
|
| 348 |
+
else:
|
| 349 |
+
gt_data = self.get_data(scene)
|
| 350 |
+
_wait_for_file_ready(result_path)
|
| 351 |
+
pred_data = Dict({k: v for k, v in np.load(result_path).items()})
|
| 352 |
+
|
| 353 |
+
# Load original images (keep original size)
|
| 354 |
+
images = []
|
| 355 |
+
orig_sizes = [] # (H, W) for each image
|
| 356 |
+
for img_path in gt_data.image_files:
|
| 357 |
+
img = cv2.imread(img_path)
|
| 358 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 359 |
+
images.append(img)
|
| 360 |
+
orig_sizes.append((img.shape[0], img.shape[1]))
|
| 361 |
+
|
| 362 |
+
# Prepare depths, intrinsics, extrinsics with resize to original size
|
| 363 |
+
if mode == "recon_unposed":
|
| 364 |
+
depths, intrinsics, extrinsics = self._prep_unposed(
|
| 365 |
+
pred_data, gt_data, orig_sizes, scene=scene
|
| 366 |
+
)
|
| 367 |
+
elif mode == "recon_posed":
|
| 368 |
+
depths, intrinsics, extrinsics = self._prep_posed(
|
| 369 |
+
pred_data, gt_data, orig_sizes, scene=scene
|
| 370 |
+
)
|
| 371 |
+
else:
|
| 372 |
+
raise ValueError(f"Invalid mode: {mode}")
|
| 373 |
+
|
| 374 |
+
images = np.stack(images, axis=0)
|
| 375 |
+
|
| 376 |
+
# Create TSDF volume and fuse
|
| 377 |
+
volume = create_tsdf_volume(
|
| 378 |
+
voxel_length=self.voxel_length,
|
| 379 |
+
sdf_trunc=self.sdf_trunc,
|
| 380 |
+
)
|
| 381 |
+
mesh = fuse_depth_to_tsdf(
|
| 382 |
+
volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth
|
| 383 |
+
)
|
| 384 |
+
|
| 385 |
+
# Sample points from mesh
|
| 386 |
+
pcd = sample_points_from_mesh(mesh, self.sampling_number)
|
| 387 |
+
|
| 388 |
+
# Save point cloud
|
| 389 |
+
os.makedirs(os.path.dirname(fuse_path), exist_ok=True)
|
| 390 |
+
o3d.io.write_point_cloud(fuse_path, pcd)
|
| 391 |
+
|
| 392 |
+
# ------------------------------
|
| 393 |
+
# Private helpers
|
| 394 |
+
# ------------------------------
|
| 395 |
+
|
| 396 |
+
def _prep_unposed(
|
| 397 |
+
self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str = None
|
| 398 |
+
) -> tuple:
|
| 399 |
+
"""
|
| 400 |
+
Prepare depths/intrinsics/extrinsics for recon_unposed mode.
|
| 401 |
+
|
| 402 |
+
Pipeline:
|
| 403 |
+
1. Umeyama scale alignment
|
| 404 |
+
2. Load GT mask for each frame
|
| 405 |
+
3. Resize depth to original image size (nearest)
|
| 406 |
+
4. Apply GT mask BEFORE scale
|
| 407 |
+
5. Apply scale
|
| 408 |
+
6. Adjust intrinsics to original image size
|
| 409 |
+
"""
|
| 410 |
+
# Scale alignment with fixed random_state for reproducibility
|
| 411 |
+
_, _, scale, extrinsics = align_poses_umeyama(
|
| 412 |
+
gt_data.extrinsics.copy(),
|
| 413 |
+
pred_data.extrinsics.copy(),
|
| 414 |
+
return_aligned=True,
|
| 415 |
+
ransac=True,
|
| 416 |
+
random_state=42,
|
| 417 |
+
)
|
| 418 |
+
|
| 419 |
+
# Get model output size
|
| 420 |
+
model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2]
|
| 421 |
+
|
| 422 |
+
# Process each frame
|
| 423 |
+
depths_out = []
|
| 424 |
+
intrinsics_out = []
|
| 425 |
+
for i in range(len(pred_data.depth)):
|
| 426 |
+
orig_h, orig_w = orig_sizes[i]
|
| 427 |
+
image_name = os.path.basename(gt_data.image_files[i])
|
| 428 |
+
|
| 429 |
+
# Resize depth to original image size (nearest interpolation)
|
| 430 |
+
depth = cv2.resize(
|
| 431 |
+
pred_data.depth[i],
|
| 432 |
+
(orig_w, orig_h),
|
| 433 |
+
interpolation=cv2.INTER_NEAREST,
|
| 434 |
+
)
|
| 435 |
+
|
| 436 |
+
# Load GT mask (apply BEFORE scale)
|
| 437 |
+
gt_zero_mask = None
|
| 438 |
+
if scene is not None:
|
| 439 |
+
gt_zero_mask = self._load_gt_mask(scene, image_name, (orig_h, orig_w))
|
| 440 |
+
|
| 441 |
+
# Mask invalid depths BEFORE scale
|
| 442 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 443 |
+
|
| 444 |
+
# Apply scale AFTER mask
|
| 445 |
+
depth = depth * scale
|
| 446 |
+
|
| 447 |
+
# Adjust intrinsics to original image size
|
| 448 |
+
h_ratio = orig_h / model_h
|
| 449 |
+
w_ratio = orig_w / model_w
|
| 450 |
+
ixt = pred_data.intrinsics[i].copy()
|
| 451 |
+
ixt[0, :] *= w_ratio # fx, 0, cx
|
| 452 |
+
ixt[1, :] *= h_ratio # 0, fy, cy
|
| 453 |
+
|
| 454 |
+
depths_out.append(depth)
|
| 455 |
+
intrinsics_out.append(ixt)
|
| 456 |
+
|
| 457 |
+
return np.stack(depths_out), np.stack(intrinsics_out), extrinsics
|
| 458 |
+
|
| 459 |
+
def _prep_posed(
|
| 460 |
+
self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str = None
|
| 461 |
+
) -> tuple:
|
| 462 |
+
"""
|
| 463 |
+
Prepare depths/intrinsics/extrinsics for recon_posed mode.
|
| 464 |
+
|
| 465 |
+
Uses GT intrinsics/extrinsics but aligns depth scale via Umeyama.
|
| 466 |
+
Depth is resized to original image size.
|
| 467 |
+
"""
|
| 468 |
+
# Scale alignment with fixed random_state for reproducibility
|
| 469 |
+
_, _, scale, _ = align_poses_umeyama(
|
| 470 |
+
gt_data.extrinsics.copy(),
|
| 471 |
+
pred_data.extrinsics.copy(),
|
| 472 |
+
return_aligned=True,
|
| 473 |
+
ransac=True,
|
| 474 |
+
random_state=42,
|
| 475 |
+
)
|
| 476 |
+
|
| 477 |
+
# Process each frame
|
| 478 |
+
depths_out = []
|
| 479 |
+
for i in range(len(pred_data.depth)):
|
| 480 |
+
orig_h, orig_w = orig_sizes[i]
|
| 481 |
+
image_name = os.path.basename(gt_data.image_files[i])
|
| 482 |
+
|
| 483 |
+
# Resize depth to original image size (nearest interpolation)
|
| 484 |
+
depth = cv2.resize(
|
| 485 |
+
pred_data.depth[i],
|
| 486 |
+
(orig_w, orig_h),
|
| 487 |
+
interpolation=cv2.INTER_NEAREST,
|
| 488 |
+
)
|
| 489 |
+
|
| 490 |
+
# Load GT mask (apply BEFORE scale)
|
| 491 |
+
gt_zero_mask = None
|
| 492 |
+
if scene is not None:
|
| 493 |
+
gt_zero_mask = self._load_gt_mask(scene, image_name, (orig_h, orig_w))
|
| 494 |
+
|
| 495 |
+
# Mask invalid depths BEFORE scale
|
| 496 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 497 |
+
|
| 498 |
+
# Apply scale AFTER mask
|
| 499 |
+
depth = depth * scale
|
| 500 |
+
|
| 501 |
+
depths_out.append(depth)
|
| 502 |
+
|
| 503 |
+
# Use GT intrinsics and extrinsics (already at original image size)
|
| 504 |
+
return np.stack(depths_out), gt_data.intrinsics.copy(), gt_data.extrinsics.copy()
|
| 505 |
+
|
| 506 |
+
def _load_gt_mask(self, scene: str, image_name: str, shape: tuple) -> np.ndarray:
|
| 507 |
+
"""
|
| 508 |
+
Load GT mask for masking invalid regions.
|
| 509 |
+
|
| 510 |
+
GT mask marks occluded or invalid regions that should be excluded
|
| 511 |
+
from depth fusion and evaluation.
|
| 512 |
+
|
| 513 |
+
Args:
|
| 514 |
+
scene: Scene identifier
|
| 515 |
+
image_name: Image filename (e.g., "DSC_0307.JPG")
|
| 516 |
+
shape: (height, width) of the image
|
| 517 |
+
|
| 518 |
+
Returns:
|
| 519 |
+
Boolean mask where True = valid region to keep
|
| 520 |
+
"""
|
| 521 |
+
h, w = shape
|
| 522 |
+
|
| 523 |
+
# GT mask file path
|
| 524 |
+
gt_mask_path = os.path.join(
|
| 525 |
+
self.data_root, scene, "masks_for_images", "dslr_images",
|
| 526 |
+
image_name.replace(".JPG", ".png")
|
| 527 |
+
)
|
| 528 |
+
|
| 529 |
+
# GT depth file path (used to determine valid depth regions)
|
| 530 |
+
gt_depth_path = os.path.join(
|
| 531 |
+
self.data_root, scene, "ground_truth_depth", "dslr_images", image_name
|
| 532 |
+
)
|
| 533 |
+
|
| 534 |
+
# Load GT depth
|
| 535 |
+
if os.path.exists(gt_depth_path):
|
| 536 |
+
gt_depth = np.fromfile(gt_depth_path, dtype=np.float32).reshape(h, w)
|
| 537 |
+
else:
|
| 538 |
+
gt_depth = np.ones((h, w), dtype=np.float32)
|
| 539 |
+
|
| 540 |
+
# Load GT mask
|
| 541 |
+
if os.path.exists(gt_mask_path):
|
| 542 |
+
gt_mask = cv2.imread(gt_mask_path, cv2.IMREAD_GRAYSCALE)
|
| 543 |
+
gt_mask = np.asarray(gt_mask)
|
| 544 |
+
else:
|
| 545 |
+
gt_mask = np.zeros((h, w), dtype=np.uint8)
|
| 546 |
+
|
| 547 |
+
# Compute zero_mask
|
| 548 |
+
# gt_mask == 1 means occluded/invalid region
|
| 549 |
+
invalid_mask_from_gt = gt_mask == 1
|
| 550 |
+
gt_depth_copy = gt_depth.copy()
|
| 551 |
+
gt_depth_copy[gt_mask == 1] = 0
|
| 552 |
+
|
| 553 |
+
invalid_mask_from_gt_depth = np.logical_or(gt_depth_copy == 0, gt_depth_copy == np.inf)
|
| 554 |
+
|
| 555 |
+
# zero_mask: valid region that should be kept
|
| 556 |
+
zero_mask = np.logical_and(
|
| 557 |
+
np.logical_not(invalid_mask_from_gt),
|
| 558 |
+
np.logical_not(invalid_mask_from_gt_depth)
|
| 559 |
+
)
|
| 560 |
+
|
| 561 |
+
return zero_mask
|
| 562 |
+
|
| 563 |
+
def _mask_invalid_depth(
|
| 564 |
+
self, depth: np.ndarray, gt_zero_mask: np.ndarray = None
|
| 565 |
+
) -> np.ndarray:
|
| 566 |
+
"""
|
| 567 |
+
Mask invalid depth values by setting them to 0.
|
| 568 |
+
|
| 569 |
+
Logic:
|
| 570 |
+
1. Apply GT mask (if provided) - marks occluded/invalid regions
|
| 571 |
+
2. Mask pred invalid values (nan, inf)
|
| 572 |
+
|
| 573 |
+
Args:
|
| 574 |
+
depth: Depth map to mask
|
| 575 |
+
gt_zero_mask: Optional GT mask (True = valid region)
|
| 576 |
+
|
| 577 |
+
Returns:
|
| 578 |
+
Masked depth map with invalid regions set to 0
|
| 579 |
+
"""
|
| 580 |
+
depth = depth.copy()
|
| 581 |
+
|
| 582 |
+
# Apply GT mask first (before scale)
|
| 583 |
+
if gt_zero_mask is not None:
|
| 584 |
+
# Also mask out invalid pred depth
|
| 585 |
+
pred_invalid = np.isnan(depth) | np.isinf(depth)
|
| 586 |
+
combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid))
|
| 587 |
+
depth = depth * combined_mask.astype(np.float32)
|
| 588 |
+
else:
|
| 589 |
+
# Fallback: only mask pred invalid values
|
| 590 |
+
invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0)
|
| 591 |
+
depth[invalid_mask] = 0.0
|
| 592 |
+
|
| 593 |
+
return depth
|
| 594 |
+
|
depth_anything_3/bench/datasets/hiroom.py
ADDED
|
@@ -0,0 +1,440 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
HiRoom Dataset implementation.
|
| 17 |
+
|
| 18 |
+
HiRoom is an indoor RGB-D dataset containing ground truth camera poses,
|
| 19 |
+
depth maps, and fused point clouds.
|
| 20 |
+
|
| 21 |
+
Evaluation metrics:
|
| 22 |
+
- 3D reconstruction: Accuracy, Completeness, F-score
|
| 23 |
+
- Camera pose estimation: AUC metrics
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
import os
|
| 27 |
+
from typing import Dict as TDict, List
|
| 28 |
+
|
| 29 |
+
import cv2
|
| 30 |
+
import numpy as np
|
| 31 |
+
import open3d as o3d
|
| 32 |
+
from addict import Dict
|
| 33 |
+
|
| 34 |
+
from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready
|
| 35 |
+
from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY
|
| 36 |
+
from depth_anything_3.bench.utils import (
|
| 37 |
+
create_tsdf_volume,
|
| 38 |
+
evaluate_3d_reconstruction,
|
| 39 |
+
fuse_depth_to_tsdf,
|
| 40 |
+
sample_points_from_mesh,
|
| 41 |
+
)
|
| 42 |
+
from depth_anything_3.utils.constants import (
|
| 43 |
+
HIROOM_DOWN_SAMPLE,
|
| 44 |
+
HIROOM_EVAL_DATA_ROOT,
|
| 45 |
+
HIROOM_EVAL_THRESHOLD,
|
| 46 |
+
HIROOM_GT_ROOT_PATH,
|
| 47 |
+
HIROOM_MAX_DEPTH,
|
| 48 |
+
HIROOM_SAMPLING_NUMBER,
|
| 49 |
+
HIROOM_SCENE_LIST_PATH,
|
| 50 |
+
HIROOM_SDF_TRUNC,
|
| 51 |
+
HIROOM_VOXEL_LENGTH,
|
| 52 |
+
)
|
| 53 |
+
from depth_anything_3.utils.pose_align import align_poses_umeyama
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _load_scene_list() -> List[str]:
|
| 57 |
+
"""Load scene list from file."""
|
| 58 |
+
if os.path.exists(HIROOM_SCENE_LIST_PATH):
|
| 59 |
+
with open(HIROOM_SCENE_LIST_PATH, "r") as f:
|
| 60 |
+
return f.read().splitlines()
|
| 61 |
+
return []
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
@MV_REGISTRY.register(name="hiroom")
|
| 65 |
+
@MONO_REGISTRY.register(name="hiroom")
|
| 66 |
+
class HiRoomDataset(Dataset):
|
| 67 |
+
"""
|
| 68 |
+
HiRoom Dataset wrapper for DepthAnything3 evaluation.
|
| 69 |
+
|
| 70 |
+
Supports:
|
| 71 |
+
- Camera pose estimation evaluation (AUC metrics)
|
| 72 |
+
- 3D reconstruction evaluation (Accuracy, Completeness, F-score)
|
| 73 |
+
- TSDF-based point cloud fusion
|
| 74 |
+
|
| 75 |
+
Dataset structure:
|
| 76 |
+
HiRoom/
|
| 77 |
+
├── {scene_path}/
|
| 78 |
+
│ ├── image/ # RGB images
|
| 79 |
+
│ ├── depth/ # GT depth maps
|
| 80 |
+
│ ├── pose/ # Camera poses (.npy)
|
| 81 |
+
│ ├── cam_K.npy # Camera intrinsics
|
| 82 |
+
│ └── aliasing_mask/ # Aliasing masks
|
| 83 |
+
|
| 84 |
+
fused_pcd/
|
| 85 |
+
└── {scene_name}.ply # Ground truth fused point cloud
|
| 86 |
+
"""
|
| 87 |
+
|
| 88 |
+
data_root = HIROOM_EVAL_DATA_ROOT
|
| 89 |
+
gt_root_path = HIROOM_GT_ROOT_PATH
|
| 90 |
+
SCENES = _load_scene_list()
|
| 91 |
+
|
| 92 |
+
# Evaluation hyperparameters from constants
|
| 93 |
+
max_depth = HIROOM_MAX_DEPTH
|
| 94 |
+
sampling_number = HIROOM_SAMPLING_NUMBER
|
| 95 |
+
voxel_length = HIROOM_VOXEL_LENGTH
|
| 96 |
+
sdf_trunc = HIROOM_SDF_TRUNC
|
| 97 |
+
eval_threshold = HIROOM_EVAL_THRESHOLD
|
| 98 |
+
down_sample = HIROOM_DOWN_SAMPLE
|
| 99 |
+
|
| 100 |
+
def __init__(self):
|
| 101 |
+
super().__init__()
|
| 102 |
+
self._scene_cache = {}
|
| 103 |
+
|
| 104 |
+
# ------------------------------
|
| 105 |
+
# Public API
|
| 106 |
+
# ------------------------------
|
| 107 |
+
|
| 108 |
+
def get_data(self, scene: str) -> Dict:
|
| 109 |
+
"""
|
| 110 |
+
Collect per-view image paths, intrinsics/extrinsics for a scene.
|
| 111 |
+
|
| 112 |
+
Args:
|
| 113 |
+
scene: Scene path (e.g., "xxx/yyy/zzz")
|
| 114 |
+
|
| 115 |
+
Returns:
|
| 116 |
+
Dict with:
|
| 117 |
+
- image_files: List[str] - paths to images
|
| 118 |
+
- extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms
|
| 119 |
+
- intrinsics: np.ndarray [N, 3, 3] - camera intrinsics
|
| 120 |
+
- aux: Dict with gt_pcd_path, gt_depth_files, aliasing_mask_files
|
| 121 |
+
"""
|
| 122 |
+
if scene in self._scene_cache:
|
| 123 |
+
return self._scene_cache[scene]
|
| 124 |
+
|
| 125 |
+
scene_dir = os.path.join(self.data_root, scene)
|
| 126 |
+
image_dir = os.path.join(scene_dir, "image")
|
| 127 |
+
|
| 128 |
+
# Get scene name for GT point cloud
|
| 129 |
+
scene_name = "-".join(scene.split("/")[-3:])
|
| 130 |
+
gt_pcd_path = os.path.join(self.gt_root_path, f"{scene_name}.ply")
|
| 131 |
+
|
| 132 |
+
# Load shared camera intrinsics
|
| 133 |
+
intrin_path = os.path.join(scene_dir, "cam_K.npy")
|
| 134 |
+
ixt_shared = np.load(intrin_path).astype(np.float32)
|
| 135 |
+
|
| 136 |
+
# Get all image names sorted
|
| 137 |
+
image_names = sorted(os.listdir(image_dir))
|
| 138 |
+
|
| 139 |
+
out = Dict({
|
| 140 |
+
"image_files": [],
|
| 141 |
+
"extrinsics": [],
|
| 142 |
+
"intrinsics": [],
|
| 143 |
+
"aux": Dict({
|
| 144 |
+
"gt_pcd_path": gt_pcd_path,
|
| 145 |
+
"gt_depth_files": [],
|
| 146 |
+
"aliasing_mask_files": [],
|
| 147 |
+
}),
|
| 148 |
+
})
|
| 149 |
+
|
| 150 |
+
for img_name in image_names:
|
| 151 |
+
img_path = os.path.join(image_dir, img_name)
|
| 152 |
+
frame_name = img_name.split(".")[0]
|
| 153 |
+
|
| 154 |
+
# Depth and pose paths
|
| 155 |
+
depth_path = os.path.join(scene_dir, "depth", f"{frame_name}.png")
|
| 156 |
+
pose_path = os.path.join(scene_dir, "pose", f"{frame_name}.npy")
|
| 157 |
+
aliasing_mask_path = os.path.join(scene_dir, "aliasing_mask", f"{frame_name}.png")
|
| 158 |
+
|
| 159 |
+
if not os.path.exists(pose_path):
|
| 160 |
+
continue
|
| 161 |
+
|
| 162 |
+
# Load extrinsics (world-to-camera)
|
| 163 |
+
ext = np.load(pose_path).astype(np.float32)
|
| 164 |
+
|
| 165 |
+
out.image_files.append(img_path)
|
| 166 |
+
out.extrinsics.append(ext)
|
| 167 |
+
out.intrinsics.append(ixt_shared.copy())
|
| 168 |
+
out.aux.gt_depth_files.append(depth_path)
|
| 169 |
+
out.aux.aliasing_mask_files.append(aliasing_mask_path)
|
| 170 |
+
|
| 171 |
+
out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32)
|
| 172 |
+
out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32)
|
| 173 |
+
|
| 174 |
+
print(f"[HiRoom] {scene}: {len(out.image_files)} images")
|
| 175 |
+
|
| 176 |
+
self._scene_cache[scene] = out
|
| 177 |
+
return out
|
| 178 |
+
|
| 179 |
+
def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]:
|
| 180 |
+
"""
|
| 181 |
+
Evaluate fused point cloud against HiRoom ground truth point cloud.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
scene: Scene identifier
|
| 185 |
+
fuse_path: Path to fused point cloud (.ply)
|
| 186 |
+
|
| 187 |
+
Returns:
|
| 188 |
+
Dict with metrics: acc, comp, overall, precision, recall, fscore
|
| 189 |
+
"""
|
| 190 |
+
gt_data = self.get_data(scene)
|
| 191 |
+
gt_pcd_path = gt_data.aux.gt_pcd_path
|
| 192 |
+
|
| 193 |
+
# Load ground truth point cloud
|
| 194 |
+
gt_pcd = o3d.io.read_point_cloud(gt_pcd_path)
|
| 195 |
+
|
| 196 |
+
# Load predicted point cloud
|
| 197 |
+
pred_pcd = o3d.io.read_point_cloud(fuse_path)
|
| 198 |
+
|
| 199 |
+
# Evaluate using shared utility function
|
| 200 |
+
metrics = evaluate_3d_reconstruction(
|
| 201 |
+
pred_pcd,
|
| 202 |
+
gt_pcd,
|
| 203 |
+
threshold=self.eval_threshold,
|
| 204 |
+
down_sample=self.down_sample,
|
| 205 |
+
)
|
| 206 |
+
|
| 207 |
+
return metrics
|
| 208 |
+
|
| 209 |
+
def _load_gt_meta(self, result_path: str) -> Dict:
|
| 210 |
+
"""Load saved GT meta for fusion."""
|
| 211 |
+
export_dir = os.path.dirname(result_path)
|
| 212 |
+
gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz")
|
| 213 |
+
|
| 214 |
+
if os.path.exists(gt_meta_path):
|
| 215 |
+
data = np.load(gt_meta_path, allow_pickle=True)
|
| 216 |
+
image_files = list(data["image_files"])
|
| 217 |
+
return Dict({
|
| 218 |
+
"extrinsics": data["extrinsics"],
|
| 219 |
+
"intrinsics": data["intrinsics"],
|
| 220 |
+
"image_files": image_files,
|
| 221 |
+
})
|
| 222 |
+
return None
|
| 223 |
+
|
| 224 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 225 |
+
"""
|
| 226 |
+
Fuse per-view depths into a point cloud using TSDF fusion.
|
| 227 |
+
|
| 228 |
+
Args:
|
| 229 |
+
scene: Scene identifier
|
| 230 |
+
result_path: Path to npz file with predicted depths/poses
|
| 231 |
+
fuse_path: Output path for fused point cloud (.ply)
|
| 232 |
+
mode: "recon_unposed" or "recon_posed"
|
| 233 |
+
"""
|
| 234 |
+
# Get full GT data
|
| 235 |
+
full_gt_data = self.get_data(scene)
|
| 236 |
+
|
| 237 |
+
# Try to load saved GT meta (handles frame sampling)
|
| 238 |
+
gt_meta = self._load_gt_meta(result_path)
|
| 239 |
+
if gt_meta is not None:
|
| 240 |
+
gt_data = gt_meta
|
| 241 |
+
image_indices = [
|
| 242 |
+
full_gt_data.image_files.index(f)
|
| 243 |
+
for f in gt_data.image_files
|
| 244 |
+
if f in full_gt_data.image_files
|
| 245 |
+
]
|
| 246 |
+
else:
|
| 247 |
+
gt_data = full_gt_data
|
| 248 |
+
image_indices = list(range(len(full_gt_data.image_files)))
|
| 249 |
+
|
| 250 |
+
_wait_for_file_ready(result_path)
|
| 251 |
+
pred_data = Dict({k: v for k, v in np.load(result_path).items()})
|
| 252 |
+
|
| 253 |
+
# Load images
|
| 254 |
+
images = []
|
| 255 |
+
orig_sizes = []
|
| 256 |
+
for img_idx in image_indices:
|
| 257 |
+
img_path = full_gt_data.image_files[img_idx]
|
| 258 |
+
img = cv2.imread(img_path)
|
| 259 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 260 |
+
images.append(img)
|
| 261 |
+
orig_sizes.append((img.shape[0], img.shape[1]))
|
| 262 |
+
|
| 263 |
+
images = np.stack(images, axis=0)
|
| 264 |
+
|
| 265 |
+
# Prepare depths, intrinsics, extrinsics
|
| 266 |
+
if mode == "recon_unposed":
|
| 267 |
+
depths, intrinsics, extrinsics = self._prep_unposed(
|
| 268 |
+
pred_data, gt_data, full_gt_data, image_indices, orig_sizes, scene=scene
|
| 269 |
+
)
|
| 270 |
+
elif mode == "recon_posed":
|
| 271 |
+
depths, intrinsics, extrinsics = self._prep_posed(
|
| 272 |
+
pred_data, gt_data, full_gt_data, image_indices, orig_sizes, scene=scene
|
| 273 |
+
)
|
| 274 |
+
else:
|
| 275 |
+
raise ValueError(f"Invalid mode: {mode}")
|
| 276 |
+
|
| 277 |
+
# Create TSDF volume and fuse
|
| 278 |
+
volume = create_tsdf_volume(
|
| 279 |
+
voxel_length=self.voxel_length,
|
| 280 |
+
sdf_trunc=self.sdf_trunc,
|
| 281 |
+
)
|
| 282 |
+
mesh = fuse_depth_to_tsdf(
|
| 283 |
+
volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth
|
| 284 |
+
)
|
| 285 |
+
|
| 286 |
+
# Sample points from mesh
|
| 287 |
+
pcd = sample_points_from_mesh(mesh, self.sampling_number)
|
| 288 |
+
|
| 289 |
+
# Save point cloud
|
| 290 |
+
os.makedirs(os.path.dirname(fuse_path), exist_ok=True)
|
| 291 |
+
o3d.io.write_point_cloud(fuse_path, pcd)
|
| 292 |
+
|
| 293 |
+
# ------------------------------
|
| 294 |
+
# Private helpers
|
| 295 |
+
# ------------------------------
|
| 296 |
+
|
| 297 |
+
def _prep_unposed(
|
| 298 |
+
self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict,
|
| 299 |
+
image_indices: list, orig_sizes: list, scene: str = None
|
| 300 |
+
) -> tuple:
|
| 301 |
+
"""Prepare depths/intrinsics/extrinsics for recon_unposed mode."""
|
| 302 |
+
# Scale alignment with fixed random_state for reproducibility
|
| 303 |
+
_, _, scale, extrinsics = align_poses_umeyama(
|
| 304 |
+
gt_data.extrinsics.copy(),
|
| 305 |
+
pred_data.extrinsics.copy(),
|
| 306 |
+
return_aligned=True,
|
| 307 |
+
ransac=True,
|
| 308 |
+
random_state=42,
|
| 309 |
+
)
|
| 310 |
+
|
| 311 |
+
model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2]
|
| 312 |
+
|
| 313 |
+
depths_out = []
|
| 314 |
+
intrinsics_out = []
|
| 315 |
+
for i in range(len(pred_data.depth)):
|
| 316 |
+
orig_h, orig_w = orig_sizes[i]
|
| 317 |
+
img_idx = image_indices[i]
|
| 318 |
+
|
| 319 |
+
# Resize depth to original image size
|
| 320 |
+
depth = cv2.resize(
|
| 321 |
+
pred_data.depth[i],
|
| 322 |
+
(orig_w, orig_h),
|
| 323 |
+
interpolation=cv2.INTER_NEAREST,
|
| 324 |
+
)
|
| 325 |
+
|
| 326 |
+
# Load GT mask
|
| 327 |
+
gt_zero_mask = self._load_gt_mask(
|
| 328 |
+
full_gt_data.aux.gt_depth_files[img_idx],
|
| 329 |
+
full_gt_data.aux.aliasing_mask_files[img_idx],
|
| 330 |
+
)
|
| 331 |
+
|
| 332 |
+
# Mask invalid depths BEFORE scale
|
| 333 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 334 |
+
|
| 335 |
+
# Apply scale AFTER mask
|
| 336 |
+
depth = depth * scale
|
| 337 |
+
|
| 338 |
+
# Adjust intrinsics to original image size
|
| 339 |
+
h_ratio = orig_h / model_h
|
| 340 |
+
w_ratio = orig_w / model_w
|
| 341 |
+
ixt = pred_data.intrinsics[i].copy()
|
| 342 |
+
ixt[0, :] *= w_ratio
|
| 343 |
+
ixt[1, :] *= h_ratio
|
| 344 |
+
|
| 345 |
+
depths_out.append(depth)
|
| 346 |
+
intrinsics_out.append(ixt)
|
| 347 |
+
|
| 348 |
+
return np.stack(depths_out), np.stack(intrinsics_out), extrinsics
|
| 349 |
+
|
| 350 |
+
def _prep_posed(
|
| 351 |
+
self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict,
|
| 352 |
+
image_indices: list, orig_sizes: list, scene: str = None
|
| 353 |
+
) -> tuple:
|
| 354 |
+
"""Prepare depths/intrinsics/extrinsics for recon_posed mode."""
|
| 355 |
+
# Scale alignment
|
| 356 |
+
_, _, scale, _ = align_poses_umeyama(
|
| 357 |
+
gt_data.extrinsics.copy(),
|
| 358 |
+
pred_data.extrinsics.copy(),
|
| 359 |
+
return_aligned=True,
|
| 360 |
+
ransac=True,
|
| 361 |
+
random_state=42,
|
| 362 |
+
)
|
| 363 |
+
|
| 364 |
+
depths_out = []
|
| 365 |
+
for i in range(len(pred_data.depth)):
|
| 366 |
+
orig_h, orig_w = orig_sizes[i]
|
| 367 |
+
img_idx = image_indices[i]
|
| 368 |
+
|
| 369 |
+
# Resize depth to original image size
|
| 370 |
+
depth = cv2.resize(
|
| 371 |
+
pred_data.depth[i],
|
| 372 |
+
(orig_w, orig_h),
|
| 373 |
+
interpolation=cv2.INTER_NEAREST,
|
| 374 |
+
)
|
| 375 |
+
|
| 376 |
+
# Load GT mask
|
| 377 |
+
gt_zero_mask = self._load_gt_mask(
|
| 378 |
+
full_gt_data.aux.gt_depth_files[img_idx],
|
| 379 |
+
full_gt_data.aux.aliasing_mask_files[img_idx],
|
| 380 |
+
)
|
| 381 |
+
|
| 382 |
+
# Mask invalid depths BEFORE scale
|
| 383 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 384 |
+
|
| 385 |
+
# Apply scale AFTER mask
|
| 386 |
+
depth = depth * scale
|
| 387 |
+
|
| 388 |
+
depths_out.append(depth)
|
| 389 |
+
|
| 390 |
+
# Use GT intrinsics and extrinsics
|
| 391 |
+
gt_intrinsics = np.stack([full_gt_data.intrinsics[idx] for idx in image_indices])
|
| 392 |
+
gt_extrinsics = np.stack([full_gt_data.extrinsics[idx] for idx in image_indices])
|
| 393 |
+
|
| 394 |
+
return np.stack(depths_out), gt_intrinsics, gt_extrinsics
|
| 395 |
+
|
| 396 |
+
def _load_gt_mask(self, gt_depth_path: str, aliasing_mask_path: str) -> np.ndarray:
|
| 397 |
+
"""
|
| 398 |
+
Load GT depth and aliasing mask to create valid mask.
|
| 399 |
+
|
| 400 |
+
For HiRoom:
|
| 401 |
+
- GT depth is stored as 16-bit PNG, scaled to 100m range
|
| 402 |
+
- Aliasing mask marks regions to exclude
|
| 403 |
+
|
| 404 |
+
Returns:
|
| 405 |
+
Boolean mask where True = valid region to keep
|
| 406 |
+
"""
|
| 407 |
+
# Load GT depth
|
| 408 |
+
if os.path.exists(gt_depth_path):
|
| 409 |
+
gt_depth = cv2.imread(gt_depth_path, -1) / 65535.0 * 100.0
|
| 410 |
+
else:
|
| 411 |
+
return None
|
| 412 |
+
|
| 413 |
+
# Load aliasing mask
|
| 414 |
+
aliasing_mask = None
|
| 415 |
+
if os.path.exists(aliasing_mask_path):
|
| 416 |
+
aliasing_mask = cv2.imread(aliasing_mask_path, -1) > 0
|
| 417 |
+
|
| 418 |
+
# Valid mask: depth > 0 and not in aliasing region
|
| 419 |
+
valid_mask = gt_depth > 0
|
| 420 |
+
if aliasing_mask is not None:
|
| 421 |
+
valid_mask = np.logical_and(valid_mask, np.logical_not(aliasing_mask))
|
| 422 |
+
|
| 423 |
+
return valid_mask
|
| 424 |
+
|
| 425 |
+
def _mask_invalid_depth(
|
| 426 |
+
self, depth: np.ndarray, gt_zero_mask: np.ndarray = None
|
| 427 |
+
) -> np.ndarray:
|
| 428 |
+
"""Mask invalid depth values by setting them to 0."""
|
| 429 |
+
depth = depth.copy()
|
| 430 |
+
|
| 431 |
+
if gt_zero_mask is not None:
|
| 432 |
+
pred_invalid = np.isnan(depth) | np.isinf(depth)
|
| 433 |
+
combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid))
|
| 434 |
+
depth = depth * combined_mask.astype(np.float32)
|
| 435 |
+
else:
|
| 436 |
+
invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0)
|
| 437 |
+
depth[invalid_mask] = 0.0
|
| 438 |
+
|
| 439 |
+
return depth
|
| 440 |
+
|
depth_anything_3/bench/datasets/scannetpp.py
ADDED
|
@@ -0,0 +1,591 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
ScanNet++ Benchmark dataset implementation.
|
| 17 |
+
|
| 18 |
+
ScanNet++ is a high-quality indoor RGB-D dataset with iPhone and DSLR images,
|
| 19 |
+
ground truth camera poses from COLMAP, and high-resolution 3D meshes.
|
| 20 |
+
Reference: https://kaldir.vc.in.tum.de/scannetpp/
|
| 21 |
+
|
| 22 |
+
Evaluation metrics:
|
| 23 |
+
- 3D reconstruction: Accuracy, Completeness, F-score
|
| 24 |
+
- Camera pose estimation: AUC metrics
|
| 25 |
+
"""
|
| 26 |
+
|
| 27 |
+
import os
|
| 28 |
+
from typing import Dict as TDict
|
| 29 |
+
|
| 30 |
+
import cv2
|
| 31 |
+
import imageio
|
| 32 |
+
import numpy as np
|
| 33 |
+
import open3d as o3d
|
| 34 |
+
from addict import Dict
|
| 35 |
+
|
| 36 |
+
from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready
|
| 37 |
+
from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY
|
| 38 |
+
from depth_anything_3.bench.utils import (
|
| 39 |
+
create_tsdf_volume,
|
| 40 |
+
fuse_depth_to_tsdf,
|
| 41 |
+
nn_correspondance,
|
| 42 |
+
sample_points_from_mesh,
|
| 43 |
+
)
|
| 44 |
+
from depth_anything_3.utils.constants import (
|
| 45 |
+
SCANNETPP_DOWN_SAMPLE,
|
| 46 |
+
SCANNETPP_EVAL_DATA_ROOT,
|
| 47 |
+
SCANNETPP_EVAL_THRESHOLD,
|
| 48 |
+
SCANNETPP_INPUT_H,
|
| 49 |
+
SCANNETPP_INPUT_W,
|
| 50 |
+
SCANNETPP_MAX_DEPTH,
|
| 51 |
+
SCANNETPP_SAMPLING_NUMBER,
|
| 52 |
+
SCANNETPP_SCENES,
|
| 53 |
+
SCANNETPP_SDF_TRUNC,
|
| 54 |
+
SCANNETPP_VOXEL_LENGTH,
|
| 55 |
+
)
|
| 56 |
+
from depth_anything_3.utils.pose_align import align_poses_umeyama
|
| 57 |
+
from depth_anything_3.utils.read_write_model import read_model
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
@MV_REGISTRY.register(name="scannetpp")
|
| 61 |
+
@MONO_REGISTRY.register(name="scannetpp")
|
| 62 |
+
class ScanNetPP(Dataset):
|
| 63 |
+
"""
|
| 64 |
+
ScanNet++ Benchmark dataset wrapper for DepthAnything3 evaluation.
|
| 65 |
+
|
| 66 |
+
Supports:
|
| 67 |
+
- Camera pose estimation evaluation (AUC metrics)
|
| 68 |
+
- 3D reconstruction evaluation (Accuracy, Completeness, F-score)
|
| 69 |
+
- TSDF-based point cloud fusion
|
| 70 |
+
|
| 71 |
+
Dataset structure:
|
| 72 |
+
scannetpp/data/
|
| 73 |
+
├── {scene_id}/
|
| 74 |
+
│ ├── merge_dslr_iphone/
|
| 75 |
+
│ │ ├── colmap/sparse_render_rgb/ # COLMAP reconstruction
|
| 76 |
+
│ │ ├── images/ # RGB images
|
| 77 |
+
│ │ └── render_depth/ # GT depth maps
|
| 78 |
+
│ └── scans/
|
| 79 |
+
│ └── mesh_aligned_0.05.ply # Ground truth mesh
|
| 80 |
+
"""
|
| 81 |
+
|
| 82 |
+
data_root = SCANNETPP_EVAL_DATA_ROOT
|
| 83 |
+
SCENES = SCANNETPP_SCENES
|
| 84 |
+
|
| 85 |
+
# Input resolution after undistortion and resize
|
| 86 |
+
input_h = SCANNETPP_INPUT_H
|
| 87 |
+
input_w = SCANNETPP_INPUT_W
|
| 88 |
+
|
| 89 |
+
# Evaluation hyperparameters from constants
|
| 90 |
+
max_depth = SCANNETPP_MAX_DEPTH
|
| 91 |
+
sampling_number = SCANNETPP_SAMPLING_NUMBER
|
| 92 |
+
voxel_length = SCANNETPP_VOXEL_LENGTH
|
| 93 |
+
sdf_trunc = SCANNETPP_SDF_TRUNC
|
| 94 |
+
eval_threshold = SCANNETPP_EVAL_THRESHOLD
|
| 95 |
+
down_sample = SCANNETPP_DOWN_SAMPLE
|
| 96 |
+
|
| 97 |
+
def __init__(self):
|
| 98 |
+
super().__init__()
|
| 99 |
+
self._scene_cache = {}
|
| 100 |
+
|
| 101 |
+
# ------------------------------
|
| 102 |
+
# Public API
|
| 103 |
+
# ------------------------------
|
| 104 |
+
|
| 105 |
+
def get_data(self, scene: str) -> Dict:
|
| 106 |
+
"""
|
| 107 |
+
Collect per-view image paths, intrinsics/extrinsics for a scene.
|
| 108 |
+
|
| 109 |
+
Only uses iPhone images (not DSLR).
|
| 110 |
+
|
| 111 |
+
Args:
|
| 112 |
+
scene: Scene identifier (e.g., "09c1414f1b")
|
| 113 |
+
|
| 114 |
+
Returns:
|
| 115 |
+
Dict with:
|
| 116 |
+
- image_files: List[str] - paths to images
|
| 117 |
+
- extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms
|
| 118 |
+
- intrinsics: np.ndarray [N, 3, 3] - camera intrinsics
|
| 119 |
+
- aux: Dict with gt_mesh_path, dist, roi, cam_hw, etc.
|
| 120 |
+
"""
|
| 121 |
+
if scene in self._scene_cache:
|
| 122 |
+
return self._scene_cache[scene]
|
| 123 |
+
|
| 124 |
+
input_path = os.path.join(self.data_root, scene, "merge_dslr_iphone")
|
| 125 |
+
colmap_path = os.path.join(input_path, "colmap/sparse_render_rgb")
|
| 126 |
+
image_path = os.path.join(input_path, "images")
|
| 127 |
+
depth_path_dir = os.path.join(input_path, "render_depth")
|
| 128 |
+
|
| 129 |
+
# Read COLMAP model
|
| 130 |
+
cams, images, points3d = read_model(colmap_path)
|
| 131 |
+
|
| 132 |
+
# Map image names to IDs
|
| 133 |
+
name2id = {image.name: k for k, image in images.items()}
|
| 134 |
+
names = sorted([image.name for k, image in images.items()])
|
| 135 |
+
# Only use iPhone images
|
| 136 |
+
names = [name for name in names if "iphone" in name]
|
| 137 |
+
|
| 138 |
+
gt_mesh_path = os.path.join(
|
| 139 |
+
input_path.replace("merge_dslr_iphone", "scans"), "mesh_aligned_0.05.ply"
|
| 140 |
+
)
|
| 141 |
+
|
| 142 |
+
out = Dict({
|
| 143 |
+
"image_files": [],
|
| 144 |
+
"extrinsics": [],
|
| 145 |
+
"intrinsics": [],
|
| 146 |
+
"aux": Dict({
|
| 147 |
+
"gt_mesh_path": gt_mesh_path,
|
| 148 |
+
"dist_list": [],
|
| 149 |
+
"roi_list": [],
|
| 150 |
+
"cam_hw_list": [],
|
| 151 |
+
"ixt_raw_list": [],
|
| 152 |
+
"gt_depth_files": [],
|
| 153 |
+
}),
|
| 154 |
+
})
|
| 155 |
+
|
| 156 |
+
for name in names:
|
| 157 |
+
image = images[name2id[name]]
|
| 158 |
+
img_path = os.path.join(image_path, name)
|
| 159 |
+
|
| 160 |
+
if not os.path.exists(img_path):
|
| 161 |
+
continue
|
| 162 |
+
|
| 163 |
+
# Build extrinsics (world-to-camera)
|
| 164 |
+
ext = np.eye(4, dtype=np.float32)
|
| 165 |
+
ext[:3, :3] = image.qvec2rotmat()
|
| 166 |
+
ext[:3, 3] = image.tvec
|
| 167 |
+
|
| 168 |
+
# Get camera parameters
|
| 169 |
+
cam_id = image.camera_id
|
| 170 |
+
camera = cams[cam_id]
|
| 171 |
+
cam_height, cam_width = camera.height, camera.width
|
| 172 |
+
|
| 173 |
+
# Build intrinsics
|
| 174 |
+
ixt = np.eye(3, dtype=np.float32)
|
| 175 |
+
ixt[0, 0], ixt[1, 1], ixt[0, 2], ixt[1, 2] = camera.params[:4]
|
| 176 |
+
ixt[:2, 2] -= 0.5 # COLMAP convention adjustment
|
| 177 |
+
ixt_raw = ixt.copy()
|
| 178 |
+
|
| 179 |
+
# Handle distortion (OPENCV model)
|
| 180 |
+
dist = np.zeros(5, dtype=np.float32)
|
| 181 |
+
roi = (0, 0, cam_width, cam_height)
|
| 182 |
+
if camera.model == "OPENCV":
|
| 183 |
+
dist[:4] = camera.params[4:]
|
| 184 |
+
ixt, roi = cv2.getOptimalNewCameraMatrix(
|
| 185 |
+
ixt, dist, (cam_width, cam_height), 1, (cam_width, cam_height)
|
| 186 |
+
)
|
| 187 |
+
|
| 188 |
+
# Depth file path
|
| 189 |
+
frame_name = os.path.basename(name)[:-4] # Remove .jpg
|
| 190 |
+
depth_file = os.path.join(depth_path_dir, f"{frame_name}.png")
|
| 191 |
+
|
| 192 |
+
out.image_files.append(img_path)
|
| 193 |
+
out.extrinsics.append(ext)
|
| 194 |
+
out.intrinsics.append(ixt)
|
| 195 |
+
out.aux.dist_list.append(dist)
|
| 196 |
+
out.aux.roi_list.append(roi)
|
| 197 |
+
out.aux.cam_hw_list.append((cam_height, cam_width))
|
| 198 |
+
out.aux.ixt_raw_list.append(ixt_raw)
|
| 199 |
+
out.aux.gt_depth_files.append(depth_file)
|
| 200 |
+
|
| 201 |
+
out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32)
|
| 202 |
+
out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32)
|
| 203 |
+
|
| 204 |
+
print(f"[ScanNet++] {scene}: {len(out.image_files)} images")
|
| 205 |
+
|
| 206 |
+
self._scene_cache[scene] = out
|
| 207 |
+
return out
|
| 208 |
+
|
| 209 |
+
def load_image(self, img_path: str, idx: int, aux: Dict) -> np.ndarray:
|
| 210 |
+
"""
|
| 211 |
+
Load and preprocess image with undistortion and cropping.
|
| 212 |
+
|
| 213 |
+
Args:
|
| 214 |
+
img_path: Path to image file
|
| 215 |
+
idx: Index of the image in the dataset
|
| 216 |
+
aux: Auxiliary data from get_data
|
| 217 |
+
|
| 218 |
+
Returns:
|
| 219 |
+
Preprocessed RGB image
|
| 220 |
+
"""
|
| 221 |
+
image = imageio.imread(img_path).astype(np.uint8)
|
| 222 |
+
ixt_raw = aux.ixt_raw_list[idx]
|
| 223 |
+
ixt = aux.intrinsics[idx] if hasattr(aux, 'intrinsics') else None
|
| 224 |
+
dist = aux.dist_list[idx]
|
| 225 |
+
roi = aux.roi_list[idx]
|
| 226 |
+
|
| 227 |
+
# Undistort using raw intrinsics
|
| 228 |
+
# Use the stored intrinsics from get_data for newCameraMatrix
|
| 229 |
+
stored_ixt = self._scene_cache.get(aux.scene, {}).get('intrinsics', [None])[idx] if hasattr(aux, 'scene') else None
|
| 230 |
+
if stored_ixt is None:
|
| 231 |
+
# Recompute optimal camera matrix for undistortion
|
| 232 |
+
cam_h, cam_w = aux.cam_hw_list[idx]
|
| 233 |
+
ixt_for_undistort = ixt_raw.copy()
|
| 234 |
+
ixt_for_undistort, _ = cv2.getOptimalNewCameraMatrix(
|
| 235 |
+
ixt_raw, dist, (cam_w, cam_h), 1, (cam_w, cam_h)
|
| 236 |
+
)
|
| 237 |
+
else:
|
| 238 |
+
ixt_for_undistort = stored_ixt
|
| 239 |
+
|
| 240 |
+
image = cv2.undistort(image, ixt_raw, dist, newCameraMatrix=ixt_for_undistort)
|
| 241 |
+
|
| 242 |
+
# Crop to ROI
|
| 243 |
+
x, y, w, h = roi
|
| 244 |
+
image = image[y:y+h, x:x+w]
|
| 245 |
+
|
| 246 |
+
# Resize to target resolution
|
| 247 |
+
image = cv2.resize(image, (self.input_w, self.input_h), interpolation=cv2.INTER_AREA)
|
| 248 |
+
|
| 249 |
+
return image
|
| 250 |
+
|
| 251 |
+
def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]:
|
| 252 |
+
"""
|
| 253 |
+
Evaluate fused point cloud against ScanNet++ ground truth mesh.
|
| 254 |
+
|
| 255 |
+
Uses AABB cropping to only evaluate points within GT bounding box.
|
| 256 |
+
|
| 257 |
+
Args:
|
| 258 |
+
scene: Scene identifier
|
| 259 |
+
fuse_path: Path to fused point cloud (.ply)
|
| 260 |
+
|
| 261 |
+
Returns:
|
| 262 |
+
Dict with metrics: acc, comp, overall, precision, recall, fscore
|
| 263 |
+
"""
|
| 264 |
+
gt_data = self.get_data(scene)
|
| 265 |
+
gt_mesh_path = gt_data.aux.gt_mesh_path
|
| 266 |
+
|
| 267 |
+
# Load ground truth mesh and sample points
|
| 268 |
+
gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path)
|
| 269 |
+
gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number)
|
| 270 |
+
|
| 271 |
+
# Load predicted point cloud
|
| 272 |
+
pred_pcd = o3d.io.read_point_cloud(fuse_path)
|
| 273 |
+
|
| 274 |
+
# Crop prediction to GT bounding box (with 0.1m margin)
|
| 275 |
+
aabb = gt_pcd.get_axis_aligned_bounding_box()
|
| 276 |
+
points = np.asarray(pred_pcd.points)
|
| 277 |
+
inside_mask = (
|
| 278 |
+
(points[:, 0] >= aabb.min_bound[0] - 0.1) &
|
| 279 |
+
(points[:, 0] <= aabb.max_bound[0] + 0.1) &
|
| 280 |
+
(points[:, 1] >= aabb.min_bound[1] - 0.1) &
|
| 281 |
+
(points[:, 1] <= aabb.max_bound[1] + 0.1) &
|
| 282 |
+
(points[:, 2] >= aabb.min_bound[2] - 0.1) &
|
| 283 |
+
(points[:, 2] <= aabb.max_bound[2] + 0.1)
|
| 284 |
+
)
|
| 285 |
+
pred_pcd = pred_pcd.select_by_index(inside_mask.nonzero()[0])
|
| 286 |
+
|
| 287 |
+
# Downsample
|
| 288 |
+
if self.down_sample > 0:
|
| 289 |
+
pred_pcd = pred_pcd.voxel_down_sample(self.down_sample)
|
| 290 |
+
gt_pcd = gt_pcd.voxel_down_sample(self.down_sample)
|
| 291 |
+
|
| 292 |
+
verts_pred = np.asarray(pred_pcd.points)
|
| 293 |
+
verts_gt = np.asarray(gt_pcd.points)
|
| 294 |
+
|
| 295 |
+
if len(verts_pred) == 0 or len(verts_gt) == 0:
|
| 296 |
+
return {
|
| 297 |
+
"acc": float("inf"),
|
| 298 |
+
"comp": float("inf"),
|
| 299 |
+
"overall": float("inf"),
|
| 300 |
+
"precision": 0.0,
|
| 301 |
+
"recall": 0.0,
|
| 302 |
+
"fscore": 0.0,
|
| 303 |
+
}
|
| 304 |
+
|
| 305 |
+
# Compute distances
|
| 306 |
+
dist_pred_to_gt = nn_correspondance(verts_gt, verts_pred)
|
| 307 |
+
dist_gt_to_pred = nn_correspondance(verts_pred, verts_gt)
|
| 308 |
+
|
| 309 |
+
# Compute metrics
|
| 310 |
+
accuracy = float(np.mean(dist_pred_to_gt))
|
| 311 |
+
completeness = float(np.mean(dist_gt_to_pred))
|
| 312 |
+
overall = (accuracy + completeness) / 2
|
| 313 |
+
|
| 314 |
+
precision = float(np.mean((dist_pred_to_gt < self.eval_threshold).astype(float)))
|
| 315 |
+
recall = float(np.mean((dist_gt_to_pred < self.eval_threshold).astype(float)))
|
| 316 |
+
|
| 317 |
+
if precision + recall > 0:
|
| 318 |
+
fscore = 2 * precision * recall / (precision + recall)
|
| 319 |
+
else:
|
| 320 |
+
fscore = 0.0
|
| 321 |
+
|
| 322 |
+
return {
|
| 323 |
+
"acc": accuracy,
|
| 324 |
+
"comp": completeness,
|
| 325 |
+
"overall": overall,
|
| 326 |
+
"precision": precision,
|
| 327 |
+
"recall": recall,
|
| 328 |
+
"fscore": fscore,
|
| 329 |
+
}
|
| 330 |
+
|
| 331 |
+
def _load_gt_meta(self, result_path: str) -> Dict:
|
| 332 |
+
"""Load saved GT meta for fusion."""
|
| 333 |
+
export_dir = os.path.dirname(result_path)
|
| 334 |
+
gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz")
|
| 335 |
+
|
| 336 |
+
if os.path.exists(gt_meta_path):
|
| 337 |
+
data = np.load(gt_meta_path, allow_pickle=True)
|
| 338 |
+
image_files = list(data["image_files"])
|
| 339 |
+
|
| 340 |
+
# Reconstruct aux data from image files
|
| 341 |
+
return Dict({
|
| 342 |
+
"extrinsics": data["extrinsics"],
|
| 343 |
+
"intrinsics": data["intrinsics"],
|
| 344 |
+
"image_files": image_files,
|
| 345 |
+
})
|
| 346 |
+
return None
|
| 347 |
+
|
| 348 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 349 |
+
"""
|
| 350 |
+
Fuse per-view depths into a point cloud using TSDF fusion.
|
| 351 |
+
|
| 352 |
+
Args:
|
| 353 |
+
scene: Scene identifier
|
| 354 |
+
result_path: Path to npz file with predicted depths/poses
|
| 355 |
+
fuse_path: Output path for fused point cloud (.ply)
|
| 356 |
+
mode: "recon_unposed" or "recon_posed"
|
| 357 |
+
"""
|
| 358 |
+
# Get GT data
|
| 359 |
+
full_gt_data = self.get_data(scene)
|
| 360 |
+
|
| 361 |
+
# Try to load saved GT meta (handles frame sampling)
|
| 362 |
+
gt_meta = self._load_gt_meta(result_path)
|
| 363 |
+
if gt_meta is not None:
|
| 364 |
+
gt_data = gt_meta
|
| 365 |
+
# Need to rebuild aux from full GT data based on image indices
|
| 366 |
+
image_indices = [
|
| 367 |
+
full_gt_data.image_files.index(f)
|
| 368 |
+
for f in gt_data.image_files
|
| 369 |
+
if f in full_gt_data.image_files
|
| 370 |
+
]
|
| 371 |
+
else:
|
| 372 |
+
gt_data = full_gt_data
|
| 373 |
+
image_indices = list(range(len(full_gt_data.image_files)))
|
| 374 |
+
|
| 375 |
+
_wait_for_file_ready(result_path)
|
| 376 |
+
pred_data = Dict({k: v for k, v in np.load(result_path).items()})
|
| 377 |
+
|
| 378 |
+
# Load and preprocess images
|
| 379 |
+
images = []
|
| 380 |
+
for idx, img_idx in enumerate(image_indices):
|
| 381 |
+
img_path = full_gt_data.image_files[img_idx]
|
| 382 |
+
image = imageio.imread(img_path).astype(np.uint8)
|
| 383 |
+
|
| 384 |
+
# Undistort and crop
|
| 385 |
+
ixt_raw = full_gt_data.aux.ixt_raw_list[img_idx]
|
| 386 |
+
ixt = full_gt_data.intrinsics[img_idx]
|
| 387 |
+
dist = full_gt_data.aux.dist_list[img_idx]
|
| 388 |
+
roi = full_gt_data.aux.roi_list[img_idx]
|
| 389 |
+
|
| 390 |
+
image = cv2.undistort(image, ixt_raw, dist, newCameraMatrix=ixt)
|
| 391 |
+
x, y, w, h = roi
|
| 392 |
+
image = image[y:y+h, x:x+w]
|
| 393 |
+
image = cv2.resize(image, (self.input_w, self.input_h), interpolation=cv2.INTER_AREA)
|
| 394 |
+
|
| 395 |
+
images.append(image)
|
| 396 |
+
|
| 397 |
+
images = np.stack(images, axis=0)
|
| 398 |
+
|
| 399 |
+
# Prepare depths, intrinsics, extrinsics
|
| 400 |
+
if mode == "recon_unposed":
|
| 401 |
+
depths, intrinsics, extrinsics = self._prep_unposed(
|
| 402 |
+
pred_data, gt_data, full_gt_data, image_indices, scene=scene
|
| 403 |
+
)
|
| 404 |
+
elif mode == "recon_posed":
|
| 405 |
+
depths, intrinsics, extrinsics = self._prep_posed(
|
| 406 |
+
pred_data, gt_data, full_gt_data, image_indices, scene=scene
|
| 407 |
+
)
|
| 408 |
+
else:
|
| 409 |
+
raise ValueError(f"Invalid mode: {mode}")
|
| 410 |
+
|
| 411 |
+
# Create TSDF volume and fuse
|
| 412 |
+
volume = create_tsdf_volume(
|
| 413 |
+
voxel_length=self.voxel_length,
|
| 414 |
+
sdf_trunc=self.sdf_trunc,
|
| 415 |
+
)
|
| 416 |
+
mesh = fuse_depth_to_tsdf(
|
| 417 |
+
volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth
|
| 418 |
+
)
|
| 419 |
+
|
| 420 |
+
# Sample points from mesh
|
| 421 |
+
pcd = sample_points_from_mesh(mesh, self.sampling_number)
|
| 422 |
+
|
| 423 |
+
# Save point cloud
|
| 424 |
+
os.makedirs(os.path.dirname(fuse_path), exist_ok=True)
|
| 425 |
+
o3d.io.write_point_cloud(fuse_path, pcd)
|
| 426 |
+
|
| 427 |
+
# ------------------------------
|
| 428 |
+
# Private helpers
|
| 429 |
+
# ------------------------------
|
| 430 |
+
|
| 431 |
+
def _prep_unposed(
|
| 432 |
+
self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict,
|
| 433 |
+
image_indices: list, scene: str = None
|
| 434 |
+
) -> tuple:
|
| 435 |
+
"""Prepare depths/intrinsics/extrinsics for recon_unposed mode."""
|
| 436 |
+
# Scale alignment with fixed random_state for reproducibility
|
| 437 |
+
_, _, scale, extrinsics = align_poses_umeyama(
|
| 438 |
+
gt_data.extrinsics.copy(),
|
| 439 |
+
pred_data.extrinsics.copy(),
|
| 440 |
+
return_aligned=True,
|
| 441 |
+
ransac=True,
|
| 442 |
+
random_state=42,
|
| 443 |
+
)
|
| 444 |
+
|
| 445 |
+
model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2]
|
| 446 |
+
|
| 447 |
+
depths_out = []
|
| 448 |
+
intrinsics_out = []
|
| 449 |
+
for i in range(len(pred_data.depth)):
|
| 450 |
+
img_idx = image_indices[i]
|
| 451 |
+
|
| 452 |
+
# Get original image size (after undistort+crop, before resize to input_h/w)
|
| 453 |
+
orig_h, orig_w = full_gt_data.aux.cam_hw_list[img_idx]
|
| 454 |
+
|
| 455 |
+
# Step 1: nearest resize to original image size
|
| 456 |
+
depth = cv2.resize(
|
| 457 |
+
pred_data.depth[i],
|
| 458 |
+
(orig_w, orig_h),
|
| 459 |
+
interpolation=cv2.INTER_NEAREST,
|
| 460 |
+
)
|
| 461 |
+
|
| 462 |
+
# Step 2: linear resize to target resolution
|
| 463 |
+
depth = cv2.resize(
|
| 464 |
+
depth,
|
| 465 |
+
(self.input_w, self.input_h),
|
| 466 |
+
interpolation=cv2.INTER_LINEAR,
|
| 467 |
+
).astype(np.float32)
|
| 468 |
+
|
| 469 |
+
# Load GT depth for masking
|
| 470 |
+
gt_zero_mask = self._load_gt_mask(full_gt_data.aux.gt_depth_files[img_idx])
|
| 471 |
+
|
| 472 |
+
# Mask invalid depths BEFORE scale
|
| 473 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 474 |
+
|
| 475 |
+
# Apply scale AFTER mask
|
| 476 |
+
depth = depth * scale
|
| 477 |
+
|
| 478 |
+
# Adjust intrinsics to target resolution
|
| 479 |
+
h_ratio = self.input_h / model_h
|
| 480 |
+
w_ratio = self.input_w / model_w
|
| 481 |
+
ixt = pred_data.intrinsics[i].copy()
|
| 482 |
+
ixt[0, :] *= w_ratio
|
| 483 |
+
ixt[1, :] *= h_ratio
|
| 484 |
+
|
| 485 |
+
depths_out.append(depth)
|
| 486 |
+
intrinsics_out.append(ixt)
|
| 487 |
+
|
| 488 |
+
return np.stack(depths_out), np.stack(intrinsics_out), extrinsics
|
| 489 |
+
|
| 490 |
+
def _prep_posed(
|
| 491 |
+
self, pred_data: Dict, gt_data: Dict, full_gt_data: Dict,
|
| 492 |
+
image_indices: list, scene: str = None
|
| 493 |
+
) -> tuple:
|
| 494 |
+
"""Prepare depths/intrinsics/extrinsics for recon_posed mode."""
|
| 495 |
+
# Scale alignment
|
| 496 |
+
_, _, scale, _ = align_poses_umeyama(
|
| 497 |
+
gt_data.extrinsics.copy(),
|
| 498 |
+
pred_data.extrinsics.copy(),
|
| 499 |
+
return_aligned=True,
|
| 500 |
+
ransac=True,
|
| 501 |
+
random_state=42,
|
| 502 |
+
)
|
| 503 |
+
|
| 504 |
+
depths_out = []
|
| 505 |
+
intrinsics_out = []
|
| 506 |
+
extrinsics_out = []
|
| 507 |
+
|
| 508 |
+
for i in range(len(pred_data.depth)):
|
| 509 |
+
img_idx = image_indices[i]
|
| 510 |
+
|
| 511 |
+
# Get original image size (after undistort+crop, before resize to input_h/w)
|
| 512 |
+
orig_h, orig_w = full_gt_data.aux.cam_hw_list[img_idx]
|
| 513 |
+
|
| 514 |
+
# Step 1: nearest resize to original image size
|
| 515 |
+
depth = cv2.resize(
|
| 516 |
+
pred_data.depth[i],
|
| 517 |
+
(orig_w, orig_h),
|
| 518 |
+
interpolation=cv2.INTER_NEAREST,
|
| 519 |
+
)
|
| 520 |
+
|
| 521 |
+
# Step 2: linear resize to target resolution
|
| 522 |
+
depth = cv2.resize(
|
| 523 |
+
depth,
|
| 524 |
+
(self.input_w, self.input_h),
|
| 525 |
+
interpolation=cv2.INTER_LINEAR,
|
| 526 |
+
).astype(np.float32)
|
| 527 |
+
|
| 528 |
+
# Load GT depth for masking
|
| 529 |
+
gt_zero_mask = self._load_gt_mask(full_gt_data.aux.gt_depth_files[img_idx])
|
| 530 |
+
|
| 531 |
+
# Mask invalid depths BEFORE scale
|
| 532 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 533 |
+
|
| 534 |
+
# Apply scale AFTER mask
|
| 535 |
+
depth = depth * scale
|
| 536 |
+
|
| 537 |
+
depths_out.append(depth)
|
| 538 |
+
|
| 539 |
+
# Get GT intrinsics and scale to target resolution
|
| 540 |
+
ixt = full_gt_data.intrinsics[img_idx].copy()
|
| 541 |
+
cam_h, cam_w = full_gt_data.aux.cam_hw_list[img_idx]
|
| 542 |
+
ixt[:2, 2] += 0.5 # Undo COLMAP convention
|
| 543 |
+
ixt[0, :] *= self.input_w / cam_w
|
| 544 |
+
ixt[1, :] *= self.input_h / cam_h
|
| 545 |
+
intrinsics_out.append(ixt)
|
| 546 |
+
|
| 547 |
+
extrinsics_out.append(full_gt_data.extrinsics[img_idx])
|
| 548 |
+
|
| 549 |
+
return np.stack(depths_out), np.stack(intrinsics_out), np.stack(extrinsics_out)
|
| 550 |
+
|
| 551 |
+
def _load_gt_mask(self, gt_depth_path: str) -> np.ndarray:
|
| 552 |
+
"""
|
| 553 |
+
Load GT depth and create valid mask.
|
| 554 |
+
|
| 555 |
+
For ScanNet++, GT depth is stored as 16-bit PNG in millimeters.
|
| 556 |
+
|
| 557 |
+
Returns:
|
| 558 |
+
Boolean mask where True = valid region to keep
|
| 559 |
+
"""
|
| 560 |
+
if not os.path.exists(gt_depth_path):
|
| 561 |
+
return None
|
| 562 |
+
|
| 563 |
+
gt_depth = imageio.imread(gt_depth_path) / 1000.0 # mm to meters
|
| 564 |
+
|
| 565 |
+
# Resize to target resolution
|
| 566 |
+
gt_depth = cv2.resize(
|
| 567 |
+
gt_depth,
|
| 568 |
+
(self.input_w, self.input_h),
|
| 569 |
+
interpolation=cv2.INTER_LINEAR,
|
| 570 |
+
).astype(np.float32)
|
| 571 |
+
|
| 572 |
+
# Valid mask: depth > 0 and not inf
|
| 573 |
+
valid_mask = np.logical_and(gt_depth > 0, gt_depth != np.inf)
|
| 574 |
+
return valid_mask
|
| 575 |
+
|
| 576 |
+
def _mask_invalid_depth(
|
| 577 |
+
self, depth: np.ndarray, gt_zero_mask: np.ndarray = None
|
| 578 |
+
) -> np.ndarray:
|
| 579 |
+
"""Mask invalid depth values by setting them to 0."""
|
| 580 |
+
depth = depth.copy()
|
| 581 |
+
|
| 582 |
+
if gt_zero_mask is not None:
|
| 583 |
+
pred_invalid = np.isnan(depth) | np.isinf(depth)
|
| 584 |
+
combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid))
|
| 585 |
+
depth = depth * combined_mask.astype(np.float32)
|
| 586 |
+
else:
|
| 587 |
+
invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0)
|
| 588 |
+
depth[invalid_mask] = 0.0
|
| 589 |
+
|
| 590 |
+
return depth
|
| 591 |
+
|
depth_anything_3/bench/datasets/sevenscenes.py
ADDED
|
@@ -0,0 +1,449 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
7Scenes Benchmark dataset implementation.
|
| 17 |
+
|
| 18 |
+
7Scenes is an indoor RGB-D dataset with ground truth camera poses and 3D meshes.
|
| 19 |
+
Reference: https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/
|
| 20 |
+
|
| 21 |
+
Evaluation metrics:
|
| 22 |
+
- 3D reconstruction: Accuracy, Completeness, F-score
|
| 23 |
+
- Camera pose estimation: AUC metrics
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
import os
|
| 27 |
+
from typing import Dict as TDict
|
| 28 |
+
|
| 29 |
+
import cv2
|
| 30 |
+
import numpy as np
|
| 31 |
+
import open3d as o3d
|
| 32 |
+
from addict import Dict
|
| 33 |
+
|
| 34 |
+
from depth_anything_3.bench.dataset import Dataset, _wait_for_file_ready
|
| 35 |
+
from depth_anything_3.bench.registries import MONO_REGISTRY, MV_REGISTRY
|
| 36 |
+
from depth_anything_3.bench.utils import (
|
| 37 |
+
create_tsdf_volume,
|
| 38 |
+
evaluate_3d_reconstruction,
|
| 39 |
+
fuse_depth_to_tsdf,
|
| 40 |
+
sample_points_from_mesh,
|
| 41 |
+
)
|
| 42 |
+
from depth_anything_3.utils.constants import (
|
| 43 |
+
SEVENSCENES_CX,
|
| 44 |
+
SEVENSCENES_CY,
|
| 45 |
+
SEVENSCENES_DOWN_SAMPLE,
|
| 46 |
+
SEVENSCENES_EVAL_DATA_ROOT,
|
| 47 |
+
SEVENSCENES_EVAL_THRESHOLD,
|
| 48 |
+
SEVENSCENES_FX,
|
| 49 |
+
SEVENSCENES_FY,
|
| 50 |
+
SEVENSCENES_MAX_DEPTH,
|
| 51 |
+
SEVENSCENES_SAMPLING_NUMBER,
|
| 52 |
+
SEVENSCENES_SCENES,
|
| 53 |
+
SEVENSCENES_SDF_TRUNC,
|
| 54 |
+
SEVENSCENES_VOXEL_LENGTH,
|
| 55 |
+
)
|
| 56 |
+
from depth_anything_3.utils.pose_align import align_poses_umeyama
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
@MV_REGISTRY.register(name="7scenes")
|
| 60 |
+
@MONO_REGISTRY.register(name="7scenes")
|
| 61 |
+
class SevenScenes(Dataset):
|
| 62 |
+
"""
|
| 63 |
+
7Scenes Benchmark dataset wrapper for DepthAnything3 evaluation.
|
| 64 |
+
|
| 65 |
+
Supports:
|
| 66 |
+
- Camera pose estimation evaluation (AUC metrics)
|
| 67 |
+
- 3D reconstruction evaluation (Accuracy, Completeness, F-score)
|
| 68 |
+
- TSDF-based point cloud fusion
|
| 69 |
+
|
| 70 |
+
Dataset structure:
|
| 71 |
+
7scenes/
|
| 72 |
+
├── 7Scenes/
|
| 73 |
+
│ ├── {scene}/
|
| 74 |
+
│ │ └── seq-01/ (or seq-02 for stairs)
|
| 75 |
+
│ │ ├── frame-XXXXXX.color.png
|
| 76 |
+
│ │ ├── frame-XXXXXX.depth.png
|
| 77 |
+
│ │ └── frame-XXXXXX.pose.txt
|
| 78 |
+
│ └── meshes/
|
| 79 |
+
│ └── {scene}.ply # Ground truth mesh
|
| 80 |
+
"""
|
| 81 |
+
|
| 82 |
+
data_root = SEVENSCENES_EVAL_DATA_ROOT
|
| 83 |
+
SCENES = SEVENSCENES_SCENES
|
| 84 |
+
|
| 85 |
+
# Evaluation hyperparameters from constants
|
| 86 |
+
max_depth = SEVENSCENES_MAX_DEPTH
|
| 87 |
+
sampling_number = SEVENSCENES_SAMPLING_NUMBER
|
| 88 |
+
voxel_length = SEVENSCENES_VOXEL_LENGTH
|
| 89 |
+
sdf_trunc = SEVENSCENES_SDF_TRUNC
|
| 90 |
+
eval_threshold = SEVENSCENES_EVAL_THRESHOLD
|
| 91 |
+
down_sample = SEVENSCENES_DOWN_SAMPLE
|
| 92 |
+
|
| 93 |
+
# Fixed camera intrinsics for all 7Scenes images
|
| 94 |
+
fx = SEVENSCENES_FX
|
| 95 |
+
fy = SEVENSCENES_FY
|
| 96 |
+
cx = SEVENSCENES_CX
|
| 97 |
+
cy = SEVENSCENES_CY
|
| 98 |
+
|
| 99 |
+
def __init__(self):
|
| 100 |
+
super().__init__()
|
| 101 |
+
self._scene_cache = {}
|
| 102 |
+
|
| 103 |
+
# ------------------------------
|
| 104 |
+
# Public API
|
| 105 |
+
# ------------------------------
|
| 106 |
+
|
| 107 |
+
def get_data(self, scene: str) -> Dict:
|
| 108 |
+
"""
|
| 109 |
+
Collect per-view image paths, intrinsics/extrinsics for a scene.
|
| 110 |
+
|
| 111 |
+
Args:
|
| 112 |
+
scene: Scene identifier (e.g., "chess")
|
| 113 |
+
|
| 114 |
+
Returns:
|
| 115 |
+
Dict with:
|
| 116 |
+
- image_files: List[str] - paths to images
|
| 117 |
+
- extrinsics: np.ndarray [N, 4, 4] - world-to-camera transforms
|
| 118 |
+
- intrinsics: np.ndarray [N, 3, 3] - camera intrinsics
|
| 119 |
+
- aux: Dict with gt_mesh_path, gt_depth_files
|
| 120 |
+
"""
|
| 121 |
+
if scene in self._scene_cache:
|
| 122 |
+
return self._scene_cache[scene]
|
| 123 |
+
|
| 124 |
+
# Different sequence for stairs scene
|
| 125 |
+
if scene == "stairs":
|
| 126 |
+
data_folder = os.path.join(self.data_root, "7Scenes", scene, "seq-02")
|
| 127 |
+
n_imgs = 500
|
| 128 |
+
else:
|
| 129 |
+
data_folder = os.path.join(self.data_root, "7Scenes", scene, "seq-01")
|
| 130 |
+
n_imgs = 1000
|
| 131 |
+
|
| 132 |
+
gt_mesh_path = os.path.join(self.data_root, "7Scenes", "meshes", f"{scene}.ply")
|
| 133 |
+
|
| 134 |
+
# Fixed intrinsics for all images
|
| 135 |
+
ixt = np.array([
|
| 136 |
+
[self.fx, 0, self.cx],
|
| 137 |
+
[0, self.fy, self.cy],
|
| 138 |
+
[0, 0, 1],
|
| 139 |
+
], dtype=np.float32)
|
| 140 |
+
|
| 141 |
+
out = Dict({
|
| 142 |
+
"image_files": [],
|
| 143 |
+
"extrinsics": [],
|
| 144 |
+
"intrinsics": [],
|
| 145 |
+
"aux": Dict({
|
| 146 |
+
"gt_mesh_path": gt_mesh_path,
|
| 147 |
+
"gt_depth_files": [],
|
| 148 |
+
}),
|
| 149 |
+
})
|
| 150 |
+
|
| 151 |
+
for i in range(0, n_imgs, 1):
|
| 152 |
+
img_path = os.path.join(data_folder, f"frame-{i:06d}.color.png")
|
| 153 |
+
pose_path = os.path.join(data_folder, f"frame-{i:06d}.pose.txt")
|
| 154 |
+
depth_path = os.path.join(data_folder, f"frame-{i:06d}.depth.png")
|
| 155 |
+
|
| 156 |
+
if not os.path.exists(img_path) or not os.path.exists(pose_path):
|
| 157 |
+
continue
|
| 158 |
+
|
| 159 |
+
# Load camera-to-world pose and convert to world-to-camera (extrinsic)
|
| 160 |
+
c2w = np.loadtxt(pose_path)
|
| 161 |
+
ext = np.linalg.inv(c2w).astype(np.float32)
|
| 162 |
+
|
| 163 |
+
out.image_files.append(img_path)
|
| 164 |
+
out.extrinsics.append(ext)
|
| 165 |
+
out.intrinsics.append(ixt.copy())
|
| 166 |
+
out.aux.gt_depth_files.append(depth_path)
|
| 167 |
+
|
| 168 |
+
out.extrinsics = np.asarray(out.extrinsics, dtype=np.float32)
|
| 169 |
+
out.intrinsics = np.asarray(out.intrinsics, dtype=np.float32)
|
| 170 |
+
|
| 171 |
+
print(f"[7Scenes] {scene}: {len(out.image_files)} images")
|
| 172 |
+
|
| 173 |
+
self._scene_cache[scene] = out
|
| 174 |
+
return out
|
| 175 |
+
|
| 176 |
+
def eval3d(self, scene: str, fuse_path: str) -> TDict[str, float]:
|
| 177 |
+
"""
|
| 178 |
+
Evaluate fused point cloud against 7Scenes ground truth mesh.
|
| 179 |
+
|
| 180 |
+
Args:
|
| 181 |
+
scene: Scene identifier
|
| 182 |
+
fuse_path: Path to fused point cloud (.ply)
|
| 183 |
+
|
| 184 |
+
Returns:
|
| 185 |
+
Dict with metrics: acc, comp, overall, precision, recall, fscore
|
| 186 |
+
"""
|
| 187 |
+
gt_data = self.get_data(scene)
|
| 188 |
+
gt_mesh_path = gt_data.aux.gt_mesh_path
|
| 189 |
+
|
| 190 |
+
# Load and sample ground truth mesh
|
| 191 |
+
gt_mesh = o3d.io.read_triangle_mesh(gt_mesh_path)
|
| 192 |
+
gt_pcd = sample_points_from_mesh(gt_mesh, self.sampling_number)
|
| 193 |
+
|
| 194 |
+
# Load predicted point cloud
|
| 195 |
+
pred_pcd = o3d.io.read_point_cloud(fuse_path)
|
| 196 |
+
|
| 197 |
+
# Evaluate using shared utility function
|
| 198 |
+
metrics = evaluate_3d_reconstruction(
|
| 199 |
+
pred_pcd,
|
| 200 |
+
gt_pcd,
|
| 201 |
+
threshold=self.eval_threshold,
|
| 202 |
+
down_sample=self.down_sample,
|
| 203 |
+
)
|
| 204 |
+
|
| 205 |
+
return metrics
|
| 206 |
+
|
| 207 |
+
def _load_gt_meta(self, result_path: str) -> Dict:
|
| 208 |
+
"""
|
| 209 |
+
Load saved GT meta (extrinsics, intrinsics, image_files) for fusion.
|
| 210 |
+
|
| 211 |
+
This is needed when frames are sampled, so fuse3d uses the correct
|
| 212 |
+
(sampled) GT instead of full dataset GT.
|
| 213 |
+
|
| 214 |
+
Args:
|
| 215 |
+
result_path: Path to npz file (used to derive gt_meta.npz path)
|
| 216 |
+
|
| 217 |
+
Returns:
|
| 218 |
+
Dict with GT data, or None if gt_meta.npz doesn't exist
|
| 219 |
+
"""
|
| 220 |
+
export_dir = os.path.dirname(result_path) # exports/mini_npz/
|
| 221 |
+
gt_meta_path = os.path.join(os.path.dirname(export_dir), "gt_meta.npz")
|
| 222 |
+
|
| 223 |
+
if os.path.exists(gt_meta_path):
|
| 224 |
+
data = np.load(gt_meta_path, allow_pickle=True)
|
| 225 |
+
# Build aux with gt_depth_files derived from image_files
|
| 226 |
+
image_files = list(data["image_files"])
|
| 227 |
+
gt_depth_files = [
|
| 228 |
+
img_path.replace("color", "depth").replace(".color.", ".depth.")
|
| 229 |
+
for img_path in image_files
|
| 230 |
+
]
|
| 231 |
+
return Dict({
|
| 232 |
+
"extrinsics": data["extrinsics"],
|
| 233 |
+
"intrinsics": data["intrinsics"],
|
| 234 |
+
"image_files": image_files,
|
| 235 |
+
"aux": Dict({"gt_depth_files": gt_depth_files}),
|
| 236 |
+
})
|
| 237 |
+
return None
|
| 238 |
+
|
| 239 |
+
def fuse3d(self, scene: str, result_path: str, fuse_path: str, mode: str) -> None:
|
| 240 |
+
"""
|
| 241 |
+
Fuse per-view depths into a point cloud using TSDF fusion.
|
| 242 |
+
|
| 243 |
+
Args:
|
| 244 |
+
scene: Scene identifier
|
| 245 |
+
result_path: Path to npz file with predicted depths/poses
|
| 246 |
+
fuse_path: Output path for fused point cloud (.ply)
|
| 247 |
+
mode: "recon_unposed" or "recon_posed"
|
| 248 |
+
"""
|
| 249 |
+
# Try to load saved GT meta (handles frame sampling)
|
| 250 |
+
gt_meta = self._load_gt_meta(result_path)
|
| 251 |
+
if gt_meta is not None:
|
| 252 |
+
gt_data = gt_meta
|
| 253 |
+
else:
|
| 254 |
+
gt_data = self.get_data(scene)
|
| 255 |
+
_wait_for_file_ready(result_path)
|
| 256 |
+
pred_data = Dict({k: v for k, v in np.load(result_path).items()})
|
| 257 |
+
|
| 258 |
+
# Load original images (keep original size)
|
| 259 |
+
images = []
|
| 260 |
+
orig_sizes = []
|
| 261 |
+
for img_path in gt_data.image_files:
|
| 262 |
+
img = cv2.imread(img_path)
|
| 263 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 264 |
+
images.append(img)
|
| 265 |
+
orig_sizes.append((img.shape[0], img.shape[1]))
|
| 266 |
+
|
| 267 |
+
# Prepare depths, intrinsics, extrinsics
|
| 268 |
+
if mode == "recon_unposed":
|
| 269 |
+
depths, intrinsics, extrinsics = self._prep_unposed(
|
| 270 |
+
pred_data, gt_data, orig_sizes, scene=scene
|
| 271 |
+
)
|
| 272 |
+
elif mode == "recon_posed":
|
| 273 |
+
depths, intrinsics, extrinsics = self._prep_posed(
|
| 274 |
+
pred_data, gt_data, orig_sizes, scene=scene
|
| 275 |
+
)
|
| 276 |
+
else:
|
| 277 |
+
raise ValueError(f"Invalid mode: {mode}")
|
| 278 |
+
|
| 279 |
+
images = np.stack(images, axis=0)
|
| 280 |
+
|
| 281 |
+
# Create TSDF volume and fuse
|
| 282 |
+
volume = create_tsdf_volume(
|
| 283 |
+
voxel_length=self.voxel_length,
|
| 284 |
+
sdf_trunc=self.sdf_trunc,
|
| 285 |
+
)
|
| 286 |
+
mesh = fuse_depth_to_tsdf(
|
| 287 |
+
volume, depths, images, intrinsics, extrinsics, max_depth=self.max_depth
|
| 288 |
+
)
|
| 289 |
+
|
| 290 |
+
# Sample points from mesh
|
| 291 |
+
pcd = sample_points_from_mesh(mesh, self.sampling_number)
|
| 292 |
+
|
| 293 |
+
# Save point cloud
|
| 294 |
+
os.makedirs(os.path.dirname(fuse_path), exist_ok=True)
|
| 295 |
+
o3d.io.write_point_cloud(fuse_path, pcd)
|
| 296 |
+
|
| 297 |
+
# ------------------------------
|
| 298 |
+
# Private helpers
|
| 299 |
+
# ------------------------------
|
| 300 |
+
|
| 301 |
+
def _prep_unposed(
|
| 302 |
+
self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str
|
| 303 |
+
) -> tuple:
|
| 304 |
+
"""
|
| 305 |
+
Prepare depths/intrinsics/extrinsics for recon_unposed mode.
|
| 306 |
+
|
| 307 |
+
Similar to ETH3D but uses GT depth for masking instead of separate mask files.
|
| 308 |
+
"""
|
| 309 |
+
# Scale alignment with fixed random_state for reproducibility
|
| 310 |
+
_, _, scale, extrinsics = align_poses_umeyama(
|
| 311 |
+
gt_data.extrinsics.copy(),
|
| 312 |
+
pred_data.extrinsics.copy(),
|
| 313 |
+
return_aligned=True,
|
| 314 |
+
ransac=True,
|
| 315 |
+
random_state=42,
|
| 316 |
+
)
|
| 317 |
+
|
| 318 |
+
model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2]
|
| 319 |
+
|
| 320 |
+
depths_out = []
|
| 321 |
+
intrinsics_out = []
|
| 322 |
+
for i in range(len(pred_data.depth)):
|
| 323 |
+
orig_h, orig_w = orig_sizes[i]
|
| 324 |
+
|
| 325 |
+
# Resize depth to original image size (nearest interpolation)
|
| 326 |
+
depth = cv2.resize(
|
| 327 |
+
pred_data.depth[i],
|
| 328 |
+
(orig_w, orig_h),
|
| 329 |
+
interpolation=cv2.INTER_NEAREST,
|
| 330 |
+
)
|
| 331 |
+
|
| 332 |
+
# Load GT depth for masking
|
| 333 |
+
gt_zero_mask = self._load_gt_mask(gt_data.aux.gt_depth_files[i])
|
| 334 |
+
|
| 335 |
+
# Mask invalid depths BEFORE scale
|
| 336 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 337 |
+
|
| 338 |
+
# Apply scale AFTER mask
|
| 339 |
+
depth = depth * scale
|
| 340 |
+
|
| 341 |
+
# Adjust intrinsics to original image size
|
| 342 |
+
h_ratio = orig_h / model_h
|
| 343 |
+
w_ratio = orig_w / model_w
|
| 344 |
+
ixt = pred_data.intrinsics[i].copy()
|
| 345 |
+
ixt[0, :] *= w_ratio
|
| 346 |
+
ixt[1, :] *= h_ratio
|
| 347 |
+
|
| 348 |
+
depths_out.append(depth)
|
| 349 |
+
intrinsics_out.append(ixt)
|
| 350 |
+
|
| 351 |
+
return np.stack(depths_out), np.stack(intrinsics_out), extrinsics
|
| 352 |
+
|
| 353 |
+
def _prep_posed(
|
| 354 |
+
self, pred_data: Dict, gt_data: Dict, orig_sizes: list, scene: str
|
| 355 |
+
) -> tuple:
|
| 356 |
+
"""
|
| 357 |
+
Prepare depths/intrinsics/extrinsics for recon_posed mode.
|
| 358 |
+
Uses GT intrinsics/extrinsics but aligns depth scale via Umeyama.
|
| 359 |
+
"""
|
| 360 |
+
# Scale alignment with fixed random_state
|
| 361 |
+
_, _, scale, _ = align_poses_umeyama(
|
| 362 |
+
gt_data.extrinsics.copy(),
|
| 363 |
+
pred_data.extrinsics.copy(),
|
| 364 |
+
return_aligned=True,
|
| 365 |
+
ransac=True,
|
| 366 |
+
random_state=42,
|
| 367 |
+
)
|
| 368 |
+
|
| 369 |
+
model_h, model_w = pred_data.depth.shape[1], pred_data.depth.shape[2]
|
| 370 |
+
|
| 371 |
+
depths_out = []
|
| 372 |
+
for i in range(len(pred_data.depth)):
|
| 373 |
+
orig_h, orig_w = orig_sizes[i]
|
| 374 |
+
|
| 375 |
+
# Resize depth to original image size
|
| 376 |
+
depth = cv2.resize(
|
| 377 |
+
pred_data.depth[i],
|
| 378 |
+
(orig_w, orig_h),
|
| 379 |
+
interpolation=cv2.INTER_NEAREST,
|
| 380 |
+
)
|
| 381 |
+
|
| 382 |
+
# Load GT depth for masking
|
| 383 |
+
gt_zero_mask = self._load_gt_mask(gt_data.aux.gt_depth_files[i])
|
| 384 |
+
|
| 385 |
+
# Mask invalid depths BEFORE scale
|
| 386 |
+
depth = self._mask_invalid_depth(depth, gt_zero_mask)
|
| 387 |
+
|
| 388 |
+
# Apply scale AFTER mask
|
| 389 |
+
depth = depth * scale
|
| 390 |
+
|
| 391 |
+
depths_out.append(depth)
|
| 392 |
+
|
| 393 |
+
# Use GT intrinsics and extrinsics
|
| 394 |
+
return np.stack(depths_out), gt_data.intrinsics.copy(), gt_data.extrinsics.copy()
|
| 395 |
+
|
| 396 |
+
def _load_gt_mask(self, gt_depth_path: str) -> np.ndarray:
|
| 397 |
+
"""
|
| 398 |
+
Load GT depth and create valid mask.
|
| 399 |
+
|
| 400 |
+
For 7Scenes, GT depth is stored as 16-bit PNG in millimeters.
|
| 401 |
+
Value 65535 indicates invalid depth.
|
| 402 |
+
|
| 403 |
+
Returns:
|
| 404 |
+
Boolean mask where True = valid region to keep
|
| 405 |
+
"""
|
| 406 |
+
if not os.path.exists(gt_depth_path):
|
| 407 |
+
return None
|
| 408 |
+
|
| 409 |
+
gt_depth = cv2.imread(gt_depth_path, -1)
|
| 410 |
+
if gt_depth is None:
|
| 411 |
+
return None
|
| 412 |
+
|
| 413 |
+
# 65535 is invalid depth marker in 7Scenes
|
| 414 |
+
gt_depth[gt_depth == 65535] = 0
|
| 415 |
+
# Convert to meters
|
| 416 |
+
gt_depth = gt_depth / 1000.0
|
| 417 |
+
|
| 418 |
+
# Valid mask: depth > 0
|
| 419 |
+
valid_mask = gt_depth > 0
|
| 420 |
+
return valid_mask
|
| 421 |
+
|
| 422 |
+
def _mask_invalid_depth(
|
| 423 |
+
self, depth: np.ndarray, gt_zero_mask: np.ndarray = None
|
| 424 |
+
) -> np.ndarray:
|
| 425 |
+
"""
|
| 426 |
+
Mask invalid depth values by setting them to 0.
|
| 427 |
+
|
| 428 |
+
Args:
|
| 429 |
+
depth: Depth map to mask
|
| 430 |
+
gt_zero_mask: Optional GT mask (True = valid region)
|
| 431 |
+
|
| 432 |
+
Returns:
|
| 433 |
+
Masked depth map with invalid regions set to 0
|
| 434 |
+
"""
|
| 435 |
+
depth = depth.copy()
|
| 436 |
+
|
| 437 |
+
if gt_zero_mask is not None:
|
| 438 |
+
# Also mask out invalid pred depth
|
| 439 |
+
pred_invalid = np.isnan(depth) | np.isinf(depth)
|
| 440 |
+
combined_mask = np.logical_and(gt_zero_mask, np.logical_not(pred_invalid))
|
| 441 |
+
depth = depth * combined_mask.astype(np.float32)
|
| 442 |
+
else:
|
| 443 |
+
# Fallback: only mask pred invalid values
|
| 444 |
+
invalid_mask = np.isnan(depth) | np.isinf(depth) | (depth <= 0)
|
| 445 |
+
depth[invalid_mask] = 0.0
|
| 446 |
+
|
| 447 |
+
return depth
|
| 448 |
+
|
| 449 |
+
|
depth_anything_3/bench/evaluator.py
ADDED
|
@@ -0,0 +1,752 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Main Evaluator class for DepthAnything3 benchmark evaluation.
|
| 17 |
+
|
| 18 |
+
Supports multiple datasets and evaluation modes:
|
| 19 |
+
- pose: Camera pose estimation (AUC metrics)
|
| 20 |
+
- recon_unposed: 3D reconstruction with predicted poses
|
| 21 |
+
- recon_posed: 3D reconstruction with GT poses
|
| 22 |
+
- view_syn: Novel view synthesis (TODO)
|
| 23 |
+
"""
|
| 24 |
+
|
| 25 |
+
import json
|
| 26 |
+
import os
|
| 27 |
+
import random
|
| 28 |
+
from typing import Dict as TDict, Iterable, List
|
| 29 |
+
|
| 30 |
+
import numpy as np
|
| 31 |
+
import torch
|
| 32 |
+
from addict import Dict
|
| 33 |
+
from tqdm import tqdm
|
| 34 |
+
|
| 35 |
+
from depth_anything_3.bench.print_metrics import MetricsPrinter
|
| 36 |
+
from depth_anything_3.utils.parallel_utils import parallel_execution
|
| 37 |
+
from depth_anything_3.bench.registries import MV_REGISTRY
|
| 38 |
+
from depth_anything_3.utils.constants import EVAL_REF_VIEW_STRATEGY
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class Evaluator:
|
| 42 |
+
"""
|
| 43 |
+
Main evaluation orchestrator for DepthAnything3 benchmarks.
|
| 44 |
+
|
| 45 |
+
Usage:
|
| 46 |
+
evaluator = Evaluator(
|
| 47 |
+
work_dir="./eval_workspace",
|
| 48 |
+
datas=["dtu"],
|
| 49 |
+
modes=["pose", "recon_unposed", "recon_posed"],
|
| 50 |
+
)
|
| 51 |
+
api = DepthAnything3.from_pretrained("...")
|
| 52 |
+
evaluator.infer(api)
|
| 53 |
+
metrics = evaluator.eval()
|
| 54 |
+
evaluator.print_metrics()
|
| 55 |
+
"""
|
| 56 |
+
|
| 57 |
+
VALID_MODES = {"pose", "recon_unposed", "recon_posed", "view_syn"}
|
| 58 |
+
|
| 59 |
+
def __init__(
|
| 60 |
+
self,
|
| 61 |
+
work_dir: str = "./eval_workspace",
|
| 62 |
+
datas: List[str] = ("dtu",),
|
| 63 |
+
modes: List[str] = ("recon_unposed",),
|
| 64 |
+
ref_view_strategy: str = EVAL_REF_VIEW_STRATEGY,
|
| 65 |
+
scenes: List[str] = None,
|
| 66 |
+
debug: bool = False,
|
| 67 |
+
num_fusion_workers: int = 4,
|
| 68 |
+
max_frames: int = 100,
|
| 69 |
+
gpu_id: int = 0,
|
| 70 |
+
total_gpus: int = 1,
|
| 71 |
+
):
|
| 72 |
+
"""
|
| 73 |
+
Initialize the evaluator.
|
| 74 |
+
|
| 75 |
+
Args:
|
| 76 |
+
work_dir: Base directory for model outputs and metric files
|
| 77 |
+
datas: List of dataset names (must be registered in MV_REGISTRY)
|
| 78 |
+
modes: List of evaluation modes to run
|
| 79 |
+
ref_view_strategy: Reference view selection strategy for inference
|
| 80 |
+
("first", "saddle_balanced", etc.)
|
| 81 |
+
scenes: Specific scenes to evaluate (None = all scenes)
|
| 82 |
+
debug: Enable verbose debug output
|
| 83 |
+
num_fusion_workers: Number of parallel workers for TSDF fusion (default: 4)
|
| 84 |
+
max_frames: Maximum number of frames per scene (default: 100).
|
| 85 |
+
If a scene has more frames, randomly sample to this limit.
|
| 86 |
+
Set to -1 to disable sampling.
|
| 87 |
+
gpu_id: GPU index for multi-GPU (0-indexed)
|
| 88 |
+
total_gpus: Total number of GPUs for task distribution
|
| 89 |
+
"""
|
| 90 |
+
self.work_dir = work_dir
|
| 91 |
+
self.datas = list(datas)
|
| 92 |
+
self.modes = set(modes)
|
| 93 |
+
self.ref_view_strategy = ref_view_strategy
|
| 94 |
+
self.scenes_filter = scenes
|
| 95 |
+
self.debug = debug
|
| 96 |
+
self.num_fusion_workers = num_fusion_workers
|
| 97 |
+
self.max_frames = max_frames
|
| 98 |
+
self.gpu_id = gpu_id
|
| 99 |
+
self.total_gpus = total_gpus
|
| 100 |
+
|
| 101 |
+
# Validate modes
|
| 102 |
+
unknown = self.modes - self.VALID_MODES
|
| 103 |
+
if unknown:
|
| 104 |
+
raise ValueError(f"Unknown modes: {unknown}. Valid: {sorted(self.VALID_MODES)}")
|
| 105 |
+
|
| 106 |
+
os.makedirs(self.work_dir, exist_ok=True)
|
| 107 |
+
|
| 108 |
+
# Initialize datasets
|
| 109 |
+
self.datasets = Dict()
|
| 110 |
+
for data in self.datas:
|
| 111 |
+
if not MV_REGISTRY.has(data):
|
| 112 |
+
available = list(MV_REGISTRY.all().keys())
|
| 113 |
+
raise ValueError(f"Dataset '{data}' not found. Available: {available}")
|
| 114 |
+
self.datasets[data] = MV_REGISTRY.get(data)()
|
| 115 |
+
|
| 116 |
+
# Initialize metrics printer
|
| 117 |
+
self._printer = MetricsPrinter()
|
| 118 |
+
|
| 119 |
+
# -------------------- Public APIs -------------------- #
|
| 120 |
+
|
| 121 |
+
def all(self, api) -> TDict[str, dict]:
|
| 122 |
+
"""
|
| 123 |
+
Run complete evaluation pipeline: inference + evaluation.
|
| 124 |
+
|
| 125 |
+
Args:
|
| 126 |
+
api: DepthAnything3 API instance
|
| 127 |
+
|
| 128 |
+
Returns:
|
| 129 |
+
Combined metrics dictionary
|
| 130 |
+
"""
|
| 131 |
+
self.infer(api)
|
| 132 |
+
return self.eval()
|
| 133 |
+
|
| 134 |
+
def _get_scenes(self, dataset) -> List[str]:
|
| 135 |
+
"""Get list of scenes to evaluate, optionally filtered."""
|
| 136 |
+
all_scenes = dataset.SCENES
|
| 137 |
+
if self.scenes_filter:
|
| 138 |
+
scenes = [s for s in all_scenes if s in self.scenes_filter]
|
| 139 |
+
if self.debug:
|
| 140 |
+
print(f"[DEBUG] Filtered scenes: {scenes} (from {len(all_scenes)} total)")
|
| 141 |
+
return scenes
|
| 142 |
+
return all_scenes
|
| 143 |
+
|
| 144 |
+
def infer(self, api, model_path: str = None) -> None:
|
| 145 |
+
"""
|
| 146 |
+
Run inference according to requested modes.
|
| 147 |
+
|
| 148 |
+
- Unposed export if 'pose' or 'recon_unposed' is in modes
|
| 149 |
+
- Posed export if 'recon_posed' or 'view_syn' is in modes
|
| 150 |
+
|
| 151 |
+
Multi-GPU: Use --gpu_id and --total_gpus to distribute tasks.
|
| 152 |
+
Example: Launch 4 processes with gpu_id=0,1,2,3 and total_gpus=4
|
| 153 |
+
|
| 154 |
+
Args:
|
| 155 |
+
api: DepthAnything3 API instance
|
| 156 |
+
model_path: Model path (unused, kept for API compatibility)
|
| 157 |
+
"""
|
| 158 |
+
need_unposed = {"pose", "recon_unposed"} & self.modes
|
| 159 |
+
need_posed = {"recon_posed", "view_syn"} & self.modes
|
| 160 |
+
export_format = "mini_npz-glb" if self.debug else "mini_npz"
|
| 161 |
+
|
| 162 |
+
# Collect all tasks
|
| 163 |
+
all_tasks = []
|
| 164 |
+
for data in self.datas:
|
| 165 |
+
dataset = self.datasets[data]
|
| 166 |
+
for scene in self._get_scenes(dataset):
|
| 167 |
+
all_tasks.append((data, scene))
|
| 168 |
+
|
| 169 |
+
# Distribute tasks across GPUs
|
| 170 |
+
if self.total_gpus > 1:
|
| 171 |
+
tasks = [t for i, t in enumerate(all_tasks) if i % self.total_gpus == self.gpu_id]
|
| 172 |
+
print(f"[INFO] GPU {self.gpu_id}/{self.total_gpus}: {len(tasks)}/{len(all_tasks)} tasks")
|
| 173 |
+
else:
|
| 174 |
+
tasks = all_tasks
|
| 175 |
+
print(f"[INFO] Total inference tasks: {len(tasks)}")
|
| 176 |
+
|
| 177 |
+
for data, scene in tqdm(tasks, desc=f"Inference (GPU {self.gpu_id})"):
|
| 178 |
+
dataset = self.datasets[data]
|
| 179 |
+
scene_data = dataset.get_data(scene)
|
| 180 |
+
scene_data = self._sample_frames(scene_data, scene)
|
| 181 |
+
|
| 182 |
+
if need_unposed:
|
| 183 |
+
export_dir = self._export_dir(data, scene, posed=False)
|
| 184 |
+
api.inference(
|
| 185 |
+
scene_data.image_files,
|
| 186 |
+
export_dir=export_dir,
|
| 187 |
+
export_format=export_format,
|
| 188 |
+
ref_view_strategy=self.ref_view_strategy,
|
| 189 |
+
)
|
| 190 |
+
self._save_gt_meta(export_dir, scene_data)
|
| 191 |
+
|
| 192 |
+
if need_posed:
|
| 193 |
+
export_dir = self._export_dir(data, scene, posed=True)
|
| 194 |
+
api.inference(
|
| 195 |
+
scene_data.image_files,
|
| 196 |
+
scene_data.extrinsics,
|
| 197 |
+
scene_data.intrinsics,
|
| 198 |
+
export_dir=export_dir,
|
| 199 |
+
export_format=export_format,
|
| 200 |
+
ref_view_strategy=self.ref_view_strategy,
|
| 201 |
+
)
|
| 202 |
+
self._save_gt_meta(export_dir, scene_data)
|
| 203 |
+
|
| 204 |
+
def eval(self) -> TDict[str, dict]:
|
| 205 |
+
"""
|
| 206 |
+
Evaluate for all configured modes and write JSON files.
|
| 207 |
+
|
| 208 |
+
Evaluation order by mode (all datasets per mode):
|
| 209 |
+
1. pose - all datasets
|
| 210 |
+
2. recon_unposed - all datasets
|
| 211 |
+
3. recon_posed - all datasets
|
| 212 |
+
|
| 213 |
+
Returns:
|
| 214 |
+
Summary mapping: {"<data>_<mode>": metrics_dict}
|
| 215 |
+
"""
|
| 216 |
+
summary: TDict[str, dict] = {}
|
| 217 |
+
|
| 218 |
+
# Evaluate by mode (all datasets per mode)
|
| 219 |
+
if "pose" in self.modes:
|
| 220 |
+
print(f"\n{'='*60}")
|
| 221 |
+
print(f"📊 Evaluating POSE for all datasets...")
|
| 222 |
+
print(f"{'='*60}")
|
| 223 |
+
for data, result in self._eval_pose():
|
| 224 |
+
summary[f"{data}_pose"] = result
|
| 225 |
+
|
| 226 |
+
if "recon_unposed" in self.modes:
|
| 227 |
+
print(f"\n{'='*60}")
|
| 228 |
+
print(f"📊 Evaluating RECON_UNPOSED for all datasets...")
|
| 229 |
+
print(f"{'='*60}")
|
| 230 |
+
for data, result in self._eval_reconstruction("recon_unposed"):
|
| 231 |
+
summary[f"{data}_recon_unposed"] = result
|
| 232 |
+
|
| 233 |
+
if "recon_posed" in self.modes:
|
| 234 |
+
print(f"\n{'='*60}")
|
| 235 |
+
print(f"📊 Evaluating RECON_POSED for all datasets...")
|
| 236 |
+
print(f"{'='*60}")
|
| 237 |
+
for data, result in self._eval_reconstruction("recon_posed"):
|
| 238 |
+
summary[f"{data}_recon_posed"] = result
|
| 239 |
+
|
| 240 |
+
if "view_syn" in self.modes:
|
| 241 |
+
# TODO: Add view synthesis metrics here when available
|
| 242 |
+
pass
|
| 243 |
+
|
| 244 |
+
return summary
|
| 245 |
+
|
| 246 |
+
def print_metrics(self, metrics: TDict[str, dict] = None) -> None:
|
| 247 |
+
"""
|
| 248 |
+
Print evaluation metrics in a beautiful tabular format.
|
| 249 |
+
|
| 250 |
+
Args:
|
| 251 |
+
metrics: Metrics dictionary. If None, loads from saved JSON files.
|
| 252 |
+
"""
|
| 253 |
+
if metrics is None:
|
| 254 |
+
metrics = self._load_metrics()
|
| 255 |
+
|
| 256 |
+
self._printer.print_results(metrics)
|
| 257 |
+
|
| 258 |
+
# -------------------- Evaluation Methods -------------------- #
|
| 259 |
+
|
| 260 |
+
def _eval_pose(self) -> Iterable[tuple]:
|
| 261 |
+
"""Compute pose-estimation metrics for each dataset and scene."""
|
| 262 |
+
os.makedirs(self._metric_dir, exist_ok=True)
|
| 263 |
+
|
| 264 |
+
for data in tqdm(self.datas, desc="Datasets (pose eval)"):
|
| 265 |
+
dataset = self.datasets[data]
|
| 266 |
+
dataset_results = Dict()
|
| 267 |
+
scenes = self._get_scenes(dataset)
|
| 268 |
+
|
| 269 |
+
for scene in tqdm(scenes, desc=f"{data} scenes", leave=False):
|
| 270 |
+
export_dir = self._export_dir(data, scene, posed=False)
|
| 271 |
+
result_path = os.path.join(export_dir, "exports", "mini_npz", "results.npz")
|
| 272 |
+
|
| 273 |
+
# Check if result file exists and is valid
|
| 274 |
+
if not os.path.exists(result_path):
|
| 275 |
+
print(f"\n[ERROR] Result file not found: {result_path}")
|
| 276 |
+
print(f"[ERROR] CWD: {os.getcwd()}")
|
| 277 |
+
print(f"[ERROR] Please run inference first (remove --eval_only)")
|
| 278 |
+
continue
|
| 279 |
+
|
| 280 |
+
try:
|
| 281 |
+
# Use saved GT meta (handles frame sampling correctly)
|
| 282 |
+
gt_meta = self._load_gt_meta(export_dir)
|
| 283 |
+
if gt_meta is not None:
|
| 284 |
+
result = self._compute_pose_with_gt(result_path, gt_meta)
|
| 285 |
+
else:
|
| 286 |
+
# Fallback to dataset GT (no sampling was done)
|
| 287 |
+
result = dataset.eval_pose(scene, result_path)
|
| 288 |
+
dataset_results[scene] = self._to_float_dict(result)
|
| 289 |
+
except Exception as e:
|
| 290 |
+
print(f"\n[ERROR] Failed to evaluate pose for {data}/{scene}: {e}")
|
| 291 |
+
print(f"[ERROR] File path: {os.path.abspath(result_path)}")
|
| 292 |
+
if self.debug:
|
| 293 |
+
import traceback
|
| 294 |
+
traceback.print_exc()
|
| 295 |
+
continue
|
| 296 |
+
|
| 297 |
+
if not dataset_results:
|
| 298 |
+
print(f"[WARNING] No valid results for {data}")
|
| 299 |
+
continue
|
| 300 |
+
|
| 301 |
+
dataset_results["mean"] = self._mean_of_dicts(dataset_results.values())
|
| 302 |
+
out_path = os.path.join(self._metric_dir, f"{data}_pose.json")
|
| 303 |
+
self._dump_json(out_path, dataset_results)
|
| 304 |
+
yield data, dataset_results
|
| 305 |
+
|
| 306 |
+
def _eval_reconstruction(self, mode: str) -> Iterable[tuple]:
|
| 307 |
+
"""
|
| 308 |
+
Compute reconstruction metrics for each dataset and scene.
|
| 309 |
+
|
| 310 |
+
Args:
|
| 311 |
+
mode: "recon_unposed" or "recon_posed"
|
| 312 |
+
"""
|
| 313 |
+
assert mode in {"recon_unposed", "recon_posed"}
|
| 314 |
+
os.makedirs(self._metric_dir, exist_ok=True)
|
| 315 |
+
|
| 316 |
+
posed_flag = mode == "recon_posed"
|
| 317 |
+
|
| 318 |
+
# Filter out datasets that don't support reconstruction (e.g., dtu64)
|
| 319 |
+
recon_datas = [d for d in self.datas if d != "dtu64"]
|
| 320 |
+
|
| 321 |
+
for data in tqdm(recon_datas, desc=f"Datasets ({mode} eval)"):
|
| 322 |
+
dataset = self.datasets[data]
|
| 323 |
+
dataset_results = Dict()
|
| 324 |
+
scenes = self._get_scenes(dataset)
|
| 325 |
+
|
| 326 |
+
# Prepare paths for all scenes
|
| 327 |
+
scene_list = []
|
| 328 |
+
result_paths = []
|
| 329 |
+
fuse_paths = []
|
| 330 |
+
for scene in scenes:
|
| 331 |
+
export_dir = self._export_dir(data, scene, posed=posed_flag)
|
| 332 |
+
result_path = os.path.join(export_dir, "exports", "mini_npz", "results.npz")
|
| 333 |
+
fuse_path = os.path.join(export_dir, "exports", "fuse", "pcd.ply")
|
| 334 |
+
scene_list.append(scene)
|
| 335 |
+
result_paths.append(result_path)
|
| 336 |
+
fuse_paths.append(fuse_path)
|
| 337 |
+
|
| 338 |
+
# Parallel fusion (default 4 workers)
|
| 339 |
+
# DTU uses CUDA operations in fusion, which doesn't work well with ThreadPool
|
| 340 |
+
use_sequential = (data == "dtu")
|
| 341 |
+
parallel_execution(
|
| 342 |
+
scene_list,
|
| 343 |
+
result_paths,
|
| 344 |
+
fuse_paths,
|
| 345 |
+
action=lambda s, rp, fp: dataset.fuse3d(s, rp, fp, mode),
|
| 346 |
+
num_processes=self.num_fusion_workers,
|
| 347 |
+
print_progress=True,
|
| 348 |
+
desc=f"{data} fusion",
|
| 349 |
+
sequential=use_sequential,
|
| 350 |
+
)
|
| 351 |
+
|
| 352 |
+
# Sequential evaluation (fast, no need to parallelize)
|
| 353 |
+
for scene, fuse_path in zip(scene_list, fuse_paths):
|
| 354 |
+
# DTU supports CPU-based evaluation
|
| 355 |
+
if data == "dtu" and hasattr(dataset, "eval3d"):
|
| 356 |
+
result = dataset.eval3d(scene, fuse_path)
|
| 357 |
+
else:
|
| 358 |
+
result = dataset.eval3d(scene, fuse_path)
|
| 359 |
+
dataset_results[scene] = self._to_float_dict(result)
|
| 360 |
+
print(f" {mode} | {data} | {scene}: {result}")
|
| 361 |
+
|
| 362 |
+
dataset_results["mean"] = self._mean_of_dicts(dataset_results.values())
|
| 363 |
+
out_path = os.path.join(self._metric_dir, f"{data}_{mode}.json")
|
| 364 |
+
self._dump_json(out_path, dataset_results)
|
| 365 |
+
yield data, dataset_results
|
| 366 |
+
|
| 367 |
+
# -------------------- Helpers -------------------- #
|
| 368 |
+
|
| 369 |
+
def _save_gt_meta(self, export_dir: str, scene_data: Dict) -> None:
|
| 370 |
+
"""
|
| 371 |
+
Save GT extrinsics/intrinsics/image_files for evaluation.
|
| 372 |
+
|
| 373 |
+
This is needed when frames are sampled, so eval_pose and fuse3d can use
|
| 374 |
+
the correct (sampled) GT instead of full dataset GT.
|
| 375 |
+
|
| 376 |
+
Args:
|
| 377 |
+
export_dir: Export directory for the scene
|
| 378 |
+
scene_data: Sampled scene data
|
| 379 |
+
"""
|
| 380 |
+
meta_path = os.path.join(export_dir, "exports", "gt_meta.npz")
|
| 381 |
+
os.makedirs(os.path.dirname(meta_path), exist_ok=True)
|
| 382 |
+
np.savez_compressed(
|
| 383 |
+
meta_path,
|
| 384 |
+
extrinsics=scene_data.extrinsics,
|
| 385 |
+
intrinsics=scene_data.intrinsics,
|
| 386 |
+
image_files=np.array(scene_data.image_files, dtype=object),
|
| 387 |
+
)
|
| 388 |
+
|
| 389 |
+
def _load_gt_meta(self, export_dir: str) -> Dict:
|
| 390 |
+
"""
|
| 391 |
+
Load saved GT extrinsics/intrinsics for evaluation.
|
| 392 |
+
|
| 393 |
+
Returns:
|
| 394 |
+
Dict with extrinsics and intrinsics, or None if not found
|
| 395 |
+
"""
|
| 396 |
+
meta_path = os.path.join(export_dir, "exports", "gt_meta.npz")
|
| 397 |
+
if os.path.exists(meta_path):
|
| 398 |
+
data = np.load(meta_path)
|
| 399 |
+
return Dict({
|
| 400 |
+
"extrinsics": data["extrinsics"],
|
| 401 |
+
"intrinsics": data["intrinsics"],
|
| 402 |
+
})
|
| 403 |
+
return None
|
| 404 |
+
|
| 405 |
+
def _compute_pose_with_gt(self, result_path: str, gt_meta: Dict) -> TDict[str, float]:
|
| 406 |
+
"""
|
| 407 |
+
Compute pose metrics using saved GT meta (handles frame sampling).
|
| 408 |
+
|
| 409 |
+
Args:
|
| 410 |
+
result_path: Path to npz with predicted extrinsics
|
| 411 |
+
gt_meta: Dict with GT extrinsics from saved meta
|
| 412 |
+
|
| 413 |
+
Returns:
|
| 414 |
+
Dict with pose metrics
|
| 415 |
+
"""
|
| 416 |
+
from depth_anything_3.bench.dataset import _wait_for_file_ready
|
| 417 |
+
from depth_anything_3.bench.utils import compute_pose
|
| 418 |
+
from depth_anything_3.utils.geometry import as_homogeneous
|
| 419 |
+
|
| 420 |
+
_wait_for_file_ready(result_path)
|
| 421 |
+
pred = np.load(result_path)
|
| 422 |
+
return compute_pose(
|
| 423 |
+
torch.from_numpy(as_homogeneous(pred["extrinsics"])),
|
| 424 |
+
torch.from_numpy(as_homogeneous(gt_meta["extrinsics"])),
|
| 425 |
+
)
|
| 426 |
+
|
| 427 |
+
def _sample_frames(self, scene_data: Dict, scene: str) -> Dict:
|
| 428 |
+
"""
|
| 429 |
+
Sample frames if scene has more than max_frames.
|
| 430 |
+
|
| 431 |
+
Uses fixed random seed (42) for reproducibility.
|
| 432 |
+
|
| 433 |
+
Args:
|
| 434 |
+
scene_data: Scene data dict with image_files, extrinsics, intrinsics, aux
|
| 435 |
+
scene: Scene name (for logging)
|
| 436 |
+
|
| 437 |
+
Returns:
|
| 438 |
+
Sampled scene_data if num_frames > max_frames, otherwise original
|
| 439 |
+
"""
|
| 440 |
+
if self.max_frames <= 0:
|
| 441 |
+
return scene_data
|
| 442 |
+
|
| 443 |
+
num_frames = len(scene_data.image_files)
|
| 444 |
+
if num_frames <= self.max_frames:
|
| 445 |
+
return scene_data
|
| 446 |
+
|
| 447 |
+
# Sample with fixed seed for reproducibility
|
| 448 |
+
random.seed(42)
|
| 449 |
+
indices = list(range(num_frames))
|
| 450 |
+
random.shuffle(indices)
|
| 451 |
+
sampled_indices = sorted(indices[:self.max_frames])
|
| 452 |
+
|
| 453 |
+
print(f" [Sampling] {scene}: {num_frames} -> {self.max_frames} frames")
|
| 454 |
+
|
| 455 |
+
# Create new scene_data with sampled frames
|
| 456 |
+
sampled = Dict()
|
| 457 |
+
sampled.image_files = [scene_data.image_files[i] for i in sampled_indices]
|
| 458 |
+
sampled.extrinsics = scene_data.extrinsics[sampled_indices]
|
| 459 |
+
sampled.intrinsics = scene_data.intrinsics[sampled_indices]
|
| 460 |
+
|
| 461 |
+
# Copy aux data, sampling lists if needed
|
| 462 |
+
sampled.aux = Dict()
|
| 463 |
+
for key, val in scene_data.aux.items():
|
| 464 |
+
if isinstance(val, list) and len(val) == num_frames:
|
| 465 |
+
sampled.aux[key] = [val[i] for i in sampled_indices]
|
| 466 |
+
elif isinstance(val, np.ndarray) and len(val) == num_frames:
|
| 467 |
+
sampled.aux[key] = val[sampled_indices]
|
| 468 |
+
else:
|
| 469 |
+
sampled.aux[key] = val
|
| 470 |
+
|
| 471 |
+
return sampled
|
| 472 |
+
|
| 473 |
+
@property
|
| 474 |
+
def _metric_dir(self) -> str:
|
| 475 |
+
"""Directory for storing metric JSON files."""
|
| 476 |
+
return os.path.join(self.work_dir, "metric_results")
|
| 477 |
+
|
| 478 |
+
def _export_dir(self, data: str, scene: str, posed: bool) -> str:
|
| 479 |
+
"""
|
| 480 |
+
Get export directory path.
|
| 481 |
+
|
| 482 |
+
Structure: .../model_results/{data}/{scene}/{posed|unposed}
|
| 483 |
+
"""
|
| 484 |
+
suffix = "posed" if posed else "unposed"
|
| 485 |
+
export_dir = os.path.join(self.work_dir, "model_results", data, scene, suffix)
|
| 486 |
+
os.makedirs(export_dir, exist_ok=True)
|
| 487 |
+
return export_dir
|
| 488 |
+
|
| 489 |
+
@staticmethod
|
| 490 |
+
def _to_float_dict(d: TDict[str, float]) -> dict:
|
| 491 |
+
"""Convert numpy scalars to plain Python floats for JSON safety."""
|
| 492 |
+
return {k: float(v) for k, v in d.items()}
|
| 493 |
+
|
| 494 |
+
@staticmethod
|
| 495 |
+
def _mean_of_dicts(dicts: Iterable[dict]) -> dict:
|
| 496 |
+
"""Compute elementwise mean across a list of homogeneous metric dicts."""
|
| 497 |
+
dicts = list(dicts)
|
| 498 |
+
if not dicts:
|
| 499 |
+
return {}
|
| 500 |
+
keys = dicts[0].keys()
|
| 501 |
+
return {k: float(np.mean([d[k] for d in dicts]).item()) for k in keys}
|
| 502 |
+
|
| 503 |
+
@staticmethod
|
| 504 |
+
def _dump_json(path: str, obj: dict, indent: int = 4) -> None:
|
| 505 |
+
"""Write JSON with UTF-8 and pretty indentation."""
|
| 506 |
+
os.makedirs(os.path.dirname(path), exist_ok=True)
|
| 507 |
+
with open(path, "w", encoding="utf-8") as f:
|
| 508 |
+
json.dump(obj, f, indent=indent, ensure_ascii=False)
|
| 509 |
+
|
| 510 |
+
def _load_metrics(self) -> TDict[str, dict]:
|
| 511 |
+
"""Load evaluation metrics from JSON files."""
|
| 512 |
+
metrics = {}
|
| 513 |
+
metric_dir = self._metric_dir
|
| 514 |
+
|
| 515 |
+
if not os.path.exists(metric_dir):
|
| 516 |
+
return metrics
|
| 517 |
+
|
| 518 |
+
for filename in os.listdir(metric_dir):
|
| 519 |
+
if filename.endswith(".json"):
|
| 520 |
+
filepath = os.path.join(metric_dir, filename)
|
| 521 |
+
try:
|
| 522 |
+
with open(filepath, encoding="utf-8") as f:
|
| 523 |
+
data = json.load(f)
|
| 524 |
+
key = filename[:-5] # Remove .json extension
|
| 525 |
+
metrics[key] = data
|
| 526 |
+
except Exception as e:
|
| 527 |
+
print(f"Warning: Failed to read metrics file: {filename} - {e}")
|
| 528 |
+
|
| 529 |
+
return metrics
|
| 530 |
+
|
| 531 |
+
|
| 532 |
+
# -------------------- CLI Entry Point -------------------- #
|
| 533 |
+
|
| 534 |
+
|
| 535 |
+
if __name__ == "__main__":
|
| 536 |
+
import sys
|
| 537 |
+
from omegaconf import OmegaConf
|
| 538 |
+
from depth_anything_3.cfg import load_config
|
| 539 |
+
|
| 540 |
+
# Get default config path (relative to this file)
|
| 541 |
+
_default_config = os.path.join(
|
| 542 |
+
os.path.dirname(__file__), "configs", "eval_bench.yaml"
|
| 543 |
+
)
|
| 544 |
+
|
| 545 |
+
# Check for help flag first (we need to handle this before OmegaConf)
|
| 546 |
+
if "--help" in sys.argv or "-h" in sys.argv:
|
| 547 |
+
pass # Will handle after config loading
|
| 548 |
+
|
| 549 |
+
# Set up argv for OmegaConf processing
|
| 550 |
+
argv = sys.argv[1:]
|
| 551 |
+
|
| 552 |
+
# Check if user provides custom config
|
| 553 |
+
config_path = _default_config
|
| 554 |
+
if "--config" in argv:
|
| 555 |
+
config_idx = argv.index("--config")
|
| 556 |
+
if config_idx + 1 < len(argv):
|
| 557 |
+
config_path = argv[config_idx + 1]
|
| 558 |
+
# Remove --config and its value
|
| 559 |
+
argv = argv[:config_idx] + argv[config_idx + 2:]
|
| 560 |
+
|
| 561 |
+
# Print help if requested
|
| 562 |
+
if "--help" in sys.argv or "-h" in sys.argv:
|
| 563 |
+
print("""
|
| 564 |
+
DepthAnything3 Benchmark Evaluation
|
| 565 |
+
|
| 566 |
+
Usage:
|
| 567 |
+
python -m depth_anything_3.bench.evaluator [OPTIONS] [KEY=VALUE ...]
|
| 568 |
+
|
| 569 |
+
Configuration:
|
| 570 |
+
--config PATH Config YAML file (default: bench/configs/eval_bench.yaml)
|
| 571 |
+
|
| 572 |
+
Config Overrides (using dotlist notation):
|
| 573 |
+
model.path=VALUE Model path or HuggingFace ID
|
| 574 |
+
workspace.work_dir=VALUE Working directory for outputs
|
| 575 |
+
eval.datasets=[dataset1,dataset2] Datasets to evaluate (eth3d,7scenes,scannetpp,hiroom,dtu,dtu64)
|
| 576 |
+
eval.modes=[mode1,mode2] Evaluation modes (pose,recon_unposed,recon_posed)
|
| 577 |
+
eval.scenes=[scene1,scene2] Specific scenes to evaluate (null=all)
|
| 578 |
+
eval.max_frames=VALUE Max frames per scene (-1=no limit, default: 100)
|
| 579 |
+
eval.ref_view_strategy=VALUE Reference view strategy (default: first)
|
| 580 |
+
eval.eval_only=VALUE Only run evaluation (skip inference) (true/false)
|
| 581 |
+
eval.print_only=VALUE Only print saved metrics (true/false)
|
| 582 |
+
inference.num_fusion_workers=VALUE Number of parallel workers (default: 4)
|
| 583 |
+
inference.debug=VALUE Enable debug mode (true/false)
|
| 584 |
+
|
| 585 |
+
Special Flags:
|
| 586 |
+
--help, -h Show this help message
|
| 587 |
+
|
| 588 |
+
Multi-GPU:
|
| 589 |
+
Use CUDA_VISIBLE_DEVICES to specify GPUs (auto-detected and distributed)
|
| 590 |
+
|
| 591 |
+
Examples:
|
| 592 |
+
# Use default config
|
| 593 |
+
python -m depth_anything_3.bench.evaluator
|
| 594 |
+
|
| 595 |
+
# Override model path
|
| 596 |
+
python -m depth_anything_3.bench.evaluator model.path=depth-anything/DA3-LARGE
|
| 597 |
+
|
| 598 |
+
# Evaluate specific datasets and modes
|
| 599 |
+
python -m depth_anything_3.bench.evaluator \\
|
| 600 |
+
eval.datasets=[eth3d,hiroom] \\
|
| 601 |
+
eval.modes=[pose]
|
| 602 |
+
|
| 603 |
+
# Use custom config with overrides
|
| 604 |
+
python -m depth_anything_3.bench.evaluator \\
|
| 605 |
+
--config my_config.yaml \\
|
| 606 |
+
model.path=/path/to/model \\
|
| 607 |
+
eval.max_frames=50
|
| 608 |
+
|
| 609 |
+
# Multi-GPU inference (auto-distributed)
|
| 610 |
+
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m depth_anything_3.bench.evaluator
|
| 611 |
+
|
| 612 |
+
# Debug specific scenes
|
| 613 |
+
python -m depth_anything_3.bench.evaluator \\
|
| 614 |
+
eval.datasets=[eth3d] \\
|
| 615 |
+
eval.scenes=[courtyard] \\
|
| 616 |
+
inference.debug=true
|
| 617 |
+
|
| 618 |
+
# Only evaluate (skip inference)
|
| 619 |
+
python -m depth_anything_3.bench.evaluator eval.eval_only=true
|
| 620 |
+
|
| 621 |
+
# Only print saved metrics
|
| 622 |
+
python -m depth_anything_3.bench.evaluator eval.print_only=true
|
| 623 |
+
|
| 624 |
+
""")
|
| 625 |
+
sys.exit(0)
|
| 626 |
+
|
| 627 |
+
# Load config with CLI overrides using OmegaConf dotlist
|
| 628 |
+
# Example: python evaluator.py model.path=/path/to/model eval.datasets=[eth3d,dtu]
|
| 629 |
+
config = load_config(config_path, argv=argv)
|
| 630 |
+
|
| 631 |
+
# Extract config values
|
| 632 |
+
work_dir = config.workspace.work_dir
|
| 633 |
+
model_path = config.model.path
|
| 634 |
+
datasets = config.eval.datasets
|
| 635 |
+
modes = config.eval.modes
|
| 636 |
+
ref_view_strategy = config.eval.ref_view_strategy
|
| 637 |
+
scenes = config.eval.scenes
|
| 638 |
+
max_frames = config.eval.max_frames
|
| 639 |
+
eval_only = config.eval.eval_only
|
| 640 |
+
print_only = config.eval.print_only
|
| 641 |
+
debug = config.inference.debug
|
| 642 |
+
num_fusion_workers = config.inference.num_fusion_workers
|
| 643 |
+
|
| 644 |
+
# GPU settings: parse from CLI dotlist args (gpu_id=X total_gpus=Y)
|
| 645 |
+
# These are passed by the main process when spawning workers
|
| 646 |
+
gpu_id = 0
|
| 647 |
+
total_gpus = 1
|
| 648 |
+
for arg in argv:
|
| 649 |
+
if arg.startswith("gpu_id="):
|
| 650 |
+
gpu_id = int(arg.split("=")[1])
|
| 651 |
+
elif arg.startswith("total_gpus="):
|
| 652 |
+
total_gpus = int(arg.split("=")[1])
|
| 653 |
+
|
| 654 |
+
# Override dataset scenes if specified
|
| 655 |
+
if scenes:
|
| 656 |
+
print(f"[INFO] Running on specific scenes: {scenes}")
|
| 657 |
+
|
| 658 |
+
evaluator = Evaluator(
|
| 659 |
+
work_dir=work_dir,
|
| 660 |
+
datas=datasets,
|
| 661 |
+
modes=modes,
|
| 662 |
+
ref_view_strategy=ref_view_strategy,
|
| 663 |
+
scenes=scenes,
|
| 664 |
+
debug=debug,
|
| 665 |
+
num_fusion_workers=num_fusion_workers,
|
| 666 |
+
max_frames=max_frames,
|
| 667 |
+
gpu_id=gpu_id,
|
| 668 |
+
total_gpus=total_gpus,
|
| 669 |
+
)
|
| 670 |
+
|
| 671 |
+
if print_only:
|
| 672 |
+
evaluator.print_metrics()
|
| 673 |
+
elif eval_only:
|
| 674 |
+
metrics = evaluator.eval()
|
| 675 |
+
evaluator.print_metrics(metrics)
|
| 676 |
+
else:
|
| 677 |
+
# Parse CUDA_VISIBLE_DEVICES to get GPU list
|
| 678 |
+
# If not set, use all available GPUs
|
| 679 |
+
cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
| 680 |
+
if cuda_devices is not None and cuda_devices.strip():
|
| 681 |
+
gpu_list = [g.strip() for g in cuda_devices.split(",") if g.strip()]
|
| 682 |
+
else:
|
| 683 |
+
# CUDA_VISIBLE_DEVICES not set, use all available GPUs
|
| 684 |
+
num_available = torch.cuda.device_count()
|
| 685 |
+
gpu_list = [str(i) for i in range(num_available)] if num_available > 0 else ["0"]
|
| 686 |
+
|
| 687 |
+
# Auto multi-GPU: if multiple GPUs and not a worker process
|
| 688 |
+
is_worker = os.environ.get("_DA3_WORKER") == "1"
|
| 689 |
+
|
| 690 |
+
if len(gpu_list) > 1 and not is_worker:
|
| 691 |
+
# Launch worker processes
|
| 692 |
+
import subprocess
|
| 693 |
+
|
| 694 |
+
num_gpus = len(gpu_list)
|
| 695 |
+
print(f"[INFO] Detected {num_gpus} GPUs: {gpu_list}")
|
| 696 |
+
print(f"[INFO] Launching {num_gpus} workers...")
|
| 697 |
+
|
| 698 |
+
# Build base command
|
| 699 |
+
base_cmd = [sys.executable, "-m", "depth_anything_3.bench.evaluator"]
|
| 700 |
+
# Pass config via dotlist instead of CLI args
|
| 701 |
+
if config_path != _default_config:
|
| 702 |
+
base_cmd += ["--config", config_path]
|
| 703 |
+
base_cmd += [f"model.path={model_path}"]
|
| 704 |
+
base_cmd += [f"workspace.work_dir={work_dir}"]
|
| 705 |
+
base_cmd += [f"eval.datasets=[{','.join(datasets)}]"]
|
| 706 |
+
base_cmd += [f"eval.modes=[{','.join(modes)}]"]
|
| 707 |
+
if scenes:
|
| 708 |
+
base_cmd += [f"eval.scenes=[{','.join(scenes)}]"]
|
| 709 |
+
base_cmd += [f"eval.max_frames={max_frames}"]
|
| 710 |
+
base_cmd += [f"eval.ref_view_strategy={ref_view_strategy}"]
|
| 711 |
+
base_cmd += [f"inference.debug={str(debug).lower()}"]
|
| 712 |
+
base_cmd += [f"inference.num_fusion_workers={num_fusion_workers}"]
|
| 713 |
+
|
| 714 |
+
# Launch workers
|
| 715 |
+
processes = []
|
| 716 |
+
for idx, gpu_id in enumerate(gpu_list):
|
| 717 |
+
env = os.environ.copy()
|
| 718 |
+
env["CUDA_VISIBLE_DEVICES"] = gpu_id
|
| 719 |
+
env["_DA3_WORKER"] = "1" # Mark as worker process
|
| 720 |
+
|
| 721 |
+
cmd = base_cmd.copy()
|
| 722 |
+
# GPU-specific worker config
|
| 723 |
+
cmd += [f"gpu_id={idx}", f"total_gpus={num_gpus}"]
|
| 724 |
+
|
| 725 |
+
print(f"[INFO] Starting worker {idx} on GPU {gpu_id}")
|
| 726 |
+
p = subprocess.Popen(cmd, env=env)
|
| 727 |
+
processes.append(p)
|
| 728 |
+
|
| 729 |
+
# Wait for all workers
|
| 730 |
+
for p in processes:
|
| 731 |
+
p.wait()
|
| 732 |
+
|
| 733 |
+
print(f"[INFO] All {num_gpus} workers completed")
|
| 734 |
+
|
| 735 |
+
# Run evaluation after all inference is done
|
| 736 |
+
metrics = evaluator.eval()
|
| 737 |
+
evaluator.print_metrics(metrics)
|
| 738 |
+
else:
|
| 739 |
+
# Single GPU or worker process
|
| 740 |
+
from depth_anything_3.api import DepthAnything3
|
| 741 |
+
|
| 742 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 743 |
+
api = DepthAnything3.from_pretrained(model_path)
|
| 744 |
+
api = api.to(device)
|
| 745 |
+
|
| 746 |
+
evaluator.infer(api, model_path=model_path)
|
| 747 |
+
|
| 748 |
+
# Only run eval if single GPU mode (workers don't eval)
|
| 749 |
+
if not is_worker:
|
| 750 |
+
metrics = evaluator.eval()
|
| 751 |
+
evaluator.print_metrics(metrics)
|
| 752 |
+
|
depth_anything_3/bench/print_metrics.py
ADDED
|
@@ -0,0 +1,618 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Beautiful metrics printing utilities for benchmark evaluation.
|
| 17 |
+
|
| 18 |
+
Provides colorized, well-formatted tabular output for evaluation results.
|
| 19 |
+
Supports highlighting best/worst values and grouping by dataset/mode.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import argparse
|
| 23 |
+
import json
|
| 24 |
+
import os
|
| 25 |
+
import re
|
| 26 |
+
from typing import Dict as TDict, List, Optional
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
# ANSI color codes for terminal output
|
| 30 |
+
class Colors:
|
| 31 |
+
"""ANSI escape codes for terminal colors."""
|
| 32 |
+
|
| 33 |
+
RESET = "\033[0m"
|
| 34 |
+
BOLD = "\033[1m"
|
| 35 |
+
RED = "\033[31m"
|
| 36 |
+
GREEN = "\033[32m"
|
| 37 |
+
YELLOW = "\033[33m"
|
| 38 |
+
BLUE = "\033[34m"
|
| 39 |
+
MAGENTA = "\033[35m"
|
| 40 |
+
CYAN = "\033[36m"
|
| 41 |
+
WHITE = "\033[37m"
|
| 42 |
+
|
| 43 |
+
# Bold variants
|
| 44 |
+
BOLD_RED = "\033[1;31m"
|
| 45 |
+
BOLD_GREEN = "\033[1;32m"
|
| 46 |
+
BOLD_YELLOW = "\033[1;33m"
|
| 47 |
+
BOLD_BLUE = "\033[1;34m"
|
| 48 |
+
BOLD_MAGENTA = "\033[1;35m"
|
| 49 |
+
BOLD_CYAN = "\033[1;36m"
|
| 50 |
+
|
| 51 |
+
# Background
|
| 52 |
+
BG_DARK = "\033[48;5;236m"
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def strip_ansi(text: str) -> str:
|
| 56 |
+
"""Remove ANSI escape sequences from string for length calculation."""
|
| 57 |
+
ansi_escape = re.compile(r"\x1b\[[0-9;]*m")
|
| 58 |
+
return ansi_escape.sub("", text)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def colorize_value(
|
| 62 |
+
value: str,
|
| 63 |
+
is_best: bool = False,
|
| 64 |
+
is_worst: bool = False,
|
| 65 |
+
lower_is_better: bool = False,
|
| 66 |
+
) -> str:
|
| 67 |
+
"""
|
| 68 |
+
Apply color to a metric value based on whether it's best/worst.
|
| 69 |
+
|
| 70 |
+
Args:
|
| 71 |
+
value: String representation of the value
|
| 72 |
+
is_best: Whether this is the best value in its column
|
| 73 |
+
is_worst: Whether this is the worst value in its column
|
| 74 |
+
lower_is_better: If True, lower values are better (e.g., error metrics)
|
| 75 |
+
|
| 76 |
+
Returns:
|
| 77 |
+
Colorized string
|
| 78 |
+
"""
|
| 79 |
+
if lower_is_better:
|
| 80 |
+
# For metrics like error/distance, lower is better
|
| 81 |
+
if is_best:
|
| 82 |
+
return f"{Colors.BOLD_GREEN}{value}{Colors.RESET}"
|
| 83 |
+
elif is_worst:
|
| 84 |
+
return f"{Colors.BOLD_RED}{value}{Colors.RESET}"
|
| 85 |
+
else:
|
| 86 |
+
# For metrics like accuracy/AUC, higher is better
|
| 87 |
+
if is_best:
|
| 88 |
+
return f"{Colors.BOLD_GREEN}{value}{Colors.RESET}"
|
| 89 |
+
elif is_worst:
|
| 90 |
+
return f"{Colors.BOLD_RED}{value}{Colors.RESET}"
|
| 91 |
+
return value
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
class MetricsPrinter:
|
| 95 |
+
"""
|
| 96 |
+
Beautiful tabular metrics printer with color support.
|
| 97 |
+
|
| 98 |
+
Features:
|
| 99 |
+
- Colorized best/worst values
|
| 100 |
+
- Grouped by dataset and evaluation mode
|
| 101 |
+
- Automatic column width calculation
|
| 102 |
+
- Support for multiple input directories comparison
|
| 103 |
+
"""
|
| 104 |
+
|
| 105 |
+
# Metrics where lower values are better
|
| 106 |
+
LOWER_IS_BETTER = {"comp", "acc", "overall", "error", "loss", "rmse", "mae"}
|
| 107 |
+
|
| 108 |
+
def __init__(self, use_color: bool = True):
|
| 109 |
+
"""
|
| 110 |
+
Initialize the printer.
|
| 111 |
+
|
| 112 |
+
Args:
|
| 113 |
+
use_color: Whether to use ANSI colors in output
|
| 114 |
+
"""
|
| 115 |
+
self.use_color = use_color
|
| 116 |
+
|
| 117 |
+
def print_results(self, metrics: TDict[str, dict], summary_only: bool = True) -> None:
|
| 118 |
+
"""
|
| 119 |
+
Print evaluation metrics in a beautiful tabular format.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
metrics: Dictionary mapping "dataset_mode" to metric results
|
| 123 |
+
summary_only: If True, only print summary table. If False, print per-dataset details too.
|
| 124 |
+
"""
|
| 125 |
+
if not metrics:
|
| 126 |
+
print(f"\n{Colors.BOLD_RED}❌ No evaluation metrics found{Colors.RESET}")
|
| 127 |
+
return
|
| 128 |
+
|
| 129 |
+
if not summary_only:
|
| 130 |
+
self._print_header()
|
| 131 |
+
grouped = self._group_by_dataset(metrics)
|
| 132 |
+
|
| 133 |
+
for dataset, modes_data in grouped.items():
|
| 134 |
+
self._print_dataset_section(dataset, modes_data)
|
| 135 |
+
|
| 136 |
+
# Print summary table with average metrics across datasets
|
| 137 |
+
self._print_summary(metrics)
|
| 138 |
+
|
| 139 |
+
self._print_footer()
|
| 140 |
+
|
| 141 |
+
def print_comparison(
|
| 142 |
+
self,
|
| 143 |
+
metrics_list: List[TDict[str, dict]],
|
| 144 |
+
labels: List[str],
|
| 145 |
+
) -> None:
|
| 146 |
+
"""
|
| 147 |
+
Print comparison table for multiple evaluation runs.
|
| 148 |
+
|
| 149 |
+
Args:
|
| 150 |
+
metrics_list: List of metrics dictionaries
|
| 151 |
+
labels: Labels for each metrics dictionary
|
| 152 |
+
"""
|
| 153 |
+
if not metrics_list or not all(metrics_list):
|
| 154 |
+
print(f"\n{Colors.BOLD_RED}❌ No metrics to compare{Colors.RESET}")
|
| 155 |
+
return
|
| 156 |
+
|
| 157 |
+
# Collect all datasets and modes
|
| 158 |
+
all_keys = set()
|
| 159 |
+
for metrics in metrics_list:
|
| 160 |
+
all_keys.update(metrics.keys())
|
| 161 |
+
|
| 162 |
+
self._print_header("COMPARISON")
|
| 163 |
+
|
| 164 |
+
for key in sorted(all_keys):
|
| 165 |
+
parts = key.rsplit("_", 1)
|
| 166 |
+
if len(parts) == 2:
|
| 167 |
+
dataset, mode = parts[0], parts[1]
|
| 168 |
+
else:
|
| 169 |
+
dataset, mode = key, "unknown"
|
| 170 |
+
|
| 171 |
+
print(f"\n{Colors.BOLD_CYAN}📊 {dataset.upper()} - {mode.upper()}{Colors.RESET}")
|
| 172 |
+
print("-" * 100)
|
| 173 |
+
|
| 174 |
+
# Collect metrics from all runs
|
| 175 |
+
all_metric_names = set()
|
| 176 |
+
for metrics in metrics_list:
|
| 177 |
+
if key in metrics and "mean" in metrics[key]:
|
| 178 |
+
all_metric_names.update(metrics[key]["mean"].keys())
|
| 179 |
+
|
| 180 |
+
if not all_metric_names:
|
| 181 |
+
continue
|
| 182 |
+
|
| 183 |
+
# Build comparison table
|
| 184 |
+
metric_width = max(15, max(len(m) for m in all_metric_names) + 2)
|
| 185 |
+
label_width = max(15, max(len(l) for l in labels) + 2)
|
| 186 |
+
|
| 187 |
+
# Header
|
| 188 |
+
header = f"{'Metric':<{metric_width}}"
|
| 189 |
+
for label in labels:
|
| 190 |
+
header += f"{label:<{label_width}}"
|
| 191 |
+
print(header)
|
| 192 |
+
print("-" * len(strip_ansi(header)))
|
| 193 |
+
|
| 194 |
+
# Collect values for highlighting
|
| 195 |
+
for metric_name in sorted(all_metric_names):
|
| 196 |
+
values = []
|
| 197 |
+
for metrics in metrics_list:
|
| 198 |
+
if key in metrics and "mean" in metrics[key]:
|
| 199 |
+
val = metrics[key]["mean"].get(metric_name)
|
| 200 |
+
values.append(val if val is not None else float("nan"))
|
| 201 |
+
else:
|
| 202 |
+
values.append(float("nan"))
|
| 203 |
+
|
| 204 |
+
# Find best/worst
|
| 205 |
+
valid_values = [v for v in values if not (v != v)] # Filter NaN
|
| 206 |
+
if valid_values:
|
| 207 |
+
lower_better = any(
|
| 208 |
+
lb in metric_name.lower() for lb in self.LOWER_IS_BETTER
|
| 209 |
+
)
|
| 210 |
+
best_val = min(valid_values) if lower_better else max(valid_values)
|
| 211 |
+
worst_val = max(valid_values) if lower_better else min(valid_values)
|
| 212 |
+
else:
|
| 213 |
+
best_val = worst_val = None
|
| 214 |
+
|
| 215 |
+
# Print row
|
| 216 |
+
row = f"{metric_name:<{metric_width}}"
|
| 217 |
+
for val in values:
|
| 218 |
+
if val != val: # NaN check
|
| 219 |
+
val_str = "N/A"
|
| 220 |
+
else:
|
| 221 |
+
val_str = f"{val:.4f}"
|
| 222 |
+
if self.use_color and len(valid_values) > 1:
|
| 223 |
+
lower_better = any(
|
| 224 |
+
lb in metric_name.lower() for lb in self.LOWER_IS_BETTER
|
| 225 |
+
)
|
| 226 |
+
is_best = abs(val - best_val) < 1e-8 if best_val else False
|
| 227 |
+
is_worst = abs(val - worst_val) < 1e-8 if worst_val else False
|
| 228 |
+
val_str_padded = f"{val_str:<{label_width}}"
|
| 229 |
+
val_str = colorize_value(
|
| 230 |
+
val_str_padded, is_best, is_worst, lower_better
|
| 231 |
+
)
|
| 232 |
+
row += val_str
|
| 233 |
+
continue
|
| 234 |
+
row += f"{val_str:<{label_width}}"
|
| 235 |
+
print(row)
|
| 236 |
+
|
| 237 |
+
self._print_footer()
|
| 238 |
+
|
| 239 |
+
def _print_header(self, title: str = "EVALUATION RESULTS") -> None:
|
| 240 |
+
"""Print report header."""
|
| 241 |
+
width = 100
|
| 242 |
+
print()
|
| 243 |
+
print("=" * width)
|
| 244 |
+
print(f"{Colors.BOLD_CYAN}📊 DEPTH ANYTHING 3 {title}{Colors.RESET}")
|
| 245 |
+
print("=" * width)
|
| 246 |
+
|
| 247 |
+
def _print_footer(self) -> None:
|
| 248 |
+
"""Print report footer."""
|
| 249 |
+
width = 100
|
| 250 |
+
print()
|
| 251 |
+
print("=" * width)
|
| 252 |
+
print(f"{Colors.BOLD_GREEN}✅ Evaluation Complete{Colors.RESET}")
|
| 253 |
+
print("=" * width)
|
| 254 |
+
print()
|
| 255 |
+
|
| 256 |
+
def _group_by_dataset(self, metrics: TDict[str, dict]) -> TDict[str, dict]:
|
| 257 |
+
"""Group metrics by dataset."""
|
| 258 |
+
grouped = {}
|
| 259 |
+
for key, data in metrics.items():
|
| 260 |
+
if not isinstance(data, dict) or "mean" not in data:
|
| 261 |
+
continue
|
| 262 |
+
# Parse key format: "dataset_mode" (e.g., "dtu_recon_unposed")
|
| 263 |
+
parts = key.split("_", 1)
|
| 264 |
+
if len(parts) == 2:
|
| 265 |
+
dataset, mode = parts
|
| 266 |
+
if dataset not in grouped:
|
| 267 |
+
grouped[dataset] = {}
|
| 268 |
+
grouped[dataset][mode] = data
|
| 269 |
+
return grouped
|
| 270 |
+
|
| 271 |
+
def _print_dataset_section(self, dataset: str, modes_data: TDict[str, dict]) -> None:
|
| 272 |
+
"""Print metrics section for a single dataset."""
|
| 273 |
+
print(f"\n{Colors.BOLD_MAGENTA}🔍 {dataset.upper()}{Colors.RESET}")
|
| 274 |
+
print("-" * 100)
|
| 275 |
+
|
| 276 |
+
# Collect all unique metrics across all modes
|
| 277 |
+
all_metrics = set()
|
| 278 |
+
for mode_data in modes_data.values():
|
| 279 |
+
all_metrics.update(mode_data["mean"].keys())
|
| 280 |
+
all_metrics = sorted(list(all_metrics))
|
| 281 |
+
|
| 282 |
+
if not all_metrics:
|
| 283 |
+
print(" No metrics available")
|
| 284 |
+
return
|
| 285 |
+
|
| 286 |
+
# Calculate column widths
|
| 287 |
+
metric_width = max(18, max(len(m) for m in all_metrics) + 2)
|
| 288 |
+
mode_width = 18
|
| 289 |
+
modes = list(modes_data.keys())
|
| 290 |
+
|
| 291 |
+
# Print header
|
| 292 |
+
header = f"{'Metric':<{metric_width}}"
|
| 293 |
+
for mode in modes:
|
| 294 |
+
header += f"{mode.upper():<{mode_width}}"
|
| 295 |
+
print(f"{Colors.BOLD}{header}{Colors.RESET}")
|
| 296 |
+
print("-" * len(header))
|
| 297 |
+
|
| 298 |
+
# Print each metric row
|
| 299 |
+
for metric in all_metrics:
|
| 300 |
+
row = f"{metric:<{metric_width}}"
|
| 301 |
+
|
| 302 |
+
# Collect values for this metric across modes
|
| 303 |
+
values = []
|
| 304 |
+
for mode in modes:
|
| 305 |
+
if metric in modes_data[mode]["mean"]:
|
| 306 |
+
values.append(modes_data[mode]["mean"][metric])
|
| 307 |
+
else:
|
| 308 |
+
values.append(None)
|
| 309 |
+
|
| 310 |
+
# Find best/worst values
|
| 311 |
+
valid_values = [v for v in values if v is not None]
|
| 312 |
+
if valid_values:
|
| 313 |
+
lower_better = any(lb in metric.lower() for lb in self.LOWER_IS_BETTER)
|
| 314 |
+
best_val = min(valid_values) if lower_better else max(valid_values)
|
| 315 |
+
worst_val = max(valid_values) if lower_better else min(valid_values)
|
| 316 |
+
else:
|
| 317 |
+
best_val = worst_val = None
|
| 318 |
+
|
| 319 |
+
# Format each value
|
| 320 |
+
for val in values:
|
| 321 |
+
if val is None:
|
| 322 |
+
row += f"{'N/A':<{mode_width}}"
|
| 323 |
+
else:
|
| 324 |
+
val_str = f"{val:.4f}"
|
| 325 |
+
if self.use_color and len(valid_values) > 1:
|
| 326 |
+
is_best = abs(val - best_val) < 1e-8 if best_val else False
|
| 327 |
+
is_worst = abs(val - worst_val) < 1e-8 if worst_val else False
|
| 328 |
+
lower_better = any(
|
| 329 |
+
lb in metric.lower() for lb in self.LOWER_IS_BETTER
|
| 330 |
+
)
|
| 331 |
+
# Pad before colorizing to maintain alignment
|
| 332 |
+
val_str_padded = f"{val_str:<{mode_width}}"
|
| 333 |
+
row += colorize_value(
|
| 334 |
+
val_str_padded, is_best, is_worst, lower_better
|
| 335 |
+
)
|
| 336 |
+
else:
|
| 337 |
+
row += f"{val_str:<{mode_width}}"
|
| 338 |
+
print(row)
|
| 339 |
+
|
| 340 |
+
# Show scene counts
|
| 341 |
+
scene_info = []
|
| 342 |
+
for mode, mode_data in modes_data.items():
|
| 343 |
+
scene_count = len([k for k in mode_data.keys() if k != "mean"])
|
| 344 |
+
scene_info.append(f"{mode}: {scene_count} scenes")
|
| 345 |
+
print(f"\n{Colors.CYAN}📈 {' | '.join(scene_info)}{Colors.RESET}")
|
| 346 |
+
|
| 347 |
+
def _print_summary(self, metrics: TDict[str, dict]) -> None:
|
| 348 |
+
"""
|
| 349 |
+
Print summary table with key metrics across all datasets.
|
| 350 |
+
|
| 351 |
+
Format: One row per metric, datasets as columns.
|
| 352 |
+
Order: HiRoom, ETH3D, DTU, 7Scenes, ScanNet++, (DTU-64 for pose only)
|
| 353 |
+
"""
|
| 354 |
+
print(f"\n{Colors.BOLD_CYAN}{'=' * 120}{Colors.RESET}")
|
| 355 |
+
print(f"{Colors.BOLD_CYAN}📊 SUMMARY{Colors.RESET}")
|
| 356 |
+
print(f"{Colors.BOLD_CYAN}{'=' * 120}{Colors.RESET}")
|
| 357 |
+
|
| 358 |
+
# Dataset display order and names
|
| 359 |
+
DATASET_ORDER = ["hiroom", "eth3d", "dtu", "7scenes", "scannetpp", "dtu64"]
|
| 360 |
+
DATASET_DISPLAY = {
|
| 361 |
+
"hiroom": "HiRoom",
|
| 362 |
+
"eth3d": "ETH3D",
|
| 363 |
+
"dtu": "DTU",
|
| 364 |
+
"7scenes": "7Scenes",
|
| 365 |
+
"scannetpp": "ScanNet++",
|
| 366 |
+
"dtu64": "DTU-64",
|
| 367 |
+
}
|
| 368 |
+
|
| 369 |
+
# Collect all metrics into a structured dict
|
| 370 |
+
# metric_data[dataset][mode] = {"Auc_3": x, "Auc_30": x, "fscore": x, "overall": x}
|
| 371 |
+
metric_data = {}
|
| 372 |
+
for key, data in metrics.items():
|
| 373 |
+
if not isinstance(data, dict) or "mean" not in data:
|
| 374 |
+
continue
|
| 375 |
+
parts = key.split("_", 1)
|
| 376 |
+
if len(parts) != 2:
|
| 377 |
+
continue
|
| 378 |
+
dataset, mode = parts
|
| 379 |
+
dataset_lower = dataset.lower()
|
| 380 |
+
if dataset_lower not in metric_data:
|
| 381 |
+
metric_data[dataset_lower] = {}
|
| 382 |
+
metric_data[dataset_lower][mode] = data["mean"]
|
| 383 |
+
|
| 384 |
+
col_width = 12
|
| 385 |
+
|
| 386 |
+
def fmt_val(val):
|
| 387 |
+
"""Format value or return N/A."""
|
| 388 |
+
if val is None:
|
| 389 |
+
return "N/A"
|
| 390 |
+
return f"{val:.4f}"
|
| 391 |
+
|
| 392 |
+
def get_metric(dataset, mode, metric_name):
|
| 393 |
+
"""Get metric value or None."""
|
| 394 |
+
if dataset not in metric_data:
|
| 395 |
+
return None
|
| 396 |
+
if mode not in metric_data[dataset]:
|
| 397 |
+
return None
|
| 398 |
+
return metric_data[dataset][mode].get(metric_name)
|
| 399 |
+
|
| 400 |
+
# ============ POSE METRICS ============
|
| 401 |
+
print(f"\n{Colors.BOLD_MAGENTA}🎯 POSE ESTIMATION{Colors.RESET}")
|
| 402 |
+
|
| 403 |
+
# Pose: show all datasets except DTU (keep DTU-64 only)
|
| 404 |
+
# Order: HiRoom, ETH3D, DTU-64, 7Scenes, ScanNet++
|
| 405 |
+
pose_datasets = ["hiroom", "eth3d", "dtu64", "7scenes", "scannetpp"]
|
| 406 |
+
|
| 407 |
+
# Header: Avg first, then datasets
|
| 408 |
+
header = f"{'Metric':<15}{'Avg':<{col_width}}"
|
| 409 |
+
for ds in pose_datasets:
|
| 410 |
+
header += f"{DATASET_DISPLAY[ds]:<{col_width}}"
|
| 411 |
+
print("-" * len(strip_ansi(header)))
|
| 412 |
+
print(f"{Colors.BOLD}{header}{Colors.RESET}")
|
| 413 |
+
print("-" * len(strip_ansi(header)))
|
| 414 |
+
|
| 415 |
+
# Helper to get metric with fallback names
|
| 416 |
+
def get_pose_metric(dataset, metric_name):
|
| 417 |
+
"""Get pose metric with fallback for different naming conventions."""
|
| 418 |
+
# Try different naming conventions
|
| 419 |
+
names = {
|
| 420 |
+
"Auc3": ["Auc_3", "auc03", "auc_3", "AUC_3", "Auc3", "auc3"],
|
| 421 |
+
"Auc30": ["Auc_30", "auc30", "auc_30", "AUC_30", "Auc30"],
|
| 422 |
+
}
|
| 423 |
+
for name in names.get(metric_name, [metric_name]):
|
| 424 |
+
val = get_metric(dataset, "pose", name)
|
| 425 |
+
if val is not None:
|
| 426 |
+
return val
|
| 427 |
+
return None
|
| 428 |
+
|
| 429 |
+
# Auc3 row
|
| 430 |
+
values = []
|
| 431 |
+
for ds in pose_datasets:
|
| 432 |
+
val = get_pose_metric(ds, "Auc3")
|
| 433 |
+
if val is not None:
|
| 434 |
+
values.append(val)
|
| 435 |
+
avg = sum(values) / len(values) if values else None
|
| 436 |
+
row = f"{'Auc3':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}"
|
| 437 |
+
for ds in pose_datasets:
|
| 438 |
+
val = get_pose_metric(ds, "Auc3")
|
| 439 |
+
row += f"{fmt_val(val):<{col_width}}"
|
| 440 |
+
print(row)
|
| 441 |
+
|
| 442 |
+
# Auc30 row
|
| 443 |
+
values = []
|
| 444 |
+
for ds in pose_datasets:
|
| 445 |
+
val = get_pose_metric(ds, "Auc30")
|
| 446 |
+
if val is not None:
|
| 447 |
+
values.append(val)
|
| 448 |
+
avg = sum(values) / len(values) if values else None
|
| 449 |
+
row = f"{'Auc30':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}"
|
| 450 |
+
for ds in pose_datasets:
|
| 451 |
+
val = get_pose_metric(ds, "Auc30")
|
| 452 |
+
row += f"{fmt_val(val):<{col_width}}"
|
| 453 |
+
print(row)
|
| 454 |
+
|
| 455 |
+
# ============ RECON_UNPOSED METRICS ============
|
| 456 |
+
print(f"\n{Colors.BOLD_MAGENTA}🏗️ RECON_UNPOSED (Pred Pose){Colors.RESET}")
|
| 457 |
+
|
| 458 |
+
# For recon, exclude dtu64 from columns
|
| 459 |
+
recon_datasets = ["hiroom", "eth3d", "dtu", "7scenes", "scannetpp"]
|
| 460 |
+
avg_datasets = ["hiroom", "eth3d", "7scenes", "scannetpp"] # Exclude DTU from avg
|
| 461 |
+
|
| 462 |
+
# Header: Avg first, then datasets
|
| 463 |
+
header = f"{'Metric':<15}{'Avg*':<{col_width}}"
|
| 464 |
+
for ds in recon_datasets:
|
| 465 |
+
header += f"{DATASET_DISPLAY[ds]:<{col_width}}"
|
| 466 |
+
print("-" * len(strip_ansi(header)))
|
| 467 |
+
print(f"{Colors.BOLD}{header}{Colors.RESET}")
|
| 468 |
+
print("-" * len(strip_ansi(header)))
|
| 469 |
+
|
| 470 |
+
# F-score row (only metric for avg)
|
| 471 |
+
values = []
|
| 472 |
+
for ds in recon_datasets:
|
| 473 |
+
val = get_metric(ds, "recon_unposed", "fscore")
|
| 474 |
+
if val is not None and ds in avg_datasets:
|
| 475 |
+
values.append(val)
|
| 476 |
+
avg = sum(values) / len(values) if values else None
|
| 477 |
+
row = f"{'F-score':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}"
|
| 478 |
+
for ds in recon_datasets:
|
| 479 |
+
val = get_metric(ds, "recon_unposed", "fscore")
|
| 480 |
+
row += f"{fmt_val(val):<{col_width}}"
|
| 481 |
+
print(row)
|
| 482 |
+
|
| 483 |
+
# Overall row (avg over 4 datasets excluding DTU)
|
| 484 |
+
values = []
|
| 485 |
+
for ds in recon_datasets:
|
| 486 |
+
val = get_metric(ds, "recon_unposed", "overall")
|
| 487 |
+
if val is not None and ds in avg_datasets:
|
| 488 |
+
values.append(val)
|
| 489 |
+
avg = sum(values) / len(values) if values else None
|
| 490 |
+
row = f"{'Overall':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}"
|
| 491 |
+
for ds in recon_datasets:
|
| 492 |
+
val = get_metric(ds, "recon_unposed", "overall")
|
| 493 |
+
row += f"{fmt_val(val):<{col_width}}"
|
| 494 |
+
print(row)
|
| 495 |
+
|
| 496 |
+
# ============ RECON_POSED METRICS ============
|
| 497 |
+
print(f"\n{Colors.BOLD_MAGENTA}🏗️ RECON_POSED (GT Pose){Colors.RESET}")
|
| 498 |
+
|
| 499 |
+
# Header: Avg first, then datasets
|
| 500 |
+
header = f"{'Metric':<15}{'Avg*':<{col_width}}"
|
| 501 |
+
for ds in recon_datasets:
|
| 502 |
+
header += f"{DATASET_DISPLAY[ds]:<{col_width}}"
|
| 503 |
+
print("-" * len(strip_ansi(header)))
|
| 504 |
+
print(f"{Colors.BOLD}{header}{Colors.RESET}")
|
| 505 |
+
print("-" * len(strip_ansi(header)))
|
| 506 |
+
|
| 507 |
+
# F-score row (only metric for avg)
|
| 508 |
+
values = []
|
| 509 |
+
for ds in recon_datasets:
|
| 510 |
+
val = get_metric(ds, "recon_posed", "fscore")
|
| 511 |
+
if val is not None and ds in avg_datasets:
|
| 512 |
+
values.append(val)
|
| 513 |
+
avg = sum(values) / len(values) if values else None
|
| 514 |
+
row = f"{'F-score':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}"
|
| 515 |
+
for ds in recon_datasets:
|
| 516 |
+
val = get_metric(ds, "recon_posed", "fscore")
|
| 517 |
+
row += f"{fmt_val(val):<{col_width}}"
|
| 518 |
+
print(row)
|
| 519 |
+
|
| 520 |
+
# Overall row (avg over 4 datasets excluding DTU)
|
| 521 |
+
values = []
|
| 522 |
+
for ds in recon_datasets:
|
| 523 |
+
val = get_metric(ds, "recon_posed", "overall")
|
| 524 |
+
if val is not None and ds in avg_datasets:
|
| 525 |
+
values.append(val)
|
| 526 |
+
avg = sum(values) / len(values) if values else None
|
| 527 |
+
row = f"{'Overall':<15}{Colors.BOLD_GREEN}{fmt_val(avg):<{col_width}}{Colors.RESET}"
|
| 528 |
+
for ds in recon_datasets:
|
| 529 |
+
val = get_metric(ds, "recon_posed", "overall")
|
| 530 |
+
row += f"{fmt_val(val):<{col_width}}"
|
| 531 |
+
print(row)
|
| 532 |
+
|
| 533 |
+
print(f"\n{Colors.CYAN}* Avg F-score / Overall = average over HiRoom, ETH3D, 7Scenes, ScanNet++ (4 datasets){Colors.RESET}")
|
| 534 |
+
|
| 535 |
+
|
| 536 |
+
def load_metrics_from_dir(metric_dir: str) -> TDict[str, dict]:
|
| 537 |
+
"""
|
| 538 |
+
Load all metrics JSON files from a directory.
|
| 539 |
+
|
| 540 |
+
Args:
|
| 541 |
+
metric_dir: Path to directory containing metric JSON files
|
| 542 |
+
|
| 543 |
+
Returns:
|
| 544 |
+
Dictionary mapping filename (without .json) to metric data
|
| 545 |
+
"""
|
| 546 |
+
metrics = {}
|
| 547 |
+
if not os.path.exists(metric_dir):
|
| 548 |
+
return metrics
|
| 549 |
+
|
| 550 |
+
for filename in os.listdir(metric_dir):
|
| 551 |
+
if filename.endswith(".json"):
|
| 552 |
+
filepath = os.path.join(metric_dir, filename)
|
| 553 |
+
try:
|
| 554 |
+
with open(filepath, encoding="utf-8") as f:
|
| 555 |
+
content = f.read()
|
| 556 |
+
# Handle trailing commas in JSON
|
| 557 |
+
content = re.sub(r",\s*([\]\}])", r"\1", content)
|
| 558 |
+
data = json.loads(content)
|
| 559 |
+
key = filename[:-5]
|
| 560 |
+
metrics[key] = data
|
| 561 |
+
except Exception as e:
|
| 562 |
+
print(f"Warning: Failed to load {filename}: {e}")
|
| 563 |
+
|
| 564 |
+
return metrics
|
| 565 |
+
|
| 566 |
+
|
| 567 |
+
def main():
|
| 568 |
+
"""Command-line interface for metrics printing."""
|
| 569 |
+
parser = argparse.ArgumentParser(
|
| 570 |
+
description="Print DepthAnything3 benchmark evaluation metrics."
|
| 571 |
+
)
|
| 572 |
+
parser.add_argument(
|
| 573 |
+
"--input_dir",
|
| 574 |
+
type=str,
|
| 575 |
+
default="./eval_workspace/metric_results",
|
| 576 |
+
help="Directory containing metric JSON files (comma-separated for comparison)",
|
| 577 |
+
)
|
| 578 |
+
parser.add_argument(
|
| 579 |
+
"--no_color",
|
| 580 |
+
action="store_true",
|
| 581 |
+
help="Disable colored output",
|
| 582 |
+
)
|
| 583 |
+
parser.add_argument(
|
| 584 |
+
"--key",
|
| 585 |
+
type=str,
|
| 586 |
+
default=None,
|
| 587 |
+
help="Specific metric key to highlight",
|
| 588 |
+
)
|
| 589 |
+
args = parser.parse_args()
|
| 590 |
+
|
| 591 |
+
# Support multiple directories for comparison
|
| 592 |
+
input_dirs = [d.strip() for d in args.input_dir.split(",") if d.strip()]
|
| 593 |
+
|
| 594 |
+
printer = MetricsPrinter(use_color=not args.no_color)
|
| 595 |
+
|
| 596 |
+
if len(input_dirs) == 1:
|
| 597 |
+
# Single directory - simple print
|
| 598 |
+
metrics = load_metrics_from_dir(input_dirs[0])
|
| 599 |
+
printer.print_results(metrics)
|
| 600 |
+
else:
|
| 601 |
+
# Multiple directories - comparison mode
|
| 602 |
+
metrics_list = []
|
| 603 |
+
labels = []
|
| 604 |
+
for d in input_dirs:
|
| 605 |
+
metrics = load_metrics_from_dir(d)
|
| 606 |
+
if metrics:
|
| 607 |
+
metrics_list.append(metrics)
|
| 608 |
+
labels.append(os.path.basename(d.rstrip("/")))
|
| 609 |
+
|
| 610 |
+
if metrics_list:
|
| 611 |
+
printer.print_comparison(metrics_list, labels)
|
| 612 |
+
else:
|
| 613 |
+
print("No metrics found in specified directories")
|
| 614 |
+
|
| 615 |
+
|
| 616 |
+
if __name__ == "__main__":
|
| 617 |
+
main()
|
| 618 |
+
|
depth_anything_3/bench/registries.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Auto-loading registry system for benchmark datasets.
|
| 17 |
+
|
| 18 |
+
This module provides registry classes that automatically discover and import
|
| 19 |
+
dataset implementations from the datasets subpackage on first access.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import importlib
|
| 23 |
+
import pkgutil
|
| 24 |
+
import threading
|
| 25 |
+
|
| 26 |
+
from depth_anything_3.utils.registry import Registry
|
| 27 |
+
|
| 28 |
+
__all__ = ["METRIC_REGISTRY", "MONO_REGISTRY", "MV_REGISTRY", "NVS_REGISTRY"]
|
| 29 |
+
|
| 30 |
+
# ---- Lazy import: Only scan and import all datasets submodules on first registry access ----
|
| 31 |
+
_loaded = False
|
| 32 |
+
_lock = threading.Lock()
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def _import_all_datasets_once():
|
| 36 |
+
"""
|
| 37 |
+
Scan and import all .py submodules under depth_anything_3.bench.datasets
|
| 38 |
+
(skip files/packages starting with underscore), to trigger @REGISTRY.register(...) in each module.
|
| 39 |
+
"""
|
| 40 |
+
global _loaded
|
| 41 |
+
if _loaded:
|
| 42 |
+
return
|
| 43 |
+
|
| 44 |
+
with _lock:
|
| 45 |
+
if _loaded:
|
| 46 |
+
return
|
| 47 |
+
|
| 48 |
+
pkg_name = "depth_anything_3.bench.datasets"
|
| 49 |
+
pkg = importlib.import_module(pkg_name)
|
| 50 |
+
pkg_paths = list(getattr(pkg, "__path__", []))
|
| 51 |
+
|
| 52 |
+
for finder, name, ispkg in pkgutil.walk_packages(pkg_paths, prefix=pkg_name + "."):
|
| 53 |
+
base = name.rsplit(".", 1)[-1]
|
| 54 |
+
if base.startswith("_"):
|
| 55 |
+
continue
|
| 56 |
+
try:
|
| 57 |
+
importlib.import_module(name)
|
| 58 |
+
except Exception as e:
|
| 59 |
+
print(f"[datasets auto-import] Failed to import {name}: {e}")
|
| 60 |
+
|
| 61 |
+
_loaded = True
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
class AutoRegistry(Registry):
|
| 65 |
+
"""Registry that ensures all datasets are auto-discovered and imported on first use."""
|
| 66 |
+
|
| 67 |
+
def get(self, name):
|
| 68 |
+
_import_all_datasets_once()
|
| 69 |
+
return super().get(name)
|
| 70 |
+
|
| 71 |
+
def all(self):
|
| 72 |
+
_import_all_datasets_once()
|
| 73 |
+
return super().all()
|
| 74 |
+
|
| 75 |
+
def has(self, name):
|
| 76 |
+
_import_all_datasets_once()
|
| 77 |
+
return name in self._map
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
# Four auto-lazy registry instances for different evaluation types
|
| 81 |
+
METRIC_REGISTRY = AutoRegistry() # For metric depth evaluation
|
| 82 |
+
MONO_REGISTRY = AutoRegistry() # For monocular depth evaluation
|
| 83 |
+
MV_REGISTRY = AutoRegistry() # For multi-view evaluation
|
| 84 |
+
NVS_REGISTRY = AutoRegistry() # For novel view synthesis evaluation
|
| 85 |
+
|
depth_anything_3/bench/utils.py
ADDED
|
@@ -0,0 +1,525 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Utility functions for benchmark evaluation.
|
| 17 |
+
|
| 18 |
+
Contains:
|
| 19 |
+
- Pose evaluation metrics (AUC) and helper functions
|
| 20 |
+
- 3D reconstruction evaluation metrics (Acc/Comp/F-score)
|
| 21 |
+
- Geometry utilities (quaternion conversion, etc.)
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
from typing import Dict as TDict, Optional, Tuple, Union
|
| 25 |
+
|
| 26 |
+
import numpy as np
|
| 27 |
+
import open3d as o3d
|
| 28 |
+
import torch
|
| 29 |
+
from addict import Dict
|
| 30 |
+
from scipy.spatial import KDTree
|
| 31 |
+
|
| 32 |
+
from depth_anything_3.utils.geometry import mat_to_quat
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
# =============================================================================
|
| 36 |
+
# Geometry Utilities
|
| 37 |
+
# =============================================================================
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def quat2rotmat(qvec: list) -> np.ndarray:
|
| 41 |
+
"""
|
| 42 |
+
Convert quaternion (WXYZ order) to rotation matrix.
|
| 43 |
+
|
| 44 |
+
Args:
|
| 45 |
+
qvec: Quaternion as [w, x, y, z]
|
| 46 |
+
|
| 47 |
+
Returns:
|
| 48 |
+
3x3 rotation matrix
|
| 49 |
+
"""
|
| 50 |
+
rotmat = np.array(
|
| 51 |
+
[
|
| 52 |
+
1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
|
| 53 |
+
2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
|
| 54 |
+
2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
|
| 55 |
+
2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
|
| 56 |
+
1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
|
| 57 |
+
2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
|
| 58 |
+
2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
|
| 59 |
+
2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
|
| 60 |
+
1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
|
| 61 |
+
]
|
| 62 |
+
)
|
| 63 |
+
rotmat = rotmat.reshape(3, 3)
|
| 64 |
+
return rotmat
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
# =============================================================================
|
| 68 |
+
# 3D Reconstruction Evaluation
|
| 69 |
+
# =============================================================================
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def nn_correspondance(verts1: np.ndarray, verts2: np.ndarray) -> np.ndarray:
|
| 73 |
+
"""
|
| 74 |
+
Compute nearest neighbor distances from verts2 to verts1 using KDTree.
|
| 75 |
+
|
| 76 |
+
Args:
|
| 77 |
+
verts1: Reference point cloud [N, 3]
|
| 78 |
+
verts2: Query point cloud [M, 3]
|
| 79 |
+
|
| 80 |
+
Returns:
|
| 81 |
+
Distance array [M,] - distance from each point in verts2 to nearest in verts1
|
| 82 |
+
"""
|
| 83 |
+
if len(verts1) == 0 or len(verts2) == 0:
|
| 84 |
+
return np.array([])
|
| 85 |
+
|
| 86 |
+
kdtree = KDTree(verts1)
|
| 87 |
+
distances, _ = kdtree.query(verts2)
|
| 88 |
+
return distances.reshape(-1)
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def evaluate_3d_reconstruction(
|
| 92 |
+
pcd_pred: Union[o3d.geometry.PointCloud, np.ndarray],
|
| 93 |
+
pcd_trgt: Union[o3d.geometry.PointCloud, np.ndarray],
|
| 94 |
+
threshold: float = 0.05,
|
| 95 |
+
down_sample: Optional[float] = None,
|
| 96 |
+
) -> TDict[str, float]:
|
| 97 |
+
"""
|
| 98 |
+
Evaluate 3D reconstruction quality using standard metrics.
|
| 99 |
+
|
| 100 |
+
This function computes:
|
| 101 |
+
- Accuracy: Mean distance from predicted points to GT surface
|
| 102 |
+
- Completeness: Mean distance from GT points to predicted surface
|
| 103 |
+
- Overall: Average of accuracy and completeness
|
| 104 |
+
- Precision: Fraction of predicted points within threshold of GT
|
| 105 |
+
- Recall: Fraction of GT points within threshold of prediction
|
| 106 |
+
- F-score: Harmonic mean of precision and recall
|
| 107 |
+
|
| 108 |
+
Args:
|
| 109 |
+
pcd_pred: Predicted point cloud (Open3D or numpy array)
|
| 110 |
+
pcd_trgt: Ground truth point cloud (Open3D or numpy array)
|
| 111 |
+
threshold: Distance threshold for precision/recall (meters)
|
| 112 |
+
down_sample: Voxel size for downsampling (None to skip)
|
| 113 |
+
|
| 114 |
+
Returns:
|
| 115 |
+
Dict with metrics: acc, comp, overall, precision, recall, fscore
|
| 116 |
+
"""
|
| 117 |
+
# Convert to Open3D if needed
|
| 118 |
+
if isinstance(pcd_pred, np.ndarray):
|
| 119 |
+
pcd_pred_o3d = o3d.geometry.PointCloud()
|
| 120 |
+
pcd_pred_o3d.points = o3d.utility.Vector3dVector(pcd_pred)
|
| 121 |
+
pcd_pred = pcd_pred_o3d
|
| 122 |
+
if isinstance(pcd_trgt, np.ndarray):
|
| 123 |
+
pcd_trgt_o3d = o3d.geometry.PointCloud()
|
| 124 |
+
pcd_trgt_o3d.points = o3d.utility.Vector3dVector(pcd_trgt)
|
| 125 |
+
pcd_trgt = pcd_trgt_o3d
|
| 126 |
+
|
| 127 |
+
# Downsample if requested
|
| 128 |
+
if down_sample is not None and down_sample > 0:
|
| 129 |
+
pcd_pred = pcd_pred.voxel_down_sample(down_sample)
|
| 130 |
+
pcd_trgt = pcd_trgt.voxel_down_sample(down_sample)
|
| 131 |
+
|
| 132 |
+
verts_pred = np.asarray(pcd_pred.points)
|
| 133 |
+
verts_trgt = np.asarray(pcd_trgt.points)
|
| 134 |
+
|
| 135 |
+
# Handle empty point clouds
|
| 136 |
+
if len(verts_pred) == 0 or len(verts_trgt) == 0:
|
| 137 |
+
return {
|
| 138 |
+
"acc": float("inf"),
|
| 139 |
+
"comp": float("inf"),
|
| 140 |
+
"overall": float("inf"),
|
| 141 |
+
"precision": 0.0,
|
| 142 |
+
"recall": 0.0,
|
| 143 |
+
"fscore": 0.0,
|
| 144 |
+
}
|
| 145 |
+
|
| 146 |
+
# Compute distances
|
| 147 |
+
dist_pred_to_gt = nn_correspondance(verts_trgt, verts_pred) # Accuracy
|
| 148 |
+
dist_gt_to_pred = nn_correspondance(verts_pred, verts_trgt) # Completeness
|
| 149 |
+
|
| 150 |
+
# Compute metrics
|
| 151 |
+
accuracy = float(np.mean(dist_pred_to_gt))
|
| 152 |
+
completeness = float(np.mean(dist_gt_to_pred))
|
| 153 |
+
overall = (accuracy + completeness) / 2
|
| 154 |
+
|
| 155 |
+
precision = float(np.mean((dist_pred_to_gt < threshold).astype(float)))
|
| 156 |
+
recall = float(np.mean((dist_gt_to_pred < threshold).astype(float)))
|
| 157 |
+
|
| 158 |
+
if precision + recall > 0:
|
| 159 |
+
fscore = 2 * precision * recall / (precision + recall)
|
| 160 |
+
else:
|
| 161 |
+
fscore = 0.0
|
| 162 |
+
|
| 163 |
+
return {
|
| 164 |
+
"acc": accuracy,
|
| 165 |
+
"comp": completeness,
|
| 166 |
+
"overall": overall,
|
| 167 |
+
"precision": precision,
|
| 168 |
+
"recall": recall,
|
| 169 |
+
"fscore": fscore,
|
| 170 |
+
}
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def create_tsdf_volume(
|
| 174 |
+
voxel_length: float = 4.0 / 512.0,
|
| 175 |
+
sdf_trunc: float = 0.04,
|
| 176 |
+
color_type: str = "RGB8",
|
| 177 |
+
) -> o3d.pipelines.integration.ScalableTSDFVolume:
|
| 178 |
+
"""
|
| 179 |
+
Create a scalable TSDF volume for depth fusion.
|
| 180 |
+
|
| 181 |
+
Args:
|
| 182 |
+
voxel_length: Size of each voxel
|
| 183 |
+
sdf_trunc: Truncation distance for SDF
|
| 184 |
+
color_type: Color integration type ("RGB8" or "Gray32")
|
| 185 |
+
|
| 186 |
+
Returns:
|
| 187 |
+
Initialized ScalableTSDFVolume
|
| 188 |
+
"""
|
| 189 |
+
if color_type == "RGB8":
|
| 190 |
+
color_enum = o3d.pipelines.integration.TSDFVolumeColorType.RGB8
|
| 191 |
+
else:
|
| 192 |
+
color_enum = o3d.pipelines.integration.TSDFVolumeColorType.Gray32
|
| 193 |
+
|
| 194 |
+
volume = o3d.pipelines.integration.ScalableTSDFVolume(
|
| 195 |
+
voxel_length=voxel_length,
|
| 196 |
+
sdf_trunc=sdf_trunc,
|
| 197 |
+
color_type=color_enum,
|
| 198 |
+
)
|
| 199 |
+
return volume
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def fuse_depth_to_tsdf(
|
| 203 |
+
volume: o3d.pipelines.integration.ScalableTSDFVolume,
|
| 204 |
+
depths: np.ndarray,
|
| 205 |
+
images: np.ndarray,
|
| 206 |
+
intrinsics: np.ndarray,
|
| 207 |
+
extrinsics: np.ndarray,
|
| 208 |
+
max_depth: float = 10.0,
|
| 209 |
+
) -> o3d.geometry.TriangleMesh:
|
| 210 |
+
"""
|
| 211 |
+
Fuse multiple depth maps into TSDF volume and extract mesh.
|
| 212 |
+
|
| 213 |
+
Args:
|
| 214 |
+
volume: TSDF volume to integrate into
|
| 215 |
+
depths: Depth maps [N, H, W]
|
| 216 |
+
images: RGB images [N, H, W, 3]
|
| 217 |
+
intrinsics: Camera intrinsics [N, 3, 3]
|
| 218 |
+
extrinsics: Camera extrinsics (world-to-camera) [N, 4, 4]
|
| 219 |
+
max_depth: Maximum depth for truncation
|
| 220 |
+
|
| 221 |
+
Returns:
|
| 222 |
+
Extracted triangle mesh
|
| 223 |
+
"""
|
| 224 |
+
for i in range(len(depths)):
|
| 225 |
+
depth = depths[i]
|
| 226 |
+
image = images[i]
|
| 227 |
+
ixt = intrinsics[i]
|
| 228 |
+
ext = extrinsics[i]
|
| 229 |
+
|
| 230 |
+
h, w = depth.shape[:2]
|
| 231 |
+
|
| 232 |
+
# Create RGBD image
|
| 233 |
+
depth_o3d = o3d.geometry.Image(depth.astype(np.float32))
|
| 234 |
+
color_o3d = o3d.geometry.Image(image.astype(np.uint8))
|
| 235 |
+
rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
|
| 236 |
+
color_o3d,
|
| 237 |
+
depth_o3d,
|
| 238 |
+
depth_trunc=max_depth,
|
| 239 |
+
convert_rgb_to_intensity=False,
|
| 240 |
+
depth_scale=1.0,
|
| 241 |
+
)
|
| 242 |
+
|
| 243 |
+
# Create camera intrinsics
|
| 244 |
+
ixt_o3d = o3d.camera.PinholeCameraIntrinsic(
|
| 245 |
+
w, h, ixt[0, 0], ixt[1, 1], ixt[0, 2], ixt[1, 2]
|
| 246 |
+
)
|
| 247 |
+
|
| 248 |
+
# Integrate into volume
|
| 249 |
+
volume.integrate(rgbd, ixt_o3d, ext)
|
| 250 |
+
|
| 251 |
+
# Extract mesh
|
| 252 |
+
mesh = volume.extract_triangle_mesh()
|
| 253 |
+
return mesh
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
def sample_points_from_mesh(
|
| 257 |
+
mesh: o3d.geometry.TriangleMesh,
|
| 258 |
+
num_points: int = 1000000,
|
| 259 |
+
) -> o3d.geometry.PointCloud:
|
| 260 |
+
"""
|
| 261 |
+
Uniformly sample points from a triangle mesh.
|
| 262 |
+
|
| 263 |
+
Args:
|
| 264 |
+
mesh: Input triangle mesh
|
| 265 |
+
num_points: Number of points to sample
|
| 266 |
+
|
| 267 |
+
Returns:
|
| 268 |
+
Sampled point cloud
|
| 269 |
+
"""
|
| 270 |
+
try:
|
| 271 |
+
pcd = mesh.sample_points_uniformly(number_of_points=num_points)
|
| 272 |
+
# Clamp colors to valid range [0, 1] for Open3D PLY export
|
| 273 |
+
if pcd.has_colors():
|
| 274 |
+
colors = np.asarray(pcd.colors)
|
| 275 |
+
colors = np.clip(colors, 0.0, 1.0)
|
| 276 |
+
pcd.colors = o3d.utility.Vector3dVector(colors)
|
| 277 |
+
except Exception:
|
| 278 |
+
# Fallback: create random points if mesh is invalid (with fixed seed for reproducibility)
|
| 279 |
+
rng = np.random.default_rng(seed=42)
|
| 280 |
+
points = rng.uniform(-1, 1, size=(num_points, 3))
|
| 281 |
+
pcd = o3d.geometry.PointCloud()
|
| 282 |
+
pcd.points = o3d.utility.Vector3dVector(points)
|
| 283 |
+
return pcd
|
| 284 |
+
|
| 285 |
+
|
| 286 |
+
# =============================================================================
|
| 287 |
+
# Pose Evaluation
|
| 288 |
+
# =============================================================================
|
| 289 |
+
|
| 290 |
+
|
| 291 |
+
def build_pair_index(N: int, B: int = 1):
|
| 292 |
+
"""
|
| 293 |
+
Build indices for all possible pairs of frames.
|
| 294 |
+
|
| 295 |
+
Args:
|
| 296 |
+
N: Number of frames
|
| 297 |
+
B: Batch size
|
| 298 |
+
|
| 299 |
+
Returns:
|
| 300 |
+
i1, i2: Indices for all possible pairs
|
| 301 |
+
"""
|
| 302 |
+
i1_, i2_ = torch.combinations(torch.arange(N), 2, with_replacement=False).unbind(-1)
|
| 303 |
+
i1, i2 = ((i[None] + torch.arange(B)[:, None] * N).reshape(-1) for i in [i1_, i2_])
|
| 304 |
+
return i1, i2
|
| 305 |
+
|
| 306 |
+
|
| 307 |
+
def compute_pose(pred_se3: torch.Tensor, gt_se3: torch.Tensor) -> Dict:
|
| 308 |
+
"""
|
| 309 |
+
Compute pose estimation metrics between predicted and ground truth trajectories.
|
| 310 |
+
|
| 311 |
+
Args:
|
| 312 |
+
pred_se3: Predicted SE(3) transformations [N, 4, 4]
|
| 313 |
+
gt_se3: Ground truth SE(3) transformations [N, 4, 4]
|
| 314 |
+
|
| 315 |
+
Returns:
|
| 316 |
+
Dict with AUC metrics at different thresholds (auc30, auc15, auc05, auc03)
|
| 317 |
+
"""
|
| 318 |
+
pred_se3 = align_to_first_camera(pred_se3)
|
| 319 |
+
gt_se3 = align_to_first_camera(gt_se3)
|
| 320 |
+
|
| 321 |
+
rel_rangle_deg, rel_tangle_deg = se3_to_relative_pose_error(pred_se3, gt_se3, len(pred_se3))
|
| 322 |
+
rError = rel_rangle_deg.cpu().numpy()
|
| 323 |
+
tError = rel_tangle_deg.cpu().numpy()
|
| 324 |
+
|
| 325 |
+
output = Dict()
|
| 326 |
+
output.auc30, _ = calculate_auc_np(rError, tError, max_threshold=30)
|
| 327 |
+
output.auc15, _ = calculate_auc_np(rError, tError, max_threshold=15)
|
| 328 |
+
output.auc05, _ = calculate_auc_np(rError, tError, max_threshold=5)
|
| 329 |
+
output.auc03, _ = calculate_auc_np(rError, tError, max_threshold=3)
|
| 330 |
+
return output
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
def align_to_first_camera(camera_poses: torch.Tensor) -> torch.Tensor:
|
| 334 |
+
"""
|
| 335 |
+
Align all camera poses to the first camera's coordinate frame.
|
| 336 |
+
|
| 337 |
+
Args:
|
| 338 |
+
camera_poses: Camera poses as SE3 transformations [N, 4, 4]
|
| 339 |
+
|
| 340 |
+
Returns:
|
| 341 |
+
Aligned camera poses [N, 4, 4]
|
| 342 |
+
"""
|
| 343 |
+
first_cam_extrinsic_inv = closed_form_inverse_se3(camera_poses[0][None])
|
| 344 |
+
aligned_poses = torch.matmul(camera_poses, first_cam_extrinsic_inv)
|
| 345 |
+
return aligned_poses
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
def rotation_angle(
|
| 349 |
+
rot_gt: torch.Tensor, rot_pred: torch.Tensor, batch_size: int = None, eps: float = 1e-15
|
| 350 |
+
) -> torch.Tensor:
|
| 351 |
+
"""
|
| 352 |
+
Calculate rotation angle error between ground truth and predicted rotations.
|
| 353 |
+
|
| 354 |
+
Args:
|
| 355 |
+
rot_gt: Ground truth rotation matrices
|
| 356 |
+
rot_pred: Predicted rotation matrices
|
| 357 |
+
batch_size: Batch size for reshaping the result
|
| 358 |
+
eps: Small value to avoid numerical issues
|
| 359 |
+
|
| 360 |
+
Returns:
|
| 361 |
+
Rotation angle error in degrees
|
| 362 |
+
"""
|
| 363 |
+
q_pred = mat_to_quat(rot_pred)
|
| 364 |
+
q_gt = mat_to_quat(rot_gt)
|
| 365 |
+
|
| 366 |
+
loss_q = (1 - (q_pred * q_gt).sum(dim=1) ** 2).clamp(min=eps)
|
| 367 |
+
err_q = torch.arccos(1 - 2 * loss_q)
|
| 368 |
+
|
| 369 |
+
rel_rangle_deg = err_q * 180 / np.pi
|
| 370 |
+
|
| 371 |
+
if batch_size is not None:
|
| 372 |
+
rel_rangle_deg = rel_rangle_deg.reshape(batch_size, -1)
|
| 373 |
+
|
| 374 |
+
return rel_rangle_deg
|
| 375 |
+
|
| 376 |
+
|
| 377 |
+
def translation_angle(
|
| 378 |
+
tvec_gt: torch.Tensor,
|
| 379 |
+
tvec_pred: torch.Tensor,
|
| 380 |
+
batch_size: int = None,
|
| 381 |
+
ambiguity: bool = True,
|
| 382 |
+
) -> torch.Tensor:
|
| 383 |
+
"""
|
| 384 |
+
Calculate translation angle error between ground truth and predicted translations.
|
| 385 |
+
|
| 386 |
+
Args:
|
| 387 |
+
tvec_gt: Ground truth translation vectors
|
| 388 |
+
tvec_pred: Predicted translation vectors
|
| 389 |
+
batch_size: Batch size for reshaping the result
|
| 390 |
+
ambiguity: Whether to handle direction ambiguity
|
| 391 |
+
|
| 392 |
+
Returns:
|
| 393 |
+
Translation angle error in degrees
|
| 394 |
+
"""
|
| 395 |
+
rel_tangle_deg = compare_translation_by_angle(tvec_gt, tvec_pred)
|
| 396 |
+
rel_tangle_deg = rel_tangle_deg * 180.0 / np.pi
|
| 397 |
+
|
| 398 |
+
if ambiguity:
|
| 399 |
+
rel_tangle_deg = torch.min(rel_tangle_deg, (180 - rel_tangle_deg).abs())
|
| 400 |
+
|
| 401 |
+
if batch_size is not None:
|
| 402 |
+
rel_tangle_deg = rel_tangle_deg.reshape(batch_size, -1)
|
| 403 |
+
|
| 404 |
+
return rel_tangle_deg
|
| 405 |
+
|
| 406 |
+
|
| 407 |
+
def compare_translation_by_angle(
|
| 408 |
+
t_gt: torch.Tensor, t: torch.Tensor, eps: float = 1e-15, default_err: float = 1e6
|
| 409 |
+
) -> torch.Tensor:
|
| 410 |
+
"""
|
| 411 |
+
Normalize the translation vectors and compute the angle between them.
|
| 412 |
+
|
| 413 |
+
Args:
|
| 414 |
+
t_gt: Ground truth translation vectors
|
| 415 |
+
t: Predicted translation vectors
|
| 416 |
+
eps: Small value to avoid division by zero
|
| 417 |
+
default_err: Default error value for invalid cases
|
| 418 |
+
|
| 419 |
+
Returns:
|
| 420 |
+
Angular error between translation vectors in radians
|
| 421 |
+
"""
|
| 422 |
+
t_norm = torch.norm(t, dim=1, keepdim=True)
|
| 423 |
+
t = t / (t_norm + eps)
|
| 424 |
+
|
| 425 |
+
t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True)
|
| 426 |
+
t_gt = t_gt / (t_gt_norm + eps)
|
| 427 |
+
|
| 428 |
+
loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps)
|
| 429 |
+
err_t = torch.acos(torch.sqrt(1 - loss_t))
|
| 430 |
+
|
| 431 |
+
err_t[torch.isnan(err_t) | torch.isinf(err_t)] = default_err
|
| 432 |
+
return err_t
|
| 433 |
+
|
| 434 |
+
|
| 435 |
+
def calculate_auc_np(
|
| 436 |
+
r_error: np.ndarray, t_error: np.ndarray, max_threshold: int = 30
|
| 437 |
+
) -> tuple:
|
| 438 |
+
"""
|
| 439 |
+
Calculate the Area Under the Curve (AUC) for the given error arrays.
|
| 440 |
+
|
| 441 |
+
Args:
|
| 442 |
+
r_error: Rotation error values in degrees
|
| 443 |
+
t_error: Translation error values in degrees
|
| 444 |
+
max_threshold: Maximum threshold value for binning
|
| 445 |
+
|
| 446 |
+
Returns:
|
| 447 |
+
Tuple of (AUC value, normalized histogram)
|
| 448 |
+
"""
|
| 449 |
+
error_matrix = np.concatenate((r_error[:, None], t_error[:, None]), axis=1)
|
| 450 |
+
max_errors = np.max(error_matrix, axis=1)
|
| 451 |
+
bins = np.arange(max_threshold + 1)
|
| 452 |
+
histogram, _ = np.histogram(max_errors, bins=bins)
|
| 453 |
+
num_pairs = float(len(max_errors))
|
| 454 |
+
normalized_histogram = histogram.astype(float) / num_pairs
|
| 455 |
+
return np.mean(np.cumsum(normalized_histogram)), normalized_histogram
|
| 456 |
+
|
| 457 |
+
|
| 458 |
+
def se3_to_relative_pose_error(
|
| 459 |
+
pred_se3: torch.Tensor, gt_se3: torch.Tensor, num_frames: int
|
| 460 |
+
) -> tuple:
|
| 461 |
+
"""
|
| 462 |
+
Compute rotation and translation errors between predicted and ground truth poses.
|
| 463 |
+
|
| 464 |
+
Args:
|
| 465 |
+
pred_se3: Predicted SE(3) transformations
|
| 466 |
+
gt_se3: Ground truth SE(3) transformations
|
| 467 |
+
num_frames: Number of frames
|
| 468 |
+
|
| 469 |
+
Returns:
|
| 470 |
+
Tuple of (rotation angle errors, translation angle errors) in degrees
|
| 471 |
+
"""
|
| 472 |
+
pair_idx_i1, pair_idx_i2 = build_pair_index(num_frames)
|
| 473 |
+
|
| 474 |
+
# Compute relative camera poses between pairs using closed-form inverse
|
| 475 |
+
relative_pose_gt = closed_form_inverse_se3(gt_se3[pair_idx_i1]).bmm(gt_se3[pair_idx_i2])
|
| 476 |
+
relative_pose_pred = closed_form_inverse_se3(pred_se3[pair_idx_i1]).bmm(pred_se3[pair_idx_i2])
|
| 477 |
+
|
| 478 |
+
# Compute the difference in rotation and translation
|
| 479 |
+
rel_rangle_deg = rotation_angle(relative_pose_gt[:, :3, :3], relative_pose_pred[:, :3, :3])
|
| 480 |
+
rel_tangle_deg = translation_angle(relative_pose_gt[:, :3, 3], relative_pose_pred[:, :3, 3])
|
| 481 |
+
|
| 482 |
+
return rel_rangle_deg, rel_tangle_deg
|
| 483 |
+
|
| 484 |
+
|
| 485 |
+
def closed_form_inverse_se3(
|
| 486 |
+
se3: torch.Tensor, R: torch.Tensor = None, T: torch.Tensor = None
|
| 487 |
+
) -> torch.Tensor:
|
| 488 |
+
"""
|
| 489 |
+
Compute the inverse of each 4x4 (or 3x4) SE3 matrix in a batch.
|
| 490 |
+
|
| 491 |
+
Uses closed-form solution instead of torch.inverse() for numerical stability.
|
| 492 |
+
|
| 493 |
+
Args:
|
| 494 |
+
se3: Nx4x4 or Nx3x4 tensor of SE3 matrices
|
| 495 |
+
R: Optional Nx3x3 rotation matrices
|
| 496 |
+
T: Optional Nx3x1 translation vectors
|
| 497 |
+
|
| 498 |
+
Returns:
|
| 499 |
+
Inverted SE3 matrices with same shape as input
|
| 500 |
+
"""
|
| 501 |
+
is_numpy = isinstance(se3, np.ndarray)
|
| 502 |
+
|
| 503 |
+
if se3.shape[-2:] != (4, 4) and se3.shape[-2:] != (3, 4):
|
| 504 |
+
raise ValueError(f"se3 must be of shape (N,4,4), got {se3.shape}.")
|
| 505 |
+
|
| 506 |
+
if R is None:
|
| 507 |
+
R = se3[:, :3, :3]
|
| 508 |
+
if T is None:
|
| 509 |
+
T = se3[:, :3, 3:]
|
| 510 |
+
|
| 511 |
+
if is_numpy:
|
| 512 |
+
R_transposed = np.transpose(R, (0, 2, 1))
|
| 513 |
+
top_right = -np.matmul(R_transposed, T)
|
| 514 |
+
inverted_matrix = np.tile(np.eye(4), (len(R), 1, 1))
|
| 515 |
+
else:
|
| 516 |
+
R_transposed = R.transpose(1, 2)
|
| 517 |
+
top_right = -torch.bmm(R_transposed, T)
|
| 518 |
+
inverted_matrix = torch.eye(4, 4)[None].repeat(len(R), 1, 1)
|
| 519 |
+
inverted_matrix = inverted_matrix.to(R.dtype).to(R.device)
|
| 520 |
+
|
| 521 |
+
inverted_matrix[:, :3, :3] = R_transposed
|
| 522 |
+
inverted_matrix[:, :3, 3:] = top_right
|
| 523 |
+
|
| 524 |
+
return inverted_matrix
|
| 525 |
+
|
depth_anything_3/cfg.py
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Configuration utility functions
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
import importlib
|
| 20 |
+
from pathlib import Path
|
| 21 |
+
from typing import Any, Callable, List, Union
|
| 22 |
+
from omegaconf import DictConfig, ListConfig, OmegaConf
|
| 23 |
+
|
| 24 |
+
try:
|
| 25 |
+
OmegaConf.register_new_resolver("eval", eval)
|
| 26 |
+
except Exception as e:
|
| 27 |
+
# if eval is not available, we can just pass
|
| 28 |
+
print(f"Error registering eval resolver: {e}")
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]:
|
| 32 |
+
"""
|
| 33 |
+
Load a configuration. Will resolve inheritance.
|
| 34 |
+
Supports both file paths and module paths (e.g., depth_anything_3.configs.giant).
|
| 35 |
+
"""
|
| 36 |
+
# Check if path is a module path (contains dots but no slashes and doesn't end with .yaml)
|
| 37 |
+
if "." in path and "/" not in path and not path.endswith(".yaml"):
|
| 38 |
+
# It's a module path, load from package resources
|
| 39 |
+
path_parts = path.split(".")[1:]
|
| 40 |
+
config_path = Path(__file__).resolve().parent
|
| 41 |
+
for part in path_parts:
|
| 42 |
+
config_path = config_path.joinpath(part)
|
| 43 |
+
config_path = config_path.with_suffix(".yaml")
|
| 44 |
+
config = OmegaConf.load(str(config_path))
|
| 45 |
+
else:
|
| 46 |
+
# It's a file path (absolute, relative, or with .yaml extension)
|
| 47 |
+
config = OmegaConf.load(path)
|
| 48 |
+
|
| 49 |
+
if argv is not None:
|
| 50 |
+
config_argv = OmegaConf.from_dotlist(argv)
|
| 51 |
+
config = OmegaConf.merge(config, config_argv)
|
| 52 |
+
config = resolve_recursive(config, resolve_inheritance)
|
| 53 |
+
return config
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def resolve_recursive(
|
| 57 |
+
config: Any,
|
| 58 |
+
resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]],
|
| 59 |
+
) -> Any:
|
| 60 |
+
config = resolver(config)
|
| 61 |
+
if isinstance(config, DictConfig):
|
| 62 |
+
for k in config.keys():
|
| 63 |
+
v = config.get(k)
|
| 64 |
+
if isinstance(v, (DictConfig, ListConfig)):
|
| 65 |
+
config[k] = resolve_recursive(v, resolver)
|
| 66 |
+
if isinstance(config, ListConfig):
|
| 67 |
+
for i in range(len(config)):
|
| 68 |
+
v = config.get(i)
|
| 69 |
+
if isinstance(v, (DictConfig, ListConfig)):
|
| 70 |
+
config[i] = resolve_recursive(v, resolver)
|
| 71 |
+
return config
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any:
|
| 75 |
+
"""
|
| 76 |
+
Recursively resolve inheritance if the config contains:
|
| 77 |
+
__inherit__: path/to/parent.yaml or a ListConfig of such paths.
|
| 78 |
+
"""
|
| 79 |
+
if isinstance(config, DictConfig):
|
| 80 |
+
inherit = config.pop("__inherit__", None)
|
| 81 |
+
|
| 82 |
+
if inherit:
|
| 83 |
+
inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit]
|
| 84 |
+
|
| 85 |
+
parent_config = None
|
| 86 |
+
for parent_path in inherit_list:
|
| 87 |
+
assert isinstance(parent_path, str)
|
| 88 |
+
parent_config = (
|
| 89 |
+
load_config(parent_path)
|
| 90 |
+
if parent_config is None
|
| 91 |
+
else OmegaConf.merge(parent_config, load_config(parent_path))
|
| 92 |
+
)
|
| 93 |
+
|
| 94 |
+
if len(config.keys()) > 0:
|
| 95 |
+
config = OmegaConf.merge(parent_config, config)
|
| 96 |
+
else:
|
| 97 |
+
config = parent_config
|
| 98 |
+
return config
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def import_item(path: str, name: str) -> Any:
|
| 102 |
+
"""
|
| 103 |
+
Import a python item. Example: import_item("path.to.file", "MyClass") -> MyClass
|
| 104 |
+
"""
|
| 105 |
+
return getattr(importlib.import_module(path), name)
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def create_object(config: DictConfig) -> Any:
|
| 109 |
+
"""
|
| 110 |
+
Create an object from config.
|
| 111 |
+
The config is expected to contains the following:
|
| 112 |
+
__object__:
|
| 113 |
+
path: path.to.module
|
| 114 |
+
name: MyClass
|
| 115 |
+
args: as_config | as_params (default to as_config)
|
| 116 |
+
"""
|
| 117 |
+
config = DictConfig(config)
|
| 118 |
+
item = import_item(
|
| 119 |
+
path=config.__object__.path,
|
| 120 |
+
name=config.__object__.name,
|
| 121 |
+
)
|
| 122 |
+
args = config.__object__.get("args", "as_config")
|
| 123 |
+
if args == "as_config":
|
| 124 |
+
return item(config)
|
| 125 |
+
if args == "as_params":
|
| 126 |
+
config = OmegaConf.to_object(config)
|
| 127 |
+
config.pop("__object__")
|
| 128 |
+
return item(**config)
|
| 129 |
+
raise NotImplementedError(f"Unknown args type: {args}")
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def create_dataset(path: str, *args, **kwargs) -> Any:
|
| 133 |
+
"""
|
| 134 |
+
Create a dataset. Requires the file to contain a "create_dataset" function.
|
| 135 |
+
"""
|
| 136 |
+
return import_item(path, "create_dataset")(*args, **kwargs)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def to_dict_recursive(config_obj):
|
| 140 |
+
if isinstance(config_obj, DictConfig):
|
| 141 |
+
return {k: to_dict_recursive(v) for k, v in config_obj.items()}
|
| 142 |
+
elif isinstance(config_obj, ListConfig):
|
| 143 |
+
return [to_dict_recursive(item) for item in config_obj]
|
| 144 |
+
return config_obj
|
depth_anything_3/cli.py
ADDED
|
@@ -0,0 +1,824 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# flake8: noqa: E402
|
| 2 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 3 |
+
#
|
| 4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 5 |
+
# you may not use this file except in compliance with the License.
|
| 6 |
+
# You may obtain a copy of the License at
|
| 7 |
+
#
|
| 8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 9 |
+
#
|
| 10 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 13 |
+
# See the License for the specific language governing permissions and
|
| 14 |
+
# limitations under the License.
|
| 15 |
+
"""
|
| 16 |
+
Refactored Depth Anything 3 CLI
|
| 17 |
+
Clean, modular command-line interface
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
from __future__ import annotations
|
| 21 |
+
|
| 22 |
+
import os
|
| 23 |
+
import typer
|
| 24 |
+
|
| 25 |
+
from depth_anything_3.services import start_server
|
| 26 |
+
from depth_anything_3.services.gallery import gallery as gallery_main
|
| 27 |
+
from depth_anything_3.services.inference_service import run_inference
|
| 28 |
+
from depth_anything_3.services.input_handlers import (
|
| 29 |
+
ColmapHandler,
|
| 30 |
+
ImageHandler,
|
| 31 |
+
ImagesHandler,
|
| 32 |
+
InputHandler,
|
| 33 |
+
VideoHandler,
|
| 34 |
+
parse_export_feat,
|
| 35 |
+
)
|
| 36 |
+
from depth_anything_3.utils.constants import (
|
| 37 |
+
DEFAULT_EXPORT_DIR,
|
| 38 |
+
DEFAULT_GALLERY_DIR,
|
| 39 |
+
DEFAULT_GRADIO_DIR,
|
| 40 |
+
DEFAULT_MODEL,
|
| 41 |
+
)
|
| 42 |
+
|
| 43 |
+
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
| 44 |
+
|
| 45 |
+
app = typer.Typer(help="Depth Anything 3 - Video depth estimation CLI", add_completion=False)
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
# ============================================================================
|
| 49 |
+
# Input type detection utilities
|
| 50 |
+
# ============================================================================
|
| 51 |
+
|
| 52 |
+
# Supported file extensions
|
| 53 |
+
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif"}
|
| 54 |
+
VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".webm", ".m4v"}
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def detect_input_type(input_path: str) -> str:
|
| 58 |
+
"""
|
| 59 |
+
Detect input type from path.
|
| 60 |
+
|
| 61 |
+
Returns:
|
| 62 |
+
- "image": Single image file
|
| 63 |
+
- "images": Directory containing images
|
| 64 |
+
- "video": Video file
|
| 65 |
+
- "colmap": COLMAP directory structure
|
| 66 |
+
- "unknown": Cannot determine type
|
| 67 |
+
"""
|
| 68 |
+
if not os.path.exists(input_path):
|
| 69 |
+
return "unknown"
|
| 70 |
+
|
| 71 |
+
# Check if it's a file
|
| 72 |
+
if os.path.isfile(input_path):
|
| 73 |
+
ext = os.path.splitext(input_path)[1].lower()
|
| 74 |
+
if ext in IMAGE_EXTENSIONS:
|
| 75 |
+
return "image"
|
| 76 |
+
elif ext in VIDEO_EXTENSIONS:
|
| 77 |
+
return "video"
|
| 78 |
+
return "unknown"
|
| 79 |
+
|
| 80 |
+
# Check if it's a directory
|
| 81 |
+
if os.path.isdir(input_path):
|
| 82 |
+
# Check for COLMAP structure
|
| 83 |
+
images_dir = os.path.join(input_path, "images")
|
| 84 |
+
sparse_dir = os.path.join(input_path, "sparse")
|
| 85 |
+
|
| 86 |
+
if os.path.isdir(images_dir) and os.path.isdir(sparse_dir):
|
| 87 |
+
return "colmap"
|
| 88 |
+
|
| 89 |
+
# Check if directory contains image files
|
| 90 |
+
for item in os.listdir(input_path):
|
| 91 |
+
item_path = os.path.join(input_path, item)
|
| 92 |
+
if os.path.isfile(item_path):
|
| 93 |
+
ext = os.path.splitext(item)[1].lower()
|
| 94 |
+
if ext in IMAGE_EXTENSIONS:
|
| 95 |
+
return "images"
|
| 96 |
+
|
| 97 |
+
return "unknown"
|
| 98 |
+
|
| 99 |
+
return "unknown"
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
# ============================================================================
|
| 103 |
+
# Common parameters and configuration
|
| 104 |
+
# ============================================================================
|
| 105 |
+
|
| 106 |
+
# ============================================================================
|
| 107 |
+
# Inference commands
|
| 108 |
+
# ============================================================================
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
@app.command()
|
| 112 |
+
def auto(
|
| 113 |
+
input_path: str = typer.Argument(
|
| 114 |
+
..., help="Path to input (image, directory, video, or COLMAP)"
|
| 115 |
+
),
|
| 116 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 117 |
+
export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"),
|
| 118 |
+
export_format: str = typer.Option("glb", help="Export format"),
|
| 119 |
+
device: str = typer.Option("cuda", help="Device to use"),
|
| 120 |
+
use_backend: bool = typer.Option(False, help="Use backend service for inference"),
|
| 121 |
+
backend_url: str = typer.Option(
|
| 122 |
+
"http://localhost:8008", help="Backend URL (default: http://localhost:8008)"
|
| 123 |
+
),
|
| 124 |
+
process_res: int = typer.Option(504, help="Processing resolution"),
|
| 125 |
+
process_res_method: str = typer.Option(
|
| 126 |
+
"upper_bound_resize", help="Processing resolution method"
|
| 127 |
+
),
|
| 128 |
+
export_feat: str = typer.Option(
|
| 129 |
+
"",
|
| 130 |
+
help="[FEAT_VIS]Export features from specified layers using comma-separated indices (e.g., '0,1,2').",
|
| 131 |
+
),
|
| 132 |
+
auto_cleanup: bool = typer.Option(
|
| 133 |
+
False, help="Automatically clean export directory if it exists (no prompt)"
|
| 134 |
+
),
|
| 135 |
+
# Video-specific options
|
| 136 |
+
fps: float = typer.Option(1.0, help="[Video] Sampling FPS for frame extraction"),
|
| 137 |
+
# COLMAP-specific options
|
| 138 |
+
sparse_subdir: str = typer.Option(
|
| 139 |
+
"", help="[COLMAP] Sparse reconstruction subdirectory (e.g., '0' for sparse/0/)"
|
| 140 |
+
),
|
| 141 |
+
align_to_input_ext_scale: bool = typer.Option(
|
| 142 |
+
True, help="[COLMAP] Align prediction to input extrinsics scale"
|
| 143 |
+
),
|
| 144 |
+
# Pose estimation options
|
| 145 |
+
use_ray_pose: bool = typer.Option(
|
| 146 |
+
False, help="Use ray-based pose estimation instead of camera decoder"
|
| 147 |
+
),
|
| 148 |
+
ref_view_strategy: str = typer.Option(
|
| 149 |
+
"saddle_balanced",
|
| 150 |
+
help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range",
|
| 151 |
+
),
|
| 152 |
+
# GLB export options
|
| 153 |
+
conf_thresh_percentile: float = typer.Option(
|
| 154 |
+
40.0, help="[GLB] Lower percentile for adaptive confidence threshold"
|
| 155 |
+
),
|
| 156 |
+
num_max_points: int = typer.Option(
|
| 157 |
+
1_000_000, help="[GLB] Maximum number of points in the point cloud"
|
| 158 |
+
),
|
| 159 |
+
show_cameras: bool = typer.Option(
|
| 160 |
+
True, help="[GLB] Show camera wireframes in the exported scene"
|
| 161 |
+
),
|
| 162 |
+
# Feat_vis export options
|
| 163 |
+
feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"),
|
| 164 |
+
):
|
| 165 |
+
"""
|
| 166 |
+
Automatically detect input type and run appropriate processing.
|
| 167 |
+
|
| 168 |
+
Supports:
|
| 169 |
+
- Single image file (.jpg, .png, etc.)
|
| 170 |
+
- Directory of images
|
| 171 |
+
- Video file (.mp4, .avi, etc.)
|
| 172 |
+
- COLMAP directory (with 'images' and 'sparse' subdirectories)
|
| 173 |
+
"""
|
| 174 |
+
# Detect input type
|
| 175 |
+
input_type = detect_input_type(input_path)
|
| 176 |
+
|
| 177 |
+
if input_type == "unknown":
|
| 178 |
+
typer.echo(f"❌ Error: Cannot determine input type for: {input_path}", err=True)
|
| 179 |
+
typer.echo("Supported inputs:", err=True)
|
| 180 |
+
typer.echo(" - Single image file (.jpg, .png, etc.)", err=True)
|
| 181 |
+
typer.echo(" - Directory containing images", err=True)
|
| 182 |
+
typer.echo(" - Video file (.mp4, .avi, etc.)", err=True)
|
| 183 |
+
typer.echo(" - COLMAP directory (with 'images/' and 'sparse/' subdirectories)", err=True)
|
| 184 |
+
raise typer.Exit(1)
|
| 185 |
+
|
| 186 |
+
# Display detected type
|
| 187 |
+
typer.echo(f"🔍 Detected input type: {input_type.upper()}")
|
| 188 |
+
typer.echo(f"📁 Input path: {input_path}")
|
| 189 |
+
typer.echo()
|
| 190 |
+
|
| 191 |
+
# Determine backend URL based on use_backend flag
|
| 192 |
+
final_backend_url = backend_url if use_backend else None
|
| 193 |
+
|
| 194 |
+
# Parse export_feat parameter
|
| 195 |
+
export_feat_layers = parse_export_feat(export_feat)
|
| 196 |
+
|
| 197 |
+
# Route to appropriate handler
|
| 198 |
+
if input_type == "image":
|
| 199 |
+
typer.echo("Processing single image...")
|
| 200 |
+
# Process input
|
| 201 |
+
image_files = ImageHandler.process(input_path)
|
| 202 |
+
|
| 203 |
+
# Handle export directory
|
| 204 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 205 |
+
|
| 206 |
+
# Run inference
|
| 207 |
+
run_inference(
|
| 208 |
+
image_paths=image_files,
|
| 209 |
+
export_dir=export_dir,
|
| 210 |
+
model_dir=model_dir,
|
| 211 |
+
device=device,
|
| 212 |
+
backend_url=final_backend_url,
|
| 213 |
+
export_format=export_format,
|
| 214 |
+
process_res=process_res,
|
| 215 |
+
process_res_method=process_res_method,
|
| 216 |
+
export_feat_layers=export_feat_layers,
|
| 217 |
+
use_ray_pose=use_ray_pose,
|
| 218 |
+
ref_view_strategy=ref_view_strategy,
|
| 219 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 220 |
+
num_max_points=num_max_points,
|
| 221 |
+
show_cameras=show_cameras,
|
| 222 |
+
feat_vis_fps=feat_vis_fps,
|
| 223 |
+
)
|
| 224 |
+
|
| 225 |
+
elif input_type == "images":
|
| 226 |
+
typer.echo("Processing directory of images...")
|
| 227 |
+
# Process input - use default extensions
|
| 228 |
+
image_files = ImagesHandler.process(input_path, "png,jpg,jpeg")
|
| 229 |
+
|
| 230 |
+
# Handle export directory
|
| 231 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 232 |
+
|
| 233 |
+
# Run inference
|
| 234 |
+
run_inference(
|
| 235 |
+
image_paths=image_files,
|
| 236 |
+
export_dir=export_dir,
|
| 237 |
+
model_dir=model_dir,
|
| 238 |
+
device=device,
|
| 239 |
+
backend_url=final_backend_url,
|
| 240 |
+
export_format=export_format,
|
| 241 |
+
process_res=process_res,
|
| 242 |
+
process_res_method=process_res_method,
|
| 243 |
+
export_feat_layers=export_feat_layers,
|
| 244 |
+
use_ray_pose=use_ray_pose,
|
| 245 |
+
ref_view_strategy=ref_view_strategy,
|
| 246 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 247 |
+
num_max_points=num_max_points,
|
| 248 |
+
show_cameras=show_cameras,
|
| 249 |
+
feat_vis_fps=feat_vis_fps,
|
| 250 |
+
)
|
| 251 |
+
|
| 252 |
+
elif input_type == "video":
|
| 253 |
+
typer.echo(f"Processing video with FPS={fps}...")
|
| 254 |
+
# Handle export directory
|
| 255 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 256 |
+
|
| 257 |
+
# Process input
|
| 258 |
+
image_files = VideoHandler.process(input_path, export_dir, fps)
|
| 259 |
+
|
| 260 |
+
# Run inference
|
| 261 |
+
run_inference(
|
| 262 |
+
image_paths=image_files,
|
| 263 |
+
export_dir=export_dir,
|
| 264 |
+
model_dir=model_dir,
|
| 265 |
+
device=device,
|
| 266 |
+
backend_url=final_backend_url,
|
| 267 |
+
export_format=export_format,
|
| 268 |
+
process_res=process_res,
|
| 269 |
+
process_res_method=process_res_method,
|
| 270 |
+
export_feat_layers=export_feat_layers,
|
| 271 |
+
use_ray_pose=use_ray_pose,
|
| 272 |
+
ref_view_strategy=ref_view_strategy,
|
| 273 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 274 |
+
num_max_points=num_max_points,
|
| 275 |
+
show_cameras=show_cameras,
|
| 276 |
+
feat_vis_fps=feat_vis_fps,
|
| 277 |
+
)
|
| 278 |
+
|
| 279 |
+
elif input_type == "colmap":
|
| 280 |
+
typer.echo(
|
| 281 |
+
f"Processing COLMAP directory (sparse subdirectory: '{sparse_subdir or 'default'}')..."
|
| 282 |
+
)
|
| 283 |
+
# Process input
|
| 284 |
+
image_files, extrinsics, intrinsics = ColmapHandler.process(input_path, sparse_subdir)
|
| 285 |
+
|
| 286 |
+
# Handle export directory
|
| 287 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 288 |
+
|
| 289 |
+
# Run inference
|
| 290 |
+
run_inference(
|
| 291 |
+
image_paths=image_files,
|
| 292 |
+
export_dir=export_dir,
|
| 293 |
+
model_dir=model_dir,
|
| 294 |
+
device=device,
|
| 295 |
+
backend_url=final_backend_url,
|
| 296 |
+
export_format=export_format,
|
| 297 |
+
process_res=process_res,
|
| 298 |
+
process_res_method=process_res_method,
|
| 299 |
+
export_feat_layers=export_feat_layers,
|
| 300 |
+
extrinsics=extrinsics,
|
| 301 |
+
intrinsics=intrinsics,
|
| 302 |
+
align_to_input_ext_scale=align_to_input_ext_scale,
|
| 303 |
+
use_ray_pose=use_ray_pose,
|
| 304 |
+
ref_view_strategy=ref_view_strategy,
|
| 305 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 306 |
+
num_max_points=num_max_points,
|
| 307 |
+
show_cameras=show_cameras,
|
| 308 |
+
feat_vis_fps=feat_vis_fps,
|
| 309 |
+
)
|
| 310 |
+
|
| 311 |
+
typer.echo()
|
| 312 |
+
typer.echo("✅ Processing completed successfully!")
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
@app.command()
|
| 316 |
+
def image(
|
| 317 |
+
image_path: str = typer.Argument(..., help="Path to input image file"),
|
| 318 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 319 |
+
export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"),
|
| 320 |
+
export_format: str = typer.Option("glb", help="Export format"),
|
| 321 |
+
device: str = typer.Option("cuda", help="Device to use"),
|
| 322 |
+
use_backend: bool = typer.Option(False, help="Use backend service for inference"),
|
| 323 |
+
backend_url: str = typer.Option(
|
| 324 |
+
"http://localhost:8008", help="Backend URL (default: http://localhost:8008)"
|
| 325 |
+
),
|
| 326 |
+
process_res: int = typer.Option(504, help="Processing resolution"),
|
| 327 |
+
process_res_method: str = typer.Option(
|
| 328 |
+
"upper_bound_resize", help="Processing resolution method"
|
| 329 |
+
),
|
| 330 |
+
export_feat: str = typer.Option(
|
| 331 |
+
"",
|
| 332 |
+
help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').",
|
| 333 |
+
),
|
| 334 |
+
auto_cleanup: bool = typer.Option(
|
| 335 |
+
False, help="Automatically clean export directory if it exists (no prompt)"
|
| 336 |
+
),
|
| 337 |
+
# Pose estimation options
|
| 338 |
+
use_ray_pose: bool = typer.Option(
|
| 339 |
+
False, help="Use ray-based pose estimation instead of camera decoder"
|
| 340 |
+
),
|
| 341 |
+
ref_view_strategy: str = typer.Option(
|
| 342 |
+
"saddle_balanced",
|
| 343 |
+
help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range",
|
| 344 |
+
),
|
| 345 |
+
# GLB export options
|
| 346 |
+
conf_thresh_percentile: float = typer.Option(
|
| 347 |
+
40.0, help="[GLB] Lower percentile for adaptive confidence threshold"
|
| 348 |
+
),
|
| 349 |
+
num_max_points: int = typer.Option(
|
| 350 |
+
1_000_000, help="[GLB] Maximum number of points in the point cloud"
|
| 351 |
+
),
|
| 352 |
+
show_cameras: bool = typer.Option(
|
| 353 |
+
True, help="[GLB] Show camera wireframes in the exported scene"
|
| 354 |
+
),
|
| 355 |
+
# Feat_vis export options
|
| 356 |
+
feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"),
|
| 357 |
+
):
|
| 358 |
+
"""Run camera pose and depth estimation on a single image."""
|
| 359 |
+
# Process input
|
| 360 |
+
image_files = ImageHandler.process(image_path)
|
| 361 |
+
|
| 362 |
+
# Handle export directory
|
| 363 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 364 |
+
|
| 365 |
+
# Parse export_feat parameter
|
| 366 |
+
export_feat_layers = parse_export_feat(export_feat)
|
| 367 |
+
|
| 368 |
+
# Determine backend URL based on use_backend flag
|
| 369 |
+
final_backend_url = backend_url if use_backend else None
|
| 370 |
+
|
| 371 |
+
# Run inference
|
| 372 |
+
run_inference(
|
| 373 |
+
image_paths=image_files,
|
| 374 |
+
export_dir=export_dir,
|
| 375 |
+
model_dir=model_dir,
|
| 376 |
+
device=device,
|
| 377 |
+
backend_url=final_backend_url,
|
| 378 |
+
export_format=export_format,
|
| 379 |
+
process_res=process_res,
|
| 380 |
+
process_res_method=process_res_method,
|
| 381 |
+
export_feat_layers=export_feat_layers,
|
| 382 |
+
use_ray_pose=use_ray_pose,
|
| 383 |
+
reference_view_strategy=reference_view_strategy,
|
| 384 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 385 |
+
num_max_points=num_max_points,
|
| 386 |
+
show_cameras=show_cameras,
|
| 387 |
+
feat_vis_fps=feat_vis_fps,
|
| 388 |
+
)
|
| 389 |
+
|
| 390 |
+
|
| 391 |
+
@app.command()
|
| 392 |
+
def images(
|
| 393 |
+
images_dir: str = typer.Argument(..., help="Path to directory containing input images"),
|
| 394 |
+
image_extensions: str = typer.Option(
|
| 395 |
+
"png,jpg,jpeg", help="Comma-separated image file extensions to process"
|
| 396 |
+
),
|
| 397 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 398 |
+
export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"),
|
| 399 |
+
export_format: str = typer.Option("glb", help="Export format"),
|
| 400 |
+
device: str = typer.Option("cuda", help="Device to use"),
|
| 401 |
+
use_backend: bool = typer.Option(False, help="Use backend service for inference"),
|
| 402 |
+
backend_url: str = typer.Option(
|
| 403 |
+
"http://localhost:8008", help="Backend URL (default: http://localhost:8008)"
|
| 404 |
+
),
|
| 405 |
+
process_res: int = typer.Option(504, help="Processing resolution"),
|
| 406 |
+
process_res_method: str = typer.Option(
|
| 407 |
+
"upper_bound_resize", help="Processing resolution method"
|
| 408 |
+
),
|
| 409 |
+
export_feat: str = typer.Option(
|
| 410 |
+
"",
|
| 411 |
+
help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').",
|
| 412 |
+
),
|
| 413 |
+
auto_cleanup: bool = typer.Option(
|
| 414 |
+
False, help="Automatically clean export directory if it exists (no prompt)"
|
| 415 |
+
),
|
| 416 |
+
# Pose estimation options
|
| 417 |
+
use_ray_pose: bool = typer.Option(
|
| 418 |
+
False, help="Use ray-based pose estimation instead of camera decoder"
|
| 419 |
+
),
|
| 420 |
+
ref_view_strategy: str = typer.Option(
|
| 421 |
+
"saddle_balanced",
|
| 422 |
+
help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range",
|
| 423 |
+
),
|
| 424 |
+
# GLB export options
|
| 425 |
+
conf_thresh_percentile: float = typer.Option(
|
| 426 |
+
40.0, help="[GLB] Lower percentile for adaptive confidence threshold"
|
| 427 |
+
),
|
| 428 |
+
num_max_points: int = typer.Option(
|
| 429 |
+
1_000_000, help="[GLB] Maximum number of points in the point cloud"
|
| 430 |
+
),
|
| 431 |
+
show_cameras: bool = typer.Option(
|
| 432 |
+
True, help="[GLB] Show camera wireframes in the exported scene"
|
| 433 |
+
),
|
| 434 |
+
# Feat_vis export options
|
| 435 |
+
feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"),
|
| 436 |
+
):
|
| 437 |
+
"""Run camera pose and depth estimation on a directory of images."""
|
| 438 |
+
# Process input
|
| 439 |
+
image_files = ImagesHandler.process(images_dir, image_extensions)
|
| 440 |
+
|
| 441 |
+
# Handle export directory
|
| 442 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 443 |
+
|
| 444 |
+
# Parse export_feat parameter
|
| 445 |
+
export_feat_layers = parse_export_feat(export_feat)
|
| 446 |
+
|
| 447 |
+
# Determine backend URL based on use_backend flag
|
| 448 |
+
final_backend_url = backend_url if use_backend else None
|
| 449 |
+
|
| 450 |
+
# Run inference
|
| 451 |
+
run_inference(
|
| 452 |
+
image_paths=image_files,
|
| 453 |
+
export_dir=export_dir,
|
| 454 |
+
model_dir=model_dir,
|
| 455 |
+
device=device,
|
| 456 |
+
backend_url=final_backend_url,
|
| 457 |
+
export_format=export_format,
|
| 458 |
+
process_res=process_res,
|
| 459 |
+
process_res_method=process_res_method,
|
| 460 |
+
export_feat_layers=export_feat_layers,
|
| 461 |
+
use_ray_pose=use_ray_pose,
|
| 462 |
+
reference_view_strategy=reference_view_strategy,
|
| 463 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 464 |
+
num_max_points=num_max_points,
|
| 465 |
+
show_cameras=show_cameras,
|
| 466 |
+
feat_vis_fps=feat_vis_fps,
|
| 467 |
+
)
|
| 468 |
+
|
| 469 |
+
|
| 470 |
+
@app.command()
|
| 471 |
+
def colmap(
|
| 472 |
+
colmap_dir: str = typer.Argument(
|
| 473 |
+
..., help="Path to COLMAP directory containing 'images' and 'sparse' subdirectories"
|
| 474 |
+
),
|
| 475 |
+
sparse_subdir: str = typer.Option(
|
| 476 |
+
"", help="Sparse reconstruction subdirectory (e.g., '0' for sparse/0/, empty for sparse/)"
|
| 477 |
+
),
|
| 478 |
+
align_to_input_ext_scale: bool = typer.Option(
|
| 479 |
+
True, help="Align prediction to input extrinsics scale"
|
| 480 |
+
),
|
| 481 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 482 |
+
export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"),
|
| 483 |
+
export_format: str = typer.Option("glb", help="Export format"),
|
| 484 |
+
device: str = typer.Option("cuda", help="Device to use"),
|
| 485 |
+
use_backend: bool = typer.Option(False, help="Use backend service for inference"),
|
| 486 |
+
backend_url: str = typer.Option(
|
| 487 |
+
"http://localhost:8008", help="Backend URL (default: http://localhost:8008)"
|
| 488 |
+
),
|
| 489 |
+
process_res: int = typer.Option(504, help="Processing resolution"),
|
| 490 |
+
process_res_method: str = typer.Option(
|
| 491 |
+
"upper_bound_resize", help="Processing resolution method"
|
| 492 |
+
),
|
| 493 |
+
export_feat: str = typer.Option(
|
| 494 |
+
"",
|
| 495 |
+
help="Export features from specified layers using comma-separated indices (e.g., '0,1,2').",
|
| 496 |
+
),
|
| 497 |
+
auto_cleanup: bool = typer.Option(
|
| 498 |
+
False, help="Automatically clean export directory if it exists (no prompt)"
|
| 499 |
+
),
|
| 500 |
+
# Pose estimation options
|
| 501 |
+
use_ray_pose: bool = typer.Option(
|
| 502 |
+
False, help="Use ray-based pose estimation instead of camera decoder"
|
| 503 |
+
),
|
| 504 |
+
ref_view_strategy: str = typer.Option(
|
| 505 |
+
"saddle_balanced",
|
| 506 |
+
help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range",
|
| 507 |
+
),
|
| 508 |
+
# GLB export options
|
| 509 |
+
conf_thresh_percentile: float = typer.Option(
|
| 510 |
+
40.0, help="[GLB] Lower percentile for adaptive confidence threshold"
|
| 511 |
+
),
|
| 512 |
+
num_max_points: int = typer.Option(
|
| 513 |
+
1_000_000, help="[GLB] Maximum number of points in the point cloud"
|
| 514 |
+
),
|
| 515 |
+
show_cameras: bool = typer.Option(
|
| 516 |
+
True, help="[GLB] Show camera wireframes in the exported scene"
|
| 517 |
+
),
|
| 518 |
+
# Feat_vis export options
|
| 519 |
+
feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"),
|
| 520 |
+
):
|
| 521 |
+
"""Run pose conditioned depth estimation on COLMAP data."""
|
| 522 |
+
# Process input
|
| 523 |
+
image_files, extrinsics, intrinsics = ColmapHandler.process(colmap_dir, sparse_subdir)
|
| 524 |
+
|
| 525 |
+
# Handle export directory
|
| 526 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 527 |
+
|
| 528 |
+
# Parse export_feat parameter
|
| 529 |
+
export_feat_layers = parse_export_feat(export_feat)
|
| 530 |
+
|
| 531 |
+
# Determine backend URL based on use_backend flag
|
| 532 |
+
final_backend_url = backend_url if use_backend else None
|
| 533 |
+
|
| 534 |
+
# Run inference
|
| 535 |
+
run_inference(
|
| 536 |
+
image_paths=image_files,
|
| 537 |
+
export_dir=export_dir,
|
| 538 |
+
model_dir=model_dir,
|
| 539 |
+
device=device,
|
| 540 |
+
backend_url=final_backend_url,
|
| 541 |
+
export_format=export_format,
|
| 542 |
+
process_res=process_res,
|
| 543 |
+
process_res_method=process_res_method,
|
| 544 |
+
export_feat_layers=export_feat_layers,
|
| 545 |
+
extrinsics=extrinsics,
|
| 546 |
+
intrinsics=intrinsics,
|
| 547 |
+
align_to_input_ext_scale=align_to_input_ext_scale,
|
| 548 |
+
use_ray_pose=use_ray_pose,
|
| 549 |
+
reference_view_strategy=reference_view_strategy,
|
| 550 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 551 |
+
num_max_points=num_max_points,
|
| 552 |
+
show_cameras=show_cameras,
|
| 553 |
+
feat_vis_fps=feat_vis_fps,
|
| 554 |
+
)
|
| 555 |
+
|
| 556 |
+
|
| 557 |
+
@app.command()
|
| 558 |
+
def video(
|
| 559 |
+
video_path: str = typer.Argument(..., help="Path to input video file"),
|
| 560 |
+
fps: float = typer.Option(1.0, help="Sampling FPS for frame extraction"),
|
| 561 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 562 |
+
export_dir: str = typer.Option(DEFAULT_EXPORT_DIR, help="Export directory"),
|
| 563 |
+
export_format: str = typer.Option("glb", help="Export format"),
|
| 564 |
+
device: str = typer.Option("cuda", help="Device to use"),
|
| 565 |
+
use_backend: bool = typer.Option(False, help="Use backend service for inference"),
|
| 566 |
+
backend_url: str = typer.Option(
|
| 567 |
+
"http://localhost:8008", help="Backend URL (default: http://localhost:8008)"
|
| 568 |
+
),
|
| 569 |
+
process_res: int = typer.Option(504, help="Processing resolution"),
|
| 570 |
+
process_res_method: str = typer.Option(
|
| 571 |
+
"upper_bound_resize", help="Processing resolution method"
|
| 572 |
+
),
|
| 573 |
+
export_feat: str = typer.Option(
|
| 574 |
+
"",
|
| 575 |
+
help="[FEAT_VIS] Export features from specified layers using comma-separated indices (e.g., '0,1,2').",
|
| 576 |
+
),
|
| 577 |
+
auto_cleanup: bool = typer.Option(
|
| 578 |
+
False, help="Automatically clean export directory if it exists (no prompt)"
|
| 579 |
+
),
|
| 580 |
+
# Pose estimation options
|
| 581 |
+
use_ray_pose: bool = typer.Option(
|
| 582 |
+
False, help="Use ray-based pose estimation instead of camera decoder"
|
| 583 |
+
),
|
| 584 |
+
ref_view_strategy: str = typer.Option(
|
| 585 |
+
"saddle_balanced",
|
| 586 |
+
help="Reference view selection strategy: empty, first, middle, saddle_balanced, saddle_sim_range",
|
| 587 |
+
),
|
| 588 |
+
# GLB export options
|
| 589 |
+
conf_thresh_percentile: float = typer.Option(
|
| 590 |
+
40.0, help="[GLB] Lower percentile for adaptive confidence threshold"
|
| 591 |
+
),
|
| 592 |
+
num_max_points: int = typer.Option(
|
| 593 |
+
1_000_000, help="[GLB] Maximum number of points in the point cloud"
|
| 594 |
+
),
|
| 595 |
+
show_cameras: bool = typer.Option(
|
| 596 |
+
True, help="[GLB] Show camera wireframes in the exported scene"
|
| 597 |
+
),
|
| 598 |
+
# Feat_vis export options
|
| 599 |
+
feat_vis_fps: int = typer.Option(15, help="[FEAT_VIS] Frame rate for output video"),
|
| 600 |
+
):
|
| 601 |
+
"""Run depth estimation on video by extracting frames and processing them."""
|
| 602 |
+
# Handle export directory
|
| 603 |
+
export_dir = InputHandler.handle_export_dir(export_dir, auto_cleanup)
|
| 604 |
+
|
| 605 |
+
# Process input
|
| 606 |
+
image_files = VideoHandler.process(video_path, export_dir, fps)
|
| 607 |
+
|
| 608 |
+
# Parse export_feat parameter
|
| 609 |
+
export_feat_layers = parse_export_feat(export_feat)
|
| 610 |
+
|
| 611 |
+
# Determine backend URL based on use_backend flag
|
| 612 |
+
final_backend_url = backend_url if use_backend else None
|
| 613 |
+
|
| 614 |
+
# Run inference
|
| 615 |
+
run_inference(
|
| 616 |
+
image_paths=image_files,
|
| 617 |
+
export_dir=export_dir,
|
| 618 |
+
model_dir=model_dir,
|
| 619 |
+
device=device,
|
| 620 |
+
backend_url=final_backend_url,
|
| 621 |
+
export_format=export_format,
|
| 622 |
+
process_res=process_res,
|
| 623 |
+
process_res_method=process_res_method,
|
| 624 |
+
export_feat_layers=export_feat_layers,
|
| 625 |
+
use_ray_pose=use_ray_pose,
|
| 626 |
+
reference_view_strategy=reference_view_strategy,
|
| 627 |
+
conf_thresh_percentile=conf_thresh_percentile,
|
| 628 |
+
num_max_points=num_max_points,
|
| 629 |
+
show_cameras=show_cameras,
|
| 630 |
+
feat_vis_fps=feat_vis_fps,
|
| 631 |
+
)
|
| 632 |
+
|
| 633 |
+
|
| 634 |
+
# ============================================================================
|
| 635 |
+
# Service management commands
|
| 636 |
+
# ============================================================================
|
| 637 |
+
|
| 638 |
+
|
| 639 |
+
@app.command()
|
| 640 |
+
def backend(
|
| 641 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 642 |
+
device: str = typer.Option("cuda", help="Device to use"),
|
| 643 |
+
host: str = typer.Option("127.0.0.1", help="Host to bind to"),
|
| 644 |
+
port: int = typer.Option(8008, help="Port to bind to"),
|
| 645 |
+
gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery directory path (optional)"),
|
| 646 |
+
api_key: str = typer.Option(
|
| 647 |
+
None,
|
| 648 |
+
help="Require this API key (X-API-Key header) on inference requests. Falls back "
|
| 649 |
+
"to the DA3_BACKEND_API_KEY env var, or an auto-generated key when binding to a "
|
| 650 |
+
"non-loopback host.",
|
| 651 |
+
),
|
| 652 |
+
allow_unauthenticated: bool = typer.Option(
|
| 653 |
+
False,
|
| 654 |
+
help="Skip authentication entirely on a non-loopback host, instead of using an "
|
| 655 |
+
"auto-generated API key.",
|
| 656 |
+
),
|
| 657 |
+
):
|
| 658 |
+
"""Start model backend service with integrated gallery."""
|
| 659 |
+
typer.echo("=" * 60)
|
| 660 |
+
typer.echo("🚀 Starting Depth Anything 3 Backend Server")
|
| 661 |
+
typer.echo("=" * 60)
|
| 662 |
+
typer.echo(f"Model directory: {model_dir}")
|
| 663 |
+
typer.echo(f"Device: {device}")
|
| 664 |
+
|
| 665 |
+
# The gallery directory is also where /inference writes exports, so make sure it
|
| 666 |
+
# exists up front rather than treating a fresh checkout as "no gallery configured".
|
| 667 |
+
if gallery_dir:
|
| 668 |
+
os.makedirs(gallery_dir, exist_ok=True)
|
| 669 |
+
typer.echo(f"Gallery directory: {gallery_dir}")
|
| 670 |
+
else:
|
| 671 |
+
gallery_dir = None
|
| 672 |
+
|
| 673 |
+
typer.echo()
|
| 674 |
+
typer.echo("📡 Server URLs (Ctrl/CMD+Click to open):")
|
| 675 |
+
typer.echo(f" 🏠 Home: http://{host}:{port}")
|
| 676 |
+
typer.echo(f" 📊 Dashboard: http://{host}:{port}/dashboard")
|
| 677 |
+
typer.echo(f" 📈 API Status: http://{host}:{port}/status")
|
| 678 |
+
|
| 679 |
+
if gallery_dir:
|
| 680 |
+
typer.echo(f" 🎨 Gallery: http://{host}:{port}/gallery/")
|
| 681 |
+
|
| 682 |
+
typer.echo("=" * 60)
|
| 683 |
+
|
| 684 |
+
try:
|
| 685 |
+
start_server(
|
| 686 |
+
model_dir,
|
| 687 |
+
device,
|
| 688 |
+
host,
|
| 689 |
+
port,
|
| 690 |
+
gallery_dir,
|
| 691 |
+
api_key=api_key,
|
| 692 |
+
allow_unauthenticated=allow_unauthenticated,
|
| 693 |
+
)
|
| 694 |
+
except KeyboardInterrupt:
|
| 695 |
+
typer.echo("\n👋 Backend server stopped.")
|
| 696 |
+
except Exception as e:
|
| 697 |
+
typer.echo(f"❌ Failed to start backend: {e}")
|
| 698 |
+
raise typer.Exit(1)
|
| 699 |
+
|
| 700 |
+
|
| 701 |
+
# ============================================================================
|
| 702 |
+
# Application launch commands
|
| 703 |
+
# ============================================================================
|
| 704 |
+
|
| 705 |
+
|
| 706 |
+
@app.command()
|
| 707 |
+
def gradio(
|
| 708 |
+
model_dir: str = typer.Option(DEFAULT_MODEL, help="Model directory path"),
|
| 709 |
+
workspace_dir: str = typer.Option(DEFAULT_GRADIO_DIR, help="Workspace directory path"),
|
| 710 |
+
gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery directory path"),
|
| 711 |
+
host: str = typer.Option("127.0.0.1", help="Host address to bind to"),
|
| 712 |
+
port: int = typer.Option(7860, help="Port number to bind to"),
|
| 713 |
+
share: bool = typer.Option(False, help="Create a public link for the app"),
|
| 714 |
+
debug: bool = typer.Option(False, help="Enable debug mode"),
|
| 715 |
+
cache_examples: bool = typer.Option(
|
| 716 |
+
False, help="Pre-cache all example scenes at startup for faster loading"
|
| 717 |
+
),
|
| 718 |
+
cache_gs_tag: str = typer.Option(
|
| 719 |
+
"",
|
| 720 |
+
help="Tag to match scene names for high-res+3DGS caching (e.g., 'dl3dv'). Scenes containing this tag will use high_res and infer_gs=True; others will use low_res only.",
|
| 721 |
+
),
|
| 722 |
+
):
|
| 723 |
+
"""Launch Depth Anything 3 Gradio interactive web application"""
|
| 724 |
+
from depth_anything_3.app.gradio_app import DepthAnything3App
|
| 725 |
+
|
| 726 |
+
# Create necessary directories
|
| 727 |
+
os.makedirs(workspace_dir, exist_ok=True)
|
| 728 |
+
os.makedirs(gallery_dir, exist_ok=True)
|
| 729 |
+
|
| 730 |
+
typer.echo("Launching Depth Anything 3 Gradio application...")
|
| 731 |
+
typer.echo(f"Model directory: {model_dir}")
|
| 732 |
+
typer.echo(f"Workspace directory: {workspace_dir}")
|
| 733 |
+
typer.echo(f"Gallery directory: {gallery_dir}")
|
| 734 |
+
typer.echo(f"Host: {host}")
|
| 735 |
+
typer.echo(f"Port: {port}")
|
| 736 |
+
typer.echo(f"Share: {share}")
|
| 737 |
+
typer.echo(f"Debug mode: {debug}")
|
| 738 |
+
typer.echo(f"Cache examples: {cache_examples}")
|
| 739 |
+
if cache_examples:
|
| 740 |
+
if cache_gs_tag:
|
| 741 |
+
typer.echo(
|
| 742 |
+
f"Cache GS Tag: '{cache_gs_tag}' (scenes matching this tag will use high-res + 3DGS)"
|
| 743 |
+
)
|
| 744 |
+
else:
|
| 745 |
+
typer.echo(f"Cache GS Tag: None (all scenes will use low-res only)")
|
| 746 |
+
|
| 747 |
+
try:
|
| 748 |
+
# Initialize and launch application
|
| 749 |
+
app = DepthAnything3App(
|
| 750 |
+
model_dir=model_dir, workspace_dir=workspace_dir, gallery_dir=gallery_dir
|
| 751 |
+
)
|
| 752 |
+
|
| 753 |
+
# Pre-cache examples if requested
|
| 754 |
+
if cache_examples:
|
| 755 |
+
typer.echo("\n" + "=" * 60)
|
| 756 |
+
typer.echo("Pre-caching mode enabled")
|
| 757 |
+
if cache_gs_tag:
|
| 758 |
+
typer.echo(f"Scenes containing '{cache_gs_tag}' will use HIGH-RES + 3DGS")
|
| 759 |
+
typer.echo(f"Other scenes will use LOW-RES only")
|
| 760 |
+
else:
|
| 761 |
+
typer.echo(f"All scenes will use LOW-RES only")
|
| 762 |
+
typer.echo("=" * 60)
|
| 763 |
+
app.cache_examples(
|
| 764 |
+
show_cam=True,
|
| 765 |
+
filter_black_bg=False,
|
| 766 |
+
filter_white_bg=False,
|
| 767 |
+
save_percentage=20.0,
|
| 768 |
+
num_max_points=1000,
|
| 769 |
+
cache_gs_tag=cache_gs_tag,
|
| 770 |
+
gs_trj_mode="smooth",
|
| 771 |
+
gs_video_quality="low",
|
| 772 |
+
)
|
| 773 |
+
|
| 774 |
+
# Prepare launch arguments
|
| 775 |
+
launch_kwargs = {"share": share, "debug": debug}
|
| 776 |
+
|
| 777 |
+
app.launch(host=host, port=port, **launch_kwargs)
|
| 778 |
+
|
| 779 |
+
except KeyboardInterrupt:
|
| 780 |
+
typer.echo("\nGradio application stopped.")
|
| 781 |
+
except Exception as e:
|
| 782 |
+
typer.echo(f"Failed to launch Gradio application: {e}")
|
| 783 |
+
raise typer.Exit(1)
|
| 784 |
+
|
| 785 |
+
|
| 786 |
+
@app.command()
|
| 787 |
+
def gallery(
|
| 788 |
+
gallery_dir: str = typer.Option(DEFAULT_GALLERY_DIR, help="Gallery root directory"),
|
| 789 |
+
host: str = typer.Option("127.0.0.1", help="Host address to bind to"),
|
| 790 |
+
port: int = typer.Option(8007, help="Port number to bind to"),
|
| 791 |
+
open_browser: bool = typer.Option(False, help="Open browser after launch"),
|
| 792 |
+
):
|
| 793 |
+
"""Launch Depth Anything 3 Gallery server"""
|
| 794 |
+
|
| 795 |
+
# Validate gallery directory
|
| 796 |
+
if not os.path.exists(gallery_dir):
|
| 797 |
+
raise typer.BadParameter(f"Gallery directory not found: {gallery_dir}")
|
| 798 |
+
|
| 799 |
+
typer.echo("Launching Depth Anything 3 Gallery server...")
|
| 800 |
+
typer.echo(f"Gallery directory: {gallery_dir}")
|
| 801 |
+
typer.echo(f"Host: {host}")
|
| 802 |
+
typer.echo(f"Port: {port}")
|
| 803 |
+
typer.echo(f"Auto-open browser: {open_browser}")
|
| 804 |
+
|
| 805 |
+
try:
|
| 806 |
+
# Set command line arguments
|
| 807 |
+
import sys
|
| 808 |
+
|
| 809 |
+
sys.argv = ["gallery", "--dir", gallery_dir, "--host", host, "--port", str(port)]
|
| 810 |
+
if open_browser:
|
| 811 |
+
sys.argv.append("--open")
|
| 812 |
+
|
| 813 |
+
# Launch gallery server
|
| 814 |
+
gallery_main()
|
| 815 |
+
|
| 816 |
+
except KeyboardInterrupt:
|
| 817 |
+
typer.echo("\nGallery server stopped.")
|
| 818 |
+
except Exception as e:
|
| 819 |
+
typer.echo(f"Failed to launch Gallery server: {e}")
|
| 820 |
+
raise typer.Exit(1)
|
| 821 |
+
|
| 822 |
+
|
| 823 |
+
if __name__ == "__main__":
|
| 824 |
+
app()
|
depth_anything_3/configs/da3-base.yaml
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: DepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
net:
|
| 7 |
+
__object__:
|
| 8 |
+
path: depth_anything_3.model.dinov2.dinov2
|
| 9 |
+
name: DinoV2
|
| 10 |
+
args: as_params
|
| 11 |
+
|
| 12 |
+
name: vitb
|
| 13 |
+
out_layers: [5, 7, 9, 11]
|
| 14 |
+
alt_start: 4
|
| 15 |
+
qknorm_start: 4
|
| 16 |
+
rope_start: 4
|
| 17 |
+
cat_token: True
|
| 18 |
+
|
| 19 |
+
head:
|
| 20 |
+
__object__:
|
| 21 |
+
path: depth_anything_3.model.dualdpt
|
| 22 |
+
name: DualDPT
|
| 23 |
+
args: as_params
|
| 24 |
+
|
| 25 |
+
dim_in: &head_dim_in 1536
|
| 26 |
+
output_dim: 2
|
| 27 |
+
features: &head_features 128
|
| 28 |
+
out_channels: &head_out_channels [96, 192, 384, 768]
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
cam_enc:
|
| 32 |
+
__object__:
|
| 33 |
+
path: depth_anything_3.model.cam_enc
|
| 34 |
+
name: CameraEnc
|
| 35 |
+
args: as_params
|
| 36 |
+
|
| 37 |
+
dim_out: 768
|
| 38 |
+
|
| 39 |
+
cam_dec:
|
| 40 |
+
__object__:
|
| 41 |
+
path: depth_anything_3.model.cam_dec
|
| 42 |
+
name: CameraDec
|
| 43 |
+
args: as_params
|
| 44 |
+
|
| 45 |
+
dim_in: 1536
|
depth_anything_3/configs/da3-giant.yaml
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: DepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
net:
|
| 7 |
+
__object__:
|
| 8 |
+
path: depth_anything_3.model.dinov2.dinov2
|
| 9 |
+
name: DinoV2
|
| 10 |
+
args: as_params
|
| 11 |
+
|
| 12 |
+
name: vitg
|
| 13 |
+
out_layers: [19, 27, 33, 39]
|
| 14 |
+
alt_start: 13
|
| 15 |
+
qknorm_start: 13
|
| 16 |
+
rope_start: 13
|
| 17 |
+
cat_token: True
|
| 18 |
+
|
| 19 |
+
head:
|
| 20 |
+
__object__:
|
| 21 |
+
path: depth_anything_3.model.dualdpt
|
| 22 |
+
name: DualDPT
|
| 23 |
+
args: as_params
|
| 24 |
+
|
| 25 |
+
dim_in: &head_dim_in 3072
|
| 26 |
+
output_dim: 2
|
| 27 |
+
features: &head_features 256
|
| 28 |
+
out_channels: &head_out_channels [256, 512, 1024, 1024]
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
cam_enc:
|
| 32 |
+
__object__:
|
| 33 |
+
path: depth_anything_3.model.cam_enc
|
| 34 |
+
name: CameraEnc
|
| 35 |
+
args: as_params
|
| 36 |
+
|
| 37 |
+
dim_out: 1536
|
| 38 |
+
|
| 39 |
+
cam_dec:
|
| 40 |
+
__object__:
|
| 41 |
+
path: depth_anything_3.model.cam_dec
|
| 42 |
+
name: CameraDec
|
| 43 |
+
args: as_params
|
| 44 |
+
|
| 45 |
+
dim_in: 3072
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
gs_head:
|
| 49 |
+
__object__:
|
| 50 |
+
path: depth_anything_3.model.gsdpt
|
| 51 |
+
name: GSDPT
|
| 52 |
+
args: as_params
|
| 53 |
+
|
| 54 |
+
dim_in: *head_dim_in
|
| 55 |
+
output_dim: 38 # should align with gs_adapter's setting, for gs params
|
| 56 |
+
features: *head_features
|
| 57 |
+
out_channels: *head_out_channels
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
gs_adapter:
|
| 61 |
+
__object__:
|
| 62 |
+
path: depth_anything_3.model.gs_adapter
|
| 63 |
+
name: GaussianAdapter
|
| 64 |
+
args: as_params
|
| 65 |
+
|
| 66 |
+
sh_degree: 2
|
| 67 |
+
pred_color: false # predict SH coefficient if false
|
| 68 |
+
pred_offset_depth: true
|
| 69 |
+
pred_offset_xy: true
|
| 70 |
+
gaussian_scale_min: 1e-5
|
| 71 |
+
gaussian_scale_max: 30.0
|
depth_anything_3/configs/da3-large.yaml
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: DepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
net:
|
| 7 |
+
__object__:
|
| 8 |
+
path: depth_anything_3.model.dinov2.dinov2
|
| 9 |
+
name: DinoV2
|
| 10 |
+
args: as_params
|
| 11 |
+
|
| 12 |
+
name: vitl
|
| 13 |
+
out_layers: [11, 15, 19, 23]
|
| 14 |
+
alt_start: 8
|
| 15 |
+
qknorm_start: 8
|
| 16 |
+
rope_start: 8
|
| 17 |
+
cat_token: True
|
| 18 |
+
|
| 19 |
+
head:
|
| 20 |
+
__object__:
|
| 21 |
+
path: depth_anything_3.model.dualdpt
|
| 22 |
+
name: DualDPT
|
| 23 |
+
args: as_params
|
| 24 |
+
|
| 25 |
+
dim_in: &head_dim_in 2048
|
| 26 |
+
output_dim: 2
|
| 27 |
+
features: &head_features 256
|
| 28 |
+
out_channels: &head_out_channels [256, 512, 1024, 1024]
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
cam_enc:
|
| 32 |
+
__object__:
|
| 33 |
+
path: depth_anything_3.model.cam_enc
|
| 34 |
+
name: CameraEnc
|
| 35 |
+
args: as_params
|
| 36 |
+
|
| 37 |
+
dim_out: 1024
|
| 38 |
+
|
| 39 |
+
cam_dec:
|
| 40 |
+
__object__:
|
| 41 |
+
path: depth_anything_3.model.cam_dec
|
| 42 |
+
name: CameraDec
|
| 43 |
+
args: as_params
|
| 44 |
+
|
| 45 |
+
dim_in: 2048
|
depth_anything_3/configs/da3-small.yaml
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: DepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
net:
|
| 7 |
+
__object__:
|
| 8 |
+
path: depth_anything_3.model.dinov2.dinov2
|
| 9 |
+
name: DinoV2
|
| 10 |
+
args: as_params
|
| 11 |
+
|
| 12 |
+
name: vits
|
| 13 |
+
out_layers: [5, 7, 9, 11]
|
| 14 |
+
alt_start: 4
|
| 15 |
+
qknorm_start: 4
|
| 16 |
+
rope_start: 4
|
| 17 |
+
cat_token: True
|
| 18 |
+
|
| 19 |
+
head:
|
| 20 |
+
__object__:
|
| 21 |
+
path: depth_anything_3.model.dualdpt
|
| 22 |
+
name: DualDPT
|
| 23 |
+
args: as_params
|
| 24 |
+
|
| 25 |
+
dim_in: &head_dim_in 768
|
| 26 |
+
output_dim: 2
|
| 27 |
+
features: &head_features 64
|
| 28 |
+
out_channels: &head_out_channels [48, 96, 192, 384]
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
cam_enc:
|
| 32 |
+
__object__:
|
| 33 |
+
path: depth_anything_3.model.cam_enc
|
| 34 |
+
name: CameraEnc
|
| 35 |
+
args: as_params
|
| 36 |
+
|
| 37 |
+
dim_out: 384
|
| 38 |
+
|
| 39 |
+
cam_dec:
|
| 40 |
+
__object__:
|
| 41 |
+
path: depth_anything_3.model.cam_dec
|
| 42 |
+
name: CameraDec
|
| 43 |
+
args: as_params
|
| 44 |
+
|
| 45 |
+
dim_in: 768
|
depth_anything_3/configs/da3metric-large.yaml
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: DepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
net:
|
| 7 |
+
__object__:
|
| 8 |
+
path: depth_anything_3.model.dinov2.dinov2
|
| 9 |
+
name: DinoV2
|
| 10 |
+
args: as_params
|
| 11 |
+
|
| 12 |
+
name: vitl
|
| 13 |
+
out_layers: [4, 11, 17, 23]
|
| 14 |
+
alt_start: -1 # -1 means disable
|
| 15 |
+
qknorm_start: -1
|
| 16 |
+
rope_start: -1
|
| 17 |
+
cat_token: False
|
| 18 |
+
|
| 19 |
+
head:
|
| 20 |
+
__object__:
|
| 21 |
+
path: depth_anything_3.model.dpt
|
| 22 |
+
name: DPT
|
| 23 |
+
args: as_params
|
| 24 |
+
|
| 25 |
+
dim_in: 1024
|
| 26 |
+
output_dim: 1
|
| 27 |
+
features: 256
|
| 28 |
+
out_channels: [256, 512, 1024, 1024]
|
depth_anything_3/configs/da3mono-large.yaml
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: DepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
net:
|
| 7 |
+
__object__:
|
| 8 |
+
path: depth_anything_3.model.dinov2.dinov2
|
| 9 |
+
name: DinoV2
|
| 10 |
+
args: as_params
|
| 11 |
+
|
| 12 |
+
name: vitl
|
| 13 |
+
out_layers: [4, 11, 17, 23]
|
| 14 |
+
alt_start: -1 # -1 means disable
|
| 15 |
+
qknorm_start: -1
|
| 16 |
+
rope_start: -1
|
| 17 |
+
cat_token: False
|
| 18 |
+
|
| 19 |
+
head:
|
| 20 |
+
__object__:
|
| 21 |
+
path: depth_anything_3.model.dpt
|
| 22 |
+
name: DPT
|
| 23 |
+
args: as_params
|
| 24 |
+
|
| 25 |
+
dim_in: 1024
|
| 26 |
+
output_dim: 1
|
| 27 |
+
features: 256
|
| 28 |
+
out_channels: [256, 512, 1024, 1024]
|
depth_anything_3/configs/da3nested-giant-large.yaml
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__object__:
|
| 2 |
+
path: depth_anything_3.model.da3
|
| 3 |
+
name: NestedDepthAnything3Net
|
| 4 |
+
args: as_params
|
| 5 |
+
|
| 6 |
+
anyview:
|
| 7 |
+
__inherit__: depth_anything_3.configs.da3-giant
|
| 8 |
+
|
| 9 |
+
metric:
|
| 10 |
+
__inherit__: depth_anything_3.configs.da3metric-large
|
depth_anything_3/model/__init__.py
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
from depth_anything_3.model.da3 import DepthAnything3Net, NestedDepthAnything3Net
|
| 16 |
+
|
| 17 |
+
__export__ = [
|
| 18 |
+
NestedDepthAnything3Net,
|
| 19 |
+
DepthAnything3Net,
|
| 20 |
+
]
|
depth_anything_3/model/cam_dec.py
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import torch
|
| 16 |
+
import torch.nn as nn
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
class CameraDec(nn.Module):
|
| 20 |
+
def __init__(self, dim_in=1536):
|
| 21 |
+
super().__init__()
|
| 22 |
+
output_dim = dim_in
|
| 23 |
+
self.backbone = nn.Sequential(
|
| 24 |
+
nn.Linear(output_dim, output_dim),
|
| 25 |
+
nn.ReLU(),
|
| 26 |
+
nn.Linear(output_dim, output_dim),
|
| 27 |
+
nn.ReLU(),
|
| 28 |
+
)
|
| 29 |
+
self.fc_t = nn.Linear(output_dim, 3)
|
| 30 |
+
self.fc_qvec = nn.Linear(output_dim, 4)
|
| 31 |
+
self.fc_fov = nn.Sequential(nn.Linear(output_dim, 2), nn.ReLU())
|
| 32 |
+
|
| 33 |
+
def forward(self, feat, camera_encoding=None, *args, **kwargs):
|
| 34 |
+
B, N = feat.shape[:2]
|
| 35 |
+
feat = feat.reshape(B * N, -1)
|
| 36 |
+
feat = self.backbone(feat)
|
| 37 |
+
out_t = self.fc_t(feat.float()).reshape(B, N, 3)
|
| 38 |
+
if camera_encoding is None:
|
| 39 |
+
out_qvec = self.fc_qvec(feat.float()).reshape(B, N, 4)
|
| 40 |
+
out_fov = self.fc_fov(feat.float()).reshape(B, N, 2)
|
| 41 |
+
else:
|
| 42 |
+
out_qvec = camera_encoding[..., 3:7]
|
| 43 |
+
out_fov = camera_encoding[..., -2:]
|
| 44 |
+
pose_enc = torch.cat([out_t, out_qvec, out_fov], dim=-1)
|
| 45 |
+
return pose_enc
|
depth_anything_3/model/cam_enc.py
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import torch.nn as nn
|
| 16 |
+
|
| 17 |
+
from depth_anything_3.model.utils.attention import Mlp
|
| 18 |
+
from depth_anything_3.model.utils.block import Block
|
| 19 |
+
from depth_anything_3.model.utils.transform import extri_intri_to_pose_encoding
|
| 20 |
+
from depth_anything_3.utils.geometry import affine_inverse
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class CameraEnc(nn.Module):
|
| 24 |
+
"""
|
| 25 |
+
CameraHead predicts camera parameters from token representations using iterative refinement.
|
| 26 |
+
|
| 27 |
+
It applies a series of transformer blocks (the "trunk") to dedicated camera tokens.
|
| 28 |
+
"""
|
| 29 |
+
|
| 30 |
+
def __init__(
|
| 31 |
+
self,
|
| 32 |
+
dim_out: int = 1024,
|
| 33 |
+
dim_in: int = 9,
|
| 34 |
+
trunk_depth: int = 4,
|
| 35 |
+
target_dim: int = 9,
|
| 36 |
+
num_heads: int = 16,
|
| 37 |
+
mlp_ratio: int = 4,
|
| 38 |
+
init_values: float = 0.01,
|
| 39 |
+
**kwargs,
|
| 40 |
+
):
|
| 41 |
+
super().__init__()
|
| 42 |
+
self.target_dim = target_dim
|
| 43 |
+
self.trunk_depth = trunk_depth
|
| 44 |
+
self.trunk = nn.Sequential(
|
| 45 |
+
*[
|
| 46 |
+
Block(
|
| 47 |
+
dim=dim_out,
|
| 48 |
+
num_heads=num_heads,
|
| 49 |
+
mlp_ratio=mlp_ratio,
|
| 50 |
+
init_values=init_values,
|
| 51 |
+
)
|
| 52 |
+
for _ in range(trunk_depth)
|
| 53 |
+
]
|
| 54 |
+
)
|
| 55 |
+
self.token_norm = nn.LayerNorm(dim_out)
|
| 56 |
+
self.trunk_norm = nn.LayerNorm(dim_out)
|
| 57 |
+
self.pose_branch = Mlp(
|
| 58 |
+
in_features=dim_in,
|
| 59 |
+
hidden_features=dim_out // 2,
|
| 60 |
+
out_features=dim_out,
|
| 61 |
+
drop=0,
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
def forward(
|
| 65 |
+
self,
|
| 66 |
+
ext,
|
| 67 |
+
ixt,
|
| 68 |
+
image_size,
|
| 69 |
+
) -> tuple:
|
| 70 |
+
c2ws = affine_inverse(ext)
|
| 71 |
+
pose_encoding = extri_intri_to_pose_encoding(
|
| 72 |
+
c2ws,
|
| 73 |
+
ixt,
|
| 74 |
+
image_size,
|
| 75 |
+
)
|
| 76 |
+
pose_tokens = self.pose_branch(pose_encoding)
|
| 77 |
+
pose_tokens = self.token_norm(pose_tokens)
|
| 78 |
+
pose_tokens = self.trunk(pose_tokens)
|
| 79 |
+
pose_tokens = self.trunk_norm(pose_tokens)
|
| 80 |
+
return pose_tokens
|
depth_anything_3/model/da3.py
ADDED
|
@@ -0,0 +1,442 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
from addict import Dict
|
| 20 |
+
from omegaconf import DictConfig, OmegaConf
|
| 21 |
+
|
| 22 |
+
from depth_anything_3.cfg import create_object
|
| 23 |
+
from depth_anything_3.model.utils.transform import pose_encoding_to_extri_intri
|
| 24 |
+
from depth_anything_3.utils.alignment import (
|
| 25 |
+
apply_metric_scaling,
|
| 26 |
+
compute_alignment_mask,
|
| 27 |
+
compute_sky_mask,
|
| 28 |
+
least_squares_scale_scalar,
|
| 29 |
+
sample_tensor_for_quantile,
|
| 30 |
+
set_sky_regions_to_max_depth,
|
| 31 |
+
)
|
| 32 |
+
from depth_anything_3.utils.geometry import affine_inverse, as_homogeneous, map_pdf_to_opacity
|
| 33 |
+
from depth_anything_3.utils.ray_utils import get_extrinsic_from_camray
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def _wrap_cfg(cfg_obj):
|
| 37 |
+
return OmegaConf.create(cfg_obj)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
class DepthAnything3Net(nn.Module):
|
| 41 |
+
"""
|
| 42 |
+
Depth Anything 3 network for depth estimation and camera pose estimation.
|
| 43 |
+
|
| 44 |
+
This network consists of:
|
| 45 |
+
- Backbone: DinoV2 feature extractor
|
| 46 |
+
- Head: DPT or DualDPT for depth prediction
|
| 47 |
+
- Optional camera decoders for pose estimation
|
| 48 |
+
- Optional GSDPT for 3DGS prediction
|
| 49 |
+
|
| 50 |
+
Args:
|
| 51 |
+
preset: Configuration preset containing network dimensions and settings
|
| 52 |
+
|
| 53 |
+
Returns:
|
| 54 |
+
Dictionary containing:
|
| 55 |
+
- depth: Predicted depth map (B, H, W)
|
| 56 |
+
- depth_conf: Depth confidence map (B, H, W)
|
| 57 |
+
- extrinsics: Camera extrinsics (B, N, 4, 4)
|
| 58 |
+
- intrinsics: Camera intrinsics (B, N, 3, 3)
|
| 59 |
+
- gaussians: 3D Gaussian Splats (world space), type: model.gs_adapter.Gaussians
|
| 60 |
+
- aux: Auxiliary features for specified layers
|
| 61 |
+
"""
|
| 62 |
+
|
| 63 |
+
# Patch size for feature extraction
|
| 64 |
+
PATCH_SIZE = 14
|
| 65 |
+
|
| 66 |
+
def __init__(self, net, head, cam_dec=None, cam_enc=None, gs_head=None, gs_adapter=None):
|
| 67 |
+
"""
|
| 68 |
+
Initialize DepthAnything3Net with given yaml-initialized configuration.
|
| 69 |
+
"""
|
| 70 |
+
super().__init__()
|
| 71 |
+
self.backbone = net if isinstance(net, nn.Module) else create_object(_wrap_cfg(net))
|
| 72 |
+
self.head = head if isinstance(head, nn.Module) else create_object(_wrap_cfg(head))
|
| 73 |
+
self.cam_dec, self.cam_enc = None, None
|
| 74 |
+
if cam_dec is not None:
|
| 75 |
+
self.cam_dec = (
|
| 76 |
+
cam_dec if isinstance(cam_dec, nn.Module) else create_object(_wrap_cfg(cam_dec))
|
| 77 |
+
)
|
| 78 |
+
self.cam_enc = (
|
| 79 |
+
cam_enc if isinstance(cam_enc, nn.Module) else create_object(_wrap_cfg(cam_enc))
|
| 80 |
+
)
|
| 81 |
+
self.gs_adapter, self.gs_head = None, None
|
| 82 |
+
if gs_head is not None and gs_adapter is not None:
|
| 83 |
+
self.gs_adapter = (
|
| 84 |
+
gs_adapter
|
| 85 |
+
if isinstance(gs_adapter, nn.Module)
|
| 86 |
+
else create_object(_wrap_cfg(gs_adapter))
|
| 87 |
+
)
|
| 88 |
+
gs_out_dim = self.gs_adapter.d_in + 1
|
| 89 |
+
if isinstance(gs_head, nn.Module):
|
| 90 |
+
assert (
|
| 91 |
+
gs_head.out_dim == gs_out_dim
|
| 92 |
+
), f"gs_head.out_dim should be {gs_out_dim}, got {gs_head.out_dim}"
|
| 93 |
+
self.gs_head = gs_head
|
| 94 |
+
else:
|
| 95 |
+
assert (
|
| 96 |
+
gs_head["output_dim"] == gs_out_dim
|
| 97 |
+
), f"gs_head output_dim should set to {gs_out_dim}, got {gs_head['output_dim']}"
|
| 98 |
+
self.gs_head = create_object(_wrap_cfg(gs_head))
|
| 99 |
+
|
| 100 |
+
def forward(
|
| 101 |
+
self,
|
| 102 |
+
x: torch.Tensor,
|
| 103 |
+
extrinsics: torch.Tensor | None = None,
|
| 104 |
+
intrinsics: torch.Tensor | None = None,
|
| 105 |
+
export_feat_layers: list[int] | None = [],
|
| 106 |
+
infer_gs: bool = False,
|
| 107 |
+
use_ray_pose: bool = False,
|
| 108 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 109 |
+
) -> Dict[str, torch.Tensor]:
|
| 110 |
+
"""
|
| 111 |
+
Forward pass through the network.
|
| 112 |
+
|
| 113 |
+
Args:
|
| 114 |
+
x: Input images (B, N, 3, H, W)
|
| 115 |
+
extrinsics: Camera extrinsics (B, N, 4, 4)
|
| 116 |
+
intrinsics: Camera intrinsics (B, N, 3, 3)
|
| 117 |
+
feat_layers: List of layer indices to extract features from
|
| 118 |
+
infer_gs: Enable Gaussian Splatting branch
|
| 119 |
+
use_ray_pose: Use ray-based pose estimation
|
| 120 |
+
ref_view_strategy: Strategy for selecting reference view
|
| 121 |
+
|
| 122 |
+
Returns:
|
| 123 |
+
Dictionary containing predictions and auxiliary features
|
| 124 |
+
"""
|
| 125 |
+
# Extract features using backbone
|
| 126 |
+
if extrinsics is not None:
|
| 127 |
+
with torch.autocast(device_type=x.device.type, enabled=False):
|
| 128 |
+
cam_token = self.cam_enc(extrinsics, intrinsics, x.shape[-2:])
|
| 129 |
+
else:
|
| 130 |
+
cam_token = None
|
| 131 |
+
|
| 132 |
+
feats, aux_feats = self.backbone(
|
| 133 |
+
x, cam_token=cam_token, export_feat_layers=export_feat_layers, ref_view_strategy=ref_view_strategy
|
| 134 |
+
)
|
| 135 |
+
# feats = [[item for item in feat] for feat in feats]
|
| 136 |
+
H, W = x.shape[-2], x.shape[-1]
|
| 137 |
+
|
| 138 |
+
# Process features through depth head
|
| 139 |
+
with torch.autocast(device_type=x.device.type, enabled=False):
|
| 140 |
+
output = self._process_depth_head(feats, H, W)
|
| 141 |
+
if use_ray_pose:
|
| 142 |
+
output = self._process_ray_pose_estimation(output, H, W)
|
| 143 |
+
else:
|
| 144 |
+
output = self._process_camera_estimation(feats, H, W, output)
|
| 145 |
+
if infer_gs:
|
| 146 |
+
output = self._process_gs_head(feats, H, W, output, x, extrinsics, intrinsics)
|
| 147 |
+
|
| 148 |
+
output = self._process_mono_sky_estimation(output)
|
| 149 |
+
|
| 150 |
+
# Extract auxiliary features if requested
|
| 151 |
+
output.aux = self._extract_auxiliary_features(aux_feats, export_feat_layers, H, W)
|
| 152 |
+
|
| 153 |
+
return output
|
| 154 |
+
|
| 155 |
+
def _process_mono_sky_estimation(
|
| 156 |
+
self, output: Dict[str, torch.Tensor]
|
| 157 |
+
) -> Dict[str, torch.Tensor]:
|
| 158 |
+
"""Process mono sky estimation."""
|
| 159 |
+
if "sky" not in output:
|
| 160 |
+
return output
|
| 161 |
+
non_sky_mask = compute_sky_mask(output.sky, threshold=0.3)
|
| 162 |
+
if non_sky_mask.sum() <= 10:
|
| 163 |
+
return output
|
| 164 |
+
if (~non_sky_mask).sum() <= 10:
|
| 165 |
+
return output
|
| 166 |
+
|
| 167 |
+
non_sky_depth = output.depth[non_sky_mask]
|
| 168 |
+
if non_sky_depth.numel() > 100000:
|
| 169 |
+
idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device)
|
| 170 |
+
sampled_depth = non_sky_depth[idx]
|
| 171 |
+
else:
|
| 172 |
+
sampled_depth = non_sky_depth
|
| 173 |
+
non_sky_max = torch.quantile(sampled_depth, 0.99)
|
| 174 |
+
|
| 175 |
+
# Set sky regions to maximum depth and high confidence
|
| 176 |
+
output.depth, _ = set_sky_regions_to_max_depth(
|
| 177 |
+
output.depth, None, non_sky_mask, max_depth=non_sky_max
|
| 178 |
+
)
|
| 179 |
+
return output
|
| 180 |
+
|
| 181 |
+
def _process_ray_pose_estimation(
|
| 182 |
+
self, output: Dict[str, torch.Tensor], height: int, width: int
|
| 183 |
+
) -> Dict[str, torch.Tensor]:
|
| 184 |
+
"""Process ray pose estimation if ray pose decoder is available."""
|
| 185 |
+
if "ray" in output and "ray_conf" in output:
|
| 186 |
+
pred_extrinsic, pred_focal_lengths, pred_principal_points = get_extrinsic_from_camray(
|
| 187 |
+
output.ray,
|
| 188 |
+
output.ray_conf,
|
| 189 |
+
output.ray.shape[-3],
|
| 190 |
+
output.ray.shape[-2],
|
| 191 |
+
)
|
| 192 |
+
pred_extrinsic = affine_inverse(pred_extrinsic) # w2c -> c2w
|
| 193 |
+
pred_extrinsic = pred_extrinsic[:, :, :3, :]
|
| 194 |
+
pred_intrinsic = torch.eye(3, 3)[None, None].repeat(pred_extrinsic.shape[0], pred_extrinsic.shape[1], 1, 1).clone().to(pred_extrinsic.device)
|
| 195 |
+
pred_intrinsic[:, :, 0, 0] = pred_focal_lengths[:, :, 0] / 2 * width
|
| 196 |
+
pred_intrinsic[:, :, 1, 1] = pred_focal_lengths[:, :, 1] / 2 * height
|
| 197 |
+
pred_intrinsic[:, :, 0, 2] = pred_principal_points[:, :, 0] * width * 0.5
|
| 198 |
+
pred_intrinsic[:, :, 1, 2] = pred_principal_points[:, :, 1] * height * 0.5
|
| 199 |
+
del output.ray
|
| 200 |
+
del output.ray_conf
|
| 201 |
+
output.extrinsics = pred_extrinsic
|
| 202 |
+
output.intrinsics = pred_intrinsic
|
| 203 |
+
return output
|
| 204 |
+
|
| 205 |
+
def _process_depth_head(
|
| 206 |
+
self, feats: list[torch.Tensor], H: int, W: int
|
| 207 |
+
) -> Dict[str, torch.Tensor]:
|
| 208 |
+
"""Process features through the depth prediction head."""
|
| 209 |
+
return self.head(feats, H, W, patch_start_idx=0)
|
| 210 |
+
|
| 211 |
+
def _process_camera_estimation(
|
| 212 |
+
self, feats: list[torch.Tensor], H: int, W: int, output: Dict[str, torch.Tensor]
|
| 213 |
+
) -> Dict[str, torch.Tensor]:
|
| 214 |
+
"""Process camera pose estimation if camera decoder is available."""
|
| 215 |
+
if self.cam_dec is not None:
|
| 216 |
+
pose_enc = self.cam_dec(feats[-1][1])
|
| 217 |
+
# Remove ray information as it's not needed for pose estimation
|
| 218 |
+
if "ray" in output:
|
| 219 |
+
del output.ray
|
| 220 |
+
if "ray_conf" in output:
|
| 221 |
+
del output.ray_conf
|
| 222 |
+
|
| 223 |
+
# Convert pose encoding to extrinsics and intrinsics
|
| 224 |
+
c2w, ixt = pose_encoding_to_extri_intri(pose_enc, (H, W))
|
| 225 |
+
output.extrinsics = affine_inverse(c2w)
|
| 226 |
+
output.intrinsics = ixt
|
| 227 |
+
|
| 228 |
+
return output
|
| 229 |
+
|
| 230 |
+
def _process_gs_head(
|
| 231 |
+
self,
|
| 232 |
+
feats: list[torch.Tensor],
|
| 233 |
+
H: int,
|
| 234 |
+
W: int,
|
| 235 |
+
output: Dict[str, torch.Tensor],
|
| 236 |
+
in_images: torch.Tensor,
|
| 237 |
+
extrinsics: torch.Tensor | None = None,
|
| 238 |
+
intrinsics: torch.Tensor | None = None,
|
| 239 |
+
) -> Dict[str, torch.Tensor]:
|
| 240 |
+
"""Process 3DGS parameters estimation if 3DGS head is available."""
|
| 241 |
+
if self.gs_head is None or self.gs_adapter is None:
|
| 242 |
+
return output
|
| 243 |
+
assert output.get("depth", None) is not None, "must provide MV depth for the GS head."
|
| 244 |
+
|
| 245 |
+
# The depth is defined in the DA3 model's camera space,
|
| 246 |
+
# so even with provided GT camera poses,
|
| 247 |
+
# we instead use the predicted camera poses for better alignment.
|
| 248 |
+
ctx_extr = output.get("extrinsics", None)
|
| 249 |
+
ctx_intr = output.get("intrinsics", None)
|
| 250 |
+
assert (
|
| 251 |
+
ctx_extr is not None and ctx_intr is not None
|
| 252 |
+
), "must process camera info first if GT is not available"
|
| 253 |
+
|
| 254 |
+
gt_extr = extrinsics
|
| 255 |
+
# homo the extr if needed
|
| 256 |
+
ctx_extr = as_homogeneous(ctx_extr)
|
| 257 |
+
if gt_extr is not None:
|
| 258 |
+
gt_extr = as_homogeneous(gt_extr)
|
| 259 |
+
|
| 260 |
+
# forward through the gs_dpt head to get 'camera space' parameters
|
| 261 |
+
gs_outs = self.gs_head(
|
| 262 |
+
feats=feats,
|
| 263 |
+
H=H,
|
| 264 |
+
W=W,
|
| 265 |
+
patch_start_idx=0,
|
| 266 |
+
images=in_images,
|
| 267 |
+
)
|
| 268 |
+
raw_gaussians = gs_outs.raw_gs
|
| 269 |
+
densities = gs_outs.raw_gs_conf
|
| 270 |
+
|
| 271 |
+
# convert to 'world space' 3DGS parameters; ready to export and render
|
| 272 |
+
# gt_extr could be None, and will be used to align the pose scale if available
|
| 273 |
+
gs_world = self.gs_adapter(
|
| 274 |
+
extrinsics=ctx_extr,
|
| 275 |
+
intrinsics=ctx_intr,
|
| 276 |
+
depths=output.depth,
|
| 277 |
+
opacities=map_pdf_to_opacity(densities),
|
| 278 |
+
raw_gaussians=raw_gaussians,
|
| 279 |
+
image_shape=(H, W),
|
| 280 |
+
gt_extrinsics=gt_extr,
|
| 281 |
+
)
|
| 282 |
+
output.gaussians = gs_world
|
| 283 |
+
|
| 284 |
+
return output
|
| 285 |
+
|
| 286 |
+
def _extract_auxiliary_features(
|
| 287 |
+
self, feats: list[torch.Tensor], feat_layers: list[int], H: int, W: int
|
| 288 |
+
) -> Dict[str, torch.Tensor]:
|
| 289 |
+
"""Extract auxiliary features from specified layers."""
|
| 290 |
+
aux_features = Dict()
|
| 291 |
+
assert len(feats) == len(feat_layers)
|
| 292 |
+
for feat, feat_layer in zip(feats, feat_layers):
|
| 293 |
+
# Reshape features to spatial dimensions
|
| 294 |
+
feat_reshaped = feat.reshape(
|
| 295 |
+
[
|
| 296 |
+
feat.shape[0],
|
| 297 |
+
feat.shape[1],
|
| 298 |
+
H // self.PATCH_SIZE,
|
| 299 |
+
W // self.PATCH_SIZE,
|
| 300 |
+
feat.shape[-1],
|
| 301 |
+
]
|
| 302 |
+
)
|
| 303 |
+
aux_features[f"feat_layer_{feat_layer}"] = feat_reshaped
|
| 304 |
+
|
| 305 |
+
return aux_features
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
class NestedDepthAnything3Net(nn.Module):
|
| 309 |
+
"""
|
| 310 |
+
Nested Depth Anything 3 network with metric scaling capabilities.
|
| 311 |
+
|
| 312 |
+
This network combines two DepthAnything3Net branches:
|
| 313 |
+
- Main branch: Standard depth estimation
|
| 314 |
+
- Metric branch: Metric depth estimation for scaling alignment
|
| 315 |
+
|
| 316 |
+
The network performs depth alignment using least squares scaling
|
| 317 |
+
and handles sky region masking for improved depth estimation.
|
| 318 |
+
|
| 319 |
+
Args:
|
| 320 |
+
preset: Configuration for the main depth estimation branch
|
| 321 |
+
second_preset: Configuration for the metric depth branch
|
| 322 |
+
"""
|
| 323 |
+
|
| 324 |
+
def __init__(self, anyview: DictConfig, metric: DictConfig):
|
| 325 |
+
"""
|
| 326 |
+
Initialize NestedDepthAnything3Net with two branches.
|
| 327 |
+
|
| 328 |
+
Args:
|
| 329 |
+
preset: Configuration for main depth estimation branch
|
| 330 |
+
second_preset: Configuration for metric depth branch
|
| 331 |
+
"""
|
| 332 |
+
super().__init__()
|
| 333 |
+
self.da3 = create_object(anyview)
|
| 334 |
+
self.da3_metric = create_object(metric)
|
| 335 |
+
|
| 336 |
+
def forward(
|
| 337 |
+
self,
|
| 338 |
+
x: torch.Tensor,
|
| 339 |
+
extrinsics: torch.Tensor | None = None,
|
| 340 |
+
intrinsics: torch.Tensor | None = None,
|
| 341 |
+
export_feat_layers: list[int] | None = [],
|
| 342 |
+
infer_gs: bool = False,
|
| 343 |
+
use_ray_pose: bool = False,
|
| 344 |
+
ref_view_strategy: str = "saddle_balanced",
|
| 345 |
+
) -> Dict[str, torch.Tensor]:
|
| 346 |
+
"""
|
| 347 |
+
Forward pass through both branches with metric scaling alignment.
|
| 348 |
+
|
| 349 |
+
Args:
|
| 350 |
+
x: Input images (B, N, 3, H, W)
|
| 351 |
+
extrinsics: Camera extrinsics (B, N, 4, 4) - unused
|
| 352 |
+
intrinsics: Camera intrinsics (B, N, 3, 3) - unused
|
| 353 |
+
feat_layers: List of layer indices to extract features from
|
| 354 |
+
infer_gs: Enable Gaussian Splatting branch
|
| 355 |
+
use_ray_pose: Use ray-based pose estimation
|
| 356 |
+
ref_view_strategy: Strategy for selecting reference view
|
| 357 |
+
|
| 358 |
+
Returns:
|
| 359 |
+
Dictionary containing aligned depth predictions and camera parameters
|
| 360 |
+
"""
|
| 361 |
+
# Get predictions from both branches
|
| 362 |
+
output = self.da3(
|
| 363 |
+
x, extrinsics, intrinsics, export_feat_layers=export_feat_layers, infer_gs=infer_gs, use_ray_pose=use_ray_pose, ref_view_strategy=ref_view_strategy
|
| 364 |
+
)
|
| 365 |
+
metric_output = self.da3_metric(x)
|
| 366 |
+
|
| 367 |
+
# Apply metric scaling and alignment
|
| 368 |
+
output = self._apply_metric_scaling(output, metric_output)
|
| 369 |
+
output = self._apply_depth_alignment(output, metric_output)
|
| 370 |
+
output = self._handle_sky_regions(output, metric_output)
|
| 371 |
+
|
| 372 |
+
return output
|
| 373 |
+
|
| 374 |
+
def _apply_metric_scaling(
|
| 375 |
+
self, output: Dict[str, torch.Tensor], metric_output: Dict[str, torch.Tensor]
|
| 376 |
+
) -> Dict[str, torch.Tensor]:
|
| 377 |
+
"""Apply metric scaling to the metric depth output."""
|
| 378 |
+
# Scale metric depth based on camera intrinsics
|
| 379 |
+
metric_output.depth = apply_metric_scaling(
|
| 380 |
+
metric_output.depth,
|
| 381 |
+
output.intrinsics,
|
| 382 |
+
)
|
| 383 |
+
return output
|
| 384 |
+
|
| 385 |
+
def _apply_depth_alignment(
|
| 386 |
+
self, output: Dict[str, torch.Tensor], metric_output: Dict[str, torch.Tensor]
|
| 387 |
+
) -> Dict[str, torch.Tensor]:
|
| 388 |
+
"""Apply depth alignment using least squares scaling."""
|
| 389 |
+
# Compute non-sky mask
|
| 390 |
+
non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3)
|
| 391 |
+
|
| 392 |
+
# Ensure we have enough non-sky pixels
|
| 393 |
+
assert non_sky_mask.sum() > 10, "Insufficient non-sky pixels for alignment"
|
| 394 |
+
|
| 395 |
+
# Sample depth confidence for quantile computation
|
| 396 |
+
depth_conf_ns = output.depth_conf[non_sky_mask]
|
| 397 |
+
depth_conf_sampled = sample_tensor_for_quantile(depth_conf_ns, max_samples=100000)
|
| 398 |
+
median_conf = torch.quantile(depth_conf_sampled, 0.5)
|
| 399 |
+
|
| 400 |
+
# Compute alignment mask
|
| 401 |
+
align_mask = compute_alignment_mask(
|
| 402 |
+
output.depth_conf, non_sky_mask, output.depth, metric_output.depth, median_conf
|
| 403 |
+
)
|
| 404 |
+
|
| 405 |
+
# Compute scale factor using least squares
|
| 406 |
+
valid_depth = output.depth[align_mask]
|
| 407 |
+
valid_metric_depth = metric_output.depth[align_mask]
|
| 408 |
+
scale_factor = least_squares_scale_scalar(valid_metric_depth, valid_depth)
|
| 409 |
+
|
| 410 |
+
# Apply scaling to depth and extrinsics
|
| 411 |
+
output.depth *= scale_factor
|
| 412 |
+
output.extrinsics[:, :, :3, 3] *= scale_factor
|
| 413 |
+
output.is_metric = 1
|
| 414 |
+
output.scale_factor = scale_factor.item()
|
| 415 |
+
|
| 416 |
+
return output
|
| 417 |
+
|
| 418 |
+
def _handle_sky_regions(
|
| 419 |
+
self,
|
| 420 |
+
output: Dict[str, torch.Tensor],
|
| 421 |
+
metric_output: Dict[str, torch.Tensor],
|
| 422 |
+
sky_depth_def: float = 200.0,
|
| 423 |
+
) -> Dict[str, torch.Tensor]:
|
| 424 |
+
"""Handle sky regions by setting them to maximum depth."""
|
| 425 |
+
non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3)
|
| 426 |
+
|
| 427 |
+
# Compute maximum depth for non-sky regions
|
| 428 |
+
# Use sampling to safely compute quantile on large tensors
|
| 429 |
+
non_sky_depth = output.depth[non_sky_mask]
|
| 430 |
+
if non_sky_depth.numel() > 100000:
|
| 431 |
+
idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device)
|
| 432 |
+
sampled_depth = non_sky_depth[idx]
|
| 433 |
+
else:
|
| 434 |
+
sampled_depth = non_sky_depth
|
| 435 |
+
non_sky_max = min(torch.quantile(sampled_depth, 0.99), sky_depth_def)
|
| 436 |
+
|
| 437 |
+
# Set sky regions to maximum depth and high confidence
|
| 438 |
+
output.depth, output.depth_conf = set_sky_regions_to_max_depth(
|
| 439 |
+
output.depth, output.depth_conf, non_sky_mask, max_depth=non_sky_max
|
| 440 |
+
)
|
| 441 |
+
|
| 442 |
+
return output
|
depth_anything_3/model/dinov2/dinov2.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
#
|
| 3 |
+
# This source code is licensed under the Apache License, Version 2.0
|
| 4 |
+
# found in the LICENSE file in the root directory of this source tree.
|
| 5 |
+
|
| 6 |
+
# References:
|
| 7 |
+
# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
|
| 8 |
+
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
from typing import List
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
|
| 14 |
+
from depth_anything_3.model.dinov2.vision_transformer import (
|
| 15 |
+
vit_base,
|
| 16 |
+
vit_giant2,
|
| 17 |
+
vit_large,
|
| 18 |
+
vit_small,
|
| 19 |
+
)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class DinoV2(nn.Module):
|
| 23 |
+
def __init__(
|
| 24 |
+
self,
|
| 25 |
+
name: str,
|
| 26 |
+
out_layers: List[int],
|
| 27 |
+
alt_start: int = -1,
|
| 28 |
+
qknorm_start: int = -1,
|
| 29 |
+
rope_start: int = -1,
|
| 30 |
+
cat_token: bool = True,
|
| 31 |
+
**kwargs,
|
| 32 |
+
):
|
| 33 |
+
super().__init__()
|
| 34 |
+
assert name in {"vits", "vitb", "vitl", "vitg"}
|
| 35 |
+
self.name = name
|
| 36 |
+
self.out_layers = out_layers
|
| 37 |
+
self.alt_start = alt_start
|
| 38 |
+
self.qknorm_start = qknorm_start
|
| 39 |
+
self.rope_start = rope_start
|
| 40 |
+
self.cat_token = cat_token
|
| 41 |
+
encoder_map = {
|
| 42 |
+
"vits": vit_small,
|
| 43 |
+
"vitb": vit_base,
|
| 44 |
+
"vitl": vit_large,
|
| 45 |
+
"vitg": vit_giant2,
|
| 46 |
+
}
|
| 47 |
+
encoder_fn = encoder_map[self.name]
|
| 48 |
+
ffn_layer = "swiglufused" if self.name == "vitg" else "mlp"
|
| 49 |
+
self.pretrained = encoder_fn(
|
| 50 |
+
img_size=518,
|
| 51 |
+
patch_size=14,
|
| 52 |
+
ffn_layer=ffn_layer,
|
| 53 |
+
alt_start=alt_start,
|
| 54 |
+
qknorm_start=qknorm_start,
|
| 55 |
+
rope_start=rope_start,
|
| 56 |
+
cat_token=cat_token,
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
def forward(self, x, **kwargs):
|
| 60 |
+
return self.pretrained.get_intermediate_layers(
|
| 61 |
+
x,
|
| 62 |
+
self.out_layers,
|
| 63 |
+
**kwargs,
|
| 64 |
+
)
|
depth_anything_3/model/dinov2/layers/__init__.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# from .attention import MemEffAttention
|
| 8 |
+
from .block import Block
|
| 9 |
+
from .layer_scale import LayerScale
|
| 10 |
+
from .mlp import Mlp
|
| 11 |
+
from .patch_embed import PatchEmbed
|
| 12 |
+
from .rope import PositionGetter, RotaryPositionEmbedding2D
|
| 13 |
+
from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
|
| 14 |
+
|
| 15 |
+
__all__ = [
|
| 16 |
+
Mlp,
|
| 17 |
+
PatchEmbed,
|
| 18 |
+
SwiGLUFFN,
|
| 19 |
+
SwiGLUFFNFused,
|
| 20 |
+
Block,
|
| 21 |
+
# MemEffAttention,
|
| 22 |
+
LayerScale,
|
| 23 |
+
PositionGetter,
|
| 24 |
+
RotaryPositionEmbedding2D,
|
| 25 |
+
]
|
depth_anything_3/model/dinov2/layers/attention.py
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# References:
|
| 8 |
+
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
| 9 |
+
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
| 10 |
+
|
| 11 |
+
import logging
|
| 12 |
+
import torch.nn.functional as F
|
| 13 |
+
from torch import Tensor, nn
|
| 14 |
+
|
| 15 |
+
logger = logging.getLogger("dinov2")
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class Attention(nn.Module):
|
| 19 |
+
def __init__(
|
| 20 |
+
self,
|
| 21 |
+
dim: int,
|
| 22 |
+
num_heads: int = 8,
|
| 23 |
+
qkv_bias: bool = False,
|
| 24 |
+
proj_bias: bool = True,
|
| 25 |
+
attn_drop: float = 0.0,
|
| 26 |
+
proj_drop: float = 0.0,
|
| 27 |
+
norm_layer: nn.Module = nn.LayerNorm,
|
| 28 |
+
qk_norm: bool = False,
|
| 29 |
+
fused_attn: bool = True, # use F.scaled_dot_product_attention or not
|
| 30 |
+
rope=None,
|
| 31 |
+
) -> None:
|
| 32 |
+
super().__init__()
|
| 33 |
+
assert dim % num_heads == 0, "dim should be divisible by num_heads"
|
| 34 |
+
self.num_heads = num_heads
|
| 35 |
+
head_dim = dim // num_heads
|
| 36 |
+
self.scale = head_dim**-0.5
|
| 37 |
+
self.fused_attn = fused_attn
|
| 38 |
+
|
| 39 |
+
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
| 40 |
+
self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
|
| 41 |
+
self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
|
| 42 |
+
self.attn_drop = nn.Dropout(attn_drop)
|
| 43 |
+
self.proj = nn.Linear(dim, dim, bias=proj_bias)
|
| 44 |
+
self.proj_drop = nn.Dropout(proj_drop)
|
| 45 |
+
self.rope = rope
|
| 46 |
+
|
| 47 |
+
def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor:
|
| 48 |
+
B, N, C = x.shape
|
| 49 |
+
qkv = (
|
| 50 |
+
self.qkv(x)
|
| 51 |
+
.reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
| 52 |
+
.permute(2, 0, 3, 1, 4)
|
| 53 |
+
)
|
| 54 |
+
q, k, v = qkv[0], qkv[1], qkv[2]
|
| 55 |
+
q, k = self.q_norm(q), self.k_norm(k)
|
| 56 |
+
if self.rope is not None and pos is not None:
|
| 57 |
+
q = self.rope(q, pos)
|
| 58 |
+
k = self.rope(k, pos)
|
| 59 |
+
if self.fused_attn:
|
| 60 |
+
x = F.scaled_dot_product_attention(
|
| 61 |
+
q,
|
| 62 |
+
k,
|
| 63 |
+
v,
|
| 64 |
+
dropout_p=self.attn_drop.p if self.training else 0.0,
|
| 65 |
+
attn_mask=(
|
| 66 |
+
(attn_mask)[:, None].repeat(1, self.num_heads, 1, 1)
|
| 67 |
+
if attn_mask is not None
|
| 68 |
+
else None
|
| 69 |
+
),
|
| 70 |
+
)
|
| 71 |
+
else:
|
| 72 |
+
q = q * self.scale
|
| 73 |
+
attn = q @ k.transpose(-2, -1)
|
| 74 |
+
attn = attn.softmax(dim=-1)
|
| 75 |
+
attn = self.attn_drop(attn)
|
| 76 |
+
x = attn @ v
|
| 77 |
+
|
| 78 |
+
x = x.transpose(1, 2).reshape(B, N, C)
|
| 79 |
+
x = self.proj(x)
|
| 80 |
+
x = self.proj_drop(x)
|
| 81 |
+
return x
|
| 82 |
+
|
| 83 |
+
def _forward(self, x: Tensor) -> Tensor:
|
| 84 |
+
B, N, C = x.shape
|
| 85 |
+
qkv = (
|
| 86 |
+
self.qkv(x)
|
| 87 |
+
.reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
| 88 |
+
.permute(2, 0, 3, 1, 4)
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
|
| 92 |
+
attn = q @ k.transpose(-2, -1)
|
| 93 |
+
|
| 94 |
+
attn = attn.softmax(dim=-1)
|
| 95 |
+
attn = self.attn_drop(attn)
|
| 96 |
+
|
| 97 |
+
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
| 98 |
+
x = self.proj(x)
|
| 99 |
+
x = self.proj_drop(x)
|
| 100 |
+
return x
|
depth_anything_3/model/dinov2/layers/block.py
ADDED
|
@@ -0,0 +1,143 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# flake8: noqa: F821
|
| 2 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 3 |
+
# All rights reserved.
|
| 4 |
+
#
|
| 5 |
+
# This source code is licensed under the license found in the
|
| 6 |
+
# LICENSE file in the root directory of this source tree.
|
| 7 |
+
|
| 8 |
+
# References:
|
| 9 |
+
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
| 10 |
+
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
| 11 |
+
|
| 12 |
+
import logging
|
| 13 |
+
from typing import Callable, Optional
|
| 14 |
+
import torch
|
| 15 |
+
from torch import Tensor, nn
|
| 16 |
+
|
| 17 |
+
from .attention import Attention
|
| 18 |
+
from .drop_path import DropPath
|
| 19 |
+
from .layer_scale import LayerScale
|
| 20 |
+
from .mlp import Mlp
|
| 21 |
+
|
| 22 |
+
logger = logging.getLogger("dinov2")
|
| 23 |
+
XFORMERS_AVAILABLE = True
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
class Block(nn.Module):
|
| 27 |
+
def __init__(
|
| 28 |
+
self,
|
| 29 |
+
dim: int,
|
| 30 |
+
num_heads: int,
|
| 31 |
+
mlp_ratio: float = 4.0,
|
| 32 |
+
qkv_bias: bool = False,
|
| 33 |
+
proj_bias: bool = True,
|
| 34 |
+
ffn_bias: bool = True,
|
| 35 |
+
drop: float = 0.0,
|
| 36 |
+
attn_drop: float = 0.0,
|
| 37 |
+
init_values=None,
|
| 38 |
+
drop_path: float = 0.0,
|
| 39 |
+
act_layer: Callable[..., nn.Module] = nn.GELU,
|
| 40 |
+
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
|
| 41 |
+
attn_class: Callable[..., nn.Module] = Attention,
|
| 42 |
+
ffn_layer: Callable[..., nn.Module] = Mlp,
|
| 43 |
+
qk_norm: bool = False,
|
| 44 |
+
rope=None,
|
| 45 |
+
ln_eps: float = 1e-6,
|
| 46 |
+
) -> None:
|
| 47 |
+
super().__init__()
|
| 48 |
+
# print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
|
| 49 |
+
self.norm1 = norm_layer(dim, eps=ln_eps)
|
| 50 |
+
self.attn = attn_class(
|
| 51 |
+
dim,
|
| 52 |
+
num_heads=num_heads,
|
| 53 |
+
qkv_bias=qkv_bias,
|
| 54 |
+
proj_bias=proj_bias,
|
| 55 |
+
attn_drop=attn_drop,
|
| 56 |
+
proj_drop=drop,
|
| 57 |
+
qk_norm=qk_norm,
|
| 58 |
+
rope=rope,
|
| 59 |
+
)
|
| 60 |
+
self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
| 61 |
+
self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
| 62 |
+
|
| 63 |
+
self.norm2 = norm_layer(dim, eps=ln_eps)
|
| 64 |
+
mlp_hidden_dim = int(dim * mlp_ratio)
|
| 65 |
+
self.mlp = ffn_layer(
|
| 66 |
+
in_features=dim,
|
| 67 |
+
hidden_features=mlp_hidden_dim,
|
| 68 |
+
act_layer=act_layer,
|
| 69 |
+
drop=drop,
|
| 70 |
+
bias=ffn_bias,
|
| 71 |
+
)
|
| 72 |
+
self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
| 73 |
+
self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
| 74 |
+
|
| 75 |
+
self.sample_drop_ratio = drop_path
|
| 76 |
+
|
| 77 |
+
def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor:
|
| 78 |
+
def attn_residual_func(x: Tensor, pos=None, attn_mask=None) -> Tensor:
|
| 79 |
+
return self.ls1(self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask))
|
| 80 |
+
|
| 81 |
+
def ffn_residual_func(x: Tensor) -> Tensor:
|
| 82 |
+
return self.ls2(self.mlp(self.norm2(x)))
|
| 83 |
+
|
| 84 |
+
if self.training and self.sample_drop_ratio > 0.1:
|
| 85 |
+
# the overhead is compensated only for a drop path rate larger than 0.1
|
| 86 |
+
x = drop_add_residual_stochastic_depth(
|
| 87 |
+
x,
|
| 88 |
+
residual_func=attn_residual_func,
|
| 89 |
+
sample_drop_ratio=self.sample_drop_ratio,
|
| 90 |
+
pos=pos,
|
| 91 |
+
)
|
| 92 |
+
x = drop_add_residual_stochastic_depth(
|
| 93 |
+
x,
|
| 94 |
+
residual_func=ffn_residual_func,
|
| 95 |
+
sample_drop_ratio=self.sample_drop_ratio,
|
| 96 |
+
)
|
| 97 |
+
elif self.training and self.sample_drop_ratio > 0.0:
|
| 98 |
+
x = x + self.drop_path1(attn_residual_func(x, pos=pos, attn_mask=attn_mask))
|
| 99 |
+
x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
|
| 100 |
+
else:
|
| 101 |
+
x = x + attn_residual_func(x, pos=pos, attn_mask=attn_mask)
|
| 102 |
+
x = x + ffn_residual_func(x)
|
| 103 |
+
return x
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def drop_add_residual_stochastic_depth(
|
| 107 |
+
x: Tensor,
|
| 108 |
+
residual_func: Callable[[Tensor], Tensor],
|
| 109 |
+
sample_drop_ratio: float = 0.0,
|
| 110 |
+
pos: Optional[Tensor] = None,
|
| 111 |
+
) -> Tensor:
|
| 112 |
+
# 1) extract subset using permutation
|
| 113 |
+
b, n, d = x.shape
|
| 114 |
+
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
| 115 |
+
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
| 116 |
+
x_subset = x[brange]
|
| 117 |
+
|
| 118 |
+
# 2) apply residual_func to get residual
|
| 119 |
+
if pos is not None:
|
| 120 |
+
# if necessary, apply rope to the subset
|
| 121 |
+
pos = pos[brange]
|
| 122 |
+
residual = residual_func(x_subset, pos=pos)
|
| 123 |
+
else:
|
| 124 |
+
residual = residual_func(x_subset)
|
| 125 |
+
|
| 126 |
+
x_flat = x.flatten(1)
|
| 127 |
+
residual = residual.flatten(1)
|
| 128 |
+
|
| 129 |
+
residual_scale_factor = b / sample_subset_size
|
| 130 |
+
|
| 131 |
+
# 3) add the residual
|
| 132 |
+
x_plus_residual = torch.index_add(
|
| 133 |
+
x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor
|
| 134 |
+
)
|
| 135 |
+
return x_plus_residual.view_as(x)
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def get_branges_scales(x, sample_drop_ratio=0.0):
|
| 139 |
+
b, n, d = x.shape
|
| 140 |
+
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
| 141 |
+
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
| 142 |
+
residual_scale_factor = b / sample_subset_size
|
| 143 |
+
return brange, residual_scale_factor
|
depth_anything_3/model/dinov2/layers/drop_path.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# References:
|
| 8 |
+
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
| 9 |
+
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
from torch import nn
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
|
| 16 |
+
if drop_prob == 0.0 or not training:
|
| 17 |
+
return x
|
| 18 |
+
keep_prob = 1 - drop_prob
|
| 19 |
+
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
|
| 20 |
+
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
| 21 |
+
if keep_prob > 0.0:
|
| 22 |
+
random_tensor.div_(keep_prob)
|
| 23 |
+
output = x * random_tensor
|
| 24 |
+
return output
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class DropPath(nn.Module):
|
| 28 |
+
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
| 29 |
+
|
| 30 |
+
def __init__(self, drop_prob=None):
|
| 31 |
+
super().__init__()
|
| 32 |
+
self.drop_prob = drop_prob
|
| 33 |
+
|
| 34 |
+
def forward(self, x):
|
| 35 |
+
return drop_path(x, self.drop_prob, self.training)
|
depth_anything_3/model/dinov2/layers/layer_scale.py
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 # noqa: E501
|
| 8 |
+
|
| 9 |
+
from typing import Union
|
| 10 |
+
import torch
|
| 11 |
+
from torch import Tensor, nn
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class LayerScale(nn.Module):
|
| 15 |
+
def __init__(
|
| 16 |
+
self,
|
| 17 |
+
dim: int,
|
| 18 |
+
init_values: Union[float, Tensor] = 1e-5,
|
| 19 |
+
inplace: bool = False,
|
| 20 |
+
) -> None:
|
| 21 |
+
super().__init__()
|
| 22 |
+
self.dim = dim
|
| 23 |
+
self.inplace = inplace
|
| 24 |
+
self.init_values = init_values
|
| 25 |
+
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
| 26 |
+
|
| 27 |
+
def forward(self, x: Tensor) -> Tensor:
|
| 28 |
+
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|
| 29 |
+
|
| 30 |
+
def extra_repr(self) -> str:
|
| 31 |
+
return f"{self.dim}, init_values={self.init_values}, inplace={self.inplace}"
|
depth_anything_3/model/dinov2/layers/mlp.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# References:
|
| 8 |
+
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
| 9 |
+
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
from typing import Callable, Optional
|
| 13 |
+
from torch import Tensor, nn
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class Mlp(nn.Module):
|
| 17 |
+
def __init__(
|
| 18 |
+
self,
|
| 19 |
+
in_features: int,
|
| 20 |
+
hidden_features: Optional[int] = None,
|
| 21 |
+
out_features: Optional[int] = None,
|
| 22 |
+
act_layer: Callable[..., nn.Module] = nn.GELU,
|
| 23 |
+
drop: float = 0.0,
|
| 24 |
+
bias: bool = True,
|
| 25 |
+
) -> None:
|
| 26 |
+
super().__init__()
|
| 27 |
+
out_features = out_features or in_features
|
| 28 |
+
hidden_features = hidden_features or in_features
|
| 29 |
+
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
|
| 30 |
+
self.act = act_layer()
|
| 31 |
+
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
|
| 32 |
+
self.drop = nn.Dropout(drop)
|
| 33 |
+
|
| 34 |
+
def forward(self, x: Tensor) -> Tensor:
|
| 35 |
+
x = self.fc1(x)
|
| 36 |
+
x = self.act(x)
|
| 37 |
+
x = self.drop(x)
|
| 38 |
+
x = self.fc2(x)
|
| 39 |
+
x = self.drop(x)
|
| 40 |
+
return x
|
depth_anything_3/model/dinov2/layers/patch_embed.py
ADDED
|
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
# References:
|
| 8 |
+
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
| 9 |
+
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
| 10 |
+
|
| 11 |
+
from typing import Callable, Optional, Tuple, Union
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
from torch import Tensor
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def make_2tuple(x):
|
| 17 |
+
if isinstance(x, tuple):
|
| 18 |
+
assert len(x) == 2
|
| 19 |
+
return x
|
| 20 |
+
|
| 21 |
+
assert isinstance(x, int)
|
| 22 |
+
return (x, x)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class PatchEmbed(nn.Module):
|
| 26 |
+
"""
|
| 27 |
+
2D image to patch embedding: (B,C,H,W) -> (B,N,D)
|
| 28 |
+
|
| 29 |
+
Args:
|
| 30 |
+
img_size: Image size.
|
| 31 |
+
patch_size: Patch token size.
|
| 32 |
+
in_chans: Number of input image channels.
|
| 33 |
+
embed_dim: Number of linear projection output channels.
|
| 34 |
+
norm_layer: Normalization layer.
|
| 35 |
+
"""
|
| 36 |
+
|
| 37 |
+
def __init__(
|
| 38 |
+
self,
|
| 39 |
+
img_size: Union[int, Tuple[int, int]] = 224,
|
| 40 |
+
patch_size: Union[int, Tuple[int, int]] = 16,
|
| 41 |
+
in_chans: int = 3,
|
| 42 |
+
embed_dim: int = 768,
|
| 43 |
+
norm_layer: Optional[Callable] = None,
|
| 44 |
+
flatten_embedding: bool = True,
|
| 45 |
+
) -> None:
|
| 46 |
+
super().__init__()
|
| 47 |
+
|
| 48 |
+
image_HW = make_2tuple(img_size)
|
| 49 |
+
patch_HW = make_2tuple(patch_size)
|
| 50 |
+
patch_grid_size = (
|
| 51 |
+
image_HW[0] // patch_HW[0],
|
| 52 |
+
image_HW[1] // patch_HW[1],
|
| 53 |
+
)
|
| 54 |
+
|
| 55 |
+
self.img_size = image_HW
|
| 56 |
+
self.patch_size = patch_HW
|
| 57 |
+
self.patches_resolution = patch_grid_size
|
| 58 |
+
self.num_patches = patch_grid_size[0] * patch_grid_size[1]
|
| 59 |
+
|
| 60 |
+
self.in_chans = in_chans
|
| 61 |
+
self.embed_dim = embed_dim
|
| 62 |
+
|
| 63 |
+
self.flatten_embedding = flatten_embedding
|
| 64 |
+
|
| 65 |
+
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
|
| 66 |
+
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
| 67 |
+
|
| 68 |
+
def forward(self, x: Tensor) -> Tensor:
|
| 69 |
+
_, _, H, W = x.shape
|
| 70 |
+
patch_H, patch_W = self.patch_size
|
| 71 |
+
|
| 72 |
+
assert (
|
| 73 |
+
H % patch_H == 0
|
| 74 |
+
), f"Input image height {H} is not a multiple of patch height {patch_H}"
|
| 75 |
+
assert (
|
| 76 |
+
W % patch_W == 0
|
| 77 |
+
), f"Input image width {W} is not a multiple of patch width: {patch_W}"
|
| 78 |
+
|
| 79 |
+
x = self.proj(x) # B C H W
|
| 80 |
+
H, W = x.size(2), x.size(3)
|
| 81 |
+
x = x.flatten(2).transpose(1, 2) # B HW C
|
| 82 |
+
x = self.norm(x)
|
| 83 |
+
if not self.flatten_embedding:
|
| 84 |
+
x = x.reshape(-1, H, W, self.embed_dim) # B H W C
|
| 85 |
+
return x
|
| 86 |
+
|
| 87 |
+
def flops(self) -> float:
|
| 88 |
+
Ho, Wo = self.patches_resolution
|
| 89 |
+
flops = (
|
| 90 |
+
Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
|
| 91 |
+
)
|
| 92 |
+
if self.norm is not None:
|
| 93 |
+
flops += Ho * Wo * self.embed_dim
|
| 94 |
+
return flops
|
depth_anything_3/model/dinov2/layers/rope.py
ADDED
|
@@ -0,0 +1,200 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
#
|
| 3 |
+
# This source code is licensed under the Apache License, Version 2.0
|
| 4 |
+
# found in the LICENSE file in the root directory of this source tree.
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
# Implementation of 2D Rotary Position Embeddings (RoPE).
|
| 8 |
+
|
| 9 |
+
# This module provides a clean implementation of 2D Rotary Position Embeddings,
|
| 10 |
+
# which extends the original RoPE concept to handle 2D spatial positions.
|
| 11 |
+
|
| 12 |
+
# Inspired by:
|
| 13 |
+
# https://github.com/meta-llama/codellama/blob/main/llama/model.py
|
| 14 |
+
# https://github.com/naver-ai/rope-vit
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
from typing import Dict, Tuple
|
| 18 |
+
import torch
|
| 19 |
+
import torch.nn as nn
|
| 20 |
+
import torch.nn.functional as F
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class PositionGetter:
|
| 24 |
+
"""Generates and caches 2D spatial positions for patches in a grid.
|
| 25 |
+
|
| 26 |
+
This class efficiently manages the generation of spatial coordinates for patches
|
| 27 |
+
in a 2D grid, caching results to avoid redundant computations.
|
| 28 |
+
|
| 29 |
+
Attributes:
|
| 30 |
+
position_cache: Dictionary storing precomputed position tensors for different
|
| 31 |
+
grid dimensions.
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
def __init__(self):
|
| 35 |
+
"""Initializes the position generator with an empty cache."""
|
| 36 |
+
self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {}
|
| 37 |
+
|
| 38 |
+
def __call__(
|
| 39 |
+
self, batch_size: int, height: int, width: int, device: torch.device
|
| 40 |
+
) -> torch.Tensor:
|
| 41 |
+
"""Generates spatial positions for a batch of patches.
|
| 42 |
+
|
| 43 |
+
Args:
|
| 44 |
+
batch_size: Number of samples in the batch.
|
| 45 |
+
height: Height of the grid in patches.
|
| 46 |
+
width: Width of the grid in patches.
|
| 47 |
+
device: Target device for the position tensor.
|
| 48 |
+
|
| 49 |
+
Returns:
|
| 50 |
+
Tensor of shape (batch_size, height*width, 2) containing y,x coordinates
|
| 51 |
+
for each position in the grid, repeated for each batch item.
|
| 52 |
+
"""
|
| 53 |
+
if (height, width) not in self.position_cache:
|
| 54 |
+
y_coords = torch.arange(height, device=device)
|
| 55 |
+
x_coords = torch.arange(width, device=device)
|
| 56 |
+
positions = torch.cartesian_prod(y_coords, x_coords)
|
| 57 |
+
self.position_cache[height, width] = positions
|
| 58 |
+
|
| 59 |
+
cached_positions = self.position_cache[height, width]
|
| 60 |
+
return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone()
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
class RotaryPositionEmbedding2D(nn.Module):
|
| 64 |
+
"""2D Rotary Position Embedding implementation.
|
| 65 |
+
|
| 66 |
+
This module applies rotary position embeddings to input tokens based on their
|
| 67 |
+
2D spatial positions. It handles the position-dependent rotation of features
|
| 68 |
+
separately for vertical and horizontal dimensions.
|
| 69 |
+
|
| 70 |
+
Args:
|
| 71 |
+
frequency: Base frequency for the position embeddings. Default: 100.0
|
| 72 |
+
scaling_factor: Scaling factor for frequency computation. Default: 1.0
|
| 73 |
+
|
| 74 |
+
Attributes:
|
| 75 |
+
base_frequency: Base frequency for computing position embeddings.
|
| 76 |
+
scaling_factor: Factor to scale the computed frequencies.
|
| 77 |
+
frequency_cache: Cache for storing precomputed frequency components.
|
| 78 |
+
"""
|
| 79 |
+
|
| 80 |
+
def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0):
|
| 81 |
+
"""Initializes the 2D RoPE module."""
|
| 82 |
+
super().__init__()
|
| 83 |
+
self.base_frequency = frequency
|
| 84 |
+
self.scaling_factor = scaling_factor
|
| 85 |
+
self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {}
|
| 86 |
+
|
| 87 |
+
def _compute_frequency_components(
|
| 88 |
+
self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype
|
| 89 |
+
) -> Tuple[torch.Tensor, torch.Tensor]:
|
| 90 |
+
"""Computes frequency components for rotary embeddings.
|
| 91 |
+
|
| 92 |
+
Args:
|
| 93 |
+
dim: Feature dimension (must be even).
|
| 94 |
+
seq_len: Maximum sequence length.
|
| 95 |
+
device: Target device for computations.
|
| 96 |
+
dtype: Data type for the computed tensors.
|
| 97 |
+
|
| 98 |
+
Returns:
|
| 99 |
+
Tuple of (cosine, sine) tensors for frequency components.
|
| 100 |
+
"""
|
| 101 |
+
cache_key = (dim, seq_len, device, dtype)
|
| 102 |
+
if cache_key not in self.frequency_cache:
|
| 103 |
+
# Compute frequency bands
|
| 104 |
+
exponents = torch.arange(0, dim, 2, device=device).float() / dim
|
| 105 |
+
inv_freq = 1.0 / (self.base_frequency**exponents)
|
| 106 |
+
|
| 107 |
+
# Generate position-dependent frequencies
|
| 108 |
+
positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
|
| 109 |
+
angles = torch.einsum("i,j->ij", positions, inv_freq)
|
| 110 |
+
|
| 111 |
+
# Compute and cache frequency components
|
| 112 |
+
angles = angles.to(dtype)
|
| 113 |
+
angles = torch.cat((angles, angles), dim=-1)
|
| 114 |
+
cos_components = angles.cos().to(dtype)
|
| 115 |
+
sin_components = angles.sin().to(dtype)
|
| 116 |
+
self.frequency_cache[cache_key] = (cos_components, sin_components)
|
| 117 |
+
|
| 118 |
+
return self.frequency_cache[cache_key]
|
| 119 |
+
|
| 120 |
+
@staticmethod
|
| 121 |
+
def _rotate_features(x: torch.Tensor) -> torch.Tensor:
|
| 122 |
+
"""Performs feature rotation by splitting and recombining feature dimensions.
|
| 123 |
+
|
| 124 |
+
Args:
|
| 125 |
+
x: Input tensor to rotate.
|
| 126 |
+
|
| 127 |
+
Returns:
|
| 128 |
+
Rotated feature tensor.
|
| 129 |
+
"""
|
| 130 |
+
feature_dim = x.shape[-1]
|
| 131 |
+
x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :]
|
| 132 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 133 |
+
|
| 134 |
+
def _apply_1d_rope(
|
| 135 |
+
self,
|
| 136 |
+
tokens: torch.Tensor,
|
| 137 |
+
positions: torch.Tensor,
|
| 138 |
+
cos_comp: torch.Tensor,
|
| 139 |
+
sin_comp: torch.Tensor,
|
| 140 |
+
) -> torch.Tensor:
|
| 141 |
+
"""Applies 1D rotary position embeddings along one dimension.
|
| 142 |
+
|
| 143 |
+
Args:
|
| 144 |
+
tokens: Input token features.
|
| 145 |
+
positions: Position indices.
|
| 146 |
+
cos_comp: Cosine components for rotation.
|
| 147 |
+
sin_comp: Sine components for rotation.
|
| 148 |
+
|
| 149 |
+
Returns:
|
| 150 |
+
Tokens with applied rotary position embeddings.
|
| 151 |
+
"""
|
| 152 |
+
# Embed positions with frequency components
|
| 153 |
+
cos = F.embedding(positions, cos_comp)[:, None, :, :]
|
| 154 |
+
sin = F.embedding(positions, sin_comp)[:, None, :, :]
|
| 155 |
+
# Apply rotation
|
| 156 |
+
return (tokens * cos) + (self._rotate_features(tokens) * sin)
|
| 157 |
+
|
| 158 |
+
def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
| 159 |
+
"""Applies 2D rotary position embeddings to input tokens.
|
| 160 |
+
|
| 161 |
+
Args:
|
| 162 |
+
tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim).
|
| 163 |
+
The feature dimension (dim) must be divisible by 4.
|
| 164 |
+
positions: Position tensor of shape (batch_size, n_tokens, 2) containing
|
| 165 |
+
the y and x coordinates for each token.
|
| 166 |
+
|
| 167 |
+
Returns:
|
| 168 |
+
Tensor of same shape as input with applied 2D rotary position embeddings.
|
| 169 |
+
|
| 170 |
+
Raises:
|
| 171 |
+
AssertionError: If input dimensions are invalid or positions are malformed.
|
| 172 |
+
"""
|
| 173 |
+
# Validate inputs
|
| 174 |
+
assert tokens.size(-1) % 2 == 0, "Feature dimension must be even"
|
| 175 |
+
assert (
|
| 176 |
+
positions.ndim == 3 and positions.shape[-1] == 2
|
| 177 |
+
), "Positions must have shape (batch_size, n_tokens, 2)"
|
| 178 |
+
|
| 179 |
+
# Compute feature dimension for each spatial direction
|
| 180 |
+
feature_dim = tokens.size(-1) // 2
|
| 181 |
+
|
| 182 |
+
# Get frequency components
|
| 183 |
+
max_position = int(positions.max()) + 1
|
| 184 |
+
cos_comp, sin_comp = self._compute_frequency_components(
|
| 185 |
+
feature_dim, max_position, tokens.device, tokens.dtype
|
| 186 |
+
)
|
| 187 |
+
|
| 188 |
+
# Split features for vertical and horizontal processing
|
| 189 |
+
vertical_features, horizontal_features = tokens.chunk(2, dim=-1)
|
| 190 |
+
|
| 191 |
+
# Apply RoPE separately for each dimension
|
| 192 |
+
vertical_features = self._apply_1d_rope(
|
| 193 |
+
vertical_features, positions[..., 0], cos_comp, sin_comp
|
| 194 |
+
)
|
| 195 |
+
horizontal_features = self._apply_1d_rope(
|
| 196 |
+
horizontal_features, positions[..., 1], cos_comp, sin_comp
|
| 197 |
+
)
|
| 198 |
+
|
| 199 |
+
# Combine processed features
|
| 200 |
+
return torch.cat((vertical_features, horizontal_features), dim=-1)
|
depth_anything_3/model/dinov2/layers/swiglu_ffn.py
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
from typing import Callable, Optional
|
| 8 |
+
import torch.nn.functional as F
|
| 9 |
+
from torch import Tensor, nn
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class SwiGLUFFN(nn.Module):
|
| 13 |
+
def __init__(
|
| 14 |
+
self,
|
| 15 |
+
in_features: int,
|
| 16 |
+
hidden_features: Optional[int] = None,
|
| 17 |
+
out_features: Optional[int] = None,
|
| 18 |
+
act_layer: Callable[..., nn.Module] = None,
|
| 19 |
+
drop: float = 0.0,
|
| 20 |
+
bias: bool = True,
|
| 21 |
+
) -> None:
|
| 22 |
+
super().__init__()
|
| 23 |
+
out_features = out_features or in_features
|
| 24 |
+
hidden_features = hidden_features or in_features
|
| 25 |
+
self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
|
| 26 |
+
self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
|
| 27 |
+
|
| 28 |
+
def forward(self, x: Tensor) -> Tensor:
|
| 29 |
+
x12 = self.w12(x)
|
| 30 |
+
x1, x2 = x12.chunk(2, dim=-1)
|
| 31 |
+
hidden = F.silu(x1) * x2
|
| 32 |
+
return self.w3(hidden)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
try:
|
| 36 |
+
from xformers.ops import SwiGLU
|
| 37 |
+
|
| 38 |
+
XFORMERS_AVAILABLE = True
|
| 39 |
+
except ImportError:
|
| 40 |
+
SwiGLU = SwiGLUFFN
|
| 41 |
+
XFORMERS_AVAILABLE = False
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
class SwiGLUFFNFused(SwiGLU):
|
| 45 |
+
def __init__(
|
| 46 |
+
self,
|
| 47 |
+
in_features: int,
|
| 48 |
+
hidden_features: Optional[int] = None,
|
| 49 |
+
out_features: Optional[int] = None,
|
| 50 |
+
act_layer: Callable[..., nn.Module] = None,
|
| 51 |
+
drop: float = 0.0,
|
| 52 |
+
bias: bool = True,
|
| 53 |
+
) -> None:
|
| 54 |
+
out_features = out_features or in_features
|
| 55 |
+
hidden_features = hidden_features or in_features
|
| 56 |
+
hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
|
| 57 |
+
super().__init__(
|
| 58 |
+
in_features=in_features,
|
| 59 |
+
hidden_features=hidden_features,
|
| 60 |
+
out_features=out_features,
|
| 61 |
+
bias=bias,
|
| 62 |
+
)
|