diff --git a/src/presa_mano/perception/work_plane.py b/src/presa_mano/perception/work_plane.py new file mode 100644 index 0000000..af9ca65 --- /dev/null +++ b/src/presa_mano/perception/work_plane.py @@ -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 diff --git a/tests/test_work_plane.py b/tests/test_work_plane.py new file mode 100644 index 0000000..af2cd0e --- /dev/null +++ b/tests/test_work_plane.py @@ -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