Remove obsolete implementation and simplify project structure
This commit is contained in:
35
src/anyimage/server/geometry/__init__.py
Normal file
35
src/anyimage/server/geometry/__init__.py
Normal file
@@ -0,0 +1,35 @@
|
||||
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",
|
||||
)
|
||||
202
src/anyimage/server/geometry/grid_mesh.py
Normal file
202
src/anyimage/server/geometry/grid_mesh.py
Normal file
@@ -0,0 +1,202 @@
|
||||
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) > 0
|
||||
|
||||
|
||||
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)
|
||||
|
||||
31
src/anyimage/server/geometry/projection.py
Normal file
31
src/anyimage/server/geometry/projection.py
Normal file
@@ -0,0 +1,31 @@
|
||||
def unproject_depth_frame(depth, intrinsics, extrinsics):
|
||||
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
|
||||
camera_points = rays * depth.reshape(-1, 1)
|
||||
camera_points = np.column_stack(
|
||||
(camera_points, np.ones(camera_points.shape[0], dtype=np.float32))
|
||||
)
|
||||
world_to_camera = np.eye(4, dtype=np.float32)
|
||||
world_to_camera[: extrinsics.shape[0], :4] = extrinsics
|
||||
world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T
|
||||
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)
|
||||
133
src/anyimage/server/geometry/surface.py
Normal file
133
src/anyimage/server/geometry/surface.py
Normal file
@@ -0,0 +1,133 @@
|
||||
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