diff --git a/assets/pano.png b/assets/pano.png new file mode 100644 index 0000000..cd39250 Binary files /dev/null and b/assets/pano.png differ diff --git a/comfy-env-root.toml b/comfy-env-root.toml index 3db49cb..be07239 100644 --- a/comfy-env-root.toml +++ b/comfy-env-root.toml @@ -1 +1,2 @@ -[node_reqs] \ No newline at end of file +[node_reqs] +ComfyUI-GeometryPack = "PozzettiAndrea/ComfyUI-GeometryPack" \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py index 8dfc766..8960c64 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -6,6 +6,16 @@ from .predict import NODE_CLASS_MAPPINGS as PREDICT_MAPPINGS from .predict import NODE_DISPLAY_NAME_MAPPINGS as PREDICT_DISPLAY_MAPPINGS from .load_image_exif import NODE_CLASS_MAPPINGS as LOAD_IMAGE_EXIF_MAPPINGS from .load_image_exif import NODE_DISPLAY_NAME_MAPPINGS as LOAD_IMAGE_EXIF_DISPLAY_MAPPINGS +from .sample_panorama import NODE_CLASS_MAPPINGS as SAMPLE_PANORAMA_MAPPINGS +from .sample_panorama import NODE_DISPLAY_NAME_MAPPINGS as SAMPLE_PANORAMA_DISPLAY_MAPPINGS +from .merge_gaussians import NODE_CLASS_MAPPINGS as MERGE_GAUSSIANS_MAPPINGS +from .merge_gaussians import NODE_DISPLAY_NAME_MAPPINGS as MERGE_GAUSSIANS_DISPLAY_MAPPINGS +from .predict_depth import NODE_CLASS_MAPPINGS as PREDICT_DEPTH_MAPPINGS +from .predict_depth import NODE_DISPLAY_NAME_MAPPINGS as PREDICT_DEPTH_DISPLAY_MAPPINGS +from .project_depth_to_panorama import NODE_CLASS_MAPPINGS as PROJECT_DEPTH_MAPPINGS +from .project_depth_to_panorama import NODE_DISPLAY_NAME_MAPPINGS as PROJECT_DEPTH_DISPLAY_MAPPINGS +from .align_depth_maps import NODE_CLASS_MAPPINGS as ALIGN_DEPTH_MAPPINGS +from .align_depth_maps import NODE_DISPLAY_NAME_MAPPINGS as ALIGN_DEPTH_DISPLAY_MAPPINGS # Aggregate all node mappings NODE_CLASS_MAPPINGS = {} @@ -14,9 +24,19 @@ NODE_DISPLAY_NAME_MAPPINGS = {} NODE_CLASS_MAPPINGS.update(LOAD_MODEL_MAPPINGS) NODE_CLASS_MAPPINGS.update(PREDICT_MAPPINGS) NODE_CLASS_MAPPINGS.update(LOAD_IMAGE_EXIF_MAPPINGS) +NODE_CLASS_MAPPINGS.update(SAMPLE_PANORAMA_MAPPINGS) +NODE_CLASS_MAPPINGS.update(MERGE_GAUSSIANS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(PREDICT_DEPTH_MAPPINGS) +NODE_CLASS_MAPPINGS.update(PROJECT_DEPTH_MAPPINGS) +NODE_CLASS_MAPPINGS.update(ALIGN_DEPTH_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_MODEL_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_IMAGE_EXIF_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(SAMPLE_PANORAMA_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(MERGE_GAUSSIANS_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_DEPTH_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(PROJECT_DEPTH_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(ALIGN_DEPTH_DISPLAY_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/align_depth_maps.py b/nodes/align_depth_maps.py new file mode 100644 index 0000000..0867fc1 --- /dev/null +++ b/nodes/align_depth_maps.py @@ -0,0 +1,664 @@ +"""AlignDepthMaps node for ComfyUI-Sharp. + +RANSAC-based scale-shift alignment to make depth maps globally consistent +across multiple perspective views before panorama projection. +""" + +import json +import math +from collections import deque + +import torch +import torch.nn.functional as F + + +def extrinsics_to_direction(extrinsics: torch.Tensor) -> torch.Tensor: + """Get the viewing direction (forward vector) from extrinsics. + + Args: + extrinsics: [4, 4] world-to-camera matrix + + Returns: + [3] normalized direction vector in world space + """ + # Camera-to-world is inverse of world-to-camera + # For pure rotation, inverse = transpose + R_c2w = extrinsics[:3, :3].T + + # Camera looks along +Z in camera space + forward = R_c2w @ torch.tensor([0., 0., 1.], device=extrinsics.device) + return forward + + +def compute_view_overlap( + ext_i: torch.Tensor, + ext_j: torch.Tensor, + half_fov: float, +) -> float: + """Compute approximate overlap fraction between two views. + + Uses angular distance between view centers as a proxy for overlap. + + Args: + ext_i, ext_j: [4, 4] extrinsics matrices + half_fov: Half field of view in radians + + Returns: + Overlap fraction [0, 1], where 1 = complete overlap + """ + dir_i = extrinsics_to_direction(ext_i) + dir_j = extrinsics_to_direction(ext_j) + + # Angular distance between view centers + dot = torch.clamp(torch.dot(dir_i, dir_j), -1, 1) + angle = torch.acos(dot).item() + + # If angle < 2*half_fov, there's overlap + # Overlap fraction decreases linearly from 1 (same direction) to 0 (2*half_fov apart) + max_angle = 2 * half_fov + if angle >= max_angle: + return 0.0 + return 1.0 - angle / max_angle + + +def compute_overlap_mask( + depth_h: int, + depth_w: int, + ext_i: torch.Tensor, + ext_j: torch.Tensor, + intrinsics: torch.Tensor, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute which pixels in view i overlap with view j and vice versa. + + Returns corresponding pixel coordinates for RANSAC matching. + + Args: + depth_h, depth_w: Depth map dimensions + ext_i, ext_j: [4, 4] extrinsics matrices + intrinsics: [4, 4] intrinsics matrix + device: Torch device + + Returns: + Tuple of (mask_i, mask_j, coords_i, coords_j) where: + - mask_i: [H, W] bool mask of pixels in i that overlap with j + - mask_j: [H, W] bool mask of pixels in j that overlap with i + - coords_i: [N, 2] pixel coordinates in view i + - coords_j: [N, 2] corresponding pixel coordinates in view j + """ + f_px = intrinsics[0, 0].item() + cx = intrinsics[0, 2].item() + cy = intrinsics[1, 2].item() + + # Create pixel grid for view i + u = torch.arange(depth_w, dtype=torch.float32, device=device) + v = torch.arange(depth_h, dtype=torch.float32, device=device) + uu, vv = torch.meshgrid(u, v, indexing='xy') + + # Convert to camera-space ray directions + dx = (uu - cx) / f_px + dy = (vv - cy) / f_px + dz = torch.ones_like(dx) + rays_cam = torch.stack([dx, dy, dz], dim=-1) + rays_cam = F.normalize(rays_cam, dim=-1) + + # Rotate to world space using view i's camera-to-world + R_i_c2w = ext_i[:3, :3].T + rays_world = torch.einsum('ij,hwj->hwi', R_i_c2w, rays_cam) + + # Transform world rays to view j's camera space + R_j_w2c = ext_j[:3, :3] # World-to-camera for view j + rays_j_cam = torch.einsum('ij,hwj->hwi', R_j_w2c, rays_world) + + # Project to view j's image plane + # Only valid if z > 0 (in front of camera) + z_j = rays_j_cam[..., 2] + valid_z = z_j > 0.01 + + # Perspective projection + x_j = rays_j_cam[..., 0] / (z_j + 1e-6) + y_j = rays_j_cam[..., 1] / (z_j + 1e-6) + + # Convert to pixel coordinates + u_j = x_j * f_px + cx + v_j = y_j * f_px + cy + + # Check if projected point is within view j's image bounds + in_bounds = (u_j >= 0) & (u_j < depth_w) & (v_j >= 0) & (v_j < depth_h) + + # Combined mask for view i + mask_i = valid_z & in_bounds + + # Get pixel coordinates for matching + coords_i_y, coords_i_x = torch.where(mask_i) + coords_i = torch.stack([coords_i_x, coords_i_y], dim=1) # [N, 2] as (x, y) + + # Corresponding coordinates in view j + u_j_valid = u_j[mask_i] + v_j_valid = v_j[mask_i] + coords_j = torch.stack([u_j_valid, v_j_valid], dim=1) # [N, 2] as (x, y) + + # Create mask_j (approximate - mark pixels that have correspondences) + mask_j = torch.zeros(depth_h, depth_w, dtype=torch.bool, device=device) + u_j_int = coords_j[:, 0].long().clamp(0, depth_w - 1) + v_j_int = coords_j[:, 1].long().clamp(0, depth_h - 1) + mask_j[v_j_int, u_j_int] = True + + return mask_i, mask_j, coords_i, coords_j + + +def median_scale_alignment( + depth1: torch.Tensor, + depth2: torch.Tensor, + coords1: torch.Tensor, + coords2: torch.Tensor, + debug: bool = True, +) -> tuple[float, float, int, dict]: + """Find scale to align depth2 to depth1 using median ratio. + + Solves: depth1 = scale * depth2 (NO SHIFT - prevents negative depths) + + Args: + depth1: [H, W] reference depth map + depth2: [H, W] depth map to align + coords1: [N, 2] pixel coordinates in depth1 (x, y) + coords2: [N, 2] corresponding pixel coordinates in depth2 (x, y) + debug: Whether to print debug info + + Returns: + Tuple of (scale, shift=0.0, num_valid, debug_info) + """ + n_points = coords1.shape[0] + debug_info = {"overlap_pixels": n_points} + + if n_points < 10: + return 1.0, 0.0, 0, debug_info + + # Sample depth values at corresponding locations + x1, y1 = coords1[:, 0].long(), coords1[:, 1].long() + x2, y2 = coords2[:, 0].long().clamp(0, depth2.shape[1] - 1), coords2[:, 1].long().clamp(0, depth2.shape[0] - 1) + + d1 = depth1[y1, x1] + d2 = depth2[y2, x2] + + # Filter out invalid depths (must be positive and finite) + valid = (d1 > 0.1) & (d2 > 0.1) & torch.isfinite(d1) & torch.isfinite(d2) + d1_valid = d1[valid] + d2_valid = d2[valid] + n_valid = len(d1_valid) + + debug_info["valid_pixels"] = n_valid + debug_info["d1_range"] = (d1_valid.min().item(), d1_valid.max().item()) if n_valid > 0 else (0, 0) + debug_info["d2_range"] = (d2_valid.min().item(), d2_valid.max().item()) if n_valid > 0 else (0, 0) + + if n_valid < 10: + return 1.0, 0.0, n_valid, debug_info + + # Compute per-pixel scale ratios + ratios = d1_valid / d2_valid + + # Debug: ratio statistics + debug_info["ratio_range"] = (ratios.min().item(), ratios.max().item()) + debug_info["ratio_median"] = ratios.median().item() + debug_info["ratio_std"] = ratios.std().item() + + # Use median ratio (robust to outliers) + scale = ratios.median().item() + + # Constraint: keep scale in reasonable range [0.5, 2.0] + # If scale is outside this range, there's likely a problem + original_scale = scale + scale = max(0.5, min(2.0, scale)) + debug_info["scale_clamped"] = (original_scale != scale) + debug_info["original_scale"] = original_scale + + # Count "inliers" - pixels where aligned depth is within 10% of reference + aligned = scale * d2_valid + relative_error = torch.abs(d1_valid - aligned) / (d1_valid + 1e-6) + inliers = (relative_error < 0.1).sum().item() + debug_info["inlier_ratio"] = inliers / n_valid if n_valid > 0 else 0 + + return scale, 0.0, n_valid, debug_info # shift is always 0! + + +def build_adjacency_graph( + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + min_overlap: float = 0.1, +) -> dict[int, list[int]]: + """Build a graph of which views overlap with each other. + + Args: + extrinsics: [N, 4, 4] batch of extrinsics + intrinsics: [4, 4] shared intrinsics + min_overlap: Minimum overlap fraction to consider views adjacent + + Returns: + Dictionary mapping view index to list of adjacent view indices + """ + num_views = extrinsics.shape[0] + f_px = intrinsics[0, 0].item() + cx = intrinsics[0, 2].item() + depth_size = int(cx * 2) # Derive from principal point (cx = width/2) + half_fov = math.atan((depth_size / 2) / f_px) + + adjacency = {i: [] for i in range(num_views)} + + for i in range(num_views): + for j in range(i + 1, num_views): + overlap = compute_view_overlap(extrinsics[i], extrinsics[j], half_fov) + if overlap >= min_overlap: + adjacency[i].append(j) + adjacency[j].append(i) + + return adjacency + + +def compute_global_alignments( + depth_maps: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + reference_view: int, + num_iterations: int, + inlier_threshold: float, +) -> dict[int, tuple[float, float]]: + """Compute alignment parameters for all views relative to reference. + + Uses BFS to chain pairwise alignments through the view graph. + DEPRECATED: Use compute_global_alignments_optimized instead. + + Args: + depth_maps: [N, H, W, C] depth maps + extrinsics: [N, 4, 4] extrinsics + intrinsics: [4, 4] intrinsics + reference_view: Index of reference view + num_iterations: RANSAC iterations + inlier_threshold: RANSAC inlier threshold + + Returns: + Dictionary mapping view index to (scale, shift) tuple + """ + num_views = depth_maps.shape[0] + depth_h, depth_w = depth_maps.shape[1:3] + device = depth_maps.device + + # Build adjacency graph + adjacency = build_adjacency_graph(extrinsics, intrinsics) + + # BFS from reference view + alignments = {reference_view: (1.0, 0.0)} # Reference is unchanged + visited = {reference_view} + queue = deque([reference_view]) + + pairwise_results = {} # Cache pairwise alignments + + while queue: + current = queue.popleft() + current_scale, current_shift = alignments[current] + + for neighbor in adjacency[current]: + if neighbor in visited: + continue + + # Compute overlap mask and coordinates + mask_c, mask_n, coords_c, coords_n = compute_overlap_mask( + depth_h, depth_w, + extrinsics[current], extrinsics[neighbor], + intrinsics, device + ) + + if coords_c.shape[0] < 10: + # Not enough overlap, skip this edge + continue + + # Get depth maps (single channel) + depth_c = depth_maps[current, :, :, 0] + depth_n = depth_maps[neighbor, :, :, 0] + + # Apply current view's alignment to get aligned reference depth + aligned_depth_c = current_scale * depth_c + current_shift + + # Find alignment from neighbor to aligned current + # aligned_depth_c = scale * depth_n (NO SHIFT!) + scale, shift, n_valid, debug_info = median_scale_alignment( + aligned_depth_c, depth_n, + coords_c, coords_n, + debug=True + ) + + # Detailed debug output + print(f"[AlignDepthMaps] View {neighbor} -> {current}:") + print(f" Overlap: {debug_info['overlap_pixels']} pixels, {debug_info['valid_pixels']} valid") + print(f" Depth1 (ref): {debug_info['d1_range'][0]:.2f} - {debug_info['d1_range'][1]:.2f}") + print(f" Depth2 (src): {debug_info['d2_range'][0]:.2f} - {debug_info['d2_range'][1]:.2f}") + print(f" Ratio range: {debug_info['ratio_range'][0]:.3f} - {debug_info['ratio_range'][1]:.3f}") + print(f" Scale: {scale:.4f} (median={debug_info['ratio_median']:.4f}, std={debug_info['ratio_std']:.4f})") + if debug_info.get('scale_clamped'): + print(f" WARNING: Scale clamped from {debug_info['original_scale']:.4f} to {scale:.4f}") + print(f" Inlier ratio: {debug_info['inlier_ratio']:.1%}") + + alignments[neighbor] = (scale, shift) + visited.add(neighbor) + queue.append(neighbor) + + # Handle disconnected views (no overlap with any aligned view) + for i in range(num_views): + if i not in alignments: + print(f"[AlignDepthMaps] Warning: View {i} has no overlap with aligned views, using identity") + alignments[i] = (1.0, 0.0) + + return alignments + + +def compute_global_misalignment( + scales: list[float], + pairwise_ratios: dict[tuple[int, int], float], +) -> tuple[float, float, dict]: + """Compute global misalignment loss given scales and measured ratios. + + The loss measures how well the solved scales satisfy the pairwise constraints. + + Args: + scales: List of scale values [s_0, s_1, ..., s_N] + pairwise_ratios: Dict mapping (i,j) to measured ratio d_i/d_j + + Returns: + Tuple of (total_loss, mean_loss, per_pair_losses) + """ + per_pair_losses = {} + total_loss = 0.0 + + for (i, j), measured_ratio in pairwise_ratios.items(): + # We want: s_i * d_i = s_j * d_j in overlap + # So: s_i / s_j should equal d_j / d_i = 1/measured_ratio + target_ratio = 1.0 / measured_ratio + actual_ratio = scales[i] / scales[j] if scales[j] > 0 else float('inf') + + # Loss in log-space (scale-invariant) + if actual_ratio > 0 and target_ratio > 0: + log_error = abs(math.log(actual_ratio) - math.log(target_ratio)) + per_pair_losses[(i, j)] = log_error + total_loss += log_error ** 2 + else: + per_pair_losses[(i, j)] = float('inf') + + num_pairs = len(pairwise_ratios) + mean_loss = math.sqrt(total_loss / num_pairs) if num_pairs > 0 else 0.0 + + return total_loss, mean_loss, per_pair_losses + + +def compute_global_alignments_optimized( + depth_maps: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, +) -> dict[int, tuple[float, float]]: + """Compute globally optimal scale alignments using least squares. + + Instead of BFS chaining (which causes cumulative drift), solve for all + scales jointly by minimizing total pairwise inconsistency. + + The formulation: + - For each pair (i,j) we measure ratio r_ij = median(d_i / d_j) + - We want scales s_i such that s_i / s_j ≈ r_ij + - In log-space: x_i - x_j = log(r_ij) where x_i = log(s_i) + - This is a linear system A @ x = b, solved with least squares + + Args: + depth_maps: [N, H, W, C] depth maps + extrinsics: [N, 4, 4] extrinsics + intrinsics: [4, 4] intrinsics + + Returns: + Dictionary mapping view index to (scale, shift=0.0) tuple + """ + num_views = depth_maps.shape[0] + depth_h, depth_w = depth_maps.shape[1:3] + device = depth_maps.device + + print(f"[GlobalOpt] Computing globally optimal alignment for {num_views} views") + + # Step 1: Compute ALL pairwise scale ratios + adjacency = build_adjacency_graph(extrinsics, intrinsics) + pairwise_ratios = {} # (i, j) -> ratio where i < j + + print(f"[GlobalOpt] Computing pairwise ratios...") + for i in range(num_views): + for j in adjacency[i]: + if i < j: # Only compute once per pair + # Compute overlap + mask_i, mask_j, coords_i, coords_j = compute_overlap_mask( + depth_h, depth_w, extrinsics[i], extrinsics[j], + intrinsics, device + ) + + if coords_i.shape[0] < 10: + print(f"[GlobalOpt] Pair ({i},{j}): insufficient overlap ({coords_i.shape[0]} pixels)") + continue + + depth_i = depth_maps[i, :, :, 0] + depth_j = depth_maps[j, :, :, 0] + + # Get median ratio: d_i / d_j + scale, _, n_valid, debug_info = median_scale_alignment( + depth_i, depth_j, coords_i, coords_j, debug=True + ) + + if n_valid > 100: + pairwise_ratios[(i, j)] = scale + print(f"[GlobalOpt] Pair ({i},{j}): ratio={scale:.4f}, " + f"valid={n_valid}, inlier_ratio={debug_info['inlier_ratio']:.1%}") + else: + print(f"[GlobalOpt] Pair ({i},{j}): insufficient valid pixels ({n_valid})") + + # Compute initial loss (with identity scales) + identity_scales = [1.0] * num_views + initial_loss, initial_rmse, _ = compute_global_misalignment(identity_scales, pairwise_ratios) + print(f"[GlobalOpt] Initial misalignment (identity scales): loss={initial_loss:.4f}, RMSE={initial_rmse:.4f}") + + # Step 2: Build linear system for log-scale optimization + # x_i - x_j = log(r_ij) for each pair + # Fix x_0 = 0 (reference) + + num_pairs = len(pairwise_ratios) + print(f"[GlobalOpt] Found {num_pairs} valid pairs") + + if num_pairs == 0: + print("[GlobalOpt] No valid pairs, using identity scales") + return {i: (1.0, 0.0) for i in range(num_views)} + + # Build A matrix and b vector + # For N views with reference fixed, we have N-1 unknowns + # A is [num_pairs, num_views-1], b is [num_pairs] + A = torch.zeros(num_pairs, num_views - 1, device=device) + b = torch.zeros(num_pairs, device=device) + + for k, ((i, j), ratio) in enumerate(pairwise_ratios.items()): + # Equation: x_i - x_j = log(ratio) + # ratio = d_i / d_j, so we want s_i * d_i = s_j * d_j + # which means s_i / s_j = d_j / d_i = 1/ratio + # So: x_i - x_j = -log(ratio) = log(1/ratio) + + # Actually, let's think about this more carefully: + # We measured: ratio = median(d_i / d_j) + # After alignment: s_i * d_i should equal s_j * d_j in overlap + # So: s_i / s_j = d_j / d_i = 1/ratio + # In log: x_i - x_j = log(1/ratio) = -log(ratio) + + log_ratio = -math.log(max(ratio, 1e-6)) # s_i / s_j = 1/ratio + + if i == 0: + # x_0 = 0, so: -x_j = log_ratio → x_j = -log_ratio + A[k, j - 1] = -1.0 + elif j == 0: + # x_0 = 0, so: x_i = log_ratio + A[k, i - 1] = 1.0 + else: + # x_i - x_j = log_ratio + A[k, i - 1] = 1.0 + A[k, j - 1] = -1.0 + + b[k] = log_ratio + + # Step 3: Solve least squares: x = (A^T A)^{-1} A^T b + print(f"[GlobalOpt] Solving least squares system ({num_pairs} equations, {num_views-1} unknowns)") + + try: + # Use torch.linalg.lstsq for numerical stability + solution = torch.linalg.lstsq(A, b).solution # [num_views-1] + + # Convert back from log-space + scales = [1.0] # x_0 = 0 → s_0 = 1 + for i in range(num_views - 1): + scales.append(math.exp(solution[i].item())) + + print(f"[GlobalOpt] Raw scales: {[f'{s:.4f}' for s in scales]}") + + # Compute loss after optimization (before normalization) + opt_loss, opt_rmse, per_pair = compute_global_misalignment(scales, pairwise_ratios) + print(f"[GlobalOpt] Optimized misalignment: loss={opt_loss:.4f}, RMSE={opt_rmse:.4f}") + + except Exception as e: + print(f"[GlobalOpt] Least squares failed: {e}, using identity scales") + scales = [1.0] * num_views + + # Step 4: Normalize so median scale = 1.0 (prevents overall drift) + median_scale = sorted(scales)[len(scales) // 2] + if median_scale > 0: + scales = [s / median_scale for s in scales] + + print(f"[GlobalOpt] Normalized scales: {[f'{s:.4f}' for s in scales]}") + + # Compute final loss after normalization + final_loss, final_rmse, per_pair = compute_global_misalignment(scales, pairwise_ratios) + print(f"[GlobalOpt] Final misalignment: loss={final_loss:.4f}, RMSE={final_rmse:.4f}") + + # Print per-pair residuals + print(f"[GlobalOpt] Per-pair log-errors:") + for (i, j), err in sorted(per_pair.items()): + status = "OK" if err < 0.1 else "WARN" if err < 0.2 else "BAD" + print(f" Pair ({i},{j}): {err:.4f} [{status}]") + + # Summary + min_scale, max_scale = min(scales), max(scales) + variation = max_scale / min_scale if min_scale > 0 else float('inf') + print(f"[GlobalOpt] Scale range: {min_scale:.4f} - {max_scale:.4f} ({variation:.2f}x variation)") + + improvement = (initial_rmse - final_rmse) / initial_rmse * 100 if initial_rmse > 0 else 0 + print(f"[GlobalOpt] Improvement: {improvement:.1f}% reduction in RMSE") + + # Warn if variation is still high + if variation > 2.0: + print(f"[GlobalOpt] WARNING: High scale variation ({variation:.2f}x) - alignment may be unreliable") + + return {i: (scales[i], 0.0) for i in range(num_views)} + + +class AlignDepthMaps: + """Global optimization-based scale alignment for multi-view depth maps. + + Aligns depth maps to remove seams caused by scale ambiguity in + monocular depth estimation. Uses least-squares optimization in log-space + to find globally consistent scales. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "depth_maps": ("IMAGE",), + "extrinsics": ("EXTRINSICS",), + "intrinsics": ("INTRINSICS",), + }, + "optional": { + "method": (["global_optimization", "bfs_chain"], { + "default": "global_optimization", + "tooltip": "global_optimization: solve all scales jointly (recommended). bfs_chain: chain pairwise alignments (can drift)." + }), + "reference_view": ("INT", { + "default": 0, + "min": 0, + "max": 999, + "step": 1, + "tooltip": "Index of the reference view (scale=1.0)" + }), + } + } + + RETURN_TYPES = ("IMAGE", "EXTRINSICS", "INTRINSICS", "STRING",) + RETURN_NAMES = ("aligned_depth_maps", "extrinsics", "intrinsics", "alignment_info",) + FUNCTION = "align" + CATEGORY = "SHARP" + DESCRIPTION = "Align depth maps using global optimization to remove seams from monocular depth estimation." + + def align( + self, + depth_maps: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + method: str = "global_optimization", + reference_view: int = 0, + ): + """Align depth maps using global optimization or BFS chain.""" + + num_views = depth_maps.shape[0] + print(f"[AlignDepthMaps] Aligning {num_views} depth maps") + print(f"[AlignDepthMaps] Method: {method}") + + # Clamp reference view to valid range + reference_view = min(reference_view, num_views - 1) + + # Compute alignments using selected method + if method == "global_optimization": + alignments = compute_global_alignments_optimized( + depth_maps, extrinsics, intrinsics + ) + else: + # Legacy BFS method (kept for comparison) + print(f"[AlignDepthMaps] Using BFS chain (reference view: {reference_view})") + alignments = compute_global_alignments( + depth_maps, extrinsics, intrinsics, + reference_view, num_iterations=1000, inlier_threshold=0.05 + ) + + # Apply alignments + aligned_depth_maps = [] + alignment_info = {"method": method} + + for i in range(num_views): + scale, shift = alignments[i] + depth = depth_maps[i] # [H, W, C] + + # Apply alignment (scale only, shift is always 0) + aligned = scale * depth + shift + + aligned_depth_maps.append(aligned) + alignment_info[f"view_{i}"] = {"scale": scale, "shift": shift} + + # Stack result + aligned_batch = torch.stack(aligned_depth_maps, dim=0) + + # Convert alignment info to JSON string + info_str = json.dumps(alignment_info, indent=2) + + scales = [alignments[i][0] for i in range(num_views)] + min_scale, max_scale = min(scales), max(scales) + variation = max_scale / min_scale if min_scale > 0 else float('inf') + + print(f"[AlignDepthMaps] Alignment complete") + print(f"[AlignDepthMaps] Final scale range: {min_scale:.4f} - {max_scale:.4f} ({variation:.2f}x variation)") + + # Check for problematic alignments + if variation > 2.0: + print(f"[AlignDepthMaps] WARNING: Large scale variation ({variation:.2f}x) - alignment may be unreliable.") + + return (aligned_batch, extrinsics, intrinsics, info_str,) + + +NODE_CLASS_MAPPINGS = { + "AlignDepthMaps": AlignDepthMaps, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AlignDepthMaps": "Align Depth Maps (Global Optimization)", +} diff --git a/nodes/load_model.py b/nodes/load_model.py index a72ee60..dc07772 100644 --- a/nodes/load_model.py +++ b/nodes/load_model.py @@ -49,7 +49,7 @@ class LoadSharpModel: def load_model(self, device: str, checkpoint_path: str = ""): """Load and cache the SHARP model.""" - from nodes.sharp.models import PredictorParams, create_predictor + from .sharp.models import PredictorParams, create_predictor # Resolve device if device == "auto": diff --git a/nodes/merge_gaussians.py b/nodes/merge_gaussians.py new file mode 100644 index 0000000..a49fe0c --- /dev/null +++ b/nodes/merge_gaussians.py @@ -0,0 +1,241 @@ +"""MergeGaussians node for ComfyUI-Sharp. + +Merges multiple PLY files (from panorama samples) into a single unified Gaussian scene. +""" + +import os +import time +from pathlib import Path + +import numpy as np +import torch +from plyfile import PlyData, PlyElement + + +def load_ply_simple(path: str) -> dict: + """Load Gaussian data from PLY file. + + Returns dict with arrays: positions, colors, scales, rotations, opacities + """ + plydata = PlyData.read(path) + vertex = plydata['vertex'] + + positions = np.stack([ + vertex['x'], + vertex['y'], + vertex['z'] + ], axis=-1) + + # Colors (SH coefficients - we take DC term) + colors = np.stack([ + vertex['f_dc_0'], + vertex['f_dc_1'], + vertex['f_dc_2'] + ], axis=-1) + + # Scales + scales = np.stack([ + vertex['scale_0'], + vertex['scale_1'], + vertex['scale_2'] + ], axis=-1) + + # Rotations (quaternion) + rotations = np.stack([ + vertex['rot_0'], + vertex['rot_1'], + vertex['rot_2'], + vertex['rot_3'] + ], axis=-1) + + # Opacity + opacities = vertex['opacity'] + + return { + 'positions': positions, + 'colors': colors, + 'scales': scales, + 'rotations': rotations, + 'opacities': opacities, + } + + +def save_merged_ply( + positions: np.ndarray, + colors: np.ndarray, + scales: np.ndarray, + rotations: np.ndarray, + opacities: np.ndarray, + output_path: str, +): + """Save merged Gaussians to PLY file.""" + + num_gaussians = len(positions) + + # Create structured array + dtype = [ + ('x', 'f4'), ('y', 'f4'), ('z', 'f4'), + ('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4'), + ('opacity', 'f4'), + ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'), + ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4'), + ] + + elements = np.empty(num_gaussians, dtype=dtype) + elements['x'] = positions[:, 0] + elements['y'] = positions[:, 1] + elements['z'] = positions[:, 2] + elements['f_dc_0'] = colors[:, 0] + elements['f_dc_1'] = colors[:, 1] + elements['f_dc_2'] = colors[:, 2] + elements['opacity'] = opacities + elements['scale_0'] = scales[:, 0] + elements['scale_1'] = scales[:, 1] + elements['scale_2'] = scales[:, 2] + elements['rot_0'] = rotations[:, 0] + elements['rot_1'] = rotations[:, 1] + elements['rot_2'] = rotations[:, 2] + elements['rot_3'] = rotations[:, 3] + + el = PlyElement.describe(elements, 'vertex') + PlyData([el]).write(output_path) + + +class MergeGaussians: + """Merge multiple Gaussian PLY files into a single scene. + + Used after running SHARP on panorama samples to combine all views. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "ply_folder": ("STRING", { + "tooltip": "Path to folder containing PLY files to merge" + }), + }, + "optional": { + "output_prefix": ("STRING", { + "default": "merged", + "tooltip": "Prefix for output merged PLY file" + }), + "max_depth": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1000.0, + "step": 0.1, + "tooltip": "Filter out Gaussians beyond this depth (0 = no filter)" + }), + "min_opacity": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Filter out Gaussians with opacity below this (0 = no filter)" + }), + } + } + + RETURN_TYPES = ("STRING", "INT",) + RETURN_NAMES = ("ply_path", "num_gaussians",) + FUNCTION = "merge" + CATEGORY = "SHARP" + OUTPUT_NODE = True + DESCRIPTION = "Merge multiple Gaussian PLY files into a single unified scene." + + def merge( + self, + ply_folder: str, + output_prefix: str = "merged", + max_depth: float = 0.0, + min_opacity: float = 0.0, + ): + """Merge all PLY files in folder.""" + + # Find all PLY files + ply_folder = Path(ply_folder) + if not ply_folder.exists(): + raise ValueError(f"PLY folder does not exist: {ply_folder}") + + ply_files = sorted(ply_folder.glob("*.ply")) + if not ply_files: + raise ValueError(f"No PLY files found in: {ply_folder}") + + print(f"[MergeGaussians] Found {len(ply_files)} PLY files to merge") + + # Load all PLY files + all_positions = [] + all_colors = [] + all_scales = [] + all_rotations = [] + all_opacities = [] + + for i, ply_path in enumerate(ply_files): + print(f"[MergeGaussians] Loading {ply_path.name} ({i+1}/{len(ply_files)})") + data = load_ply_simple(str(ply_path)) + + positions = data['positions'] + colors = data['colors'] + scales = data['scales'] + rotations = data['rotations'] + opacities = data['opacities'] + + # Apply filters + mask = np.ones(len(positions), dtype=bool) + + if max_depth > 0: + # Filter by depth (distance from origin) + depths = np.linalg.norm(positions, axis=-1) + mask &= depths <= max_depth + + if min_opacity > 0: + mask &= opacities >= min_opacity + + filtered_count = (~mask).sum() + if filtered_count > 0: + print(f"[MergeGaussians] Filtered out {filtered_count:,} Gaussians") + + all_positions.append(positions[mask]) + all_colors.append(colors[mask]) + all_scales.append(scales[mask]) + all_rotations.append(rotations[mask]) + all_opacities.append(opacities[mask]) + + # Concatenate all + merged_positions = np.concatenate(all_positions, axis=0) + merged_colors = np.concatenate(all_colors, axis=0) + merged_scales = np.concatenate(all_scales, axis=0) + merged_rotations = np.concatenate(all_rotations, axis=0) + merged_opacities = np.concatenate(all_opacities, axis=0) + + num_gaussians = len(merged_positions) + print(f"[MergeGaussians] Total Gaussians after merge: {num_gaussians:,}") + + # Save merged PLY + timestamp = int(time.time() * 1000) + output_filename = f"{output_prefix}_{timestamp}.ply" + output_path = ply_folder.parent / output_filename + + print(f"[MergeGaussians] Saving to {output_path}") + save_merged_ply( + merged_positions, + merged_colors, + merged_scales, + merged_rotations, + merged_opacities, + str(output_path), + ) + + print(f"[MergeGaussians] Done! Merged {len(ply_files)} files into {num_gaussians:,} Gaussians") + + return (str(output_path), num_gaussians,) + + +NODE_CLASS_MAPPINGS = { + "MergeGaussians": MergeGaussians, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MergeGaussians": "Merge Gaussians (PLY Files)", +} diff --git a/nodes/predict.py b/nodes/predict.py index 440a8c8..cd746dc 100644 --- a/nodes/predict.py +++ b/nodes/predict.py @@ -49,12 +49,18 @@ class SharpPredict: "min": 0.0, "max": 500.0, "step": 0.1, - "tooltip": "Focal length in mm (35mm equivalent). 0 = auto (defaults to 30mm)." + "tooltip": "Focal length in mm (35mm equivalent). 0 = auto (defaults to 30mm). Ignored if intrinsics provided." }), "output_prefix": ("STRING", { "default": "sharp", "tooltip": "Prefix for output PLY filename or folder name for batches." }), + "extrinsics": ("EXTRINSICS", { + "tooltip": "Camera extrinsics (from SamplePanorama). If batched, must match image batch size." + }), + "intrinsics": ("INTRINSICS", { + "tooltip": "Camera intrinsics (from SamplePanorama). Overrides focal_length_mm if provided." + }), } } @@ -72,6 +78,8 @@ class SharpPredict: image: torch.Tensor, focal_length_mm: float = 0.0, output_prefix: str = "sharp", + extrinsics: torch.Tensor = None, + intrinsics: torch.Tensor = None, ): """Run SHARP inference and save PLY file(s). @@ -79,8 +87,11 @@ class SharpPredict: For batch: creates folder {prefix}_{timestamp}/ with 001.ply, 002.ply, etc. Features are cached per image - changing focal_length with same image is instant. + + If extrinsics/intrinsics are provided (from SamplePanorama), Gaussians are + unprojected into world coordinates using those camera parameters. """ - from nodes.sharp.utils.gaussians import save_ply, unproject_gaussians + from .sharp.utils.gaussians import save_ply, unproject_gaussians predictor = model["predictor"] device = torch.device(model["device"]) @@ -90,7 +101,17 @@ class SharpPredict: image = image.unsqueeze(0) batch_size = image.shape[0] - print(f"[SHARP] Processing {batch_size} image(s)") + + # Validate extrinsics batch size if provided + has_camera_params = extrinsics is not None and intrinsics is not None + if has_camera_params: + if extrinsics.dim() == 2: + extrinsics = extrinsics.unsqueeze(0) + if extrinsics.shape[0] != batch_size: + raise ValueError(f"Extrinsics batch size ({extrinsics.shape[0]}) must match image batch size ({batch_size})") + print(f"[SHARP] Processing {batch_size} image(s) with provided camera parameters (panorama mode)") + else: + print(f"[SHARP] Processing {batch_size} image(s)") # Ensure output directory exists os.makedirs(OUTPUT_DIR, exist_ok=True) @@ -125,15 +146,28 @@ class SharpPredict: if i == 0: print(f"[SHARP] Image size: {width}x{height}") - # Determine focal length in pixels - if focal_length_mm > 0: - f_px = convert_focallength(width, height, focal_length_mm) + # Get camera parameters for this image + if has_camera_params: + # Use provided intrinsics (extract focal length) + img_intrinsics = intrinsics.to(device) + img_extrinsics = extrinsics[i].to(device) + f_px = img_intrinsics[0, 0].item() # fx from intrinsics matrix else: - f_px = convert_focallength(width, height, 30.0) + # Use focal_length_mm parameter + if focal_length_mm > 0: + f_px = convert_focallength(width, height, focal_length_mm) + else: + f_px = convert_focallength(width, height, 30.0) + img_extrinsics = None + img_intrinsics = None # Run inference with caching print(f"[SHARP] Running inference on image {i+1}/{batch_size}...") - gaussians = self._predict_image_cached(predictor, image_np, f_px, device) + gaussians = self._predict_image_cached( + predictor, image_np, f_px, device, + extrinsics=img_extrinsics, + intrinsics=img_intrinsics, + ) # Determine output filename if is_batch: @@ -169,14 +203,24 @@ class SharpPredict: image: np.ndarray, f_px: float, device: torch.device, + extrinsics: torch.Tensor = None, + intrinsics: torch.Tensor = None, ): """Predict Gaussians with caching of encoded features. The expensive encode step is cached per image. Changing focal_length reuses cached features (instant). + + Args: + predictor: SHARP predictor model + image: Input image as numpy array [H, W, 3] + f_px: Focal length in pixels + device: Torch device + extrinsics: Optional 4x4 camera extrinsics (world-to-camera) + intrinsics: Optional 4x4 camera intrinsics """ global _encode_cache - from nodes.sharp.utils.gaussians import unproject_gaussians + from .sharp.utils.gaussians import unproject_gaussians internal_shape = (1536, 1536) height, width = image.shape[:2] @@ -231,26 +275,37 @@ class SharpPredict: decode_time = time.time() - decode_start print(f"[SHARP] Decode time: {decode_time:.2f}s") - # Build intrinsics for unprojection - intrinsics = ( - torch.tensor( - [ - [f_px, 0, width / 2, 0], - [0, f_px, height / 2, 0], - [0, 0, 1, 0], - [0, 0, 0, 1], - ] + # Build intrinsics for unprojection (use provided or construct from f_px) + if intrinsics is not None: + unproj_intrinsics = intrinsics.clone() + else: + unproj_intrinsics = ( + torch.tensor( + [ + [f_px, 0, width / 2, 0], + [0, f_px, height / 2, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + ] + ) + .float() + .to(device) ) - .float() - .to(device) - ) - intrinsics_resized = intrinsics.clone() + + # Scale intrinsics to internal resolution + intrinsics_resized = unproj_intrinsics.clone() intrinsics_resized[0] *= internal_shape[0] / width intrinsics_resized[1] *= internal_shape[1] / height - # Convert Gaussians to metric space + # Use provided extrinsics or identity + if extrinsics is not None: + unproj_extrinsics = extrinsics + else: + unproj_extrinsics = torch.eye(4, device=device) + + # Convert Gaussians to world/metric space gaussians = unproject_gaussians( - gaussians_ndc, torch.eye(4).to(device), intrinsics_resized, internal_shape + gaussians_ndc, unproj_extrinsics, intrinsics_resized, internal_shape ) return gaussians diff --git a/nodes/predict_depth.py b/nodes/predict_depth.py new file mode 100644 index 0000000..6271310 --- /dev/null +++ b/nodes/predict_depth.py @@ -0,0 +1,255 @@ +"""SharpPredictDepth node for ComfyUI-Sharp. + +Runs SHARP inference but outputs depth maps instead of PLY files. +Used for panorama depth visualization and verification before Gaussian creation. +""" + +import time + +import numpy as np +import torch +import torch.nn.functional as F + +from .utils.image import comfy_to_numpy_rgb + + +class SharpPredictDepth: + """Run SHARP inference to generate depth maps from images. + + Unlike SharpPredict which outputs PLY files, this node outputs + depth maps as images for visualization and further processing. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("SHARP_MODEL",), + "image": ("IMAGE",), + }, + "optional": { + "extrinsics": ("EXTRINSICS", { + "tooltip": "Camera extrinsics (from SamplePanorama). Passed through for pipeline." + }), + "intrinsics": ("INTRINSICS", { + "tooltip": "Camera intrinsics (from SamplePanorama). Used for depth scale." + }), + "reference_depth": ("IMAGE", { + "tooltip": "Reference depth maps (e.g., from DepthAnythingV3) for dense alignment. Must match image batch size." + }), + } + } + + RETURN_TYPES = ("IMAGE", "EXTRINSICS", "INTRINSICS", "IMAGE",) + RETURN_NAMES = ("depth_maps", "extrinsics", "intrinsics", "alignment_maps",) + FUNCTION = "predict_depth" + CATEGORY = "SHARP" + DESCRIPTION = "Generate depth maps from images using SHARP. Optionally align to reference depth (e.g., DepthAnythingV3) using learned dense alignment." + + @torch.no_grad() + def predict_depth( + self, + model: dict, + image: torch.Tensor, + extrinsics: torch.Tensor = None, + intrinsics: torch.Tensor = None, + reference_depth: torch.Tensor = None, + ): + """Run SHARP inference and extract depth maps. + + Returns depth maps as [N, H, W, 1] tensor (grayscale images). + If reference_depth is provided, uses SHARP's learned dense alignment. + """ + predictor = model["predictor"] + device = torch.device(model["device"]) + + # Handle batch dimension + if image.dim() == 3: + image = image.unsqueeze(0) + + batch_size = image.shape[0] + print(f"[SharpPredictDepth] Processing {batch_size} image(s)") + + # Check if we have reference depth for alignment + use_alignment = reference_depth is not None + if use_alignment: + if reference_depth.dim() == 3: + reference_depth = reference_depth.unsqueeze(0) + if reference_depth.shape[0] != batch_size: + print(f"[SharpPredictDepth] WARNING: reference_depth batch size ({reference_depth.shape[0]}) " + f"doesn't match image batch size ({batch_size}). Disabling alignment.") + use_alignment = False + else: + # Check if scale_map_estimator is available + scale_map_estimator = predictor.depth_alignment.scale_map_estimator + if scale_map_estimator is None: + print(f"[SharpPredictDepth] WARNING: scale_map_estimator not available in model. Disabling alignment.") + use_alignment = False + else: + print(f"[SharpPredictDepth] Using learned dense alignment with reference depth") + + # SHARP processes at 1536x1536 internally + # We output depth at native disparity resolution (1536x1536) + internal_shape = (1536, 1536) + + all_depth_maps = [] + all_alignment_maps = [] + + inference_start = time.time() + + for i in range(batch_size): + # Extract single image from batch + single_image = image[i:i+1] + image_np = comfy_to_numpy_rgb(single_image) + height, width = image_np.shape[:2] + + if i == 0: + print(f"[SharpPredictDepth] Image size: {width}x{height}") + + # Get focal length from intrinsics or use default + if intrinsics is not None: + f_px = intrinsics[0, 0].item() + else: + # Default: assume 65° FOV for the input image size + import math + fov_rad = math.radians(65) + f_px = (width / 2) / math.tan(fov_rad / 2) + + # Convert to tensor and normalize + image_pt = torch.from_numpy(image_np.copy()).float().to(device).permute(2, 0, 1) / 255.0 + + # Resize to internal resolution + image_resized_pt = F.interpolate( + image_pt[None], + size=(internal_shape[1], internal_shape[0]), + mode="bilinear", + align_corners=True, + ) + + # Run SHARP encode/decode + print(f"[SharpPredictDepth] Running inference on image {i+1}/{batch_size}...") + + encode_start = time.time() + monodepth_output, depth_decoder_features = predictor.encode(image_resized_pt) + encode_time = time.time() - encode_start + + # Compute disparity factor (consistent across all views with same FOV) + disparity_factor = f_px / width + + print(f"[SharpPredictDepth] Encode: {encode_time:.2f}s, disparity_factor: {disparity_factor:.4f}") + + # Extract depth from RAW DISPARITY (not Gaussians!) + # monodepth_output.disparity is [1, 1, 1536, 1536] at internal resolution + raw_disparity = monodepth_output.disparity # [1, 1, 1536, 1536] + + # Apply learned dense alignment if reference depth is provided + if use_alignment: + # Get reference depth for this image + ref_depth_single = reference_depth[i] # [H, W, C] + + # Convert to [1, 1, H, W] and resize to internal resolution + ref_depth_gray = ref_depth_single[:, :, 0] # Take first channel + ref_depth_tensor = ref_depth_gray.unsqueeze(0).unsqueeze(0).to(device) # [1, 1, H, W] + ref_depth_resized = F.interpolate( + ref_depth_tensor, + size=internal_shape, + mode="bilinear", + align_corners=True, + ) + + # Convert reference depth to same scale as SHARP disparity + # Reference depth is likely in [0, 1] range, convert to depth-like values + # Assuming reference is normalized: 1 = close, 0 = far (like our output) + # Convert to actual depth: larger value = farther + ref_depth_scaled = 1.0 / (ref_depth_resized.clamp(min=0.01) + 1e-4) + + # SHARP's alignment expects depth (not disparity), so convert + sharp_depth = disparity_factor / raw_disparity.clamp(min=1e-4) + + # Run scale_map_estimator + # It expects: tensor_src (SHARP depth), tensor_tgt (reference depth) + alignment_map = scale_map_estimator( + sharp_depth, + ref_depth_scaled, + depth_decoder_features if predictor.depth_alignment.scale_map_estimator.depth_decoder_features else None, + ) + + # Apply alignment: aligned_depth = alignment_map * sharp_depth + aligned_depth = alignment_map * sharp_depth + + # Convert back to disparity for consistent output + aligned_disparity = disparity_factor / aligned_depth.clamp(min=1e-4) + + print(f"[SharpPredictDepth] Alignment map range: {alignment_map.min():.3f} - {alignment_map.max():.3f}") + + # Store alignment map for visualization + alignment_map_vis = alignment_map[0, 0].cpu() # [H, W] + # Normalize for visualization + align_min, align_max = alignment_map_vis.min(), alignment_map_vis.max() + if align_max > align_min: + alignment_map_vis = (alignment_map_vis - align_min) / (align_max - align_min) + alignment_map_rgb = alignment_map_vis.unsqueeze(-1).repeat(1, 1, 3) + all_alignment_maps.append(alignment_map_rgb) + + # Use aligned disparity + disparity_resized = aligned_disparity[0, 0] # [1536, 1536] + else: + # No alignment - use raw disparity + disparity_resized = raw_disparity[0, 0] # [1536, 1536] + + # Create dummy alignment map (all ones) + alignment_map_rgb = torch.ones(internal_shape[0], internal_shape[1], 3) + all_alignment_maps.append(alignment_map_rgb) + + # Convert disparity to depth: depth = disparity_factor / disparity + depth_map_raw = disparity_factor / disparity_resized.clamp(min=1e-4) # [H, W] + + if i == 0: + print(f"[SharpPredictDepth] Depth range: {depth_map_raw.min():.3f} - {depth_map_raw.max():.3f}") + + # Convert to [H, W, 1] format for ComfyUI + depth_map = depth_map_raw.unsqueeze(-1).cpu() # [H, W, 1] + + # Store as single channel (will be converted to RGB later if needed) + depth_map_rgb = depth_map.repeat(1, 1, 3) # [H, W, 3] + + all_depth_maps.append(depth_map_rgb) + + inference_time = time.time() - inference_start + print(f"[SharpPredictDepth] Total inference time: {inference_time:.2f}s") + + # Stack all depth maps + depth_maps_batch = torch.stack(all_depth_maps, dim=0) # [N, H, W, 3] + alignment_maps_batch = torch.stack(all_alignment_maps, dim=0) # [N, H, W, 3] + + print(f"[SharpPredictDepth] Output shape: {depth_maps_batch.shape}") + + # Scale intrinsics to match output depth map size (1536x1536 internal resolution) + # Input intrinsics are for the original image size + if intrinsics is not None: + # Get scale factor: output_size / input_size + input_size = width # Original image width + output_size = internal_shape[0] # 1536 (native disparity resolution) + scale_factor = output_size / input_size + + if abs(scale_factor - 1.0) > 0.01: # Only scale if significantly different + print(f"[SharpPredictDepth] Scaling intrinsics by {scale_factor:.3f} " + f"(input {input_size} -> output {output_size})") + # Scale f_px, f_py, cx, cy + intrinsics_scaled = intrinsics.clone() + intrinsics_scaled[0, 0] *= scale_factor # f_px + intrinsics_scaled[1, 1] *= scale_factor # f_py + intrinsics_scaled[0, 2] *= scale_factor # cx + intrinsics_scaled[1, 2] *= scale_factor # cy + intrinsics = intrinsics_scaled + + return (depth_maps_batch, extrinsics, intrinsics, alignment_maps_batch,) + + +NODE_CLASS_MAPPINGS = { + "SharpPredictDepth": SharpPredictDepth, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SharpPredictDepth": "SHARP Predict Depth", +} diff --git a/nodes/project_depth_to_panorama.py b/nodes/project_depth_to_panorama.py new file mode 100644 index 0000000..247bc33 --- /dev/null +++ b/nodes/project_depth_to_panorama.py @@ -0,0 +1,713 @@ +"""ProjectDepthToPanorama node for ComfyUI-Sharp. + +Projects depth maps from perspective views back to equirectangular panorama +with blending and debug visualization. +""" + +import math + +import numpy as np +import torch +import torch.nn.functional as F + + +def extrinsics_to_yaw_pitch(extrinsics: torch.Tensor) -> tuple[float, float]: + """Extract yaw and pitch from extrinsics matrix. + + The extrinsics is world-to-camera, so we need the inverse (camera-to-world) + to get the camera's viewing direction. + """ + # Camera-to-world is inverse of world-to-camera + # For pure rotation, inverse = transpose + R_c2w = extrinsics[:3, :3].T + + # Camera looks along +Z in camera space + # Transform to world space + forward = R_c2w @ torch.tensor([0., 0., 1.], device=extrinsics.device) + + # Extract yaw (rotation around Y) and pitch (rotation around X) + yaw = math.atan2(forward[0].item(), forward[2].item()) + pitch = math.asin(torch.clamp(forward[1], -1, 1).item()) + + return yaw, pitch + + +def compute_blend_weight( + uu: torch.Tensor, + vv: torch.Tensor, + depth_w: int, + depth_h: int, + half_fov: float, + blend_mode: str, +) -> torch.Tensor: + """Compute blending weights based on distance from view center. + + Args: + uu, vv: Pixel coordinate grids + depth_w, depth_h: Depth map dimensions + half_fov: Half field of view in radians + blend_mode: Blending algorithm to use + + Returns: + Weight tensor [H, W] + """ + # Normalized distance from center (0 at center, 1 at edge) + dist_from_center = torch.sqrt( + ((uu - depth_w / 2) / (depth_w / 2)) ** 2 + + ((vv - depth_h / 2) / (depth_h / 2)) ** 2 + ) + + if blend_mode == "cosine": + # Cosine falloff (original) + angle_from_center = dist_from_center * half_fov + weight = torch.cos(torch.clamp(angle_from_center, 0, half_fov)) + + elif blend_mode == "gaussian": + # Gaussian falloff (smoother) + sigma = 0.5 # Standard deviation in normalized units + weight = torch.exp(-0.5 * (dist_from_center / sigma) ** 2) + + elif blend_mode == "linear": + # Linear falloff + weight = 1.0 - torch.clamp(dist_from_center, 0, 1) + + elif blend_mode == "quadratic": + # Quadratic falloff (smoother than linear) + weight = 1.0 - torch.clamp(dist_from_center, 0, 1) ** 2 + + elif blend_mode == "feather": + # Hard center with feathered edges (feather starts at 70% from center) + feather_start = 0.7 + feather_width = 1.0 - feather_start + weight = torch.where( + dist_from_center < feather_start, + torch.ones_like(dist_from_center), + 1.0 - (dist_from_center - feather_start) / feather_width + ) + + else: # "none" or fallback + weight = torch.ones_like(dist_from_center) + + return torch.clamp(weight, 0.01, 1.0) + + +def project_depth_maps_to_panorama_with_disagreement( + depth_maps: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + output_width: int, + blend_mode: str = "gaussian", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Project depth maps back to equirectangular panorama with disagreement tracking. + + Args: + depth_maps: [N, H, W, C] depth maps from each view + extrinsics: [N, 4, 4] camera extrinsics for each view + intrinsics: [4, 4] shared camera intrinsics + output_width: Width of output panorama (height = width/2) + blend_mode: Blending algorithm ("none", "cosine", "gaussian", "linear", "quadratic", "feather") + + Returns: + Tuple of: + - Panoramic depth map [output_height, output_width, 1] + - Disagreement map [output_height, output_width, 1] (std dev of depths) + - Overlap count [output_height, output_width] (number of views per pixel) + """ + device = depth_maps.device + num_views = depth_maps.shape[0] + depth_h, depth_w = depth_maps.shape[1:3] + output_height = output_width // 2 + + # Get FOV from intrinsics + f_px = intrinsics[0, 0].item() + cx = intrinsics[0, 2].item() + half_fov = math.atan((depth_w / 2) / f_px) + + # Create output panorama accumulators + # For Welford's online variance algorithm + panorama_depth = torch.zeros(output_height, output_width, device=device) # weighted mean + panorama_weight = torch.zeros(output_height, output_width, device=device) + panorama_m2 = torch.zeros(output_height, output_width, device=device) # sum of squared differences + panorama_count = torch.zeros(output_height, output_width, device=device) # view count + + # Store per-view contributions for disagreement calculation + # We'll use a simpler approach: track sum, sum of squares, and count + panorama_sum = torch.zeros(output_height, output_width, device=device) + panorama_sum_sq = torch.zeros(output_height, output_width, device=device) + + # For each view, project its depth to panorama + for view_idx in range(num_views): + depth_map = depth_maps[view_idx, :, :, 0] # [H, W] + ext = extrinsics[view_idx] + + # Get view direction + view_yaw, view_pitch = extrinsics_to_yaw_pitch(ext) + + # Create pixel grid for depth map + u = torch.arange(depth_w, dtype=torch.float32, device=device) + v = torch.arange(depth_h, dtype=torch.float32, device=device) + uu, vv = torch.meshgrid(u, v, indexing='xy') + + # Convert to camera-space ray directions + dx = (uu - cx) / f_px + dy = (vv - cx) / f_px # Assuming square pixels, cy = cx + dz = torch.ones_like(dx) + + rays_cam = torch.stack([dx, dy, dz], dim=-1) + rays_cam = F.normalize(rays_cam, dim=-1) + + # Rotate to world space (camera-to-world) + R_c2w = ext[:3, :3].T + rays_world = torch.einsum('ij,hwj->hwi', R_c2w, rays_cam) + + # Convert world rays to spherical (yaw, pitch) + rx, ry, rz = rays_world[..., 0], rays_world[..., 1], rays_world[..., 2] + ray_yaw = torch.atan2(rx, rz) + ray_pitch = torch.asin(torch.clamp(ry, -1, 1)) + + # Map to panorama pixel coordinates + pano_x = ((ray_yaw / math.pi + 1) * (output_width - 1) / 2).long() + pano_y = ((0.5 - ray_pitch / math.pi) * (output_height - 1)).long() + + # Clamp to valid range + pano_x = torch.clamp(pano_x, 0, output_width - 1) + pano_y = torch.clamp(pano_y, 0, output_height - 1) + + # Compute weight for blending + weight = compute_blend_weight(uu, vv, depth_w, depth_h, half_fov, blend_mode) + + # Scatter depth values to panorama + flat_pano_idx = pano_y * output_width + pano_x + flat_depth = depth_map.flatten() + flat_weight = weight.flatten() + + # Accumulate weighted depth + panorama_depth.view(-1).scatter_add_(0, flat_pano_idx.flatten(), flat_depth * flat_weight) + panorama_weight.view(-1).scatter_add_(0, flat_pano_idx.flatten(), flat_weight.flatten()) + + # Accumulate for variance calculation (unweighted for simplicity) + panorama_sum.view(-1).scatter_add_(0, flat_pano_idx.flatten(), flat_depth) + panorama_sum_sq.view(-1).scatter_add_(0, flat_pano_idx.flatten(), flat_depth ** 2) + panorama_count.view(-1).scatter_add_(0, flat_pano_idx.flatten(), torch.ones_like(flat_depth)) + + # Normalize by accumulated weights + panorama_depth = panorama_depth / (panorama_weight + 1e-8) + + # Compute variance: var = E[X^2] - E[X]^2 + # Only meaningful where count >= 2 + mean_depth = panorama_sum / (panorama_count + 1e-8) + mean_sq_depth = panorama_sum_sq / (panorama_count + 1e-8) + variance = mean_sq_depth - mean_depth ** 2 + variance = torch.clamp(variance, min=0) # Numerical stability + std_dev = torch.sqrt(variance) + + # Compute relative disagreement (normalized by mean depth) + relative_disagreement = std_dev / (mean_depth + 1e-8) + + # Print statistics + overlap_mask = panorama_count >= 2 + if overlap_mask.any(): + overlap_std = std_dev[overlap_mask] + overlap_rel = relative_disagreement[overlap_mask] + print(f"[ProjectDepthToPanorama] Overlap regions ({overlap_mask.sum().item()} pixels):") + print(f" Absolute disagreement: mean={overlap_std.mean():.4f}, max={overlap_std.max():.4f}") + print(f" Relative disagreement: mean={overlap_rel.mean():.2%}, max={overlap_rel.max():.2%}") + + # Count pixels with high disagreement (>10% relative) + high_disagreement = (relative_disagreement > 0.1) & overlap_mask + print(f" High disagreement (>10%): {high_disagreement.sum().item()} pixels " + f"({high_disagreement.sum().item() / overlap_mask.sum().item() * 100:.1f}% of overlap)") + + # Global normalization for visualization + valid_mask = panorama_weight > 0.01 + depth_min, depth_max = 0.0, 1.0 + if valid_mask.any(): + valid_depths = panorama_depth[valid_mask] + depth_min = valid_depths.min() + depth_max = valid_depths.max() + print(f"[ProjectDepthToPanorama] Depth range: {depth_min:.4f} to {depth_max:.4f}") + + if depth_max > depth_min: + panorama_depth = torch.where( + valid_mask, + 1.0 - (panorama_depth - depth_min) / (depth_max - depth_min), + torch.zeros_like(panorama_depth) + ) + else: + panorama_depth = torch.where( + valid_mask, + torch.full_like(panorama_depth, 0.5), + torch.zeros_like(panorama_depth) + ) + + return ( + panorama_depth.unsqueeze(-1), + relative_disagreement.unsqueeze(-1), + panorama_count, + ) + + +def project_depth_maps_to_panorama( + depth_maps: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + output_width: int, + blend_mode: str = "gaussian", +) -> torch.Tensor: + """Project depth maps back to equirectangular panorama. + + Args: + depth_maps: [N, H, W, C] depth maps from each view + extrinsics: [N, 4, 4] camera extrinsics for each view + intrinsics: [4, 4] shared camera intrinsics + output_width: Width of output panorama (height = width/2) + blend_mode: Blending algorithm ("none", "cosine", "gaussian", "linear", "quadratic", "feather") + + Returns: + Panoramic depth map [output_height, output_width, 1] + """ + depth, _, _ = project_depth_maps_to_panorama_with_disagreement( + depth_maps, extrinsics, intrinsics, output_width, blend_mode + ) + return depth + + +def get_disagreement_color(d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert normalized disagreement [0,1] to RGB color. + + Color gradient: blue (0) -> cyan (0.25) -> green (0.5) -> yellow (0.75) -> red (1.0) + """ + # Red channel: ramps up from 0.5 to 1.0 + r = torch.where(d < 0.5, 2 * d, torch.ones_like(d)) + + # Green channel: ramps up from 0 to 0.5, then down from 0.5 to 1.0 + g = torch.where(d < 0.5, 2 * d, 2 * (1 - d)) + + # Blue channel: ramps down from 0 to 0.5 + b = torch.where(d < 0.5, 1 - 2 * d, torch.zeros_like(d)) + + return r, g, b + + +def draw_legend( + heatmap: torch.Tensor, + max_disagreement: float, + legend_width: int = 30, + margin: int = 10, +) -> torch.Tensor: + """Draw a color legend bar on the right side of the heatmap. + + Args: + heatmap: [H, W, 3] input heatmap + max_disagreement: Maximum disagreement value for scale + legend_width: Width of the legend bar in pixels + margin: Margin around the legend + + Returns: + [H, W + legend_width + margin*2, 3] heatmap with legend + """ + device = heatmap.device + H, W = heatmap.shape[:2] + + # Create extended canvas + total_width = W + legend_width + margin * 3 + canvas = torch.zeros(H, total_width, 3, device=device) + + # Copy original heatmap + canvas[:, :W, :] = heatmap + + # Legend area + legend_x_start = W + margin + legend_x_end = legend_x_start + legend_width + + # Color bar region (leave space for labels at top and bottom) + label_height = 40 + bar_y_start = margin + label_height + bar_y_end = H - margin - label_height + bar_height = bar_y_end - bar_y_start + + if bar_height > 0: + # Draw gradient color bar + for y in range(bar_y_start, bar_y_end): + # Map y to disagreement value (top = max, bottom = 0) + t = 1.0 - (y - bar_y_start) / bar_height + d = torch.tensor([t], device=device) + r, g, b = get_disagreement_color(d) + + canvas[y, legend_x_start:legend_x_end, 0] = r.item() + canvas[y, legend_x_start:legend_x_end, 1] = g.item() + canvas[y, legend_x_start:legend_x_end, 2] = b.item() + + # Draw border around color bar + canvas[bar_y_start:bar_y_end, legend_x_start, :] = 1.0 # Left + canvas[bar_y_start:bar_y_end, legend_x_end-1, :] = 1.0 # Right + canvas[bar_y_start, legend_x_start:legend_x_end, :] = 1.0 # Top + canvas[bar_y_end-1, legend_x_start:legend_x_end, :] = 1.0 # Bottom + + # Draw tick marks and approximate label positions + # We'll draw simple markers since we can't render text + tick_positions = [0.0, 0.25, 0.5, 0.75, 1.0] + tick_width = 5 + + for t in tick_positions: + y = int(bar_y_start + (1 - t) * (bar_height - 1)) + # Draw tick mark + canvas[y, legend_x_end:legend_x_end + tick_width, :] = 1.0 + + # Draw "single view" indicator (dark gray square) + indicator_y = H - margin - 20 + indicator_size = 15 + canvas[indicator_y:indicator_y+indicator_size, + legend_x_start:legend_x_start+indicator_size, :] = 0.2 + # Border + canvas[indicator_y, legend_x_start:legend_x_start+indicator_size, :] = 0.5 + canvas[indicator_y+indicator_size-1, legend_x_start:legend_x_start+indicator_size, :] = 0.5 + canvas[indicator_y:indicator_y+indicator_size, legend_x_start, :] = 0.5 + canvas[indicator_y:indicator_y+indicator_size, legend_x_start+indicator_size-1, :] = 0.5 + + # Add text labels using simple block letters + # Top label: max value (e.g., "20%") + pct = int(max_disagreement * 100) + draw_text_simple(canvas, f"{pct}%", legend_x_start, margin + 5, device) + draw_text_simple(canvas, "BAD", legend_x_start, margin + 20, device) + + # Bottom label + draw_text_simple(canvas, "0%", legend_x_start, bar_y_end + 5, device) + draw_text_simple(canvas, "GOOD", legend_x_start, bar_y_end + 20, device) + + # Single view label + draw_text_simple(canvas, "1 VIEW", legend_x_start + 18, indicator_y + 3, device) + + return canvas + + +def draw_text_simple(canvas: torch.Tensor, text: str, x: int, y: int, device: torch.device): + """Draw simple blocky text on canvas. Very basic 3x5 pixel font.""" + # Simple 3x5 pixel font for basic characters + font = { + '0': [[1,1,1], [1,0,1], [1,0,1], [1,0,1], [1,1,1]], + '1': [[0,1,0], [1,1,0], [0,1,0], [0,1,0], [1,1,1]], + '2': [[1,1,1], [0,0,1], [1,1,1], [1,0,0], [1,1,1]], + '3': [[1,1,1], [0,0,1], [1,1,1], [0,0,1], [1,1,1]], + '4': [[1,0,1], [1,0,1], [1,1,1], [0,0,1], [0,0,1]], + '5': [[1,1,1], [1,0,0], [1,1,1], [0,0,1], [1,1,1]], + '6': [[1,1,1], [1,0,0], [1,1,1], [1,0,1], [1,1,1]], + '7': [[1,1,1], [0,0,1], [0,0,1], [0,0,1], [0,0,1]], + '8': [[1,1,1], [1,0,1], [1,1,1], [1,0,1], [1,1,1]], + '9': [[1,1,1], [1,0,1], [1,1,1], [0,0,1], [1,1,1]], + '%': [[1,0,1], [0,0,1], [0,1,0], [1,0,0], [1,0,1]], + 'B': [[1,1,0], [1,0,1], [1,1,0], [1,0,1], [1,1,0]], + 'A': [[0,1,0], [1,0,1], [1,1,1], [1,0,1], [1,0,1]], + 'D': [[1,1,0], [1,0,1], [1,0,1], [1,0,1], [1,1,0]], + 'G': [[1,1,1], [1,0,0], [1,0,1], [1,0,1], [1,1,1]], + 'O': [[1,1,1], [1,0,1], [1,0,1], [1,0,1], [1,1,1]], + 'V': [[1,0,1], [1,0,1], [1,0,1], [0,1,0], [0,1,0]], + 'I': [[1,1,1], [0,1,0], [0,1,0], [0,1,0], [1,1,1]], + 'E': [[1,1,1], [1,0,0], [1,1,0], [1,0,0], [1,1,1]], + 'W': [[1,0,1], [1,0,1], [1,1,1], [1,1,1], [1,0,1]], + ' ': [[0,0,0], [0,0,0], [0,0,0], [0,0,0], [0,0,0]], + } + + H, W = canvas.shape[:2] + cursor_x = x + scale = 2 # Scale up the tiny font + + for char in text.upper(): + if char in font: + glyph = font[char] + for row_idx, row in enumerate(glyph): + for col_idx, pixel in enumerate(row): + if pixel: + px = cursor_x + col_idx * scale + py = y + row_idx * scale + # Draw scaled pixel + for dy in range(scale): + for dx in range(scale): + if 0 <= py + dy < H and 0 <= px + dx < W: + canvas[py + dy, px + dx, :] = 1.0 + cursor_x += (len(glyph[0]) + 1) * scale + + +def create_disagreement_heatmap( + disagreement: torch.Tensor, + overlap_count: torch.Tensor, + max_disagreement: float = 0.2, + add_legend: bool = True, +) -> torch.Tensor: + """Create a color heatmap showing depth disagreement in overlap regions. + + Args: + disagreement: [H, W, 1] relative disagreement (std/mean) + overlap_count: [H, W] number of views per pixel + max_disagreement: Maximum disagreement for color scale (clips above) + add_legend: Whether to add a color legend bar + + Returns: + [H, W, 3] or [H, W+legend_width, 3] RGB heatmap where: + - Black = no overlap (single view or no coverage) + - Blue = low disagreement (good alignment) + - Green = medium disagreement + - Yellow = high disagreement + - Red = very high disagreement (bad alignment) + """ + device = disagreement.device + H, W = disagreement.shape[:2] + + # Normalize disagreement to [0, 1] + disagreement_norm = (disagreement[:, :, 0] / max_disagreement).clamp(0, 1) + + # Create RGB output + heatmap = torch.zeros(H, W, 3, device=device) + + # Mask: only show overlap regions (2+ views) + overlap_mask = overlap_count >= 2 + + # Get colors + r, g, b = get_disagreement_color(disagreement_norm) + + # Apply mask - only show in overlap regions + heatmap[:, :, 0] = torch.where(overlap_mask, r, torch.zeros_like(r)) + heatmap[:, :, 1] = torch.where(overlap_mask, g, torch.zeros_like(g)) + heatmap[:, :, 2] = torch.where(overlap_mask, b, torch.zeros_like(b)) + + # Make non-overlap regions dark gray so you can see coverage + non_overlap = (overlap_count == 1) + heatmap[:, :, 0] = torch.where(non_overlap, torch.full_like(r, 0.2), heatmap[:, :, 0]) + heatmap[:, :, 1] = torch.where(non_overlap, torch.full_like(g, 0.2), heatmap[:, :, 1]) + heatmap[:, :, 2] = torch.where(non_overlap, torch.full_like(b, 0.2), heatmap[:, :, 2]) + + # Add legend + if add_legend: + heatmap = draw_legend(heatmap, max_disagreement) + + return heatmap + + +def draw_sample_borders( + panorama: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + border_color: tuple = (1.0, 0.0, 0.0), + border_width: int = 2, +) -> torch.Tensor: + """Draw borders around sample regions on the panorama. + + Args: + panorama: [H, W, C] panorama image + extrinsics: [N, 4, 4] camera extrinsics for each view + intrinsics: [4, 4] shared camera intrinsics + border_color: RGB color for borders + border_width: Width of border lines in pixels + + Returns: + Panorama with borders drawn [H, W, 3] + """ + device = panorama.device + H, W = panorama.shape[:2] + num_views = extrinsics.shape[0] + + # Convert grayscale to RGB if needed + if panorama.shape[-1] == 1: + panorama_rgb = panorama.repeat(1, 1, 3) + else: + panorama_rgb = panorama.clone() + + # Get FOV from intrinsics + # The intrinsics encode the image size via principal point (cx = width/2) + f_px = intrinsics[0, 0].item() + cx = intrinsics[0, 2].item() + depth_size = int(cx * 2) # Derive from principal point (cx = width/2) + half_fov = math.atan((depth_size / 2) / f_px) + + border_color_t = torch.tensor(border_color, device=device) + + for view_idx in range(num_views): + ext = extrinsics[view_idx] + view_yaw, view_pitch = extrinsics_to_yaw_pitch(ext) + + # Draw border of this view's coverage area + # The view covers a square region of angular size 2*half_fov + + # Sample points along the border + num_border_points = 100 + + for edge in ['top', 'bottom', 'left', 'right']: + if edge == 'top': + us = torch.linspace(-1, 1, num_border_points) + vs = torch.ones(num_border_points) * -1 + elif edge == 'bottom': + us = torch.linspace(-1, 1, num_border_points) + vs = torch.ones(num_border_points) * 1 + elif edge == 'left': + us = torch.ones(num_border_points) * -1 + vs = torch.linspace(-1, 1, num_border_points) + else: # right + us = torch.ones(num_border_points) * 1 + vs = torch.linspace(-1, 1, num_border_points) + + # Convert to camera ray directions + dx = us * math.tan(half_fov) + dy = vs * math.tan(half_fov) + dz = torch.ones_like(dx) + + rays_cam = torch.stack([dx, dy, dz], dim=-1).to(device) + rays_cam = F.normalize(rays_cam, dim=-1) + + # Rotate to world space + R_c2w = ext[:3, :3].T + rays_world = rays_cam @ R_c2w.T + + # Convert to spherical + rx, ry, rz = rays_world[:, 0], rays_world[:, 1], rays_world[:, 2] + ray_yaw = torch.atan2(rx, rz) + ray_pitch = torch.asin(torch.clamp(ry, -1, 1)) + + # Map to panorama pixels + pano_x = ((ray_yaw / math.pi + 1) * (W - 1) / 2).long() + pano_y = ((0.5 - ray_pitch / math.pi) * (H - 1)).long() + + # Draw border pixels + for bw in range(-border_width // 2, border_width // 2 + 1): + px = torch.clamp(pano_x + bw, 0, W - 1) + py = torch.clamp(pano_y, 0, H - 1) + panorama_rgb[py, px] = border_color_t + + px = torch.clamp(pano_x, 0, W - 1) + py = torch.clamp(pano_y + bw, 0, H - 1) + panorama_rgb[py, px] = border_color_t + + return panorama_rgb + + +class ProjectDepthToPanorama: + """Project depth maps from perspective views back to equirectangular panorama. + + Combines depth maps with blending and optional debug visualization. + Outputs a disagreement heatmap showing where overlapping views have mismatching depths. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "depth_maps": ("IMAGE",), + "extrinsics": ("EXTRINSICS",), + "intrinsics": ("INTRINSICS",), + }, + "optional": { + "panorama_width": ("INT", { + "default": 2048, + "min": 512, + "max": 8192, + "step": 64, + "tooltip": "Output panorama width (height = width/2)" + }), + "blend_mode": (["gaussian", "cosine", "quadratic", "linear", "feather", "none"], { + "default": "gaussian", + "tooltip": "Blending algorithm: gaussian (smoothest), cosine, quadratic, linear, feather (hard center), none" + }), + "show_borders": ("BOOLEAN", { + "default": True, + "tooltip": "Draw borders around sample regions for debugging" + }), + "border_color": (["red", "green", "blue", "white", "yellow"], { + "default": "red", + "tooltip": "Color for sample region borders" + }), + "disagreement_scale": ("FLOAT", { + "default": 0.2, + "min": 0.01, + "max": 1.0, + "step": 0.01, + "tooltip": "Max relative disagreement for heatmap color scale (0.2 = 20%)" + }), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE",) + RETURN_NAMES = ("panoramic_depth", "debug_overlay", "disagreement_heatmap",) + FUNCTION = "project" + CATEGORY = "SHARP" + DESCRIPTION = "Project depth maps to panorama. Outputs disagreement heatmap showing where views mismatch (blue=good, red=bad)." + + def project( + self, + depth_maps: torch.Tensor, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + panorama_width: int = 2048, + blend_mode: str = "gaussian", + show_borders: bool = True, + border_color: str = "red", + disagreement_scale: float = 0.2, + ): + """Project depth maps to panorama with disagreement visualization.""" + + print(f"[ProjectDepthToPanorama] Input: {depth_maps.shape[0]} depth maps") + print(f"[ProjectDepthToPanorama] Output panorama: {panorama_width} x {panorama_width // 2}") + print(f"[ProjectDepthToPanorama] Blend mode: {blend_mode}") + + # Project depth maps with disagreement tracking + panoramic_depth, disagreement, overlap_count = project_depth_maps_to_panorama_with_disagreement( + depth_maps, + extrinsics, + intrinsics, + panorama_width, + blend_mode, + ) + + # Convert to RGB for output + panoramic_depth_rgb = panoramic_depth.repeat(1, 1, 3) + + # Create debug overlay with borders + if show_borders: + color_map = { + "red": (1.0, 0.0, 0.0), + "green": (0.0, 1.0, 0.0), + "blue": (0.0, 0.0, 1.0), + "white": (1.0, 1.0, 1.0), + "yellow": (1.0, 1.0, 0.0), + } + debug_overlay = draw_sample_borders( + panoramic_depth_rgb.clone(), + extrinsics, + intrinsics, + border_color=color_map.get(border_color, (1.0, 0.0, 0.0)), + ) + else: + debug_overlay = panoramic_depth_rgb.clone() + + # Create disagreement heatmap + disagreement_heatmap = create_disagreement_heatmap( + disagreement, + overlap_count, + max_disagreement=disagreement_scale, + ) + + # Add borders to heatmap too for reference + if show_borders: + disagreement_heatmap = draw_sample_borders( + disagreement_heatmap, + extrinsics, + intrinsics, + border_color=(1.0, 1.0, 1.0), # White borders on heatmap + border_width=1, + ) + + # Add batch dimension for ComfyUI + panoramic_depth_out = panoramic_depth_rgb.unsqueeze(0) # [1, H, W, 3] + debug_overlay_out = debug_overlay.unsqueeze(0) # [1, H, W, 3] + disagreement_out = disagreement_heatmap.unsqueeze(0) # [1, H, W, 3] + + print(f"[ProjectDepthToPanorama] Output shape: {panoramic_depth_out.shape}") + + return (panoramic_depth_out, debug_overlay_out, disagreement_out,) + + +NODE_CLASS_MAPPINGS = { + "ProjectDepthToPanorama": ProjectDepthToPanorama, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ProjectDepthToPanorama": "Project Depth to Panorama", +} diff --git a/nodes/sample_panorama.py b/nodes/sample_panorama.py new file mode 100644 index 0000000..14189de --- /dev/null +++ b/nodes/sample_panorama.py @@ -0,0 +1,279 @@ +"""SamplePanorama node for ComfyUI-Sharp. + +Samples perspective cutouts from an equirectangular panorama for use with SHARP. +""" + +import math + +import numpy as np +import torch +import torch.nn.functional as F + + +def create_rotation_matrix(yaw: float, pitch: float) -> torch.Tensor: + """Create a rotation matrix from yaw and pitch angles. + + Args: + yaw: Horizontal angle in radians (0 = forward, positive = right) + pitch: Vertical angle in radians (0 = horizon, positive = up) + + Returns: + 3x3 rotation matrix (camera-to-world) + """ + # Rotation around Y axis (yaw - looking left/right) + cy, sy = math.cos(yaw), math.sin(yaw) + R_yaw = torch.tensor([ + [cy, 0, sy], + [0, 1, 0], + [-sy, 0, cy] + ], dtype=torch.float32) + + # Rotation around X axis (pitch - looking up/down) + cp, sp = math.cos(pitch), math.sin(pitch) + R_pitch = torch.tensor([ + [1, 0, 0], + [0, cp, -sp], + [0, sp, cp] + ], dtype=torch.float32) + + # Combined: first pitch, then yaw + R = R_yaw @ R_pitch + return R + + +def sample_perspective_from_equirectangular( + panorama: torch.Tensor, + yaw: float, + pitch: float, + fov_radians: float, + output_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample a perspective view from an equirectangular panorama. + + Args: + panorama: Equirectangular image [H, W, 3] or [B, H, W, 3] + yaw: Horizontal angle in radians (0 = center, positive = right) + pitch: Vertical angle in radians (0 = horizon, positive = up) + fov_radians: Field of view in radians + output_size: Output image size (square) + + Returns: + Tuple of (perspective_image, extrinsics_4x4, intrinsics_4x4) + """ + if panorama.dim() == 4: + panorama = panorama[0] # Take first if batched + + H, W, C = panorama.shape + device = panorama.device + + # Compute focal length in pixels + f_px = (output_size / 2) / math.tan(fov_radians / 2) + + # Principal point at center + cx = (output_size - 1) / 2 + cy = (output_size - 1) / 2 + + # Create pixel grid for output image + u = torch.arange(output_size, dtype=torch.float32, device=device) + v = torch.arange(output_size, dtype=torch.float32, device=device) + uu, vv = torch.meshgrid(u, v, indexing='xy') # [H, W] + + # Convert to camera-space ray directions + # d_cam = [(u - cx) / f, (v - cy) / f, 1] + dx = (uu - cx) / f_px + dy = (vv - cy) / f_px + dz = torch.ones_like(dx) + + # Stack and normalize + rays_cam = torch.stack([dx, dy, dz], dim=-1) # [H, W, 3] + rays_cam = F.normalize(rays_cam, dim=-1) + + # Rotate rays to world space + R = create_rotation_matrix(yaw, pitch).to(device) # [3, 3] + rays_world = torch.einsum('ij,hwj->hwi', R, rays_cam) # [H, W, 3] + + # Convert world rays to spherical coordinates + # x = right, y = up, z = forward + rx, ry, rz = rays_world[..., 0], rays_world[..., 1], rays_world[..., 2] + + # Yaw: angle around Y axis (horizontal) + ray_yaw = torch.atan2(rx, rz) # [-pi, pi] + + # Pitch: angle from horizon (vertical) + ray_pitch = torch.asin(torch.clamp(ry, -1, 1)) # [-pi/2, pi/2] + + # Map to equirectangular pixel coordinates + # Yaw: -pi..pi -> 0..W + # Pitch: -pi/2..pi/2 -> H..0 (note: inverted, top of image is up) + eq_x = (ray_yaw / math.pi + 1) * (W - 1) / 2 + eq_y = (0.5 - ray_pitch / math.pi) * (H - 1) + + # Normalize to [-1, 1] for grid_sample + grid_x = eq_x / (W - 1) * 2 - 1 + grid_y = eq_y / (H - 1) * 2 - 1 + grid = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0) # [1, H, W, 2] + + # Sample from panorama + panorama_nchw = panorama.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W] + sampled = F.grid_sample( + panorama_nchw, + grid, + mode='bilinear', + padding_mode='border', + align_corners=True + ) + perspective = sampled[0].permute(1, 2, 0) # [H, W, C] + + # Build intrinsics matrix (4x4) + intrinsics = torch.tensor([ + [f_px, 0, cx, 0], + [0, f_px, cy, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + ], dtype=torch.float32, device=device) + + # Build extrinsics matrix (4x4, world-to-camera) + # World-to-camera is the inverse (transpose for rotation) + R_w2c = R.T + extrinsics = torch.eye(4, dtype=torch.float32, device=device) + extrinsics[:3, :3] = R_w2c + # No translation (camera at origin) + + return perspective, extrinsics, intrinsics + + +class SamplePanorama: + """Sample perspective cutouts from an equirectangular panorama. + + Automatically calculates the number of samples needed to cover the full + 360° x 180° panorama based on FOV and overlap settings. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "panorama": ("IMAGE",), + "fov_degrees": ("FLOAT", { + "default": 65.0, + "min": 30.0, + "max": 120.0, + "step": 1.0, + "tooltip": "Field of view in degrees for each perspective cutout" + }), + "overlap_percent": ("FLOAT", { + "default": 10.0, + "min": 0.0, + "max": 50.0, + "step": 1.0, + "tooltip": "Overlap between adjacent samples as percentage of FOV" + }), + "output_size": ("INT", { + "default": 1536, + "min": 256, + "max": 2048, + "step": 64, + "tooltip": "Output resolution for each perspective image (square). 1536 is SHARP's native resolution." + }), + }, + "optional": { + "skip_poles": ("BOOLEAN", { + "default": True, + "tooltip": "Skip samples pointing straight up/down (often low quality in panoramas)" + }), + } + } + + RETURN_TYPES = ("IMAGE", "EXTRINSICS", "INTRINSICS", "INT", "INT") + RETURN_NAMES = ("images", "extrinsics", "intrinsics", "num_horizontal", "num_vertical") + FUNCTION = "sample" + CATEGORY = "SHARP" + DESCRIPTION = "Sample perspective views from a 360° equirectangular panorama for SHARP processing." + + def sample( + self, + panorama: torch.Tensor, + fov_degrees: float = 65.0, + overlap_percent: float = 10.0, + output_size: int = 1536, + skip_poles: bool = True, + ): + """Sample perspective cutouts covering the full panorama.""" + + fov_radians = math.radians(fov_degrees) + + # Calculate step angle based on overlap + step_degrees = fov_degrees * (1 - overlap_percent / 100) + step_radians = math.radians(step_degrees) + + # Calculate number of samples needed + num_horizontal = math.ceil(360 / step_degrees) + + # For vertical, we cover -90° to +90° (180° total) + # But if skip_poles, we limit to avoid extreme angles + if skip_poles: + # Skip top and bottom 15° + vertical_range = 150 # degrees + vertical_start = -75 # degrees + else: + vertical_range = 180 + vertical_start = -90 + + num_vertical = max(1, math.ceil(vertical_range / step_degrees)) + + print(f"[SamplePanorama] FOV: {fov_degrees}°, Overlap: {overlap_percent}%") + print(f"[SamplePanorama] Step angle: {step_degrees:.1f}°") + print(f"[SamplePanorama] Samples: {num_horizontal} horizontal × {num_vertical} vertical = {num_horizontal * num_vertical} total") + + # Handle batch dimension + if panorama.dim() == 3: + panorama = panorama.unsqueeze(0) + + # Take first image if batch + pano = panorama[0] # [H, W, C] + + all_images = [] + all_extrinsics = [] + intrinsics = None + + for v_idx in range(num_vertical): + # Calculate pitch angle + pitch_degrees = vertical_start + (v_idx + 0.5) * step_degrees + pitch_radians = math.radians(pitch_degrees) + + for h_idx in range(num_horizontal): + # Calculate yaw angle + yaw_degrees = -180 + (h_idx + 0.5) * step_degrees + yaw_radians = math.radians(yaw_degrees) + + # Sample perspective view + perspective, extrinsics, intr = sample_perspective_from_equirectangular( + pano, + yaw_radians, + pitch_radians, + fov_radians, + output_size, + ) + + all_images.append(perspective) + all_extrinsics.append(extrinsics) + + if intrinsics is None: + intrinsics = intr + + # Stack results + images_batch = torch.stack(all_images, dim=0) # [N, H, W, C] + extrinsics_batch = torch.stack(all_extrinsics, dim=0) # [N, 4, 4] + + print(f"[SamplePanorama] Output shape: {images_batch.shape}") + + return (images_batch, extrinsics_batch, intrinsics, num_horizontal, num_vertical) + + +NODE_CLASS_MAPPINGS = { + "SamplePanorama": SamplePanorama, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SamplePanorama": "Sample Panorama (Equirect → Perspective)", +} diff --git a/nodes/sharp/cli/predict.py b/nodes/sharp/cli/predict.py index 597d8d9..cd52e42 100644 --- a/nodes/sharp/cli/predict.py +++ b/nodes/sharp/cli/predict.py @@ -15,14 +15,14 @@ import torch import torch.nn.functional as F import torch.utils.data -from nodes.sharp.models import ( +from ..models import ( PredictorParams, RGBGaussianPredictor, create_predictor, ) -from nodes.sharp.utils import io -from nodes.sharp.utils import logging as logging_utils -from nodes.sharp.utils.gaussians import ( +from ..utils import io +from ..utils import logging as logging_utils +from ..utils.gaussians import ( Gaussians3D, SceneMetaData, save_ply, diff --git a/nodes/sharp/cli/render.py b/nodes/sharp/cli/render.py index ff6d191..29e8792 100644 --- a/nodes/sharp/cli/render.py +++ b/nodes/sharp/cli/render.py @@ -13,9 +13,9 @@ import click import torch import torch.utils.data -from nodes.sharp.utils import camera, gsplat, io -from nodes.sharp.utils import logging as logging_utils -from nodes.sharp.utils.gaussians import Gaussians3D, SceneMetaData, load_ply +from ..utils import camera, gsplat, io +from ..utils import logging as logging_utils +from ..utils.gaussians import Gaussians3D, SceneMetaData, load_ply LOGGER = logging.getLogger(__name__) diff --git a/nodes/sharp/models/__init__.py b/nodes/sharp/models/__init__.py index d7e20ef..4b9a047 100644 --- a/nodes/sharp/models/__init__.py +++ b/nodes/sharp/models/__init__.py @@ -6,7 +6,7 @@ Copyright (C) 2025 Apple Inc. All Rights Reserved. from __future__ import annotations -from nodes.sharp.models.monodepth import ( +from .monodepth import ( create_monodepth_adaptor, create_monodepth_dpt, ) diff --git a/nodes/sharp/models/alignment.py b/nodes/sharp/models/alignment.py index 22eedb1..f6e495c 100644 --- a/nodes/sharp/models/alignment.py +++ b/nodes/sharp/models/alignment.py @@ -12,9 +12,9 @@ import torch import torch.nn.functional as F from torch import nn -from nodes.sharp.models.decoders import UNetDecoder -from nodes.sharp.models.encoders import UNetEncoder -from nodes.sharp.utils import math as math_utils +from .decoders import UNetDecoder +from .encoders import UNetEncoder +from ..utils import math as math_utils from .params import AlignmentParams diff --git a/nodes/sharp/models/composer.py b/nodes/sharp/models/composer.py index efacc1a..2751566 100644 --- a/nodes/sharp/models/composer.py +++ b/nodes/sharp/models/composer.py @@ -10,10 +10,10 @@ import torch from torch import nn from torch.nn import functional as F -from nodes.sharp.models.initializer import GaussianBaseValues -from nodes.sharp.utils import math as math_utils -from nodes.sharp.utils.color_space import ColorSpace, sRGB2linearRGB -from nodes.sharp.utils.gaussians import Gaussians3D +from .initializer import GaussianBaseValues +from ..utils import math as math_utils +from ..utils.color_space import ColorSpace, sRGB2linearRGB +from ..utils.gaussians import Gaussians3D from .params import DeltaFactor diff --git a/nodes/sharp/models/decoders/monodepth_decoder.py b/nodes/sharp/models/decoders/monodepth_decoder.py index 4f46e1b..7790bd7 100644 --- a/nodes/sharp/models/decoders/monodepth_decoder.py +++ b/nodes/sharp/models/decoders/monodepth_decoder.py @@ -7,7 +7,7 @@ Copyright (C) 2025 Apple Inc. All Rights Reserved. from __future__ import annotations -from nodes.sharp.models.presets import ( +from ..presets import ( MONODEPTH_ENCODER_DIMS_MAP, ViTPreset, ) diff --git a/nodes/sharp/models/decoders/multires_conv_decoder.py b/nodes/sharp/models/decoders/multires_conv_decoder.py index d992c71..197246e 100644 --- a/nodes/sharp/models/decoders/multires_conv_decoder.py +++ b/nodes/sharp/models/decoders/multires_conv_decoder.py @@ -13,8 +13,8 @@ from typing import Iterable import torch import torch.nn as nn -from nodes.sharp.models.blocks import FeatureFusionBlock2d, UpsamplingMode -from nodes.sharp.utils.training import checkpoint_wrapper +from ..blocks import FeatureFusionBlock2d, UpsamplingMode +from ...utils.training import checkpoint_wrapper from .base_decoder import BaseDecoder diff --git a/nodes/sharp/models/decoders/unet_decoder.py b/nodes/sharp/models/decoders/unet_decoder.py index f8a7d0c..f1594f8 100644 --- a/nodes/sharp/models/decoders/unet_decoder.py +++ b/nodes/sharp/models/decoders/unet_decoder.py @@ -11,7 +11,7 @@ from typing import List import torch from torch import nn -from nodes.sharp.models.blocks import ( +from ..blocks import ( NormLayerName, norm_layer_2d, residual_block_2d, diff --git a/nodes/sharp/models/encoders/__init__.py b/nodes/sharp/models/encoders/__init__.py index d96f17a..6735bed 100644 --- a/nodes/sharp/models/encoders/__init__.py +++ b/nodes/sharp/models/encoders/__init__.py @@ -4,7 +4,7 @@ For licensing see accompanying LICENSE file. Copyright (C) 2025 Apple Inc. All Rights Reserved. """ -from nodes.sharp.models.encoders.base_encoder import BaseEncoder +from .base_encoder import BaseEncoder from .monodepth_encoder import ( MonodepthFeatureEncoder, diff --git a/nodes/sharp/models/encoders/monodepth_encoder.py b/nodes/sharp/models/encoders/monodepth_encoder.py index 33ad86a..25d7efe 100644 --- a/nodes/sharp/models/encoders/monodepth_encoder.py +++ b/nodes/sharp/models/encoders/monodepth_encoder.py @@ -11,7 +11,7 @@ from __future__ import annotations import torch import torch.nn as nn -from nodes.sharp.models.presets import ( +from ..presets import ( MONODEPTH_ENCODER_DIMS_MAP, MONODEPTH_HOOK_IDS_MAP, ViTPreset, diff --git a/nodes/sharp/models/encoders/spn_encoder.py b/nodes/sharp/models/encoders/spn_encoder.py index 2c2790d..2d517a9 100644 --- a/nodes/sharp/models/encoders/spn_encoder.py +++ b/nodes/sharp/models/encoders/spn_encoder.py @@ -14,7 +14,7 @@ import torch.fx import torch.nn as nn import torch.nn.functional as F -from nodes.sharp.utils.training import checkpoint_wrapper +from ...utils.training import checkpoint_wrapper from .base_encoder import BaseEncoder from .vit_encoder import TimmViT diff --git a/nodes/sharp/models/encoders/unet_encoder.py b/nodes/sharp/models/encoders/unet_encoder.py index 6b15da0..7bdf199 100644 --- a/nodes/sharp/models/encoders/unet_encoder.py +++ b/nodes/sharp/models/encoders/unet_encoder.py @@ -11,7 +11,7 @@ from typing import List import torch from torch import nn -from nodes.sharp.models.blocks import ( +from ..blocks import ( NormLayerName, norm_layer_2d, residual_block_2d, diff --git a/nodes/sharp/models/encoders/vit_encoder.py b/nodes/sharp/models/encoders/vit_encoder.py index 218a193..eb36379 100644 --- a/nodes/sharp/models/encoders/vit_encoder.py +++ b/nodes/sharp/models/encoders/vit_encoder.py @@ -11,7 +11,7 @@ import logging import timm import torch -from nodes.sharp.models.presets.vit import VIT_CONFIG_DICT, ViTConfig, ViTPreset +from ..presets.vit import VIT_CONFIG_DICT, ViTConfig, ViTPreset LOGGER = logging.getLogger(__name__) diff --git a/nodes/sharp/models/gaussian_decoder.py b/nodes/sharp/models/gaussian_decoder.py index 77e5f1b..eb982d2 100644 --- a/nodes/sharp/models/gaussian_decoder.py +++ b/nodes/sharp/models/gaussian_decoder.py @@ -13,13 +13,13 @@ from typing import NamedTuple import torch import torch.nn as nn -from nodes.sharp.models.blocks import ( +from .blocks import ( FeatureFusionBlock2d, NormLayerName, residual_block_2d, ) -from nodes.sharp.models.decoders import BaseDecoder, MultiresConvDecoder -from nodes.sharp.models.params import DPTImageEncoderType, GaussianDecoderParams +from .decoders import BaseDecoder, MultiresConvDecoder +from .params import DPTImageEncoderType, GaussianDecoderParams def create_gaussian_decoder( diff --git a/nodes/sharp/models/monodepth.py b/nodes/sharp/models/monodepth.py index cc6ea71..1e9f6cf 100644 --- a/nodes/sharp/models/monodepth.py +++ b/nodes/sharp/models/monodepth.py @@ -14,13 +14,13 @@ from typing import NamedTuple, Tuple import torch import torch.nn as nn -from nodes.sharp.models import normalizers -from nodes.sharp.models.decoders import MultiresConvDecoder, create_monodepth_decoder -from nodes.sharp.models.encoders import ( +from . import normalizers +from .decoders import MultiresConvDecoder, create_monodepth_decoder +from .encoders import ( SlidingPyramidNetwork, create_monodepth_encoder, ) -from nodes.sharp.utils import module_surgery +from ..utils import module_surgery from .params import MonodepthAdaptorParams, MonodepthParams diff --git a/nodes/sharp/models/params.py b/nodes/sharp/models/params.py index 12da284..21afbfb 100644 --- a/nodes/sharp/models/params.py +++ b/nodes/sharp/models/params.py @@ -7,10 +7,10 @@ Copyright (C) 2025 Apple Inc. All Rights Reserved. import dataclasses from typing import Literal -import nodes.sharp.utils.math as math_utils -from nodes.sharp.models.blocks import NormLayerName, UpsamplingMode -from nodes.sharp.models.presets import ViTPreset -from nodes.sharp.utils.color_space import ColorSpace +from ..utils import math as math_utils +from .blocks import NormLayerName, UpsamplingMode +from .presets import ViTPreset +from ..utils.color_space import ColorSpace DimsDecoder = tuple[int, int, int, int, int] DPTImageEncoderType = Literal["skip_conv", "skip_conv_kernel2"] diff --git a/nodes/sharp/models/predictor.py b/nodes/sharp/models/predictor.py index a90a2ae..8d5fed4 100644 --- a/nodes/sharp/models/predictor.py +++ b/nodes/sharp/models/predictor.py @@ -11,8 +11,8 @@ import logging import torch from torch import nn -from nodes.sharp.models.monodepth import MonodepthWithEncodingAdaptor -from nodes.sharp.utils.gaussians import Gaussians3D +from .monodepth import MonodepthWithEncodingAdaptor +from ..utils.gaussians import Gaussians3D from .composer import GaussianComposer diff --git a/nodes/sharp/utils/color_space.py b/nodes/sharp/utils/color_space.py index 356e298..d1cb1ae 100644 --- a/nodes/sharp/utils/color_space.py +++ b/nodes/sharp/utils/color_space.py @@ -11,7 +11,7 @@ from typing import Literal import torch -from nodes.sharp.utils.robust import robust_where +from .robust import robust_where LOGGER = logging.getLogger(__name__) diff --git a/nodes/sharp/utils/gaussians.py b/nodes/sharp/utils/gaussians.py index 05796b4..be550b0 100644 --- a/nodes/sharp/utils/gaussians.py +++ b/nodes/sharp/utils/gaussians.py @@ -14,8 +14,8 @@ import numpy as np import torch from plyfile import PlyData, PlyElement -from nodes.sharp.utils import color_space as cs_utils -from nodes.sharp.utils import linalg +from . import color_space as cs_utils +from . import linalg LOGGER = logging.getLogger(__name__) diff --git a/nodes/sharp/utils/gsplat.py b/nodes/sharp/utils/gsplat.py index c8fbdbc..f5e6d8b 100644 --- a/nodes/sharp/utils/gsplat.py +++ b/nodes/sharp/utils/gsplat.py @@ -13,9 +13,9 @@ import gsplat import torch from torch import nn -from nodes.sharp.utils import color_space as cs_utils -from nodes.sharp.utils import io, vis -from nodes.sharp.utils.gaussians import BackgroundColor, Gaussians3D +from . import color_space as cs_utils +from . import io, vis +from .gaussians import BackgroundColor, Gaussians3D class RenderingOutputs(NamedTuple): diff --git a/workflows/pano.json b/workflows/pano.json new file mode 100644 index 0000000..3e21f50 --- /dev/null +++ b/workflows/pano.json @@ -0,0 +1 @@ +{"id":"sharp-basic-workflow","revision":0,"last_node_id":25,"last_link_id":52,"nodes":[{"id":9,"type":"LoadSharpModel","pos":[-177.95877396853095,-1213.308620911364],"size":[270,82],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"checkpoint_path","name":"checkpoint_path","shape":7,"type":"STRING","widget":{"name":"checkpoint_path"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"SHARP_MODEL","links":[20]}],"properties":{"cnr_id":"comfyui-sharp","ver":"51939f0d3bdcd9f4965978445ad0fa31a6222ef0","Node name for S&R":"LoadSharpModel"},"widgets_values":["auto",""]},{"id":21,"type":"PreviewImage","pos":[1387.5905422433161,-1871.477198804017],"size":[636,398],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":38}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":20,"type":"PreviewImage","pos":[1394.930870035046,-2300.6935283639473],"size":[627.4232386014363,358.3867594737071],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":37}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":7,"type":"LoadImage","pos":[-910.0234950238612,-940.0676692572397],"size":[662.2282800153919,367.6329862374266],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[11]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"LoadImage"},"widgets_values":["pano.png","image"]},{"id":11,"type":"SharpPredictDepth","pos":[167.9465603665549,-1182.3544511142193],"size":[181.1509765625,86],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"SHARP_MODEL","link":20},{"localized_name":"image","name":"image","type":"IMAGE","link":14},{"localized_name":"extrinsics","name":"extrinsics","shape":7,"type":"EXTRINSICS","link":18},{"localized_name":"intrinsics","name":"intrinsics","shape":7,"type":"INTRINSICS","link":19}],"outputs":[{"localized_name":"depth_maps","name":"depth_maps","type":"IMAGE","links":[25,46]},{"localized_name":"extrinsics","name":"extrinsics","type":"EXTRINSICS","links":[26,47]},{"localized_name":"intrinsics","name":"intrinsics","type":"INTRINSICS","links":[27,48]}],"properties":{"cnr_id":"comfyui-sharp","ver":"51939f0d3bdcd9f4965978445ad0fa31a6222ef0","Node name for S&R":"SharpPredictDepth"},"widgets_values":[]},{"id":19,"type":"ProjectDepthToPanorama","pos":[809.5589801635481,-1818.007378585998],"size":[277.4771484375,194],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"depth_maps","name":"depth_maps","type":"IMAGE","link":49},{"localized_name":"extrinsics","name":"extrinsics","type":"EXTRINSICS","link":50},{"localized_name":"intrinsics","name":"intrinsics","type":"INTRINSICS","link":51},{"localized_name":"panorama_width","name":"panorama_width","shape":7,"type":"INT","widget":{"name":"panorama_width"},"link":null},{"localized_name":"blend_mode","name":"blend_mode","shape":7,"type":"COMBO","widget":{"name":"blend_mode"},"link":null},{"localized_name":"show_borders","name":"show_borders","shape":7,"type":"BOOLEAN","widget":{"name":"show_borders"},"link":null},{"localized_name":"border_color","name":"border_color","shape":7,"type":"COMBO","widget":{"name":"border_color"},"link":null},{"localized_name":"disagreement_scale","name":"disagreement_scale","shape":7,"type":"FLOAT","widget":{"name":"disagreement_scale"},"link":null}],"outputs":[{"localized_name":"panoramic_depth","name":"panoramic_depth","type":"IMAGE","links":[37]},{"localized_name":"debug_overlay","name":"debug_overlay","type":"IMAGE","links":[38]},{"localized_name":"disagreement_heatmap","name":"disagreement_heatmap","type":"IMAGE","links":null}],"properties":{"cnr_id":"comfyui-sharp","ver":"51939f0d3bdcd9f4965978445ad0fa31a6222ef0","Node name for S&R":"ProjectDepthToPanorama"},"widgets_values":[2048,"none",true,"red",0.2]},{"id":24,"type":"AlignDepthMaps","pos":[383.48454847695496,-1605.533120374793],"size":[360.65,142],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"depth_maps","name":"depth_maps","type":"IMAGE","link":46},{"localized_name":"extrinsics","name":"extrinsics","type":"EXTRINSICS","link":47},{"localized_name":"intrinsics","name":"intrinsics","type":"INTRINSICS","link":48},{"localized_name":"method","name":"method","shape":7,"type":"COMBO","widget":{"name":"method"},"link":null},{"localized_name":"reference_view","name":"reference_view","shape":7,"type":"INT","widget":{"name":"reference_view"},"link":null}],"outputs":[{"localized_name":"aligned_depth_maps","name":"aligned_depth_maps","type":"IMAGE","links":[49]},{"localized_name":"extrinsics","name":"extrinsics","type":"EXTRINSICS","links":[50]},{"localized_name":"intrinsics","name":"intrinsics","type":"INTRINSICS","links":[51]},{"localized_name":"alignment_info","name":"alignment_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-sharp","ver":"51939f0d3bdcd9f4965978445ad0fa31a6222ef0","Node name for S&R":"AlignDepthMaps"},"widgets_values":["global_optimization",10]},{"id":16,"type":"ProjectDepthToPanorama","pos":[422.29970963740055,-1189.3127313610928],"size":[277.4771484375,194],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"depth_maps","name":"depth_maps","type":"IMAGE","link":25},{"localized_name":"extrinsics","name":"extrinsics","type":"EXTRINSICS","link":26},{"localized_name":"intrinsics","name":"intrinsics","type":"INTRINSICS","link":27},{"localized_name":"panorama_width","name":"panorama_width","shape":7,"type":"INT","widget":{"name":"panorama_width"},"link":null},{"localized_name":"blend_mode","name":"blend_mode","shape":7,"type":"COMBO","widget":{"name":"blend_mode"},"link":null},{"localized_name":"show_borders","name":"show_borders","shape":7,"type":"BOOLEAN","widget":{"name":"show_borders"},"link":null},{"localized_name":"border_color","name":"border_color","shape":7,"type":"COMBO","widget":{"name":"border_color"},"link":null},{"localized_name":"disagreement_scale","name":"disagreement_scale","shape":7,"type":"FLOAT","widget":{"name":"disagreement_scale"},"link":null}],"outputs":[{"localized_name":"panoramic_depth","name":"panoramic_depth","type":"IMAGE","links":[29]},{"localized_name":"debug_overlay","name":"debug_overlay","type":"IMAGE","links":[30]},{"localized_name":"disagreement_heatmap","name":"disagreement_heatmap","type":"IMAGE","links":[52]}],"properties":{"cnr_id":"comfyui-sharp","ver":"51939f0d3bdcd9f4965978445ad0fa31a6222ef0","Node name for S&R":"ProjectDepthToPanorama"},"widgets_values":[2048,"gaussian",true,"red",0.2]},{"id":6,"type":"SamplePanorama","pos":[-200.75380822431256,-934.3551031268805],"size":[388.2466796875,210],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"panorama","name":"panorama","type":"IMAGE","link":11},{"localized_name":"fov_degrees","name":"fov_degrees","type":"FLOAT","widget":{"name":"fov_degrees"},"link":null},{"localized_name":"overlap_percent","name":"overlap_percent","type":"FLOAT","widget":{"name":"overlap_percent"},"link":null},{"localized_name":"output_size","name":"output_size","type":"INT","widget":{"name":"output_size"},"link":null},{"localized_name":"skip_poles","name":"skip_poles","shape":7,"type":"BOOLEAN","widget":{"name":"skip_poles"},"link":null}],"outputs":[{"localized_name":"images","name":"images","type":"IMAGE","links":[12,14]},{"localized_name":"extrinsics","name":"extrinsics","type":"EXTRINSICS","links":[18]},{"localized_name":"intrinsics","name":"intrinsics","type":"INTRINSICS","links":[19]},{"localized_name":"num_horizontal","name":"num_horizontal","type":"INT","links":null},{"localized_name":"num_vertical","name":"num_vertical","type":"INT","links":null}],"properties":{"cnr_id":"comfyui-sharp","ver":"51939f0d3bdcd9f4965978445ad0fa31a6222ef0","Node name for S&R":"SamplePanorama","aux_id":"PozzettiAndrea/ComfyUI-Sharp"},"widgets_values":[90,20,1536,true]},{"id":25,"type":"PreviewImage","pos":[837.7108350995264,-452.772224935793],"size":[636,398],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":52}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":8,"type":"PreviewImage","pos":[294.0726257921507,-862.2110152531175],"size":[456.7585111624088,246],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":12}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":13,"type":"PreviewImage","pos":[827.8445250599017,-891.9706860455786],"size":[636,398],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":30}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":12,"type":"PreviewImage","pos":[823.5595711961479,-1305.772046037818],"size":[548.9009283914478,321.74301470904584],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":29}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.2","Node name for S&R":"PreviewImage"},"widgets_values":[]}],"links":[[11,7,0,6,0,"IMAGE"],[12,6,0,8,0,"IMAGE"],[14,6,0,11,1,"IMAGE"],[18,6,1,11,2,"EXTRINSICS"],[19,6,2,11,3,"INTRINSICS"],[20,9,0,11,0,"SHARP_MODEL"],[25,11,0,16,0,"IMAGE"],[26,11,1,16,1,"EXTRINSICS"],[27,11,2,16,2,"INTRINSICS"],[29,16,0,12,0,"IMAGE"],[30,16,1,13,0,"IMAGE"],[37,19,0,20,0,"IMAGE"],[38,19,1,21,0,"IMAGE"],[46,11,0,24,0,"IMAGE"],[47,11,1,24,1,"EXTRINSICS"],[48,11,2,24,2,"INTRINSICS"],[49,24,0,19,0,"IMAGE"],[50,24,1,19,1,"EXTRINSICS"],[51,24,2,19,2,"INTRINSICS"],[52,16,2,25,0,"IMAGE"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.573085533011683,"offset":[733.1530148998932,1856.7715714261788]},"workflowRendererVersion":"LG","frontendVersion":"1.34.9"},"version":0.4} \ No newline at end of file