feat(perception): parse STL + uniform surface sampling
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
876224a71d
commit
d1b27433d5
2 files changed
+156
No files matched your search
@@ -0,0 +1,68 @@
|
||||
"""Parse binary/ASCII STL bytes and uniformly sample the surface.
|
||||
|
||||
Output unit follows the STL convention (no unit conversion - STL files for
|
||||
robotics are by convention in millimetres, but we don't enforce it here).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
|
||||
import numpy as np
|
||||
from stl import mesh as stl_mesh
|
||||
|
||||
|
||||
def parse_stl_bytes(stl_bytes: bytes) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Parse STL bytes (binary or ASCII) -> (vertices (M,3) float64,
|
||||
triangles (T,3) int32).
|
||||
|
||||
Vertices are deduplicated so the same point used by multiple triangles
|
||||
appears once. Raises numpy_stl exceptions on malformed input.
|
||||
"""
|
||||
m = stl_mesh.Mesh.from_file(filename="dummy.stl", fh=io.BytesIO(stl_bytes))
|
||||
# m.vectors shape: (T, 3, 3) - three vertices per triangle, raw.
|
||||
if m.vectors is None or len(m.vectors) == 0:
|
||||
raise ValueError("STL parse yielded no triangles (empty or malformed)")
|
||||
tris_raw = np.asarray(m.vectors, dtype=np.float64).reshape(-1, 3)
|
||||
unique, inverse = np.unique(tris_raw, axis=0, return_inverse=True)
|
||||
triangles = inverse.reshape(-1, 3).astype(np.int32)
|
||||
return unique, triangles
|
||||
|
||||
|
||||
def sample_surface(
|
||||
vertices: np.ndarray,
|
||||
triangles: np.ndarray,
|
||||
n: int,
|
||||
seed: int = 0,
|
||||
) -> np.ndarray:
|
||||
"""Sample n points uniformly on the mesh surface, weighted by triangle area.
|
||||
|
||||
Algorithm:
|
||||
1. Compute area for each triangle.
|
||||
2. Cumulative area -> pick triangle index via inverse CDF.
|
||||
3. For each pick, sample a uniform barycentric coordinate:
|
||||
(1-sqrt(r1)) A + sqrt(r1)(1-r2) B + sqrt(r1) r2 C with r1, r2 ~ U(0,1).
|
||||
"""
|
||||
if triangles.size == 0:
|
||||
return np.empty((0, 3), dtype=np.float64)
|
||||
A = vertices[triangles[:, 0]]
|
||||
B = vertices[triangles[:, 1]]
|
||||
C = vertices[triangles[:, 2]]
|
||||
areas = 0.5 * np.linalg.norm(np.cross(B - A, C - A), axis=1)
|
||||
total = areas.sum()
|
||||
if total <= 0.0:
|
||||
# degenerate mesh: fall back to vertex sampling
|
||||
rng = np.random.default_rng(seed)
|
||||
idx = rng.integers(0, vertices.shape[0], size=n)
|
||||
return vertices[idx]
|
||||
cdf = np.cumsum(areas) / total
|
||||
rng = np.random.default_rng(seed)
|
||||
u = rng.random(n)
|
||||
tri_idx = np.searchsorted(cdf, u, side="right")
|
||||
tri_idx = np.clip(tri_idx, 0, triangles.shape[0] - 1)
|
||||
r1 = rng.random(n)
|
||||
r2 = rng.random(n)
|
||||
sqrt_r1 = np.sqrt(r1)
|
||||
w_a = (1.0 - sqrt_r1)[:, None]
|
||||
w_b = (sqrt_r1 * (1.0 - r2))[:, None]
|
||||
w_c = (sqrt_r1 * r2)[:, None]
|
||||
return (w_a * A[tri_idx] + w_b * B[tri_idx] + w_c * C[tri_idx]).astype(np.float64)
|
||||
Reference in new issue
Block a user