439 lines
17 KiB
Python
439 lines
17 KiB
Python
"""Shared utilities for DepthAnythingV3 nodes."""
|
|
import json
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import logging
|
|
|
|
import comfy.model_management as mm
|
|
|
|
logger = logging.getLogger("DepthAnythingV3")
|
|
|
|
# Constants
|
|
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
|
IMAGENET_STD = [0.229, 0.224, 0.225]
|
|
DEFAULT_PATCH_SIZE = 14
|
|
|
|
|
|
def imagenet_normalize(images_pt):
|
|
"""Apply ImageNet normalization without torchvision dependency.
|
|
|
|
Args:
|
|
images_pt: Tensor with shape [B, C, H, W] in [0, 1] range
|
|
|
|
Returns:
|
|
Normalized tensor
|
|
"""
|
|
mean = torch.tensor(IMAGENET_MEAN, device=images_pt.device, dtype=images_pt.dtype).view(1, 3, 1, 1)
|
|
std = torch.tensor(IMAGENET_STD, device=images_pt.device, dtype=images_pt.dtype).view(1, 3, 1, 1)
|
|
return (images_pt - mean) / std
|
|
|
|
|
|
def format_camera_params(param_list, param_name):
|
|
"""Format camera parameters as JSON string.
|
|
|
|
Args:
|
|
param_list: List of camera parameter tensors (or None values)
|
|
param_name: Name of the parameter type (e.g., 'intrinsics', 'extrinsics')
|
|
|
|
Returns:
|
|
JSON string with formatted parameters
|
|
"""
|
|
if all(p is None for p in param_list):
|
|
return json.dumps({param_name: "Not available (mono/metric model)"})
|
|
|
|
formatted = []
|
|
for i, param in enumerate(param_list):
|
|
if param is not None:
|
|
# Convert tensor to list for JSON serialization
|
|
formatted.append({
|
|
f"image_{i}": param.squeeze().tolist()
|
|
})
|
|
else:
|
|
formatted.append({
|
|
f"image_{i}": None
|
|
})
|
|
|
|
return json.dumps({param_name: formatted}, indent=2)
|
|
|
|
|
|
def check_model_capabilities(model):
|
|
"""Check what capabilities a model has.
|
|
|
|
Args:
|
|
model: The DA3 model
|
|
|
|
Returns:
|
|
Dictionary of capabilities
|
|
"""
|
|
has_camera = (
|
|
hasattr(model, 'cam_enc') and model.cam_enc is not None and
|
|
hasattr(model, 'cam_dec') and model.cam_dec is not None
|
|
)
|
|
|
|
# Main series models have camera support, Mono/Metric don't
|
|
# Sky is available on Mono/Metric (DPT head), not on Main series (DualDPT head)
|
|
has_sky = not has_camera # Inverse relationship
|
|
|
|
# Nested model has both (camera from main branch, sky from metric branch)
|
|
is_nested = hasattr(model, 'da3') and hasattr(model, 'da3_metric')
|
|
if is_nested:
|
|
has_sky = True
|
|
has_camera = True
|
|
|
|
has_gs = (
|
|
hasattr(model, 'gs_head') and model.gs_head is not None and
|
|
hasattr(model, 'gs_adapter') and model.gs_adapter is not None
|
|
)
|
|
|
|
return {
|
|
"has_camera_conditioning": has_camera,
|
|
"has_sky_segmentation": has_sky,
|
|
"has_multiview_attention": has_camera,
|
|
"has_3d_gaussians": has_gs,
|
|
"is_nested": is_nested,
|
|
}
|
|
|
|
|
|
def process_tensor_to_image(tensor_list, orig_H, orig_W, normalize_output=False, skip_resize=False):
|
|
"""Convert list of depth/conf tensors to ComfyUI IMAGE format.
|
|
|
|
Args:
|
|
tensor_list: List of tensors with shape [1, H, W] or [H, W]
|
|
orig_H: Original image height
|
|
orig_W: Original image width
|
|
normalize_output: If True, clamp output to 0-1 range
|
|
skip_resize: If True, keep model's native output size instead of resizing back
|
|
|
|
Returns:
|
|
Tensor with shape [B, H, W, 3] in ComfyUI IMAGE format
|
|
"""
|
|
# Concatenate all tensors
|
|
out = torch.cat(tensor_list, dim=0) # [B, 1, H, W] or [B, H, W]
|
|
|
|
# Ensure 4D: [B, 1, H, W]
|
|
if out.dim() == 3:
|
|
out = out.unsqueeze(1)
|
|
|
|
# Convert to 3-channel image [B, H, W, 3]
|
|
out = out.squeeze(1) # [B, H, W]
|
|
out = out.unsqueeze(-1).repeat(1, 1, 1, 3).cpu().float() # [B, H, W, 3]
|
|
|
|
# Resize back to original dimensions (with even constraint) unless skip_resize is True
|
|
if not skip_resize:
|
|
final_H = (orig_H // 2) * 2
|
|
final_W = (orig_W // 2) * 2
|
|
|
|
if out.shape[1] != final_H or out.shape[2] != final_W:
|
|
out = F.interpolate(
|
|
out.permute(0, 3, 1, 2),
|
|
size=(final_H, final_W),
|
|
mode="bilinear"
|
|
).permute(0, 2, 3, 1)
|
|
|
|
if normalize_output:
|
|
return torch.clamp(out, 0, 1)
|
|
return out
|
|
|
|
|
|
def process_tensor_to_mask(tensor_list, orig_H, orig_W, skip_resize=False):
|
|
"""Convert list of tensors to ComfyUI MASK format.
|
|
|
|
Args:
|
|
tensor_list: List of tensors with shape [1, H, W] or [H, W]
|
|
orig_H: Original image height
|
|
orig_W: Original image width
|
|
skip_resize: If True, keep model's native output size instead of resizing back
|
|
|
|
Returns:
|
|
Tensor with shape [B, H, W] in ComfyUI MASK format
|
|
"""
|
|
# Concatenate all tensors
|
|
out = torch.cat(tensor_list, dim=0) # [B, 1, H, W] or [B, H, W]
|
|
|
|
# Ensure 3D: [B, H, W]
|
|
if out.dim() == 4:
|
|
out = out.squeeze(1) # [B, H, W]
|
|
|
|
out = out.cpu().float()
|
|
|
|
# Resize back to original dimensions (with even constraint) unless skip_resize is True
|
|
if not skip_resize:
|
|
final_H = (orig_H // 2) * 2
|
|
final_W = (orig_W // 2) * 2
|
|
|
|
if out.shape[1] != final_H or out.shape[2] != final_W:
|
|
out = F.interpolate(
|
|
out.unsqueeze(1), # [B, 1, H, W] for interpolation
|
|
size=(final_H, final_W),
|
|
mode="bilinear"
|
|
).squeeze(1) # Back to [B, H, W]
|
|
|
|
return torch.clamp(out, 0, 1)
|
|
|
|
|
|
def resize_to_patch_multiple(images_pt, patch_size=DEFAULT_PATCH_SIZE, method="resize"):
|
|
"""Resize images to be divisible by patch size.
|
|
|
|
Args:
|
|
images_pt: Tensor with shape [B, C, H, W]
|
|
patch_size: Patch size to align to (default 14)
|
|
method: How to handle non-divisible sizes:
|
|
- "resize": Resize to nearest patch multiple (default, preserves content)
|
|
- "crop": Center crop to floor patch multiple (loses edges)
|
|
- "pad": Pad to ceiling patch multiple (adds black padding)
|
|
|
|
Returns:
|
|
Tuple of (resized_images, original_H, original_W)
|
|
"""
|
|
_, _, H, W = images_pt.shape
|
|
orig_H, orig_W = H, W
|
|
|
|
if H % patch_size == 0 and W % patch_size == 0:
|
|
return images_pt, orig_H, orig_W
|
|
|
|
if method == "crop":
|
|
# Center crop to floor of patch multiple
|
|
new_H = (H // patch_size) * patch_size
|
|
new_W = (W // patch_size) * patch_size
|
|
|
|
if new_H == 0 or new_W == 0:
|
|
raise ValueError(f"Image too small for patch size {patch_size}. Min size: {patch_size}x{patch_size}")
|
|
|
|
# Calculate crop offsets (center crop)
|
|
top = (H - new_H) // 2
|
|
left = (W - new_W) // 2
|
|
|
|
images_pt = images_pt[:, :, top:top+new_H, left:left+new_W]
|
|
logger.debug(f"Cropped from {orig_H}x{orig_W} to {new_H}x{new_W} (center crop)")
|
|
|
|
elif method == "pad":
|
|
# Pad to ceiling of patch multiple
|
|
new_H = ((H + patch_size - 1) // patch_size) * patch_size
|
|
new_W = ((W + patch_size - 1) // patch_size) * patch_size
|
|
|
|
# Calculate padding (pad bottom and right)
|
|
pad_bottom = new_H - H
|
|
pad_right = new_W - W
|
|
|
|
# F.pad expects (left, right, top, bottom) for 4D tensor
|
|
images_pt = F.pad(images_pt, (0, pad_right, 0, pad_bottom), mode='constant', value=0)
|
|
logger.debug(f"Padded from {orig_H}x{orig_W} to {new_H}x{new_W} (zero padding)")
|
|
|
|
else: # method == "resize" (default)
|
|
# Resize to nearest patch multiple
|
|
def nearest_multiple(x, p):
|
|
down = (x // p) * p
|
|
up = down + p
|
|
return up if abs(up - x) <= abs(x - down) else down
|
|
|
|
new_H = nearest_multiple(H, patch_size)
|
|
new_W = nearest_multiple(W, patch_size)
|
|
|
|
if new_H == 0:
|
|
new_H = patch_size
|
|
if new_W == 0:
|
|
new_W = patch_size
|
|
|
|
images_pt = F.interpolate(images_pt, size=(new_H, new_W), mode="bilinear", align_corners=False)
|
|
logger.debug(f"Resized from {orig_H}x{orig_W} to {new_H}x{new_W} (nearest multiple)")
|
|
|
|
return images_pt, orig_H, orig_W
|
|
|
|
|
|
def save_gaussians_to_ply(gaussians, save_path, depth=None, extrinsics=None,
|
|
shift_and_scale=False, save_sh_dc_only=True,
|
|
prune_border=True, prune_depth_percent=0.9):
|
|
"""Save Gaussians to PLY file in standard 3DGS format.
|
|
|
|
Ported from original DA3 gsply_helpers.export_ply + save_gaussian_ply.
|
|
|
|
Args:
|
|
gaussians: Object with .means, .scales, .rotations, .harmonics, .opacities
|
|
All tensors shape (B, N, ...) where B=1, N = V*H*W
|
|
save_path: Output PLY file path (str or Path)
|
|
depth: Optional depth tensor (V, H, W) or (V, H, W, 1) for spatial pruning.
|
|
Provides shape info for border pruning + depth values for far pruning.
|
|
extrinsics: Optional world-to-camera extrinsics tensor (4x4). If provided,
|
|
positions are transformed from world space to camera space via rigid
|
|
transform, preserving scale/position relationship.
|
|
shift_and_scale: If True, normalize positions to ~[-1, 1] (default: False)
|
|
save_sh_dc_only: If True, save only DC band SH coefficients (default: True)
|
|
prune_border: If True, remove border Gaussians (requires depth for spatial shape)
|
|
prune_depth_percent: Keep closest N% by depth (0.9 = keep 90%). Set 1.0 to skip.
|
|
"""
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from plyfile import PlyData, PlyElement
|
|
except ImportError:
|
|
logger.warning("plyfile not installed - cannot save Gaussians to PLY. "
|
|
"Install with: pip install plyfile")
|
|
return ""
|
|
|
|
# Extract tensors from Gaussians object (handle both object and dict access)
|
|
if isinstance(gaussians, dict):
|
|
means = gaussians['means']
|
|
scales = gaussians['scales']
|
|
rotations = gaussians['rotations']
|
|
harmonics = gaussians['harmonics']
|
|
opacities = gaussians['opacities']
|
|
else:
|
|
means = gaussians.means
|
|
scales = gaussians.scales
|
|
rotations = gaussians.rotations
|
|
harmonics = gaussians.harmonics
|
|
opacities = gaussians.opacities
|
|
|
|
# Ensure batch dim exists, then take batch 0
|
|
if means.dim() == 2:
|
|
means = means.unsqueeze(0)
|
|
scales = scales.unsqueeze(0)
|
|
rotations = rotations.unsqueeze(0)
|
|
harmonics = harmonics.unsqueeze(0)
|
|
opacities = opacities.unsqueeze(0)
|
|
|
|
# Work with batch=0: shapes are (N, 3), (N, 3), (N, 4), (N, 3, d_sh), (N,)
|
|
world_means = means[0]
|
|
gs_scales = scales[0]
|
|
world_rotations = rotations[0]
|
|
world_shs = harmonics[0]
|
|
raw_opacities = opacities[0]
|
|
if raw_opacities.dim() > 1:
|
|
raw_opacities = raw_opacities.squeeze(-1)
|
|
|
|
# Convert opacity to logit space: inverse_sigmoid(opacity)
|
|
raw_opacities = raw_opacities.clamp(1e-6, 1.0 - 1e-6)
|
|
gs_opacities = torch.log(raw_opacities / (1.0 - raw_opacities))
|
|
|
|
# --- Pruning (requires depth for spatial layout) ---
|
|
if depth is not None:
|
|
depth_t = depth
|
|
if depth_t.dim() == 4:
|
|
depth_t = depth_t.squeeze(-1) # (V, H, W)
|
|
if depth_t.dim() == 2:
|
|
depth_t = depth_t.unsqueeze(0) # (1, H, W)
|
|
src_v, out_h, out_w = depth_t.shape
|
|
|
|
N = world_means.shape[0]
|
|
expected_N = src_v * out_h * out_w
|
|
if N == expected_N:
|
|
# Reshape from flat (V*H*W, ...) to spatial (V, H, W, ...)
|
|
spatial_means = world_means.reshape(src_v, out_h, out_w, -1)
|
|
spatial_scales = gs_scales.reshape(src_v, out_h, out_w, -1)
|
|
spatial_rots = world_rotations.reshape(src_v, out_h, out_w, -1)
|
|
spatial_shs = world_shs.reshape(src_v, out_h, out_w, *world_shs.shape[1:])
|
|
spatial_opac = gs_opacities.reshape(src_v, out_h, out_w)
|
|
|
|
# --- Transform positions from world space to camera space ---
|
|
# The GaussianAdapter computes positions AND scales in the same world
|
|
# coordinate system. Applying extrinsics (world-to-camera rigid transform)
|
|
# preserves the scale/position relationship — scales don't change under
|
|
# rotation (same approach as Sharp's covariance transform).
|
|
if extrinsics is not None:
|
|
# Parse extrinsics: squeeze to (4, 4)
|
|
E = extrinsics
|
|
while E.dim() > 2:
|
|
E = E.squeeze(0)
|
|
E = E.to(spatial_means.device, dtype=spatial_means.dtype)
|
|
R = E[:3, :3] # rotation (world-to-camera)
|
|
t = E[:3, 3] # translation
|
|
|
|
# Transform all positions: cam_pos = R @ world_pos + t
|
|
flat_means = spatial_means.reshape(-1, 3)
|
|
cam_means = (flat_means @ R.T) + t
|
|
# Flip Y and Z: OpenCV (Y-down, Z-forward) -> viewer (Y-up, Z-backward)
|
|
cam_means[:, 1] *= -1
|
|
cam_means[:, 2] *= -1
|
|
spatial_means = cam_means.reshape(src_v, out_h, out_w, 3)
|
|
|
|
# Scales are unchanged — rigid transform preserves singular values
|
|
logger.info(f"Transformed Gaussian positions to camera space via extrinsics")
|
|
|
|
# Build mask
|
|
if prune_border:
|
|
mask = torch.zeros(src_v, out_h, out_w, dtype=torch.bool,
|
|
device=world_means.device)
|
|
gstrim_h = max(int(8 / 256 * out_h), 1)
|
|
gstrim_w = max(int(8 / 256 * out_w), 1)
|
|
mask[:, gstrim_h:-gstrim_h, gstrim_w:-gstrim_w] = True
|
|
else:
|
|
mask = torch.ones(src_v, out_h, out_w, dtype=torch.bool,
|
|
device=world_means.device)
|
|
|
|
# Depth pruning
|
|
if prune_depth_percent < 1.0:
|
|
d_percentile = torch.quantile(
|
|
depth_t.reshape(src_v, -1).float(),
|
|
q=prune_depth_percent, dim=1
|
|
).reshape(-1, 1, 1)
|
|
d_mask = depth_t <= d_percentile
|
|
mask = mask & d_mask
|
|
|
|
# Apply mask
|
|
world_means = spatial_means[mask]
|
|
gs_scales = spatial_scales[mask]
|
|
world_rotations = spatial_rots[mask]
|
|
world_shs = spatial_shs[mask]
|
|
gs_opacities = spatial_opac[mask]
|
|
|
|
logger.info(f"Pruned Gaussians: {N} -> {world_means.shape[0]} "
|
|
f"(border={prune_border}, depth_pct={prune_depth_percent})")
|
|
else:
|
|
logger.warning(f"Cannot prune: N={N} != V*H*W={expected_N}. Saving all.")
|
|
|
|
# --- Optional shift and scale ---
|
|
if shift_and_scale:
|
|
median = world_means.median(dim=0).values
|
|
world_means = world_means - median
|
|
scale_factor = world_means.abs().quantile(0.95, dim=0).max()
|
|
if scale_factor > 0:
|
|
world_means = world_means / scale_factor
|
|
gs_scales = gs_scales / scale_factor
|
|
|
|
# --- Build PLY attributes (matching original DA3 format) ---
|
|
f_dc = world_shs[..., 0] # (N, 3) — DC band
|
|
f_rest = world_shs[..., 1:].flatten(start_dim=1) if world_shs.shape[-1] > 1 else None
|
|
|
|
# Build property list (includes nx/ny/nz dummy normals to match original DA3 format)
|
|
attr_names = ["x", "y", "z", "nx", "ny", "nz"]
|
|
if save_sh_dc_only:
|
|
attr_names += ["f_dc_0", "f_dc_1", "f_dc_2"]
|
|
else:
|
|
attr_names += ["f_dc_0", "f_dc_1", "f_dc_2"]
|
|
if f_rest is not None:
|
|
for i in range(f_rest.shape[1]):
|
|
attr_names.append(f"f_rest_{i}")
|
|
attr_names.append("opacity")
|
|
attr_names += ["scale_0", "scale_1", "scale_2"]
|
|
attr_names += ["rot_0", "rot_1", "rot_2", "rot_3"]
|
|
|
|
dtype_full = [(name, "f4") for name in attr_names]
|
|
|
|
# Concatenate all attribute arrays
|
|
means_np = world_means.detach().cpu().numpy()
|
|
attr_arrays = [
|
|
means_np, # x, y, z
|
|
np.zeros_like(means_np), # nx, ny, nz (dummy normals)
|
|
f_dc.detach().cpu().contiguous().numpy(), # f_dc_0, f_dc_1, f_dc_2
|
|
]
|
|
if not save_sh_dc_only and f_rest is not None:
|
|
attr_arrays.append(f_rest.detach().cpu().contiguous().numpy())
|
|
attr_arrays += [
|
|
gs_opacities[..., None].detach().cpu().numpy(), # opacity (logit)
|
|
gs_scales.log().detach().cpu().numpy(), # scale_0,1,2 (log)
|
|
world_rotations.detach().cpu().numpy(), # rot_0,1,2,3
|
|
]
|
|
|
|
attributes = np.concatenate(attr_arrays, axis=1)
|
|
|
|
N = world_means.shape[0]
|
|
elements = np.empty(N, dtype=dtype_full)
|
|
elements[:] = list(map(tuple, attributes))
|
|
|
|
save_path = Path(save_path)
|
|
save_path.parent.mkdir(exist_ok=True, parents=True)
|
|
PlyData([PlyElement.describe(elements, "vertex")]).write(str(save_path))
|
|
logger.info(f"Saved Gaussians ({N} splats, {len(attr_names)} properties) to: {save_path}")
|
|
return str(save_path)
|