mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-12 01:07:30 +08:00
682 lines
31 KiB
Python
682 lines
31 KiB
Python
"""ComfyUI nodes for Depth Anything 3.
|
|
Model capability matrix:
|
|
|
|
Variant head_type has_sky has_conf cam_dec
|
|
DA3-Small dualdpt False True yes
|
|
DA3-Base dualdpt False True yes
|
|
DA3-Mono-Large dpt True False no
|
|
DA3-Metric-Large dpt True False no (raw output is metres)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing_extensions import override
|
|
|
|
import torch
|
|
|
|
import comfy.model_management as mm
|
|
import comfy.sd
|
|
import folder_paths
|
|
from comfy.ldm.colormap import turbo as _turbo
|
|
from comfy.ldm.depth_anything_3 import preprocess as da3_preprocess
|
|
from comfy_api.latest import ComfyExtension, Types, io
|
|
from comfy.ldm.moge.geometry import triangulate_grid_mesh
|
|
|
|
DA3ModelType = io.Custom("DA3_MODEL")
|
|
DA3Geometry = io.Custom("DA3_GEOMETRY")
|
|
DA3PointCloud = io.Custom("DA3_POINT_CLOUD")
|
|
|
|
# DA3_GEOMETRY is a dict with these optional keys (absent when the upstream model didn't produce them):
|
|
#
|
|
# Per-frame tensors - B = batch size in mono mode; B = S (number of views) in multi-view mode.
|
|
# "depth": torch.Tensor (B, H, W) -- raw model depth (always present; matches MoGe convention)
|
|
# "image": torch.Tensor (B, H, W, 3) -- source image in [0, 1], CPU (always present)
|
|
# "mode": str -- "mono" or "multiview" (always present)
|
|
# "sky": torch.Tensor (B, H, W) -- sky probability in [0, 1] (Mono/Metric variants only)
|
|
# "confidence": torch.Tensor (B, H, W) -- raw model confidence output (Small/Base variants only)
|
|
#
|
|
# Multi-view only - S = number of views; the leading 1 is the scene dimension from the model.
|
|
# "extrinsics": torch.Tensor (1, S, 3, 4) -- world-to-camera [R|t] matrices
|
|
# "intrinsics": torch.Tensor (1, S, 3, 3) -- pixel-space intrinsics
|
|
#
|
|
# DA3_POINT_CLOUD is a dict:
|
|
# "points": torch.Tensor (N, 3) -- 3-D coords in glTF convention (Y-up, Z-back)
|
|
# "colors": torch.Tensor (N, 3) -- RGB in [0, 1], or None
|
|
# "confidence": torch.Tensor (N,) -- raw confidence per point, or None
|
|
|
|
|
|
def _da3_unproject(depth: torch.Tensor, K: torch.Tensor) -> torch.Tensor:
|
|
"""Pixel-space K⁻¹ unprojection: (H,W) depth → (H,W,3) point map in OpenCV space."""
|
|
H, W = depth.shape
|
|
u = torch.arange(W, dtype=torch.float32, device=depth.device)
|
|
v = torch.arange(H, dtype=torch.float32, device=depth.device)
|
|
u, v = torch.meshgrid(u, v, indexing='xy') # both (H, W)
|
|
pix = torch.stack([u, v, torch.ones_like(u)], dim=-1) # (H, W, 3)
|
|
rays = torch.einsum('ij,hwj->hwi', torch.linalg.inv(K.to(depth.device)), pix)
|
|
return rays * depth.unsqueeze(-1) # (H, W, 3)
|
|
|
|
|
|
def _da3_default_K(H: int, W: int) -> torch.Tensor:
|
|
"""Fallback ~60° FOV pinhole K for mono-mode DA3 (no intrinsics in geometry)."""
|
|
fx = fy = float(W) * 0.7
|
|
return torch.tensor([[fx, 0.0, (W - 1) / 2.0],
|
|
[0.0, fy, (H - 1) / 2.0],
|
|
[0.0, 0.0, 1.0]], dtype=torch.float32)
|
|
|
|
|
|
def _da3_get_K(geometry: dict, b: int, H: int, W: int) -> torch.Tensor:
|
|
"""Return pixel-space K for batch element b, falling back to a default estimate."""
|
|
if "intrinsics" in geometry:
|
|
# shape (1, S, 3, 3) - leading scene dimension from the multiview head
|
|
return geometry["intrinsics"][0, b].float()
|
|
logging.getLogger("comfy").warning(
|
|
"DA3_GEOMETRY has no intrinsics (mono-mode model). "
|
|
"Using a ~60° FOV estimate; 3-D reconstruction may be inaccurate."
|
|
)
|
|
return _da3_default_K(H, W)
|
|
|
|
|
|
def _da3_get_extrinsic(geometry: dict, b: int) -> torch.Tensor | None:
|
|
"""Return the world-to-camera extrinsic for batch element b, or None in mono mode.
|
|
|
|
The model outputs (1, S, 3, 4) [R|t] matrices; the fallback identity is (4, 4).
|
|
_da3_apply_extrinsic handles both shapes via [:3, :3] / [:3, 3] slicing.
|
|
"""
|
|
if "extrinsics" not in geometry:
|
|
return None
|
|
return geometry["extrinsics"][0, b].float()
|
|
|
|
|
|
def _da3_apply_extrinsic(points_cam: torch.Tensor, E: torch.Tensor) -> torch.Tensor:
|
|
"""Transform (H,W,3) OpenCV camera-space points to world space."""
|
|
E = E.to(points_cam.device).float()
|
|
if not torch.isfinite(E).all():
|
|
logging.getLogger("comfy").warning(
|
|
"DA3 extrinsic matrix contains non-finite values (pose estimation may have failed). "
|
|
"Falling back to camera-space coordinates."
|
|
)
|
|
return points_cam
|
|
H, W, _ = points_cam.shape
|
|
R = E[:3, :3] # (3, 3) rotation
|
|
t = E[:3, 3] # (3,) translation
|
|
R_inv = R.T # rotation inverse = transpose for orthogonal R
|
|
t_inv = -(R_inv @ t) # (3,)
|
|
pts = points_cam.reshape(-1, 3) # (N, 3)
|
|
pts_world = pts @ R_inv.T + t_inv # (N, 3)
|
|
return pts_world.reshape(H, W, 3)
|
|
|
|
|
|
def _normalize_confidence(conf: torch.Tensor) -> torch.Tensor:
|
|
"""Map raw confidence to [0, 1] per image."""
|
|
B = conf.shape[0]
|
|
out = []
|
|
for i in range(B):
|
|
c = conf[i]
|
|
c_min, c_max = c.min(), c.max()
|
|
out.append((c - c_min) / (c_max - c_min) if c_max > c_min else torch.ones_like(c))
|
|
return torch.stack(out, dim=0)
|
|
|
|
|
|
def _da3_build_mask(geometry: dict, b: int, H: int, W: int, confidence_threshold: float, use_sky_mask: bool) -> torch.Tensor:
|
|
"""Build (H,W) bool keep-mask from sky probability and confidence."""
|
|
mask = torch.ones(H, W, dtype=torch.bool)
|
|
if use_sky_mask and "sky" in geometry:
|
|
mask = mask & (geometry["sky"][b] < 0.5)
|
|
if "confidence" in geometry and confidence_threshold > 0.0:
|
|
conf_norm = _normalize_confidence(geometry["confidence"][b:b + 1])[0]
|
|
mask = mask & (conf_norm >= confidence_threshold)
|
|
return mask
|
|
|
|
|
|
class LoadDA3Model(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LoadDA3Model",
|
|
display_name="Load Depth Anything 3",
|
|
category="model/loaders",
|
|
inputs=[
|
|
io.Combo.Input(
|
|
"model_name",
|
|
options=folder_paths.get_filename_list("geometry_estimation"),
|
|
),
|
|
io.Combo.Input(
|
|
"weight_dtype",
|
|
options=["default", "fp16", "bf16", "fp32"],
|
|
default="default",
|
|
),
|
|
],
|
|
outputs=[DA3ModelType.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_name, weight_dtype) -> io.NodeOutput:
|
|
model_options = {}
|
|
if weight_dtype == "fp16":
|
|
model_options["dtype"] = torch.float16
|
|
elif weight_dtype == "bf16":
|
|
model_options["dtype"] = torch.bfloat16
|
|
elif weight_dtype == "fp32":
|
|
model_options["dtype"] = torch.float32
|
|
|
|
path = folder_paths.get_full_path_or_raise("geometry_estimation", model_name)
|
|
model = comfy.sd.load_diffusion_model(path, model_options=model_options)
|
|
return io.NodeOutput(model)
|
|
|
|
|
|
def _run_da3(model_patcher, image: torch.Tensor, process_res: int, method: str = "upper_bound_resize"):
|
|
"""Run DA3 on (B,H,W,3), returns depth/conf/sky at original resolution (or None)."""
|
|
assert image.ndim == 4 and image.shape[-1] == 3, f"expected (B,H,W,3) IMAGE; got {tuple(image.shape)}"
|
|
|
|
B, H, W, _ = image.shape
|
|
mm.load_model_gpu(model_patcher)
|
|
diffusion = model_patcher.model.diffusion_model
|
|
device = mm.get_torch_device()
|
|
dtype = diffusion.dtype if diffusion.dtype is not None else torch.float32
|
|
|
|
depths, confs, skies = [], [], []
|
|
for i in range(B):
|
|
single = image[i:i + 1].to(device)
|
|
x = da3_preprocess.preprocess_image(single, process_res=process_res, method=method)
|
|
x = x.to(dtype=dtype)
|
|
with torch.no_grad():
|
|
out = diffusion(x)
|
|
|
|
depth_lr = out["depth"]
|
|
depth_full = torch.nn.functional.interpolate(
|
|
depth_lr.unsqueeze(1).float(), size=(H, W),
|
|
mode="bilinear", align_corners=False,
|
|
).squeeze(1).cpu()
|
|
depths.append(depth_full)
|
|
|
|
if "depth_conf" in out:
|
|
conf_full = torch.nn.functional.interpolate(
|
|
out["depth_conf"].unsqueeze(1).float(), size=(H, W),
|
|
mode="bilinear", align_corners=False,
|
|
).squeeze(1).cpu()
|
|
confs.append(conf_full)
|
|
if "sky" in out:
|
|
sky_full = torch.nn.functional.interpolate(
|
|
out["sky"].unsqueeze(1).float(), size=(H, W),
|
|
mode="bilinear", align_corners=False,
|
|
).squeeze(1).cpu()
|
|
skies.append(sky_full)
|
|
|
|
depth = torch.cat(depths, dim=0)
|
|
confidence = torch.cat(confs, dim=0) if confs else None
|
|
sky = torch.cat(skies, dim=0) if skies else None
|
|
return depth, confidence, sky
|
|
|
|
|
|
class DA3Inference(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DA3Inference",
|
|
search_aliases=["depth", "geometry", "da3", "depth anything", "monocular", "pointmap", "sky", "3d", "metric depth", "disparity"],
|
|
display_name="Run Depth Anything 3",
|
|
category="image/geometry estimation",
|
|
description="Run Depth Anything 3 on an image. In multi-view mode each image is treated as a separate view of the same scene.",
|
|
inputs=[
|
|
DA3ModelType.Input("da3_model"),
|
|
io.Image.Input("image"),
|
|
io.Int.Input("resolution", default=504, min=140, max=2520, step=14,
|
|
tooltip="Resolution the model runs at (longest side, multiple of 14).\n"
|
|
"Lower = faster / less VRAM.\n"
|
|
"Higher = more detail.\n"
|
|
"Output is upsampled back to the original size."),
|
|
io.Combo.Input("resize_method", options=["upper_bound_resize", "lower_bound_resize"], default="upper_bound_resize",
|
|
tooltip="upper_bound_resize: scale so the longest side = resolution (caps memory, default).\n"
|
|
"lower_bound_resize: scale so the shortest side = resolution (preserves more detail on tall/wide images, uses more memory)."),
|
|
io.DynamicCombo.Input("mode", tooltip="mono: single view image (works with any model variant).\n"
|
|
"multiview: all images processed together for geometric consistency + camera pose (for Small/Base models only).",
|
|
options=[
|
|
io.DynamicCombo.Option("mono", []),
|
|
io.DynamicCombo.Option("multiview", [
|
|
io.Combo.Input("ref_view_strategy", options=["saddle_balanced", "saddle_sim_range", "first", "middle"], default="saddle_balanced",
|
|
tooltip="Which view acts as the geometric anchor.\n"
|
|
"- saddle_balanced: the view most 'average' across all others (best general choice).\n"
|
|
"- saddle_sim_range: the view most visually distinct from the others.\n"
|
|
"- first / middle: fixed positional picks."),
|
|
io.Combo.Input("pose_method", options=["cam_dec", "ray_pose"], default="cam_dec",
|
|
tooltip="How the camera field-of-view is estimated (for Small/Base models only).\n"
|
|
"- cam_dec: learned from image features.\n"
|
|
"- ray_pose: derived geometrically from the model's 3D ray output.\n"
|
|
"Affects perspective correctness of the 3D output. Try both if results look distorted."),
|
|
]),
|
|
]),
|
|
],
|
|
outputs=[
|
|
DA3Geometry.Output("da3_geometry", tooltip="Dictionary of non-normalized tensors.\n"
|
|
"Always has the keys: depth, image, mode.\n"
|
|
"Optional keys: sky (for Mono/Metric), confidence (for Small/Base), extrinsics + intrinsics (for multi-view)."),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, da3_model, image, resolution, resize_method, mode) -> io.NodeOutput:
|
|
mode_val = mode["mode"] # "mono" or "multiview"
|
|
|
|
if mode_val == "mono":
|
|
return cls._execute_mono(da3_model, image, resolution, resize_method)
|
|
|
|
# Capability checks for multi-view mode.
|
|
diffusion = da3_model.model.diffusion_model
|
|
pose_method = mode["pose_method"]
|
|
ref_view_strategy = mode["ref_view_strategy"]
|
|
|
|
has_cam_dec = diffusion.cam_dec is not None
|
|
has_dualdpt = diffusion.head_type == "dualdpt"
|
|
|
|
if not has_cam_dec and not has_dualdpt:
|
|
raise ValueError(
|
|
"multi-view mode requires Small or Base model. The loaded model "
|
|
f"(head_type='{diffusion.head_type}') does not support cross-view "
|
|
"attention or camera pose estimation. Switch mode to 'mono', or "
|
|
"load Small or Base model for mult-view."
|
|
)
|
|
|
|
if pose_method == "cam_dec" and not has_cam_dec:
|
|
raise ValueError(
|
|
"pose_method='cam_dec' requires a camera decoder, but the loaded "
|
|
f"model (head_type='{diffusion.head_type}') does not have one. "
|
|
"Use pose_method='ray_pose' instead."
|
|
)
|
|
if pose_method == "ray_pose" and not has_dualdpt:
|
|
raise ValueError(
|
|
"pose_method='ray_pose' requires a DualDPT head, but the loaded "
|
|
f"model has a '{diffusion.head_type}' head. "
|
|
"Use pose_method='cam_dec' instead."
|
|
)
|
|
|
|
return cls._execute_multiview(
|
|
da3_model, image, resolution, resize_method,
|
|
ref_view_strategy, pose_method,
|
|
)
|
|
|
|
@classmethod
|
|
def _execute_mono(cls, model, image, resolution, resize_method) -> io.NodeOutput:
|
|
depth, confidence, sky = _run_da3(model, image, resolution, method=resize_method)
|
|
|
|
geometry: dict = {
|
|
"depth": depth.contiguous(),
|
|
"image": image[..., :3].cpu(),
|
|
"mode": "mono",
|
|
}
|
|
if sky is not None:
|
|
geometry["sky"] = sky.contiguous()
|
|
if confidence is not None:
|
|
geometry["confidence"] = confidence.contiguous()
|
|
return io.NodeOutput(geometry)
|
|
|
|
@classmethod
|
|
def _execute_multiview(cls, model, image, resolution, resize_method, ref_view_strategy, pose_method) -> io.NodeOutput:
|
|
assert image.ndim == 4 and image.shape[-1] == 3, \
|
|
f"expected (B,H,W,3) IMAGE; got {tuple(image.shape)}"
|
|
S, H, W, _ = image.shape
|
|
|
|
mm.load_model_gpu(model)
|
|
diffusion = model.model.diffusion_model
|
|
device = mm.get_torch_device()
|
|
dtype = diffusion.dtype if diffusion.dtype is not None else torch.float32
|
|
|
|
# All views in a single forward pass: (1, S, 3, H', W').
|
|
x = image.to(device)
|
|
x = da3_preprocess.preprocess_image(x, process_res=resolution, method=resize_method)
|
|
x = x.to(dtype=dtype).unsqueeze(0)
|
|
|
|
use_ray_pose = (pose_method == "ray_pose")
|
|
with torch.no_grad():
|
|
out = diffusion(x, use_ray_pose=use_ray_pose, ref_view_strategy=ref_view_strategy)
|
|
|
|
depth = torch.nn.functional.interpolate(
|
|
out["depth"].float().unsqueeze(1), size=(H, W),
|
|
mode="bilinear", align_corners=False,
|
|
).squeeze(1).cpu()
|
|
|
|
sky = None
|
|
if "sky" in out:
|
|
sky = torch.nn.functional.interpolate(
|
|
out["sky"].unsqueeze(1).float(), size=(H, W),
|
|
mode="bilinear", align_corners=False,
|
|
).squeeze(1).cpu()
|
|
|
|
if "extrinsics" in out and "intrinsics" in out:
|
|
extrinsics = out["extrinsics"].float().cpu()
|
|
intrinsics = out["intrinsics"].float().cpu()
|
|
else:
|
|
extrinsics = torch.eye(4)[None, None].expand(1, S, 4, 4).clone()
|
|
intrinsics = torch.eye(3)[None, None].expand(1, S, 3, 3).clone()
|
|
|
|
geometry: dict = {
|
|
"depth": depth.contiguous(),
|
|
"image": image[..., :3].cpu(),
|
|
"mode": "multiview",
|
|
"extrinsics": extrinsics.contiguous(),
|
|
"intrinsics": intrinsics.contiguous(),
|
|
}
|
|
if sky is not None:
|
|
geometry["sky"] = sky.contiguous()
|
|
if "depth_conf" in out:
|
|
conf = torch.nn.functional.interpolate(
|
|
out["depth_conf"].unsqueeze(1).float(), size=(H, W),
|
|
mode="bilinear", align_corners=False,
|
|
).squeeze(1).cpu()
|
|
geometry["confidence"] = conf.contiguous()
|
|
return io.NodeOutput(geometry)
|
|
|
|
|
|
class DA3Render(io.ComfyNode):
|
|
"""Render a visualization from a DA3_GEOMETRY packet."""
|
|
|
|
_DEPTH_RENDER_INPUTS = [
|
|
io.Combo.Input("normalization",
|
|
options=["v2_style", "min_max", "raw"],
|
|
default="v2_style",
|
|
tooltip="- v2_style: mean/std normalisation for perceptually balanced results (default).\n"
|
|
"- min_max: stretches the full depth range to [0, 1] for maximum contrast.\n"
|
|
"- raw: no scaling,preserves metric units for Metric model."),
|
|
io.Boolean.Input("apply_sky_clip", default=False,
|
|
tooltip="Clip sky-region depth to the 99th percentile of foreground depth before normalisation. "
|
|
"Requires a sky key in the da3_geometry input (for Mono/Metric models only)."),
|
|
]
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DA3Render",
|
|
display_name="Render Depth Anything 3",
|
|
category="image/geometry estimation",
|
|
description="Render a depth map, confidence map, or sky mask from Depth Anything 3 geometry data.",
|
|
inputs=[
|
|
DA3Geometry.Input("da3_geometry"),
|
|
io.DynamicCombo.Input("output",
|
|
tooltip="- depth: normalised greyscale depth image.\n"
|
|
"- depth_colored: depth mapped through the Turbo colormap.\n"
|
|
"- sky_mask: sky probability in [0, 1] (for Mono/Metric models only).\n"
|
|
"- confidence: normalised depth confidence (for Small/Base models only).",
|
|
options=[
|
|
io.DynamicCombo.Option("depth", cls._DEPTH_RENDER_INPUTS),
|
|
io.DynamicCombo.Option("depth_colored", cls._DEPTH_RENDER_INPUTS),
|
|
io.DynamicCombo.Option("sky_mask", [
|
|
io.Boolean.Input("colored", default=False, tooltip="Apply the Turbo colormap to the sky mask."),
|
|
]),
|
|
io.DynamicCombo.Option("confidence", [
|
|
io.Boolean.Input("colored", default=False, tooltip="Apply the Turbo colormap to the confidence map."),
|
|
]),
|
|
]),
|
|
],
|
|
outputs=[io.Image.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, da3_geometry, output) -> io.NodeOutput:
|
|
output_val = output["output"]
|
|
|
|
if output_val in ("depth", "depth_colored"):
|
|
normalization = output["normalization"]
|
|
apply_sky_clip = output["apply_sky_clip"]
|
|
if apply_sky_clip and "sky" not in da3_geometry:
|
|
raise ValueError(
|
|
"apply_sky_clip=True requires a sky tensor in the da3_geometry input, but none is present. "
|
|
"Run with Mono/Metric models or set apply_sky_clip=False."
|
|
)
|
|
depth = da3_geometry["depth"]
|
|
sky = da3_geometry.get("sky")
|
|
if apply_sky_clip and sky is not None:
|
|
depth = torch.stack([
|
|
da3_preprocess.apply_sky_aware_clip(depth[i], sky[i])
|
|
for i in range(depth.shape[0])
|
|
], dim=0)
|
|
grey = cls._depth_to_image(depth, sky, normalization) # (B,H,W,3) greyscale
|
|
result = _turbo(grey[..., 0]) if output_val == "depth_colored" else grey
|
|
|
|
elif output_val == "sky_mask":
|
|
if "sky" not in da3_geometry:
|
|
raise ValueError("geometry has no sky output; run with Mono/Metric models.")
|
|
sky = da3_geometry["sky"]
|
|
if output["colored"]:
|
|
result = _turbo(sky)
|
|
else:
|
|
result = sky.unsqueeze(-1).expand(*sky.shape, 3).contiguous()
|
|
|
|
elif output_val == "confidence":
|
|
if "confidence" not in da3_geometry:
|
|
raise ValueError("da3_geometry has no confidence output; run with Small/Base models.")
|
|
conf = _normalize_confidence(da3_geometry["confidence"])
|
|
if output["colored"]:
|
|
result = _turbo(conf)
|
|
else:
|
|
result = conf.unsqueeze(-1).expand(*conf.shape, 3).contiguous()
|
|
|
|
else:
|
|
raise ValueError(f"Unknown output mode: {output_val}")
|
|
|
|
return io.NodeOutput(result.float())
|
|
|
|
@staticmethod
|
|
def _depth_to_image(depth: torch.Tensor, sky_for_norm: torch.Tensor | None, normalization: str) -> torch.Tensor:
|
|
"""Normalise depth and pack as an (B,H,W,3) image tensor."""
|
|
|
|
N = depth.shape[0]
|
|
if normalization == "v2_style":
|
|
norm = torch.stack([
|
|
da3_preprocess.normalize_depth_v2_style(
|
|
depth[i], sky_for_norm[i] if sky_for_norm is not None else None)
|
|
for i in range(N)
|
|
], dim=0)
|
|
elif normalization == "min_max":
|
|
norm = da3_preprocess.normalize_depth_min_max(depth)
|
|
else:
|
|
norm = depth
|
|
|
|
out = norm.unsqueeze(-1).repeat(1, 1, 1, 3)
|
|
if normalization != "raw":
|
|
out = out.clamp(0.0, 1.0)
|
|
return out.contiguous()
|
|
|
|
|
|
class DA3GeometryToMesh(io.ComfyNode):
|
|
"""Convert a DA3_GEOMETRY packet into a Types.MESH by unprojecting depth and triangulating."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DA3GeometryToMesh",
|
|
search_aliases=["da3", "depth anything", "mesh", "geometry", "3d", "triangulate"],
|
|
display_name="Convert DA3 Geometry to Mesh",
|
|
category="image/geometry estimation",
|
|
description="Convert a depth map into a triangulated 3D mesh.",
|
|
inputs=[
|
|
DA3Geometry.Input("da3_geometry"),
|
|
io.Int.Input("batch_index", default=0, min=0, max=4096, tooltip="Which image of a batch to convert. Per-image vertex counts differ so batches cannot be stacked."),
|
|
io.Int.Input("decimation", default=1, min=1, max=8, tooltip="Vertex stride. 1 = full resolution, 2 = half, etc."),
|
|
io.Float.Input("discontinuity_threshold", default=0.04, min=0.0, max=1.0, step=0.01, tooltip="Drop triangles whose 3x3 depth span exceeds this fraction. 0 = off."),
|
|
io.Float.Input("confidence_threshold", default=0.1, min=0.0, max=1.0, step=0.01,
|
|
tooltip="Exclude pixels whose per-image normalised confidence is below this value (0 = keep all, 1 = keep only the single most confident pixel). "
|
|
"Used when the geometry has a confidence map (Small/Base models)."),
|
|
io.Boolean.Input("use_sky_mask", default=True, tooltip="Exclude sky-probability pixels (sky >= 0.5) from the mesh. Used when the geometry has a sky map (Mono/Metric models)."),
|
|
io.Boolean.Input("texture", default=True, tooltip="Use the source image as a base color texture."),
|
|
],
|
|
outputs=[io.Mesh.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, da3_geometry, batch_index, decimation, discontinuity_threshold, confidence_threshold, use_sky_mask, texture) -> io.NodeOutput:
|
|
depth_all = da3_geometry["depth"] # (B, H, W)
|
|
B = depth_all.shape[0]
|
|
if batch_index >= B:
|
|
raise ValueError(f"batch_index {batch_index} is out of range; DA3_GEOMETRY has batch size {B}.")
|
|
|
|
depth = depth_all[batch_index] # (H, W)
|
|
H, W = depth.shape
|
|
|
|
# NaN/inf depth would propagate silently through unproject and produce an
|
|
# empty mesh; replace them with 0 here so those pixels are later excluded
|
|
# by the isfinite check inside triangulate_grid_mesh.
|
|
depth = depth.clone()
|
|
n_bad = (~torch.isfinite(depth)).sum().item()
|
|
if n_bad:
|
|
logging.getLogger("comfy").warning(
|
|
f"DA3GeometryToMesh: depth[{batch_index}] has {n_bad} non-finite pixels "
|
|
f"({100*n_bad/(H*W):.1f}%) - zeroed before unproject."
|
|
)
|
|
depth[~torch.isfinite(depth)] = 0.0
|
|
logging.getLogger("comfy").debug(
|
|
f"DA3GeometryToMesh: depth[{batch_index}] range "
|
|
f"[{depth.min():.4g}, {depth.max():.4g}], mean={depth.mean():.4g}"
|
|
)
|
|
|
|
K = _da3_get_K(da3_geometry, batch_index, H, W)
|
|
points = _da3_unproject(depth, K) # (H, W, 3) in OpenCV camera space
|
|
|
|
# Apply world-to-camera inverse so multi-view frames share a common world frame.
|
|
E = _da3_get_extrinsic(da3_geometry, batch_index)
|
|
if E is not None:
|
|
points = _da3_apply_extrinsic(points, E)
|
|
|
|
# Mask invalid pixels by setting them to inf so triangulate_grid_mesh skips them.
|
|
mask = _da3_build_mask(da3_geometry, batch_index, H, W, confidence_threshold, use_sky_mask)
|
|
# Also exclude pixels where depth was invalid.
|
|
mask = mask & (depth_all[batch_index] > 0) & torch.isfinite(depth_all[batch_index])
|
|
points = points.clone()
|
|
points[~mask] = float('inf')
|
|
|
|
verts, faces, uvs = triangulate_grid_mesh(
|
|
points,
|
|
decimation=decimation,
|
|
discontinuity_threshold=discontinuity_threshold,
|
|
depth=depth,
|
|
)
|
|
if verts.shape[0] == 0 or faces.shape[0] == 0:
|
|
raise ValueError(
|
|
"DA3GeometryToMesh produced an empty mesh. "
|
|
"Try raising discontinuity_threshold, lowering confidence_threshold, "
|
|
"or disabling use_sky_mask."
|
|
)
|
|
|
|
# OpenCV (X right, Y down, Z forward) → glTF (X right, Y up, Z back).
|
|
# Same transform as MoGePointMapToMesh perspective branch.
|
|
verts = verts * torch.tensor([1.0, -1.0, -1.0], dtype=verts.dtype)
|
|
faces = faces[:, [0, 2, 1]].contiguous()
|
|
|
|
tex = da3_geometry["image"][batch_index:batch_index + 1] if texture else None
|
|
mesh = Types.MESH(
|
|
vertices=verts.unsqueeze(0),
|
|
faces=faces.unsqueeze(0),
|
|
uvs=uvs.unsqueeze(0),
|
|
texture=tex,
|
|
)
|
|
return io.NodeOutput(mesh)
|
|
|
|
|
|
class DA3GeometryToPointCloud(io.ComfyNode):
|
|
"""Unproject a DA3_GEOMETRY depth map into a filtered DA3_POINT_CLOUD."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DA3GeometryToPointCloud",
|
|
search_aliases=["da3", "depth anything", "point cloud", "pointcloud", "3d", "geometry"],
|
|
display_name="Convert DA3 Geometry to Point Cloud",
|
|
category="image/geometry estimation",
|
|
description="Convert a depth map into a 3D point cloud.",
|
|
inputs=[
|
|
DA3Geometry.Input("da3_geometry"),
|
|
io.Int.Input("batch_index", default=0, min=0, max=4096, tooltip="Which image of a batch to convert."),
|
|
io.Float.Input("confidence_threshold", default=0.1, min=0.0, max=1.0, step=0.01,
|
|
tooltip="Exclude pixels whose per-image normalised confidence is below this value (0 = keep all). Used when the geometry has a confidence map (Small/Base models)."),
|
|
io.Boolean.Input("use_sky_mask", default=True,
|
|
tooltip="Exclude sky-probability pixels (sky >= 0.5). Used when the geometry has a sky map (Mono/Metric models)."),
|
|
io.Int.Input("downsample", default=1, min=1, max=16,
|
|
tooltip="Take every Nth pixel (1 = full resolution). Higher values give fewer points and faster processing."),
|
|
],
|
|
# TODO: add a proper PointCloud output type
|
|
outputs=[DA3PointCloud.Output(display_name="point_cloud")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, da3_geometry, batch_index, confidence_threshold, use_sky_mask, downsample) -> io.NodeOutput:
|
|
depth_all = da3_geometry["depth"] # (B, H, W)
|
|
B = depth_all.shape[0]
|
|
if batch_index >= B:
|
|
raise ValueError(f"batch_index {batch_index} is out of range; DA3_GEOMETRY has batch size {B}.")
|
|
|
|
depth = depth_all[batch_index].clone() # (H, W)
|
|
depth[~torch.isfinite(depth)] = 0.0
|
|
H, W = depth.shape
|
|
|
|
K = _da3_get_K(da3_geometry, batch_index, H, W)
|
|
|
|
if downsample > 1:
|
|
depth = depth[::downsample, ::downsample].contiguous()
|
|
# Scale intrinsics to the downsampled grid.
|
|
K = K.clone()
|
|
K[0, :] /= downsample
|
|
K[1, :] /= downsample
|
|
|
|
H_ds, W_ds = depth.shape
|
|
points = _da3_unproject(depth, K) # (H_ds, W_ds, 3) in OpenCV camera space
|
|
|
|
# Apply world-to-camera inverse so multi-view frames share a common world frame.
|
|
E = _da3_get_extrinsic(da3_geometry, batch_index)
|
|
if E is not None:
|
|
points = _da3_apply_extrinsic(points, E)
|
|
|
|
# Rebuild mask at downsampled resolution.
|
|
mask = _da3_build_mask(da3_geometry, batch_index, H, W, confidence_threshold, use_sky_mask)
|
|
if downsample > 1:
|
|
mask = mask[::downsample, ::downsample]
|
|
|
|
mask = mask & torch.isfinite(depth)
|
|
|
|
# OpenCV → glTF: flip Y and Z.
|
|
points_gltf = points.clone()
|
|
points_gltf[..., 1] *= -1.0
|
|
points_gltf[..., 2] *= -1.0
|
|
|
|
pts_flat = points_gltf.reshape(-1, 3)[mask.reshape(-1)]
|
|
|
|
colors_flat = None
|
|
if "image" in da3_geometry:
|
|
img = da3_geometry["image"][batch_index] # (H, W, 3)
|
|
if downsample > 1:
|
|
img = img[::downsample, ::downsample]
|
|
colors_flat = img.reshape(-1, 3)[mask.reshape(-1)]
|
|
|
|
conf_flat = None
|
|
if "confidence" in da3_geometry:
|
|
conf = da3_geometry["confidence"][batch_index] # (H, W)
|
|
if downsample > 1:
|
|
conf = conf[::downsample, ::downsample]
|
|
conf_flat = conf.reshape(-1)[mask.reshape(-1)]
|
|
|
|
if pts_flat.shape[0] == 0:
|
|
raise ValueError(
|
|
"DA3GeometryToPointCloud produced zero points after filtering. "
|
|
"Try lowering confidence_threshold or disabling use_sky_mask."
|
|
)
|
|
|
|
return io.NodeOutput({
|
|
"points": pts_flat,
|
|
"colors": colors_flat,
|
|
"confidence": conf_flat,
|
|
})
|
|
|
|
|
|
class DA3Extension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
LoadDA3Model,
|
|
DA3Inference,
|
|
DA3Render,
|
|
DA3GeometryToMesh,
|
|
# DA3GeometryToPointCloud, # Keep this commented out for now until we have a proper PointCloud output type
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> DA3Extension:
|
|
return DA3Extension()
|