path: root/src/argaze/utils/tobii_stream_aruco_aoi_display.py

#!/usr/bin/env python

import argparse
import os, json

from argaze import DataStructures
from argaze import GazeFeatures
from argaze.TobiiGlassesPro2 import *
from argaze.ArUcoMarkers import *
from argaze.AreaOfInterest import *
from argaze.utils import MiscFeatures

import cv2 as cv
import numpy

def main():
    """
    Track any ArUco marker into Tobii Glasses Pro 2 camera video stream.
    For each loaded AOI scene .obj file, position the scene virtually relatively to each detected ArUco markers and project the scene into camera frame. 
    """

    # Manage arguments
    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
    parser.add_argument('-t', '--tobii_ip', metavar='TOBII_IP', type=str, default=None, help='tobii glasses ip')
    parser.add_argument('-c', '--camera_calibration', metavar='CAM_CALIB', type=str, default=None, help='json camera calibration filepath')
    parser.add_argument('-p', '--aruco_tracker_configuration', metavar='TRACK_CONFIG', type=str, default=None, help='json aruco tracker configuration filepath')
    parser.add_argument('-md', '--marker_dictionary', metavar='MARKER_DICT', type=ArUcoMarkersDictionary.ArUcoMarkersDictionary, default='DICT_ARUCO_ORIGINAL', help='aruco marker dictionnary (DICT_4X4_50, DICT_4X4_100, DICT_4X4_250, DICT_4X4_1000, DICT_5X5_50, DICT_5X5_100, DICT_5X5_250, DICT_5X5_1000, DICT_6X6_50, DICT_6X6_100, DICT_6X6_250, DICT_6X6_1000, DICT_7X7_50, DICT_7X7_100, DICT_7X7_250, DICT_7X7_1000, DICT_ARUCO_ORIGINAL,DICT_APRILTAG_16h5, DICT_APRILTAG_25h9, DICT_APRILTAG_36h10, DICT_APRILTAG_36h11)')
    parser.add_argument('-ms', '--marker_size', metavar='MARKER_SIZE', type=float, default=6, help='aruco marker size (cm)')
    parser.add_argument('-mi', '--marker_id_scene', metavar='MARKER_ID_SCENE', type=json.loads, help='{"marker": "aoi scene filepath"} dictionary')
    parser.add_argument('-w', '--window', metavar='DISPLAY', type=bool, default=True, help='enable window display', action=argparse.BooleanOptionalAction)
    args = parser.parse_args()

    # Manage markers id to track
    if args.marker_id_scene == None:
        print(f'Track any Aruco markers from the {args.marker_dictionary.name} dictionary')
    else:
        print(f'Track Aruco markers {list(args.marker_id_scene.keys())} from the {args.marker_dictionary.name} dictionary')
    
    # Create tobii controller (with auto discovery network process if no ip argument is provided)
    print("Looking for a Tobii Glasses Pro 2 device ...")

    try:

        tobii_controller = TobiiController.TobiiController(args.tobii_ip)
        print(f'Tobii Glasses Pro 2 device found at {tobii_controller.address} address.')

    except ConnectionError as e:

        print(e)
        exit()

    # Setup camera at 25 fps to work on Full HD video stream
    tobii_controller.set_scene_camera_freq_25()

    # Print current confirugration
    print(f'Tobii Glasses Pro 2 configuration:')
    for key, value in tobii_controller.get_configuration().items():
        print(f'\t{key}: {value}')

    # Enable tobii data stream 
    tobii_data_stream = tobii_controller.enable_data_stream()

    # Enable tobii video stream
    tobii_video_stream = tobii_controller.enable_video_stream()

    # Create aruco camera
    aruco_camera = ArUcoCamera.ArUcoCamera()
    
     # Load calibration file
    if args.camera_calibration != None:

        aruco_camera.load_calibration_file(args.camera_calibration)

    else:

        raise UserWarning('.json camera calibration filepath required. Use -c option.')

    # Create aruco tracker
    aruco_tracker = ArUcoTracker.ArUcoTracker(args.marker_dictionary, args.marker_size, aruco_camera)

    # Load specific configuration file
    if args.aruco_tracker_configuration != None:

        aruco_tracker.load_configuration_file(args.aruco_tracker_configuration)

        print(f'ArUcoTracker configuration for {args.marker_dictionary.name} markers detection:')
        aruco_tracker.print_configuration()

    # Load AOI 3D scene for each marker and create a AOI 2D scene and frame when a 'Visualisation_Plan' AOI exist
    aoi3D_scenes = {}
    aoi2D_visu_scenes = {}

    for marker_id, aoi_scene_filepath in args.marker_id_scene.items():

        marker_id = int(marker_id)
        
        aoi3D_scenes[marker_id] = AOI3DScene.AOI3DScene()
        aoi3D_scenes[marker_id].load(aoi_scene_filepath)

        print(f'AOI in {os.path.basename(aoi_scene_filepath)} scene related to marker #{marker_id}:')
        for aoi in aoi3D_scenes[marker_id].keys():

            print(f'\t{aoi}')

    def aoi3D_scene_selector(marker_id):
        return aoi3D_scenes.get(marker_id, None)

    # Create timestamped buffer to store AOIs scene in time
    ts_aois_scenes = AOIFeatures.TimeStampedAOIScenes()

    # Init head movement
    head_movement_px = numpy.array((0, 0))
    head_movement_norm = 0

    # Init data timestamped in millisecond
    data_ts_ms = 0
    
    # Assess temporal performance
    loop_chrono = MiscFeatures.TimeProbe()
    gyroscope_chrono = MiscFeatures.TimeProbe()

    loop_ps = 0
    gyroscope_ps = 0

    def data_stream_callback(data_ts, data_object, data_object_type):

        nonlocal head_movement_px
        nonlocal head_movement_norm
        nonlocal data_ts_ms
        nonlocal gyroscope_chrono

        data_ts_ms = data_ts / 1e3

        match data_object_type:

            case 'Gyroscope':
                
                # Assess gyroscope stream performance
                gyroscope_chrono.lap()

                # Calculate head movement considering only head yaw and pitch
                head_movement = numpy.array(data_object.value)
                head_movement_px = head_movement.astype(int)
                head_movement_norm = numpy.linalg.norm(head_movement[0:2])

    tobii_data_stream.reading_callback = data_stream_callback

    # Start streaming
    tobii_controller.start_streaming()

    # Live video stream capture loop
    try:

        # Assess loop performance
        loop_chrono = MiscFeatures.TimeProbe()
        fps = 0

        # Detect head movement
        head_moving = False
        head_movement_last = 0.

        while tobii_video_stream.is_alive():

            # Read video stream
            video_ts, video_frame = tobii_video_stream.read()
            video_ts_ms = video_ts / 1e3

            # Copy video frame to edit visualisation on it without disrupting aruco tracking
            visu_frame = video_frame.copy()

            # Process video and data frame
            try:

                # Head movement detection hysteresis
                # TODO : pass the threshold value as argument
                if not head_moving and head_movement_norm > 50:
                    head_moving = True
                    
                if head_moving and head_movement_norm < 10:
                    head_moving = False

                # When head is moving, ArUco tracking could return bad pose estimation and so bad AOI scene projection
                if head_moving:

                    ts_aois_scenes[round(video_ts_ms)] = AOIScene2D.AOI2DScene()
                    
                    raise UserWarning('Head is moving')

                # Hide frame left and right borders before tracking to ignore markers outside focus area
                cv.rectangle(video_frame.matrix, (0, 0), (int(video_frame.width/6), int(video_frame.height)), (0, 0, 0), -1)
                cv.rectangle(video_frame.matrix, (int(video_frame.width*(1 - 1/6)), 0), (int(video_frame.width), int(video_frame.height)), (0, 0, 0), -1)

                # Track markers with pose estimation and draw them
                aruco_tracker.track(video_frame.matrix)
                aruco_tracker.draw_tracked_markers(visu_frame.matrix)

                # When no marker is detected, no AOI scene projection can't be done
                if aruco_tracker.tracked_markers_number == 0:

                    ts_aois_scenes[round(video_ts_ms)] = AOIScene2D.AOI2DScene()

                    raise UserWarning('No marker detected')

                # Store aoi 2D video for further scene merging
                aoi2D_dict = {}

                # Project 3D scene on each video frame and the visualisation frame
                for marker_id, marker in aruco_tracker.tracked_markers.items():

                    # Copy 3D scene related to detected marker
                    aoi3D_scene = aoi3D_scene_selector(marker_id)
                    
                    if aoi3D_scene == None:
                        continue

                    # Transform scene into camera referential
                    aoi3D_camera = aoi3D_scene.transform(marker.translation, marker.rotation)

                    # Get aoi inside vision cone field 
                    cone_vision_height_cm = 200 # cm
                    cone_vision_radius_cm = numpy.tan(numpy.deg2rad(TobiiSpecifications.VISUAL_HFOV / 2)) * cone_vision_height_cm

                    aoi3D_inside, aoi3D_outside = aoi3D_camera.vision_cone(cone_vision_radius_cm, cone_vision_height_cm)

                    # Keep only aoi inside vision cone field
                    aoi3D_scene = aoi3D_scene.copy(exclude=aoi3D_outside.keys())

                    # DON'T APPLY CAMERA DISTORSION : it projects points which are far from the frame into it
                    # This hack isn't realistic but as the gaze will mainly focus on centered AOI, where the distorsion is low, it is acceptable.
                    aoi2D_video_scene = aoi3D_scene.project(marker.translation, marker.rotation, aruco_camera.K)

                    # Store each 2D aoi for further scene merging
                    for name, aoi in aoi2D_video_scene.items():

                        if name not in aoi2D_dict.keys():
                            aoi2D_dict[name] = []

                        aoi2D_dict[name].append(aoi.clockwise())

                # Merge all 2D aoi into a single 2D scene
                aoi2D_merged_scene = AOI2DScene.AOI2DScene()
                for name, aoi_array in aoi2D_dict.items():
                    aoi2D_merged_scene[name] = numpy.sum(aoi_array, axis=0) / len(aoi_array)

                aoi2D_merged_scene.draw(visu_frame.matrix)

                # Store 2D merged scene at this time in millisecond
                ts_aois_scenes[round(video_ts_ms)] = aoi2D_merged_scene

                # Warn user when the merged scene is empty
                if len(aoi2D_merged_scene.keys()) == 0:

                    raise UserWarning('Scene is empty')

            # Write warning
            except UserWarning as w:

                cv.rectangle(visu_frame.matrix, (0, 100), (500, 150), (127, 127, 127), -1)
                cv.putText(visu_frame.matrix, str(w), (20, 140), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

            # Assess loop performance
            lap_time, lap_counter, elapsed_time = loop_chrono.lap()

            # Update fps each 10 loops
            if lap_counter >= 10:

                loop_ps = 1e3 * lap_counter / elapsed_time
                loop_chrono.restart()
            
                # Assess gyroscope streaming performance
                elapsed_time, lap_counter = gyroscope_chrono.end()
                gyroscope_ps = 1e3 * lap_counter / elapsed_time
                gyroscope_chrono.restart()

            # Draw head movement vector
            cv.line(visu_frame.matrix, (int(visu_frame.width/2), int(visu_frame.height/2)), (int(visu_frame.width/2) + head_movement_px[1], int(visu_frame.height/2) - head_movement_px[0]), (150, 150, 150), 3)
            
            # Draw center
            cv.line(visu_frame.matrix, (int(visu_frame.width/2) - 50, int(visu_frame.height/2)), (int(visu_frame.width/2) + 50, int(visu_frame.height/2)), (255, 150, 150), 1)
            cv.line(visu_frame.matrix, (int(visu_frame.width/2), int(visu_frame.height/2) - 50), (int(visu_frame.width/2), int(visu_frame.height/2) + 50), (255, 150, 150), 1)

            # Write stream timing
            cv.rectangle(visu_frame.matrix, (0, 0), (1100, 50), (63, 63, 63), -1)
            cv.putText(visu_frame.matrix, f'Data stream time: {int(data_ts_ms)} ms', (20, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
            cv.putText(visu_frame.matrix, f'Video delay: {int(data_ts_ms - video_ts_ms)} ms', (550, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
            cv.putText(visu_frame.matrix, f'Fps: {int(loop_ps)}', (950, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)

            cv.rectangle(visu_frame.matrix, (0, 50), (500, 100), (127, 127, 127), -1)
            cv.putText(visu_frame.matrix, f'Gyroscope fps: {int(gyroscope_ps)}', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

            cv.imshow(f'Stream ArUco AOI', visu_frame.matrix)
            
            # Close window using 'Esc' key
            if cv.waitKey(1) == 27:
                break

    # Exit on 'ctrl+C' interruption
    except KeyboardInterrupt:
        pass

    # Stop frame display
    cv.destroyAllWindows()

    # Stop streaming
    tobii_controller.stop_streaming()
    
if __name__ == '__main__':

    main()