src/surface.py - master - koda/FiVI

src/surface.py
560 строк21 KB
koda
Prepare production-ready grasp planner
18 июл 2026, 09:39
18 июл 2026, 09:39f540458
from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Protocol, cast

import numpy as np
import numpy.typing as npt
from scipy.spatial import cKDTree

from .ik import FloatArray


def _array(value: npt.ArrayLike, shape: tuple[int, ...], name: str) -> FloatArray:
    result = np.array(value, dtype=np.float64, copy=True)
    if result.shape != shape or not bool(np.all(np.isfinite(result))):
        raise ValueError(f"{name} must be finite with shape {shape}")
    result.flags.writeable = False
    return result


def _rotation(value: npt.ArrayLike) -> FloatArray:
    result = _array(value, (3, 3), "rotation")
    if not np.allclose(result.T @ result, np.eye(3), atol=1.0e-7):
        raise ValueError("rotation must be orthonormal")
    if float(np.linalg.det(result)) < 0.0:
        raise ValueError("rotation must be proper")
    return result


def _count(value: int, minimum: int, name: str) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
        raise ValueError(f"{name} must be an integer of at least {minimum}")
    return value


class Shape(str, Enum):
    """Supported surface models."""

    BOX = "box"
    SPHERE = "sphere"
    CYLINDER = "cylinder"
    CLOUD = "point_cloud"


@dataclass(frozen=True, slots=True)
class Contact:
    """Projected contact in the hand frame, measured in metres."""

    position: FloatArray
    normal: FloatArray
    error_m: float
    confidence: float

    def __post_init__(self) -> None:
        position = _array(self.position, (3,), "position")
        normal = _array(self.normal, (3,), "normal")
        length = float(np.linalg.norm(normal))
        if length <= np.finfo(np.float64).eps:
            raise ValueError("normal must be non-zero")
        normal = np.array(normal / length, dtype=np.float64, copy=True)
        normal.flags.writeable = False
        if not np.isfinite(self.error_m) or self.error_m < 0.0:
            raise ValueError("error_m must be finite and non-negative")
        if not np.isfinite(self.confidence) or not 0.0 <= self.confidence <= 1.0:
            raise ValueError("confidence must be in [0, 1]")
        object.__setattr__(self, "position", position)
        object.__setattr__(self, "normal", normal)


class Surface(Protocol):
    """Hand-frame projection and deterministic sampling contract."""

    @property
    def shape(self) -> Shape:
        """Return the fitted surface type."""

    @property
    def center(self) -> FloatArray:
        """Return the hand-frame surface center in metres."""

    def project(self, point: npt.ArrayLike) -> Contact:
        """Return the closest valid surface contact."""

    def sample(self, count: int) -> tuple[Contact, ...]:
        """Return deterministic surface contacts."""

    def distance(self, point: npt.ArrayLike) -> float:
        """Return signed distance: positive outside, negative inside."""

    def distances(self, points: npt.ArrayLike) -> FloatArray:
        """Return signed distances for hand-frame points with shape (N, 3)."""


@dataclass(frozen=True, slots=True)
class Box:
    """Oriented box surface in the hand frame."""

    center: FloatArray
    rotation: FloatArray
    dimensions_m: FloatArray
    confidence: float = 1.0

    def __post_init__(self) -> None:
        center = _array(self.center, (3,), "center")
        rotation = _rotation(self.rotation)
        dimensions = _array(self.dimensions_m, (3,), "dimensions_m")
        if bool(np.any(dimensions <= 0.0)):
            raise ValueError("dimensions_m must be positive")
        if not 0.0 <= self.confidence <= 1.0:
            raise ValueError("confidence must be in [0, 1]")
        object.__setattr__(self, "center", center)
        object.__setattr__(self, "rotation", rotation)
        object.__setattr__(self, "dimensions_m", dimensions)

    @property
    def shape(self) -> Shape:
        """Return box."""

        return Shape.BOX

    def _local(self, point: npt.ArrayLike) -> FloatArray:
        return (_array(point, (3,), "point") - self.center) @ self.rotation

    def _hand(self, point: FloatArray) -> FloatArray:
        return point @ self.rotation.T + self.center

    def project(self, point: npt.ArrayLike) -> Contact:
        """Project to the nearest fitted face."""

        local = self._local(point)
        half = 0.5 * self.dimensions_m
        candidates = np.empty((6, 3), dtype=np.float64)
        normals = np.zeros((6, 3), dtype=np.float64)
        for index, (axis, sign) in enumerate(
            ((0, -1.0), (0, 1.0), (1, -1.0), (1, 1.0), (2, -1.0), (2, 1.0))
        ):
            candidate = np.clip(local, -half, half)
            candidate[axis] = sign * half[axis]
            candidates[index] = candidate
            normals[index, axis] = sign
        errors = np.linalg.norm(candidates - local, axis=1)
        selected = int(np.argmin(errors))
        return Contact(
            self._hand(candidates[selected]),
            cast(FloatArray, normals[selected] @ self.rotation.T),
            float(errors[selected]),
            self.confidence,
        )

    def sample(self, count: int) -> tuple[Contact, ...]:
        """Sample inset grids on all faces."""

        count = _count(count, 6, "box sample count")
        values: list[Contact] = []
        half = 0.5 * self.dimensions_m
        faces = ((0, -1.0), (0, 1.0), (1, -1.0), (1, 1.0), (2, -1.0), (2, 1.0))
        base, remainder = divmod(count, len(faces))
        for face_index, (axis, sign) in enumerate(faces):
            face_count = base + (1 if face_index < remainder else 0)
            side = int(np.ceil(np.sqrt(face_count)))
            grid = np.linspace(-0.72, 0.72, side)
            other = tuple(index for index in range(3) if index != axis)
            added = 0
            for first in grid:
                for second in grid:
                    if added >= face_count:
                        break
                    local = np.zeros(3, dtype=np.float64)
                    local[axis] = sign * half[axis]
                    local[other[0]] = first * half[other[0]]
                    local[other[1]] = second * half[other[1]]
                    normal = np.zeros(3, dtype=np.float64)
                    normal[axis] = sign
                    values.append(
                        Contact(
                            self._hand(local),
                            normal @ self.rotation.T,
                            0.0,
                            self.confidence,
                        )
                    )
                    added += 1
                if added >= face_count:
                    break
        return tuple(values)

    def distance(self, point: npt.ArrayLike) -> float:
        """Return signed distance to the box boundary."""

        return float(self.distances(np.asarray(point, dtype=np.float64)[None, :])[0])

    def distances(self, points: npt.ArrayLike) -> FloatArray:
        values = np.asarray(points, dtype=np.float64)
        if values.ndim != 2 or values.shape[1:] != (3,) or not bool(np.all(np.isfinite(values))):
            raise ValueError("points must be finite with shape (N, 3)")
        local = (values - self.center) @ self.rotation
        half = 0.5 * self.dimensions_m
        delta = np.abs(local) - half
        outside = np.linalg.norm(np.maximum(delta, 0.0), axis=1)
        inside = np.minimum(np.max(delta, axis=1), 0.0)
        return cast(FloatArray, outside + inside)


@dataclass(frozen=True, slots=True)
class Sphere:
    """Sphere surface in the hand frame."""

    center: FloatArray
    radius_m: float
    confidence: float = 1.0

    def __post_init__(self) -> None:
        center = _array(self.center, (3,), "center")
        if not np.isfinite(self.radius_m) or self.radius_m <= 0.0:
            raise ValueError("radius_m must be finite and positive")
        if not 0.0 <= self.confidence <= 1.0:
            raise ValueError("confidence must be in [0, 1]")
        object.__setattr__(self, "center", center)

    @property
    def shape(self) -> Shape:
        """Return sphere."""

        return Shape.SPHERE

    def project(self, point: npt.ArrayLike) -> Contact:
        """Project radially to the fitted sphere."""

        value = _array(point, (3,), "point")
        radial = value - self.center
        length = float(np.linalg.norm(radial))
        normal = np.array((1.0, 0.0, 0.0)) if length <= 1.0e-12 else radial / length
        return Contact(
            self.center + self.radius_m * normal,
            normal,
            abs(length - self.radius_m),
            self.confidence,
        )

    def sample(self, count: int) -> tuple[Contact, ...]:
        """Sample the sphere with a deterministic Fibonacci lattice."""

        count = _count(count, 4, "sphere sample count")
        index = np.arange(count, dtype=np.float64)
        z = 1.0 - 2.0 * (index + 0.5) / count
        angle = np.pi * (3.0 - np.sqrt(5.0)) * index
        radius = np.sqrt(np.maximum(1.0 - z * z, 0.0))
        normals = np.column_stack((radius * np.cos(angle), radius * np.sin(angle), z))
        return tuple(
            Contact(self.center + self.radius_m * normal, normal, 0.0, self.confidence)
            for normal in normals
        )

    def distance(self, point: npt.ArrayLike) -> float:
        """Return signed radial distance to the sphere."""

        return float(self.distances(np.asarray(point, dtype=np.float64)[None, :])[0])

    def distances(self, points: npt.ArrayLike) -> FloatArray:
        values = np.asarray(points, dtype=np.float64)
        if values.ndim != 2 or values.shape[1:] != (3,) or not bool(np.all(np.isfinite(values))):
            raise ValueError("points must be finite with shape (N, 3)")
        return cast(FloatArray, np.linalg.norm(values - self.center, axis=1) - self.radius_m)


@dataclass(frozen=True, slots=True)
class Cylinder:
    """Oriented cylinder surface in the hand frame; local Z is its axis."""

    center: FloatArray
    rotation: FloatArray
    radius_m: float
    half_height_m: float
    confidence: float = 1.0

    def __post_init__(self) -> None:
        center = _array(self.center, (3,), "center")
        rotation = _rotation(self.rotation)
        values = np.array((self.radius_m, self.half_height_m, self.confidence))
        if not bool(np.all(np.isfinite(values))) or self.radius_m <= 0.0 or self.half_height_m <= 0.0:
            raise ValueError("cylinder dimensions must be finite and positive")
        if not 0.0 <= self.confidence <= 1.0:
            raise ValueError("confidence must be in [0, 1]")
        object.__setattr__(self, "center", center)
        object.__setattr__(self, "rotation", rotation)

    @property
    def shape(self) -> Shape:
        """Return cylinder."""

        return Shape.CYLINDER

    def _local(self, point: npt.ArrayLike) -> FloatArray:
        return (_array(point, (3,), "point") - self.center) @ self.rotation

    def _hand(self, point: FloatArray) -> FloatArray:
        return point @ self.rotation.T + self.center

    def project(self, point: npt.ArrayLike) -> Contact:
        """Project to the closest side or cap."""

        local = self._local(point)
        radial = local[:2]
        length = float(np.linalg.norm(radial))
        direction = np.array((1.0, 0.0)) if length <= 1.0e-12 else radial / length
        side = np.array(
            (self.radius_m * direction[0], self.radius_m * direction[1], np.clip(local[2], -self.half_height_m, self.half_height_m)),
            dtype=np.float64,
        )
        cap_radius = min(length, self.radius_m)
        cap_xy = cap_radius * direction
        lower = np.array((cap_xy[0], cap_xy[1], -self.half_height_m))
        upper = np.array((cap_xy[0], cap_xy[1], self.half_height_m))
        candidates = np.stack((side, lower, upper))
        selected = int(np.argmin(np.linalg.norm(candidates - local, axis=1)))
        normals = np.array(
            ((direction[0], direction[1], 0.0), (0.0, 0.0, -1.0), (0.0, 0.0, 1.0)),
            dtype=np.float64,
        )
        return Contact(
            self._hand(candidates[selected]),
            cast(FloatArray, normals[selected] @ self.rotation.T),
            float(np.linalg.norm(candidates[selected] - local)),
            self.confidence,
        )

    def sample(self, count: int) -> tuple[Contact, ...]:
        """Sample the side and caps."""

        count = _count(count, 8, "cylinder sample count")
        side_count = max(6, int(0.8 * count))
        angles = np.linspace(0.0, 2.0 * np.pi, side_count, endpoint=False)
        axial_levels = max(2, int(np.ceil(side_count / 16.0)))
        z_values = np.linspace(-0.75, 0.75, axial_levels) * self.half_height_m
        values: list[Contact] = []
        for index, angle in enumerate(angles):
            normal = np.array((np.cos(angle), np.sin(angle), 0.0))
            z = z_values[index % axial_levels]
            local = np.array((self.radius_m * normal[0], self.radius_m * normal[1], z))
            values.append(
                Contact(
                    self._hand(local),
                    cast(FloatArray, normal @ self.rotation.T),
                    0.0,
                    self.confidence,
                )
            )
        remaining = count - len(values)
        for index in range(remaining):
            angle = 2.0 * np.pi * index / max(remaining, 1)
            radial = 0.65 * self.radius_m * np.array((np.cos(angle), np.sin(angle)))
            sign = -1.0 if index % 2 == 0 else 1.0
            local = np.array((radial[0], radial[1], sign * self.half_height_m))
            normal = np.array((0.0, 0.0, sign))
            values.append(
                Contact(
                    self._hand(local),
                    cast(FloatArray, normal @ self.rotation.T),
                    0.0,
                    self.confidence,
                )
            )
        return tuple(values[:count])

    def distance(self, point: npt.ArrayLike) -> float:
        """Return signed distance to side or cap."""

        return float(self.distances(np.asarray(point, dtype=np.float64)[None, :])[0])

    def distances(self, points: npt.ArrayLike) -> FloatArray:
        values = np.asarray(points, dtype=np.float64)
        if values.ndim != 2 or values.shape[1:] != (3,) or not bool(np.all(np.isfinite(values))):
            raise ValueError("points must be finite with shape (N, 3)")
        local = (values - self.center) @ self.rotation
        delta = np.column_stack(
            (
                np.linalg.norm(local[:, :2], axis=1) - self.radius_m,
                np.abs(local[:, 2]) - self.half_height_m,
            )
        )
        return cast(
            FloatArray,
            np.linalg.norm(np.maximum(delta, 0.0), axis=1)
            + np.minimum(np.max(delta, axis=1), 0.0),
        )


@dataclass(frozen=True, slots=True)
class Cloud:
    """Observed point-cloud surface with local PCA normals."""

    points: FloatArray
    camera: FloatArray
    neighbor_count: int = 16
    confidence: float = 1.0
    _normals: FloatArray = field(init=False, repr=False)
    _confidence: FloatArray = field(init=False, repr=False)
    _support_radius: FloatArray = field(init=False, repr=False)
    _tree: cKDTree = field(init=False, repr=False, compare=False)

    def __post_init__(self) -> None:
        points = np.array(self.points, dtype=np.float64, copy=True)
        if points.ndim != 2 or points.shape[1:] != (3,) or points.shape[0] < 32:
            raise ValueError("points must have shape (N, 3), N >= 32")
        if not bool(np.all(np.isfinite(points))):
            raise ValueError("points must be finite")
        camera = _array(self.camera, (3,), "camera")
        if (
            isinstance(self.neighbor_count, bool)
            or not isinstance(self.neighbor_count, int)
            or self.neighbor_count < 8
            or self.neighbor_count > points.shape[0]
        ):
            raise ValueError("neighbor_count is outside the point count")
        if not np.isfinite(self.confidence) or not 0.0 <= self.confidence <= 1.0:
            raise ValueError("confidence must be in [0, 1]")
        if points.shape[0] > 1024:
            index = np.linspace(0, points.shape[0] - 1, 1024, dtype=np.intp)
            points = points[index]
        tree = cKDTree(points)
        neighbor_distance, neighbors = tree.query(
            points,
            k=min(self.neighbor_count, points.shape[0]),
        )
        local = points[neighbors]
        centered = local - np.mean(local, axis=1, keepdims=True)
        covariance = np.einsum("nki,nkj->nij", centered, centered) / max(local.shape[1] - 1, 1)
        eigenvalues, eigenvectors = np.linalg.eigh(covariance)
        normals = eigenvectors[:, :, 0]
        view = camera - points
        normals *= np.where(np.sum(normals * view, axis=1) >= 0.0, 1.0, -1.0)[:, None]
        confidence = self.confidence * np.clip(
            1.0
            - eigenvalues[:, 0]
            / np.maximum(np.sum(eigenvalues, axis=1), 1.0e-12),
            0.0,
            1.0,
        )
        distances = np.asarray(neighbor_distance, dtype=np.float64)
        support_radius = np.maximum(distances[:, -1], 1.0e-9)
        points.flags.writeable = False
        normals.flags.writeable = False
        confidence.flags.writeable = False
        support_radius.flags.writeable = False
        object.__setattr__(self, "points", points)
        object.__setattr__(self, "camera", camera)
        object.__setattr__(self, "_normals", normals)
        object.__setattr__(self, "_confidence", confidence)
        object.__setattr__(self, "_support_radius", support_radius)
        object.__setattr__(self, "_tree", tree)

    @property
    def shape(self) -> Shape:
        """Return point_cloud."""

        return Shape.CLOUD

    @property
    def center(self) -> FloatArray:
        """Return the visible-cloud centroid."""

        return cast(FloatArray, np.mean(self.points, axis=0))

    def project(self, point: npt.ArrayLike) -> Contact:
        """Project to the nearest observed cloud sample."""

        value = _array(point, (3,), "point")
        distance, index = self._tree.query(value, k=1)
        selected = int(index)
        return Contact(
            self.points[selected],
            self._normals[selected],
            float(distance),
            float(self._confidence[selected]),
        )

    def sample(self, count: int) -> tuple[Contact, ...]:
        """Return spatially separated observed samples."""

        count = min(_count(count, 2, "cloud sample count"), self.points.shape[0])
        selected = np.empty(count, dtype=np.intp)
        available = np.ones(self.points.shape[0], dtype=np.bool_)
        selected[0] = int(np.argmax(np.linalg.norm(self.points - self.center, axis=1)))
        available[selected[0]] = False
        nearest = np.linalg.norm(self.points - self.points[selected[0]], axis=1)
        for index in range(1, count):
            score = nearest * np.maximum(self._confidence, 1.0e-12)
            score[~available] = -np.inf
            selected[index] = int(np.argmax(score))
            available[selected[index]] = False
            nearest = np.minimum(
                nearest,
                np.linalg.norm(self.points - self.points[selected[index]], axis=1),
            )
        return tuple(
            Contact(self.points[index], self._normals[index], 0.0, float(self._confidence[index]))
            for index in selected
        )

    def distance(self, point: npt.ArrayLike) -> float:
        """Return local signed point-to-plane distance."""

        return float(self.distances(np.asarray(point, dtype=np.float64)[None, :])[0])

    def distances(self, points: npt.ArrayLike) -> FloatArray:
        values = np.asarray(points, dtype=np.float64)
        if values.ndim != 2 or values.shape[1:] != (3,) or not bool(np.all(np.isfinite(values))):
            raise ValueError("points must be finite with shape (N, 3)")
        _, index = self._tree.query(values, k=1)
        selected = np.asarray(index, dtype=np.intp)
        delta = values - self.points[selected]
        normal_distance = np.sum(delta * self._normals[selected], axis=1)
        tangent = delta - normal_distance[:, None] * self._normals[selected]
        support_excess = np.maximum(
            np.linalg.norm(tangent, axis=1) - self._support_radius[selected],
            0.0,
        )
        return cast(
            FloatArray,
            normal_distance - support_excess,
        )


def build_surface(
    shape: Shape,
    center: npt.ArrayLike,
    rotation: npt.ArrayLike,
    dimensions_m: npt.ArrayLike,
    points: npt.ArrayLike,
    camera: npt.ArrayLike,
    confidence: float,
) -> Surface:
    """Construct a hand-frame surface from one RGB-D observation."""

    if not isinstance(shape, Shape):
        raise ValueError("shape must be a Shape")
    dimensions = _array(dimensions_m, (3,), "dimensions_m")
    if shape is Shape.BOX:
        return Box(_array(center, (3,), "center"), _rotation(rotation), dimensions, confidence)
    if shape is Shape.SPHERE:
        return Sphere(_array(center, (3,), "center"), 0.5 * float(dimensions[0]), confidence)
    if shape is Shape.CYLINDER:
        return Cylinder(
            _array(center, (3,), "center"),
            _rotation(rotation),
            0.25 * float(dimensions[0] + dimensions[1]),
            0.5 * float(dimensions[2]),
            confidence,
        )
    if shape is Shape.CLOUD:
        return Cloud(
            np.asarray(points, dtype=np.float64),
            _array(camera, (3,), "camera"),
            confidence=confidence,
        )
    raise ValueError(f"unsupported shape: {shape}")


__all__ = ["Box", "Cloud", "Contact", "Cylinder", "Shape", "Sphere", "Surface", "build_surface"]