Presa-Mano/src/presa_mano/perception/work_plane.py

"""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