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)
|
||||||
@@ -0,0 +1,88 @@
|
|||||||
|
"""Unit tests for stl_sampler."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import io
|
||||||
|
import struct
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from presa_mano.perception.stl_sampler import parse_stl_bytes, sample_surface
|
||||||
|
|
||||||
|
|
||||||
|
def _make_cube_stl_bytes(side_mm: float = 10.0) -> bytes:
|
||||||
|
"""Synthesize a binary STL of an axis-aligned cube centered at origin,
|
||||||
|
from -side/2 to +side/2 on each axis. 12 triangles, no normals needed."""
|
||||||
|
s = side_mm / 2.0
|
||||||
|
v = [
|
||||||
|
(-s, -s, -s), ( s, -s, -s), ( s, s, -s), (-s, s, -s),
|
||||||
|
(-s, -s, s), ( s, -s, s), ( s, s, s), (-s, s, s),
|
||||||
|
]
|
||||||
|
# 6 faces x 2 tris (vertex order: any orientation works for our use)
|
||||||
|
tris = [
|
||||||
|
(0, 1, 2), (0, 2, 3), # z=-s
|
||||||
|
(4, 6, 5), (4, 7, 6), # z=+s
|
||||||
|
(0, 4, 5), (0, 5, 1), # y=-s
|
||||||
|
(3, 2, 6), (3, 6, 7), # y=+s
|
||||||
|
(0, 3, 7), (0, 7, 4), # x=-s
|
||||||
|
(1, 5, 6), (1, 6, 2), # x=+s
|
||||||
|
]
|
||||||
|
buf = io.BytesIO()
|
||||||
|
buf.write(b"\x00" * 80) # header
|
||||||
|
buf.write(struct.pack("<I", len(tris))) # n_triangles
|
||||||
|
for (i, j, k) in tris:
|
||||||
|
buf.write(struct.pack("<fff", 0.0, 0.0, 0.0)) # normal (unused)
|
||||||
|
for idx in (i, j, k):
|
||||||
|
buf.write(struct.pack("<fff", *v[idx]))
|
||||||
|
buf.write(struct.pack("<H", 0)) # attribute byte count
|
||||||
|
return buf.getvalue()
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_stl_cube_bbox():
|
||||||
|
stl_bytes = _make_cube_stl_bytes(side_mm=10.0)
|
||||||
|
verts, tris = parse_stl_bytes(stl_bytes)
|
||||||
|
assert verts.ndim == 2 and verts.shape[1] == 3
|
||||||
|
assert tris.ndim == 2 and tris.shape[1] == 3
|
||||||
|
assert tris.dtype.kind == "i"
|
||||||
|
# BB del cubo: [-5, +5] su ogni asse
|
||||||
|
np.testing.assert_allclose(verts.min(axis=0), [-5.0, -5.0, -5.0], atol=1e-6)
|
||||||
|
np.testing.assert_allclose(verts.max(axis=0), [ 5.0, 5.0, 5.0], atol=1e-6)
|
||||||
|
|
||||||
|
|
||||||
|
def test_sample_surface_in_bbox():
|
||||||
|
stl_bytes = _make_cube_stl_bytes(side_mm=10.0)
|
||||||
|
verts, tris = parse_stl_bytes(stl_bytes)
|
||||||
|
pts = sample_surface(verts, tris, n=1000, seed=0)
|
||||||
|
assert pts.shape == (1000, 3)
|
||||||
|
# tutti i punti devono cadere dentro la BB del cubo (più epsilon
|
||||||
|
# per il rumore di floating point bary)
|
||||||
|
np.testing.assert_array_less(pts.max(axis=0), 5.0 + 1e-4)
|
||||||
|
np.testing.assert_array_less(-pts.min(axis=0), 5.0 + 1e-4)
|
||||||
|
|
||||||
|
|
||||||
|
def test_sample_surface_distribution_uniform():
|
||||||
|
"""8 octanti del cubo, ognuno >= 8% dei punti (perfetto = 100/6 ~= 16.7%
|
||||||
|
per faccia, distribuita su 4 octanti adiacenti)."""
|
||||||
|
stl_bytes = _make_cube_stl_bytes(side_mm=10.0)
|
||||||
|
verts, tris = parse_stl_bytes(stl_bytes)
|
||||||
|
pts = sample_surface(verts, tris, n=2000, seed=1)
|
||||||
|
counts = np.zeros(8, dtype=np.int64)
|
||||||
|
for p in pts:
|
||||||
|
# bin (x>0, y>0, z>0) -> 0..7
|
||||||
|
idx = (int(p[0] > 0) << 2) | (int(p[1] > 0) << 1) | int(p[2] > 0)
|
||||||
|
counts[idx] += 1
|
||||||
|
frac = counts / pts.shape[0]
|
||||||
|
assert (frac > 0.08).all(), f"octant fractions too uneven: {frac}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_sample_surface_seed_deterministic():
|
||||||
|
stl_bytes = _make_cube_stl_bytes(side_mm=10.0)
|
||||||
|
verts, tris = parse_stl_bytes(stl_bytes)
|
||||||
|
a = sample_surface(verts, tris, n=100, seed=42)
|
||||||
|
b = sample_surface(verts, tris, n=100, seed=42)
|
||||||
|
np.testing.assert_array_equal(a, b)
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_stl_bad_payload_raises():
|
||||||
|
with pytest.raises(Exception): # numpy-stl raises various; we just want it to not silently pass
|
||||||
|
parse_stl_bytes(b"not a real STL")
|
||||||
Reference in new issue
Block a user