266 lines
11 KiB
Python
266 lines
11 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import os
|
|
import sys
|
|
from typing import Dict, Any, Tuple
|
|
from tqdm import tqdm # Added tqdm import
|
|
|
|
# Import existing pointcloud nodes and projection definitions
|
|
from .pointcloud_nodes import DepthToPointCloud, TransformPointCloud, ProjectPointCloud, Projection, PointCloudCleaner
|
|
import folder_paths
|
|
|
|
# Ensure video_depth_anything is on path
|
|
video_depth_path = os.path.join("/root/Video-Depth-Anything", "metric_depth")
|
|
if video_depth_path not in sys.path:
|
|
sys.path.append(video_depth_path)
|
|
try:
|
|
from video_depth_anything.video_depth import VideoDepthAnything
|
|
print("video_depth_anything module loaded successfully.")
|
|
except ImportError:
|
|
VideoDepthAnything = None
|
|
print("Warning: video_depth_anything module not found. Ensure it is installed correctly.")
|
|
print("error: ", sys.exc_info()[1])
|
|
|
|
|
|
class VideoCameraMotionSequence:
|
|
"""
|
|
Takes a sequence of RGB frames and corresponding depth maps,
|
|
converts each frame+depth to a pointcloud, interpolates a camera
|
|
trajectory to match video length, cleans the pointcloud if needed,
|
|
and outputs reprojected images, masks, and depth maps per frame.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
# Sequence of frames: Tensor [T, H, W, 3]
|
|
"frames": ("IMAGE", {"shape_hint": [None, None, None, 3]}),
|
|
# Sequence of depth maps: Tensor [T, H, W] or [T, H, W, 1]
|
|
"depth_seq": ("TENSOR", {"shape_hint": [None, None, None]}),
|
|
# Camera trajectory waypoints: Tensor [K, 4, 4]
|
|
"trajectory": ("TENSOR", {"shape_hint": [None, 4, 4]}),
|
|
# Input projection parameters
|
|
"input_projection": (Projection.PROJECTIONS, {}),
|
|
"input_horizontal_fov": ("FLOAT", {"default": 90.0}),
|
|
"depth_scale": ("FLOAT", {"default": 1.0}),
|
|
"invert_depth": ("BOOLEAN", {"default": False}),
|
|
# Output projection parameters
|
|
"output_projection": (Projection.PROJECTIONS, {}),
|
|
"output_horizontal_fov": ("FLOAT", {"default": 90.0}),
|
|
"output_width": ("INT", {"default": 512, "min": 1}),
|
|
"output_height": ("INT", {"default": 512, "min": 1}),
|
|
"point_size": ("INT", {"default": 1, "min": 1}),
|
|
# Cleaning parameters
|
|
"voxel_size": ("FLOAT", {"default": 1.0, "min": 1e-3}),
|
|
"min_points_per_voxel": ("INT", {"default": 3, "min": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "TENSOR")
|
|
RETURN_NAMES = ("video_frames", "mask_frames", "depths")
|
|
FUNCTION = "process_sequence"
|
|
CATEGORY = "Camera/Video"
|
|
|
|
def process_sequence(
|
|
self,
|
|
frames: torch.Tensor,
|
|
depth_seq: torch.Tensor,
|
|
trajectory: torch.Tensor,
|
|
input_projection: str,
|
|
input_horizontal_fov: float,
|
|
depth_scale: float,
|
|
invert_depth: bool,
|
|
output_projection: str,
|
|
output_horizontal_fov: float,
|
|
output_width: int,
|
|
output_height: int,
|
|
point_size: int,
|
|
voxel_size: float,
|
|
min_points_per_voxel: int,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
# frames: [T, H, W, 3]
|
|
# depth_seq: [T, H, W] or [T, H, W, 1]
|
|
T, H, W, _ = frames.shape
|
|
|
|
# Interpolate trajectory to match T
|
|
K = trajectory.shape[0]
|
|
if K < 2:
|
|
interp_traj = trajectory.expand(T, 4, 4).clone()
|
|
else:
|
|
idxs = torch.linspace(0, K - 1, T, device=trajectory.device)
|
|
lower = idxs.floor().long().clamp(max=K - 2)
|
|
upper = lower + 1
|
|
alpha = (idxs - lower.float()).unsqueeze(-1).unsqueeze(-1)
|
|
traj_lower = trajectory[lower]
|
|
traj_upper = trajectory[upper]
|
|
interp_traj = traj_lower * (1 - alpha) + traj_upper * alpha
|
|
|
|
out_frames = []
|
|
out_masks = []
|
|
out_depths = []
|
|
|
|
# Add tqdm progress bar for the sequence
|
|
for frame, depth, pose in tqdm(zip(frames, depth_seq, interp_traj), total=T, desc="Processing video frames"):
|
|
if depth.dim() == 3 and depth.shape[-1] == 1:
|
|
depth = depth.squeeze(-1)
|
|
# to pointcloud
|
|
pc, = DepthToPointCloud().depth_to_pointcloud(
|
|
image=frame.permute(2, 0, 1),
|
|
input_projection=input_projection,
|
|
input_horizontal_fov=input_horizontal_fov,
|
|
depth_scale=depth_scale,
|
|
invert_depth=invert_depth,
|
|
depthmap=depth,
|
|
mask=None,
|
|
)
|
|
# optional cleaning
|
|
if min_points_per_voxel > 1:
|
|
pc, = PointCloudCleaner().clean_pointcloud(
|
|
pointcloud=pc,
|
|
width=output_width,
|
|
height=output_height,
|
|
voxel_size=voxel_size,
|
|
min_points_per_voxel=min_points_per_voxel,
|
|
)
|
|
# transform and project
|
|
pc_t, = TransformPointCloud().transform_pointcloud(pc, pose)
|
|
img_t, mask_t, depth_t = ProjectPointCloud().project_pointcloud(
|
|
pointcloud=pc_t,
|
|
output_projection=output_projection,
|
|
output_horizontal_fov=output_horizontal_fov,
|
|
output_width=output_width,
|
|
output_height=output_height,
|
|
point_size=point_size,
|
|
)
|
|
|
|
out_frames.append(img_t[0])
|
|
out_masks.append(mask_t)
|
|
out_depths.append(depth_t)
|
|
|
|
return (
|
|
torch.stack(out_frames, dim=0), # [T, 3, H, W]
|
|
torch.stack(out_masks, dim=0), # [T, H, W]
|
|
torch.stack(out_depths, dim=0), # [T, H, W]
|
|
)
|
|
|
|
|
|
class DepthFramesToVideo:
|
|
"""
|
|
Converts a sequence of depth maps into video frame tensors for saving.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"depth_seq": ("TENSOR", {"shape_hint": [None, None, None]}),
|
|
"mask_seq": ("MASK", {"shape_hint": [None, None, None]}),
|
|
"normalize": ("BOOLEAN", {"default": True}),
|
|
"invert_depth": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TENSOR", "IMAGE")
|
|
RETURN_NAMES = ("video_frames", "depth_video")
|
|
FUNCTION = "depth_to_video_frames"
|
|
CATEGORY = "Camera/Video"
|
|
|
|
def depth_to_video_frames(
|
|
self,
|
|
depth_seq: torch.Tensor,
|
|
normalize: bool,
|
|
invert_depth: bool,
|
|
mask_seq: torch.Tensor,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
ds = depth_seq.clone().squeeze()
|
|
if ds.dim() == 2:
|
|
ds = ds.unsqueeze(0) # [H, W] -> [1, H, W]
|
|
if ds.dim() != 3:
|
|
raise ValueError(f"Expected ds to be 3D [T, H, W], got shape {ds.shape}")
|
|
if invert_depth:
|
|
ds= 1.0 / (ds + 1e-8) # Avoid division by zero
|
|
if normalize:
|
|
# Mask: only normalize where depth > 0
|
|
mask = mask_seq>0.5
|
|
if mask.any():
|
|
#percentile first 10 percent min
|
|
# sample
|
|
minv = ds[mask]
|
|
# sample 10000 and find 10% quantile
|
|
if minv.numel() > 10000:
|
|
minv = minv[torch.randperm(minv.numel())[:10000]]
|
|
|
|
minv = minv.quantile(0.2)
|
|
minv = minv if minv > 0.1 else 0.1 # Avoid division by zero
|
|
#percentile last 10 percent max
|
|
maxv = ds[mask]
|
|
if maxv.numel() > 10000:
|
|
maxv = maxv[torch.randperm(maxv.numel())[:10000]]
|
|
maxv = maxv.quantile(0.98)
|
|
maxv = maxv if maxv < 100 else 100
|
|
print(f"Normalizing depth: min={minv}, max={maxv}")
|
|
ds_norm = (ds - minv) / (maxv - minv + 1e-8)
|
|
ds = ds_norm.clamp(0, 1) # torch.where(mask, ds_norm, ds) # Only normalize valid values
|
|
else:
|
|
print("Warning: No valid depth values for normalization.")
|
|
# expand to 3 channels: [T, H, W] -> [T, 3, H, W]
|
|
raw = depth_seq.clone().squeeze()
|
|
ds_u8 = (ds * 255.0).round().to(torch.uint8)
|
|
raw_u8 = (raw.clamp(0, 255)).to(torch.uint8) # if raw is already in a displayable range
|
|
|
|
# expand to 3 channels and permute to HWC
|
|
ds_color = ds_u8.unsqueeze(1).repeat(1, 3, 1, 1).permute(0, 2, 3, 1)
|
|
raw_color = raw_u8.unsqueeze(1).repeat(1, 3, 1, 1).permute(0, 2, 3, 1)
|
|
return raw_color, ds_color # [T, 3, H, W] -> [T, H, W, 3]
|
|
|
|
class VideoMetricDepthEstimate:
|
|
"""
|
|
Estimates metric depth for a sequence of frames using VideoDepthAnything.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
# model files (.pth) in input directory
|
|
model_dir = os.path.join(os.getcwd(), "models", "checkpoints")
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
files = [f for f in os.listdir(model_dir) if f.lower().endswith(('.pth', '.ckpt', '.safetensors'))]
|
|
return {
|
|
"required": {
|
|
"frames": ("IMAGE", {"shape_hint": [None, None, None, 3]}),
|
|
"model_checkpoint": (files, {"file_chooser": True}),
|
|
"input_size": ("INT", {"default": 518, "min": 64, "max": 2048}),
|
|
"max_fps": ("INT", {"default": 60, "min": 1}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("TENSOR", "FLOAT")
|
|
RETURN_NAMES = ("metric_depths", "fps")
|
|
FUNCTION = "estimate_metric_depth"
|
|
CATEGORY = "Camera/Video"
|
|
|
|
def estimate_metric_depth(
|
|
self,
|
|
frames: torch.Tensor,
|
|
model_checkpoint: str,
|
|
input_size: int,
|
|
max_fps: int,
|
|
) -> Tuple[torch.Tensor, float]:
|
|
if VideoDepthAnything is None:
|
|
raise ImportError("VideoDepthAnything library not found")
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
# if max input<1.5 normalize to 0-255
|
|
if frames.max() < 1.5:
|
|
frames = (frames * 255)
|
|
model = VideoDepthAnything(**{"encoder": "vitl", "features": 256, "out_channels": [256,512,1024,1024]})
|
|
state = torch.load("/root/ComfyUI/models/checkpoints/{}".format(model_checkpoint), map_location='cpu')
|
|
model.load_state_dict(state, strict=True)
|
|
model = model.to(device).eval()
|
|
np_frames = frames.cpu().numpy().astype(np.uint8)
|
|
metric_depths, fps = model.infer_video_depth(np_frames, max_fps, input_size=input_size, device=device.type, fp32=False)
|
|
return (torch.from_numpy(metric_depths), float(fps))
|
|
|
|
# Register nodes
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VideoCameraMotionSequence": VideoCameraMotionSequence,
|
|
"VideoMetricDepthEstimate": VideoMetricDepthEstimate,
|
|
"DepthFramesToVideo": DepthFramesToVideo,
|
|
}
|