diff --git a/nodes/__init__.py b/nodes/__init__.py index 8960c64..d58e273 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -16,6 +16,16 @@ from .project_depth_to_panorama import NODE_CLASS_MAPPINGS as PROJECT_DEPTH_MAPP 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 +from .panorama_cube_split import NODE_CLASS_MAPPINGS as PANORAMA_CUBE_SPLIT_MAPPINGS +from .panorama_cube_split import NODE_DISPLAY_NAME_MAPPINGS as PANORAMA_CUBE_SPLIT_DISPLAY_MAPPINGS +from .depth_merge import NODE_CLASS_MAPPINGS as DEPTH_MERGE_MAPPINGS +from .depth_merge import NODE_DISPLAY_NAME_MAPPINGS as DEPTH_MERGE_DISPLAY_MAPPINGS +from .refined_depth_from_gaussians import NODE_CLASS_MAPPINGS as REFINED_DEPTH_MAPPINGS +from .refined_depth_from_gaussians import NODE_DISPLAY_NAME_MAPPINGS as REFINED_DEPTH_DISPLAY_MAPPINGS +from .predict_gaussian_attrs import NODE_CLASS_MAPPINGS as PREDICT_ATTRS_MAPPINGS +from .predict_gaussian_attrs import NODE_DISPLAY_NAME_MAPPINGS as PREDICT_ATTRS_DISPLAY_MAPPINGS +from .panorama_icosahedron_split import NODE_CLASS_MAPPINGS as ICOSAHEDRON_SPLIT_MAPPINGS +from .panorama_icosahedron_split import NODE_DISPLAY_NAME_MAPPINGS as ICOSAHEDRON_SPLIT_DISPLAY_MAPPINGS # Aggregate all node mappings NODE_CLASS_MAPPINGS = {} @@ -29,6 +39,11 @@ 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_CLASS_MAPPINGS.update(PANORAMA_CUBE_SPLIT_MAPPINGS) +NODE_CLASS_MAPPINGS.update(DEPTH_MERGE_MAPPINGS) +NODE_CLASS_MAPPINGS.update(REFINED_DEPTH_MAPPINGS) +NODE_CLASS_MAPPINGS.update(PREDICT_ATTRS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(ICOSAHEDRON_SPLIT_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_MODEL_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_DISPLAY_MAPPINGS) @@ -38,5 +53,10 @@ 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) +NODE_DISPLAY_NAME_MAPPINGS.update(PANORAMA_CUBE_SPLIT_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(DEPTH_MERGE_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(REFINED_DEPTH_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_ATTRS_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(ICOSAHEDRON_SPLIT_DISPLAY_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/_vendor/__init__.py b/nodes/_vendor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/_vendor/moge_panorama.py b/nodes/_vendor/moge_panorama.py new file mode 100644 index 0000000..6241ae7 --- /dev/null +++ b/nodes/_vendor/moge_panorama.py @@ -0,0 +1,354 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from pathlib import Path +from typing import * +import itertools +import json +import warnings + +import cv2 +import numpy as np +import torch +from numpy import ndarray +from tqdm import tqdm, trange +from scipy.sparse import csr_array, hstack, vstack +from scipy.ndimage import convolve +from scipy.sparse.linalg import lsmr + +import utils3d + + +def get_panorama_cameras(): + vertices, _ = utils3d.np.create_icosahedron_mesh() + intrinsics = utils3d.np.intrinsics_from_fov(fov_x=np.deg2rad(90), fov_y=np.deg2rad(90)) + extrinsics = utils3d.np.extrinsics_look_at([0, 0, 0], vertices, [0, 0, 1]).astype(np.float32) + return extrinsics, [intrinsics] * len(vertices) + + +def spherical_uv_to_directions(uv: np.ndarray): + theta, phi = (1 - uv[..., 0]) * (2 * np.pi), uv[..., 1] * np.pi + directions = np.stack([np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)], axis=-1) + return directions + + +def directions_to_spherical_uv(directions: np.ndarray): + directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True) + u = 1 - np.arctan2(directions[..., 1], directions[..., 0]) / (2 * np.pi) % 1.0 + v = np.arccos(directions[..., 2]) / np.pi + return np.stack([u, v], axis=-1) + + +def split_panorama_image(image: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, resolution: int): + height, width = image.shape[:2] + uv = utils3d.np.uv_map((resolution, resolution)) + splitted_images = [] + for i in range(len(extrinsics)): + spherical_uv = directions_to_spherical_uv(utils3d.np.unproject_cv(uv, np.ones_like(uv[..., 0]), extrinsics=extrinsics[i], intrinsics=intrinsics[i])) + pixels = utils3d.np.uv_to_pixel(spherical_uv, (height, width)).astype(np.float32) + + splitted_image = cv2.remap(image, pixels[..., 0], pixels[..., 1], interpolation=cv2.INTER_LINEAR) + splitted_images.append(splitted_image) + return splitted_images + + +def split_panorama_image_gpu(image: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, resolution: int): + """Batched-GPU equivalent of `split_panorama_image`. + + Same math: for each of N face cameras, build a (resolution, resolution) + lookup of spherical UVs from the panorama, and bilinearly resample. + Difference: the lookup is computed once for all N faces as a batched + (N, R, R, 2) grid (np vectorized over the face dim) and the per-face + resampling collapses to one F.grid_sample call. cv2.remap is single- + threaded per call and was the dominant cost on N=42 face splits. + + Bilinear matches cv2.INTER_LINEAR to within fp32 round-off when we use + `align_corners=False` and convert the cv2 pixel coordinate `p` (where + `p=0` is the centre of pixel 0) to the grid_sample coordinate + `(2p + 1) / W - 1`. Panoramic wrap in U is handled by the + `(% 1.0)` inside `directions_to_spherical_uv` so the lookups never + exit the image in U; in V we clip to [0, 1] which is what cv2 does + by default (BORDER_CONSTANT 0 below the south pole, but those samples + are unreachable because the icosahedron forwards are constrained to + the sphere). + + Falls back to the CPU implementation if CUDA isn't available. + """ + if not torch.cuda.is_available(): + return split_panorama_image(image, extrinsics, intrinsics, resolution) + + height, width = image.shape[:2] + N = len(extrinsics) + device = torch.device("cuda") + + # Build the (R, R, 2) UV map once on CPU (cheap, sub-ms). + uv = utils3d.np.uv_map((resolution, resolution)) + + # Per-face lookup pixels in panorama pixel coordinates (cv2 convention: + # pixel centre = integer). Stack into (N, R, R, 2) before going to GPU. + pixel_lookups = np.empty((N, resolution, resolution, 2), dtype=np.float32) + for i in range(N): + spherical_uv = directions_to_spherical_uv( + utils3d.np.unproject_cv( + uv, np.ones_like(uv[..., 0]), + extrinsics=extrinsics[i], intrinsics=intrinsics[i], + ) + ) + pixel_lookups[i] = utils3d.np.uv_to_pixel( + spherical_uv, (height, width) + ).astype(np.float32) + + # Convert pixel coords -> grid_sample normalized coords. Use the + # half-pixel-centre convention so this lines up with cv2.remap's + # INTER_LINEAR. + grid = np.empty_like(pixel_lookups) + grid[..., 0] = (2.0 * pixel_lookups[..., 0] + 1.0) / float(width) - 1.0 + grid[..., 1] = (2.0 * pixel_lookups[..., 1] + 1.0) / float(height) - 1.0 + + grid_t = torch.from_numpy(grid).to(device) # (N, R, R, 2) fp32 + + # Image to (1, C, H, W) → broadcast to (N, C, H, W) for batched sample. + img_chw = torch.from_numpy(np.ascontiguousarray(image)).to(device) + if img_chw.ndim == 2: + img_chw = img_chw.unsqueeze(-1) # (H, W) → (H, W, 1) + img_chw = img_chw.permute(2, 0, 1).contiguous().to(torch.float32) # (C, H, W) + img_batched = img_chw.unsqueeze(0).expand(N, -1, -1, -1) # (N, C, H, W) + + sampled = torch.nn.functional.grid_sample( + img_batched, grid_t, + mode="bilinear", padding_mode="zeros", align_corners=False, + ) # (N, C, R, R) + + sampled = sampled.permute(0, 2, 3, 1).contiguous() # (N, R, R, C) + arr = sampled.cpu().numpy() + if image.dtype == np.uint8: + arr = np.clip(np.rint(arr), 0, 255).astype(np.uint8) + elif image.dtype == np.float32: + arr = arr.astype(np.float32) + else: + arr = arr.astype(image.dtype) + # Drop the C dim if input was 2-D (single-channel mask etc.) + if image.ndim == 2: + arr = arr[..., 0] + return [arr[i] for i in range(N)] + + +def poisson_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False) -> Tuple[csr_array, ndarray]: + grid_index = np.arange(height * width).reshape(height, width) + grid_index = np.pad(grid_index, ((0, 0), (1, 1)), mode='wrap' if wrap_x else 'edge') + grid_index = np.pad(grid_index, ((1, 1), (0, 0)), mode='wrap' if wrap_y else 'edge') + + data = np.array([[-4, 1, 1, 1, 1]], dtype=np.float32).repeat(height * width, axis=0).reshape(-1) + indices = np.stack([ + grid_index[1:-1, 1:-1], + grid_index[:-2, 1:-1], # up + grid_index[2:, 1:-1], # down + grid_index[1:-1, :-2], # left + grid_index[1:-1, 2:] # right + ], axis=-1).reshape(-1) + indptr = np.arange(0, height * width * 5 + 1, 5) + A = csr_array((data, indices, indptr), shape=(height * width, height * width)) + + return A + + +def grad_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False) -> Tuple[csr_array, np.ndarray]: + grid_index = np.arange(width * height).reshape(height, width) + if wrap_x: + grid_index = np.pad(grid_index, ((0, 0), (0, 1)), mode='wrap') + if wrap_y: + grid_index = np.pad(grid_index, ((0, 1), (0, 0)), mode='wrap') + + data = np.concatenate([ + np.concatenate([ + np.ones((grid_index.shape[0], grid_index.shape[1] - 1), dtype=np.float32).reshape(-1, 1), # x[i,j] + -np.ones((grid_index.shape[0], grid_index.shape[1] - 1), dtype=np.float32).reshape(-1, 1), # x[i,j-1] + ], axis=1).reshape(-1), + np.concatenate([ + np.ones((grid_index.shape[0] - 1, grid_index.shape[1]), dtype=np.float32).reshape(-1, 1), # x[i,j] + -np.ones((grid_index.shape[0] - 1, grid_index.shape[1]), dtype=np.float32).reshape(-1, 1), # x[i-1,j] + ], axis=1).reshape(-1), + ]) + indices = np.concatenate([ + np.concatenate([ + grid_index[:, :-1].reshape(-1, 1), + grid_index[:, 1:].reshape(-1, 1), + ], axis=1).reshape(-1), + np.concatenate([ + grid_index[:-1, :].reshape(-1, 1), + grid_index[1:, :].reshape(-1, 1), + ], axis=1).reshape(-1), + ]) + indptr = np.arange(0, grid_index.shape[0] * (grid_index.shape[1] - 1) * 2 + (grid_index.shape[0] - 1) * grid_index.shape[1] * 2 + 1, 2) + A = csr_array((data, indices, indptr), shape=(grid_index.shape[0] * (grid_index.shape[1] - 1) + (grid_index.shape[0] - 1) * grid_index.shape[1], height * width)) + + return A + + +def merge_panorama_depth(width: int, height: int, distance_maps: List[np.ndarray], pred_masks: List[np.ndarray], extrinsics: List[np.ndarray], intrinsics: List[np.ndarray], *, pbar=None, _top_size=None): + """Recursive multi-resolution merge of N rectilinear face depth maps → equirect depth. + + Optional `pbar` is a ComfyUI ProgressBar-like object (any obj with + `update_absolute(value, total)`). Progress is weighted by `pixels^1.5` + per level to match LSMR's asymptotic cost — the top level dominates, + so the bar advances slowly until the unwind reaches it. + """ + import sys + import time + + def _p(msg): + print(f"[merge_panorama_depth {width:>4}x{height:<4}] {msg}", file=sys.stderr, flush=True) + + # --- Progress-bar setup (one-time on top-level call) --------------------- + if _top_size is None: + _top_size = (width, height) + if pbar is not None and not hasattr(pbar, "_wn_levels"): + levels = [] + w_, h_ = _top_size + while True: + levels.append((w_, h_)) + if max(w_, h_) <= 256: + break + w_, h_ = w_ // 2, h_ // 2 + levels.reverse() # smallest → largest, matches recursion unwind order + weights = np.array([(lw * lh) ** 1.5 for lw, lh in levels], dtype=np.float64) + weights = weights / weights.sum() + pbar._wn_levels = levels + pbar._wn_weights = weights + _p(f"progress bar: {len(levels)} levels — weights {[f'{x:.1%}' for x in weights]}") + + def _bump(level_frac: float): + """Advance the global progress bar to (cum % of prior levels) + (level_frac × this level's share).""" + if pbar is None or not hasattr(pbar, "_wn_levels"): + return + try: + i = pbar._wn_levels.index((width, height)) + except ValueError: + return + pct_floor = float(pbar._wn_weights[:i].sum()) + pct_ceil = float(pbar._wn_weights[:i + 1].sum()) + pct = pct_floor + (pct_ceil - pct_floor) * max(0.0, min(1.0, level_frac)) + pbar.update_absolute(int(min(99, pct * 100)), 100) + + t_lvl_start = time.perf_counter() + _bump(0.0) + + if max(width, height) > 256: + _p(f"recurse into {width // 2}x{height // 2}…") + panorama_depth_init, _ = merge_panorama_depth( + width // 2, height // 2, distance_maps, pred_masks, extrinsics, intrinsics, + pbar=pbar, _top_size=_top_size, + ) + _p(f"recursion returned; upsampling {width // 2}x{height // 2} → {width}x{height}") + panorama_depth_init = cv2.resize(panorama_depth_init, (width, height), cv2.INTER_LINEAR) + else: + _p(f"base case (≤256 on each side) → no recursion, no init x0 for LSMR") + panorama_depth_init = None + _bump(0.05) + + _p(f"building sphere uv + ray directions ({width}×{height})…") + uv = utils3d.np.uv_map(height, width) + spherical_directions = spherical_uv_to_directions(uv) + + # Warp each view to the panorama + _p(f"warping {len(distance_maps)} face depth maps onto the equirect…") + t_warp_start = time.perf_counter() + panorama_log_distance_grad_maps, panorama_grad_masks = [], [] + panorama_log_distance_laplacian_maps, panorama_laplacian_masks = [], [] + panorama_pred_masks = [] + for i in range(len(distance_maps)): + projected_uv, projected_depth = utils3d.np.project_cv(spherical_directions, extrinsics=extrinsics[i], intrinsics=intrinsics[i]) + projection_valid_mask = (projected_depth > 0) & (projected_uv > 0).all(axis=-1) & (projected_uv < 1).all(axis=-1) + + projected_pixels = utils3d.np.uv_to_pixel(np.clip(projected_uv, 0, 1), distance_maps[i].shape).astype(np.float32) + + log_splitted_distance = np.log(distance_maps[i]) + panorama_log_distance_map = np.where(projection_valid_mask, cv2.remap(log_splitted_distance, projected_pixels[..., 0], projected_pixels[..., 1], cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE), 0) + panorama_pred_mask = projection_valid_mask & (cv2.remap(pred_masks[i].astype(np.uint8), projected_pixels[..., 0], projected_pixels[..., 1], cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE) > 0) + + # calculate gradient map + padded = np.pad(panorama_log_distance_map, ((0, 0), (0, 1)), mode='wrap') + grad_x, grad_y = padded[:, :-1] - padded[:, 1:], padded[:-1, :] - padded[1:, :] + + padded = np.pad(panorama_pred_mask, ((0, 0), (0, 1)), mode='wrap') + mask_x, mask_y = padded[:, :-1] & padded[:, 1:], padded[:-1, :] & padded[1:, :] + + panorama_log_distance_grad_maps.append((grad_x, grad_y)) + panorama_grad_masks.append((mask_x, mask_y)) + + # calculate laplacian map + padded = np.pad(panorama_log_distance_map, ((1, 1), (0, 0)), mode='edge') + padded = np.pad(padded, ((0, 0), (1, 1)), mode='wrap') + laplacian = convolve(padded, np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=np.float32))[1:-1, 1:-1] + + padded = np.pad(panorama_pred_mask, ((1, 1), (0, 0)), mode='edge') + padded = np.pad(padded, ((0, 0), (1, 1)), mode='wrap') + mask = convolve(padded.astype(np.uint8), np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.uint8))[1:-1, 1:-1] == 5 + + panorama_log_distance_laplacian_maps.append(laplacian) + panorama_laplacian_masks.append(mask) + + panorama_pred_masks.append(panorama_pred_mask) + if (i + 1) % 10 == 0 or i == len(distance_maps) - 1: + _p(f" warped {i + 1}/{len(distance_maps)} faces") + + _bump(0.45) + _p(f"warp stage done in {time.perf_counter() - t_warp_start:.2f}s; " + f"stacking + averaging gradient/laplacian maps…") + panorama_log_distance_grad_x = np.stack([grad_map[0] for grad_map in panorama_log_distance_grad_maps], axis=0) + panorama_log_distance_grad_y = np.stack([grad_map[1] for grad_map in panorama_log_distance_grad_maps], axis=0) + panorama_grad_mask_x = np.stack([mask_map[0] for mask_map in panorama_grad_masks], axis=0) + panorama_grad_mask_y = np.stack([mask_map[1] for mask_map in panorama_grad_masks], axis=0) + + panorama_log_distance_grad_x = np.sum(panorama_log_distance_grad_x * panorama_grad_mask_x, axis=0) / np.sum(panorama_grad_mask_x, axis=0).clip(1e-3) + panorama_log_distance_grad_y = np.sum(panorama_log_distance_grad_y * panorama_grad_mask_y, axis=0) / np.sum(panorama_grad_mask_y, axis=0).clip(1e-3) + + panorama_laplacian_maps = np.stack(panorama_log_distance_laplacian_maps, axis=0) + panorama_laplacian_masks = np.stack(panorama_laplacian_masks, axis=0) + panorama_laplacian_map = np.sum(panorama_laplacian_maps * panorama_laplacian_masks, axis=0) / np.sum(panorama_laplacian_masks, axis=0).clip(1e-3) + + grad_x_mask = np.any(panorama_grad_mask_x, axis=0).reshape(-1) + grad_y_mask = np.any(panorama_grad_mask_y, axis=0).reshape(-1) + grad_mask = np.concatenate([grad_x_mask, grad_y_mask]) + laplacian_mask = np.any(panorama_laplacian_masks, axis=0).reshape(-1) + + _bump(0.55) + # Solve overdetermined system + _p(f"building sparse linear system (grad_equation + poisson_equation)…") + t_build = time.perf_counter() + A = vstack([ + grad_equation(width, height, wrap_x=True, wrap_y=False)[grad_mask], + poisson_equation(width, height, wrap_x=True, wrap_y=False)[laplacian_mask], + ]) + b = np.concatenate([ + panorama_log_distance_grad_x.reshape(-1)[grad_x_mask], + panorama_log_distance_grad_y.reshape(-1)[grad_y_mask], + panorama_laplacian_map.reshape(-1)[laplacian_mask] + ]) + _p(f"linear system: A {A.shape} ({A.nnz} nnz), b {b.shape} — built in " + f"{time.perf_counter() - t_build:.2f}s") + _bump(0.65) + + _p(f"LSMR solve (atol=1e-5, btol=1e-5, " + f"x0={'provided' if panorama_depth_init is not None else 'None'})…") + t_solve = time.perf_counter() + x, istop, itn, normr, normar, norma, conda, normx = lsmr( + A, b, + atol=1e-5, btol=1e-5, + x0=np.log(panorama_depth_init).reshape(-1) if panorama_depth_init is not None else None, + show=False, + ) + _p(f"LSMR done in {time.perf_counter() - t_solve:.2f}s: " + f"{itn} iterations, residual={normr:.4f}, istop={istop}") + _bump(0.95) + + panorama_depth = np.exp(x).reshape(height, width).astype(np.float32) + panorama_mask = np.any(panorama_pred_masks, axis=0) + + _p(f"level done in {time.perf_counter() - t_lvl_start:.2f}s; " + f"depth: min={float(panorama_depth.min()):.3f}, " + f"median={float(np.median(panorama_depth)):.3f}, " + f"max={float(panorama_depth.max()):.3f}; " + f"valid pixels: {int(panorama_mask.sum())}/{panorama_mask.size}") + _bump(1.0) + return panorama_depth, panorama_mask + diff --git a/nodes/_vendor/worldgen/__init__.py b/nodes/_vendor/worldgen/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/_vendor/worldgen/src/__init__.py b/nodes/_vendor/worldgen/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/_vendor/worldgen/src/panorama_utils.py b/nodes/_vendor/worldgen/src/panorama_utils.py new file mode 100644 index 0000000..a5e9b7c --- /dev/null +++ b/nodes/_vendor/worldgen/src/panorama_utils.py @@ -0,0 +1,2074 @@ +from collections import defaultdict +from typing import Optional, Literal, List + +# cupy stripped — see solve_lsmr_gpu below (now CPU via scipy). +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +import utils3d +from scipy.sparse import csr_matrix as cp_csr_matrix +from scipy.sparse.linalg import lsmr as cp_lsmr + +# Heavy / optional deps. The depth-merge path in Sharp only needs +# `utils3d.np.uv_map` (used by merge_panorama_depth_gpu) and scipy/torch +# (lsmr solver). open3d / trimesh / PIL are only touched by the +# point-cloud / mesh / image helpers further down in this file — none +# of which the Sharp depth_merge node calls. Make them optional so +# Sharp's nodes package imports cleanly even when those libs aren't in +# the env. +try: + import open3d as o3d # noqa: F401 +except ImportError: + o3d = None +try: + import trimesh # noqa: F401 +except ImportError: + trimesh = None +try: + from PIL import Image # noqa: F401 +except ImportError: + Image = None +# `convolve`, `vstack`, `grad_equation`, `poisson_equation` are lazy-imported +# inside merge_panorama_depth_gpu — that function is the cube-face stitcher +# used by upstream's pred_pano_depth, which WorldNav doesn't call (MoGe is +# externalized to a sibling ComfyUI node). The moge dep stays optional. +# `List` was imported from moge.utils.panorama but is actually re-exported +# from typing — sourcing directly from typing above. +from scipy.sparse import csr_array +from tqdm import tqdm + + +def subdivide_icosahedron(subdivisions: int = 1) -> np.ndarray: + """ + Subdivide an icosahedron to generate denser spherical sample points. + + Args: + subdivisions: Number of subdivisions. + - 0: 12 vertices + - 1: 42 vertices + - 2: 162 vertices + - 3: 642 vertices + Formula: V = 10 * 4^n + 2 + + Returns: + vertices: (N, 3) Subdivided vertex coordinates on the unit sphere. + """ + # utils3d renamed `icosahedron()` → `create_icosahedron_mesh()`. + vertices, faces = utils3d.numpy.create_icosahedron_mesh() + + # Convert to a list so new vertices can be appended dynamically. + vertices_list = [v for v in vertices] + + for _ in range(subdivisions): + edge_midpoint_cache = {} # Cache edge midpoint indices to avoid duplicates. + new_faces = [] + + def get_or_create_midpoint(idx1: int, idx2: int) -> int: + """ + Get the midpoint index for the edge between two vertices. + Create it if missing, otherwise return the cached index. + """ + # Use the sorted tuple as the key so (a, b) and (b, a) are the same edge. + edge_key = (min(idx1, idx2), max(idx1, idx2)) + + if edge_key in edge_midpoint_cache: + return edge_midpoint_cache[edge_key] + + # Create a new midpoint. + v1 = vertices_list[idx1] + v2 = vertices_list[idx2] + midpoint = (v1 + v2) / 2.0 + + # Project onto the unit sphere. + midpoint = midpoint / np.linalg.norm(midpoint) + + # Add it to the vertex list. + new_idx = len(vertices_list) + vertices_list.append(midpoint) + edge_midpoint_cache[edge_key] = new_idx + + return new_idx + + # Subdivide each triangle. + for face in faces: + v0, v1, v2 = face + + # Get the midpoints of the three edges. + # v0 + # / \ + # a----c + # / \ / \ + # v1--b----v2 + a = get_or_create_midpoint(v0, v1) + b = get_or_create_midpoint(v1, v2) + c = get_or_create_midpoint(v2, v0) + + # Split the original triangle into four smaller triangles. + new_faces.append([v0, a, c]) + new_faces.append([a, v1, b]) + new_faces.append([c, b, v2]) + new_faces.append([a, b, c]) + + faces = np.array(new_faces, dtype=np.int32) + + return np.array(vertices_list, dtype=np.float32) + + +def get_panorama_cameras_v2(subdivisions=0): + vertices = subdivide_icosahedron(subdivisions=subdivisions) + intrinsics = utils3d.numpy.intrinsics_from_fov(fov_x=np.deg2rad(90), fov_y=np.deg2rad(90)) + extrinsics = utils3d.numpy.extrinsics_look_at([0, 0, 0], vertices, [0, 0, 1]).astype(np.float32) + return extrinsics, [intrinsics] * len(vertices) + + +def rotate_around_z_axis(points, angle_deg): + """ + Rotate 3D points clockwise around the Z axis. + + Args: + points: 3D point array with shape (N, 3), where each row is (x, y, z). + angle_deg: Rotation angle in degrees. Positive values rotate clockwise. + + Returns: + rotated_points: Rotated 3D point array with shape (N, 3). + """ + # Convert the angle to radians. + angle_rad = np.radians(angle_deg) + + # Build the clockwise rotation matrix around the Z axis. + cos_theta = np.cos(angle_rad) + sin_theta = np.sin(angle_rad) + + # In a right-handed coordinate system, clockwise rotation equals negative counterclockwise rotation. + rotation_matrix = np.array([ + [cos_theta, sin_theta, 0], # X component + [-sin_theta, cos_theta, 0], # Y component + [0, 0, 1] # Z component, unchanged + ]) + + # Apply the rotation to each point by matrix multiplication. + rotated_points = np.dot(points, rotation_matrix.T) + + return rotated_points + + +def directions_to_spherical_uv(directions: np.ndarray): + directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True) + u = 1 - np.arctan2(directions[..., 1], directions[..., 0]) / (2 * np.pi) % 1.0 + v = np.arccos(directions[..., 2]) / np.pi + return np.stack([u, v], axis=-1) + + +def split_panorama_image(image: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, h: int, w: int, interp): + height, width = image.shape[:2] + safe_height = height // 2 + safe_width = int(round(safe_height / h * w)) + if interp == cv2.INTER_AREA: # remap does not support area downsampling; remap to a safe resolution first to avoid frequency artifacts. + uv = utils3d.numpy.image_uv(width=safe_width, height=safe_height) + else: + uv = utils3d.numpy.image_uv(width=w, height=h) + splitted_images = [] + for i in range(len(extrinsics)): + spherical_uv = directions_to_spherical_uv(utils3d.numpy.unproject_cv(uv, extrinsics=extrinsics[i], intrinsics=intrinsics[i])) + pixels = utils3d.numpy.uv_to_pixel(spherical_uv, width=width, height=height).astype(np.float32) + if interp == cv2.INTER_AREA: # remap does not support area downsampling; remap to a safe resolution first to avoid frequency artifacts. + splitted_image = cv2.remap(image, pixels[..., 0], pixels[..., 1], interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP) + splitted_image = cv2.resize(splitted_image, (w, h), interpolation=interp) + else: + splitted_image = cv2.remap(image, pixels[..., 0], pixels[..., 1], interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP) + splitted_images.append(splitted_image) + return splitted_images + + +def split_panorama_depth(depth: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, h: int, w: int, distance_to_depth=False): + height, width = depth.shape[:2] + depth = torch.tensor(depth, dtype=torch.float32)[None, None] + uv = utils3d.numpy.image_uv(width=w, height=h) + u_grid, v_grid = np.meshgrid(np.arange(w), np.arange(h)) + splitted_depths = [] + for i in range(len(extrinsics)): + spherical_uv = directions_to_spherical_uv(utils3d.numpy.unproject_cv(uv, extrinsics=extrinsics[i], intrinsics=intrinsics[i])) + pixels = utils3d.numpy.uv_to_pixel(spherical_uv, width=width, height=height).astype(np.float32) + pixels = torch.tensor(pixels, dtype=torch.float32)[None, ...] # [1,h,w,2] + pixels[..., 0] /= width + pixels[..., 1] /= height + pixels = pixels * 2 - 1.0 + splitted_depth = F.grid_sample(depth, grid=pixels, mode="nearest", align_corners=True) + + if distance_to_depth: + fx = intrinsics[i][0, 0] * w + fy = intrinsics[i][1, 1] * h + cx = intrinsics[i][0, 2] * w + cy = intrinsics[i][1, 2] * h + x_cam = (u_grid - cx) / fx + y_cam = (v_grid - cy) / fy + z_cam = np.ones_like(x_cam) + rays_cam = np.stack([x_cam, y_cam, z_cam], axis=-1).astype(np.float32) # (H, W, 3) + ray_length = np.linalg.norm(rays_cam, axis=-1).astype(np.float32) + splitted_depth = splitted_depth * (z_cam[None, None] / ray_length[None, None]) + + splitted_depths.append(splitted_depth) + return torch.cat(splitted_depths, dim=0).float() + + +def smooth_south_pole_depth(depth_map, smooth_height_ratio=0.03): + """ + Smooth depth near the panorama south pole (bottom region) to fix left-right inconsistencies. + + Args: + depth_map: Depth map (H, W). + smooth_height_ratio: Height ratio of the smoothing region. The default 0.03 means the bottom 0.03 * H region. + + Returns: + Smoothed depth map. + """ + height, width = depth_map.shape + smooth_height = int(height * smooth_height_ratio) + + if smooth_height == 0: + return depth_map + + # Copy the depth map to avoid modifying the input. + smoothed_depth = depth_map.copy() + + # Compute the reference depth from the last 3 rows when possible, otherwise from the bottom row. + if smooth_height > 3: + # Use the last 3 rows to compute the average depth. + reference_rows = depth_map[-3:, :] + reference_data = reference_rows.flatten() + else: + # Use the bottom row. + reference_data = depth_map[-1, :] + + # Filter outliers, including invalid, overly large, or overly small depth values. + valid_mask = np.isfinite(reference_data) & (reference_data > 0) + + if np.any(valid_mask): + valid_depths = reference_data[valid_mask] + + # Use quantiles to filter extreme outliers. + lower_bound, upper_bound = np.quantile(valid_depths, [0.1, 0.9]) + + # Further remove overly large or small depth values. + depth_filter_mask = (valid_depths >= lower_bound) & (valid_depths <= upper_bound) + + if np.any(depth_filter_mask): + avg_depth = np.mean(valid_depths[depth_filter_mask]) + else: + # Fall back to the median if all values are filtered out. + avg_depth = np.median(valid_depths) + else: + avg_depth = np.nanmean(reference_data) + + # Set the bottom row to the average value. + smoothed_depth[-1, :] = avg_depth + + # Smooth upward to the specified height. + for i in range(1, smooth_height): + y_idx = height - 1 - i # Index moving upward from the bottom. + if y_idx < 0: + break + + # The closer to the bottom, the stronger the smoothing. + weight = (smooth_height - i) / smooth_height + + # Smooth the current row. + current_row = depth_map[y_idx, :] + valid_mask = np.isfinite(current_row) & (current_row > 0) + + if np.any(valid_mask): + valid_row_depths = current_row[valid_mask] + + # Apply outlier filtering to the current row as well. + if len(valid_row_depths) > 1: + q25, q75 = np.quantile(valid_row_depths, [0.25, 0.75]) + iqr = q75 - q25 + lower_bound = q25 - 1.5 * iqr + upper_bound = q75 + 1.5 * iqr + depth_filter_mask = (valid_row_depths >= lower_bound) & (valid_row_depths <= upper_bound) + + if np.any(depth_filter_mask): + row_avg = np.mean(valid_row_depths[depth_filter_mask]) + else: + row_avg = np.median(valid_row_depths) + else: + row_avg = valid_row_depths[0] if len(valid_row_depths) > 0 else avg_depth + + # Linearly interpolate between the original depth and the row average. + smoothed_depth[y_idx, :] = (1 - weight) * current_row + weight * row_avg + + return smoothed_depth + + +def solve_lsmr_gpu(A, b, x0=None, atol=1e-5, btol=1e-5, conlim=0.0, + maxiter=None): + """On-device LSMR for sparse least-squares ``min ||Ax - b||_2``. + + Same Golub-Kahan bidiagonalization + Givens-rotation update steps as + scipy.sparse.linalg.lsmr (Fong & Saunders 2011), but the entire iteration + runs through torch tensors — the big work vectors u, v, h, hbar, x live + on the GPU; only the ~6 scalar Givens reductions per iteration come back + to the host. Eliminates the per-iteration host↔device round-trip that + the old `_TorchSparseLinearOperator` paid (matvec → .cpu().numpy() → + scipy numpy ops → as_tensor → next matvec). For a 1920×960 panorama + solve that's ~600 iterations × 4 transfers × 7 MB ≈ 30 GB of avoided + PCIe traffic. + + Numerical behaviour vs scipy fp64 CPU: SpMV runs in fp32 on cuSPARSE + (same as the previous GPU path), all 1.84M-element vector ops in fp32; + scalar reductions cast to fp64 on the host so the Givens math is + bit-equivalent to scipy's. Same converged answer up to fp32 SpMV + round-off, which is the existing baseline. + + Args: + A: scipy CSR-like sparse matrix (csr_array, csr_matrix, or + anything `csr_matrix(...)` accepts). + b: numpy 1-D vector, length A.shape[0]. + x0: optional initial guess, length A.shape[1]. + atol, stopping tolerances — same semantics as scipy.sparse.linalg.lsmr. + btol: `atol` controls ||A^T r||/(||A|| ||r||) test, `btol` the + ||r||/||b|| test. + conlim: if > 0, also stop when an estimate of cond(A) exceeds + conlim. Default 0 = disabled — the old GPU path with a + host-side LinearOperator triggered scipy's `condA` cast + to overflow at lsmr.py:407 because fp32 SpMV outputs + didn't match scipy's fp64 expectations; skipping the + tracking removes the warning AND saves a handful of host + ops per iteration. The same solution is returned (cond + estimate is an early-stop, never affects the iterate). + maxiter: cap on iterations. Default min(m, n). + + Returns: numpy 1-D vector, length A.shape[1]. + """ + from math import sqrt, inf + from scipy.sparse import csr_matrix + from scipy.sparse.linalg._isolve.lsqr import _sym_ortho + + # --- CPU fallback: pure scipy lsmr. Same answer, slower. --- + if not torch.cuda.is_available(): + A_csr = cp_csr_matrix(A) + x, *_ = cp_lsmr(A_csr, b, atol=atol, btol=btol, x0=x0, + conlim=conlim if conlim > 0 else 1e8, + maxiter=maxiter) + return x + + device = torch.device("cuda") + # Normalize input to scipy CSR so we can read indptr / indices / data. + A_csr = A if isinstance(A, csr_matrix) else csr_matrix(A) + m, n = A_csr.shape + if maxiter is None: + maxiter = min(m, n) + + # --- Upload A and A^T as torch CSR (one-time per solve). --- + # int32 indices are safe (m, n both < 2^31); halves SpMV index bandwidth + # vs the int64 the old operator used. + A_gpu = torch.sparse_csr_tensor( + torch.from_numpy(A_csr.indptr.astype(np.int32)).to(device), + torch.from_numpy(A_csr.indices.astype(np.int32)).to(device), + torch.from_numpy(A_csr.data.astype(np.float32)).to(device), + size=A_csr.shape, + ) + AT_csr = A_csr.T.tocsr() + AT_gpu = torch.sparse_csr_tensor( + torch.from_numpy(AT_csr.indptr.astype(np.int32)).to(device), + torch.from_numpy(AT_csr.indices.astype(np.int32)).to(device), + torch.from_numpy(AT_csr.data.astype(np.float32)).to(device), + size=AT_csr.shape, + ) + + def _matvec(vec): # vec: (n,) → (m,) + return torch.sparse.mm(A_gpu, vec.unsqueeze(-1)).squeeze(-1) + + def _rmatvec(vec): # vec: (m,) → (n,) + return torch.sparse.mm(AT_gpu, vec.unsqueeze(-1)).squeeze(-1) + + # --- Upload b (and x0 if given). --- + b_gpu = torch.as_tensor(np.asarray(b, dtype=np.float32), device=device) + u = b_gpu.clone() + normb = float(torch.linalg.vector_norm(b_gpu).item()) + + if x0 is None: + x = torch.zeros(n, device=device, dtype=torch.float32) + beta = normb + else: + x = torch.as_tensor(np.asarray(x0, dtype=np.float32), + device=device).clone() + u.sub_(_matvec(x)) + beta = float(torch.linalg.vector_norm(u).item()) + + if beta > 0.0: + u.mul_(1.0 / beta) + v = _rmatvec(u) + alpha = float(torch.linalg.vector_norm(v).item()) + else: + v = torch.zeros(n, device=device, dtype=torch.float32) + alpha = 0.0 + + if alpha > 0.0: + v.mul_(1.0 / alpha) + + # --- Givens / convergence scalar state (fp64 on host). --- + itn = 0 + zetabar = alpha * beta + alphabar = alpha + rho = 1.0 + rhobar = 1.0 + cbar = 1.0 + sbar = 0.0 + + h = v.clone() + hbar = torch.zeros(n, device=device, dtype=torch.float32) + + betadd = beta + betad = 0.0 + rhodold = 1.0 + tautildeold = 0.0 + thetatilde = 0.0 + zeta = 0.0 + d = 0.0 + + normA2 = alpha * alpha + maxrbar = 0.0 + minrbar = 1e100 + normA = sqrt(normA2) + condA = 1.0 + normx = 0.0 + damp = 0.0 # we never use damping in this codepath + + istop = 0 + ctol = (1.0 / conlim) if conlim > 0 else 0.0 + normr = beta + normar = alpha * beta + + if normar == 0.0: + return x.cpu().numpy() + if normb == 0.0: + x.zero_() + return x.cpu().numpy() + + while itn < maxiter: + itn += 1 + + # Bidiagonalization step. u <- A v - alpha u; v <- A^T u - beta v. + u.mul_(-alpha).add_(_matvec(v)) + beta = float(torch.linalg.vector_norm(u).item()) + if beta > 0.0: + u.mul_(1.0 / beta) + v.mul_(-beta).add_(_rmatvec(u)) + alpha = float(torch.linalg.vector_norm(v).item()) + if alpha > 0.0: + v.mul_(1.0 / alpha) + + # Givens rotations (host scalar math, ~1us each). + chat, shat, alphahat = _sym_ortho(alphabar, damp) + rhoold = rho + c, s, rho = _sym_ortho(alphahat, beta) + thetanew = s * alpha + alphabar = c * alpha + + rhobarold = rhobar + zetaold = zeta + thetabar = sbar * rho + rhotemp = cbar * rho + cbar, sbar, rhobar = _sym_ortho(cbar * rho, thetanew) + zeta = cbar * zetabar + zetabar = -sbar * zetabar + + # Update h, hbar, x — fused in-place AXPYs on GPU. + hbar.mul_(-(thetabar * rho / (rhoold * rhobarold))).add_(h) + x.add_(hbar, alpha=(zeta / (rho * rhobar))) + h.mul_(-(thetanew / rho)).add_(v) + + # Estimate ||r||. + betaacute = chat * betadd + betacheck = -shat * betadd + betahat = c * betaacute + betadd = -s * betaacute + + thetatildeold = thetatilde + ctildeold, stildeold, rhotildeold = _sym_ortho(rhodold, thetabar) + thetatilde = stildeold * rhobar + rhodold = ctildeold * rhobar + betad = -stildeold * betad + ctildeold * betahat + + tautildeold = (zetaold - thetatildeold * tautildeold) / rhotildeold + taud = (zeta - thetatilde * tautildeold) / rhodold + d = d + betacheck * betacheck + normr = sqrt(d + (betad - taud) ** 2 + betadd * betadd) + + # Estimate ||A||. + normA2 = normA2 + beta * beta + normA = sqrt(normA2) + normA2 = normA2 + alpha * alpha + + # Estimate cond(A) — only if user asked (conlim > 0). Skipping the + # whole block saves a handful of host ops AND avoids scipy's + # `condA = max(...)/min(...)` fp32 cast overflow that polluted the + # old GPU log with RuntimeWarnings. + if conlim > 0.0: + maxrbar = max(maxrbar, rhobarold) + if itn > 1: + minrbar = min(minrbar, rhobarold) + condA = max(maxrbar, rhotemp) / min(minrbar, rhotemp) + + # Convergence tests. + normar = abs(zetabar) + normx = float(torch.linalg.vector_norm(x).item()) + + test1 = normr / normb + if (normA * normr) != 0.0: + test2 = normar / (normA * normr) + else: + test2 = inf + test3 = 1.0 / condA + t1 = test1 / (1.0 + normA * normx / normb) + rtol = btol + atol * normA * normx / normb + + if itn >= maxiter: + istop = 7 + if 1.0 + test3 <= 1.0: + istop = 6 + if 1.0 + test2 <= 1.0: + istop = 5 + if 1.0 + t1 <= 1.0: + istop = 4 + + if conlim > 0.0 and test3 <= ctol: + istop = 3 + if test2 <= atol: + istop = 2 + if test1 <= rtol: + istop = 1 + + if istop > 0: + break + + return x.cpu().numpy() + + +def merge_panorama_depth_gpu(width: int, height: int, distance_maps: List[np.ndarray], pred_masks: List[np.ndarray], extrinsics: List[np.ndarray], intrinsics: List[np.ndarray]): + # `grad_equation` / `poisson_equation` defined in moge_panorama.py — pure + # geometry sparse skeletons; cheap to rebuild per level (sub-0.4s at top + # res). `vstack` stacks them after the per-face validity rows are + # selected. + import sys + import time + from scipy.sparse import vstack + from ...moge_panorama import grad_equation, poisson_equation + + def _p(msg): + print(f"[merge_panorama_depth_gpu {width:>4}x{height:<4}] {msg}", + file=sys.stderr, flush=True) + + t_lvl_start = time.perf_counter() + + if max(width, height) > 256: + panorama_depth_init, _ = merge_panorama_depth_gpu(width // 2, height // 2, distance_maps, pred_masks, extrinsics, intrinsics) + panorama_depth_init = cv2.resize(panorama_depth_init, (width, height), cv2.INTER_LINEAR) + else: + panorama_depth_init = None + + # Phase 2 (commit after 78380f4): the per-face cv2.remap loop + + # scipy.ndimage.convolve calls + np.stack/np.sum reductions are now + # one batched torch pipeline. All faces' projections, warps, gradients, + # Laplacians, and mask reductions go through a single set of CUDA + # kernels operating on (N, H, W) tensors. Math is fp32-equivalent to + # the original numpy/cv2 path: + # - grid_sample mode='bilinear' + padding_mode='border' matches + # cv2.INTER_LINEAR + BORDER_REPLICATE (verified by split helper). + # - Gradient via x-wrap pad then forward-diff is bit-identical to the + # np.pad + slice pair. + # - Laplacian via direct 5-point stencil (up + down + left + right - + # 4*center) is the same kernel as scipy.ndimage.convolve with the + # symmetric [[0,1,0],[1,-4,1],[0,1,0]] kernel. + # - Reductions (sum * mask / sum mask) reproduce the exact arithmetic + # of np.sum(maps * masks, axis=0) / np.sum(masks, axis=0).clip(1e-3). + + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + N = len(distance_maps) + fh, fw = distance_maps[0].shape + + # Sphere ray directions for the equirect grid. Tiny (~7 MB at 1920x960), + # CPU build is fine — utils3d.np.uv_map + spherical_uv_to_directions is + # vectorized numpy. + uv = utils3d.np.uv_map((height, width)) + spherical_directions_np = spherical_uv_to_directions(uv).astype(np.float32) + + # --- One-shot upload of the per-face data + camera params. --- + t_warp = time.perf_counter() + + dirs = torch.from_numpy(spherical_directions_np).to(device).reshape(-1, 3) # (H*W, 3) + dist_stack = np.stack([d.astype(np.float32, copy=False) for d in distance_maps]) # (N, fh, fw) + dist_gpu = torch.from_numpy(dist_stack).to(device) + log_dist_gpu = torch.log(dist_gpu.clamp(min=1e-30)) # (N, fh, fw) + mask_stack = np.stack([m.astype(np.uint8, copy=False) for m in pred_masks]) # (N, fh, fw) + pred_mask_gpu = torch.from_numpy(mask_stack).to(device).bool() + ext_gpu = torch.from_numpy(np.stack(extrinsics).astype(np.float32)).to(device) # (N, 4, 4) + intr_gpu = torch.from_numpy(np.stack(intrinsics).astype(np.float32)).to(device) # (N, 3, 3) + + # --- Batched projection: equirect rays → per-face image plane. --- + # p_cam = R @ dirs + t. Then K @ p_cam, perspective divide. + R = ext_gpu[:, :3, :3] # (N, 3, 3) + t = ext_gpu[:, :3, 3] # (N, 3) + p_cam = torch.einsum('nij,pj->npi', R, dirs) + t.unsqueeze(1) # (N, H*W, 3) + p_proj = torch.einsum('nij,npj->npi', intr_gpu, p_cam) # (N, H*W, 3) + depth = p_cam[..., 2] # (N, H*W) camera-Z + + # Safe divide: points behind the camera get garbage uv but the validity + # mask catches them. Avoids NaN poisoning grid_sample below. + safe_w = torch.where( + p_proj[..., 2:3] > 1e-12, + p_proj[..., 2:3], torch.ones_like(p_proj[..., 2:3]), + ) + projected_uv = p_proj[..., :2] / safe_w # (N, H*W, 2) in [0, 1] when valid + projection_valid = ( + (depth > 0) + & (projected_uv[..., 0] >= 0) & (projected_uv[..., 0] <= 1) + & (projected_uv[..., 1] >= 0) & (projected_uv[..., 1] <= 1) + ).reshape(N, height, width) # (N, H, W) bool + + # grid_sample convention: normalized [-1, +1] with align_corners=False + # = half-pixel-centre (centre of pixel 0 maps to -1 + 1/W). For + # utils3d-normalized uv (uv = (pixel + 0.5) / N), the conversion is + # the simple `grid = 2*uv - 1`. + grid = (projected_uv.clamp(0.0, 1.0) * 2.0 - 1.0).reshape(N, height, width, 2) + + # --- Batched bilinear sample of log-distance + nearest sample of mask. --- + sampled_log_dist = F.grid_sample( + log_dist_gpu.unsqueeze(1), # (N, 1, fh, fw) + grid, + mode='bilinear', padding_mode='border', align_corners=False, + ).squeeze(1) # (N, H, W) + sampled_pred_mask = F.grid_sample( + pred_mask_gpu.float().unsqueeze(1), + grid, + mode='nearest', padding_mode='border', align_corners=False, + ).squeeze(1) > 0.5 # (N, H, W) bool + + # Zero-out log-distance where projection is invalid (matches the CPU + # `np.where(projection_valid_mask, cv2.remap(...), 0)`). Same with the + # final per-face mask: AND of projection-valid and resampled-pred-mask. + panorama_log_distance = torch.where( + projection_valid, sampled_log_dist, + torch.zeros_like(sampled_log_dist), + ) # (N, H, W) + panorama_pred_mask_per_face = projection_valid & sampled_pred_mask # (N, H, W) + + # --- Batched gradient computation (wrap in x, no wrap in y). --- + # CPU equivalent: np.pad(((0,0),(0,1)), mode='wrap'); diff. Result + # shapes (H, W) for grad_x, (H-1, W+1) for grad_y. + log_dist_xpad = torch.cat( + [panorama_log_distance, panorama_log_distance[:, :, :1]], dim=-1 + ) # (N, H, W+1) + grad_x = log_dist_xpad[:, :, :-1] - log_dist_xpad[:, :, 1:] # (N, H, W) + grad_y = log_dist_xpad[:, :-1, :] - log_dist_xpad[:, 1:, :] # (N, H-1, W+1) + + mask_xpad = torch.cat( + [panorama_pred_mask_per_face, panorama_pred_mask_per_face[:, :, :1]], dim=-1 + ) # (N, H, W+1) + mask_x = mask_xpad[:, :, :-1] & mask_xpad[:, :, 1:] # (N, H, W) + mask_y = mask_xpad[:, :-1, :] & mask_xpad[:, 1:, :] # (N, H-1, W+1) + + # --- Batched 5-point Laplacian + 5-neighbour validity. --- + # CPU does np.pad y-edge, x-wrap then scipy.ndimage.convolve with the + # symmetric Laplacian kernel. We do the same padding then direct + # neighbour slicing — equivalent up to fp32 round-off and avoids the + # conv2d kernel launch overhead. + log_dist_4d = panorama_log_distance.unsqueeze(1) # (N, 1, H, W) + log_dist_pad = F.pad(log_dist_4d, (0, 0, 1, 1), mode='replicate') + log_dist_pad = F.pad(log_dist_pad, (1, 1, 0, 0), mode='circular') # (N, 1, H+2, W+2) + center = log_dist_pad[:, 0, 1:-1, 1:-1] # (N, H, W) + up = log_dist_pad[:, 0, :-2, 1:-1] + down = log_dist_pad[:, 0, 2:, 1:-1] + left = log_dist_pad[:, 0, 1:-1, :-2] + right = log_dist_pad[:, 0, 1:-1, 2:] + laplacian = up + down + left + right - 4.0 * center # (N, H, W) + + mask_4d = panorama_pred_mask_per_face.unsqueeze(1).float() # (N, 1, H, W) + mask_pad = F.pad(mask_4d, (0, 0, 1, 1), mode='replicate') + mask_pad = F.pad(mask_pad, (1, 1, 0, 0), mode='circular') # (N, 1, H+2, W+2) + mc = mask_pad[:, 0, 1:-1, 1:-1] > 0.5 + mu = mask_pad[:, 0, :-2, 1:-1] > 0.5 + md = mask_pad[:, 0, 2:, 1:-1] > 0.5 + ml = mask_pad[:, 0, 1:-1, :-2] > 0.5 + mr = mask_pad[:, 0, 1:-1, 2:] > 0.5 + mask_laplacian = mc & mu & md & ml & mr # (N, H, W) — bit-equiv to convolve(==5) + + # --- Weighted reductions across the face dimension. --- + # CPU: np.sum(maps * mask, axis=0) / np.sum(mask, axis=0).clip(1e-3). + # GPU equivalent — fused into 3 reductions, one each for grad_x, + # grad_y, laplacian. + mask_x_f = mask_x.float() + mask_y_f = mask_y.float() + mask_lap_f = mask_laplacian.float() + avg_grad_x = (grad_x * mask_x_f).sum(0) / mask_x_f.sum(0).clamp(min=1e-3) # (H, W) + avg_grad_y = (grad_y * mask_y_f).sum(0) / mask_y_f.sum(0).clamp(min=1e-3) # (H-1, W+1) + avg_lap = (laplacian * mask_lap_f).sum(0) / mask_lap_f.sum(0).clamp(min=1e-3) # (H, W) + + any_mask_x = mask_x.any(0) # (H, W) + any_mask_y = mask_y.any(0) # (H-1, W+1) + any_mask_lap = mask_laplacian.any(0) # (H, W) + panorama_pred_mask_final = panorama_pred_mask_per_face.any(0) # (H, W) + + # --- Single device→host transfer of the data we need on CPU for the + # sparse-system build below. --- + avg_grad_x_np = avg_grad_x.cpu().numpy() + avg_grad_y_np = avg_grad_y.cpu().numpy() + avg_lap_np = avg_lap.cpu().numpy() + grad_x_mask_np = any_mask_x.cpu().numpy().reshape(-1) + grad_y_mask_np = any_mask_y.cpu().numpy().reshape(-1) + grad_mask_np = np.concatenate([grad_x_mask_np, grad_y_mask_np]) + laplacian_mask_np = any_mask_lap.cpu().numpy().reshape(-1) + panorama_pred_mask_final_np = panorama_pred_mask_final.cpu().numpy() + + t_warp_done = time.perf_counter() + + # --- Solve overdetermined system (scipy CSR + on-device LSMR). --- + t_build = time.perf_counter() + A = vstack([ + grad_equation(width, height, wrap_x=True, wrap_y=False)[grad_mask_np], + poisson_equation(width, height, wrap_x=True, wrap_y=False)[laplacian_mask_np], + ]) + b = np.concatenate([ + avg_grad_x_np.reshape(-1)[grad_x_mask_np], + avg_grad_y_np.reshape(-1)[grad_y_mask_np], + avg_lap_np.reshape(-1)[laplacian_mask_np], + ]) + t_build_done = time.perf_counter() + + t_solve = time.perf_counter() + x = solve_lsmr_gpu( + A, b, + x0=np.log(panorama_depth_init).reshape(-1) if panorama_depth_init is not None else None, + ) + t_solve_done = time.perf_counter() + + panorama_depth = np.exp(x).reshape(height, width).astype(np.float32) + panorama_mask = panorama_pred_mask_final_np + + _p(f"done {time.perf_counter() - t_lvl_start:.2f}s " + f"(warp {t_warp_done - t_warp:.2f}, build {t_build_done - t_build:.2f}, " + f"LSMR {t_solve_done - t_solve:.2f}); " + f"depth min/med/max = {float(panorama_depth.min()):.3f}/" + f"{float(np.median(panorama_depth)):.3f}/" + f"{float(panorama_depth.max()):.3f}; " + f"valid {int(panorama_mask.sum())}/{panorama_mask.size}") + return panorama_depth, panorama_mask + + +def pred_pano_depth(model, image: Image.Image, scale=1.0, resize_to=1920, remove_pano_depth_nan=True): + """ + last_layer_mask: Previous-layer object mask. + last_layer_depth: Previous-layer depth. + """ + print("\t - Predicting pano depth with moge") + image_origin = np.array(image) + height_origin, width_origin = image_origin.shape[:2] + + image, height, width = image_origin, height_origin, width_origin + if resize_to is not None: + _height, _width = min(resize_to, int(resize_to * height_origin / width_origin)), min(resize_to, int(resize_to * width_origin / height_origin)) + if _height < height_origin: + print(f"\t - Resizing image from {width_origin}x{height_origin} to {_width}x{_height} for pano depth prediction") + image = cv2.resize(image_origin, (_width, _height), cv2.INTER_AREA) + height, width = _height, _width + + splitted_extrinsics, splitted_intriniscs = get_panorama_cameras_v2(subdivisions=1) + splitted_resolution = 512 + splitted_images = split_panorama_image(image, splitted_extrinsics, splitted_intriniscs, splitted_resolution, splitted_resolution, interp=cv2.INTER_AREA) + + # infer moge depth + num_splitted_images = len(splitted_images) + splitted_distance_maps = [None] * num_splitted_images + splitted_masks = [None] * num_splitted_images + + indices_to_process_model = [] + skipped_count = 0 + + for i in range(num_splitted_images): + indices_to_process_model.append(i) + + pred_count = 0 + # Process images that require model inference in batches + inference_batch_size = 1 + for i in range(0, len(indices_to_process_model), inference_batch_size): + batch_indices = indices_to_process_model[i: i + inference_batch_size] + if not batch_indices: + continue + + current_batch_images = [splitted_images[k] for k in batch_indices] + current_batch_intrinsics = [splitted_intriniscs[k] for k in batch_indices] + + image_tensor = torch.tensor( + np.stack(current_batch_images) / 255, + dtype=torch.float32, + device=next(model.parameters()).device, + ).permute(0, 3, 1, 2) + + fov_x, _ = np.rad2deg( # fov_y is not used by model.infer + utils3d.numpy.intrinsics_to_fov(np.array(current_batch_intrinsics)) + ) + fov_x_tensor = torch.tensor( + fov_x, dtype=torch.float32, device=next(model.parameters()).device + ) + + with torch.no_grad(), torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): + output = model.infer(image_tensor, fov_x=fov_x_tensor, apply_mask=False) + + batch_distance_maps = output["points"].norm(dim=-1).cpu().numpy() + batch_masks = output["mask"].cpu().numpy() + + for batch_idx, original_idx in enumerate(batch_indices): + splitted_distance_maps[original_idx] = batch_distance_maps[batch_idx] + splitted_masks[original_idx] = batch_masks[batch_idx] + pred_count += 1 + + if (pred_count + skipped_count) == 0: # Avoid division by zero if num_splitted_images is 0 + skip_ratio_info = "N/A (no images to process)" + else: + skip_ratio_info = f"{skipped_count / (pred_count + skipped_count):.2%}" + + print(f"\t 🔍 Predicted {pred_count} splitted images, skipped {skipped_count} splitted images. Skip ratio: {skip_ratio_info}") + + # merge moge depth + merging_width, merging_height = width, height + panorama_depth, panorama_mask = merge_panorama_depth_gpu( + merging_width, + merging_height, + splitted_distance_maps, + splitted_masks, + splitted_extrinsics, + splitted_intriniscs, + ) + + panorama_depth = panorama_depth.astype(np.float32) + # Align the left and right depths in the bottom region of pano depth. + if remove_pano_depth_nan: + # for depth inpainting, remove nan + panorama_depth[~panorama_mask] = 1.0 * np.nanquantile(panorama_depth, 0.999) # Sky depth. + panorama_depth = cv2.resize(panorama_depth, (width_origin, height_origin), cv2.INTER_LINEAR) + panorama_mask = cv2.resize(panorama_mask.astype(np.uint8), (width_origin, height_origin), cv2.INTER_NEAREST) > 0 + + # Smooth south-pole (bottom region) depth to fix left-right inconsistencies. + print("\t - Smoothing south pole depth for consistency") + panorama_depth = smooth_south_pole_depth(panorama_depth, smooth_height_ratio=0.05) + + rays = torch.from_numpy(spherical_uv_to_directions(utils3d.numpy.image_uv(width=width_origin, height=height_origin))).to(next(model.parameters()).device) + + panorama_depth = ( + torch.from_numpy(panorama_depth).to(next(model.parameters()).device) * scale + ) + + return { + "rgb": torch.from_numpy(image_origin).to(next(model.parameters()).device), + "distance": panorama_depth, + "rays": rays, + "mask": panorama_mask, + "splitted_masks": splitted_masks, + "splitted_distance_maps": splitted_distance_maps, + } + + +# Panorama depth stitching based on normal constraints. +def compute_spherical_ray_derivatives(spherical_directions: np.ndarray): + """ + Compute partial derivatives of panorama ray directions with respect to spherical coordinates. + + Args: + spherical_directions: (H, W, 3) Ray direction for each panorama pixel in world coordinates as unit vectors. + + Returns: + dray_dtheta: (H, W, 3) Partial derivative of rays with respect to azimuth theta. + dray_dphi: (H, W, 3) Partial derivative of rays with respect to elevation phi. + """ + H, W, _ = spherical_directions.shape + + # Recover spherical coordinates from ray directions. + # Assumes x-right, y-up, z-forward coordinates; adjust if the actual convention differs. + ray = spherical_directions + + # Azimuth theta and elevation phi. + # ray = (cos(φ)sin(θ), sin(φ), cos(φ)cos(θ)) + # Or adapt to the coordinate-system definition in use. + + # Method 1: numerical computation, which is more robust. + u = (np.arange(W) + 0.5) / W # [0, 1] + v = (np.arange(H) + 0.5) / H # [0, 1] + u_grid, v_grid = np.meshgrid(u, v) + + theta = (u_grid - 0.5) * 2 * np.pi # [-π, π] + phi = (0.5 - v_grid) * np.pi # [π/2, -π/2] + + cos_phi = np.cos(phi) + sin_phi = np.sin(phi) + cos_theta = np.cos(theta) + sin_theta = np.sin(theta) + + # ∂ray/∂θ, assuming standard equirectangular mapping. + # Adjust according to how spherical_directions is actually computed. + dray_dtheta = np.stack([ + cos_phi * cos_theta, + np.zeros_like(theta), + -cos_phi * sin_theta + ], axis=-1) + + # ∂ray/∂φ + dray_dphi = np.stack([ + -sin_phi * sin_theta, + cos_phi, + -sin_phi * cos_theta + ], axis=-1) + + return dray_dtheta, dray_dphi + + +def normal_to_log_distance_gradient( + panorama_normal: np.ndarray, + spherical_directions: np.ndarray, + width: int, + height: int +) -> tuple: + """ + Compute log-distance gradients from panorama normals. + + Derivation: + Surface point P = d * ray + Tangent vector ∂P/∂θ = (∂d/∂θ) * ray + d * (∂ray/∂θ) + The normal is perpendicular to the tangent: n · ∂P/∂θ = 0 + => ∂(log d)/∂θ = -(n · ∂ray/∂θ) / (n · ray) + + Args: + panorama_normal: (H, W, 3) Panorama normal map in world coordinates. + spherical_directions: (H, W, 3) Ray directions in world coordinates. + width, height: Panorama dimensions. + + Returns: + grad_x: (H, W) x-direction gradient corresponding to log_d[j] - log_d[j+1]. + grad_y: (H-1, W) y-direction gradient corresponding to log_d[i] - log_d[i+1]. + valid_mask: (H, W) Valid region. + """ + H, W = height, width + + # Get the ray-direction partial derivatives. + dray_dtheta, dray_dphi = compute_spherical_ray_derivatives(spherical_directions) + + # Compute dot products. + n_dot_ray = np.sum(panorama_normal * spherical_directions, axis=-1) # (H, W) + n_dot_dray_dtheta = np.sum(panorama_normal * dray_dtheta, axis=-1) # (H, W) + n_dot_dray_dphi = np.sum(panorama_normal * dray_dphi, axis=-1) # (H, W) + + # Validity check: normals cannot be perpendicular to the view direction. + eps = 1e-4 + valid_mask = np.abs(n_dot_ray) > eps + + # Safe division. + n_dot_ray_safe = np.where(valid_mask, n_dot_ray, 1.0) + + # Continuous-space log(d) gradients with respect to angles. + # ∂(log d)/∂θ = -(n · ∂ray/∂θ) / (n · ray) + dlogd_dtheta = -n_dot_dray_dtheta / n_dot_ray_safe # (H, W) + dlogd_dphi = -n_dot_dray_dphi / n_dot_ray_safe # (H, W) + + # Convert to discrete pixel gradients. + # In this code, grad_x = log_d[j] - log_d[j+1] = -∂(log d)/∂θ * Δθ. + # Δθ = 2π / W, the θ change per pixel. + # grad_y = log_d[i] - log_d[i+1] = -∂(log d)/∂φ * Δφ = ∂(log d)/∂φ * (π/H) + + delta_theta = 2 * np.pi / W + delta_phi = np.pi / H + + # Pixel-scale continuous gradients. + pixel_dlogd_dx = -dlogd_dtheta * delta_theta # Corresponds to log_d[j] - log_d[j+1]. + pixel_dlogd_dy = dlogd_dphi * delta_phi # Corresponds to log_d[i] - log_d[i+1]. + + # Discrete gradients, averaging both sides at boundaries. + # X direction (wrap). + padded_grad_x = np.pad(pixel_dlogd_dx, ((0, 0), (0, 1)), mode='wrap') + grad_x = (padded_grad_x[:, :-1] + padded_grad_x[:, 1:]) / 2 # (H, W) + + # Y direction (no wrap). + grad_y = (pixel_dlogd_dy[:-1, :] + pixel_dlogd_dy[1:, :]) / 2 # (H-1, W) + + # Valid masks for gradients + padded_valid = np.pad(valid_mask, ((0, 0), (0, 1)), mode='wrap') + mask_x = padded_valid[:, :-1] & padded_valid[:, 1:] # (H, W) + mask_y = valid_mask[:-1, :] & valid_mask[1:, :] # (H-1, W) + + return grad_x, grad_y, mask_x, mask_y, valid_mask + + +def grad_equation_separate(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False): + """ + Return gradient equation matrices for the x and y directions separately. + + Returns: + grad_eq_x: x-direction gradient matrix, with H * W rows if wrap_x, otherwise H * (W - 1). + grad_eq_y: y-direction gradient matrix, with (H - 1) * W rows if wrap_y is false. + """ + grid_index = np.arange(width * height).reshape(height, width) + + # X direction. + if wrap_x: + grid_x = np.pad(grid_index, ((0, 0), (0, 1)), mode='wrap') + else: + grid_x = grid_index + + n_grad_x = grid_x.shape[0] * (grid_x.shape[1] - 1) + data_x = np.concatenate([ + np.ones((grid_x.shape[0], grid_x.shape[1] - 1), dtype=np.float32).reshape(-1, 1), + -np.ones((grid_x.shape[0], grid_x.shape[1] - 1), dtype=np.float32).reshape(-1, 1), + ], axis=1).reshape(-1) + indices_x = np.concatenate([ + grid_x[:, :-1].reshape(-1, 1), + grid_x[:, 1:].reshape(-1, 1), + ], axis=1).reshape(-1) + indptr_x = np.arange(0, n_grad_x * 2 + 1, 2) + grad_eq_x = csr_array((data_x, indices_x, indptr_x), shape=(n_grad_x, height * width)) + + # Y direction, usually without wrapping. + if wrap_y: + grid_y = np.pad(grid_index, ((0, 1), (0, 0)), mode='wrap') + else: + grid_y = grid_index + + n_grad_y = (grid_y.shape[0] - 1) * grid_y.shape[1] + data_y = np.concatenate([ + np.ones((grid_y.shape[0] - 1, grid_y.shape[1]), dtype=np.float32).reshape(-1, 1), + -np.ones((grid_y.shape[0] - 1, grid_y.shape[1]), dtype=np.float32).reshape(-1, 1), + ], axis=1).reshape(-1) + indices_y = np.concatenate([ + grid_y[:-1, :].reshape(-1, 1), + grid_y[1:, :].reshape(-1, 1), + ], axis=1).reshape(-1) + indptr_y = np.arange(0, n_grad_y * 2 + 1, 2) + grad_eq_y = csr_array((data_y, indices_y, indptr_y), shape=(n_grad_y, height * width)) + + return grad_eq_x, grad_eq_y + + + + + + + + + + +def spherical_uv_to_directions(uv: np.ndarray): + theta, phi = (1 - uv[..., 0]) * (2 * np.pi), uv[..., 1] * np.pi + directions = np.stack([np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)], axis=-1) + return directions + + +def convert_rgbd2pcd_panorama( + rgb: torch.Tensor, # (H, W, 3) RGB image, values [0, 1] + distance: torch.Tensor, # (H, W) Distance map + rays: torch.Tensor, # (H, W, 3) Ray directions (unit vectors ideally) + excluded_region_mask: Optional[torch.Tensor] = None, # (H, W) Optional boolean mask + max_size: int = 4096, # Max dimension for resizing + device: Literal["cuda", "cpu"] = "cuda", # Computation device + dropout_pcd=False +): + """ + Converts panoramic RGBD data (image, distance, rays) into an Open3D mesh. + + Args: + image: Input RGB image tensor (H, W, 3), uint8 or float [0, 255]. + distance: Input distance map tensor (H, W). + rays: Input ray directions tensor (H, W, 3). Assumed to originate from (0,0,0). + excluded_region_mask: Optional boolean mask tensor (H, W). True values indicate regions to potentially exclude. + max_size: Maximum size (height or width) to resize inputs to. + device: The torch device ('cuda' or 'cpu') to use for computations. + + Returns: + An Open3D TriangleMesh object. + """ + assert rgb.ndim == 3 and rgb.shape[2] == 3, "Image must be HxWx3" + assert distance.ndim == 2, "Distance must be HxW" + assert rays.ndim == 3 and rays.shape[2] == 3, "Rays must be HxWx3" + assert ( + rgb.shape[:2] == distance.shape[:2] == rays.shape[:2] + ), "Input shapes must match" + + mask = excluded_region_mask + + if mask is not None: + assert ( + mask.ndim == 2 and mask.shape[:2] == rgb.shape[:2] + ), "Mask shape must match" + assert mask.dtype == torch.bool, "Mask must be a boolean tensor" + + rgb = rgb.to(device) + distance = distance.to(device) + rays = rays.to(device) + if mask is not None: + mask = mask.to(device) + + H, W = distance.shape + if max(H, W) > max_size: + scale = max_size / max(H, W) + else: + scale = 1.0 + + rgb_nchw = rgb.permute(2, 0, 1).unsqueeze(0) + distance_nchw = distance.unsqueeze(0).unsqueeze(0) + rays_nchw = rays.permute(2, 0, 1).unsqueeze(0) + + rgb_resized = ( + F.interpolate( + rgb_nchw, + scale_factor=scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + .squeeze(0) + .permute(1, 2, 0) + ) + + distance_resized = ( + F.interpolate( + distance_nchw, + scale_factor=scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + .squeeze(0) + .squeeze(0) + ) + + rays_resized_nchw = F.interpolate( + rays_nchw, + scale_factor=scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + + # IMPORTANT: Renormalize ray directions after interpolation + rays_resized = rays_resized_nchw.squeeze(0).permute(1, 2, 0) + rays_norm = torch.linalg.norm(rays_resized, dim=-1, keepdim=True) + rays_resized = rays_resized / (rays_norm + 1e-8) + + if mask is not None: + mask_resized = ( + F.interpolate( + mask.unsqueeze(0).unsqueeze(0).float(), # Needs float for interpolation + scale_factor=scale, + mode="nearest", # Or 'nearest' if sharp boundaries are critical + # align_corners=False, + recompute_scale_factor=False, + ) + .squeeze(0) + .squeeze(0) + ) + mask_resized = mask_resized > 0.5 # Convert back to boolean + else: + mask_resized = None + + # --- Calculate 3D Vertices --- + # Vertex position = origin + distance * ray_direction + # Assuming origin is (0, 0, 0) + distance_flat = distance_resized.reshape(-1, 1) # (H*W, 1) + rays_flat = rays_resized.reshape(-1, 3) # (H*W, 3) + vertices = distance_flat * rays_flat # (H*W, 3) + vertex_colors = rgb_resized.reshape(-1, 3) # (H*W, 3) + if mask_resized is not None: + mask_resized = mask_resized.reshape(-1, ) + vertices = vertices[~mask_resized] + vertex_colors = vertex_colors[~mask_resized] + + # downsample + if dropout_pcd and vertices.shape[0] > 1_000_000: + rdx = np.arange(vertices.shape[0]) + np.random.shuffle(rdx) + rdx = rdx[:1_000_000] + vertices = vertices[rdx] + vertex_colors = vertex_colors[rdx] + + pcd = trimesh.PointCloud(vertices=vertices.cpu().numpy(), colors=vertex_colors.cpu().numpy()) + + return pcd + + +def convert_rgbd2pcd_multi_scale_panorama( + rgb: torch.Tensor, # (H, W, 3) RGB image, values [0, 1] + distance: torch.Tensor, # (H, W) Distance map + rays: torch.Tensor, # (H, W, 3) Ray directions (unit vectors ideally) + excluded_region_mask: Optional[torch.Tensor] = None, # (H, W) Optional boolean mask + device: Literal["cuda", "cpu"] = "cuda", # Computation device + depth_intervals=[0, 1, 2, 4, 8] +): + """ + Converts panoramic RGBD data (image, distance, rays) into an Open3D mesh. + + Args: + image: Input RGB image tensor (H, W, 3), uint8 or float [0, 255]. + distance: Input distance map tensor (H, W). + rays: Input ray directions tensor (H, W, 3). Assumed to originate from (0,0,0). + excluded_region_mask: Optional boolean mask tensor (H, W). True values indicate regions to potentially exclude. + device: The torch device ('cuda' or 'cpu') to use for computations. + + Returns: + An Open3D TriangleMesh object. + """ + assert rgb.ndim == 3 and rgb.shape[2] == 3, "Image must be HxWx3" + assert distance.ndim == 2, "Distance must be HxW" + assert rays.ndim == 3 and rays.shape[2] == 3, "Rays must be HxWx3" + assert ( + rgb.shape[:2] == distance.shape[:2] == rays.shape[:2] + ), "Input shapes must match" + + mask = excluded_region_mask + + if mask is not None: + assert ( + mask.ndim == 2 and mask.shape[:2] == rgb.shape[:2] + ), "Mask shape must match" + assert mask.dtype == torch.bool, "Mask must be a boolean tensor" + + rgb = rgb.to(device) + distance = distance.to(device) + rays = rays.to(device) + if mask is not None: + mask = mask.to(device) + + rgb_nchw = rgb.permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W] + distance_nchw = distance.unsqueeze(0).unsqueeze(0) # [1, 1, H, W] + rays_nchw = rays.permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W] + + median_distance = torch.median(distance).item() + + total_points = [] + total_colors = [] + + for i in tqdm(range(1, len(depth_intervals)), desc="Processing depth intervals"): + if i == len(depth_intervals) - 1: + interval_mask = distance_nchw > (median_distance * depth_intervals[i - 1]) + else: + interval_mask = ((median_distance * depth_intervals[i - 1]) < distance_nchw) & (distance_nchw <= (median_distance * depth_intervals[i])) + + # pointclouds number ∝ depth^2 + resize_scale = depth_intervals[i] + if interval_mask.sum() == 0: + continue + + rgb_resized = ( + F.interpolate( + rgb_nchw, + scale_factor=resize_scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + .squeeze(0) + .permute(1, 2, 0) + ) + + distance_resized = ( + F.interpolate( + distance_nchw, + scale_factor=resize_scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + .squeeze(0) + .squeeze(0) + ) + + rays_resized_nchw = F.interpolate( + rays_nchw, + scale_factor=resize_scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + + interval_mask_resized = F.interpolate(interval_mask.float(), + scale_factor=resize_scale, + mode="nearest", + recompute_scale_factor=False).bool().squeeze(0).squeeze(0) + + # IMPORTANT: Renormalize ray directions after interpolation + rays_resized = rays_resized_nchw.squeeze(0).permute(1, 2, 0) + rays_norm = torch.linalg.norm(rays_resized, dim=-1, keepdim=True) + rays_resized = rays_resized / (rays_norm + 1e-8) + + if mask is not None: + mask_resized = ( + F.interpolate( + mask.unsqueeze(0).unsqueeze(0).float(), # Needs float for interpolation + scale_factor=resize_scale, + mode="nearest", # Or 'nearest' if sharp boundaries are critical + # align_corners=False, + recompute_scale_factor=False, + ) + .squeeze(0) + .squeeze(0) + ) + mask_resized = mask_resized > 0.5 # Convert back to boolean + else: + mask_resized = None + + mask_resized = mask_resized & interval_mask_resized + + # --- Calculate 3D Vertices --- + # Vertex position = origin + distance * ray_direction + # Assuming origin is (0, 0, 0) + distance_flat = distance_resized.reshape(-1, 1) # (H*W, 1) + rays_flat = rays_resized.reshape(-1, 3) # (H*W, 3) + vertices = distance_flat * rays_flat # (H*W, 3) + vertex_colors = rgb_resized.reshape(-1, 3) # (H*W, 3) + if mask_resized is not None: + mask_resized = mask_resized.reshape(-1, ) + vertices = vertices[~mask_resized] + vertex_colors = vertex_colors[~mask_resized] + + total_points.append(vertices) + total_colors.append(vertex_colors) + print(f"Depth interval: {depth_intervals[i]}, Number of points: {vertices.shape[0]}") + + vertices = torch.cat(total_points, dim=0) + vertex_colors = torch.cat(total_colors, dim=0) + pcd = trimesh.PointCloud(vertices=vertices.cpu().numpy(), colors=vertex_colors.cpu().numpy()) + + return pcd + + +def convert_rgbd2pcd_panorama_da360(depth, rgb, mask=None, dropout_pcd=True): + h, w = depth.shape + Theta = np.arange(h).reshape(h, 1) * np.pi / h + np.pi / h / 2 + Theta = np.repeat(Theta, w, axis=1) + Phi = np.arange(w).reshape(1, w) * 2 * np.pi / w + np.pi / w - np.pi + Phi = -np.repeat(Phi, h, axis=0) + + X = depth * np.sin(Theta) * np.sin(Phi) + Y = depth * np.cos(Theta) + Z = depth * np.sin(Theta) * np.cos(Phi) + + if mask is None: + X = X.flatten() + Y = Y.flatten() + Z = Z.flatten() + R = rgb[:, :, 0].flatten() + G = rgb[:, :, 1].flatten() + B = rgb[:, :, 2].flatten() + else: + X = X[mask] + Y = Y[mask] + Z = Z[mask] + R = rgb[:, :, 0][mask] + G = rgb[:, :, 1][mask] + B = rgb[:, :, 2][mask] + + XYZ = np.stack([X, Y, Z], axis=1) + RGB = np.stack([R, G, B], axis=1) + + # downsample + if dropout_pcd and XYZ.shape[0] > 1_000_000: + rdx = np.arange(XYZ.shape[0]) + np.random.shuffle(rdx) + rdx = rdx[:1_000_000] + XYZ = XYZ[rdx] + RGB = RGB[rdx] + + pcd = trimesh.PointCloud(vertices=XYZ, colors=RGB) + return pcd + + +def _generate_faces_numpy(H: int, W: int, mask: Optional[torch.Tensor]) -> np.ndarray: + """ + Pure NumPy implementation, 2-3x faster than the PyTorch version. + """ + # Precompute all vertex indices. + idx = np.arange(H * W, dtype=np.int32).reshape(H, W) + + # Four corners of each quad, with horizontal wrapping. + tl = idx[:-1, :] # top-left + tr = idx[:-1, :].copy() + tr[:, :-1] = idx[:-1, 1:] + tr[:, -1] = idx[:-1, 0] # wrap + bl = idx[1:, :] # bottom-left + br = idx[1:, :].copy() + br[:, :-1] = idx[1:, 1:] + br[:, -1] = idx[1:, 0] # wrap + + # Apply mask. + if mask is not None: + mask_np = mask.cpu().numpy() + # Check whether any of the four corners is masked. + m_tl = mask_np[:-1, :] + m_tr = np.roll(mask_np[:-1, :], -1, axis=1) + m_bl = mask_np[1:, :] + m_br = np.roll(mask_np[1:, :], -1, axis=1) + + keep = ~(m_tl | m_tr | m_bl | m_br) + + tl, tr, bl, br = tl[keep], tr[keep], bl[keep], br[keep] + else: + tl, tr, bl, br = tl.ravel(), tr.ravel(), bl.ravel(), br.ravel() + + # Build triangles, two per quad. + n = len(tl) + faces = np.empty((2 * n, 3), dtype=np.int32) + faces[0::2] = np.column_stack([tl, tr, bl]) + faces[1::2] = np.column_stack([tr, br, bl]) + + return faces + + +def convert_rgbd2mesh_panorama( + rgb: torch.Tensor, # (H, W, 3) RGB image, values [0, 1] + distance: torch.Tensor, # (H, W) Distance map + rays: torch.Tensor, # (H, W, 3) Ray directions (unit vectors ideally) + excluded_region_mask: Optional[torch.Tensor] = None, # (H, W) Optional boolean mask + max_size: int = 4096, # Max dimension for resizing + device: Literal["cuda", "cpu"] = "cuda", # Computation device + connect_boundary_max_dist: Optional[float] = 0.5, # Max distance to bridge boundary vertices + connect_boundary_repeat_times: int = 2 +) -> o3d.geometry.TriangleMesh: + """ + Converts panoramic RGBD data (image, distance, rays) into an Open3D mesh. + + Args: + image: Input RGB image tensor (H, W, 3), uint8 or float [0, 255]. + distance: Input distance map tensor (H, W). + rays: Input ray directions tensor (H, W, 3). Assumed to originate from (0,0,0). + excluded_region_mask: Optional boolean mask tensor (H, W). True values indicate regions to potentially exclude. + max_size: Maximum size (height or width) to resize inputs to. + device: The torch device ('cuda' or 'cpu') to use for computations. + + Returns: + An Open3D TriangleMesh object. + """ + """Optimized version: about 3-5x faster.""" + H, W = distance.shape + scale = min(1.0, max_size / max(H, W)) + need_resize = scale < 1.0 + + # ========== 1. Data preparation with asynchronous transfers. ========== + rgb = rgb.to(device, non_blocking=True) + distance = distance.to(device, non_blocking=True) + rays = rays.to(device, non_blocking=True) + if excluded_region_mask is not None: + mask = excluded_region_mask.to(device, non_blocking=True) + else: + mask = None + + # ========== 2. Scaling. ========== + if need_resize: + H_new, W_new = int(H * scale), int(W * scale) + + # Combine rgb and rays to reduce interpolation calls. + combined = torch.cat([rgb, rays], dim=-1).permute(2, 0, 1).unsqueeze(0) + combined_resized = F.interpolate(combined, size=(H_new, W_new), mode='bilinear', align_corners=False) + combined_resized = combined_resized.squeeze(0).permute(1, 2, 0) + + rgb_resized = combined_resized[..., :3] + rays_resized = F.normalize(combined_resized[..., 3:], dim=-1) + + distance_resized = F.interpolate( + distance[None, None], size=(H_new, W_new), mode='bilinear', align_corners=False + ).squeeze() + + if mask is not None: + mask_resized = F.max_pool2d( + mask[None, None].float(), + kernel_size=int(1 / scale), + stride=int(1 / scale) + ).squeeze().bool() + # Ensure the size matches. + if mask_resized.shape != (H_new, W_new): + mask_resized = F.interpolate( + mask[None, None].float(), size=(H_new, W_new), mode='nearest' + ).squeeze().bool() + else: + mask_resized = None + else: + H_new, W_new = H, W + rgb_resized = rgb + rays_resized = F.normalize(rays, dim=-1) + distance_resized = distance + mask_resized = mask + + # ========== 3. Compute vertices on GPU. ========== + vertices = (distance_resized.unsqueeze(-1) * rays_resized).reshape(-1, 3) + vertex_colors = rgb_resized.reshape(-1, 3) + + # Synchronize and convert to NumPy. + torch.cuda.synchronize() if device == 'cuda' else None + vertices_np = vertices.cpu().numpy().astype(np.float64) + colors_np = vertex_colors.cpu().numpy().astype(np.float64) + + # ========== 4. Generate faces on CPU/NumPy, which is faster. ========== + faces_np = _generate_faces_numpy(H_new, W_new, mask_resized) + + # ========== 5. Create Open3D Mesh. ========== + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(vertices_np) + mesh.triangles = o3d.utility.Vector3iVector(faces_np) + mesh.vertex_colors = o3d.utility.Vector3dVector(colors_np) + + mesh.remove_unreferenced_vertices() + mesh.remove_degenerate_triangles() + + # ========== 6. Boundary handling. ========== + if connect_boundary_max_dist is not None and connect_boundary_max_dist > 0: + mesh = _fill_small_boundary_spikes(mesh, connect_boundary_max_dist, connect_boundary_repeat_times) + # Recompute normals after potential modification, if mesh still valid + if mesh.has_triangles() and mesh.has_vertices(): + mesh.compute_vertex_normals() + mesh.compute_triangle_normals() # Also computes triangle normals if vertex normals are computed + + return mesh + + +def _fill_small_boundary_spikes( + mesh: o3d.geometry.TriangleMesh, + max_bridge_dist: float, + repeat_times: int = 3 +) -> o3d.geometry.TriangleMesh: + print(f"\t - DEBUG: Filling small boundary spikes with max_bridge_dist: {max_bridge_dist} and repeat_times: {repeat_times}") + for iteration in range(repeat_times): + if not mesh.has_triangles() or not mesh.has_vertices(): + return mesh + + vertices = np.asarray(mesh.vertices) + triangles = np.asarray(mesh.triangles) + + # 1. Identify boundary edges + edge_to_triangle_count = defaultdict(int) + + for tri_idx, tri in enumerate(triangles): + for i in range(3): + v1_idx, v2_idx = tri[i], tri[(i + 1) % 3] + edge = tuple(sorted((v1_idx, v2_idx))) + edge_to_triangle_count[edge] += 1 + + boundary_edges = [edge for edge, count in edge_to_triangle_count.items() if count == 1] + + if not boundary_edges: + return mesh + + # 2. Create an adjacency list for boundary vertices using only boundary edges + boundary_adj = defaultdict(list) + for v1_idx, v2_idx in boundary_edges: + boundary_adj[v1_idx].append(v2_idx) + boundary_adj[v2_idx].append(v1_idx) + + # 3. Process boundary vertices with new smooth filling algorithm + new_triangles_list = [] + edge_added = defaultdict(bool) + + # print(f"DEBUG: Found {len(boundary_edges)} boundary edges.") + # print(f"DEBUG: Max bridge distance set to: {max_bridge_dist}") + + new_triangles_added_count = 0 + + for v_curr_idx, neighbors in boundary_adj.items(): + if len(neighbors) != 2: # Only process vertices with exactly 2 boundary neighbors + continue + + v_a_idx, v_b_idx = neighbors[0], neighbors[1] + + # Skip if these vertices already form a triangle + potential_edge = tuple(sorted((v_a_idx, v_b_idx))) + if edge_to_triangle_count[potential_edge] > 0 or edge_added[potential_edge]: + continue + + # Calculate distances + v_curr_coord = vertices[v_curr_idx] + v_a_coord = vertices[v_a_idx] + v_b_coord = vertices[v_b_idx] + + dist_a_b = np.linalg.norm(v_a_coord - v_b_coord) + + # Skip if distance exceeds threshold + if dist_a_b > max_bridge_dist: + continue + + # Create simple triangle (v_a, v_b, v_curr) + new_triangles_list.append([v_a_idx, v_b_idx, v_curr_idx]) + new_triangles_added_count += 1 + edge_added[potential_edge] = True + + # Mark edges as processed + edge_added[tuple(sorted((v_curr_idx, v_a_idx)))] = True + edge_added[tuple(sorted((v_curr_idx, v_b_idx)))] = True + + # 4. Now process multi-step connections for better smoothing + # First build boundary chains for multi-step connections + boundary_loops = [] + visited_vertices = set() + + # Find boundary vertices with exactly 2 neighbors (part of continuous chains) + chain_starts = [v for v in boundary_adj if len(boundary_adj[v]) == 2 and v not in visited_vertices] + + for start_vertex in chain_starts: + if start_vertex in visited_vertices: + continue + + chain = [] + curr_vertex = start_vertex + + # Follow the chain in one direction + while curr_vertex not in visited_vertices: + visited_vertices.add(curr_vertex) + chain.append(curr_vertex) + + next_candidates = [n for n in boundary_adj[curr_vertex] if n not in visited_vertices] + if not next_candidates: + break + + curr_vertex = next_candidates[0] + + if len(chain) >= 3: + boundary_loops.append(chain) + + # print(f"DEBUG: Found {len(boundary_loops)} boundary chains for smoothing.") + + # Process each boundary chain for multi-step smoothing + for chain in boundary_loops: + chain_length = len(chain) + + # Skip very small chains + if chain_length < 3: + continue + + # Compute multi-step connections + max_step = min(8, chain_length - 1) + + for i in range(chain_length): + anchor_idx = chain[i] + anchor_coord = vertices[anchor_idx] + + for step in range(3, max_step + 1): + if i + step >= chain_length: + break + + far_idx = chain[i + step] + far_coord = vertices[far_idx] + + # Check distance criteria + dist_anchor_far = np.linalg.norm(anchor_coord - far_coord) + if dist_anchor_far > max_bridge_dist * step: + continue + + # Check if anchor and far are already connected + edge_anchor_far = tuple(sorted((anchor_idx, far_idx))) + if edge_to_triangle_count[edge_anchor_far] > 0 or edge_added[edge_anchor_far]: + continue + + # Create fan triangles + fan_valid = True + fan_triangles = [] + + prev_mid_idx = anchor_idx + + for j in range(1, step): + mid_idx = chain[i + j] + + if prev_mid_idx != anchor_idx: + tri_edge1 = tuple(sorted((anchor_idx, mid_idx))) + tri_edge2 = tuple(sorted((prev_mid_idx, mid_idx))) + + # Check if edges already exist (not created by our fan) + if (edge_to_triangle_count[tri_edge1] > 0 and not edge_added[tri_edge1]) or \ + (edge_to_triangle_count[tri_edge2] > 0 and not edge_added[tri_edge2]): + fan_valid = False + break + + fan_triangles.append([anchor_idx, prev_mid_idx, mid_idx]) + + prev_mid_idx = mid_idx + + # Add final triangle to connect to far_idx + if fan_valid: + fan_triangles.append([anchor_idx, prev_mid_idx, far_idx]) + + # Add all fan triangles if valid + if fan_valid and fan_triangles: + for triangle in fan_triangles: + v_a, v_b, v_c = triangle + edge_ab = tuple(sorted((v_a, v_b))) + edge_bc = tuple(sorted((v_b, v_c))) + edge_ac = tuple(sorted((v_a, v_c))) + + new_triangles_list.append(triangle) + new_triangles_added_count += 1 + + edge_added[edge_ab] = True + edge_added[edge_bc] = True + edge_added[edge_ac] = True + + # Once we've added a fan, move to the next anchor + break + + # print(f"DEBUG: Total new triangles added in iteration {iteration}: {new_triangles_added_count}") + + if new_triangles_added_count == 0: + break + + # Update the mesh with new triangles + if new_triangles_list: + all_triangles_np = np.vstack((triangles, np.array(new_triangles_list, dtype=np.int32))) + + final_mesh = o3d.geometry.TriangleMesh() + final_mesh.vertices = o3d.utility.Vector3dVector(vertices) + final_mesh.triangles = o3d.utility.Vector3iVector(all_triangles_np) + + if mesh.has_vertex_colors(): + final_mesh.vertex_colors = mesh.vertex_colors + + # Clean up the mesh + final_mesh.remove_degenerate_triangles() + final_mesh.remove_unreferenced_vertices() + mesh = final_mesh + + return mesh + + +def get_view_point_from_panorama_point(global_pcd, w2c, K, image_h, image_w): + # Get valid points corresponding to the current view. + projected_uv, projected_depth = utils3d.numpy.project_cv(global_pcd.vertices, extrinsics=w2c, intrinsics=K) + projection_valid_mask = (projected_depth > 0) & (projected_uv > 0).all(axis=-1) & (projected_uv < 1).all(axis=-1) + projected_uv = projected_uv[projection_valid_mask] + projected_uv[:, 0] = (projected_uv[:, 0] * image_w).round() + projected_uv[:, 1] = (projected_uv[:, 1] * image_h).round() + projected_uv = projected_uv.astype(np.int64) + projected_uv[:, 0] = np.clip(projected_uv[:, 0], 0, image_w - 1) + projected_uv[:, 1] = np.clip(projected_uv[:, 1], 0, image_h - 1) + projected_depth = projected_depth[projection_valid_mask] + + projected_uv_1d = projected_uv[:, 1] * image_w + projected_uv[:, 0] # Convert coordinates to 1D. + original_indices = np.arange(projected_uv_1d.shape[0]) + # Double sort: for equal 1D coordinates, sort by depth ascending so only the nearest point is kept for each u, v. + sorted_indices = np.lexsort((-projected_depth, projected_uv_1d)) + projected_uv_1d_sorted = projected_uv_1d[sorted_indices] + sub_uv_1d = projected_uv_1d_sorted[:-1] - projected_uv_1d_sorted[1:] # Offset subtraction; nonzero entries have the minimum depth. + final_valid_indices = original_indices[sorted_indices][:-1][(sub_uv_1d != 0)] + + projected_points = global_pcd.vertices[projection_valid_mask][final_valid_indices] + projected_uv = projected_uv[final_valid_indices] + projected_colors = global_pcd.colors[projection_valid_mask][final_valid_indices, :3] + + return projected_points, projected_colors, projected_uv + + +def smooth_sky_depth_boundary( + depth: torch.Tensor, + sky_mask: torch.Tensor, + transition_width: int = 50, + depth_max: float = None, + method: str = 'mean' # 'mean', 'median', 'gaussian' +) -> torch.Tensor: + """ + Smooth the depth transition between sky and foreground. + + Args: + depth: Depth map, with sky regions already set to depth_max. + sky_mask: Sky mask, where True indicates sky. + transition_width: Transition-region width in pixels. + depth_max: Sky depth value. + method: Boundary diffusion method. + - 'mean': Weighted mean diffusion, recommended. + - 'median': Boundary median. + - 'gaussian': Gaussian-blur diffusion. + + Returns: + Depth map after boundary smoothing. + """ + # 1. Dimension handling. + original_dim = depth.dim() + if original_dim == 2: + depth = depth.unsqueeze(0).unsqueeze(0) + sky_mask = sky_mask.unsqueeze(0).unsqueeze(0) + elif original_dim == 3: + depth = depth.unsqueeze(0) + sky_mask = sky_mask.unsqueeze(0) + + device = depth.device + dtype = depth.dtype + + if depth_max is None: + depth_max = depth.max().item() + + sky_mask_np = sky_mask.squeeze().cpu().numpy().astype(np.uint8) + depth_np = depth.squeeze().cpu().numpy() + + if sky_mask_np.sum() == 0 or sky_mask_np.sum() == sky_mask_np.size: + return _restore_dim(depth, original_dim) + + # --------------------------------------------------------- + # 2. Compute the distance from each sky pixel to the boundary. + # --------------------------------------------------------- + dist_to_foreground = cv2.distanceTransform(sky_mask_np, cv2.DIST_L2, 5) + + # --------------------------------------------------------- + # 3. Get boundary depth values according to the selected method. + # --------------------------------------------------------- + sky_mask_tensor = sky_mask.bool() + + if method == 'mean': + boundary_depth = _mean_diffusion(depth, sky_mask_tensor, transition_width) + elif method == 'median': + boundary_depth = _median_boundary(depth_np, sky_mask_np, depth_max, device, dtype, depth.shape) + elif method == 'gaussian': + boundary_depth = _gaussian_diffusion(depth, sky_mask_tensor, depth_max, transition_width) + else: + raise ValueError(f"Unknown method: {method}") + + # --------------------------------------------------------- + # 4. Linear interpolation: use diffused depth at the boundary and depth_max deeper in the sky. + # --------------------------------------------------------- + t = np.clip(dist_to_foreground / max(transition_width, 1), 0, 1) + t = t * t * (3 - 2 * t) # smoothstep + t = torch.from_numpy(t).to(device=device, dtype=dtype).view_as(depth) + + interpolated = boundary_depth * (1 - t) + depth_max * t + + # --------------------------------------------------------- + # 5. Clamp the range and modify only sky regions. + # --------------------------------------------------------- + interpolated = torch.clamp(interpolated, max=depth_max) + result = torch.where(sky_mask_tensor, interpolated, depth) + + return _restore_dim(result, original_dim) + + +def _mean_diffusion(depth: torch.Tensor, sky_mask: torch.Tensor, transition_width: int) -> torch.Tensor: + """ + Weighted mean diffusion, recommended. + Preserves local depth characteristics while avoiding extreme values. + """ + device = depth.device + dtype = depth.dtype + + # Valid-region mask, where non-sky equals 1. + valid_mask = (~sky_mask).float() + + # Initialize sky regions to 0. + flood_depth = depth.clone() + flood_depth[sky_mask] = 0 + + kernel_size = 31 + pad = kernel_size // 2 + max_iterations = (transition_width // pad) + 10 + + for _ in range(max_iterations): + # Compute the sum of neighboring depths. + depth_sum = F.avg_pool2d(flood_depth, kernel_size, stride=1, padding=pad) * (kernel_size ** 2) + # Compute the number of valid neighboring pixels. + valid_count = F.avg_pool2d(valid_mask, kernel_size, stride=1, padding=pad) * (kernel_size ** 2) + + # Mean = total sum / valid count. + mean_depth = depth_sum / (valid_count + 1e-8) + + # Update only sky pixels whose neighborhoods contain valid pixels. + can_update = sky_mask & (valid_count > 0) + + # Check whether any pixels still need updates. + newly_filled = can_update & (valid_mask < 0.5) + if not newly_filled.any(): + break + + # Update. + flood_depth = torch.where(can_update, mean_depth, flood_depth) + valid_mask = torch.where(can_update, torch.ones_like(valid_mask), valid_mask) + + # Fallback: fill any remaining zero regions with the global foreground mean. + still_zero = (flood_depth == 0) & sky_mask + if still_zero.any(): + fg_mean = depth[~sky_mask].mean() + flood_depth = torch.where(still_zero, fg_mean, flood_depth) + + return flood_depth + + +def _median_boundary(depth_np: np.ndarray, sky_mask_np: np.ndarray, + depth_max: float, device, dtype, shape) -> torch.Tensor: + """ + Boundary median method. + Extract boundary depths and use their median as the uniform transition value. + """ + # Extract the boundary: non-sky pixels adjacent to sky. + kernel = np.ones((3, 3), np.uint8) + dilated_sky = cv2.dilate(sky_mask_np, kernel) + boundary_mask = (dilated_sky == 1) & (sky_mask_np == 0) + + boundary_depths = depth_np[boundary_mask] + + if len(boundary_depths) > 0: + # Filter out values that are already depth_max. + valid_depths = boundary_depths[boundary_depths < depth_max * 0.99] + if len(valid_depths) > 0: + median_depth = np.median(valid_depths) + else: + median_depth = np.median(boundary_depths) + else: + median_depth = depth_np[sky_mask_np == 0].mean() + + return torch.full(shape, median_depth, device=device, dtype=dtype) + + +def _gaussian_diffusion(depth: torch.Tensor, sky_mask: torch.Tensor, + depth_max: float, transition_width: int) -> torch.Tensor: + """ + Gaussian-blur diffusion. + Blur foreground depth first, then diffuse into sky regions. + """ + device = depth.device + dtype = depth.dtype + + # Fill sky with the foreground mean to avoid depth_max affecting the blur. + fg_mean = depth[~sky_mask].mean() + temp_depth = depth.clone() + temp_depth[sky_mask] = fg_mean + + # Gaussian blur. + blur_size = max(transition_width // 2, 3) + if blur_size % 2 == 0: + blur_size += 1 + + from torchvision.transforms.functional import gaussian_blur + blurred = gaussian_blur(temp_depth, kernel_size=[blur_size, blur_size], sigma=[blur_size / 4]) + + # Iterative diffusion using min-pooling on blurred values, after extreme values have been smoothed. + flood_depth = blurred.clone() + flood_depth[sky_mask] = float('inf') + + kernel_size = 31 + pad = kernel_size // 2 + + for _ in range(transition_width // pad + 5): + if not torch.isinf(flood_depth[sky_mask]).any(): + break + next_step = -F.max_pool2d(-flood_depth, kernel_size, stride=1, padding=pad) + is_inf = torch.isinf(flood_depth) + flood_depth = torch.where(is_inf, next_step, flood_depth) + + flood_depth = torch.where(torch.isinf(flood_depth), fg_mean, flood_depth) + + return flood_depth + + +def _restore_dim(tensor: torch.Tensor, original_dim: int) -> torch.Tensor: + if original_dim == 2: + return tensor.squeeze(0).squeeze(0) + elif original_dim == 3: + return tensor.squeeze(0) + return tensor + + +def erp_distance_ray_to_normal(distance_map, ray_directions, + smooth_sigma=0.0, + facing_camera=True): + """ + Convert an ERP distance map and ray directions to a world-coordinate normal map. + + Args: + distance_map: (H, W) Distance map. + ray_directions: (H, W, 3) Ray direction for each pixel, as unit vectors emitted from the origin. + smooth_sigma: Smoothing sigma. 0 means no smoothing. + facing_camera: True makes normals face the camera (origin); False makes them face outward. + + Returns: + normal_map: (H, W, 3) World-coordinate normals in the range [-1, 1]. + normal_rgb: (H, W, 3) RGB visualization in the range [0, 255]. + + Coordinate system (OpenCV): + X: right + Y: down + Z: forward + """ + # 0. Optional: smooth the distance map. + if smooth_sigma > 0: + distance_map = cv2.GaussianBlur( + distance_map.astype(np.float32), (0, 0), smooth_sigma + ) + + # 1. Compute 3D point coordinates: P = ray * distance. + d = distance_map.astype(np.float64) + points = ray_directions * d[..., np.newaxis] # (H, W, 3) + + # Split x, y, and z. + x = points[..., 0] + y = points[..., 1] + z = points[..., 2] + + # 2. Compute tangent vectors. + # Horizontal direction Tu uses cyclic boundaries because ERP wraps left to right. + x_right = np.roll(x, -1, axis=1) + y_right = np.roll(y, -1, axis=1) + z_right = np.roll(z, -1, axis=1) + + x_left = np.roll(x, 1, axis=1) + y_left = np.roll(y, 1, axis=1) + z_left = np.roll(z, 1, axis=1) + + Tu_x = x_right - x_left + Tu_y = y_right - y_left + Tu_z = z_right - z_left + + # Vertical direction Tv does not wrap, so boundaries need special handling. + x_down = np.roll(x, -1, axis=0) + y_down = np.roll(y, -1, axis=0) + z_down = np.roll(z, -1, axis=0) + + x_up = np.roll(x, 1, axis=0) + y_up = np.roll(y, 1, axis=0) + z_up = np.roll(z, 1, axis=0) + + Tv_x = x_down - x_up + Tv_y = y_down - y_up + Tv_z = z_down - z_up + + # Handle top and bottom boundaries near the poles. + # Top (v=0): one-sided difference. + Tv_x[0, :] = x[1, :] - x[0, :] + Tv_y[0, :] = y[1, :] - y[0, :] + Tv_z[0, :] = z[1, :] - z[0, :] + + # Bottom (v=H-1): one-sided difference. + Tv_x[-1, :] = x[-1, :] - x[-2, :] + Tv_y[-1, :] = y[-1, :] - y[-2, :] + Tv_z[-1, :] = z[-1, :] - z[-2, :] + + # 3. Compute normals with a cross product: N = Tu × Tv. + normal_x = Tu_y * Tv_z - Tu_z * Tv_y + normal_y = Tu_z * Tv_x - Tu_x * Tv_z + normal_z = Tu_x * Tv_y - Tu_y * Tv_x + + # 4. Normalize. + norm = np.sqrt(normal_x ** 2 + normal_y ** 2 + normal_z ** 2) + norm = np.where(norm > 1e-10, norm, 1e-10) + + normal_x = normal_x / norm + normal_y = normal_y / norm + normal_z = normal_z / norm + + # 5. Ensure normals face the correct direction. + if facing_camera: + # Normals should face the camera (origin). + # Check dot(normal, -ray) > 0, equivalent to dot(normal, ray) < 0. + dot_product = (normal_x * ray_directions[..., 0] + + normal_y * ray_directions[..., 1] + + normal_z * ray_directions[..., 2]) + + # If dot > 0, flip the normals. + flip_mask = dot_product > 0 + normal_x = np.where(flip_mask, -normal_x, normal_x) + normal_y = np.where(flip_mask, -normal_y, normal_y) + normal_z = np.where(flip_mask, -normal_z, normal_z) + + # 6. Stack into (H, W, 3). + normal_map = np.stack([normal_x, normal_y, normal_z], axis=-1) + + # 7. Convert to an RGB visualization. + # [-1, 1] → [0, 255] + normal_rgb = ((normal_map + 1.0) / 2.0 * 255).astype(np.uint8) + + return normal_map, normal_rgb diff --git a/nodes/comfy-env.toml b/nodes/comfy-env.toml index b8a4c9a..d011a05 100644 --- a/nodes/comfy-env.toml +++ b/nodes/comfy-env.toml @@ -11,4 +11,13 @@ safetensors = "*" einops = "*" comfy-aimdo = "*" scipy = "*" -transformers = "*" \ No newline at end of file +transformers = "*" +tqdm = "*" +opencv-python = "*" + +# utils3d: used by the vendored MoGe LSMR depth-merge solver +# (_vendor/worldgen/src/panorama_utils.py + _vendor/moge_panorama.py). +# Same pin as WorldNav/WorldStereo so behavior is identical. +[pypi-dependencies.utils3d] +git = "https://github.com/EasternJournalist/utils3d.git" +rev = "3fab839f0be9931dac7c8488eb0e1600c236e183" \ No newline at end of file diff --git a/nodes/depth_merge.py b/nodes/depth_merge.py new file mode 100644 index 0000000..db788b7 --- /dev/null +++ b/nodes/depth_merge.py @@ -0,0 +1,275 @@ +"""SharpDepthMerge node for ComfyUI-Sharp. + +Per-face rectilinear depth maps -> a single equirectangular distance map, +via MoGe's sparse-LSMR formulation in log-distance space (Laplacian + +gradient terms). 1:1 schema mirror of WorldNavDepthMerge — wiring is +identical so workflows port over by editing only the node-class name. + +Backed by the same solver as WorldNav: vendored at +`nodes/_vendor/moge_panorama.py` (CPU LSMR via scipy) + vendored at +`nodes/_vendor/worldgen/src/panorama_utils.py` (GPU LSMR via +solve_lsmr_gpu + post-merge sky inpaint + south-pole smooth). +""" + +from __future__ import annotations + +import logging +import sys + +import numpy as np +import torch +from comfy_api.latest import io + +log = logging.getLogger("sharp") + + +def _p(msg: str) -> None: + print(f"[SharpDepthMerge] {msg}", file=sys.stderr, flush=True) + + +class SharpDepthMerge(io.ComfyNode): + """Per-face depth maps + per-face extrinsics/intrinsics -> equirect depth IMAGE. + + Schema is a 1:1 mirror of WorldNavDepthMerge so workflows port over + without rewiring. Internally uses the same vendored MoGe LSMR solver. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SharpDepthMerge", + display_name="Sharp Depth Merge", + category="SHARP", + description=( + "Stitch per-face rectilinear depth maps from a panorama " + "split (SharpPanoramaIcosahedronSplit etc.) back into a " + "single equirect distance map. Sparse-LSMR formulation in " + "log-distance space (Laplacian + gradient terms in MoGe " + "style). 1:1 schema mirror of WorldNavDepthMerge." + ), + inputs=[ + io.Image.Input( + "face_depths", + tooltip="Per-face depth maps from a rectilinear depth " + "model (MoGe2 / SHARP / etc.). Shape " + "(N, h, w, C); only the first channel is used " + "as the distance value."), + io.Mask.Input( + "face_valid_masks", + optional=True, + tooltip="Per-face valid masks (typically the depth " + "model's valid_mask output). Shape (N, h, w) " + "bool/float. Load-bearing: upstream " + "merge_panorama_depth uses this to gate which " + "pixels enter the gradient + laplacian LSMR " + "system. Without it, sky / invalid pixels " + "corrupt the solve. If unwired, all pixels " + "are treated as valid (degraded fallback)."), + io.Custom("EXTRINSICS").Input( + "extrinsics", + tooltip="From the panorama-split node. (N, 4, 4) " + "world-to-camera per face."), + io.Custom("INTRINSICS").Input( + "intrinsics", + tooltip="From the panorama-split node. (N, 3, 3) " + "pixel-space per face."), + io.Int.Input( + "out_width", default=1920, min=512, max=4096, step=64, + tooltip="Output equirect width. 2:1 aspect (height = " + "width / 2 is the convention; provide " + "out_height to override)."), + io.Int.Input( + "out_height", default=960, min=256, max=2048, step=64, + tooltip="Output equirect height. Should be width/2 " + "for 2:1 aspect."), + io.Boolean.Input( + "use_gpu", default=True, + tooltip="DEPRECATED. Kept for back-compat. Equivalent " + "to strategy='cpu_fallback' when False. " + "Otherwise respects strategy."), + io.Combo.Input( + "strategy", + options=["auto", "full", "vcycle", "cpu_fallback"], + default="auto", + tooltip=( + "Solver strategy for the equirect depth output:\n" + " auto: full if out<=1920x960 else vcycle (recommended).\n" + " full: solve LSMR at out_resolution (original). OOMs above ~2 Mpix on 24 GB.\n" + " vcycle: solve LSMR at 1920x960 then bilinear-upsample depth+mask " + "to out_resolution. ~4x faster + fits memory at 3840x1920. Per-face " + "info saturates at 1920x960 so quality is indistinguishable.\n" + " cpu_fallback: scipy LSMR on CPU end-to-end. Slow (~30-60s at " + "1920x960) but works at any resolution." + )), + ], + outputs=[ + io.Image.Output(display_name="depth"), + io.Mask.Output(display_name="valid_mask"), + ], + ) + + @classmethod + def execute(cls, face_depths, extrinsics, intrinsics, + face_valid_masks=None, out_width=1920, out_height=960, + use_gpu=True, strategy="auto"): + import cv2 + # Vendored MoGe LSMR solvers — imported here so depth_merge.py + # loads cleanly even if scipy / utils3d / etc. aren't installed. + # (They ARE in Sharp's pixi env per comfy-env.toml; this lazy + # import keeps the node-class registration unconditional.) + from ._vendor.moge_panorama import merge_panorama_depth + from ._vendor.worldgen.src.panorama_utils import ( + merge_panorama_depth_gpu, + smooth_south_pole_depth, + ) + + # --- face_depths: (N, h, w, C) → list of (h, w) float --- + d = (face_depths.detach().cpu().numpy() + if isinstance(face_depths, torch.Tensor) else np.asarray(face_depths)) + if d.ndim != 4: + raise ValueError( + f"SharpDepthMerge: face_depths must be (N, h, w, C), got {d.shape}" + ) + N, fh, fw, _ = d.shape + distance_maps = [d[i, ..., 0].astype(np.float32) for i in range(N)] + + # --- extrinsics / intrinsics: (N, 4, 4) and (N, 3, 3) --- + ex = (extrinsics.detach().cpu().numpy() + if isinstance(extrinsics, torch.Tensor) else np.asarray(extrinsics)) + intr = (intrinsics.detach().cpu().numpy() + if isinstance(intrinsics, torch.Tensor) else np.asarray(intrinsics)) + if ex.shape != (N, 4, 4) or intr.shape != (N, 3, 3): + raise ValueError( + f"SharpDepthMerge: shape mismatch — extrinsics {ex.shape}, " + f"intrinsics {intr.shape}, expected ({N},4,4) and ({N},3,3)" + ) + extr_list = [ex[i].astype(np.float32) for i in range(N)] + intr_list = [intr[i].astype(np.float32) for i in range(N)] + + # --- valid_masks: (N, h, w, C) or None → list of (h, w) bool --- + if face_valid_masks is None: + pred_masks = [np.ones((fh, fw), dtype=bool) for _ in range(N)] + else: + v = (face_valid_masks.detach().cpu().numpy() + if isinstance(face_valid_masks, torch.Tensor) + else np.asarray(face_valid_masks)) + if v.ndim == 4: + v = v[..., 0] + if v.shape != (N, fh, fw): + raise ValueError( + f"SharpDepthMerge: face_valid_masks shape {v.shape} " + f"doesn't match face_depths ({N}, {fh}, {fw})" + ) + pred_masks = [(v[i] > 0.5).astype(bool) for i in range(N)] + + _p(f"merging {N} faces @ {fh}x{fw} -> equirect {out_height}x{out_width} " + f"(use_gpu={use_gpu})") + + # Progress bar — visible in ComfyUI if available. + try: + import comfy.utils + pbar = comfy.utils.ProgressBar(100) + except Exception: + class _Noop: + def update_absolute(self, *_a, **_k): + pass + pbar = _Noop() + pbar.update_absolute(1, 100) + + # Resolve strategy. use_gpu=False overrides → cpu_fallback (back-compat + # for the equivalent toggle in WorldNavDepthMerge). + out_w, out_h = int(out_width), int(out_height) + resolved = strategy + if not use_gpu: + resolved = "cpu_fallback" + if resolved == "auto": + resolved = "vcycle" if (out_w * out_h > 1920 * 960) else "full" + _p(f"resolved strategy={resolved} (requested={strategy}, " + f"use_gpu={use_gpu}, out={out_w}x{out_h})") + + if resolved == "full": + # Recursive LSMR all the way up to out_resolution. Reference; + # OOMs above ~2 Mpix on 24 GB cards. + depth_np, mask_np = merge_panorama_depth_gpu( + out_w, out_h, + distance_maps, pred_masks, + extr_list, intr_list, + ) + elif resolved == "vcycle": + # V-cycle: solve the recursive pyramid up to 1920x960 (the + # information-content ceiling for typical 42x ~512² per-face + # data), then bilinear-upsample depth + nearest-upsample mask + # to out_resolution. Skips the OOM-prone outermost LSMR. + depth_low, mask_low = merge_panorama_depth_gpu( + 1920, 960, + distance_maps, pred_masks, + extr_list, intr_list, + ) + if (out_w, out_h) == (1920, 960): + depth_np, mask_np = depth_low, mask_low + else: + depth_np = cv2.resize( + depth_low.astype(np.float32), + (out_w, out_h), + interpolation=cv2.INTER_LINEAR, + ) + mask_np = cv2.resize( + mask_low.astype(np.uint8), + (out_w, out_h), + interpolation=cv2.INTER_NEAREST, + ).astype(bool) + _p(f"vcycle: upsampled 1920x960 -> {out_w}x{out_h} " + f"(bilinear depth, nearest mask)") + elif resolved == "cpu_fallback": + # scipy LSMR end-to-end on CPU. Slow but no VRAM ceiling. + depth_np, mask_np = merge_panorama_depth( + out_w, out_h, + distance_maps, pred_masks, + extr_list, intr_list, + pbar=pbar, + ) + else: + raise ValueError(f"SharpDepthMerge: unknown strategy {resolved!r}") + + # Match upstream HY-World 2.0 post-merge processing + # (panorama_utils.py:572-580). Two steps: + # 1. Inpaint invalid pixels with the 99.9th-percentile (sky) + # depth instead of zero. Zero-depth puts those pixels right + # on the camera origin and breaks downstream point-cloud / + # mesh geometry. + # 2. Smooth the south-pole strip (bottom 5%) to fix left-right + # seams the LSMR solver can't enforce due to the equirect + # parameterization stretching infinitely at the pole. + depth_np = depth_np.astype(np.float32) + valid = mask_np.astype(bool) + if valid.any() and (~valid).any(): + sky_depth = float(np.nanquantile(depth_np[valid], 0.999)) + depth_np[~valid] = sky_depth + _p(f" inpainted {int((~valid).sum())} invalid pixels with " + f"sky_depth={sky_depth:.3f}") + depth_np = smooth_south_pole_depth(depth_np, smooth_height_ratio=0.05) + # Belt-and-braces: trap any residual non-finite values that escaped. + depth_np = np.nan_to_num(depth_np, nan=0.0, posinf=0.0, neginf=0.0) + _p(f"merged depth: shape {depth_np.shape}, min={float(depth_np.min()):.3f}, " + f"median={float(np.median(depth_np)):.3f}, max={float(depth_np.max()):.3f}; " + f"valid pixels: {int(mask_np.sum())} / {mask_np.size}") + + # IMAGE convention: (B, H, W, C). Broadcast depth across 3 channels so + # it composes with regular depth-viz nodes. + depth_img = ( + torch.from_numpy(depth_np) + .unsqueeze(-1).expand(-1, -1, 3).unsqueeze(0).contiguous() + ) + valid_mask = torch.from_numpy(mask_np.astype(np.float32)).unsqueeze(0) + + pbar.update_absolute(100, 100) + return io.NodeOutput(depth_img, valid_mask) + + +NODE_CLASS_MAPPINGS = { + "SharpDepthMerge": SharpDepthMerge, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SharpDepthMerge": "Sharp Depth Merge", +} diff --git a/nodes/panorama_cube_split.py b/nodes/panorama_cube_split.py new file mode 100644 index 0000000..2f3e5aa --- /dev/null +++ b/nodes/panorama_cube_split.py @@ -0,0 +1,306 @@ +"""SharpPanoramaCubeSplit node for ComfyUI-Sharp. + +Equirect panorama -> N perspective crops at a fixed pitch x yaw grid. + +Parallel to WorldStereo's WorldStereoPanoramaCubeSplit: + 3 pitch levels x N yaw steps at configurable rectangular FOV / aspect. + Default: 3 x 9 = 27 views at 832x480, 120 deg horizontal x 90 deg + vertical FOV. + +Different from Sharp's existing SamplePanorama (sample_panorama.py): + - SamplePanorama: square output, FOV-derived dense grid, blends to + cover the full sphere — flexible for arbitrary FOV/overlap. + - SharpPanoramaCubeSplit: fixed 3-pitch grid, rectangular aspect, + matches WorldStereo's anchor-view geometry — for feeding a + rectilinear depth model (MoGe2 etc.) before merging back to equirect + via SharpDepthMerge. + +Output convention matches the rest of ComfyUI-Sharp: + face_images (IMAGE [N, H, W, 3]) + extrinsics (EXTRINSICS [N, 4, 4]) + + intrinsics (INTRINSICS [N, 3, 3]). +""" + +import logging +import math + +import numpy as np +import torch +import torch.nn.functional as F + +from comfy_api.latest import io + +log = logging.getLogger("sharp") + + +def _rot_around_z(deg: float) -> torch.Tensor: + """3x3 rotation matrix around the Z axis. World up = +Z convention. + + Matches WorldStereo's `rotate_around_z_axis` in panorama_utils.py so + the pitch-up / pitch-down hemispheres land at the same world points. + """ + rad = math.radians(deg) + c, s = math.cos(rad), math.sin(rad) + return torch.tensor([ + [c, -s, 0], + [s, c, 0], + [0, 0, 1], + ], dtype=torch.float32) + + +def _look_at_w2c(eye: torch.Tensor, target: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """world-to-camera 4x4 matrix from a look-at triple (numpy-free). + + Convention: camera looks down +Z in camera space (matches utils3d + extrinsics_look_at). Returns a w2c matrix. + """ + # forward (in world): from eye to target, normalized + fwd = (target - eye) + fwd = fwd / (fwd.norm() + 1e-12) + # right = forward x up + right = torch.linalg.cross(fwd, up) + right = right / (right.norm() + 1e-12) + # recomputed up = right x forward + new_up = torch.linalg.cross(right, fwd) + # Camera-to-world rotation: columns are the camera basis in world coords. + # We want world-to-camera, which is the transpose for pure rotation. + # In camera space: +X right, +Y down (image convention), +Z forward. + # So R_c2w has columns = [right, -new_up, fwd]; w2c = transpose. + R_c2w = torch.stack([right, -new_up, fwd], dim=1) # [3, 3] + R_w2c = R_c2w.T + + ext = torch.eye(4, dtype=torch.float32) + ext[:3, :3] = R_w2c + ext[:3, 3] = -R_w2c @ eye + return ext + + +def _intrinsics_from_fov(fov_x_rad: float, fov_y_rad: float, w: int, h: int) -> torch.Tensor: + """Pixel-space 3x3 intrinsics from horizontal + vertical FOV.""" + fx = (w / 2.0) / math.tan(fov_x_rad / 2.0) + fy = (h / 2.0) / math.tan(fov_y_rad / 2.0) + cx = (w - 1) / 2.0 + cy = (h - 1) / 2.0 + K = torch.tensor([ + [fx, 0, cx], + [0, fy, cy], + [0, 0, 1.0], + ], dtype=torch.float32) + return K + + +def _sample_perspective_from_equirect( + panorama: torch.Tensor, + ext_w2c: torch.Tensor, + K_pixel: torch.Tensor, + out_h: int, + out_w: int, +) -> torch.Tensor: + """Backward warp: equirect panorama -> a single perspective view. + + Mirrors sample_panorama.py's geometry but with rectangular output + and an explicit extrinsics/intrinsics input (so the same look-at + targets used to build ext_w2c here can be passed downstream + unchanged). + """ + H_pano, W_pano, _ = panorama.shape + device = panorama.device + + # Pixel grid in the output perspective image. + u = torch.arange(out_w, dtype=torch.float32, device=device) + v = torch.arange(out_h, dtype=torch.float32, device=device) + uu, vv = torch.meshgrid(u, v, indexing='xy') + + fx, fy = K_pixel[0, 0].item(), K_pixel[1, 1].item() + cx, cy = K_pixel[0, 2].item(), K_pixel[1, 2].item() + + # Camera-space rays for each output pixel. + dx = (uu - cx) / fx + dy = (vv - cy) / fy + dz = torch.ones_like(dx) + rays_cam = torch.stack([dx, dy, dz], dim=-1) # [H, W, 3] + rays_cam = F.normalize(rays_cam, dim=-1) + + # World rays via camera-to-world rotation. + R_w2c = ext_w2c[:3, :3].to(device) + R_c2w = R_w2c.T + rays_world = torch.einsum('ij,hwj->hwi', R_c2w, rays_cam) # [H, W, 3] + + # World rays -> equirect sample coords. + # WorldStereo / HY-World convention: world up = +Z, panorama columns + # = yaw around Z. Match that here so this node's extrinsics line up + # with WorldStereo's pano_bank geometry. + rx, ry, rz = rays_world[..., 0], rays_world[..., 1], rays_world[..., 2] + yaw = torch.atan2(ry, rx) # [-pi, pi] + pitch = torch.asin(torch.clamp(rz, -1.0, 1.0)) # [-pi/2, pi/2] + + eq_x = (yaw / math.pi + 1.0) * (W_pano - 1) / 2.0 # 0..W-1 + eq_y = (0.5 - pitch / math.pi) * (H_pano - 1) # 0..H-1 + grid_x = eq_x / (W_pano - 1) * 2.0 - 1.0 + grid_y = eq_y / (H_pano - 1) * 2.0 - 1.0 + grid = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0) # [1, H, W, 2] + + pano_nchw = panorama.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W] + sampled = F.grid_sample( + pano_nchw, grid, + mode='bilinear', padding_mode='border', align_corners=True, + ) + return sampled[0].permute(1, 2, 0) # [H, W, C] + + +class SharpPanoramaCubeSplit(io.ComfyNode): + """Equirect panorama -> N perspective crops at a fixed pitch x yaw grid.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SharpPanoramaCubeSplit", + display_name="Sharp Panorama Cube Split", + category="SHARP", + description=( + "Sample N perspective crops from a 2:1 equirectangular panorama " + "at a fixed pitch x yaw grid (3 pitch x 9 yaw = 27 views by " + "default). Matches WorldStereoPanoramaCubeSplit's geometry — " + "useful as the input stage for a per-face depth model whose " + "outputs will be stitched back via SharpDepthMerge.\n\n" + "Different from Sharp's SamplePanorama (which derives the grid " + "from FOV+overlap, square output): this node has a fixed pitch " + "grid with rectangular aspect (default 832x480), and matches " + "the camera/intrinsics convention WorldStereo's pano_bank uses." + ), + inputs=[ + io.Image.Input( + "panorama", + tooltip="Equirectangular RGB panorama, 2:1 aspect. " + "[B, H, W, 3] or [H, W, 3] in [0, 1]. First image " + "of batch is used."), + io.Int.Input( + "image_w", default=832, min=224, max=2048, step=16, + tooltip="Crop width. 832 matches WorldStereo's anchor width."), + io.Int.Input( + "image_h", default=480, min=224, max=2048, step=16, + tooltip="Crop height. 480 matches WorldStereo's anchor height."), + io.Float.Input( + "fov_x_deg", default=120.0, min=30.0, max=170.0, step=1.0, + tooltip="Horizontal FOV. 120 matches WorldStereo upstream."), + io.Float.Input( + "fov_y_deg", default=90.0, min=30.0, max=170.0, step=1.0, + tooltip="Vertical FOV. 90 matches WorldStereo upstream."), + io.Float.Input( + "rot_deg", default=40.0, min=10.0, max=120.0, step=5.0, + tooltip="Yaw step between adjacent views. 40 -> 9 views " + "per pitch -> 27 total at 3 pitches."), + io.Float.Input( + "pitch_up", default=0.5, min=0.0, max=1.0, step=0.05, + tooltip="Z component of the upper-hemisphere look-at " + "targets. 0.5 matches WorldStereo upstream."), + io.Float.Input( + "pitch_down", default=-0.5, min=-1.0, max=0.0, step=0.05, + tooltip="Z component of the lower-hemisphere look-at " + "targets. -0.5 matches WorldStereo upstream."), + ], + outputs=[ + io.Image.Output( + display_name="face_images", + tooltip="N perspective crops [N, H, W, 3] in [0, 1]. " + "Wire to a depth model (SharpPredictDepth / MoGe2 " + "/ etc.)."), + io.Custom("EXTRINSICS").Output( + display_name="extrinsics", + tooltip="[N, 4, 4] world-to-camera matrices per face. " + "Camera at world origin, look-at along the pitch/" + "yaw direction. Pass to SharpDepthMerge."), + io.Custom("INTRINSICS").Output( + display_name="intrinsics", + tooltip="[N, 3, 3] pixel-space intrinsics per face. " + "Identical across faces (same FOV + size). Pass " + "to SharpDepthMerge."), + io.Float.Output( + display_name="fov_x_deg", + tooltip="Pass-through. Wire into MoGe2 / depth model's " + "fov_x widget so it doesn't have to estimate."), + io.Int.Output( + display_name="num_faces", + tooltip="N = 3 * (360 / rot_deg). 27 at defaults."), + ], + ) + + @classmethod + def execute( + cls, panorama: torch.Tensor, + image_w: int = 832, image_h: int = 480, + fov_x_deg: float = 120.0, fov_y_deg: float = 90.0, + rot_deg: float = 40.0, + pitch_up: float = 0.5, pitch_down: float = -0.5, + ): + # Normalize input shape -> [H, W, 3], float in [0, 1]. + if panorama.dim() == 4: + panorama = panorama[0] + if panorama.dtype != torch.float32: + panorama = panorama.float() + if panorama.shape[-1] == 4: + panorama = panorama[..., :3] + # Sharp convention is float [0,1]; guard against accidental uint8 input. + if panorama.max() > 2.0: + panorama = panorama / 255.0 + + device = panorama.device + H_pano, W_pano, _ = panorama.shape + + # Build the (3 pitch x N yaw) direction grid. + # WorldStereo convention: world up = +Z, base directions look at +X + # rotated around Z. (Eye is world origin; targets are unit-ish + # points on a sphere.) + # NOTE: WorldStereo uses [-1, 0, 0] as the base "forward" with +Z up; + # we mirror that exactly so panorama columns line up the same way. + base_targets = [ + torch.tensor([-1.0, 0.0, 0.0], dtype=torch.float32), + torch.tensor([-1.0, 0.0, float(pitch_up)], dtype=torch.float32), + torch.tensor([-1.0, 0.0, float(pitch_down)], dtype=torch.float32), + ] + N_yaw = max(1, int(round(360.0 / max(1.0, float(rot_deg))))) + targets = [] + for base in base_targets: + targets.append(base) + for i in range(1, N_yaw): + R = _rot_around_z(rot_deg * i) + targets.append((R @ base).contiguous()) + targets = torch.stack(targets, dim=0) # [N, 3] + N = targets.shape[0] + + # Same intrinsics for every face (same FOV + same output size). + fov_x_rad = math.radians(fov_x_deg) + fov_y_rad = math.radians(fov_y_deg) + K = _intrinsics_from_fov(fov_x_rad, fov_y_rad, int(image_w), int(image_h)) + + # Sample each face. + eye = torch.zeros(3, dtype=torch.float32) + up = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32) + face_images = [] + extrinsics = [] + for i in range(N): + ext_w2c = _look_at_w2c(eye, targets[i], up) + img = _sample_perspective_from_equirect( + panorama, ext_w2c, K, int(image_h), int(image_w), + ) + face_images.append(img) + extrinsics.append(ext_w2c) + + face_images_t = torch.stack(face_images, dim=0).contiguous() # [N, H, W, 3] + ext_t = torch.stack(extrinsics, dim=0).contiguous() # [N, 4, 4] + intr_t = K.unsqueeze(0).expand(N, -1, -1).contiguous() # [N, 3, 3] + + log.info( + f"[SharpPanoramaCubeSplit] {N} faces ({N_yaw}/pitch x 3 pitches) " + f"@ {image_w}x{image_h}, fov=({fov_x_deg:.0f},{fov_y_deg:.0f})" + ) + + return io.NodeOutput(face_images_t, ext_t, intr_t, float(fov_x_deg), int(N)) + + +NODE_CLASS_MAPPINGS = { + "SharpPanoramaCubeSplit": SharpPanoramaCubeSplit, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SharpPanoramaCubeSplit": "Sharp Panorama Cube Split", +} diff --git a/nodes/panorama_icosahedron_split.py b/nodes/panorama_icosahedron_split.py new file mode 100644 index 0000000..aa1f3a9 --- /dev/null +++ b/nodes/panorama_icosahedron_split.py @@ -0,0 +1,296 @@ +"""SharpPanoramaIcosahedronSplit node for ComfyUI-Sharp. + +Equirect panorama -> N square perspective crops at icosahedron vertices, +90° FOV each. Schema mirrors WorldNavPanoramaSplit 1:1 (modulo +`panorama` being plain IMAGE here instead of WorldNav's +WORLDSTEREO_PANORAMA custom type) so workflows port over by editing +only the node-class name. + +Geometry: + - icosahedron_12: 12 base vertices via golden-ratio coords + - icosahedron_42: subdivide once (add edge midpoints, normalize to + sphere) -> 12 + 30 = 42 vertices + +Why square / icosahedron for SHARP: SHARP's internal resolution is +1536x1536 square and was trained on cube-face-like inputs. The +icosahedron tiling gives uniform sphere coverage (no pole over/under- +sampling like a pitch/yaw grid) at 90° FOV which matches cube-face +geometry exactly. + +Implementation is pure numpy + torch. Reuses three helpers from +`panorama_cube_split.py` (`_look_at_w2c`, `_intrinsics_from_fov`, +`_sample_perspective_from_equirect`) so no math is duplicated. + +Pole-singularity fix (mirrors WorldNav's): two vertices of the +subdivided icosahedron land at (0, 0, ±1). With the default world-up +of +Z, `forward × up = 0` for those, producing NaN extrinsics. We +detect `|forward · up| > 0.999` per vertex and swap to up=+Y just for +those. +""" + +from __future__ import annotations + +import logging +import math + +import numpy as np +import torch +import torch.nn.functional as F + +from comfy_api.latest import io + +# Re-use the math helpers from the sibling cube-split node — no copy. +from .panorama_cube_split import ( + _look_at_w2c, + _intrinsics_from_fov, + _sample_perspective_from_equirect, +) + +log = logging.getLogger("sharp") + + +def _icosahedron_vertices() -> np.ndarray: + """12 unit-sphere vertices of a regular icosahedron. + + Standard golden-ratio coords. After normalizing to the unit sphere, + edge length squared between adjacent vertices = (2 / phi) ≈ 1.236. + """ + phi = (1.0 + math.sqrt(5.0)) / 2.0 + verts = np.array( + [ + ( 0, +1, +phi), ( 0, +1, -phi), ( 0, -1, +phi), ( 0, -1, -phi), + (+1, +phi, 0), (+1, -phi, 0), (-1, +phi, 0), (-1, -phi, 0), + (+phi, 0, +1), (+phi, 0, -1), (-phi, 0, +1), (-phi, 0, -1), + ], + dtype=np.float32, + ) + verts /= np.linalg.norm(verts, axis=1, keepdims=True) + return verts + + +def _icosahedron_edges(verts: np.ndarray) -> np.ndarray: + """30 edge (i, j) pairs of the regular icosahedron. + + Derived by picking the 30 shortest pairwise distances out of C(12,2)=66 + — equivalently, the pairs with squared distance ≈ 2/phi after the + unit-sphere normalization. No hardcoded indices, so this stays correct + regardless of the vertex ordering in `_icosahedron_vertices`. + """ + N = len(verts) + pairs = [] + for i in range(N): + for j in range(i + 1, N): + d2 = float(((verts[i] - verts[j]) ** 2).sum()) + pairs.append((d2, i, j)) + pairs.sort() + # The icosahedron has 30 edges; threshold by taking exactly the 30 + # closest pairs. Sanity-check: the gap between edge-distance and + # non-edge-distance should be large (>1.5x). + edge_d2 = pairs[29][0] + non_edge_d2 = pairs[30][0] + assert non_edge_d2 > edge_d2 * 1.5, ( + f"icosahedron edge classification ambiguous: " + f"edge_d2={edge_d2:.4f}, non_edge_d2={non_edge_d2:.4f}" + ) + return np.array([(i, j) for _, i, j in pairs[:30]], dtype=np.int32) + + +def _subdivide_icosahedron(verts: np.ndarray, edges: np.ndarray) -> np.ndarray: + """Add edge midpoints (normalized to the unit sphere) → 42 vertices.""" + midpoints = (verts[edges[:, 0]] + verts[edges[:, 1]]) / 2.0 + midpoints /= np.linalg.norm(midpoints, axis=1, keepdims=True) + return np.concatenate([verts, midpoints.astype(np.float32)], axis=0) + + +def _icosahedron_split_vertices(subdivision: str) -> np.ndarray: + """Return [N, 3] unit-sphere target points for the chosen subdivision.""" + base = _icosahedron_vertices() + if subdivision == "icosahedron_12": + return base + if subdivision == "icosahedron_42": + edges = _icosahedron_edges(base) + return _subdivide_icosahedron(base, edges) + raise ValueError(f"unknown subdivision={subdivision!r}") + + +def _build_extrinsics_with_pole_fix(targets: torch.Tensor) -> torch.Tensor: + """Build per-vertex world-to-camera matrices, swapping up at poles. + + `targets`: [N, 3] unit-sphere vectors (look-at targets from world + origin). World up default = +Z. For vertices within ~2.5° of ±Z, + swap up to +Y to avoid the cross-product singularity. + """ + N = targets.shape[0] + eye = torch.zeros(3, dtype=torch.float32) + up_default = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32) + up_fallback = torch.tensor([0.0, 1.0, 0.0], dtype=torch.float32) + cos_thresh = 0.999 + + out = torch.empty(N, 4, 4, dtype=torch.float32) + cos_with_up = torch.abs(targets @ up_default) + for i in range(N): + up = up_fallback if float(cos_with_up[i]) > cos_thresh else up_default + out[i] = _look_at_w2c(eye, targets[i], up) + return out + + +def _make_debug_grid(face_images: torch.Tensor, tile_size: int = 96) -> torch.Tensor: + """Tile N face crops into a single [H, W, 3] preview image. + + Output is a roughly-square grid (cols = ceil(sqrt(N)) by default). + Each tile is `tile_size`x`tile_size`, downsampled via bilinear. Used + to satisfy WorldNavPanoramaSplit's `debug_image` output convention. + """ + N, H, W, _ = face_images.shape + cols = int(math.ceil(math.sqrt(N))) + rows = int(math.ceil(N / cols)) + + # Bilinear-resize each face to (tile_size, tile_size). + nchw = face_images.permute(0, 3, 1, 2) # [N, 3, H, W] + tiled = F.interpolate( + nchw, size=(tile_size, tile_size), mode="bilinear", align_corners=False, + ) # [N, 3, T, T] + tiled = tiled.permute(0, 2, 3, 1) # [N, T, T, 3] + + grid = torch.zeros(rows * tile_size, cols * tile_size, 3, dtype=face_images.dtype) + for i in range(N): + r, c = divmod(i, cols) + grid[ + r * tile_size:(r + 1) * tile_size, + c * tile_size:(c + 1) * tile_size, + :, + ] = tiled[i] + return grid + + +class SharpPanoramaIcosahedronSplit(io.ComfyNode): + """Equirect panorama -> N square perspective crops at icosahedron vertices.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SharpPanoramaIcosahedronSplit", + display_name="Sharp Panorama Icosahedron Split", + category="SHARP", + description=( + "Split an equirectangular panorama into N rectilinear views " + "(90° fov each) at icosahedron vertices. Pipe the views " + "through any rectilinear depth model (e.g. SHARP / MoGe2), " + "then merge the per-view depths back into an equirect depth " + "map via SharpDepthMerge.\n\n" + "1:1 schema mirror of WorldNavPanoramaSplit (modulo " + "panorama input being plain IMAGE here, not WorldStereo's " + "WORLDSTEREO_PANORAMA custom type)." + ), + inputs=[ + io.Image.Input( + "panorama", + tooltip="Equirectangular RGB panorama (2:1). [B, H, W, 3] " + "or [H, W, 3] in [0, 1]."), + io.Int.Input( + "resolution", default=512, min=128, max=2048, step=64, + tooltip="Per-face image resolution (square). Default 512 " + "matches WorldNavPanoramaSplit."), + io.Combo.Input( + "subdivision", + options=["icosahedron_12", "icosahedron_42"], + default="icosahedron_42", + tooltip="Tiling density. 12 = base icosahedron (faster), " + "42 = subdivided (better polar coverage)."), + io.Boolean.Input( + "use_gpu", default=True, optional=True, + tooltip="Batched panorama->face resampling via " + "torch.nn.functional.grid_sample (bilinear) on " + "CUDA. Falls back to CPU if CUDA isn't " + "available. Math identical within fp32 round-off."), + ], + outputs=[ + io.Image.Output(display_name="face_images"), + io.Custom("EXTRINSICS").Output(display_name="extrinsics"), + io.Custom("INTRINSICS").Output(display_name="intrinsics"), + io.Float.Output(display_name="fov_x_deg"), + io.Image.Output(display_name="debug_image"), + ], + ) + + @classmethod + def execute( + cls, panorama: torch.Tensor, + resolution: int = 512, + subdivision: str = "icosahedron_42", + use_gpu: bool = True, + ): + # Normalize input shape -> [H, W, 3], float in [0, 1]. + pano = panorama + if pano.dim() == 4: + pano = pano[0] + if pano.dtype != torch.float32: + pano = pano.float() + if pano.shape[-1] == 4: + pano = pano[..., :3] + if pano.max() > 2.0: + pano = pano / 255.0 + + H_pano, W_pano, _ = pano.shape + # Hard 2:1 aspect check — same as WorldNavPanoramaSplit. + aspect = W_pano / H_pano + if abs(aspect - 2.0) > 0.01: + raise ValueError( + f"SharpPanoramaIcosahedronSplit: panorama must be 2:1 " + f"equirect, got {W_pano}x{H_pano} (ratio {aspect:.3f}). " + f"Crop / pad to 2:1 first." + ) + + device = pano.device + if use_gpu and torch.cuda.is_available() and device.type != "cuda": + pano = pano.cuda() + device = pano.device + + targets_np = _icosahedron_split_vertices(subdivision) + targets = torch.from_numpy(targets_np).to(device=device, dtype=torch.float32) + N = targets.shape[0] + + # Per-vertex extrinsics with pole-singularity fix. + extrinsics = _build_extrinsics_with_pole_fix(targets).to(device) + + # Same intrinsics for every face (90° FOV, square). + fov_x_deg = 90.0 + fov_y_deg = 90.0 + K = _intrinsics_from_fov( + math.radians(fov_x_deg), math.radians(fov_y_deg), + int(resolution), int(resolution), + ).to(device) + + face_images = [] + for i in range(N): + img = _sample_perspective_from_equirect( + pano, extrinsics[i], K, int(resolution), int(resolution), + ) + face_images.append(img) + + face_images_t = torch.stack(face_images, dim=0).contiguous() # [N, R, R, 3] + intrinsics_t = K.unsqueeze(0).expand(N, -1, -1).contiguous() # [N, 3, 3] + extrinsics_t = extrinsics.contiguous() # [N, 4, 4] + + # Debug grid — small thumbnails so the preview node renders fast. + debug_grid = _make_debug_grid(face_images_t.cpu(), tile_size=96) + debug_image = debug_grid.unsqueeze(0) # [1, H, W, 3] + + log.info( + f"[SharpPanoramaIcosahedronSplit] {N} faces " + f"({subdivision}) @ {resolution}x{resolution}, fov=90°" + ) + + return io.NodeOutput( + face_images_t, extrinsics_t, intrinsics_t, + float(fov_x_deg), debug_image, + ) + + +NODE_CLASS_MAPPINGS = { + "SharpPanoramaIcosahedronSplit": SharpPanoramaIcosahedronSplit, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SharpPanoramaIcosahedronSplit": "Sharp Panorama Icosahedron Split", +} diff --git a/nodes/predict_gaussian_attrs.py b/nodes/predict_gaussian_attrs.py new file mode 100644 index 0000000..a767dc0 --- /dev/null +++ b/nodes/predict_gaussian_attrs.py @@ -0,0 +1,394 @@ +"""SharpPredictGaussianAttrs node for ComfyUI-Sharp. + +Runs SHARP inference and exposes the per-pixel gaussian decoder output as +structured per-layer tensors instead of flattening to a PLY. + +Why: SHARP outputs `num_layers = 2` gaussians per stride-pixel (layer 0 = +visible/near surface, layer 1 = back/occluded surface). The gaussian +composer's pre-flatten shape is: + + mean_vectors [B, 3, num_layers, H_grid, W_grid] + singular_values [B, 3, num_layers, H_grid, W_grid] + quaternions [B, 4, num_layers, H_grid, W_grid] + colors [B, 3, num_layers, H_grid, W_grid] + opacities [B, num_layers, H_grid, W_grid] + +Where H_grid = W_grid = `predictor.output_resolution` (typically 768 when +SHARP internal_resolution = 1536). The `gaussian_composer` then flattens +to [B, num_layers*H*W, C] for save_ply; the post-flatten order is +contiguous in (layer, h, w) — layer 0 first, layer 1 second. + +This node un-flattens that, separates layers, and emits: + + first_layer_depth_refined IMAGE [B, H, W, 3] -- ndc-z of layer 0 + first_layer_gaussian_attrs MULTIBAND_IMAGE [B, 14, H, W] + second_layer_depth_refined IMAGE [B, H, W, 3] -- ndc-z of layer 1 + second_layer_gaussian_attrs MULTIBAND_IMAGE [B, 14, H, W] + +The "refined depth" is the gaussian's ndc-z. Since +`ndc_matrix[2, 2] = 1.0` in `unproject_gaussians`, ndc-z equals +camera-space z. That's the post-decoder refined depth — tighter than +the disparity head's pre-refinement output exposed by SharpPredictDepth. + +Multiband channel order (per layer): + + 0: position_x (ndc, in [-1, 1]) + 1: position_y (ndc, in [-1, 1]) + 2: position_z (ndc / camera z, same as depth_refined output; + duplicated here so multiband is self-contained) + 3: scale_x (Σ-singular value) + 4: scale_y + 5: scale_z + 6: quaternion_w + 7: quaternion_x + 8: quaternion_y + 9: quaternion_z + 10: color_r (linear or SH-DC; see gaussians.py) + 11: color_g + 12: color_b + 13: opacity (post-activation, in [0, 1]) + +Downstream: feed `first_layer_depth_refined` (across N faces) into +AlignDepthMaps to recover per-face scales, then multiply the scales into +each multiband's position channels before reconstructing a PLY. +""" + +from __future__ import annotations + +import hashlib +import logging +import time + +import numpy as np +import torch +import torch.nn.functional as F + +from comfy_api.latest import io + +log = logging.getLogger("sharp") + + +# Channel order for the multiband output. Mirrors the gaussian composer's +# concat order in the un-flattened tensor we assemble. +GAUSS_ATTR_CHANNEL_NAMES = [ + "position_x", "position_y", "position_z", + "scale_x", "scale_y", "scale_z", + "quaternion_w", "quaternion_x", "quaternion_y", "quaternion_z", + "color_r", "color_g", "color_b", + "opacity", +] +NUM_GAUSS_ATTR_CHANNELS = len(GAUSS_ATTR_CHANNEL_NAMES) # 14 + + +# Same encode cache as SharpPredict — feature-extraction is the heaviest +# step, and a user iterating on focal length should hit the cache. +_encode_cache: dict[str, object] = { + "image_hash": None, + "monodepth_output": None, + "image_resized": None, + "original_shape": None, +} + + +def _compute_image_hash(image_np: np.ndarray) -> str: + return hashlib.sha256(image_np.tobytes()).hexdigest()[:16] + + +def _monodepth_to(output, device): + """Move all tensors in a MonodepthOutput to the given device.""" + from .sharp.model import MonodepthOutput + return MonodepthOutput( + disparity=output.disparity.to(device), + encoder_features=[t.to(device) for t in output.encoder_features], + decoder_features=output.decoder_features.to(device), + output_features=[t.to(device) for t in output.output_features], + intermediate_features=[t.to(device) for t in output.intermediate_features], + ) + + +def _split_layers_from_flat_gaussians( + gaussians, + num_layers: int, + H_grid: int, + W_grid: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Flatten Gaussians3D -> per-layer per-pixel attribute stack. + + The composer's flatten was: + x.permute(0, 2, 3, 4, 1).flatten(1, 3) + starting from [B, C, num_layers, H, W]. The flatten dim ordering is + (num_layers -> H -> W), so reshaping back to [B, num_layers, H, W, C] + recovers the original grid layout. + + Returns: + layer0_attrs: [B, H, W, 14] + layer1_attrs: [B, H, W, 14] (for num_layers=2; if num_layers + differs, raises) + """ + mean = gaussians.mean_vectors # [B, N, 3] + scale = gaussians.singular_values # [B, N, 3] + quat = gaussians.quaternions # [B, N, 4] + color = gaussians.colors # [B, N, 3] + opac = gaussians.opacities # [B, N] + + B, N, _ = mean.shape + expected_N = num_layers * H_grid * W_grid + if N != expected_N: + raise ValueError( + f"gaussian count mismatch: got N={N}, expected " + f"num_layers={num_layers} * H={H_grid} * W={W_grid} = {expected_N}. " + f"Check predictor.output_resolution / num_layers." + ) + + # Un-flatten: [B, N, C] -> [B, num_layers, H, W, C] + mean5 = mean.view(B, num_layers, H_grid, W_grid, 3) + scale5 = scale.view(B, num_layers, H_grid, W_grid, 3) + quat5 = quat.view(B, num_layers, H_grid, W_grid, 4) + color5 = color.view(B, num_layers, H_grid, W_grid, 3) + opac5 = opac.view(B, num_layers, H_grid, W_grid, 1) + + # Concat into a single [B, num_layers, H, W, 14] tensor in the order + # the multiband channel names declare. + stacked = torch.cat([mean5, scale5, quat5, color5, opac5], dim=-1) + if stacked.shape[-1] != NUM_GAUSS_ATTR_CHANNELS: + raise RuntimeError( + f"channel count mismatch: stacked={stacked.shape[-1]}, " + f"expected {NUM_GAUSS_ATTR_CHANNELS}" + ) + if num_layers < 2: + raise ValueError( + f"expected num_layers>=2 for SHARP's two-layer output, got {num_layers}" + ) + return stacked[:, 0], stacked[:, 1] + + +def _build_multiband(attrs: torch.Tensor) -> dict: + """Convert [B, H, W, 14] -> MULTIBAND_IMAGE dict ([B, 14, H, W]).""" + if attrs.dim() != 4 or attrs.shape[-1] != NUM_GAUSS_ATTR_CHANNELS: + raise ValueError( + f"expected [B, H, W, {NUM_GAUSS_ATTR_CHANNELS}], got {tuple(attrs.shape)}" + ) + samples = attrs.permute(0, 3, 1, 2).contiguous().float() # [B, C, H, W] + return { + "samples": samples, + "channel_names": list(GAUSS_ATTR_CHANNEL_NAMES), + "metadata": { + "source": "SharpPredictGaussianAttrs", + "ndc_space": True, + "note": ( + "position channels are in NDC space (ndc_matrix[2,2]=1.0 " + "so position_z == camera-space z). Apply per-face scale " + "to positions and run unproject_gaussians to recover " + "world coords." + ), + }, + } + + +class SharpPredictGaussianAttrs(io.ComfyNode): + """SHARP inference -> per-layer depth + multiband gaussian attributes.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SharpPredictGaussianAttrs", + display_name="SHARP Predict (Depth + Gaussian Attrs)", + category="SHARP", + description=( + "Run SHARP inference and expose the per-pixel gaussian " + "decoder output as structured per-layer tensors (depth + " + "MULTIBAND_IMAGE of all gaussian attributes) instead of " + "flattening to a PLY file.\n\n" + "Use the depth outputs to drive AlignDepthMaps for per-" + "face scale recovery, then multiply scales into the " + "multiband position channels before reconstructing PLYs." + ), + inputs=[ + io.Custom("SHARP_MODEL_CONFIG").Input("model"), + io.Image.Input("image"), + io.Float.Input( + "focal_length_mm", default=30.0, min=0.0, max=500.0, + step=0.1, optional=True, + tooltip="Focal length in mm (35mm equiv). 0 = 30mm. " + "Ignored if intrinsics provided."), + io.Custom("EXTRINSICS").Input( + "extrinsics", optional=True, + tooltip="Pass-through (e.g. from SamplePanorama / " + "SharpPanoramaCubeSplit). Re-emitted on the " + "extrinsics output so downstream nodes " + "(AlignDepthMaps) can wire them."), + io.Custom("INTRINSICS").Input( + "intrinsics", optional=True, + tooltip="Pass-through. If provided, overrides " + "focal_length_mm. Re-emitted on the intrinsics " + "output."), + ], + outputs=[ + io.Image.Output( + display_name="first_layer_depth_refined", + tooltip="[B, H, W, 3] depth broadcast across 3 channels. " + "Layer-0 (visible surface) post-decoder z."), + io.Custom("MULTIBAND_IMAGE").Output( + display_name="first_layer_gaussian_attrs", + tooltip="[B, 14, H, W] all layer-0 gaussian attributes. " + "Channels: position_xyz, scale_xyz, quat_wxyz, " + "color_rgb, opacity."), + io.Image.Output( + display_name="second_layer_depth_refined", + tooltip="[B, H, W, 3] layer-1 (back/occluded surface) " + "post-decoder z."), + io.Custom("MULTIBAND_IMAGE").Output( + display_name="second_layer_gaussian_attrs", + tooltip="[B, 14, H, W] all layer-1 gaussian attributes."), + io.Custom("EXTRINSICS").Output(display_name="extrinsics"), + io.Custom("INTRINSICS").Output(display_name="intrinsics"), + ], + ) + + @classmethod + @torch.no_grad() + def execute( + cls, + model, + image: torch.Tensor, + focal_length_mm: float = 30.0, + extrinsics: torch.Tensor | None = None, + intrinsics: torch.Tensor | None = None, + ): + global _encode_cache + import comfy.model_management + from .load_model import _load_sharp_model + + patcher = _load_sharp_model(model) + predictor = patcher.model + device = patcher.load_device + + if image.dim() == 3: + image = image.unsqueeze(0) + B = image.shape[0] + + log.info( + f"[SharpPredictGaussianAttrs] processing {B} image(s); " + f"output_resolution={predictor.output_resolution}, " + f"num_layers=2 (SHARP convention)" + ) + + # Process per-image (matches existing SharpPredict structure). + first_depths = [] + first_attrs = [] + second_depths = [] + second_attrs = [] + + # Load model to GPU with the budget the existing predict path uses. + internal_shape = (1536, 1536) + input_shape = [1, 3, internal_shape[0], internal_shape[1]] + memory_required = patcher.memory_required(input_shape) + comfy.model_management.load_models_gpu( + [patcher], memory_required=memory_required, + ) + + H_grid = W_grid = int(predictor.output_resolution) + num_layers = 2 # SHARP fixed; see model.py:__init__ default + + for b in range(B): + img_np = image[b].cpu().numpy() if isinstance(image, torch.Tensor) else np.asarray(image[b]) + if img_np.dtype != np.uint8: + img_np = (np.clip(img_np, 0, 1) * 255 + 0.5).astype(np.uint8) + height, width = img_np.shape[:2] + image_hash = _compute_image_hash(img_np) + + # Cache logic — same as SharpPredict. + if _encode_cache["image_hash"] == image_hash: + log.info(f" [{b}] cache hit") + monodepth_output = _monodepth_to(_encode_cache["monodepth_output"], device) + image_resized_pt = _encode_cache["image_resized"].to(device) + else: + log.info(f" [{b}] encoding...") + _encode_cache["image_hash"] = None + image_pt = ( + torch.from_numpy(img_np.copy()).float().to(device).permute(2, 0, 1) / 255.0 + ) + image_resized_pt = F.interpolate( + image_pt[None], + size=(internal_shape[1], internal_shape[0]), + mode="bilinear", align_corners=True, + ) + t0 = time.time() + monodepth_output, _ = predictor.encode(image_resized_pt) + log.info(f" [{b}] encode time: {time.time() - t0:.2f}s") + _encode_cache["image_hash"] = image_hash + _encode_cache["monodepth_output"] = _monodepth_to(monodepth_output, "cpu") + _encode_cache["image_resized"] = image_resized_pt.cpu() + _encode_cache["original_shape"] = (height, width) + comfy.model_management.soft_empty_cache() + + # f_px on the original image, projected to internal_shape's coords. + # Matches SharpPredict's disparity_factor convention so the decoder + # produces the same gaussians. + if intrinsics is not None: + intr_b = intrinsics[b] if intrinsics.dim() == 3 else intrinsics + f_px = float(intr_b[0, 0]) * (internal_shape[0] / width) + else: + # 30mm 35mm-equiv focal length: f_px = (width / 36) * fl_mm + f_px = (width / 36.0) * max(0.1, float(focal_length_mm or 30.0)) + + disparity_factor = torch.tensor([f_px / width]).float().to(device) + t1 = time.time() + gaussians_ndc = predictor.decode(monodepth_output, image_resized_pt, disparity_factor) + log.info(f" [{b}] decode time: {time.time() - t1:.2f}s") + + # Un-flatten to per-pixel grid + split layers. + l0, l1 = _split_layers_from_flat_gaussians( + gaussians_ndc, num_layers, H_grid, W_grid, + ) + # l0, l1 shape: [1, H, W, 14] + + # Depth = position_z = channel index 2. + d0 = l0[..., 2] # [1, H, W] + d1 = l1[..., 2] + + first_depths.append(d0.cpu()) + second_depths.append(d1.cpu()) + first_attrs.append(l0.cpu()) + second_attrs.append(l1.cpu()) + + # Stack across batch. + d0_batch = torch.cat(first_depths, dim=0) # [B, H, W] + d1_batch = torch.cat(second_depths, dim=0) + attrs0_batch = torch.cat(first_attrs, dim=0) # [B, H, W, 14] + attrs1_batch = torch.cat(second_attrs, dim=0) + + # IMAGE format: [B, H, W, 3] depth broadcast over channels. + d0_img = d0_batch.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous() + d1_img = d1_batch.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous() + + # MULTIBAND_IMAGE: dict with samples [B, 14, H, W] + channel_names. + l0_mb = _build_multiband(attrs0_batch) + l1_mb = _build_multiband(attrs1_batch) + + n_valid_l0 = int(d0_batch.gt(0).sum()) + n_valid_l1 = int(d1_batch.gt(0).sum()) + log.info( + f"[SharpPredictGaussianAttrs] done. " + f"shape: [{B}, {H_grid}, {W_grid}], " + f"layer0 depth: min={float(d0_batch.min()):.3f} " + f"median={float(d0_batch.median()):.3f} " + f"max={float(d0_batch.max()):.3f} (valid={n_valid_l0}), " + f"layer1 depth: min={float(d1_batch.min()):.3f} " + f"median={float(d1_batch.median()):.3f} " + f"max={float(d1_batch.max()):.3f} (valid={n_valid_l1})" + ) + + return io.NodeOutput( + d0_img, l0_mb, d1_img, l1_mb, + extrinsics, intrinsics, + ) + + +NODE_CLASS_MAPPINGS = { + "SharpPredictGaussianAttrs": SharpPredictGaussianAttrs, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SharpPredictGaussianAttrs": "SHARP Predict (Depth + Gaussian Attrs)", +} diff --git a/nodes/refined_depth_from_gaussians.py b/nodes/refined_depth_from_gaussians.py new file mode 100644 index 0000000..3af9cd9 --- /dev/null +++ b/nodes/refined_depth_from_gaussians.py @@ -0,0 +1,291 @@ +"""SharpRefinedDepthFromGaussians node for ComfyUI-Sharp. + +Recovers a per-face "refined depth map" from the gaussian PLYs produced by +SharpPredict, by reprojecting each face's layer-0 (visible-surface) +gaussians back into that face's camera frustum. + +Why: SHARP's `SharpPredictDepth` exposes the pre-refinement disparity head +output. The full gaussian decoder produces a refined per-gaussian z that's +typically tighter than the disparity head's prediction. Those refined z's +ARE in the saved PLY (in world coords after `unproject_gaussians`); this +node inverts that unproject to recover the camera-space depth they +represent. + +Pipeline shape: + SamplePanorama / SharpPanoramaCubeSplit + -> images, ext, intr + SharpPredict (batch) + -> ply_path (folder with N PLYs), ext, intr + SharpRefinedDepthFromGaussians (this node) + -> depth_maps [N, H, W, 3] # refined per-face depth + AlignDepthMaps + -> globally scaled depth_maps + (multiply scales back into the gaussians + MergeGaussians) + +Algorithm per face i: + g_world = layer-0 gaussian positions in world coords (loaded from PLY) + cam = R_w2c[i] @ g_world + t_w2c[i] # world -> camera + z = cam[:, 2] # depth along view axis + valid = z > 0 # in front of camera + u = fx[i] * cam[:, 0] / z + cx[i] # face pixel coords + v = fy[i] * cam[:, 1] / z + cy[i] + depth_map[v, u] = min(depth_map[v, u], z) # nearest gaussian per pixel + # = layer-0 by construction + # (layer-1 is behind layer-0) + +Output is sparse — only pixels where a layer-0 gaussian landed get a depth. +Other pixels are 0 + masked invalid. +""" + +import logging +import os +from pathlib import Path + +import numpy as np +import torch + +from comfy_api.latest import io + +log = logging.getLogger("sharp") + + +def _load_gaussian_positions(ply_path: Path) -> np.ndarray: + """Read just the x/y/z fields from a Sharp-format PLY. Returns [N, 3].""" + from plyfile import PlyData + plydata = PlyData.read(str(ply_path)) + vertex = plydata['vertex'] + return np.stack([ + np.asarray(vertex['x']), + np.asarray(vertex['y']), + np.asarray(vertex['z']), + ], axis=-1).astype(np.float32) + + +def _list_ply_paths(ply_path_str: str, expected_n: int) -> list[Path]: + """Resolve the ply_path input (SharpPredict output) to a list of PLYs. + + SharpPredict writes: + - single image: _.ply + - batch: _/001.ply, 002.ply, ... + """ + p = Path(ply_path_str) + if p.is_file() and p.suffix.lower() == ".ply": + if expected_n != 1: + raise ValueError( + f"ply_path is a single file ({p.name}) but extrinsics has " + f"N={expected_n}. Pass the SharpPredict batch folder instead." + ) + return [p] + if p.is_dir(): + plies = sorted( + child for child in p.iterdir() + if child.is_file() and child.suffix.lower() == ".ply" + ) + if len(plies) != expected_n: + raise ValueError( + f"ply_path directory {p} has {len(plies)} PLYs but extrinsics " + f"has N={expected_n}. Did SharpPredict run on the same batch?" + ) + return plies + raise FileNotFoundError(f"ply_path {ply_path_str!r} is not a file or directory") + + +def _project_gaussians_to_depth( + gauss_world: torch.Tensor, + ext_w2c: torch.Tensor, + K_pixel: torch.Tensor, + out_h: int, + out_w: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Project N gaussian positions into one face camera, return depth map. + + Args: + gauss_world: [G, 3] world-space gaussian positions + ext_w2c: [4, 4] world-to-camera matrix for this face + K_pixel: [3, 3] pixel-space intrinsics + out_h, out_w: output depth map size + + Returns: + depth_map: [H, W] float; depth at each pixel where a gaussian landed, + 0 elsewhere. Per-pixel = min over all gaussians that hit + that pixel (so layer-0 is taken automatically since + layer-1 sits behind layer-0). + valid: [H, W] bool; True where at least one gaussian landed. + """ + device = gauss_world.device + + # world -> camera (homogeneous) + R = ext_w2c[:3, :3] + t = ext_w2c[:3, 3] + cam = gauss_world @ R.T + t # [G, 3] + z = cam[:, 2] + + # Project to pixel coords. Only consider gaussians in front of camera. + fx = K_pixel[0, 0] + fy = K_pixel[1, 1] + cx = K_pixel[0, 2] + cy = K_pixel[1, 2] + + in_front = z > 1e-6 + z_safe = torch.where(in_front, z, torch.ones_like(z)) + u = fx * cam[:, 0] / z_safe + cx + v = fy * cam[:, 1] / z_safe + cy + + # Round to nearest pixel; require in-bounds. + u_i = u.round().long() + v_i = v.round().long() + in_bounds = (u_i >= 0) & (u_i < out_w) & (v_i >= 0) & (v_i < out_h) & in_front + + # Scatter: keep MIN z per pixel (-> layer-0, visible surface). + # 1-D flat pixel index for scatter_reduce. + flat_idx = v_i * out_w + u_i + # For invalid gaussians, point at index 0 with z = +inf so they lose every reduce. + safe_idx = torch.where(in_bounds, flat_idx, torch.zeros_like(flat_idx)) + safe_z = torch.where(in_bounds, z, torch.full_like(z, float("inf"))) + + depth_flat = torch.full((out_h * out_w,), float("inf"), dtype=torch.float32, device=device) + depth_flat.scatter_reduce_(0, safe_idx, safe_z, reduce="amin", include_self=True) + + depth_map = depth_flat.view(out_h, out_w) + valid = torch.isfinite(depth_map) + depth_map = torch.where(valid, depth_map, torch.zeros_like(depth_map)) + + return depth_map, valid + + +class SharpRefinedDepthFromGaussians(io.ComfyNode): + """Per-face gaussian PLYs -> per-face refined depth maps (layer-0).""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SharpRefinedDepthFromGaussians", + display_name="SHARP Refined Depth from Gaussians", + category="SHARP", + description=( + "Reproject SharpPredict's PLY gaussians back into per-face " + "camera space to recover the refined per-pixel depth (the " + "actual z's the decoder produced, NOT the pre-refinement " + "disparity head). Min-depth per pixel = layer 0 (visible " + "surface).\n\n" + "Feed depth_maps into AlignDepthMaps for log-scale " + "alignment, then apply the resulting per-face scales back " + "to the gaussians before MergeGaussians." + ), + inputs=[ + io.String.Input( + "ply_path", + tooltip="From SharpPredict.ply_path. Single .ply (for " + "batch_size=1) OR a folder with one .ply per " + "face (for batch). Folder file count must match " + "extrinsics N."), + io.Custom("EXTRINSICS").Input( + "extrinsics", + tooltip="[N, 4, 4] world-to-camera per face. Same matrix " + "you passed to SharpPredict — typically from " + "SamplePanorama or SharpPanoramaCubeSplit."), + io.Custom("INTRINSICS").Input( + "intrinsics", + tooltip="[N, 3, 3] (or [N, 4, 4]) pixel-space intrinsics " + "per face. The 3x3 K is used; any extra row/col " + "is ignored."), + io.Int.Input( + "out_size", default=1536, min=256, max=2048, step=64, + tooltip="Output depth-map resolution per face (square). " + "1536 matches SHARP's internal resolution."), + ], + outputs=[ + io.Image.Output( + display_name="depth_maps", + tooltip="[N, H, W, 3] float depth maps. Depth value " + "broadcast across 3 channels so it composes " + "with regular IMAGE-consuming nodes. 0 where " + "no gaussian landed."), + io.Mask.Output( + display_name="valid_mask", + tooltip="[N, H, W] float 0/1; 1 where a gaussian landed."), + io.Custom("EXTRINSICS").Output( + display_name="extrinsics", + tooltip="Pass-through, so this node can sit upstream of " + "AlignDepthMaps without rewiring."), + io.Custom("INTRINSICS").Output( + display_name="intrinsics", + tooltip="Pass-through."), + ], + ) + + @classmethod + def execute( + cls, ply_path: str, + extrinsics: torch.Tensor, intrinsics: torch.Tensor, + out_size: int = 1536, + ): + # Normalize ext/intr shapes. + ext = extrinsics.float() + if ext.dim() == 2: + ext = ext.unsqueeze(0) + N = ext.shape[0] + if ext.shape[1:] != (4, 4): + raise ValueError(f"extrinsics must be [N, 4, 4], got {tuple(ext.shape)}") + + intr = intrinsics.float() + if intr.dim() == 2: + intr = intr.unsqueeze(0) + if intr.shape[0] != N: + raise ValueError( + f"intrinsics N ({intr.shape[0]}) doesn't match extrinsics N ({N})" + ) + # Accept 4x4 or 3x3. + K_pixel = intr[:, :3, :3].contiguous() + + # Resolve to per-face PLY paths. + ply_files = _list_ply_paths(ply_path, expected_n=N) + log.info( + f"[SharpRefinedDepthFromGaussians] loading {len(ply_files)} PLY(s) " + f"-> projecting to {out_size}x{out_size} per face" + ) + + device = ext.device + H_out = W_out = int(out_size) + + depth_maps = [] + valid_masks = [] + for i, ply in enumerate(ply_files): + pos_np = _load_gaussian_positions(ply) + pos = torch.from_numpy(pos_np).to(device=device, dtype=torch.float32) + + depth_i, valid_i = _project_gaussians_to_depth( + pos, ext[i].to(device), K_pixel[i].to(device), + out_h=H_out, out_w=W_out, + ) + depth_maps.append(depth_i) + valid_masks.append(valid_i) + + n_g = pos.shape[0] + n_v = int(valid_i.sum()) + d_min = float(depth_i[valid_i].min()) if n_v > 0 else 0.0 + d_med = float(torch.median(depth_i[valid_i])) if n_v > 0 else 0.0 + d_max = float(depth_i[valid_i].max()) if n_v > 0 else 0.0 + log.info( + f" face {i}: {n_g} gaussians -> {n_v}/{H_out*W_out} pixels covered " + f"({100*n_v/(H_out*W_out):.1f}%), depth min={d_min:.3f} " + f"median={d_med:.3f} max={d_max:.3f}" + ) + + depth_stack = torch.stack(depth_maps, dim=0) # [N, H, W] + valid_stack = torch.stack(valid_masks, dim=0) # [N, H, W] + + # IMAGE convention: [N, H, W, 3], broadcast depth across channels. + depth_img = depth_stack.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous() + valid_img = valid_stack.float().contiguous() + + return io.NodeOutput(depth_img, valid_img, extrinsics, intrinsics) + + +NODE_CLASS_MAPPINGS = { + "SharpRefinedDepthFromGaussians": SharpRefinedDepthFromGaussians, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SharpRefinedDepthFromGaussians": "SHARP Refined Depth from Gaussians", +}