aboutsummaryrefslogtreecommitdiff
path: root/src/argaze/utils/aruco_markers_scene_export.py
blob: 4518e48b636847190385f737f4ad02e4b39a09a3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
#!/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 image into it, detect ArUco markers belonging to a given dictionary and size into the selected image 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)
    image_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    image_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 image selection
    current_image_index = -1
    _, current_image = video_capture.read()
    next_image_index = int(args.start * video_fps)
    refresh = False

    # Hide help
    draw_help = False

    while not exit.status():

        # Select a new image and detect markers once
        if next_image_index != current_image_index or refresh:

            video_capture.set(cv2.CAP_PROP_POS_FRAMES, next_image_index)

            success, video_image = video_capture.read()

            if success:

                # Refresh once
                refresh = False

                current_image_index = video_capture.get(cv2.CAP_PROP_POS_FRAMES) - 1
                current_image_time = video_capture.get(cv2.CAP_PROP_POS_MSEC)

                # Detect markers
                aruco_detector.detect_markers(video_image)

                # 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)

                # Write scene detected markers
                cv2.putText(video_image, f'{list(aruco_detector.detected_markers.keys())}', (20, image_height-80), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                # Write timing
                cv2.putText(video_image, f'Time: {int(current_image_time)} ms', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                # Copy image
                current_image = video_image.copy()

        # Keep last image
        else:

            video_image = current_image.copy()

        # Draw detected markers
        aruco_detector.draw_detected_markers(video_image)

        # Write documentation
        cv2.putText(video_image, 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_image, (0, 50), (500, 300), (127, 127, 127), -1)
            cv2.putText(video_image, f'> Left arrow: previous image', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_image, f'> Right arrow: next image', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_image, 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 image with left arrow
        if key_pressed == 2:
            next_image_index -= 1

        # Select next image with right arrow
        if key_pressed == 3:
            next_image_index += 1

        # Clip image index
        if next_image_index < 0:
            next_image_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_image_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_image)

    # Close movie capture
    video_capture.release()

    # Stop image display
    cv2.destroyAllWindows()

if __name__ == '__main__':

    main()