diff options
Diffstat (limited to 'src/argaze/ArUcoMarkers/ArUcoTracker.py')
| -rw-r--r-- | src/argaze/ArUcoMarkers/ArUcoTracker.py | 229 |
1 files changed, 54 insertions, 175 deletions
diff --git a/src/argaze/ArUcoMarkers/ArUcoTracker.py b/src/argaze/ArUcoMarkers/ArUcoTracker.py index b887812..f5b1715 100644 --- a/src/argaze/ArUcoMarkers/ArUcoTracker.py +++ b/src/argaze/ArUcoMarkers/ArUcoTracker.py @@ -3,7 +3,7 @@ import json from collections import Counter -from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoCamera +from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoCamera import numpy import cv2 as cv @@ -45,14 +45,14 @@ ArUcoTrackerParameters = [ class ArUcoTracker(): """Track ArUco markers into a frame.""" - def __init__(self, aruco_dictionary_name: str, marker_length: float, camera: ArUcoCamera.ArUcoCamera): + def __init__(self, dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary, marker_size: float, camera: ArUcoCamera.ArUcoCamera): """Define which markers library to track and their size""" # load ArUco markers dictionary - self.__aruco_dict = ArUcoMarkersDictionary.ArUcoMarkersDictionary(aruco_dictionary_name) + self.__dictionary = dictionary # define marker length in centimeter - self.__marker_length = marker_length + self.__marker_size = marker_size # define camera self.__camera = camera @@ -62,19 +62,8 @@ class ArUcoTracker(): self.__detector_parameters.cornerRefinementMethod = aruco.CORNER_REFINE_CONTOUR # to get a better pose estimation self.__detector_parameters_loaded = {} - # define tracked markers data - self.__markers_corners = [] - self.__markers_ids = [] - self.__rvecs = [] - self.__tvecs = [] - self.__points = [] - - # define rejected markers data - self.__rejected_markers_corners = [] - self.__rejected_markers_ids = [] - self.__rejected_rvecs = [] - self.__rejected_tvecs = [] - self.__rejected_points = [] + # init tracked markers data + self.__tracked_markers = {} # define tracked board data self.__board = None @@ -84,7 +73,7 @@ class ArUcoTracker(): # define track metrics data self.__track_count = 0 - self.__tracked_markers = [] + self.__tracked_ids = [] self.__rejected_markers = [] def load_configuration_file(self, configuration_filepath): @@ -115,98 +104,56 @@ class ArUcoTracker(): def track(self, frame, estimate_pose = True, check_rotation = False): """Track ArUco markers in frame.""" + self.__tracked_markers = {} + markers_corners, markers_ids, markers_rvecs, markers_tvecs, markers_points = [], [], [], [], [] + # DON'T MIRROR FRAME : it makes the markers detection to fail - # detect markers from gray picture - gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) - self.__markers_corners, self.__markers_ids, rejectedPoints = aruco.detectMarkers(gray, self.__aruco_dict.get_markers(), parameters = self.__detector_parameters) - self.__rejected_markers_corners, __rejected_markers_ids = [], [] - - if len(self.__markers_corners) > 0 and estimate_pose: - - # markers pose estimation - self.__rvecs, self.__tvecs, self.__points = aruco.estimatePoseSingleMarkers(self.__markers_corners, self.__marker_length, self.__camera.get_K(), self.__camera.get_D()) - - # optional: check marker rotation as described in [this issue](https://github.com/opencv/opencv/issues/8813) - if check_rotation: - - valid_rotation_markers = [] - bad_rotation_markers = [] - for i, rvec in enumerate(self.__rvecs): - - tvec = self.__tvecs[i][0] - R, _ = cv.Rodrigues(rvec) - - zAxisPoint = (tvec.dot(R) + numpy.array([0., 0., 1.])).dot(R.T) - zAxis = zAxisPoint - tvec - - # TODO: How to describe the expected Z axis orientation ? - # In some situations, you can't provide such information. - - # !!! Here a description specific to SimOne cockpit !!! - zAxisExpectedDict = { - 1: numpy.array([0.5, 0.5, -1]), - 2: numpy.array([0.5, 0.5, -1]), - 3: numpy.array([1, -1, -1]), - 4: numpy.array([1, -1, -1]), - 5: numpy.array([1, -1, -1]), - 6: numpy.array([1, 1, -1]), - 7: numpy.array([1, 1, -1]), - 8: numpy.array([1, -1, -1]) - } - - zAxisExpected = zAxisExpectedDict[self.__markers_ids[i][0]] - - cosine_angle = numpy.dot(zAxis/numpy.linalg.norm(zAxis), zAxisExpected) - degree_angle = numpy.rad2deg(numpy.arccos(cosine_angle)) - ''' - print(self.__markers_ids[i][0]) - print('marker position: ', tvec) - print('zAxisPoint: ', zAxisPoint) - print('zAxis: ', zAxis) - print('zAxisExpected: ', zAxisExpected) - print('cosine_angle: ', cosine_angle) - print('degree_angle: ', degree_angle) - ''' - # Is the marker oriented as expected ? - if cosine_angle < 0 or cosine_angle > 1: - - #print('valid') - valid_rotation_markers.append(i) - - else: - - #print('bad') - bad_rotation_markers.append(i) - - # update track metrics - self.__rejected_markers.append(self.__markers_ids[i][0]) - - # keep markers with bad rotation - self.__rejected_markers_corners = tuple([self.__markers_corners[i] for i in bad_rotation_markers]) - self.__rejected_markers_ids = self.__markers_ids[bad_rotation_markers] - self.__rejected_rvecs = self.__rvecs[bad_rotation_markers] - self.__rejected_tvecs = self.__tvecs[bad_rotation_markers] - self.__rejected_points = self.__points[bad_rotation_markers] - - # keep markers with valid rotation - self.__markers_corners = tuple([self.__markers_corners[i] for i in valid_rotation_markers]) - self.__markers_ids = self.__markers_ids[valid_rotation_markers] - self.__rvecs = self.__rvecs[valid_rotation_markers] - self.__tvecs = self.__tvecs[valid_rotation_markers] - self.__points = self.__points[valid_rotation_markers] + # Track markers into gray picture + markers_corners, markers_ids, _ = aruco.detectMarkers(cv.cvtColor(frame, cv.COLOR_BGR2GRAY), self.__dictionary.get_markers(), parameters = self.__detector_parameters) - else: + if len(markers_corners) > 0: + + # Pose estimation is optional + if estimate_pose: + + markers_rvecs, markers_tvecs, markers_points = aruco.estimatePoseSingleMarkers(markers_corners, self.__marker_size, self.__camera.get_K(), self.__camera.get_D()) + + # Gather tracked markers data and update metrics + self.__track_count += 1 + + for i, marker_id in enumerate(markers_ids.T[0]): + + marker = ArUcoMarker.ArUcoMarker(self.__dictionary, marker_id, self.__marker_size) - self.__rvecs = [] - self.__tvecs = [] - self.__points = [] + marker.corners = markers_corners[i] - # update track metrics - self.__track_count += 1 - for marker_id in self.get_markers_ids(): - self.__tracked_markers.append(marker_id) + if estimate_pose: + marker.translation = markers_tvecs[i][0] + marker.rotation = markers_rvecs[i] + marker.points = markers_points[i] + self.__tracked_markers[marker_id] = marker + + self.__tracked_ids.append(marker_id) + + def get_tracked_markers(self): + """Access to tracked markers dictionary.""" + + return self.__tracked_markers + + def get_tracked_markers_number(self): + """Return tracked markers number.""" + + return len(list(self.__tracked_markers.keys())) + + def draw_tracked_markers(self, frame): + """Draw traked markers.""" + + for marker_id, marker in self.__tracked_markers.items(): + + marker.draw(frame, self.__camera.get_K(), self.__camera.get_D()) + def track_board(self, frame, board, expected_markers_number): """Track ArUco markers board in frame setting up the number of detected markers needed to agree detection.""" @@ -214,7 +161,7 @@ class ArUcoTracker(): # detect markers from gray picture gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) - self.__markers_corners, self.__markers_ids, rejectedPoints = aruco.detectMarkers(gray, self.__aruco_dict.get_markers(), parameters = self.__detector_parameters) + self.__markers_corners, self.__markers_ids, rejectedPoints = aruco.detectMarkers(gray, self.__dictionary.get_markers(), parameters = self.__detector_parameters) # if all board markers are detected if self.get_markers_number() == expected_markers_number: @@ -229,33 +176,6 @@ class ArUcoTracker(): self.__board_corners = [] self.__board_corners_ids = [] - def draw(self, frame): - """Draw tracked markers in frame.""" - - # draw detected markers square - if len(self.__markers_corners) > 0: - - # draw marker axis if pose has been estimated - if len(self.__rvecs) > 0: - - for (i, marker_id) in enumerate(self.__markers_ids): - - cv.drawFrameAxes(frame, self.__camera.get_K(), self.__camera.get_D(), self.__rvecs[i], self.__tvecs[i], self.__marker_length) - - aruco.drawDetectedMarkers(frame, self.__markers_corners, self.__markers_ids) - - # draw rejected markers square - if len(self.__rejected_markers_corners) > 0: - - # draw marker axis if pose has been estimated - if len(self.__rejected_rvecs) > 0: - - for (i, marker_id) in enumerate(self.__rejected_markers_ids): - - cv.drawFrameAxes(frame, self.__camera.get_K(), self.__camera.get_D(), self.__rejected_rvecs[i], self.__rejected_tvecs[i], self.__marker_length) - - aruco.drawDetectedMarkers(frame, self.__rejected_markers_corners, self.__rejected_markers_ids, borderColor=(0, 255, 255)) - def draw_board(self, frame): """Draw tracked board corners in frame.""" @@ -266,52 +186,11 @@ class ArUcoTracker(): def reset_track_metrics(self): """Enable marker tracking metrics.""" self.__track_count = 0 - self.__tracked_markers = [] + self.__tracked_ids = [] def get_track_metrics(self): """Get marker tracking metrics.""" - return self.__track_count, Counter(self.__tracked_markers), Counter(self.__rejected_markers) - - def get_markers_dictionay(self): - """Get tracked aruco markers dictionary.""" - return self.__aruco_dict - - def get_markers_number(self): - """Get tracked markers number.""" - return len(self.__markers_corners) - - def get_markers_ids(self): - """Get tracked markers identifers.""" - if self.__markers_ids is not None: - return [i[0] for i in self.__markers_ids] - else: - return [] - - def get_marker_index(self, marker_id): - """Retreive marker index of a given merker id. - Raise ValueError if not found.""" - - return list(self.__markers_ids).index(marker_id) - - def get_marker_corners(self, i): - """Get marker i corners.""" - return self.__markers_corners[i][0] - - def get_marker_center(self, i): - """Get marker i center coordinates.""" - return self.__markers_corners[i][0].mean(axis=0) - - def get_marker_rotation(self, i): - """Get marker i rotation vector.""" - return self.__rvecs[i] - - def get_marker_translation(self, i): - """Get marker i translation vector.""" - return self.__tvecs[i] - - def get_marker_points(self, i): - """Get marker i points.""" - return self.__points[i] + return self.__track_count, Counter(self.__tracked_ids), Counter(self.__rejected_markers) def get_board_corners_number(self): """Get tracked board corners number.""" |