From ee43fb09a97a787317405d7c85936319c6cac7e5 Mon Sep 17 00:00:00 2001 From: Ipuch Date: Fri, 3 May 2024 18:49:57 -0400 Subject: [PATCH] feat: add depth_images --- examples/biorbd/msk_model_plus_depth_image.py | 55 +++++++++++++++++++ pyorerun/multi_phase_rerun.py | 3 + pyorerun/phase_rerun.py | 20 +++++++ pyorerun/xp_components/depth_image.py | 40 ++++++++++++++ 4 files changed, 118 insertions(+) create mode 100644 examples/biorbd/msk_model_plus_depth_image.py create mode 100644 pyorerun/xp_components/depth_image.py diff --git a/examples/biorbd/msk_model_plus_depth_image.py b/examples/biorbd/msk_model_plus_depth_image.py new file mode 100644 index 0000000..5985223 --- /dev/null +++ b/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() diff --git a/pyorerun/multi_phase_rerun.py b/pyorerun/multi_phase_rerun.py index 195be6c..4615757 100644 --- a/pyorerun/multi_phase_rerun.py +++ b/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) diff --git a/pyorerun/phase_rerun.py b/pyorerun/phase_rerun.py index eac5edb..6735792 100644 --- a/pyorerun/phase_rerun.py +++ b/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) diff --git a/pyorerun/xp_components/depth_image.py b/pyorerun/xp_components/depth_image.py new file mode 100644 index 0000000..312eb80 --- /dev/null +++ b/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))