Merge pull request #25 from Ipuch/main

feat: forceplates and floor displayed with c3d

.gitignore

@@ -160,3 +160,5 @@ cython_debug/
 #.idea/
 
 .idea/
+
+sandbox/

examples/c3d/gait.py

@@ -0,0 +1,6 @@
+import os
+
+import pyorerun as prr
+
+print(os.getcwd())
+prr.c3d("gait.c3d")

pyorerun/phase_rerun.py

@@ -5,7 +5,7 @@
 from .abstract.q import QProperties
 from .biorbd_components.model_interface import BiorbdModel
 from .biorbd_phase import BiorbdRerunPhase
-from .timeless.gravity import Gravity
+from .timeless import Gravity, Floor, ForcePlate
 from .timeless_components import TimelessRerunPhase
 from .xp_components.markers import MarkersXp
 from .xp_components.timeseries_q import TimeSeriesQ
@@ -157,6 +157,16 @@ def add_q(
             TimeSeriesQ(name=f"{self.name}/{name}", q=q, properties=QProperties(joint_names=dof_names, ranges=ranges))
         )
 
+    def add_floor(self, square_width: float = None, height_offset: float = None, subsquares: int = None) -> None:
+        """Add a floor to the phase."""
+        self.timeless_components.add_component(
+            Floor(name=f"{self.name}", square_width=square_width, height_offset=height_offset, subsquares=subsquares)
+        )
+
+    def add_force_plate(self, num: int, corners: np.ndarray) -> None:
+        """Add a force plate to the phase."""
+        self.timeless_components.add_component(ForcePlate(name=f"{self.name}", num=num, corners=corners))
+
     def rerun(self, name: str = "animation_phase", init: bool = True, clear_last_node: bool = False) -> None:
         if init:
             rr.init(f"{name}_{self.phase}", spawn=True)

pyorerun/rrc3d.py

@@ -1,36 +1,75 @@
 from pathlib import Path
+from typing import Any
 
 import ezc3d
+import numpy as np
 import rerun as rr
 from pyomeca import Markers as PyoMarkers
 
 from .phase_rerun import PhaseRerun
 
 
-def rrc3d(c3d_file: str) -> None:
+def rrc3d(
+    c3d_file: str,
+    show_floor: bool = True,
+    show_force_plates: bool = True,
+    marker_trajectories: bool = False,
+) -> None:
     """
     Display a c3d file in rerun.
 
     Parameters
     ----------
     cd3_file: str
         The c3d file to display.
+    show_floor: bool
+        If True, show the floor.
+    show_force_plates: bool
+        If True, show the force plates.
+    show_camera: bool
+        If True, show the camera.
+    marker_trajectories: bool
+        If True, show the marker trajectories.
     """
 
     # Load a c3d file
     pyomarkers = PyoMarkers.from_c3d(c3d_file)
+    units = pyomarkers.units
     pyomarkers = adjust_position_unit_to_meters(pyomarkers, pyomarkers.units)
     t_span = pyomarkers.time.to_numpy()
     filename = Path(c3d_file).name
 
+    force_plates_corners = get_force_plates(c3d_file, units=units)
+    lowest_corner = get_lowest_corner(c3d_file, units=units)
+
     phase_rerun = PhaseRerun(t_span)
     phase_rerun.add_xp_markers(filename, pyomarkers)
+
+    if show_force_plates:
+        for i, corners in enumerate(force_plates_corners):
+            phase_rerun.add_force_plate(f"force_plate{i}", corners["corners"])
+
+    if show_floor:
+        square_width = max_xy_coordinate_span_by_markers(pyomarkers)
+        phase_rerun.add_floor(square_width, height_offset=lowest_corner)
+
     phase_rerun.rerun(filename)
 
-    # todo: find a better way to display curves but hacky way ok for now
-    for frame, t in enumerate(t_span):
-        rr.set_time_seconds("stable_time", t)
-        phase_rerun.xp_data.xp_data[0].to_rerun_curve(frame)
+    if marker_trajectories:
+        # todo: find a better way to display curves but hacky way ok for now
+        for frame, t in enumerate(t_span):
+            rr.set_time_seconds("stable_time", t)
+            phase_rerun.xp_data.xp_data[0].to_rerun_curve(frame)
+
+
+def max_xy_coordinate_span_by_markers(pyomarkers: PyoMarkers) -> float:
+    """Return the max span of the x and y coordinates of the markers."""
+    min_pyomarkers = np.min(np.min(pyomarkers.to_numpy(), axis=2), axis=1)
+    max_pyomarkers = np.max(np.max(pyomarkers.to_numpy(), axis=2), axis=1)
+    x_absolute_max = np.max(np.abs([min_pyomarkers[0], max_pyomarkers[0]]))
+    y_absolute_max = np.max(np.abs([min_pyomarkers[1], max_pyomarkers[1]]))
+
+    return np.max([x_absolute_max, y_absolute_max])
 
 
 def c3d_file_format(cd3_file) -> ezc3d.c3d:
@@ -41,15 +80,53 @@ def c3d_file_format(cd3_file) -> ezc3d.c3d:
     return cd3_file
 
 
-def adjust_position_unit_to_meters(pyomarkers: PyoMarkers, unit: str) -> PyoMarkers:
+def adjust_pyomarkers_unit_to_meters(pyomarkers: PyoMarkers, unit: str) -> PyoMarkers:
     """Adjust the positions to meters for displaying purposes."""
+    pyomarkers = adjust_position_unit_to_meters(pyomarkers, unit)
+    pyomarkers.attrs["units"] = "m"
+    return pyomarkers
+
+
+def adjust_position_unit_to_meters(array: Any, unit: str) -> PyoMarkers:
     conversion_factors = {"mm": 1000, "cm": 100, "m": 1}
     for u, factor in conversion_factors.items():
         if u in unit:
-            pyomarkers /= factor
+            array /= factor
             break
     else:
         raise ValueError("The unit of the c3d file is not in meters, mm or cm.")
+    return array
+
+
+def get_force_plates(c3d_file, units) -> list[dict[str, np.ndarray]]:
+    c3d_file = c3d_file_format(c3d_file)
+    force_plates = []
+    nb_force_plates = c3d_file["parameters"]["FORCE_PLATFORM"]["USED"]["value"][0]
+    for i in range(nb_force_plates):
+        force_plates.append(
+            {
+                "corners": adjust_position_unit_to_meters(
+                    c3d_file["parameters"]["FORCE_PLATFORM"]["CORNERS"]["value"][:, :, i],
+                    unit=units,
+                ),
+            }
+        )
+
+    return force_plates
+
+
+def get_lowest_corner(c3d_file, units) -> float:
+    c3d_file = c3d_file_format(c3d_file)
+    corners = c3d_file["parameters"]["FORCE_PLATFORM"]["CORNERS"]["value"][2, :, :]
+
+    if corners.shape[1] == 0:
+        return 0
+
+    return np.min(
+        adjust_position_unit_to_meters(
+            c3d_file["parameters"]["FORCE_PLATFORM"]["CORNERS"]["value"][2, :, :],
+            unit=units,
+        )
+    )
+
 
-    pyomarkers.attrs["units"] = "m"
-    return pyomarkers

pyorerun/timeless/__init__.py

@@ -0,0 +1,3 @@
+from .floor import Floor
+from .force_plate import ForcePlate
+from .gravity import Gravity

pyorerun/timeless/floor.py

@@ -0,0 +1,85 @@
+import numpy as np
+import rerun as rr
+
+from ..abstract.abstract_class import TimelessComponent
+
+
+class Floor(TimelessComponent):
+
+    def __init__(self, name, square_width: float, height_offset: float, subsquares: int):
+        self.name = name + "/floor"
+        self.vertices, self.faces = floor_mesh(
+            square_width=square_width,
+            height_offset=height_offset if height_offset is not None else 0,
+            subsquares=subsquares if subsquares is not None else 10,
+        )
+
+    @property
+    def nb_components(self):
+        return 1
+
+    def to_rerun(self) -> None:
+        rr.log(
+            "floor",
+            rr.Mesh3D(
+                vertex_positions=self.vertices,
+                vertex_normals=np.tile([0.0, 0.0, 1.0], reps=(self.vertices.shape[0], 1)),
+                vertex_colors=np.tile([150, 150, 150], reps=(self.vertices.shape[0], 1)),
+                indices=self.faces,
+            ),
+        )
+
+
+def floor_mesh(square_width: float, height_offset: float, subsquares: int) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Display a floor in rerun.
+
+    Parameters
+    ----------
+    square_width: float
+        The width of the floor in meters centered in zero.
+    height_offset: float
+        The height offset of the floor.
+    subsquares: int
+        The number of subsquares for each side of the floor. The total number of subsquares is subsquares^2.
+
+    Returns
+    -------
+    np.ndarray
+        The vertices of the floor.
+    np.ndarray
+        The faces of the floor.
+    """
+    x, y = np.meshgrid(
+        np.linspace(-square_width, square_width, subsquares),
+        np.linspace(-square_width, square_width, subsquares),
+    )
+
+    vertices = []
+    faces = []
+    base_index = 0
+
+    for i in range(0, subsquares - 1, 1):
+        starting_j = 0 if i % 2 == 0 else 1
+        for j in range(starting_j, subsquares - 1, 2):
+            triangle1_vertices = [
+                (x[i, j], y[i, j], height_offset),
+                (x[i + 1, j], y[i + 1, j], height_offset),
+                (x[i, j + 1], y[i, j + 1], height_offset),
+            ]
+            triangle2_vertices = [
+                (x[i + 1, j + 1], y[i + 1, j + 1], height_offset),
+            ]
+
+            vertices.extend(triangle1_vertices)
+            vertices.extend(triangle2_vertices)
+
+            triangle1_faces = [(base_index, base_index + 1, base_index + 2)]
+            triangle2_faces = [(base_index + 3, base_index + 1, base_index + 2)]
+
+            faces.extend(triangle1_faces)
+            faces.extend(triangle2_faces)
+
+            base_index += 4
+
+    return np.array(vertices), np.array(faces)

pyorerun/timeless/force_plate.py

@@ -0,0 +1,110 @@
+import numpy as np
+import rerun as rr
+
+from ..abstract.abstract_class import TimelessComponent
+
+
+class ForcePlate(TimelessComponent):
+
+    def __init__(self, name, num: int, corners: np.ndarray):
+        self.entity = f"force_plate{num}"
+        self.name = name + f"/{self.entity}"
+        self.vertices, self.faces = rectangle_mesh_from_corners(corners)
+
+    @property
+    def nb_components(self):
+        return 1
+
+    def to_rerun(self) -> None:
+        rr.log(
+            self.entity,
+            rr.Mesh3D(
+                vertex_positions=self.vertices,
+                vertex_normals=np.tile([0.0, 0.0, 1.0], reps=(self.vertices.shape[0], 1)),
+                vertex_colors=np.tile([1, 150, 150], reps=(self.vertices.shape[0], 1)),
+                indices=self.faces,
+            ),
+        )
+
+
+def rectangle_mesh_from_corners(corners: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Create a rectangle mesh from the corners.
+
+    Parameters
+    ----------
+    corners: np.ndarray
+        The corners of the rectangle.
+
+    Returns
+    -------
+    np.ndarray
+        The vertices of the floor.
+    np.ndarray
+        The faces of the floor.
+    """
+
+    vertices = corners.T
+
+    triangle1_faces = [0, 1, 2]
+    triangle2_faces = [2, 3, 0]
+
+    faces = []
+    faces.extend(triangle1_faces)
+    faces.extend(triangle2_faces)
+
+    return vertices, np.array(faces)
+
+
+def rectangle_mesh(length: float, width: float, height: float) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Create a rectangle mesh centered in zero.
+
+    Parameters
+    ----------
+    length: float
+        The length of the rectangle in meters (x-axis).
+    width: float
+        The width of the rectangle in meters (y-axis).
+    height: float
+        The height of the rectangle in meters (z-axis).
+
+    Returns
+    -------
+    np.ndarray
+        The vertices of the floor.
+    np.ndarray
+        The faces of the floor.
+    """
+    x, y = np.meshgrid(
+        np.linspace(-length / 2, length / 2, 2),
+        np.linspace(-width / 2, width / 2, 2),
+    )
+
+    vertices = []
+    faces = []
+    base_index = 0
+
+    for i in range(0, 1, 1):
+        for j in range(0, 1, 1):
+            triangle1_vertices = [
+                (x[i, j], y[i, j], height),
+                (x[i + 1, j], y[i + 1, j], height),
+                (x[i, j + 1], y[i, j + 1], height),
+            ]
+            triangle2_vertices = [
+                (x[i + 1, j + 1], y[i + 1, j + 1], height),
+            ]
+
+            vertices.extend(triangle1_vertices)
+            vertices.extend(triangle2_vertices)
+
+            triangle1_faces = [(base_index, base_index + 1, base_index + 2)]
+            triangle2_faces = [(base_index + 3, base_index + 1, base_index + 2)]
+
+            faces.extend(triangle1_faces)
+            faces.extend(triangle2_faces)
+
+            base_index += 4
+
+    return np.array(vertices), np.array(faces)