Source code for nicetoolbox.visual.media.main

"""
Main module for initializing and running the visualizer.
"""

import argparse
import os
from pathlib import Path

import cv2

from ...utils import visual_utils as vis_utils
from ...utils.system import check_long_path_support
from .. import config_handler as vis_cfg
from ..in_out import IO
from .components import (
    BodyJointsComponent,
    EmotionIndividualComponent,
    FaceLandmarksComponent,
    GazeIndividualComponent,
    GazeInteractionComponent,
    HandJointsComponent,
    HeadOrientationComponent,
    KinematicsComponent,
    ProximityComponent,
)
from .viewer import Viewer


[docs]def main(project_folder_path: Path, machine_specifics_file: Path, visualizer_config_file: Path): """ Main function to run the visualizer. This function sets up the configuration, initializes the input/output handlers, loads calibration data, and initializes the viewer for visualizing the components. Args: project_folder_path (Path): Path to the project folder containing nice_project.toml. machine_specifics_file (Path): Path to machine_specific_paths.toml. visualizer_config_file (Path): Path to visualizer_config.toml, may contain placeholders. """ # SYSTEM CHECK check_long_path_support() # CONFIGURATION - IO config_handler = vis_cfg.Configuration(project_folder_path, machine_specifics_file, visualizer_config_file) visualizer_config = config_handler.get_updated_visualizer_config() # IO io = IO(visualizer_config) nice_tool_input_folder = io.get_component_nice_tool_input_folder( visualizer_config["video"], visualizer_config["io"]["dataset_name"] ) # load calibration for the video calibration_file = io.get_calibration_file(visualizer_config["video"]) calib = None if calibration_file: calib = vis_utils.load_calibration( calibration_file, visualizer_config["video"], camera_names=config_handler.get_camera_names(), ) if not calib: print( "WARNING: User did not provide a valid calibration file. " "Visualization of camera positions, 3d pose estimation, and gaze results " "requires calibration data." ) if visualizer_config["media"]["multi_view"]: raise ValueError( "ERROR: Calibration file was not provided. Cannot visualize multi-view " "Set it False in visualizer_config \n " ) # INITIALIZE VIEWER viewer = Viewer(visualizer_config) # CHECK CONFIGURATION all_cameras = config_handler.get_camera_names() config_handler.check_config(calibration_file) for cam in all_cameras: config_handler.check_calibration(calib, cam) viewer.check_multiview() # LOAD COMPONENTS DATA components_list = visualizer_config["media"]["visualize"]["components"] components = [] for component in components_list: if component not in os.listdir(io.get_experiment_video_folder()): print( f"WARNING: {component} Component is not found in video output. " f"It will not be visualized.\n To avoid this warning, consider " f"removing '{component}' from the components list in the " "visualizer_config.toml file" ) continue components.append(component) if "body_joints" in components: body_joints_component = BodyJointsComponent(visualizer_config, io, viewer, "body_joints") eyes_middle_2d_data = body_joints_component.calculate_middle_eyes(dimension=2) eyes_middle_3d_data = ( body_joints_component.calculate_middle_eyes(dimension=3) if visualizer_config["media"]["multi_view"] is True else (None, None) ) else: body_joints_component = None eyes_middle_2d_data = None, None eyes_middle_3d_data = None, None hand_joints_component = ( HandJointsComponent(visualizer_config, io, viewer, "hand_joints") if "hand_joints" in components else None ) face_landmarks_component = ( FaceLandmarksComponent(visualizer_config, io, viewer, "face_landmarks") if "face_landmarks" in components else None ) gaze_interaction_component = ( GazeInteractionComponent(visualizer_config, io, viewer, "gaze_interaction") if "gaze_interaction" in components else None ) look_at_data_tuple = ( gaze_interaction_component.get_lookat_data() if "gaze_interaction" in components else None ) # returns (data, data_labels) gaze_ind_component = ( GazeIndividualComponent( visualizer_config, io, viewer, "gaze_individual", calib, eyes_middle_3d_data, look_at_data_tuple, ) if "gaze_individual" in components else None ) emotion_ind_component = ( EmotionIndividualComponent(visualizer_config, io, viewer, "emotion_individual") if "emotion_individual" in components else None ) head_orientation_component = ( HeadOrientationComponent(visualizer_config, io, viewer, "head_orientation") if "head_orientation" in components else None ) proximity_component = ( ProximityComponent( visualizer_config, io, viewer, "proximity", eyes_middle_3d_data, eyes_middle_2d_data, ) if "proximity" in components else None ) kinematics_component = ( KinematicsComponent(visualizer_config, io, viewer, "kinematics") if "kinematics" in components else None ) instances = [ body_joints_component, hand_joints_component, face_landmarks_component, gaze_ind_component, emotion_ind_component, proximity_component, kinematics_component, head_orientation_component, ] # VISUALIZATION # initialize rerun visualizer viewer.spawn() for camera in all_cameras: if viewer.get_is_camera_position(): entity_path_cams = viewer.get_camera_pos_entity_path(camera) viewer.log_camera(calib[camera], entity_path_cams) frame_idx = viewer.get_start_frame() end_frame = viewer.get_end_frame() while True: if end_frame != -1 and frame_idx > end_frame: break viewer.go_to_timestamp(frame_idx) frame_no = viewer.get_video_start() + frame_idx + config_handler.get_dataset_starting_index() image_name = f"{frame_no:09}.png" for camera in all_cameras: # log camera into 3d canvas image_path = os.path.join(nice_tool_input_folder, camera, "frames", image_name).replace("\\", "/") # TODO: we will stop if we will just stop when there is no frame exist # is it a good idea? probably not, but we can fix it latter if not os.path.exists(image_path): return image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB) entity_path_imgs = viewer.get_images_entity_path(camera) viewer.log_image(image, entity_path_imgs, img_quality=75) for instance in instances: if instance is not None: instance.visualize(frame_idx) frame_idx += viewer.get_step()
[docs]def entry_point(): """Entry point for running NICE toolbox rerun visualizations.""" parser = argparse.ArgumentParser() parser.add_argument( "--project_folder_path", default=Path("."), type=Path, required=False, help="Path to the NICE Toolbox project folder containing nice_project.toml config", ) parser.add_argument( "--machine_specifics", default=Path("machine_specific_paths.toml"), type=Path, required=False, help="Path to machine_specific_paths.toml config", ) parser.add_argument( "--visual_config", default=Path("<configs_folder_path>/visualizer_config.toml"), type=Path, required=False, help="Path to visualizer_config.toml, supports placeholders", ) args = parser.parse_args() main(args.project_folder_path, args.machine_specifics, args.visual_config)
if __name__ == "__main__": entry_point()