From 13bc2c352c6fe94163b7868a697b6257fcfba3fa Mon Sep 17 00:00:00 2001 From: Théo de la Hogue Date: Wed, 13 Apr 2022 16:56:11 +0200 Subject: Adding TobiiVideoOutput to export data visualisation into a video file --- src/argaze/TobiiGlassesPro2/TobiiVideo.py | 42 ++++++++++++++++++++-- src/argaze/utils/README.md | 2 +- .../utils/export_tobii_segment_aruco_rois.py | 25 ++++++++++--- 3 files changed, 62 insertions(+), 7 deletions(-) diff --git a/src/argaze/TobiiGlassesPro2/TobiiVideo.py b/src/argaze/TobiiGlassesPro2/TobiiVideo.py index 4d684ec..b3e11f3 100644 --- a/src/argaze/TobiiGlassesPro2/TobiiVideo.py +++ b/src/argaze/TobiiGlassesPro2/TobiiVideo.py @@ -49,6 +49,9 @@ class TobiiVideoSegment(): def get_height(self): return self.__height + def get_stream(self): + return self.__stream + def frames(self, vts_data_buffer = None): """Access to frame iterator and optionnaly setup vide / data timestamp synchronisation through vts data buffer.""" @@ -88,7 +91,7 @@ class TobiiVideoSegment(): video_ts += self.__vts_offset # return micro second timestamp and frame data - return video_ts, TobiiVideoFrame(frame.to_ndarray(format='bgr24'), frame.width, frame.height) + return video_ts, TobiiVideoFrame(frame.to_ndarray(format='rgb24'), frame.width, frame.height) class TobiiVideoStream(threading.Thread): """Capture Tobii Glasses Pro 2 video camera stream.""" @@ -158,7 +161,7 @@ class TobiiVideoStream(threading.Thread): self.__read_lock.acquire() # store frame time, matrix, width, height and pts into a tuple - self.__frame_tuple = (frame.time, frame.to_ndarray(format='bgr24'), frame.width, frame.height) + self.__frame_tuple = (frame.time, frame.to_ndarray(format='rgb24'), frame.width, frame.height) # unlock frame access self.__read_lock.release() @@ -179,3 +182,38 @@ class TobiiVideoStream(threading.Thread): self.__read_lock.release() return int(frame_tuple[0] * 1000000), TobiiVideoFrame(frame_tuple[1], frame_tuple[2], frame_tuple[3]) + +class TobiiVideoOutput(): + """Export a video file at the same format than a given referent stream.""" + # TODO : Make a generic video managment to handle video from any device (not only Tobii) + + def __init__(self, output_video_path: str, referent_stream: av.stream.Stream): + """Create a video file""" + + self.__output_video_path = output_video_path + self.__container = av.open(self.__output_video_path, 'w') + self.__stream = self.__container.add_stream(\ + referent_stream.codec_context.name, \ + width=referent_stream.codec_context.width, \ + height=referent_stream.codec_context.height, \ + rate=referent_stream.codec_context.framerate, \ + gop_size=referent_stream.codec_context.gop_size, \ + pix_fmt=referent_stream.codec_context.pix_fmt, \ + bit_rate=referent_stream.codec_context.bit_rate) + + def get_path(self): + return self.__output_video_path + + def write(self, frame): + """Write a frame into the output video file""" + + formated_frame = av.VideoFrame.from_ndarray(frame, format='rgb24') + formated_frame.reformat(format=self.__stream.codec_context.pix_fmt, interpolation=None) + self.__container.mux(self.__stream.encode(formated_frame)) + + def close(self): + """End the writing of the video file""" + + self.__container.mux(self.__stream.encode()) + self.__container.close() + diff --git a/src/argaze/utils/README.md b/src/argaze/utils/README.md index 99b8717..dc40f30 100644 --- a/src/argaze/utils/README.md +++ b/src/argaze/utils/README.md @@ -72,7 +72,7 @@ python ./src/argaze/utils/replay_tobii_session.py -s SEGMENT_PATH python ./src/argaze/utils/export_tobii_segment_fixations.py -s SEGMENT_PATH ``` -- Track ArUco markers into a Tobii camera video segment (replace SEGMENT_PATH). Load an roi scene (replace ROI_SCENE) .obj file, position it virtually relatively to any detected ArUco markers and project the scene into camera frame. Then, detect if Tobii gaze point is inside any ROI. +- Track ArUco markers into a Tobii camera video segment (replace SEGMENT_PATH). Load an roi scene (replace ROI_SCENE) .obj file, position it virtually relatively to any detected ArUco markers and project the scene into camera frame. Then, detect if Tobii gaze point is inside any ROI. Export ROIs video and data. ``` python ./src/argaze/utils/export_tobii_segment_aruco_rois.py -s SEGMENT_PATH -c export/tobii_camera.json -m 7.5 -r ROI_SCENE diff --git a/src/argaze/utils/export_tobii_segment_aruco_rois.py b/src/argaze/utils/export_tobii_segment_aruco_rois.py index 343a107..5314141 100644 --- a/src/argaze/utils/export_tobii_segment_aruco_rois.py +++ b/src/argaze/utils/export_tobii_segment_aruco_rois.py @@ -15,17 +15,20 @@ import cv2 as cv def main(): """ - Replay Tobii segment video + Track any ArUco marker into Tobii Glasses Pro 2 segment video file. + From a loaded ROI scene .obj file, position the scene virtually relatively to any detected ArUco markers and project the scene into camera frame. + Then, detect if Tobii gaze point is inside any ROI. + Export ROIs video and data. """ - # manage arguments + # Manage arguments parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0]) parser.add_argument('-s', '--segment_path', metavar='SEGMENT_PATH', type=str, default=None, help='segment path') parser.add_argument('-c', '--camera_calibration', metavar='CAM_CALIB', type=str, default='tobii_camera.json', help='json camera calibration filepath') parser.add_argument('-r', '--roi_scene', metavar='ROI_SCENE', type=str, default='roi3D_scene.obj', help='obj roi scene filepath') parser.add_argument('-d', '--dictionary', metavar='DICT', type=str, default='DICT_ARUCO_ORIGINAL', help='aruco marker dictionnary') parser.add_argument('-m', '--marker_size', metavar='MKR', type=float, default=6, help='aruco marker size (cm)') - parser.add_argument('-o', '--output', metavar='OUT', type=str, default=None, help='destination path (segment folder by default)') + parser.add_argument('-o', '--output', metavar='OUT', type=str, default=None, help='destination folder path (segment folder by default)') args = parser.parse_args() if args.segment_path != None: @@ -38,11 +41,13 @@ def main(): os.makedirs(os.path.dirname(args.output)) print(f'{os.path.dirname(args.output)} folder created') - rois_filepath = args.output + rois_filepath = f'{args.output}/rois.json' + video_filepath = f'{args.output}/fullstream+visu.mp4' else: rois_filepath = f'{args.segment_path}/rois.json' + video_filepath = f'{args.segment_path}/fullstream+visu.mp4' # Load a tobii segment tobii_segment = TobiiEntities.TobiiSegment(args.segment_path) @@ -59,6 +64,9 @@ def main(): tobii_ts_gaze_positions = tobii_segment_data.gidx_l_gp print(f'{len(tobii_ts_gaze_positions)} gaze positions loaded') + # Prepare video exportation at the same format than segment video + output_video = TobiiVideo.TobiiVideoOutput(video_filepath, tobii_segment_video.get_stream()) + # Create aruco camera aruco_camera = ArUcoCamera.ArUcoCamera() aruco_camera.load_calibration_file(args.camera_calibration) @@ -132,8 +140,12 @@ def main(): if cv.waitKey(1) == 27: break + # Display video cv.imshow(f'Segment {tobii_segment.get_id()} video', video_frame.matrix) + # Write video + output_video.write(video_frame.matrix) + # Exit on 'ctrl+C' interruption except KeyboardInterrupt: pass @@ -141,6 +153,11 @@ def main(): # Stop frame display cv.destroyAllWindows() + # End output video file + output_video.close() + + print(f'ROIs video saved into {video_filepath}') + # Export 2D rois roi2D_timestamped_buffer.export_as_json(rois_filepath) -- cgit v1.1