diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/argaze/utils/aruco_markers_scene_export.py | 178 |
1 files changed, 178 insertions, 0 deletions
diff --git a/src/argaze/utils/aruco_markers_scene_export.py b/src/argaze/utils/aruco_markers_scene_export.py new file mode 100644 index 0000000..8045d24 --- /dev/null +++ b/src/argaze/utils/aruco_markers_scene_export.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python + +""" """ + +__author__ = "Théo de la Hogue" +__credits__ = [] +__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)" +__license__ = "BSD" + +import argparse +import time +import itertools + +from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoOpticCalibrator, ArUcoDetector, ArUcoScene +from argaze.utils import MiscFeatures + +import cv2 +import numpy + +def main(): + """ + Load a movie with ArUco markers inside and select a frame into it, detect ArUco markers belonging to a given dictionary and size into the selected frame thanks to given optic parameters and detector parameters then, export detected ArUco scene as .obj file. + """ + + # Manage arguments + parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0]) + parser.add_argument('movie', metavar='MOVIE', type=str, default=None, help='movie path') + parser.add_argument('dictionary', metavar='DICTIONARY', type=str, default=None, help='ArUco dictionary to detect') + parser.add_argument('marker_size', metavar='MARKER_SIZE', type=int, default=3, help='marker size in cm') + parser.add_argument('optic_parameters', metavar='OPTIC_PARAMETERS', type=str, default=None, help='Optic parameters from camera calibration process') + parser.add_argument('detector_parameters', metavar='DETECTOR_PARAMETERS', type=str, default=None, help='ArUco detector parameters') + + parser.add_argument('-s','--start', metavar='START', type=float, default=0., help='start time in second') + parser.add_argument('-o', '--output', metavar='OUT', type=str, default='.', help='export scene folder path') + args = parser.parse_args() + + # Load movie + video_capture = cv2.VideoCapture(args.movie) + + video_fps = video_capture.get(cv2.CAP_PROP_FPS) + frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + # Load ArUco dictionary + aruco_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(args.dictionary) + + # Load optic parameters + optic_parameters = ArUcoOpticCalibrator.OpticParameters.from_json(args.optic_parameters) + + # Load detector parameters + detector_parameters = ArUcoDetector.DetectorParameters.from_json(args.detector_parameters) + + # Create ArUco detector + aruco_detector = ArUcoDetector.ArUcoDetector(dictionary=aruco_dictionary, marker_size=args.marker_size, optic_parameters=optic_parameters, parameters=detector_parameters) + + # Create empty ArUco scene + aruco_scene = None + + # Create a window to display AR environment + window_name = "Export ArUco scene" + cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE) + + # Enable exit signal handler + exit = MiscFeatures.ExitSignalHandler() + + # Init frame selection + current_frame_index = -1 + _, current_frame = video_capture.read() + next_frame_index = int(args.start * video_fps) + refresh = False + + # Hide help + draw_help = False + + while not exit.status(): + + # Select a new frame and detect markers once + if next_frame_index != current_frame_index or refresh: + + video_capture.set(cv2.CAP_PROP_POS_FRAMES, next_frame_index) + + success, video_frame = video_capture.read() + + if success: + + # Refresh once + refresh = False + + current_frame_index = video_capture.get(cv2.CAP_PROP_POS_FRAMES) - 1 + current_frame_time = video_capture.get(cv2.CAP_PROP_POS_MSEC) + + # Detect markers + aruco_detector.detect_markers(video_frame) + + # Estimate markers pose + aruco_detector.estimate_markers_pose() + + # Build aruco scene from detected markers + aruco_scene = ArUcoScene.ArUcoScene(args.marker_size, aruco_dictionary, aruco_detector.detected_markers) + + print(aruco_scene) + + # Write scene detected markers + cv2.putText(video_frame, f'{list(aruco_detector.detected_markers.keys())}', (20, frame_height-80), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) + + # Write timing + cv2.putText(video_frame, f'Time: {int(current_frame_time)} ms', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) + + # Copy frame + current_frame = video_frame.copy() + + # Keep last frame + else: + + video_frame = current_frame.copy() + + # Draw detected markers + aruco_detector.draw_detected_markers(video_frame) + + # Write documentation + cv2.putText(video_frame, f'Press \'h\' for help', (950, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) + + if draw_help: + + cv2.rectangle(video_frame, (0, 50), (500, 300), (127, 127, 127), -1) + cv2.putText(video_frame, f'> Left arrow: previous frame', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) + cv2.putText(video_frame, f'> Right arrow: next frame', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) + cv2.putText(video_frame, f'> Ctrl+s: export ArUco scene', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) + + key_pressed = cv2.waitKey(10) + + #if key_pressed != -1: + # print(key_pressed) + + # Select previous frame with left arrow + if key_pressed == 2: + next_frame_index -= 1 + + # Select next frame with right arrow + if key_pressed == 3: + next_frame_index += 1 + + # Clip frame index + if next_frame_index < 0: + next_frame_index = 0 + + # Switch help mode with h key + if key_pressed == 104: + draw_help = not draw_help + + # Save selected marker edition using 'Ctrl + s' + if key_pressed == 19: + + if aruco_scene: + + aruco_scene.to_obj(f'{args.output}/{int(current_frame_time)}-aruco_scene.obj') + print(f'ArUco scene saved into {args.output}') + + else: + + print(f'No ArUco scene to export') + + # Close window using 'Esc' key + if key_pressed == 27: + break + + # Display video + cv2.imshow(window_name, video_frame) + + # Close movie capture + video_capture.release() + + # Stop frame display + cv2.destroyAllWindows() + +if __name__ == '__main__': + + main()
\ No newline at end of file |