From 403e8ea3f01fbf6b9c1df255f6c7a732dc91d458 Mon Sep 17 00:00:00 2001 From: Alex Sadleir Date: Mon, 29 Dec 2025 13:09:30 +1100 Subject: [PATCH] Fix geometric streaks: implement pole snapping, seam stitching, and long-edge face removal in DAP_Panoramic_Mesh --- dap_modules/dap_geometry.py | 118 ++++++++++++++++++++--------------- tests/unit/test_dap_nodes.py | 11 ++-- 2 files changed, 76 insertions(+), 53 deletions(-) diff --git a/dap_modules/dap_geometry.py b/dap_modules/dap_geometry.py index 3803041..277afd6 100644 --- a/dap_modules/dap_geometry.py +++ b/dap_modules/dap_geometry.py @@ -4,7 +4,6 @@ import trimesh import os import sys -# Add dap_core to sys.path so imports work current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) dap_core_path = os.path.join(current_dir, "dap_core") if dap_core_path not in sys.path: @@ -12,10 +11,6 @@ if dap_core_path not in sys.path: def spherical_uv_to_directions_np(uv: np.ndarray): - """ - Equirectangular UV to Unit Sphere Directions - uv: [H, W, 2] in range [0, 1] - """ theta = (1.0 - uv[..., 0]) * (2.0 * np.pi) phi = uv[..., 1] * np.pi directions = np.stack( @@ -34,11 +29,16 @@ class DAP_Panoramic_Mesh: "FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.1}, ), + "downsample": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}), }, "optional": { "image": ("IMAGE",), "mask": ("MASK",), - "downsample": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}), + "remove_long_edges": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.1}, + ), + "stitch_seam": ("BOOLEAN", {"default": True}), }, } @@ -47,86 +47,106 @@ class DAP_Panoramic_Mesh: FUNCTION = "generate_mesh" CATEGORY = "DAP/Geometry" - def generate_mesh(self, depth, mesh_scale, image=None, mask=None, downsample=1): - # ComfyUI Image is [B, H, W, C] - # We take the first image in batch + def generate_mesh( + self, + depth, + mesh_scale, + downsample, + image=None, + mask=None, + remove_long_edges=0.5, + stitch_seam=True, + ): d_tensor = depth[0] H, W, C = d_tensor.shape - - # 1. Grayscale depth d_np = d_tensor.cpu().numpy() if C > 1: d_np = np.mean(d_np, axis=2) else: d_np = d_np.squeeze() - # Prevent vertices from collapsing to 0,0,0 by clamping min depth d_np = np.clip(d_np, 0.01, 1.0) - # 2. Downsample for performance if requested - if downsample > 1: d_np = d_np[::downsample, ::downsample] - h_new, w_new = d_np.shape - else: - h_new, w_new = H, W - # 3. Create UV grid + h_new, w_new = d_np.shape + + # Pole Snapping: Average depth at poles + d_np[0, :] = np.mean(d_np[0, :]) + d_np[-1, :] = np.mean(d_np[-1, :]) + y, x = np.mgrid[0:h_new, 0:w_new] u = x / (w_new - 1) v = y / (h_new - 1) uv = np.stack([u, v], axis=-1) - - # 4. Unproject to 3D directions dirs = spherical_uv_to_directions_np(uv) - - # 5. Scale by depth - # Assume depth is 0..1, we multiply by scale vertices = (dirs * d_np[..., None] * mesh_scale).reshape(-1, 3) - # 6. Generate faces for a grid - # vertices are ordered row-major - # [0, 1, ..., w-1] - # [w, w+1, ..., 2w-1] indices = np.arange(h_new * w_new).reshape(h_new, w_new) - # Two triangles per grid cell + # Grid faces f1 = np.stack([indices[:-1, :-1], indices[:-1, 1:], indices[1:, :-1]], axis=-1) f2 = np.stack([indices[:-1, 1:], indices[1:, 1:], indices[1:, :-1]], axis=-1) - faces = np.vstack([f1.reshape(-1, 3), f2.reshape(-1, 3)]) + faces_list = [f1.reshape(-1, 3), f2.reshape(-1, 3)] - # 7. Apply mask if provided - if mask is not None: - m_np = mask[0].cpu().numpy() - if downsample > 1: - m_np = m_np[::downsample, ::downsample] + # Seam Stitching (Right column back to Left column) + if stitch_seam: + s_f1 = np.stack( + [indices[:-1, -1], indices[:-1, 0], indices[1:, -1]], axis=-1 + ) + s_f2 = np.stack([indices[:-1, 0], indices[1:, 0], indices[1:, -1]], axis=-1) + faces_list.append(s_f1.reshape(-1, 3)) + faces_list.append(s_f2.reshape(-1, 3)) - valid_verts = m_np.flatten() > 0.5 - # We keep all vertices for indexing simplicity but only valid faces - # Actually, standard trimesh cleanup is better - face_mask = m_np[:-1, :-1] > 0.5 - # This is complex for indexed grids. - # Simpler: just create mesh and remove invalid vertices/faces - pass + faces = np.vstack(faces_list) - # 8. Colors vertex_colors = None if image is not None: i_np = image[0].cpu().numpy() if downsample > 1: - # Basic nearest neighbor downsampling for colors i_np = i_np[::downsample, ::downsample] - - # ComfyUI image is 0..1, trimesh expects 0..255 or 0..1? - # Trimesh usually takes 0..255 uint8 for visual. vertex_colors = (i_np.reshape(-1, 3) * 255).astype(np.uint8) - # 9. Build Mesh + # Build initial mesh mesh = trimesh.Trimesh( vertices=vertices, faces=faces, vertex_colors=vertex_colors, process=False ) - # Clean up if a mask was used or to remove long edge artifacts at the equirectangular seam - # TODO: Seam stitching + # Filter by Mask + if mask is not None: + m_np = mask[0].cpu().numpy() + if downsample > 1: + m_np = m_np[::downsample, ::downsample] + m_flat = m_np.flatten() + + # Remove faces where ANY vertex is masked (value < 0.5) + face_mask = (m_flat[mesh.faces] > 0.5).all(axis=1) + mesh.update_faces(face_mask) + mesh.remove_unreferenced_vertices() + + # Remove Long Edges (Streaks) + if remove_long_edges > 0: + edges = mesh.vertices[mesh.edges_unique] + edge_lengths = np.linalg.norm(edges[:, 0] - edges[:, 1], axis=1) + # Find vertices involved in long edges and remove faces? + # Better: find faces with long edges + v0 = mesh.vertices[mesh.faces[:, 0]] + v1 = mesh.vertices[mesh.faces[:, 1]] + v2 = mesh.vertices[mesh.faces[:, 2]] + + e1 = np.linalg.norm(v0 - v1, axis=1) + e2 = np.linalg.norm(v1 - v2, axis=1) + e3 = np.linalg.norm(v2 - v0, axis=1) + + # Threshold relative to average edge or fixed + # 0.5 is huge for a 1.0 radius sphere + bad_faces = ( + (e1 > remove_long_edges) + | (e2 > remove_long_edges) + | (e3 > remove_long_edges) + ) + mesh.update_faces(~bad_faces) + mesh.remove_unreferenced_vertices() return (mesh,) diff --git a/tests/unit/test_dap_nodes.py b/tests/unit/test_dap_nodes.py index bfb97da..ef64986 100644 --- a/tests/unit/test_dap_nodes.py +++ b/tests/unit/test_dap_nodes.py @@ -6,7 +6,6 @@ from dap_modules.dap_geometry import DAP_Panoramic_Mesh from dap_modules.dap_conversions import DAP_ERP_to_Cubemap from dap_modules.dap_normal import DAP_Normal_Map - from unittest.mock import MagicMock, patch @@ -59,8 +58,13 @@ def test_dap_geometry_process(dummy_image): mesh_gen = DAP_Panoramic_Mesh() # Run process - # depth is [1, 512, 1024, 3] from dummy_image fixture - result = mesh_gen.generate_mesh(depth=dummy_image, mesh_scale=1.0, downsample=4) + result = mesh_gen.generate_mesh( + depth=dummy_image, + mesh_scale=1.0, + downsample=4, + remove_long_edges=0.5, + stitch_seam=True, + ) mesh = result[0] assert hasattr(mesh, "vertices") @@ -115,6 +119,5 @@ def test_dap_normal_process(dummy_image): assert torch.is_tensor(result[0]) assert result[0].shape == (1, 512, 1024, 3) - # Check normalization (values should be around 0.5 for flat areas) assert result[0].min() >= 0.0 assert result[0].max() <= 1.0