primlib 0.3.0__py3-none-any.whl

primlib/utils.py

@@ -0,0 +1,53 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+def scalar_to_color(
+    scalars: ArrayLike,
+    cmap: str = "viridis",
+    min_max: tuple[float, float] | None = None,
+    nan_color: tuple[float, float, float] = (1, 0, 0),
+) -> np.ndarray:
+    scalars = np.asarray(scalars)
+    mask_valid = ~np.isnan(scalars)
+    if min_max is None:
+        min_max = (np.min(scalars[mask_valid]), np.max(scalars[mask_valid]))
+    scalars_norm = np.clip(
+        (scalars[mask_valid] - min_max[0]) / (min_max[1] - min_max[0]), 0, 1
+    )
+    cm = plt.get_cmap(cmap)
+    colors = np.empty((scalars.size, 3))
+    colors[mask_valid, :] = cm(scalars_norm)[:, 0:3]
+    colors[~mask_valid, :] = nan_color
+    return colors
+
+
+def scalar_to_bicolor(
+    scalars: ArrayLike,
+    cmaps: tuple[str, str] = ("Blues", "Reds"),
+    vmax: float | None = None,
+    nan_color: tuple[float, float, float] = (0, 1, 0),
+) -> np.ndarray:
+    scalars = np.asarray(scalars)
+    mask_valid = ~np.isnan(scalars)
+    mask_pos = scalars >= 0
+    mask_neg = ~mask_pos
+    neg = scalars[mask_valid & mask_neg]
+    pos = scalars[mask_valid & mask_pos]
+    if vmax is None:
+        if neg.size == 0:
+            vmax = np.max(pos)
+        elif pos.size == 0:
+            vmax = np.max(-neg)
+        else:
+            vmax = np.max((np.max(pos), np.max(-neg)))
+    scalars_norm_neg = np.clip(-neg / vmax, 0, 1)
+    scalars_norm_pos = np.clip(pos / vmax, 0, 1)
+    cm_neg = plt.get_cmap(cmaps[0])
+    cm_pos = plt.get_cmap(cmaps[1])
+    colors = np.empty((scalars.size, 3))
+    colors[mask_valid & mask_neg, :] = cm_neg(scalars_norm_neg)[:, 0:3]
+    colors[mask_valid & mask_pos, :] = cm_pos(scalars_norm_pos)[:, 0:3]
+    colors[~mask_valid, :] = nan_color
+    return colors

primlib/vertices.py

@@ -0,0 +1,604 @@
+import copy
+import logging
+import os
+from typing import Any, Self
+
+import numpy as np
+from numpy.typing import ArrayLike
+from plyfile import PlyData, PlyElement
+from scipy.spatial import KDTree
+from trafolib.trafo3d import Trafo3d
+
+from primlib.utils import scalar_to_bicolor, scalar_to_color
+
+
+class Vertices:
+    """
+    A class representing a set of vertices, normals and colors
+
+    Can be used as a base class for point clouds, line sets or triangle meshes.
+    """
+
+    __slots__ = ("_colors", "_normals", "_vertices")
+
+    def __init__(
+        self,
+        vertices: ArrayLike | None = None,
+        normals: ArrayLike | None = None,
+        colors: ArrayLike | None = None,
+    ) -> None:
+        # Define types and initialize
+        self._vertices: np.ndarray = np.zeros((0, 3))
+        self._normals: np.ndarray = np.zeros((0, 3))
+        self._colors: np.ndarray = np.zeros((0, 3))
+        # Initialize using parameters and setters
+        self.vertices = vertices
+        self.normals = normals
+        self.colors = colors
+
+    def __array__(
+        self,
+        dtype: Any = None,
+        copy: bool | None = None,
+    ) -> np.ndarray:
+        """Convert vertices to a NumPy array.
+
+        :param dtype: Desired data type of the resulting array.
+            If None, the data type is inferred from the vertices.
+        :param copy: Whether to return a copy of the array.
+            If None, the copy behaviour is determined by NumPy.
+        :returns: NumPy array containing the vertices.
+        """
+        if copy is True:
+            return np.array(self._vertices, dtype=dtype, copy=True)
+        return np.asarray(self._vertices, dtype=dtype)
+
+    def __getitem__(self, item: Any) -> np.ndarray:
+        """Get vertices by index or slice.
+
+        Supports integer indexing, slicing, and advanced NumPy indexing.
+
+        :param item: Index, slice, or array of indices used to select vertices.
+        :returns: Selected vertex or vertices as a NumPy array.
+        """
+        return self._vertices[item]
+
+    def __eq__(self, other: object) -> bool:
+        """Check two vertices sets for equality.
+
+        :param other: Other object
+        :returns: True if vertices sets are equal
+        """
+        if not isinstance(other, Vertices):
+            return NotImplemented
+        return (
+            # Vertices
+            self._vertices.shape == other._vertices.shape
+            and np.allclose(self._vertices, other._vertices)
+            # Normals
+            and self._normals.shape == other._normals.shape
+            and np.allclose(self._normals, other._normals)
+            # Colors
+            and self._colors.shape == other._colors.shape
+            and np.allclose(
+                self._colors,
+                other._colors,
+                atol=1.0 / 255.0,
+            )
+        )
+
+    def __ne__(self, other: object) -> bool:
+        """Check two vertices sets for inequality
+
+        :param other: Other object
+        :returns: True, if vertices sets are not equal
+        """
+        return not self.__eq__(other)
+
+    def __iadd__(self, other: "Vertices") -> Self:
+        """Concatenate another vertices set in-place
+
+        :param other: Other vertices set
+        :returns: This vertices set (self)
+        """
+        if not isinstance(other, Vertices):
+            return NotImplemented
+        if self.has_normals != other.has_normals or self.has_colors != other.has_colors:
+            raise ValueError("Both vertices sets must have the same properties")
+        self._vertices = np.vstack((self._vertices, other._vertices))
+        if self.has_normals:
+            self._normals = np.vstack((self._normals, other._normals))
+        if self.has_colors:
+            self._colors = np.vstack((self._colors, other._colors))
+        return self
+
+    def __add__(self, other: "Vertices") -> Self:
+        """Concatenate two vertices sets
+
+        :param other: Other vertices set
+        :returns: New vertices set
+        """
+        if not isinstance(other, Vertices):
+            return NotImplemented
+        result = self.copy()
+        result += other
+        return result
+
+    def copy(self) -> Self:
+        """Get copy of vertices
+
+        :returns: Copy
+        """
+        return copy.deepcopy(self)
+
+    def __copy__(self) -> Self:
+        """Return a shallow copy of the object
+
+        Required by the `copy` module to implement custom shallow copy behavior.
+
+        :returns: A new instance containing references to the original
+            instance's attributes.
+        """
+        cls = type(self)
+        result = cls.__new__(cls)
+        for base in cls.__mro__:
+            slots = getattr(base, "__slots__", ())
+            if isinstance(slots, str):
+                slots = (slots,)
+            for name in slots:
+                if name in ("__dict__", "__weakref__"):
+                    continue
+                if hasattr(self, name):
+                    setattr(result, name, copy.copy(getattr(self, name)))
+        return result
+
+    def __deepcopy__(self, memo) -> "Vertices":
+        """Return a deep copy of the object
+
+        Required by the `copy` module to implement custom deep copy behavior.
+
+        :param memo: A dictionary tracking already copied objects to prevent
+            infinite loops from recursive references
+        :returns: A new, completely independent copy of the instance.
+        """
+        cls = type(self)
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for base in cls.__mro__:
+            slots = getattr(base, "__slots__", ())
+            if isinstance(slots, str):
+                slots = (slots,)
+            for name in slots:
+                if name in ("__dict__", "__weakref__"):
+                    continue
+                if hasattr(self, name):
+                    setattr(result, name, copy.deepcopy(getattr(self, name), memo))
+        return result
+
+    @property
+    def vertices(self) -> np.ndarray:
+        """Get vertices
+
+        Vertices are points in 3d.
+
+        :returns: Vertices, shape (n, 3), type float
+        """
+        return self._vertices
+
+    @vertices.setter
+    def vertices(self, value: ArrayLike | None) -> None:
+        """Set vertices
+
+        Vertices are points in 3d.
+
+        If there are normals or colors defined in the vertices set,
+        then the number of vertices has to match their number.
+
+        :param value: Vertices, shape (n, 3), type float
+        """
+        if value is None:
+            value = np.zeros((0, 3))
+        else:
+            value = np.asarray(value, dtype=float)
+        if value.ndim != 2 or value.shape[1] != 3:
+            raise Exception("Provide vertices of shape (n, 3)")
+        if self.has_normals and value.shape[0] != self._normals.shape[0]:
+            raise Exception(
+                "vertices set already has normals, provide proper amount of vertices"
+            )
+        if self.has_colors and value.shape[0] != self._colors.shape[0]:
+            raise Exception(
+                "vertices set already has colors, provide proper amount of vertices"
+            )
+        self._vertices = value
+
+    @property
+    def normals(self) -> np.ndarray:
+        """Get normals
+
+        Normals are 3d vectors that are perpendicular to neighborhood of the
+        vertex. For a point cloud that can mean perpendicular to a plane
+        sampled by the vertex and its neighbors. For a triangle mesh the
+        vertex normal can be the mean of the face normals of all faces
+        adjacent to the vector. Normals are normalized, i.e. they have a
+        length of 1.0.
+
+        :returns: Normals, shape (n, 3)
+        """
+        return self._normals
+
+    @normals.setter
+    def normals(self, value: ArrayLike | None) -> None:
+        """Set normals
+
+        If there are vertices defined in the vertices set, then the number
+        of normals has to match the number of vertices.
+
+        Normals should have a length > 0, otherwise an exception is
+        thrown. Normals should have unit length. If not, normals will
+        be normalized and a warning is given.
+
+        :param value: Normals, shape (n, 3)
+        """
+        if value is None:
+            value = np.zeros((0, 3))
+        else:
+            value = np.asarray(value, dtype=float)
+        if value.ndim != 2 or value.shape[1] != 3:
+            raise Exception("Provide normals of shape (n, 3)")
+        if value.shape[0] > 0 and value.shape[0] != self._vertices.shape[0]:
+            raise Exception("Provide as many normals as vertices")
+        lengths = np.sqrt(value[:, 0] ** 2 + value[:, 1] ** 2 + value[:, 2] ** 2)
+        if np.any(np.isclose(lengths, 0.0)):
+            raise Exception("Some normals have zero lengths")
+        if not np.allclose(lengths, 1.0):
+            logging.warning("Some normals are not normalized. Normalizing...")
+            value = value / lengths[:, None]
+        self._normals = value
+
+    @property
+    def has_normals(self) -> bool:
+        """Check if vertices set has normals
+
+        :returns: True, if vertices set has normals
+        """
+        return self._normals.shape[0] > 0
+
+    @property
+    def colors(self) -> np.ndarray:
+        """Get colors
+
+        Color are 3-tuples of RGB values in range [0..1] of type float.
+
+        :returns: Colors
+        """
+        return self._colors
+
+    @colors.setter
+    def colors(self, value: ArrayLike | None) -> None:
+        """Set colors
+
+        If there are vertices defined in the vertices set, then the number
+        of colors has to match the number of the vertices.
+
+        Colors are RGB values in range [0..1]. If values outside that
+        range are provided, values will be clipped and a warning is given.
+
+        :param value: Colors, shape (n, 3)
+        """
+        if value is None:
+            colors = np.zeros((0, 3))
+        else:
+            colors = np.asarray(value, dtype=float)
+        if colors.ndim != 2 or colors.shape[1] != 3:
+            raise ValueError("Provide colors of shape (n, 3)")
+        if colors.shape[0] > 0 and colors.shape[0] != self._vertices.shape[0]:
+            raise ValueError("Provide as many colors as vertices")
+        if not np.all(np.logical_and(colors >= 0.0, colors <= 1.0)):
+            logging.warning("Some colors are out of bounds. Clipping to [0,1] ...")
+            colors = np.clip(colors, 0.0, 1.0)
+        self._colors = colors
+
+    @property
+    def has_colors(self) -> bool:
+        """Check if vertices set has colors
+
+        :returns: True, if vertices set has colors
+        """
+        return self._colors.shape[0] > 0
+
+    @staticmethod
+    def _read_ply_data(path: str | os.PathLike) -> PlyData:
+        """Helper method to read PlyData from file
+
+        :param path: Path to file to read from
+        :returns: PlyData
+        """
+        try:
+            return PlyData.read(path)
+        except Exception as ex:
+            raise Exception(str(type(ex))) from ex
+
+    @staticmethod
+    def _write_ply_data(path: str | os.PathLike, elements: list[PlyElement]):
+        """Helper method to write PlyData to file
+
+        :param path: Path to file to write to
+        :param elements: PlyData
+        """
+        plydata = PlyData(elements, text=False)
+        try:
+            plydata.write(path)
+        except Exception as ex:
+            raise Exception(str(type(ex))) from ex
+
+    @staticmethod
+    def _load_vertex_data(
+        plydata: PlyData,
+    ) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None]:
+        """Helper method to read vertex data
+
+        :param plydata: PlyData
+        :returns: Vertices, normals and colors
+        """
+        vertex_data = plydata["vertex"]
+        properties = {p.name: p.val_dtype for p in vertex_data.properties}
+        # Vertices
+        vertices = (
+            np.vstack([vertex_data["x"], vertex_data["y"], vertex_data["z"]])
+            .astype(float)
+            .T
+        )
+        # Normals (if existing)
+        normals = None
+        normal_keys = ["nx", "ny", "nz"]
+        if all(k in properties for k in normal_keys):
+            normals = (
+                np.vstack([vertex_data["nx"], vertex_data["ny"], vertex_data["nz"]])
+                .astype(float)
+                .T
+            )
+        # Colors (if existing)
+        colors = None
+        color_keys_rgb = ["red", "green", "blue"]
+        if all(k in properties for k in color_keys_rgb):
+            if properties["red"] == "u1":
+                colors = np.vstack(
+                    [vertex_data["red"], vertex_data["green"], vertex_data["blue"]]
+                ).T
+                colors = colors.astype(float) / 255.0
+            elif properties["red"] == "f4":
+                colors = np.vstack(
+                    [vertex_data["red"], vertex_data["green"], vertex_data["blue"]]
+                ).T
+            else:
+                raise Exception(f"Unknown color format: {properties['red']}")
+        return vertices, normals, colors
+
+    def _save_vertex_data(self) -> PlyElement:
+        """Helper method to save vertex data"""
+        # Collect dtypes
+        dtypes = [("x", "f8"), ("y", "f8"), ("z", "f8")]
+        if self.has_normals:
+            dtypes.extend([("nx", "f8"), ("ny", "f8"), ("nz", "f8")])
+        if self.has_colors:
+            dtypes.extend([("red", "u1"), ("green", "u1"), ("blue", "u1")])
+        # Collect data
+        element_data = np.empty(self._vertices.shape[0], dtype=dtypes)
+        element_data["x"] = self._vertices[:, 0]
+        element_data["y"] = self._vertices[:, 1]
+        element_data["z"] = self._vertices[:, 2]
+        if self.has_normals:
+            element_data["nx"] = self._normals[:, 0]
+            element_data["ny"] = self._normals[:, 1]
+            element_data["nz"] = self._normals[:, 2]
+        if self.has_colors:
+            element_data["red"] = np.round(self._colors[:, 0] * 255.0).astype(np.uint8)
+            element_data["green"] = np.round(self._colors[:, 1] * 255.0).astype(
+                np.uint8
+            )
+            element_data["blue"] = np.round(self._colors[:, 2] * 255.0).astype(np.uint8)
+        return PlyElement.describe(element_data, "vertex")
+
+    @classmethod
+    def from_ply(cls, path: str | os.PathLike) -> "Vertices":
+        """Create vertices set from PLY file
+
+        :param path: Path to PLY file to read
+        :returns: Resulting vertices set
+        """
+        plydata = Vertices._read_ply_data(path)
+        vertices, normals, colors = Vertices._load_vertex_data(plydata)
+        return cls(vertices, normals, colors)
+
+    def save_ply(self, path: str | os.PathLike) -> None:
+        """Save vertices set to PLY file
+
+        :param path: Path to PLY file to write
+        """
+        vertex_element = self._save_vertex_data()
+        Vertices._write_ply_data(path, [vertex_element])
+
+    def min(self) -> np.ndarray:
+        """Get minimum of all coordinates
+
+        :raises Exception: Raised if vertices set empty
+        :returns: Minimum of all coordinates, shape (3,)
+        """
+        if self._vertices.shape[0] == 0:
+            raise Exception("Cannot compute min for empty vertices set")
+        return np.min(self._vertices, axis=0)
+
+    def max(self) -> np.ndarray:
+        """Get maximum of all coordinates
+
+        :raises Exception: Raised if vertices set empty
+        :returns: Maximum of all coordinates, shape (3,)
+        """
+        if self._vertices.shape[0] == 0:
+            raise Exception("Cannot compute max for empty vertices set")
+        return np.max(self._vertices, axis=0)
+
+    def bounds(self) -> np.ndarray:
+        """Get bounds (min and max) of vertices set
+
+        :raises Exception: Raised if vertices set empty
+        :returns: Min and max values for all coordinates, shape (2, 3)
+        """
+        if self._vertices.shape[0] == 0:
+            raise Exception("Cannot compute bounds for empty vertices set")
+        return np.vstack(
+            (np.min(self._vertices, axis=0), np.max(self._vertices, axis=0))
+        )
+
+    def range(self) -> np.ndarray:
+        """Get ranges for all coordinates
+
+        :raises Exception: Raised if vertices set empty
+        :returns: Ranges (max-min) for all coordinates
+        """
+        if self._vertices.shape[0] == 0:
+            raise Exception("Cannot compute range for empty vertices set")
+        return self.max() - self.min()
+
+    def mean(self) -> np.ndarray:
+        """Get mean for all coordinates
+        :raises Exception: Raised if vertices set empty
+        :returns: Mean for all coordinates
+        """
+        if self._vertices.shape[0] == 0:
+            raise Exception("Cannot compute mean for empty vertices set")
+        return np.mean(self._vertices, axis=0)
+
+    def center(self) -> np.ndarray:
+        """Get center of gravity
+
+        Same as self.mean.
+
+        :raises Exception: Raised if vertices set empty
+        :returns: Center of gravity
+        """
+        return self.mean()
+
+    def demean(self) -> None:
+        """Translate vertices set that center of gravity is at (0, 0, 0)."""
+        if self._vertices.shape[0] == 0:
+            return
+        self._vertices -= np.mean(self._vertices, axis=0)
+
+    def scale(self, factor: float | ArrayLike, center: ArrayLike | None = None) -> None:
+        """Scale vertices
+
+        :param factor: Scaling factor, either scalar or shape (3,)
+        :param center: Center, shape (3,)
+        """
+        if center is None:
+            self._vertices = np.asarray(factor) * self._vertices
+        else:
+            center = np.asarray(center)
+            if not (center.ndim == 1 and center.size == 3):
+                raise Exception("Provide center of shape (3, )")
+            self._vertices = (np.asarray(factor) * (self._vertices - center)) + center
+
+    def translate(self, offset: ArrayLike) -> None:
+        """Translate vertices
+
+        :param offset: Translation offset, shape (3,)
+        """
+        offset = np.asarray(offset)
+        if not (offset.ndim == 1 and offset.size == 3):
+            raise Exception("Provide offset of shape (3, )")
+        self._vertices += offset
+
+    def transform(self, trafo: Trafo3d) -> None:
+        """Transform vertices
+
+        :param trafo: 6d transformation
+        """
+        self._vertices = trafo * self._vertices
+        if self.has_normals:
+            self._normals = self._normals @ trafo.rotation_matrix.T
+
+    def kdtree(self) -> KDTree:
+        """Get KDTree of vertices
+
+        Use for fast next-neighbor queries.
+
+        :returns: KDTree of type scipy.spatial.KDTree
+        """
+        return KDTree(self._vertices)
+
+    def filter_region_of_interest(
+        self,
+        bounds: ArrayLike,
+        trafo: Trafo3d | None = None,
+    ) -> np.ndarray:
+        """Filter vertices set
+
+        :param bounds: Min and max values for all coordinates, shape (2, 3)
+        :param trafo: 6d Pose of the bounding box
+        :returns: Inlier mask
+        """
+        bounds = np.array(bounds)
+        if bounds.shape != (2, 3):
+            raise Exception(
+                "Provide bounds as ((xmin, ymin, zmin), "
+                "(xmax, ymax, zmax)), shape (2, 3)"
+            )
+        P = self._vertices
+        if trafo is not None:
+            P = trafo.inverse * P
+        return np.all(P >= bounds[0], axis=1) & np.all(P <= bounds[1], axis=1)
+
+    def paint_uniform_color(self, color: ArrayLike) -> None:
+        """Colorize vertices set with uniform color
+
+        Sets colors to a single user defined color.
+
+        :param color: Color, RGB values in range [0..1], shape (3,), type float
+        """
+        color = np.asarray(color)
+        if not (color.ndim == 1 and color.size == 3):
+            raise Exception("Provide color of shape (3, )")
+        if not np.all(np.logical_and(color >= 0.0, color <= 1.0)):
+            logging.warning("Color are out of bounds. Clipping to [0,1] ...")
+            color = np.clip(color, 0.0, 1.0)
+        self._colors = np.tile(color, (self._vertices.shape[0], 1))
+
+    def colorize_vertices_by_scalar(
+        self,
+        scalars: ArrayLike,
+        cmap: str = "viridis",
+        min_max: tuple[float, float] | None = None,
+        nan_color: tuple[float, float, float] = (1, 0, 0),
+    ) -> None:
+        """Colorizes vertices by scalar
+
+        :param scalars: Scalar values, length has to match number of vertices
+        :param cmap: Matplotlib color map string
+        :param min_max: Minimum and maximum
+        :param nan_color: Color for NaN values, RGB values [0..1]
+        """
+        scalars = np.asarray(scalars)
+        if scalars.ndim != 1 or scalars.size != self._vertices.shape[0]:
+            raise Exception("Provide corrent number of values to colorize vertices")
+        self._colors = scalar_to_color(scalars, cmap, min_max, nan_color)
+
+    def colorize_vertices_by_scalar_bicolor(
+        self,
+        scalars: ArrayLike,
+        cmaps: tuple[str, str] = ("Blues", "Reds"),
+        vmax: float | None = None,
+        nan_color: tuple[float, float, float] = (0, 1, 0),
+    ) -> None:
+        """Colorizes mesh vertices by scalar, pos/neg in separate colors
+
+        :param scalars: Scalar values, length has to match number of vertices
+        :param cmaps: Matplotlib color map strings for pos/neg
+        :param vmax: Maximum (pos/neg) value
+        :param nan_color: Color for NaN values, RGB values [0..1]
+        """
+        scalars = np.asarray(scalars)
+        if scalars.ndim != 1 or scalars.size != self._vertices.shape[0]:
+            raise Exception("Provide corrent number of values to colorize vertices")
+        self._colors = scalar_to_bicolor(scalars, cmaps, vmax, nan_color)