Fix geometric streaks: implement pole snapping, seam stitching, and long-edge face removal in DAP_Panoramic_Mesh
This commit is contained in:
Alex Sadleir
+69
-49
@@ -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,)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user