feat: add depth_images

examples/biorbd/msk_model_plus_depth_image.py

@@ -0,0 +1,55 @@
+"""
+This example shows how to run multiple models in the same window. It also shows how to run multiple phases in different
+windows. It also shows how to add experimental markers to the animation.
+"""
+
+import numpy as np
+from numpy import random
+
+from pyorerun import BiorbdModel, MultiPhaseRerun
+
+
+def building_some_q_and_t_span(nb_frames: int, nb_seconds: int) -> tuple[np.ndarray, np.ndarray]:
+    t_span = np.linspace(0, nb_seconds, nb_frames)
+
+    # building some generalized coordinates
+    q = np.zeros((2, nb_frames))
+    q[0, :] = np.linspace(0, 0.1, nb_frames)
+    q[1, :] = np.linspace(0, 0.3, nb_frames)
+    return q, t_span
+
+
+def main():
+    biorbd_model_path = "models/double_pendulum.bioMod"
+
+    # building some time components
+    nb_frames = 200
+    nb_seconds = 1
+    q0, t_span0 = building_some_q_and_t_span(nb_frames, nb_seconds)
+    q1, t_span1 = building_some_q_and_t_span(20, 0.5)
+
+    # loading biorbd model
+    biorbd_model = BiorbdModel(biorbd_model_path)
+    noisy_markers = np.zeros((3, biorbd_model.nb_markers, nb_frames))
+    for i in range(nb_frames):
+        noisy_markers[:, :, i] = biorbd_model.markers(q0[:, i]).T + random.random((3, biorbd_model.nb_markers)) * 0.1
+
+    depth_image = 65535 * np.ones((200, 300), dtype=np.int64)
+    depth_image[50:150, 50:150] = 20000
+    depth_image[130:180, 100:280] = 45000
+    # add noise and repear on nb_frames
+    depth_image = np.tile(depth_image[:, :, np.newaxis], (1, 1, nb_frames))
+    depth_image += random.random_integers(low=0, high=1000, size=depth_image.shape)
+
+    # running the animation
+    multi_rerun = MultiPhaseRerun()
+
+    multi_rerun.add_phase(t_span=t_span0, phase=0, window="animation")
+    multi_rerun.add_animated_model(biorbd_model, q0, phase=0, window="animation")
+    multi_rerun.add_depth_image(name="depth", depth_image=depth_image, phase=0, window="animation")
+
+    multi_rerun.rerun("multi_model_test")
+
+
+if __name__ == "__main__":
+    main()

pyorerun/multi_phase_rerun.py

@@ -33,6 +33,9 @@ def add_animated_model(self, biomod: BiorbdModel, q: np.ndarray, phase: int = 0,
     def add_xp_markers(self, name: str, markers: PyoMarkers, phase: int = 0, window: str = "animation") -> None:
         self.rerun_biorbd_phases[phase][window].add_xp_markers(name, markers)
 
+    def add_depth_image(self, name: str, depth_image: np.ndarray, phase: int = 0, window: str = "animation") -> None:
+        self.rerun_biorbd_phases[phase][window].add_depth_image(name, depth_image)
+
     @property
     def nb_phase(self) -> int:
         return len(self.rerun_biorbd_phases)

pyorerun/phase_rerun.py

@@ -4,6 +4,7 @@
 
 from .biorbd_components.model_interface import BiorbdModel
 from .biorbd_phase import BiorbdRerunPhase
+from .xp_components.depth_image import DepthImage
 from .xp_components.markers import MarkersXp
 from .xp_phase import XpRerunPhase
 
@@ -93,6 +94,25 @@ def add_xp_markers(self, name, markers: PyoMarkers) -> None:
             )
         self.xp_data.add_data(MarkersXp(name=f"{self.name}/{name}", markers=markers))
 
+    def add_depth_image(self, name, depth_image: np.ndarray) -> None:
+        """
+        Add an animated model to the phase.
+
+        Parameters
+        ----------
+        name: str
+            The name of the markers set.
+        markers: PyoMarkers
+            The experimental data to display.
+        """
+        if depth_image.shape[2] != self.t_span.shape[0]:
+            raise ValueError(
+                f"The shapes of q and tspan are inconsistent. "
+                f"They must have the same length."
+                f"Current shapes are: {depth_image.shape[2]} and tspan: {self.t_span.shape}."
+            )
+        self.xp_data.add_data(DepthImage(name=f"{self.name}/{name}", depth_image=depth_image))
+
     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/xp_components/depth_image.py

@@ -0,0 +1,40 @@
+import numpy as np
+import rerun as rr
+
+from ..abstract.abstract_class import ExperimentalData
+
+
+class DepthImage(ExperimentalData):
+    def __init__(self, name, depth_image: np.ndarray):
+        self.name: str = name + "/depth_image"
+        self.depth_image: np.ndarray = depth_image
+
+    @property
+    def size_x(self):
+        return self.depth_image.shape[0]
+
+    @property
+    def size_y(self):
+        return self.depth_image.shape[1]
+
+    @property
+    def nb_frames(self):
+        return self.depth_image.shape[2]
+
+    @property
+    def nb_components(self):
+        return 1
+
+    def to_rerun(self, frame: int) -> None:
+        depth_image_frame = self.depth_image[:, :, frame]
+        rr.log(
+            self.name,
+            rr.Pinhole(
+                width=depth_image_frame.shape[1],
+                height=depth_image_frame.shape[0],
+                focal_length=200,
+            ),
+        )
+
+        # Log the tensor.
+        rr.log(f"{self.name}/depth", rr.DepthImage(depth_image_frame, meter=10_000.0))