Source code for nicetoolbox.visual.media.viewer

"""
Viewer module for visualizing the components.

This module defines the Viewer class, which is responsible for visualizing the
components of the NICE toolbox.

Classes:
    Viewer: Class for visualizing the components.
"""

import numpy as np
import rerun as rr


[docs]class Viewer: """ Class for visualizing the components. Attributes: config (dict): Configuration settings for the viewer. """ def __init__(self, visualizer_config: dict): """ Initializes the Viewer class with the given configuration. Args: config (dict): Configuration settings for the viewer. """ self.visualizer_config = visualizer_config canvas_list = [] for component in self.visualizer_config["media"]["visualize"]["components"]: for canvases in self.visualizer_config["media"][component]["canvas"].values(): canvas_list.extend(canvases) self.canvas_list = list(set(canvas_list)) self.is_camera_position = visualizer_config["media"]["visualize"]["camera_position"] self.fps = self.visualizer_config["dataset_properties"]["fps"] self._create_canvas_roots()
[docs] def spawn(self, app_id: str = "NICE Toolbox Visualization") -> None: """ Spawns a rerun application with the given app_id. Args: app_id (str): The ID of the visualization application. Defaults to "NICE Toolbox Visualization". """ rr.init(app_id, spawn=self.visualizer_config["spawn_viewer"]) rr.set_time_seconds("time", 0)
[docs] def go_to_timestamp(self, frame_idx: int) -> None: """ Go to the specified frame index in the rerun viewer. Args: frame_idx (int): The index of the frame to go to. """ frame_time = frame_idx * (1 / self.fps) if frame_time == 0: frame_time = 0.00001 rr.set_time_seconds("time", frame_time)
[docs] def get_start_frame(self): """ Returns the start frame for visualization specified in the visualizer config. Returns: int: The start frame for visualization. """ return self.visualizer_config["media"]["visualize"]["start_frame"]
[docs] def get_end_frame(self) -> int: """ Returns the end frame for visualization specified in the visualizer config. Returns: int: The end frame for visualization. """ end_frame = self.visualizer_config["media"]["visualize"]["end_frame"] if end_frame == -1: end_frame = self.visualizer_config["video"]["video_length"] return end_frame
[docs] def get_step(self) -> int: """ Returns the frame step size for visualization specified in the visualizer config. Only every step-th frame will be visualized. Returns: int: The step for visualization. """ return self.visualizer_config["media"]["visualize"]["visualize_interval"]
[docs] def get_video_start(self) -> int: """ Returns the start time of the video in seconds. Returns: int: The start time of the video in seconds. """ return self.visualizer_config["video"]["video_start"]
def _create_canvas_roots(self) -> None: """ Creates the root paths for the 3D canvas, cameras, and images. The root paths are used to log the cameras and images in the rerun viewer. """ if "3D_Canvas" in self.canvas_list: self.ROOT3D = "3D_Canvas" if self.is_camera_position: self.CAMERAS_ROOT = "3D_Canvas/cameras" self.IMAGES_ROOT = "3D_Canvas/cameras" else: self.CAMERAS_ROOT = None self.IMAGES_ROOT = "cameras" elif self.is_camera_position: self.CAMERAS_ROOT = "3D_Canvas/cameras" self.IMAGES_ROOT = "3D_Canvas/cameras" self.ROOT3D = None else: self.IMAGES_ROOT = "cameras" self.ROOT3D = None self.CAMERAS_ROOT = None
[docs] def get_camera_pos_entity_path(self, camera_name: str) -> str: """ Returns the entity path for the camera position. Args: camera_name (str): The name of the camera. Returns: str: The entity path for the camera position. """ return f"{self.CAMERAS_ROOT}/{camera_name}"
[docs] def get_images_entity_path(self, camera_name): """ Returns the entity path for the images of a specific camera. Args: camera_name (str): The name of the camera. Returns: str: The entity path for the images of the specified camera. """ return f"{self.IMAGES_ROOT}/{camera_name}"
[docs] def get_is_camera_position(self): """ Returns a boolean indicating whether to display the camera position in the viewer. Returns: bool: Whether to display the camera position in the viewer. """ return self.is_camera_position
[docs] def generate_component_entity_path( self, component: str, is_3d: bool = True, cam_name: str = None, alg_name: str = None, subject_name: str = None, bodypart: str = None, ) -> str: """ Generates the entity path for a given component. Args: component (str): The name of the component. is_3d (bool, optional): Flag indicating if the component is 3D. Defaults to True. cam_name (str, optional): The name of the camera. Defaults to None. alg_name (str, optional): The name of the algorithm. Defaults to None. subject_name (str, optional): The name of the subject. Defaults to None. bodypart (str, optional): The name of the body part. Defaults to None. Returns: str: The generated entity path. Raises: ValueError: If the component is not implemented. """ if ( (component == "body_joints") | (component == "hand_joints") | (component == "face_landmarks") | (component == "gaze_individual") | (component == "emotion_individual") | (component == "head_orientation") ): if is_3d: entity_path = f"{self.ROOT3D}/{component}/{alg_name}/{subject_name}" else: entity_path = f"{self.IMAGES_ROOT}/{cam_name}/{component}/{alg_name}/" f"{subject_name}" elif component == "proximity": if is_3d: entity_path = f"{self.ROOT3D}/{component}/{alg_name}" else: entity_path = f"{self.IMAGES_ROOT}/{cam_name}/{component}/{alg_name}" elif component == "kinematics": if is_3d: entity_path = f"{alg_name}_{subject_name}/{bodypart}" else: entity_path = f"{alg_name}_{subject_name}_{cam_name}/{bodypart}" else: raise ValueError(f"ERROR in generate_component_entity_path(): Component {component} " "did not implemented") return entity_path
[docs] def log_camera(self, camera_calibration, entity_path) -> None: """ Logs the camera calibration in the viewer. Args: camera_calibration (dict): The camera calibration parameters. entity_path (str): The entity path for the camera. """ # intrinsic camera matrix K = np.array(camera_calibration["intrinsic_matrix"]) rr.log( entity_path, rr.Pinhole( resolution=camera_calibration["image_size"], image_from_camera=K[:3, :3].flatten(), ), ) translation = np.array(camera_calibration["translation"]).reshape(3) R = camera_calibration["rotation_matrix"] R_inv = np.linalg.inv(R) center = np.matmul(R_inv, -translation) # center *= 1000 rr.log(entity_path, rr.Transform3D(mat3x3=R_inv, translation=center.flatten()))
[docs] def log_image(self, image: np.array, entity_path: str, img_quality: int = 75) -> None: """ Logs an image in the viewer. Args: image (np.array): The image to log. entity_path (str): The entity path for the image. img_quality (int, optional): The quality of the image. Defaults to 75. """ rr.log(entity_path, rr.Image(image).compress(jpeg_quality=img_quality))
[docs] def check_multiview(self) -> None: """ Checks the consistency of the multi-view parameter in the visualizer config. Raises: ValueError: If the multi-view parameter is set to False but a 3D canvas is present. """ if (self.visualizer_config["media"]["multi_view"] is False) and ("3D_Canvas" in self.canvas_list): raise ValueError( "ERROR: multi-view parameter in Visualizer_config set false,\n " "But 3D_Canvas found in components, canvas lists.\n" "If you don't have multiple cameras, delete 3D_Canvas in all " "canvases\nIf you have multiple cameras, change multi-view " "parameter as true\n" )