feat(perception): work_plane module with below-plane filter
Owns WorkPlane (with new offset_mm field), fit_plane_svd, disk-backed get/set, and the apply_below_plane_filter helper. WorkPlane.normal lives in OpenCV RGB-camera frame and is oriented toward the camera. The filter keeps points where signed_distance >= -offset_mm/1000. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Adriano Dal PastroandClaude Opus 4.7 committed
1 parent
2338ee7ea2
commit
16e32bf1c1
2 files changed
+291
No files matched your search
@@ -0,0 +1,146 @@
|
||||
"""Work-plane fit, persistence, and below-plane filter.
|
||||
|
||||
The work plane is fitted from a user-drawn rectangle on the RGB image. The
|
||||
fit uses SVD on the bbox depth cloud (after `nuwa_to_cv_frame`), so the
|
||||
plane normal lives in OpenCV RGB-camera frame and is oriented toward the
|
||||
camera (negative Z component in OpenCV convention).
|
||||
|
||||
The filter `apply_below_plane_filter` removes points more than `offset_mm`
|
||||
below the plane: keep `(normal · X + offset) >= -offset_mm/1000`. A
|
||||
positive `offset_mm` is permissive (preserves a buffer below the plane).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ..config import data_dir
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WorkPlaneError(RuntimeError):
|
||||
"""Bad input to `fit_plane_svd` (too few points, degenerate cloud)."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkPlane:
|
||||
"""Plane in CV camera frame. `normal · X + offset = 0`; `normal` is
|
||||
unit-length and oriented toward the camera (negative Z in OpenCV)."""
|
||||
normal: list[float]
|
||||
offset: float
|
||||
centroid: list[float]
|
||||
n_inliers: int
|
||||
fit_rms_mm: float
|
||||
roi: list[int]
|
||||
offset_mm: float
|
||||
created_at: float
|
||||
|
||||
|
||||
_WORK_PLANE: WorkPlane | None = None
|
||||
_WORK_PLANE_LOADED = False
|
||||
|
||||
|
||||
def _work_plane_path() -> Path:
|
||||
return data_dir() / "work_plane.json"
|
||||
|
||||
|
||||
def _save_work_plane_to_disk(wp: WorkPlane | None) -> None:
|
||||
path = _work_plane_path()
|
||||
try:
|
||||
if wp is None:
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
return
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(asdict(wp), indent=2))
|
||||
except Exception as e:
|
||||
log.warning("failed to persist work plane: %s", e)
|
||||
|
||||
|
||||
def _load_work_plane_from_disk() -> WorkPlane | None:
|
||||
path = _work_plane_path()
|
||||
try:
|
||||
with open(path, "r") as f:
|
||||
d = json.load(f)
|
||||
d.setdefault("offset_mm", 5.0) # legacy files
|
||||
return WorkPlane(**d)
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except Exception as e:
|
||||
log.warning("failed to load persisted work plane (%s)", e)
|
||||
return None
|
||||
|
||||
|
||||
def get_work_plane() -> WorkPlane | None:
|
||||
global _WORK_PLANE, _WORK_PLANE_LOADED
|
||||
if not _WORK_PLANE_LOADED:
|
||||
_WORK_PLANE = _load_work_plane_from_disk()
|
||||
_WORK_PLANE_LOADED = True
|
||||
return _WORK_PLANE
|
||||
|
||||
|
||||
def set_work_plane(wp: WorkPlane | None) -> None:
|
||||
global _WORK_PLANE, _WORK_PLANE_LOADED
|
||||
_WORK_PLANE = wp
|
||||
_WORK_PLANE_LOADED = True
|
||||
_save_work_plane_to_disk(wp)
|
||||
|
||||
|
||||
def fit_plane_svd(points_xyz: np.ndarray
|
||||
) -> tuple[np.ndarray, float, np.ndarray, float]:
|
||||
"""Least-squares plane via SVD. Returns (normal, offset, centroid, rms_mm).
|
||||
|
||||
Plane equation: `normal · X + offset = 0`. Normal is oriented toward
|
||||
the camera origin (so `signed_distance = normal · X + offset` is
|
||||
positive for points on the camera side of the plane)."""
|
||||
pts = np.asarray(points_xyz, dtype=np.float64)
|
||||
if pts.ndim != 2 or pts.shape[1] < 3:
|
||||
raise WorkPlaneError(f"unexpected points shape for plane fit: {pts.shape}")
|
||||
xyz = pts[:, :3]
|
||||
if len(xyz) < 10:
|
||||
raise WorkPlaneError(f"too few points to fit plane: {len(xyz)}")
|
||||
c = xyz.mean(axis=0)
|
||||
centered = xyz - c
|
||||
_, _, Vt = np.linalg.svd(centered, full_matrices=False)
|
||||
normal = Vt[-1]
|
||||
n = float(np.linalg.norm(normal))
|
||||
if n < 1e-9:
|
||||
raise WorkPlaneError("degenerate plane (zero normal)")
|
||||
normal = normal / n
|
||||
# Orient toward camera origin: signed distance of origin from plane must
|
||||
# be positive, i.e. `-n·c > 0`, i.e. `n·c < 0`.
|
||||
if float(normal @ c) > 0:
|
||||
normal = -normal
|
||||
offset = -float(normal @ c)
|
||||
signed = centered @ normal
|
||||
rms_mm = float(np.sqrt(np.mean(signed * signed)) * 1000.0)
|
||||
return normal, offset, c, rms_mm
|
||||
|
||||
|
||||
def apply_below_plane_filter(
|
||||
pts_cv: np.ndarray,
|
||||
wp: WorkPlane | None,
|
||||
) -> tuple[np.ndarray, dict]:
|
||||
"""Drop points more than `offset_mm` below the plane.
|
||||
|
||||
Returns `(filtered_pts, info)`. `info` always carries:
|
||||
plane_applied (bool): true iff a plane was provided AND there were
|
||||
points to filter
|
||||
n_points_before_plane (int): len(pts_cv) on entry
|
||||
"""
|
||||
info = {"plane_applied": False,
|
||||
"n_points_before_plane": int(len(pts_cv))}
|
||||
if wp is None or len(pts_cv) == 0:
|
||||
return pts_cv, info
|
||||
n_vec = np.asarray(wp.normal, dtype=np.float64)
|
||||
d_off = float(wp.offset)
|
||||
offset_m = float(wp.offset_mm) / 1000.0
|
||||
signed = pts_cv @ n_vec + d_off
|
||||
keep = signed >= -offset_m
|
||||
info["plane_applied"] = True
|
||||
return pts_cv[keep], info
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Unit tests for perception.work_plane."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from presa_mano.perception.work_plane import (
|
||||
WorkPlane,
|
||||
WorkPlaneError,
|
||||
apply_below_plane_filter,
|
||||
fit_plane_svd,
|
||||
_load_work_plane_from_disk,
|
||||
_save_work_plane_to_disk,
|
||||
_work_plane_path,
|
||||
)
|
||||
|
||||
|
||||
def _flat_plane_cloud(z: float = 1.0, n: int = 200, noise_m: float = 0.0):
|
||||
"""Cloud sampled on the plane z=z (camera frame, OpenCV: camera looks +Z)."""
|
||||
rng = np.random.default_rng(0)
|
||||
xy = rng.uniform(-0.2, 0.2, size=(n, 2))
|
||||
zs = np.full((n, 1), z) + rng.normal(0, noise_m, size=(n, 1))
|
||||
return np.hstack([xy, zs])
|
||||
|
||||
|
||||
def test_fit_plane_svd_horizontal_plane():
|
||||
"""Plane at z=1.0, normal must point toward camera (negative Z)."""
|
||||
pts = _flat_plane_cloud(z=1.0, n=200, noise_m=0.0)
|
||||
normal, offset, centroid, rms_mm = fit_plane_svd(pts)
|
||||
assert abs(normal[2]) > 0.99
|
||||
assert normal[2] < 0
|
||||
assert abs(offset - 1.0) < 1e-3
|
||||
assert rms_mm < 1.0
|
||||
|
||||
|
||||
def test_fit_plane_svd_raises_on_too_few_points():
|
||||
with pytest.raises(WorkPlaneError):
|
||||
fit_plane_svd(np.zeros((5, 3)))
|
||||
|
||||
|
||||
def test_apply_below_plane_filter_noop_without_plane():
|
||||
pts = _flat_plane_cloud(z=1.0)
|
||||
out, info = apply_below_plane_filter(pts, None)
|
||||
assert info["plane_applied"] is False
|
||||
assert info["n_points_before_plane"] == len(pts)
|
||||
assert out is pts
|
||||
|
||||
|
||||
def test_apply_below_plane_filter_keeps_above_drops_below():
|
||||
"""Plane at z=1.0, normal toward camera (-Z): points with z<1.0 are above."""
|
||||
above = _flat_plane_cloud(z=0.95, n=100)
|
||||
below = _flat_plane_cloud(z=1.10, n=100)
|
||||
pts = np.vstack([above, below])
|
||||
wp = WorkPlane(
|
||||
normal=[0.0, 0.0, -1.0], offset=1.0,
|
||||
centroid=[0.0, 0.0, 1.0],
|
||||
n_inliers=100, fit_rms_mm=0.0,
|
||||
roi=[0, 0, 10, 10],
|
||||
offset_mm=0.0,
|
||||
created_at=0.0,
|
||||
)
|
||||
out, info = apply_below_plane_filter(pts, wp)
|
||||
assert info["plane_applied"] is True
|
||||
assert info["n_points_before_plane"] == 200
|
||||
assert len(out) == 100
|
||||
assert np.all(out[:, 2] <= 1.0 + 1e-9)
|
||||
|
||||
|
||||
def test_apply_below_plane_filter_offset_buffer_keeps_below():
|
||||
"""offset_mm=50 keeps points up to 50mm below the plane."""
|
||||
above = _flat_plane_cloud(z=0.95, n=50)
|
||||
just_below = _flat_plane_cloud(z=1.03, n=50)
|
||||
far_below = _flat_plane_cloud(z=1.20, n=50)
|
||||
pts = np.vstack([above, just_below, far_below])
|
||||
wp = WorkPlane(
|
||||
normal=[0.0, 0.0, -1.0], offset=1.0,
|
||||
centroid=[0.0, 0.0, 1.0],
|
||||
n_inliers=50, fit_rms_mm=0.0,
|
||||
roi=[0, 0, 10, 10],
|
||||
offset_mm=50.0,
|
||||
created_at=0.0,
|
||||
)
|
||||
out, info = apply_below_plane_filter(pts, wp)
|
||||
assert info["plane_applied"] is True
|
||||
assert len(out) == 100 # above (50) + just_below (50); far_below dropped
|
||||
|
||||
|
||||
def test_apply_below_plane_filter_empty_input():
|
||||
wp = WorkPlane(
|
||||
normal=[0.0, 0.0, -1.0], offset=1.0,
|
||||
centroid=[0.0, 0.0, 1.0],
|
||||
n_inliers=10, fit_rms_mm=0.0,
|
||||
roi=[0, 0, 10, 10],
|
||||
offset_mm=5.0, created_at=0.0,
|
||||
)
|
||||
out, info = apply_below_plane_filter(np.empty((0, 3)), wp)
|
||||
assert info["plane_applied"] is False # short-circuit: empty input never "ran"
|
||||
assert info["n_points_before_plane"] == 0
|
||||
assert len(out) == 0
|
||||
|
||||
|
||||
def test_work_plane_disk_roundtrip(tmp_path, monkeypatch):
|
||||
"""Save then load preserves all fields including offset_mm."""
|
||||
monkeypatch.setenv("PRESA_MANO_DATA", str(tmp_path))
|
||||
wp = WorkPlane(
|
||||
normal=[0.0, 0.0, -1.0], offset=1.0,
|
||||
centroid=[0.1, 0.2, 1.0],
|
||||
n_inliers=42, fit_rms_mm=1.5,
|
||||
roi=[10, 20, 100, 200],
|
||||
offset_mm=7.5,
|
||||
created_at=1234567890.0,
|
||||
)
|
||||
_save_work_plane_to_disk(wp)
|
||||
loaded = _load_work_plane_from_disk()
|
||||
assert loaded is not None
|
||||
assert loaded.normal == wp.normal
|
||||
assert loaded.offset == wp.offset
|
||||
assert loaded.centroid == wp.centroid
|
||||
assert loaded.n_inliers == wp.n_inliers
|
||||
assert loaded.fit_rms_mm == wp.fit_rms_mm
|
||||
assert loaded.roi == wp.roi
|
||||
assert loaded.offset_mm == 7.5
|
||||
assert loaded.created_at == wp.created_at
|
||||
|
||||
|
||||
def test_work_plane_load_legacy_file_defaults_offset(tmp_path, monkeypatch):
|
||||
"""A pre-existing work_plane.json without offset_mm gets the 5.0 default."""
|
||||
monkeypatch.setenv("PRESA_MANO_DATA", str(tmp_path))
|
||||
path = _work_plane_path()
|
||||
legacy = {
|
||||
"normal": [0.0, 0.0, -1.0], "offset": 1.0,
|
||||
"centroid": [0.0, 0.0, 1.0],
|
||||
"n_inliers": 10, "fit_rms_mm": 1.0,
|
||||
"roi": [0, 0, 10, 10],
|
||||
"created_at": 0.0,
|
||||
}
|
||||
path.write_text(json.dumps(legacy))
|
||||
loaded = _load_work_plane_from_disk()
|
||||
assert loaded is not None
|
||||
assert loaded.offset_mm == 5.0
|
||||
Reference in new issue
Block a user