Merge pull request #42 from Ipuch/chunks

feat: send chunks

README.md

@@ -56,7 +56,7 @@ animation.rerun()
 ```
 
 ## From source
-```conda install -c conda-forge ezc3d rerun-sdk=0.16.1 trimesh numpy biorbd pyomeca tk imageio imageio-ffmpeg```
+```conda install -c conda-forge ezc3d rerun-sdk=0.18.2 trimesh numpy biorbd pyomeca tk imageio imageio-ffmpeg```
 
 Then, ensure it is accessible in your Python environment by installing the package:
 

environment.yml

@@ -4,7 +4,7 @@ channels:
 - default
 dependencies:
 - ezc3d
-- rerun-sdk=0.16.1
+- rerun-sdk>=0.20.1
 - numpy
 - biorbd>=1.10.5
 - trimesh

examples/biorbd/msk_model.py

@@ -21,9 +21,20 @@ def main():
     q[14, :] = np.linspace(0, np.pi / 8, nb_frames)
     q[15, :] = np.linspace(0, np.pi / 8, nb_frames)
 
+    import time
+
     viz = PhaseRerun(t_span)
-    viz.add_animated_model(model, q, display_q=True)
+    # viz.add_animated_model(model, q, display_q=True)
+    viz.add_animated_model(model, q, display_q=False)
+    tic = time.time()
     viz.rerun("msk_model")
+    toc = time.time()
+    print(f"Time to run: {toc - tic}")
+
+    tic = time.time()
+    viz.rerun_with_chunks("msk_model with chunk")
+    toc = time.time()
+    print(f"Time to run with chunks: {toc - tic}")
 
 
 if __name__ == "__main__":

examples/c3d/gait.py

@@ -1,3 +1,10 @@
 import pyorerun as prr
 
-prr.c3d("gait.c3d", show_floor=True, show_force_plates=True, show_forces=True, down_sampled_forces=False)
+prr.c3d(
+    "gait.c3d",
+    show_floor=True,
+    show_force_plates=True,
+    show_forces=True,
+    down_sampled_forces=True,
+    show_events=False,
+)

examples/c3d/main.py

@@ -1,3 +1,3 @@
 import pyorerun as prr
 
-prr.c3d("example.c3d", show_forces=False)
+prr.c3d("example.c3d", show_forces=False, show_events=False, marker_trajectories=True)

examples/c3d/running_gait.py

@@ -6,5 +6,6 @@
     show_force_plates=True,
     show_forces=True,
     down_sampled_forces=True,
+    show_events=False,
     video=("Running_0002_Oqus_6_15004.avi", "Running_0002_Oqus_9_15003.avi"),
 )

pyorerun/__init__.py

@@ -2,6 +2,7 @@
 from .biorbd_components.model_interface import BiorbdModel, BiorbdModelNoMesh
 from .biorbd_components.model_updapter import ModelUpdater
 from .live_animation import LiveModelAnimation
+from .live_integration import LiveModelIntegration
 from .multi_phase_rerun import MultiPhaseRerun
 
 # from .biorbd_phase import rr_biorbd as rrbiorbd, RerunBiorbdPhase

pyorerun/abstract/abstract_class.py

@@ -22,6 +22,10 @@ def to_rerun(self, q: np.ndarray):
     def nb_components(self):
         pass
 
+    @abstractmethod
+    def to_chunk(self, q: np.ndarray):
+        pass
+
 
 class Components(ABC):
     @abstractmethod
@@ -41,3 +45,7 @@ class ExperimentalData(Component):
     @abstractmethod
     def nb_frames(self):
         pass
+
+    @abstractmethod
+    def to_chunk(self, **kwargs) -> dict[str, list]:
+        pass

pyorerun/abstract/empty_updater.py

@@ -7,3 +7,9 @@ def __init__(self, name: str):
 
     def to_rerun(self, q: np.ndarray) -> None:
         pass
+
+    def to_chunk(self, q: np.ndarray) -> dict[str, list]:
+        return {"empty": None}
+
+    def initialize(self):
+        pass

pyorerun/abstract/markers.py

@@ -35,15 +35,17 @@ class MarkerProperties:
         Returns a numpy array with the color of each marker.
     """
 
-    def __init__(self, markers_names: list[str, ...] | tuple[str, ...], radius: float, color: np.ndarray):
+    def __init__(
+        self, markers_names: list[str, ...] | tuple[str, ...], radius: float | tuple[float, ...], color: np.ndarray
+    ):
         """
         Constructs all the necessary attributes for the MarkerProperties object.
 
         Parameters
         ----------
             markers_names : list[str, ...] | tuple[str, ...]
                 a list of names for the markers
-            radius : float
+            radius : float | tuple[float, ...]
                 the radius of the markers
             color : np.ndarray
                 the color of the markers
@@ -73,7 +75,9 @@ def radius_to_rerun(self) -> None:
         np.ndarray
             A numpy array with the radius of each marker.
         """
-        return np.ones(self.nb_markers) * self.radius
+        if isinstance(self.radius, float):
+            return np.ones(self.nb_markers) * self.radius
+        return np.array(self.radius)
 
     def color_to_rerun(self) -> None:
         """

pyorerun/biorbd_components/axisupdater.py

@@ -34,12 +34,18 @@ def color(self):
             return [0, 0, 255]
 
     def to_rerun(self, q: np.ndarray) -> None:
-        homogenous_matrices = self.transform_callable(q)
         rr.log(
             self.name,
-            rr.Arrows3D(
-                origins=homogenous_matrices[:3, 3],
-                vectors=homogenous_matrices[:3, self.axis] * self.scale,
-                colors=np.array(self.color),
-            ),
+            self.to_component(q),
         )
+
+    def to_component(self, q: np.ndarray) -> rr.Arrows3D:
+        homogenous_matrices = self.transform_callable(q)
+        return rr.Arrows3D(
+            origins=homogenous_matrices[:3, 3],
+            vectors=homogenous_matrices[:3, self.axis] * self.scale,
+            colors=np.array(self.color),
+        )
+
+    def to_chunk(self, q: np.ndarray) -> dict[str, list]:
+        pass

pyorerun/biorbd_components/ligaments.py

@@ -21,14 +21,41 @@ def nb_components(self) -> int:
     def to_rerun(self, q: np.ndarray) -> None:
         rr.log(
             self.name,
-            rr.LineStrips3D(
-                strips=self.update_callable(q),
-                radii=self.properties.radius_to_rerun(),
-                colors=self.properties.color_to_rerun(),
-                labels=self.properties.strip_names,
-            ),
+            self.to_component(q),
+        )
+
+    def to_component(self, q: np.ndarray) -> rr.LineStrips3D:
+        return rr.LineStrips3D(
+            strips=self.update_callable(q),
+            radii=self.properties.radius_to_rerun(),
+            colors=self.properties.color_to_rerun(),
+            labels=self.properties.strip_names,
         )
 
+    def compute_strips(self, q: np.ndarray) -> list:
+        nb_frames = q.shape[1]
+        return [self.update_callable(q[:, f]) for f in range(nb_frames)]
+
+    def to_chunk(self, q: np.ndarray) -> dict[str, list]:
+        nb_frames = q.shape[1]
+        # partition = [self.nb_strips for _ in range(nb_frames)]
+
+        strips_by_frame = self.compute_strips(q)
+        strips_output = {}
+        for s in range(self.nb_strips):
+            strips_output.update(
+                {
+                    f"{self.name}_{s}": [
+                        rr.LineStrips3D.indicator(),
+                        rr.components.LineStrip3DBatch([strips_by_frame[f][s] for f in range(nb_frames)]),
+                        rr.components.ColorBatch([self.properties.color for _ in range(nb_frames)]),
+                        rr.components.RadiusBatch([self.properties.radius for _ in range(nb_frames)]),
+                    ]
+                }
+            )
+
+        return strips_output
+
 
 class LigamentsUpdater(LineStripUpdater):
     def __init__(self, name, properties: LineStripProperties, update_callable: callable):
@@ -80,6 +107,34 @@ def to_rerun(self, q: np.ndarray = None, markers: np.ndarray = None) -> None:
             ),
         )
 
+    def compute_all_strips(self, q: np.ndarray, markers: np.ndarray) -> np.ndarray:
+        nb_frames = q.shape[1]
+        strips = np.zeros((self.nb_strips, 2, 3, nb_frames))
+        for f in range(nb_frames):
+            strips[:, :, :, f] = self.line_strips(q[:, f], markers[:, :, f])
+
+        return strips
+
+    def to_chunk(self, q: np.ndarray, markers: np.ndarray) -> dict[str, list]:
+        nb_frames = q.shape[1]
+        strips_by_frame = self.compute_all_strips(q, markers)
+
+        strips_output = {}
+        for s in range(self.nb_strips):
+            print(self.name, s)
+            strips_output.update(
+                {
+                    f"{self.name}_{s}": [
+                        rr.LineStrips3D.indicator(),
+                        rr.components.LineStrip3DBatch([strips_by_frame[s, :, :, f] for f in range(nb_frames)]),
+                        rr.components.ColorBatch([self.properties.color for _ in range(nb_frames)]),
+                        rr.components.RadiusBatch([self.properties.radius for _ in range(nb_frames)]),
+                    ]
+                }
+            )
+
+        return strips_output
+
 
 class LineStripUpdaterFromGlobalTransform(LineStripUpdater):
     def __init__(self, name, properties: LineStripProperties, strips: np.ndarray, transform_callable: callable):
@@ -90,6 +145,12 @@ def __init__(self, name, properties: LineStripProperties, strips: np.ndarray, tr
             colors=self.properties.color_to_rerun(),
         )
 
+    def initialize(self):
+        rr.log(
+            self.name,
+            self.rerun_mesh,
+        )
+
     def to_rerun(self, q: np.ndarray) -> None:
         homogenous_matrices = self.update_callable(q)
         rr.log(
@@ -103,3 +164,24 @@ def to_rerun(self, q: np.ndarray) -> None:
             self.name,
             self.rerun_mesh,
         )
+
+    def compute_all_transforms(self, q: np.ndarray) -> np.ndarray:
+        nb_frames = q.shape[1]
+        homogenous_matrices = np.zeros((4, 4, nb_frames))
+        for f in range(nb_frames):
+            homogenous_matrices[:, :, f] = self.update_callable(q[:, f])
+
+        return homogenous_matrices
+
+    def to_chunk(self, q: np.ndarray) -> dict[str, list]:
+        homogenous_matrices = self.compute_all_transforms(q)
+
+        return {
+            self.name: [
+                rr.InstancePoses3D.indicator(),
+                rr.components.PoseTranslation3DBatch(homogenous_matrices[:3, 3, :].T),
+                rr.components.PoseTransformMat3x3Batch(
+                    [homogenous_matrices[:3, :3, f] for f in range(homogenous_matrices.shape[2])]
+                ),
+            ]
+        }

pyorerun/biorbd_components/local_frame.py

@@ -22,12 +22,45 @@ def __init__(self, name, transform_callable: callable):
     def nb_components(self):
         return 1
 
-    def to_rerun(self, q: np.ndarray) -> None:
-        homogenous_matrices = self.transform_callable(q)
+    def initialize(self):
         rr.log(
             self.name,
             rr.Transform3D(
-                translation=homogenous_matrices[:3, 3],
-                mat3x3=homogenous_matrices[:3, :3] * self.scale,
+                translation=np.zeros(3),
+                mat3x3=np.eye(3),
             ),
         )
+
+    def to_rerun(self, q: np.ndarray) -> None:
+        rr.log(
+            self.name,
+            self.to_component(q),
+        )
+
+    def to_component(self, q: np.ndarray) -> rr.Transform3D:
+        homogenous_matrices = self.transform_callable(q)
+        return rr.Transform3D(
+            translation=homogenous_matrices[:3, 3],
+            mat3x3=homogenous_matrices[:3, :3] * self.scale,
+        )
+
+    def compute_all_transforms(self, q: np.ndarray) -> np.ndarray:
+        nb_frames = q.shape[1]
+        homogenous_matrices = np.zeros((4, 4, nb_frames))
+        for f in range(nb_frames):
+            homogenous_matrices[:, :, f] = self.transform_callable(q[:, f])
+
+        return homogenous_matrices
+
+    def to_chunk(self, q: np.ndarray) -> dict[str, list]:
+        homogenous_matrices = self.compute_all_transforms(q)
+
+        return {
+            self.name: [
+                rr.InstancePoses3D.indicator(),
+                rr.components.PoseTranslation3DBatch(homogenous_matrices[:3, 3, :].T),
+                rr.components.PoseTransformMat3x3Batch(
+                    [homogenous_matrices[:3, :3, f] for f in range(homogenous_matrices.shape[2])]
+                ),
+            ]
+        }

pyorerun/biorbd_components/mesh.py

@@ -60,6 +60,11 @@ def from_file(cls, name, file_path: str, transform_callable) -> "TransformableMe
             return cls(name, mesh, transform_callable)
         if file_path.endswith(".vtp"):
             output = read_vtp_file(file_path)
+            is_not_a_trimesh = output["polygons"].shape[1] > 3
+            if is_not_a_trimesh:
+                raise ValueError(
+                    f"The file {file_path} is not a triangular-only mesh. It has polygons with more than 3 vertices."
+                )
             mesh = Trimesh(
                 vertices=output["nodes"],
                 faces=output["polygons"],
@@ -91,6 +96,12 @@ def name(self):
     def nb_components(self):
         return 1
 
+    def initialize(self):
+        rr.log(
+            self.name,
+            self.rerun_mesh,
+        )
+
     def to_rerun(self, q: np.ndarray) -> None:
         homogenous_matrices = self.transform_callable(q)
         rr.log(
@@ -100,11 +111,39 @@ def to_rerun(self, q: np.ndarray) -> None:
                 mat3x3=homogenous_matrices[:3, :3],
             ),
         )
-        rr.log(
-            self.name,
-            self.rerun_mesh,
+
+    def to_component(self, q: np.ndarray) -> rr.Transform3D:
+        homogenous_matrices = self.transform_callable(q)
+        return self.to_component_from_homogenous_mat(homogenous_matrices)
+
+    @staticmethod
+    def to_component_from_homogenous_mat(mat: np.ndarray) -> rr.Transform3D:
+        return rr.Transform3D(
+            translation=mat[:3, 3],
+            mat3x3=mat[:3, :3],
         )
 
     @property
     def component_names(self):
         return [self.name]
+
+    def compute_all_transforms(self, q: np.ndarray) -> np.ndarray:
+        nb_frames = q.shape[1]
+        homogenous_matrices = np.zeros((4, 4, nb_frames))
+        for f in range(nb_frames):
+            homogenous_matrices[:, :, f] = self.transform_callable(q[:, f])
+
+        return homogenous_matrices
+
+    def to_chunk(self, q: np.ndarray) -> dict[str, list]:
+        homogenous_matrices = self.compute_all_transforms(q)
+
+        return {
+            self.name: [
+                rr.InstancePoses3D.indicator(),
+                rr.components.PoseTranslation3DBatch(homogenous_matrices[:3, 3, :].T),
+                rr.components.PoseTransformMat3x3Batch(
+                    [homogenous_matrices[:3, :3, f] for f in range(homogenous_matrices.shape[2])]
+                ),
+            ]
+        }