Add new node asset scripts for geometry and shader operations
- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters. - Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections. - Added `validate.py` to validate the integrity of the newly created node groups and their inputs. - Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies. - Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
This commit is contained in:
@@ -1,35 +1,7 @@
|
||||
from .grid_mesh import (
|
||||
compact_mesh,
|
||||
confidence_percentile_map,
|
||||
confidence_percentile_threshold,
|
||||
depth_edge_mask,
|
||||
depth_discontinuity_ratio,
|
||||
edge_culling_stretch_limit,
|
||||
foreground_mesh_mask,
|
||||
mesh_triangle_faces,
|
||||
remove_background_faces,
|
||||
triangle_edge_stretch_values,
|
||||
triangle_edge_stretch_mask,
|
||||
)
|
||||
from .projection import unproject_camera_depth, unproject_depth_frame
|
||||
from .surface import GeometryFrame, SurfaceMesh, build_surface
|
||||
|
||||
|
||||
__all__ = (
|
||||
"GeometryFrame",
|
||||
"SurfaceMesh",
|
||||
"build_surface",
|
||||
"compact_mesh",
|
||||
"confidence_percentile_map",
|
||||
"confidence_percentile_threshold",
|
||||
"depth_edge_mask",
|
||||
"depth_discontinuity_ratio",
|
||||
"edge_culling_stretch_limit",
|
||||
"foreground_mesh_mask",
|
||||
"mesh_triangle_faces",
|
||||
"remove_background_faces",
|
||||
"triangle_edge_stretch_values",
|
||||
"triangle_edge_stretch_mask",
|
||||
"unproject_camera_depth",
|
||||
"unproject_depth_frame",
|
||||
)
|
||||
|
||||
@@ -1,211 +0,0 @@
|
||||
def _choose_boundary_edge(start, direction, candidates):
|
||||
right = (-direction[1], direction[0])
|
||||
left = (direction[1], -direction[0])
|
||||
back = (-direction[0], -direction[1])
|
||||
preference = {right: 0, direction: 1, left: 2, back: 3}
|
||||
return min(
|
||||
candidates,
|
||||
key=lambda end: preference.get(
|
||||
(end[0] - start[0], end[1] - start[1]),
|
||||
4,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _simplify_axis_aligned_contour(points):
|
||||
simplified = []
|
||||
for index, current in enumerate(points):
|
||||
previous = points[index - 1]
|
||||
following = points[(index + 1) % len(points)]
|
||||
first = (current[0] - previous[0], current[1] - previous[1])
|
||||
second = (following[0] - current[0], following[1] - current[1])
|
||||
if first[0] * second[1] - first[1] * second[0] != 0:
|
||||
simplified.append(current)
|
||||
return simplified
|
||||
|
||||
|
||||
def trace_mask_contours(mask):
|
||||
import numpy as np
|
||||
|
||||
mask = np.asarray(mask, dtype=bool)
|
||||
if mask.ndim != 2:
|
||||
raise ValueError("Cutout Alpha must be two-dimensional")
|
||||
height, width = mask.shape
|
||||
edges = {}
|
||||
|
||||
def add_edge(start, end):
|
||||
edges.setdefault(start, []).append(end)
|
||||
|
||||
top = mask.copy()
|
||||
top[1:] &= ~mask[:-1]
|
||||
for y, x in np.argwhere(top):
|
||||
add_edge((int(x), int(y)), (int(x) + 1, int(y)))
|
||||
|
||||
right = mask.copy()
|
||||
right[:, :-1] &= ~mask[:, 1:]
|
||||
for y, x in np.argwhere(right):
|
||||
add_edge((int(x) + 1, int(y)), (int(x) + 1, int(y) + 1))
|
||||
|
||||
bottom = mask.copy()
|
||||
bottom[:-1] &= ~mask[1:]
|
||||
for y, x in np.argwhere(bottom):
|
||||
add_edge((int(x) + 1, int(y) + 1), (int(x), int(y) + 1))
|
||||
|
||||
left = mask.copy()
|
||||
left[:, 1:] &= ~mask[:, :-1]
|
||||
for y, x in np.argwhere(left):
|
||||
add_edge((int(x), int(y) + 1), (int(x), int(y)))
|
||||
if not edges:
|
||||
raise ValueError("The Cutout region has no usable boundary")
|
||||
|
||||
contours = []
|
||||
edge_count = sum(len(values) for values in edges.values())
|
||||
while edges:
|
||||
start = min(edges, key=lambda point: (point[1], point[0]))
|
||||
current = start
|
||||
direction = (1, 0)
|
||||
contour = []
|
||||
for _index in range(edge_count + 1):
|
||||
contour.append(current)
|
||||
candidates = edges.get(current)
|
||||
if not candidates:
|
||||
raise ValueError("The Cutout boundary is open")
|
||||
following = _choose_boundary_edge(current, direction, candidates)
|
||||
candidates.remove(following)
|
||||
if not candidates:
|
||||
del edges[current]
|
||||
direction = (following[0] - current[0], following[1] - current[1])
|
||||
current = following
|
||||
if current == start:
|
||||
break
|
||||
else:
|
||||
raise ValueError("The Cutout boundary could not be closed")
|
||||
contour = _simplify_axis_aligned_contour(contour)
|
||||
if len(contour) >= 3:
|
||||
contours.append(contour)
|
||||
if not contours:
|
||||
raise ValueError("The Cutout region is too small")
|
||||
return contours
|
||||
|
||||
|
||||
def _signed_contour_area(contour):
|
||||
return 0.5 * sum(
|
||||
start[0] * end[1] - end[0] * start[1]
|
||||
for start, end in zip(contour, contour[1:] + contour[:1])
|
||||
)
|
||||
|
||||
|
||||
def _point_inside_contour(point, contour):
|
||||
inside = False
|
||||
previous = contour[-1]
|
||||
for current in contour:
|
||||
if (current[1] > point[1]) != (previous[1] > point[1]):
|
||||
crossing_x = current[0] + (
|
||||
(point[1] - current[1])
|
||||
* (previous[0] - current[0])
|
||||
/ (previous[1] - current[1])
|
||||
)
|
||||
inside ^= point[0] < crossing_x
|
||||
previous = current
|
||||
return inside
|
||||
|
||||
|
||||
def mask_regions(mask):
|
||||
contours = trace_mask_contours(mask)
|
||||
outer = [contour for contour in contours if _signed_contour_area(contour) > 0.0]
|
||||
holes = [contour for contour in contours if _signed_contour_area(contour) < 0.0]
|
||||
if not outer:
|
||||
raise ValueError("The Cutout region has no outer boundary")
|
||||
regions = [[contour] for contour in outer]
|
||||
outer_areas = [abs(_signed_contour_area(contour)) for contour in outer]
|
||||
for hole in holes:
|
||||
start = hole[0]
|
||||
end = hole[1]
|
||||
edge_x = end[0] - start[0]
|
||||
edge_y = end[1] - start[1]
|
||||
length = max((edge_x * edge_x + edge_y * edge_y) ** 0.5, 1e-12)
|
||||
sample = (
|
||||
(start[0] + end[0]) * 0.5 - edge_y * 0.25 / length,
|
||||
(start[1] + end[1]) * 0.5 + edge_x * 0.25 / length,
|
||||
)
|
||||
owners = [
|
||||
index
|
||||
for index, contour in enumerate(outer)
|
||||
if _point_inside_contour(sample, contour)
|
||||
]
|
||||
if not owners:
|
||||
raise ValueError("The Cutout hole has no enclosing region")
|
||||
owner = min(owners, key=outer_areas.__getitem__)
|
||||
regions[owner].append(hole)
|
||||
return regions
|
||||
|
||||
|
||||
def _resize_alpha(values, size):
|
||||
import numpy as np
|
||||
|
||||
target_width, target_height = (int(value) for value in size)
|
||||
source_height, source_width = values.shape
|
||||
if (target_width, target_height) == (source_width, source_height):
|
||||
return values
|
||||
x = (np.arange(target_width, dtype=np.float32) + 0.5) * (
|
||||
source_width / target_width
|
||||
) - 0.5
|
||||
y = (np.arange(target_height, dtype=np.float32) + 0.5) * (
|
||||
source_height / target_height
|
||||
) - 0.5
|
||||
x0_unclipped = np.floor(x).astype(np.int32)
|
||||
y0_unclipped = np.floor(y).astype(np.int32)
|
||||
x_weight = x - x0_unclipped
|
||||
y_weight = y - y0_unclipped
|
||||
x0 = np.clip(x0_unclipped, 0, source_width - 1)
|
||||
y0 = np.clip(y0_unclipped, 0, source_height - 1)
|
||||
x1 = np.clip(x0_unclipped + 1, 0, source_width - 1)
|
||||
y1 = np.clip(y0_unclipped + 1, 0, source_height - 1)
|
||||
horizontal = (
|
||||
values[:, x0] * (1.0 - x_weight)[None, :]
|
||||
+ values[:, x1] * x_weight[None, :]
|
||||
)
|
||||
return (
|
||||
horizontal[y0] * (1.0 - y_weight)[:, None]
|
||||
+ horizontal[y1] * y_weight[:, None]
|
||||
)
|
||||
|
||||
|
||||
def alpha_regions(alpha, max_resolution=512, threshold=0.5):
|
||||
import numpy as np
|
||||
|
||||
values = np.asarray(alpha)
|
||||
if values.ndim != 2:
|
||||
raise ValueError("Cutout Alpha must be two-dimensional")
|
||||
if np.issubdtype(values.dtype, np.floating):
|
||||
values = np.clip(values, 0.0, 1.0) * 255.0
|
||||
else:
|
||||
values = np.clip(values, 0, 255).astype(np.float32)
|
||||
source_size = (values.shape[1], values.shape[0])
|
||||
scale = min(1.0, float(max_resolution) / max(source_size))
|
||||
sampled_size = tuple(max(1, round(value * scale)) for value in source_size)
|
||||
if sampled_size != source_size:
|
||||
values = _resize_alpha(values, sampled_size)
|
||||
mask = values >= 255.0 * threshold
|
||||
regions = mask_regions(mask)
|
||||
scale_x = source_size[0] / sampled_size[0]
|
||||
scale_y = source_size[1] / sampled_size[1]
|
||||
return (
|
||||
[
|
||||
[
|
||||
[[float(x) * scale_x, float(y) * scale_y] for x, y in contour]
|
||||
for contour in region
|
||||
]
|
||||
for region in regions
|
||||
],
|
||||
source_size,
|
||||
)
|
||||
|
||||
|
||||
def image_alpha_regions(path, max_resolution=512, threshold=0.5):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(path) as source:
|
||||
alpha = np.asarray(source.convert("RGBA").getchannel("A"), dtype=np.uint8)
|
||||
return alpha_regions(alpha, max_resolution=max_resolution, threshold=threshold)
|
||||
@@ -1,202 +0,0 @@
|
||||
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 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 depth_discontinuity_ratio(depth):
|
||||
import numpy as np
|
||||
|
||||
padded = np.pad(depth, 1, mode="edge")
|
||||
windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3))
|
||||
difference = np.nanmax(windows, axis=(-2, -1))
|
||||
difference -= np.nanmin(windows, axis=(-2, -1))
|
||||
return difference / np.maximum(depth, 1e-6)
|
||||
|
||||
|
||||
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 remove_background_faces(faces, foreground_mask):
|
||||
"""Keep only faces whose vertices are all foreground."""
|
||||
foreground_points = foreground_mask.reshape(-1)
|
||||
return faces[foreground_points[faces].all(axis=1)]
|
||||
|
||||
|
||||
def depth_edge_mask(depth, threshold=12.0 / 255.0):
|
||||
import numpy as np
|
||||
|
||||
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) >= round(255.0 * 0.5)
|
||||
|
||||
|
||||
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")
|
||||
return float(np.percentile(confidence[valid], percentile))
|
||||
|
||||
|
||||
def mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
confidence,
|
||||
confidence_threshold,
|
||||
filter_edges,
|
||||
max_edge_length,
|
||||
edge_threshold=12.0 / 255.0,
|
||||
edge_stretch_limit=None,
|
||||
vertex_mask=None,
|
||||
):
|
||||
import numpy as np
|
||||
|
||||
height, width = depth.shape
|
||||
indices = np.arange(height * width, dtype=np.int32).reshape(height, width)
|
||||
quads = np.stack(
|
||||
(
|
||||
indices[:-1, :-1],
|
||||
indices[1:, :-1],
|
||||
indices[1:, 1:],
|
||||
indices[:-1, 1:],
|
||||
),
|
||||
axis=-1,
|
||||
).reshape(-1, 4)
|
||||
valid = (
|
||||
np.isfinite(points).all(axis=-1)
|
||||
& np.isfinite(depth)
|
||||
& (depth > 0.0)
|
||||
)
|
||||
if confidence_threshold is not None:
|
||||
valid &= np.isfinite(confidence) & (
|
||||
confidence >= confidence_threshold
|
||||
)
|
||||
if filter_edges:
|
||||
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)
|
||||
|
||||
60
src/anyimage/server/geometry/normal_texture.py
Normal file
60
src/anyimage/server/geometry/normal_texture.py
Normal file
@@ -0,0 +1,60 @@
|
||||
"""Encode camera-space normals as a texture."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def resized_selection_mask(selection_mask, shape):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
selection_mask = np.asarray(selection_mask, dtype=bool)
|
||||
if selection_mask.shape == tuple(shape):
|
||||
return selection_mask
|
||||
resized = Image.fromarray(selection_mask.astype(np.uint8) * 255, mode="L")
|
||||
resized = resized.resize(
|
||||
(int(shape[1]), int(shape[0])),
|
||||
Image.Resampling.NEAREST,
|
||||
)
|
||||
return np.asarray(resized, dtype=np.uint8) > 0
|
||||
|
||||
|
||||
def write_camera_normal_texture(
|
||||
normal,
|
||||
valid,
|
||||
origin,
|
||||
crop_size,
|
||||
source_size,
|
||||
output_path,
|
||||
):
|
||||
"""Encode MoGe camera-axis normals into a full-image RGB texture."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
normal = np.asarray(normal, dtype=np.float32)
|
||||
valid = np.asarray(valid, dtype=bool)
|
||||
if normal.ndim != 3 or normal.shape[2] != 3:
|
||||
raise ValueError("Depth Balloon normal data must have three components")
|
||||
if valid.shape != normal.shape[:2]:
|
||||
valid = resized_selection_mask(valid, normal.shape[:2])
|
||||
encoded = np.clip(
|
||||
(normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32) + 1.0)
|
||||
* 127.5,
|
||||
0.0,
|
||||
255.0,
|
||||
).astype(np.uint8)
|
||||
encoded[~valid] = (128, 128, 255)
|
||||
crop_width, crop_height = (int(value) for value in crop_size)
|
||||
crop = Image.fromarray(encoded, mode="RGB")
|
||||
if crop.size != (crop_width, crop_height):
|
||||
crop = crop.resize(
|
||||
(crop_width, crop_height),
|
||||
Image.Resampling.BILINEAR,
|
||||
)
|
||||
source_width, source_height = (int(value) for value in source_size)
|
||||
full = Image.new("RGB", (source_width, source_height), (128, 128, 255))
|
||||
full.paste(crop, tuple(int(value) for value in origin))
|
||||
output_path = Path(output_path)
|
||||
full.save(output_path)
|
||||
return output_path
|
||||
|
||||
|
||||
@@ -1,133 +0,0 @@
|
||||
from dataclasses import dataclass, field
|
||||
from .grid_mesh import (
|
||||
confidence_percentile_map,
|
||||
depth_discontinuity_ratio,
|
||||
mesh_triangle_faces,
|
||||
remove_background_faces,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GeometryFrame:
|
||||
depth: object
|
||||
points: object
|
||||
intrinsics: object
|
||||
valid_mask: object
|
||||
confidence: object | None = None
|
||||
model_mask: object | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SurfaceMesh:
|
||||
vertices: object
|
||||
faces: object
|
||||
uv: object
|
||||
intrinsics: object
|
||||
image_size: tuple[int, int]
|
||||
reference_depth: float
|
||||
point_attributes: dict = field(default_factory=dict)
|
||||
face_attributes: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
def _validate_frame(frame):
|
||||
import numpy as np
|
||||
|
||||
depth = np.asarray(frame.depth, dtype=np.float32)
|
||||
points = np.asarray(frame.points, dtype=np.float32)
|
||||
intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
|
||||
valid_mask = np.asarray(frame.valid_mask, dtype=bool)
|
||||
if depth.ndim != 2:
|
||||
raise ValueError(f"Unsupported depth shape: {depth.shape}")
|
||||
if points.shape != (*depth.shape, 3):
|
||||
raise ValueError("Point grid and depth dimensions do not match")
|
||||
if intrinsics.shape != (3, 3):
|
||||
raise ValueError(f"Unsupported intrinsics shape: {intrinsics.shape}")
|
||||
if valid_mask.shape != depth.shape:
|
||||
raise ValueError("Valid mask and depth dimensions do not match")
|
||||
if frame.confidence is not None:
|
||||
confidence = np.asarray(frame.confidence)
|
||||
if confidence.shape != depth.shape:
|
||||
raise ValueError("Confidence and depth dimensions do not match")
|
||||
if frame.model_mask is not None:
|
||||
model_mask = np.asarray(frame.model_mask)
|
||||
if model_mask.shape != depth.shape:
|
||||
raise ValueError("Model mask and depth dimensions do not match")
|
||||
return depth, points, intrinsics, valid_mask
|
||||
|
||||
|
||||
def _compact_grid(points, faces, depth):
|
||||
import numpy as np
|
||||
|
||||
height, width = depth.shape
|
||||
used = np.unique(faces)
|
||||
index_map = np.full(height * width, -1, dtype=np.int32)
|
||||
index_map[used] = np.arange(len(used), dtype=np.int32)
|
||||
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
|
||||
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
|
||||
uv = np.stack(np.meshgrid(u, v), axis=-1).reshape(-1, 2)[used]
|
||||
return (
|
||||
points.reshape(-1, 3)[used].astype(np.float32),
|
||||
index_map[faces],
|
||||
uv,
|
||||
used,
|
||||
)
|
||||
|
||||
|
||||
def build_surface(frame, foreground_mask=None):
|
||||
import numpy as np
|
||||
|
||||
depth, points, intrinsics, valid_mask = _validate_frame(frame)
|
||||
faces = mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
np.ones(depth.shape, dtype=np.float32),
|
||||
None,
|
||||
False,
|
||||
0.0,
|
||||
vertex_mask=valid_mask,
|
||||
)
|
||||
if foreground_mask is not None:
|
||||
foreground_mask = np.asarray(foreground_mask, dtype=bool)
|
||||
if foreground_mask.shape != depth.shape:
|
||||
raise ValueError("Foreground mask and depth dimensions do not match")
|
||||
faces = remove_background_faces(faces, foreground_mask)
|
||||
if not len(faces):
|
||||
raise ValueError("Model output does not contain any valid surface faces")
|
||||
|
||||
face_attributes = {
|
||||
"face_depth_discontinuity": (
|
||||
depth_discontinuity_ratio(depth).reshape(-1)[faces].max(axis=1)
|
||||
).astype(np.float32),
|
||||
}
|
||||
if frame.confidence is not None:
|
||||
confidence = confidence_percentile_map(
|
||||
depth,
|
||||
np.asarray(frame.confidence, dtype=np.float32),
|
||||
)
|
||||
face_attributes["face_confidence"] = (
|
||||
confidence.reshape(-1)[faces].min(axis=1).astype(np.float32)
|
||||
)
|
||||
if frame.model_mask is not None:
|
||||
model_mask = np.asarray(frame.model_mask, dtype=bool)
|
||||
face_attributes["face_mask"] = model_mask.reshape(-1)[faces].all(axis=1)
|
||||
|
||||
vertices, compact_faces, uv, used = _compact_grid(points, faces, depth)
|
||||
point_depth = depth.reshape(-1)[used].astype(np.float32)
|
||||
reference_values = point_depth
|
||||
if frame.model_mask is not None:
|
||||
masked = point_depth[
|
||||
np.asarray(frame.model_mask, dtype=bool).reshape(-1)[used]
|
||||
]
|
||||
if len(masked):
|
||||
reference_values = masked
|
||||
return SurfaceMesh(
|
||||
vertices=vertices,
|
||||
faces=compact_faces.astype(np.int32),
|
||||
uv=uv.astype(np.float32),
|
||||
intrinsics=intrinsics.astype(np.float32),
|
||||
image_size=(depth.shape[1], depth.shape[0]),
|
||||
reference_depth=float(np.median(reference_values)),
|
||||
point_attributes={"point_depth": point_depth},
|
||||
face_attributes=face_attributes,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user