massive-cloud/alg/delivery_v1.2.1/underframe_deflection/cad_io.py

"""CAD reference loader.

Accepts pre-tessellated CAD geometry — either a sampled point cloud
(``.ply / .xyz / .txt / .csv``) or a triangle mesh (``.stl`` ASCII or
binary, ``.obj``). For meshes, only the vertex positions are kept; face
connectivity is discarded because the downstream registration is
point-to-point ICP.

STEP / IGES files are not parsed here — convert them offline to STL or a
sampled PLY first.
"""

from __future__ import annotations

import struct
from pathlib import Path

import numpy as np

from .exceptions import InvalidPointCloudError, UnsupportedFormatError
from .io import _read_ascii_ply


_SUPPORTED_SUFFIXES = {".ply", ".xyz", ".txt", ".csv", ".stl", ".obj"}


def read_cad_reference(path: str | Path) -> np.ndarray:
    """Load a CAD reference and return an ``(N, 3) float64`` array.

    The result is a unique set of vertex / point coordinates in
    millimetres. Faces / normals / colors in the source file are
    ignored.
    """
    path = Path(path)
    suffix = path.suffix.lower()
    if suffix not in _SUPPORTED_SUFFIXES:
        raise UnsupportedFormatError(
            f"Unsupported CAD reference format: {path.suffix} "
            f"(supported: {sorted(_SUPPORTED_SUFFIXES)})"
        )
    if suffix == ".ply":
        points = _read_ascii_ply(path)
    elif suffix in {".xyz", ".txt", ".csv"}:
        try:
            data = np.loadtxt(path, dtype=np.float64)
        except Exception as exc:  # numpy raises a few different things
            raise InvalidPointCloudError(
                f"Failed to read CAD reference text file {path}: {exc}"
            ) from exc
        if data.ndim == 1:
            data = data[None, :]
        if data.shape[1] < 3:
            raise InvalidPointCloudError(
                f"CAD reference {path} does not contain XYZ columns."
            )
        points = data[:, :3]
    elif suffix == ".stl":
        points = _read_stl_vertices(path)
    elif suffix == ".obj":
        points = _read_obj_vertices(path)
    else:  # pragma: no cover - guarded above
        raise UnsupportedFormatError(f"Unsupported CAD format: {suffix}")

    if points.ndim != 2 or points.shape[1] != 3 or len(points) == 0:
        raise InvalidPointCloudError(
            f"CAD reference {path} produced an empty or malformed point set."
        )
    return np.ascontiguousarray(points, dtype=np.float64)


# ---------------------------------------------------------------------------
# STL / OBJ readers
# ---------------------------------------------------------------------------


def _read_stl_vertices(path: Path) -> np.ndarray:
    """Read vertex positions from an STL file (ASCII or binary)."""
    # ASCII STL must start with "solid" and contain "facet" keywords.
    # Binary STL has an 80-byte header followed by a uint32 triangle
    # count, so we sniff the first 6 bytes for "solid " AND verify by
    # scanning a few more lines for "facet" to avoid false positives.
    with path.open("rb") as fh:
        head = fh.read(256)
    if head[:5].lower() == b"solid" and b"facet" in head.lower():
        return _read_stl_ascii(path)
    return _read_stl_binary(path)


def _read_stl_ascii(path: Path) -> np.ndarray:
    vertices: list[tuple[float, float, float]] = []
    try:
        with path.open("r", encoding="utf-8", errors="replace") as fh:
            for line in fh:
                stripped = line.strip()
                if stripped.startswith("vertex"):
                    parts = stripped.split()
                    if len(parts) >= 4:
                        vertices.append(
                            (float(parts[1]), float(parts[2]), float(parts[3]))
                        )
    except (OSError, ValueError) as exc:
        raise InvalidPointCloudError(f"Malformed ASCII STL {path}: {exc}") from exc
    if not vertices:
        raise InvalidPointCloudError(f"ASCII STL {path} contains no vertices.")
    return _unique_rows(np.asarray(vertices, dtype=np.float64))


def _read_stl_binary(path: Path) -> np.ndarray:
    try:
        with path.open("rb") as fh:
            fh.read(80)  # header
            count_bytes = fh.read(4)
            if len(count_bytes) < 4:
                raise InvalidPointCloudError(
                    f"Binary STL {path} truncated before triangle count."
                )
            (triangle_count,) = struct.unpack("<I", count_bytes)
            if triangle_count == 0:
                raise InvalidPointCloudError(f"Binary STL {path} has zero triangles.")
            payload = fh.read(triangle_count * 50)
            if len(payload) < triangle_count * 50:
                raise InvalidPointCloudError(
                    f"Binary STL {path} truncated (expected "
                    f"{triangle_count * 50} bytes, got {len(payload)})."
                )
    except (OSError, struct.error) as exc:
        raise InvalidPointCloudError(f"Failed to read binary STL {path}: {exc}") from exc

    # Each 50-byte record: 12 floats (normal + 3 vertices) + 2-byte attr.
    arr = np.frombuffer(payload, dtype=np.uint8).reshape(triangle_count, 50)
    # vertices occupy bytes 12..48 -> 36 bytes -> 9 floats.
    floats = np.frombuffer(arr[:, 12:48].tobytes(), dtype="<f4").reshape(-1, 3)
    if not np.isfinite(floats).all():
        raise InvalidPointCloudError(f"Binary STL {path} contains non-finite values.")
    return _unique_rows(floats.astype(np.float64))


def _read_obj_vertices(path: Path) -> np.ndarray:
    vertices: list[tuple[float, float, float]] = []
    try:
        with path.open("r", encoding="utf-8", errors="replace") as fh:
            for line in fh:
                if not line.startswith("v "):
                    continue
                parts = line.split()
                if len(parts) >= 4:
                    vertices.append(
                        (float(parts[1]), float(parts[2]), float(parts[3]))
                    )
    except (OSError, ValueError) as exc:
        raise InvalidPointCloudError(f"Malformed OBJ {path}: {exc}") from exc
    if not vertices:
        raise InvalidPointCloudError(f"OBJ {path} contains no vertex positions.")
    return np.asarray(vertices, dtype=np.float64)


def _unique_rows(points: np.ndarray) -> np.ndarray:
    if len(points) == 0:
        return points
    # Vertices from triangle meshes have ~3x duplication; collapse them.
    return np.unique(points, axis=0)


__all__ = ["read_cad_reference"]