feat(trc reader)

Author: ipuch

environment.yml

@@ -14,4 +14,6 @@ dependencies:
 - imageio
 - imageio-ffmpeg
 - opensim
+# from pip
+# pip install trc-data-reader
 

examples/osim/ABD01.trc

@@ -0,0 +1,5 @@
+import pyorerun as prr
+
+prr.trc(
+    "ABD01.trc",
+)

examples/osim/trc_reader.py

@@ -38,5 +38,6 @@
 from .pyoemg import PyoMuscles
 from .rrbiomod import rr_biorbd as animate
 from .rrc3d import rrc3d as c3d
+from .rrtrc import rrtrc as trc
 from .xp_components.timeseries_q import OsimTimeSeries
 from .xp_components.persistent_marker_options import PersistentMarkerOptions

pyorerun/__init__.py

@@ -5,6 +5,7 @@
 from typing import Optional, List
 
 import ezc3d
+from trc import TRCData
 import numpy as np
 
 
@@ -200,6 +201,58 @@ def from_c3d(
 
         return cls(data=points, time=time, marker_names=marker_names, show_labels=show_labels, attrs=attrs)
 
+    @classmethod
+    def from_trc(cls, filename: str, show_labels: bool = True) -> "PyoMarkers":
+        """
+        Create PyoMarkers from a TRC file.
+
+        Parameters
+        ----------
+        filename : str
+            Path to the TRC file
+        show_labels : bool, default True
+            Whether to show marker labels
+
+        Returns
+        -------
+        PyoMarkers
+            A new PyoMarkers instance
+        """
+        trc_data = TRCData()
+        trc_data.load(filename)
+
+        # Get marker names
+        marker_names = trc_data["Markers"]
+        n_markers = len(marker_names)
+        n_frames = int(trc_data["NumFrames"])
+
+        # Initialize data array (3, n_markers, n_frames)
+        data = np.zeros((3, n_markers, n_frames))
+
+        # Fill data array with marker coordinates
+        for frame_idx, frame_num in enumerate(trc_data["Frame#"]):
+            _, frame_data = trc_data[frame_num]
+            for marker_idx, marker_coords in enumerate(frame_data):
+                # marker_coords is [X, Y, Z]
+                data[:, marker_idx, frame_idx] = marker_coords
+
+        # Get time vector
+        time = np.array(trc_data["Time"])
+
+        # Get units and rate
+        units = trc_data.get("Units", "mm")
+        rate = trc_data.get("DataRate", None)
+
+        attrs = {
+            "units": units,
+            "rate": rate,
+            "filename": filename,
+            "first_frame": trc_data["Frame#"][0] if trc_data["Frame#"] else 0,
+            "last_frame": trc_data["Frame#"][-1] if trc_data["Frame#"] else n_frames - 1,
+        }
+
+        return cls(data=data, time=time, marker_names=marker_names, show_labels=show_labels, attrs=attrs)
+
 
 class MockChannel:
     """

pyorerun/pyomarkers.py

@@ -0,0 +1,96 @@
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import rerun as rr
+
+from .phase_rerun import PhaseRerun
+from .pyomarkers import PyoMarkers
+from .multi_frame_rate_phase_rerun import MultiFrameRatePhaseRerun
+
+
+def rrtrc(
+    trc_filename: str,
+    marker_trajectories: bool = False,
+    show_floor: bool = True,
+    notebook: bool = False,
+) -> None:
+    """
+    Display a c3d file in rerun.
+
+    Parameters
+    ----------
+    trc_filename: str
+        The path to the trc file.
+    marker_trajectories: bool
+        If True, show the marker trajectories.
+    show_floor: bool
+        If True, show the floor.
+    notebook: bool
+        If True, display the animation in the notebook.
+    """
+
+    # Load a c3d file
+    pyomarkers = PyoMarkers.from_trc(trc_filename)
+    units = pyomarkers.units
+    pyomarkers = adjust_position_unit_to_meters(pyomarkers, pyomarkers.units)
+    pyomarkers.show_labels = False
+
+    t_span = pyomarkers.time
+    filename = Path(trc_filename).name
+
+    phase_reruns = []
+    phase_rerun = PhaseRerun(t_span)
+    phase_reruns.append(phase_rerun)
+    phase_rerun.add_xp_markers(filename, pyomarkers)
+
+    if show_floor:
+        square_width = max_xy_coordinate_span_by_markers(pyomarkers)
+        lowest_corner = 0
+        phase_rerun.add_floor(square_width, height_offset=lowest_corner - 0.0005)
+
+    multi_phase_rerun = MultiFrameRatePhaseRerun(phase_reruns)
+    multi_phase_rerun.rerun(filename, notebook=notebook)
+
+    if marker_trajectories:
+        # # todo: find a better way to display curves but hacky way ok for now
+        marker_names = phase_rerun.xp_data.xp_data[0].marker_names
+        for m in marker_names:
+            for j, axis in enumerate(["X", "Y", "Z"]):
+                rr.send_columns(
+                    f"markers_graphs/{m}/{axis}",
+                    indexes=[rr.TimeColumn("stable_time", duration=t_span)],
+                    columns=[
+                        *rr.Scalars.columns(
+                            scalars=phase_rerun.xp_data.xp_data[0].markers_numpy[j, marker_names.index(m), :]
+                        )
+                    ],
+                )
+
+
+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.nanmin(np.nanmin(pyomarkers.to_numpy(), axis=2), axis=1)
+    max_pyomarkers = np.nanmax(np.nanmax(pyomarkers.to_numpy(), axis=2), axis=1)
+    x_absolute_max = np.nanmax(np.abs([min_pyomarkers[0], max_pyomarkers[0]]))
+    y_absolute_max = np.nanmax(np.abs([min_pyomarkers[1], max_pyomarkers[1]]))
+
+    return np.max([x_absolute_max, y_absolute_max])
+
+
+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) -> Any:
+    conversion_factors = {"mm": 1000, "cm": 100, "m": 1}
+    for u, factor in conversion_factors.items():
+        if u in unit:
+            array /= factor
+            break
+    else:
+        raise ValueError("The unit of the c3d file is not in meters, mm or cm.")
+    return array

pyorerun/rrtrc.py

@@ -28,6 +28,7 @@ dependencies = [
     "imageio",
     "imageio-ffmpeg",
     "matplotlib",
+    "trc-data-reader"
 #    "opensim", # Not yet available on pypi, use `conda install opensim-org opensim=4.5.1`
 #    "biorbd" # Not yet available on pypi, use `conda install -c conda-forge biorbd`
     ]