feat: add support for splines (#11)

* feat(wip): add support for splines

* fix: properly load the splines & add color

* feat: update visualization example

* tests: add tests for splines

* fix: weird merge conflicts

* fix: consistent naming

* docs: update readme

Author: HeadTriXz

README.md

@@ -19,17 +19,10 @@ HPS is a compressed 3D mesh format commonly used in dental scanning applications
 - Read mesh geometry (vertices and faces) from HPS files.
 - Supports CA, CC, and CE schemas, including encrypted files.
 - Extract mesh colors, texture coordinates (UVs) and texture images.
+- Extract splines and curves in the scan data.
 - Export the mesh to OBJ, PLY and STL.
 - Command-line tool for exporting meshes.
 
-<details>
-<summary>Planned Features</summary>
-
-- Extract splines and curves in the scan data.
-
-</details>
-
-
 ## Getting Started
 
 ### Installation

examples/view_hps.py

@@ -3,11 +3,13 @@
 import pathlib
 
 import numpy as np
+import numpy.typing as npt
 import trimesh
 import trimesh.visual
 from PIL import Image
 
 from hpsdecode import load_hps
+from hpsdecode.mesh import Spline
 
 
 def convert_bgr_to_rgb(image_data: bytes) -> Image.Image:
@@ -24,6 +26,75 @@ def convert_bgr_to_rgb(image_data: bytes) -> Image.Image:
     return Image.merge("RGB", (b, g, r))
 
 
+def create_cylinder_mesh(
+    start_point: npt.NDArray[np.floating],
+    end_point: npt.NDArray[np.floating],
+    radius: float,
+    slices: int = 32,
+) -> trimesh.Trimesh:
+    """Create a cylinder mesh between two points.
+
+    :param start_point: The starting point of the cylinder.
+    :param end_point: The ending point of the cylinder.
+    :param radius: The radius of the cylinder.
+    :param slices: The number of slices around the cylinder circumference.
+    :return: A Trimesh representing the cylinder.
+    """
+    direction = end_point - start_point
+    height = np.linalg.norm(direction)
+    if height < 1e-6:
+        return trimesh.Trimesh()
+
+    cylinder = trimesh.creation.cylinder(radius=radius, height=height, sections=slices)
+
+    z_axis = np.array([0.0, 0.0, 1.0])
+    direction_normalized = direction / height
+
+    rotation_matrix = trimesh.geometry.align_vectors(z_axis, direction_normalized)
+    cylinder.apply_transform(rotation_matrix)
+
+    center = (start_point + end_point) / 2.0
+    cylinder.apply_translation(center)
+
+    return cylinder
+
+
+def create_spline_mesh(spline: Spline, slices: int = 32) -> trimesh.Trimesh:
+    """Create a mesh representation of a spline as connected cylinders.
+
+    :param spline: The spline to convert to a mesh.
+    :param slices: The number of slices around each cylinder circumference.
+    :return: A trimesh representing the spline.
+    """
+    if spline.num_control_points < 2:
+        return trimesh.Trimesh()
+
+    meshes = []
+    control_points = spline.control_points
+
+    for i in range(len(control_points) - 1):
+        cylinder = create_cylinder_mesh(control_points[i], control_points[i + 1], spline.radius, slices)
+        if cylinder.vertices.size > 0:
+            meshes.append(cylinder)
+
+    if spline.is_cyclic and len(control_points) >= 2:
+        cylinder = create_cylinder_mesh(control_points[-1], control_points[0], spline.radius, slices)
+        if cylinder.vertices.size > 0:
+            meshes.append(cylinder)
+
+    if not meshes:
+        return trimesh.Trimesh()
+
+    r = (spline.color >> 16) & 0xFF
+    g = (spline.color >> 8) & 0xFF
+    b = spline.color & 0xFF
+
+    combined = trimesh.util.concatenate(meshes)
+    combined.visual.vertex_colors = np.array([r, g, b, 255], dtype=np.uint8)
+
+    return combined
+
+
 def create_texture_visual(
     vertices: np.ndarray,
     faces: np.ndarray,
@@ -90,6 +161,11 @@ def main() -> None:
         action="store_true",
         help="Disable texture rendering even if available.",
     )
+    parser.add_argument(
+        "--show-splines",
+        action="store_true",
+        help="Show splines associated with the scan, if available.",
+    )
 
     args = parser.parse_args()
     if not args.input.exists():
@@ -125,7 +201,18 @@ def main() -> None:
             process=False,
         )
 
-    t_mesh.show(caption=args.input.name, smooth=False)
+    scene = trimesh.Scene()
+    scene.add_geometry(t_mesh, node_name="mesh")
+
+    if args.show_splines and mesh.has_splines:
+        print(f"Creating spline visualization ({len(mesh.splines)} splines)...")
+
+        for idx, spline in enumerate(mesh.splines):
+            spline_mesh = create_spline_mesh(spline)
+            if spline_mesh.vertices.size > 0:
+                scene.add_geometry(spline_mesh, node_name=f"spline_{idx}")
+
+    scene.show(caption=args.input.name, smooth=False)
 
 
 if __name__ == "__main__":

src/hpsdecode/export/obj.py

@@ -4,9 +4,9 @@
 
 __all__ = ["MaterialConfig", "OBJExporter"]
 
+import dataclasses
 import io
 import typing as t
-from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
@@ -23,7 +23,7 @@
     from hpsdecode.mesh import HPSMesh
 
 
-@dataclass
+@dataclasses.dataclass
 class MaterialConfig:
     """Material properties for OBJ MTL files."""
 

src/hpsdecode/loader.py

@@ -8,13 +8,17 @@
 import typing as t
 import xml.etree.ElementTree as ET
 
+import numpy as np
+
 from hpsdecode.exceptions import HPSParseError, HPSSchemaError
-from hpsdecode.mesh import HPSMesh, HPSPackedScan, SchemaType
+from hpsdecode.mesh import HPSMesh, HPSPackedScan, SchemaType, Spline
 from hpsdecode.schemas import SUPPORTED_SCHEMAS, EncryptedData, ParseContext, get_parser
 
 if t.TYPE_CHECKING:
     import os
 
+    import numpy.typing as npt
+
     from hpsdecode.encryption import EncryptionKeyProvider
 
 #: List of XML paths to search for texture images that may be encrypted.
@@ -126,6 +130,143 @@ def get_required_text(element: ET.Element) -> str:
     return text
 
 
+def get_property_value(element: ET.Element, property_name: str) -> str:
+    """Get the value attribute from a 'Property' element.
+
+    :param element: The XML element.
+    :param property_name: The name of the property to find.
+    :return: The property value.
+    :raises HPSParseError: If the property or its value is missing.
+    """
+    property_element = element.find(f".//Property[@name='{property_name}']")
+    if property_element is None:
+        raise HPSParseError(f"Missing 'Property' element with name='{property_name}'")
+
+    value = property_element.get("value")
+    if value is None:
+        raise HPSParseError(f"Missing 'value' attribute on Property[@name='{property_name}']")
+
+    return value
+
+
+def extract_control_points_packed(data: bytes) -> npt.NDArray[np.floating]:
+    """Extract 3D control points from packed binary data (base64-encoded).
+
+    :param data: The binary data containing packed float coordinates.
+    :return: A numpy array of shape (N, 3) containing the control points.
+    :raises HPSParseError: If the data length is not divisible by 4 (size of float) or if no valid points are found.
+    """
+    if len(data) % 4 != 0:
+        raise HPSParseError(f"Packed control points data length {len(data)} is not divisible by 4 (sizeof float)")
+
+    floats = np.frombuffer(data, dtype=np.float32)
+    num_points = len(floats) // 3
+    if num_points == 0:
+        raise HPSParseError("No complete control points found in packed data")
+
+    return floats[: num_points * 3].reshape(num_points, 3)
+
+
+def extract_control_points_xml(element: ET.Element) -> npt.NDArray[np.floating]:
+    """Extract 3D control points from XML elements.
+
+    :param element: The XML element containing the control point objects.
+    :return: A numpy array of shape (N, 3) containing the control points.
+    :raises HPSParseError: If any control point is missing required attributes or if no valid points are found.
+    """
+    points: list[tuple[float, float, float]] = []
+
+    for obj in element.findall("Object"):
+        vector = obj.find("Vector[@name='p']")
+        if vector is None:
+            raise HPSParseError("Object in ControlPoints is missing Vector[@name='p'] element")
+
+        x_str = vector.get("x")
+        y_str = vector.get("y")
+        z_str = vector.get("z")
+        if x_str is None or y_str is None or z_str is None:
+            raise HPSParseError("Vector element is missing x, y, or z attribute")
+
+        try:
+            x = float(x_str)
+            y = float(y_str)
+            z = float(z_str)
+        except ValueError as e:
+            raise HPSParseError(f"Failed to parse vector coordinates: {e}") from e
+
+        points.append((x, y, z))
+
+    if not points:
+        raise HPSParseError("ControlPoints element contains no valid control points")
+
+    return np.array(points, dtype=np.float32)
+
+
+def parse_spline(element: ET.Element) -> Spline:
+    """Parse a single spline from an XML Object element.
+
+    :param element: The XML element representing the spline object.
+    :return: A Spline object containing the parsed data.
+    :raises HPSParseError: If required elements or attributes are missing.
+    """
+    name = get_property_value(element, "Name")
+    radius_str = get_property_value(element, "Radius")
+    closed_str = get_property_value(element, "Closed")
+    color_str = get_property_value(element, "Color")
+    misc_str = get_property_value(element, "iMisc1")
+
+    try:
+        radius = float(radius_str)
+        is_cyclic = closed_str.lower() == "true"
+        color = int(color_str)
+        misc = int(misc_str)
+    except ValueError as e:
+        raise HPSParseError(f"Failed to parse spline property values: {e}") from e
+
+    control_points_packed_element = element.find(".//ControlPointsPacked")
+    control_points_xml_element = element.find(".//ControlPoints")
+
+    if control_points_packed_element is not None:
+        control_points_text = control_points_packed_element.text
+        if control_points_text is None or control_points_text.strip() == "":
+            raise HPSParseError("ControlPointsPacked element has no content")
+
+        control_points_data = base64.b64decode(control_points_text.strip())
+        control_points = extract_control_points_packed(control_points_data)
+    elif control_points_xml_element is not None:
+        control_points = extract_control_points_xml(control_points_xml_element)
+    else:
+        raise HPSParseError("Spline object is missing control points")
+
+    return Spline(
+        name=name,
+        control_points=control_points,
+        radius=radius,
+        is_cyclic=is_cyclic,
+        color=color,
+        misc=misc,
+    )
+
+
+def parse_splines(root: ET.Element) -> list[Spline]:
+    """Parse all splines from the XML root element.
+
+    :param root: The root XML element of the HPS file.
+    :return: A list of parsed Spline objects.
+    """
+    splines: list[Spline] = []
+
+    splines_container = root.find(".//Splines")
+    if splines_container is None:
+        return splines
+
+    for obj in splines_container.findall(".//Object[@name='Spline']"):
+        spline = parse_spline(obj)
+        splines.append(spline)
+
+    return splines
+
+
 def parse_xml(file: str | os.PathLike[str] | t.IO[bytes] | bytes) -> ET.ElementTree:
     """Parse an HPS XML file.
 
@@ -217,6 +358,8 @@ def load_hps(
                 is_encrypted=is_encrypted and encryptable,
             )
 
+    splines = parse_splines(root)
+
     properties: dict[str, t.Any] = {}
     properties_element = root.find("Properties")
     if properties_element is not None:
@@ -237,6 +380,7 @@ def load_hps(
         vertex_colors_data=vertex_colors_data,
         texture_coords_data=texture_coords_data,
         texture_images=list(texture_images.values()),
+        splines=splines,
         check_value=int(check_value) if check_value else None,
         properties=properties,
     )
@@ -261,6 +405,7 @@ def load_hps(
         vertex_colors_data=context.vertex_colors_data,
         texture_coords_data=context.texture_coords_data,
         texture_images=context.texture_images,
+        splines=context.splines,
         vertex_commands=result.vertex_commands,
         face_commands=result.face_commands,
         check_value=context.check_value,