Add geometry nodes asset and validation script
- Implemented a new geometry nodes asset in `build_geometry_nodes_asset.py` to handle mesh refinement and topology fixing. - Created a validation script in `validate_geometry_nodes_asset.py` to ensure the integrity and functionality of the geometry node group. - Added input handling for various parameters including resolution, iterations, and smoothing options. - Included checks for node group structure and attributes to validate the geometry processing pipeline.
This commit is contained in:
@@ -24,7 +24,14 @@ VIDEO_SUFFIXES = {
|
||||
".mpeg",
|
||||
}
|
||||
SCENE_FILENAME = "scene.glb"
|
||||
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz", "mesh")
|
||||
EXPORT_FORMATS = (
|
||||
"glb",
|
||||
"gs_ply",
|
||||
"depth_vis",
|
||||
"mini_npz",
|
||||
"mesh",
|
||||
"mesh_data",
|
||||
)
|
||||
HF_ENDPOINT = "https://huggingface.co"
|
||||
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
|
||||
MODELSCOPE_REVISION = "master"
|
||||
@@ -368,41 +375,357 @@ def unproject_depth_frame(depth, intrinsics, extrinsics):
|
||||
return world_points[:, :3].reshape(height, width, 3)
|
||||
|
||||
|
||||
def unproject_camera_depth(depth, intrinsics):
|
||||
import numpy as np
|
||||
|
||||
height, width = depth.shape
|
||||
x, y = np.meshgrid(
|
||||
np.arange(width, dtype=np.float32),
|
||||
np.arange(height, dtype=np.float32),
|
||||
)
|
||||
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
|
||||
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
|
||||
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)
|
||||
|
||||
|
||||
def smooth_depth_map(depth, smoothing):
|
||||
import numpy as np
|
||||
|
||||
factor = min(max(float(smoothing) / 100.0, 0.0), 1.0)
|
||||
if factor <= 0.0:
|
||||
return np.asarray(depth, dtype=np.float32).copy()
|
||||
import cv2
|
||||
|
||||
values = np.asarray(depth, dtype=np.float32)
|
||||
valid = np.isfinite(values) & (values > 0.0)
|
||||
if not valid.any():
|
||||
raise ValueError("DA3 returned no valid depth values")
|
||||
low, high = np.percentile(values[valid], (5.0, 95.0))
|
||||
depth_range = max(float(high - low), 1e-6)
|
||||
diameter = 3 + 2 * int(round(3.0 * factor))
|
||||
working_values = values.copy()
|
||||
working_values[~valid] = float(np.median(values[valid]))
|
||||
smoothed = cv2.bilateralFilter(
|
||||
working_values,
|
||||
diameter,
|
||||
depth_range * (0.01 + 0.09 * factor),
|
||||
1.0 + 4.0 * factor,
|
||||
)
|
||||
smoothed[~valid] = values[~valid]
|
||||
return smoothed
|
||||
|
||||
|
||||
def compact_mesh(vertices, faces):
|
||||
import numpy as np
|
||||
|
||||
vertices = np.asarray(vertices, dtype=np.float32)
|
||||
faces = np.asarray(faces, dtype=np.int32)
|
||||
used_vertices = np.unique(faces)
|
||||
index_map = np.full(len(vertices), -1, dtype=np.int32)
|
||||
index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32)
|
||||
vertices = vertices[used_vertices]
|
||||
faces = index_map[faces]
|
||||
return vertices, faces
|
||||
|
||||
|
||||
def project_camera_uv(vertices, intrinsics, image_size):
|
||||
import numpy as np
|
||||
|
||||
width, height = image_size
|
||||
z = vertices[:, 2]
|
||||
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
|
||||
raise ValueError("Simplified mesh contains invalid camera depth")
|
||||
u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2]
|
||||
v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2]
|
||||
denominator_x = max(width - 1, 1)
|
||||
denominator_y = max(height - 1, 1)
|
||||
return np.column_stack(
|
||||
(
|
||||
np.clip(u / denominator_x, 0.0, 1.0),
|
||||
np.clip(1.0 - v / denominator_y, 0.0, 1.0),
|
||||
)
|
||||
).astype(np.float32)
|
||||
|
||||
|
||||
def project_camera_pixels(vertices, intrinsics):
|
||||
import numpy as np
|
||||
|
||||
z = vertices[:, 2]
|
||||
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
|
||||
raise ValueError("Mesh contains invalid camera depth")
|
||||
return np.column_stack(
|
||||
(
|
||||
intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2],
|
||||
intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def sample_image_nearest(values, pixels):
|
||||
import numpy as np
|
||||
|
||||
height, width = values.shape
|
||||
x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1)
|
||||
y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1)
|
||||
return np.asarray(values)[y, x]
|
||||
|
||||
|
||||
def confidence_percentile_map(depth, confidence):
|
||||
import numpy as np
|
||||
|
||||
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
|
||||
result = np.zeros(depth.shape, dtype=np.float32)
|
||||
if not valid.any():
|
||||
return result
|
||||
sorted_values = np.sort(np.asarray(confidence)[valid])
|
||||
ranks = np.searchsorted(
|
||||
sorted_values,
|
||||
np.asarray(confidence)[valid],
|
||||
side="right",
|
||||
)
|
||||
result[valid] = ranks.astype(np.float32) / len(sorted_values)
|
||||
return result
|
||||
|
||||
|
||||
def triangle_edge_stretch_values(points, triangles):
|
||||
import numpy as np
|
||||
|
||||
flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3)
|
||||
triangle_points = flat_points[triangles]
|
||||
triangle_depth = triangle_points[..., 2]
|
||||
safe_depth = np.where(
|
||||
np.abs(triangle_depth) > 1e-8,
|
||||
triangle_depth,
|
||||
np.nan,
|
||||
)
|
||||
triangle_rays = triangle_points / safe_depth[..., None]
|
||||
next_points = np.roll(triangle_points, -1, axis=1)
|
||||
next_rays = np.roll(triangle_rays, -1, axis=1)
|
||||
next_depth = np.roll(triangle_depth, -1, axis=1)
|
||||
local_depth = (triangle_depth + next_depth) * 0.5
|
||||
flat_edges = (
|
||||
triangle_rays * local_depth[..., None]
|
||||
- next_rays * local_depth[..., None]
|
||||
)
|
||||
flat_lengths = np.linalg.norm(flat_edges, axis=-1)
|
||||
actual_lengths = np.linalg.norm(
|
||||
triangle_points - next_points,
|
||||
axis=-1,
|
||||
)
|
||||
stretch = actual_lengths / np.maximum(flat_lengths, 1e-8)
|
||||
stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0
|
||||
return stretch
|
||||
|
||||
|
||||
def image_mesh_attributes(
|
||||
vertices,
|
||||
faces,
|
||||
intrinsics,
|
||||
depth,
|
||||
confidence,
|
||||
confidence_percentiles,
|
||||
foreground_mask,
|
||||
):
|
||||
import numpy as np
|
||||
|
||||
pixels = project_camera_pixels(vertices, intrinsics)
|
||||
point_confidence = sample_image_nearest(confidence, pixels).astype(
|
||||
np.float32
|
||||
)
|
||||
point_confidence = np.nan_to_num(point_confidence, nan=0.0)
|
||||
point_percentile = sample_image_nearest(
|
||||
confidence_percentiles,
|
||||
pixels,
|
||||
).astype(np.float32)
|
||||
point_background = ~sample_image_nearest(foreground_mask, pixels).astype(
|
||||
bool
|
||||
)
|
||||
face_confidence = point_percentile[faces].min(axis=1)
|
||||
face_edge_stretch = triangle_edge_stretch_values(
|
||||
vertices,
|
||||
faces,
|
||||
).max(axis=1)
|
||||
face_background = point_background[faces].any(axis=1)
|
||||
return {
|
||||
"point_depth": vertices[:, 2].astype(np.float32),
|
||||
"point_confidence": point_confidence,
|
||||
"point_confidence_percentile": point_percentile,
|
||||
"point_background": point_background,
|
||||
"face_confidence_percentile": face_confidence.astype(np.float32),
|
||||
"face_edge_stretch": face_edge_stretch.astype(np.float32),
|
||||
"face_background": face_background,
|
||||
}
|
||||
|
||||
|
||||
def export_prediction_attributes(prediction, output_directory):
|
||||
import numpy as np
|
||||
|
||||
values = {}
|
||||
for name in ("depth", "conf", "intrinsics", "extrinsics"):
|
||||
value = getattr(prediction, name, None)
|
||||
if value is not None:
|
||||
values[name] = np.asarray(value)
|
||||
if "depth" in values and "conf" in values:
|
||||
values["confidence_percentile"] = np.stack(
|
||||
[
|
||||
confidence_percentile_map(depth, confidence)
|
||||
for depth, confidence in zip(
|
||||
values["depth"],
|
||||
values["conf"],
|
||||
)
|
||||
]
|
||||
)
|
||||
output_path = Path(output_directory) / "prediction.attributes.npz"
|
||||
np.savez(output_path, **values)
|
||||
return output_path
|
||||
|
||||
|
||||
def export_image_mesh_data(
|
||||
prediction,
|
||||
image_path,
|
||||
output_directory,
|
||||
foreground_path=None,
|
||||
progress_callback=None,
|
||||
):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
total_steps = 6
|
||||
|
||||
def progress(step, label):
|
||||
if progress_callback is not None:
|
||||
progress_callback(step, total_steps, label)
|
||||
|
||||
progress(1, "Validating DA3 data")
|
||||
required = ("depth", "conf", "intrinsics")
|
||||
missing = [
|
||||
name
|
||||
for name in required
|
||||
if getattr(prediction, name, None) is None
|
||||
]
|
||||
if missing:
|
||||
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
|
||||
depth_frames = np.asarray(prediction.depth, dtype=np.float32)
|
||||
confidence_frames = np.asarray(prediction.conf, dtype=np.float32)
|
||||
if depth_frames.shape[0] != 1:
|
||||
raise ValueError("Convert to Mesh currently supports one image")
|
||||
raw_depth = depth_frames[0]
|
||||
confidence = confidence_frames[0]
|
||||
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
|
||||
progress(2, "Unprojecting camera geometry")
|
||||
points = unproject_camera_depth(raw_depth, intrinsics)
|
||||
progress(3, "Preparing DA3 and BEN2 attributes")
|
||||
foreground_mask = np.ones(raw_depth.shape, dtype=bool)
|
||||
if foreground_path is not None:
|
||||
foreground_mask = foreground_mesh_mask(
|
||||
foreground_path,
|
||||
raw_depth.shape,
|
||||
)
|
||||
confidence_percentiles = confidence_percentile_map(
|
||||
raw_depth,
|
||||
confidence,
|
||||
)
|
||||
progress(4, "Building editable source topology")
|
||||
faces = mesh_triangle_faces(
|
||||
points,
|
||||
raw_depth,
|
||||
confidence,
|
||||
None,
|
||||
False,
|
||||
0.0,
|
||||
)
|
||||
if not len(faces):
|
||||
raise ValueError(
|
||||
"DA3 output does not contain any valid mesh faces"
|
||||
)
|
||||
progress(5, "Projecting UV and structured attributes")
|
||||
vertices, faces = compact_mesh(
|
||||
points.reshape(-1, 3),
|
||||
faces,
|
||||
)
|
||||
height, width = raw_depth.shape
|
||||
uv = project_camera_uv(vertices, intrinsics, (width, height))
|
||||
attributes = image_mesh_attributes(
|
||||
vertices,
|
||||
faces,
|
||||
intrinsics,
|
||||
raw_depth,
|
||||
confidence,
|
||||
confidence_percentiles,
|
||||
foreground_mask,
|
||||
)
|
||||
valid_depth = raw_depth[np.isfinite(raw_depth) & (raw_depth > 0.0)]
|
||||
reference_depth = float(np.median(valid_depth))
|
||||
output_directory = Path(output_directory)
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
progress(6, "Writing mesh data")
|
||||
result_path = output_directory / "mesh.npz"
|
||||
np.savez_compressed(
|
||||
result_path,
|
||||
vertices=vertices,
|
||||
faces=faces,
|
||||
uv=uv,
|
||||
intrinsics=intrinsics,
|
||||
image_size=np.asarray((width, height), dtype=np.int32),
|
||||
reference_depth=np.asarray(reference_depth, dtype=np.float32),
|
||||
**attributes,
|
||||
)
|
||||
texture_path = output_directory / "texture.png"
|
||||
with Image.open(image_path) as image:
|
||||
image.convert("RGBA").save(texture_path, format="PNG", compress_level=1)
|
||||
return result_path, texture_path
|
||||
|
||||
|
||||
def depth_edge_mask(depth, threshold=12.0 / 255.0):
|
||||
import numpy as np
|
||||
|
||||
padded = np.pad(depth.astype(np.float32), 1, mode="edge")
|
||||
grad_x = (
|
||||
-padded[:-2, :-2]
|
||||
+ padded[:-2, 2:]
|
||||
- 2.0 * padded[1:-1, :-2]
|
||||
+ 2.0 * padded[1:-1, 2:]
|
||||
- padded[2:, :-2]
|
||||
+ padded[2:, 2:]
|
||||
)
|
||||
grad_y = (
|
||||
-padded[:-2, :-2]
|
||||
- 2.0 * padded[:-2, 1:-1]
|
||||
- padded[:-2, 2:]
|
||||
+ padded[2:, :-2]
|
||||
+ 2.0 * padded[2:, 1:-1]
|
||||
+ padded[2:, 2:]
|
||||
)
|
||||
magnitude = np.hypot(grad_x, grad_y)
|
||||
finite = np.isfinite(magnitude)
|
||||
if not finite.any():
|
||||
return np.ones(depth.shape, dtype=bool)
|
||||
minimum = float(magnitude[finite].min())
|
||||
maximum = float(magnitude[finite].max())
|
||||
if maximum <= minimum:
|
||||
return np.zeros(depth.shape, dtype=bool)
|
||||
normalized = (magnitude - minimum) / (maximum - minimum)
|
||||
return normalized >= threshold
|
||||
values = np.asarray(depth, dtype=np.float32)
|
||||
edges = np.zeros(values.shape, dtype=bool)
|
||||
for axis in (0, 1):
|
||||
first = np.take(values, range(values.shape[axis] - 1), axis=axis)
|
||||
second = np.take(values, range(1, values.shape[axis]), axis=axis)
|
||||
valid = (
|
||||
np.isfinite(first)
|
||||
& np.isfinite(second)
|
||||
& (first > 0.0)
|
||||
& (second > 0.0)
|
||||
)
|
||||
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
|
||||
discontinuity = valid & (np.abs(first - second) / scale > threshold)
|
||||
lower = [slice(None), slice(None)]
|
||||
upper = [slice(None), slice(None)]
|
||||
lower[axis] = slice(0, -1)
|
||||
upper[axis] = slice(1, None)
|
||||
edges[tuple(lower)] |= discontinuity
|
||||
edges[tuple(upper)] |= discontinuity
|
||||
return edges
|
||||
|
||||
|
||||
def edge_culling_stretch_limit(culling):
|
||||
factor = min(max(float(culling) / 100.0, 0.0001), 1.0)
|
||||
return 1.0 / factor
|
||||
|
||||
|
||||
def foreground_mesh_mask(foreground_path, depth_shape):
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
|
||||
height, width = depth_shape
|
||||
with Image.open(foreground_path) as source:
|
||||
alpha = source.convert("RGBA").getchannel("A")
|
||||
alpha = alpha.resize(
|
||||
(width, height),
|
||||
resample=Image.Resampling.BILINEAR,
|
||||
)
|
||||
alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75))
|
||||
return np.asarray(alpha) >= 128
|
||||
|
||||
|
||||
def confidence_percentile_threshold(depth, confidence, percentile):
|
||||
import numpy as np
|
||||
|
||||
if float(percentile) <= 0.0:
|
||||
return None
|
||||
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
|
||||
if not valid.any():
|
||||
raise ValueError("DA3 returned no valid depth confidence values")
|
||||
@@ -416,6 +739,9 @@ def mesh_triangle_faces(
|
||||
confidence_threshold,
|
||||
filter_edges,
|
||||
max_edge_length,
|
||||
edge_threshold=12.0 / 255.0,
|
||||
edge_stretch_limit=None,
|
||||
vertex_mask=None,
|
||||
):
|
||||
import numpy as np
|
||||
|
||||
@@ -434,24 +760,47 @@ def mesh_triangle_faces(
|
||||
np.isfinite(points).all(axis=-1)
|
||||
& np.isfinite(depth)
|
||||
& (depth > 0.0)
|
||||
& (confidence >= confidence_threshold)
|
||||
)
|
||||
if confidence_threshold is not None:
|
||||
valid &= np.isfinite(confidence) & (
|
||||
confidence >= confidence_threshold
|
||||
)
|
||||
if filter_edges:
|
||||
valid &= ~depth_edge_mask(depth)
|
||||
face_valid = valid.reshape(-1)[quads].all(axis=1)
|
||||
if max_edge_length > 0.0:
|
||||
flat_points = points.reshape(-1, 3)
|
||||
quad_points = flat_points[quads]
|
||||
next_points = np.roll(quad_points, -1, axis=1)
|
||||
edge_lengths = np.linalg.norm(quad_points - next_points, axis=-1)
|
||||
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
|
||||
quads = quads[face_valid]
|
||||
if not len(quads):
|
||||
return np.empty((0, 3), dtype=np.int32)
|
||||
return np.concatenate(
|
||||
valid &= ~depth_edge_mask(depth, edge_threshold)
|
||||
if vertex_mask is not None:
|
||||
if np.shape(vertex_mask) != depth.shape:
|
||||
raise ValueError("Foreground mask and depth dimensions do not match")
|
||||
valid &= np.asarray(vertex_mask, dtype=bool)
|
||||
triangles = np.concatenate(
|
||||
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
|
||||
axis=0,
|
||||
)
|
||||
face_valid = valid.reshape(-1)[triangles].all(axis=1)
|
||||
if edge_stretch_limit is not None:
|
||||
face_valid &= triangle_edge_stretch_mask(
|
||||
points,
|
||||
triangles,
|
||||
edge_stretch_limit,
|
||||
)
|
||||
if max_edge_length > 0.0:
|
||||
flat_points = points.reshape(-1, 3)
|
||||
triangle_points = flat_points[triangles]
|
||||
next_points = np.roll(triangle_points, -1, axis=1)
|
||||
edge_lengths = np.linalg.norm(
|
||||
triangle_points - next_points,
|
||||
axis=-1,
|
||||
)
|
||||
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
|
||||
triangles = triangles[face_valid]
|
||||
if not len(triangles):
|
||||
return np.empty((0, 3), dtype=np.int32)
|
||||
return triangles
|
||||
|
||||
|
||||
def triangle_edge_stretch_mask(points, triangles, stretch_limit):
|
||||
stretch = triangle_edge_stretch_values(points, triangles)
|
||||
tolerance = float(stretch_limit) * (1.0 + 1e-5)
|
||||
return (stretch <= tolerance).all(axis=1)
|
||||
|
||||
|
||||
def export_textured_mesh_glb(
|
||||
@@ -512,6 +861,16 @@ def export_textured_mesh_glb(
|
||||
visual=visual,
|
||||
process=False,
|
||||
)
|
||||
mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1)
|
||||
mesh.vertex_attributes["o_confidence"] = confidence[
|
||||
index
|
||||
].reshape(-1)
|
||||
mesh.vertex_attributes["o_confidence_percentile"] = (
|
||||
confidence_percentile_map(
|
||||
depth[index],
|
||||
confidence[index],
|
||||
).reshape(-1)
|
||||
)
|
||||
mesh.remove_unreferenced_vertices()
|
||||
meshes.append(mesh)
|
||||
referenced_points.append(mesh.vertices)
|
||||
@@ -1160,6 +1519,9 @@ def exported_paths(output_directory, export_format):
|
||||
if export_format == "mesh":
|
||||
result_path = output_directory / "mesh.glb"
|
||||
return [result_path] if result_path.is_file() else []
|
||||
if export_format == "mesh_data":
|
||||
result_path = output_directory / "mesh.npz"
|
||||
return [result_path] if result_path.is_file() else []
|
||||
if export_format == "mini_npz":
|
||||
result_path = output_directory / "results.npz"
|
||||
if result_path.is_file():
|
||||
@@ -1401,6 +1763,23 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
bool(getattr(args, "mesh_filter_edges", True)),
|
||||
float(getattr(args, "mesh_max_edge_length", 0.1)),
|
||||
)
|
||||
elif args.export_format == "mesh_data":
|
||||
def mesh_progress(step, total, label):
|
||||
write_status(
|
||||
args.status,
|
||||
0.88 + 0.09 * step / total,
|
||||
f"Post-processing {step}/{total}: {label}",
|
||||
)
|
||||
|
||||
export_image_mesh_data(
|
||||
prediction,
|
||||
images[0],
|
||||
output_directory,
|
||||
foreground_path=(
|
||||
foreground_images[0] if foreground_images else None
|
||||
),
|
||||
progress_callback=mesh_progress,
|
||||
)
|
||||
else:
|
||||
export(
|
||||
prediction,
|
||||
@@ -1408,6 +1787,13 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
str(output_directory),
|
||||
**export_kwargs,
|
||||
)
|
||||
if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
|
||||
write_status(
|
||||
args.status,
|
||||
0.97,
|
||||
"Writing structured debug attributes",
|
||||
)
|
||||
export_prediction_attributes(prediction, output_directory)
|
||||
check_cancelled(cancel_check)
|
||||
|
||||
result_paths = wait_for_exported_paths(
|
||||
|
||||
Reference in New Issue
Block a user