diff options
Diffstat (limited to 'src/argaze/ArUcoMarkers/ArUcoSet.py')
| -rw-r--r-- | src/argaze/ArUcoMarkers/ArUcoSet.py | 448 |
1 files changed, 448 insertions, 0 deletions
diff --git a/src/argaze/ArUcoMarkers/ArUcoSet.py b/src/argaze/ArUcoMarkers/ArUcoSet.py new file mode 100644 index 0000000..2eeea32 --- /dev/null +++ b/src/argaze/ArUcoMarkers/ArUcoSet.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python + +from typing import Tuple +from dataclasses import dataclass, field +import json +import math +import itertools + +from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoCamera + +import numpy +import cv2 as cv +import cv2.aruco as aruco + +@dataclass +class ArUcoSetPlace(): + """Define set place pose and marker.""" + + translation: numpy.array + """Position in set referential.""" + + rotation: numpy.array + """Rotation in set referential.""" + + marker: dict + """ArUco marker linked to the place.""" + +@dataclass +class ArUcoSet(): + """Define abstract class to handle specific ArUco markers set and estimate its pose.""" + + dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary + """ArUco dictionary of set markers.""" + + marker_size: int = field(init=False) + """Size of markers in centimeter.""" + + places: dict = field(init=False, default_factory=dict) + """All named places of the set and their ArUco markers.""" + + angle_tolerance: float = field(init=False) + """Angle error tolerance allowed to validate place pose in degree.""" + + distance_tolerance: float = field(init=False) + """Distance error tolerance allowed to validate place pose in centimeter.""" + + def __init__(self, configuration_filepath): + """Define set from a .json file.""" + + with open(configuration_filepath) as configuration_file: + + # Deserialize .json + # TODO find a better way + configuration = json.load(configuration_file) + + # Load dictionary + self.dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(configuration['dictionary']) + + # Load marker size + self.marker_size = configuration['marker_size'] + + # Load places + self.places = {} + for name, place in configuration['places'].items(): + marker = ArUcoMarker.ArUcoMarker(self.dictionary, place['marker'], self.marker_size) + self.places[name] = ArUcoSetPlace(numpy.array(place['translation']).astype(numpy.float32), numpy.array(place['rotation']).astype(numpy.float32), marker) + + # Load angle tolerance + self.angle_tolerance = configuration['angle_tolerance'] + + # Load distance tolerance + self.distance_tolerance = configuration['distance_tolerance'] + + # Init pose data + self._translation = numpy.zeros(3) + self._rotation = numpy.zeros(3) + self._succeded = False + self._validity = 0 + + # Process markers ids to speed up further calculations + self.__identifier_cache = {} + for name, place in self.places.items(): + self.__identifier_cache[place.marker.identifier] = name + + # Process each place pose to speed up further calculations + self.__translation_cache = {} + for name, place in self.places.items(): + self.__translation_cache[name] = place.translation + + # Process each place rotation matrix to speed up further calculations + self.__rotation_cache = {} + for name, place in self.places.items(): + + # Create intrinsic rotation matrix + R = self.__make_rotation_matrix(*place.rotation) + + assert(self.__is_rotation_matrix(R)) + + # Store rotation matrix + self.__rotation_cache[name] = R + + # Process axis-angle between place combination to speed up further calculations + self.__angle_cache = {} + for (A_name, A_place), (B_name, B_place) in itertools.combinations(self.places.items(), 2): + + A = self.__rotation_cache[A_name] + B = self.__rotation_cache[B_name] + + if numpy.array_equal(A, B): + + print('A.all() == B.all()') + angle = 0. + + else: + + # Rotation matrix from A place to B place + AB = B.dot(A.T) + + assert(self.__is_rotation_matrix(AB)) + + # Calculate axis-angle representation of AB rotation matrix + angle = numpy.rad2deg(numpy.arccos((numpy.trace(AB) - 1) / 2)) + + try: + self.__angle_cache[A_name][B_name] = angle + except: + self.__angle_cache[A_name] = {B_name: angle} + + try: + self.__angle_cache[B_name][A_name] = angle + except: + self.__angle_cache[B_name] = {A_name: angle} + + # Process distance between each place combination to speed up further calculations + self.__distance_cache = {} + for (A_name, A_place), (B_name, B_place) in itertools.combinations(self.places.items(), 2): + + A = self.__translation_cache[A_name] + B = self.__translation_cache[B_name] + + # Calculate axis-angle representation of AB rotation matrix + distance = numpy.linalg.norm(B - A) + + try: + self.__distance_cache[A_name][B_name] = distance + except: + self.__distance_cache[A_name] = {B_name: distance} + + try: + self.__distance_cache[B_name][A_name] = distance + except: + self.__distance_cache[B_name] = {A_name: distance} + + def print_cache(self): + """Print pre-processed data.""" + + print('\nIdentifier cache:') + for i, name in self.__identifier_cache.items(): + print(f'- {i}: {name}') + + print('\nTranslation cache:') + for name, item in self.__translation_cache.items(): + print(f'- {name}: {item}') + + print('\nRotation cache:') + for name, item in self.__rotation_cache.items(): + print(f'- {name}:\n{item}') + + print('\nAngle cache:') + for A_name, A_angle_cache in self.__angle_cache.items(): + for B_name, angle in A_angle_cache.items(): + print(f'- {A_name}/{B_name}: {angle:3f}') + + print('\nDistance cache:') + for A_name, A_distance_cache in self.__distance_cache.items(): + for B_name, distance in A_distance_cache.items(): + print(f'- {A_name}/{B_name}: {distance:3f}') + + def __make_rotation_matrix(self, x, y, z): + + # Create rotation matrix around x axis + c = numpy.cos(numpy.deg2rad(x)) + s = numpy.sin(numpy.deg2rad(x)) + Rx = numpy.array([[1, 0, 0], [0, c, -s], [0, s, c]]) + + # Create rotation matrix around y axis + c = numpy.cos(numpy.deg2rad(y)) + s = numpy.sin(numpy.deg2rad(y)) + Ry = numpy.array([[c, 0, s], [0, 1, 0], [-s, 0, c]]) + + # Create rotation matrix around z axis + c = numpy.cos(numpy.deg2rad(z)) + s = numpy.sin(numpy.deg2rad(z)) + Rz = numpy.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) + + # Return intrinsic rotation matrix + return Rx.dot(Ry.dot(Rz)) + + def __is_rotation_matrix(self, R): + """Checks if a matrix is a valid rotation matrix.""" + + I = numpy.identity(3, dtype = R.dtype) + return numpy.linalg.norm(I - numpy.dot(R.T, R)) < 1e-6 + + def __normalise_place_pose(self, name, place, F): + + # Transform place rotation into set rotation vector + R = self.__rotation_cache[name] + rvec, _ = cv.Rodrigues(F.dot(R)) + + #print(f'{name} rotation vector: {rvec[0][0]:3f} {rvec[1][0]:3f} {rvec[2][0]:3f}') + + # Transform place translation into set translation vector + OF = place.translation + T = self.__translation_cache[name] + FC = R.dot(F.dot(-T)) + + tvec = OF + FC + + #print(f'{name} translation vector: {tvec[0]:3f} {tvec[1]:3f} {tvec[2]:3f}') + + return rvec, tvec + + def estimate_pose(self, tracked_markers) -> Tuple[numpy.array, numpy.array, bool, int]: + """Estimate set pose from tracked markers (cf ArUcoTracker.track()) + + * **Returns:** + - translation vector + - rotation vector + - pose estimation success status + - the number of places used to estimate the pose as validity score + """ + + # Init pose data + self._translation = numpy.zeros(3) + self._rotation = numpy.zeros(3) + self._succeded = False + self._validity = 0 + + # Don't try to estimate pose if there is no tracked markers + if len(tracked_markers) == 0: + + return self._translation, self._rotation, self._succeded, self._validity + + # Look for places related to tracked markers + tracked_places = {} + for (marker_id, marker) in tracked_markers.items(): + + try: + name = self.__identifier_cache[marker_id] + tracked_places[name] = marker + + except KeyError: + continue + + #print('-------------- ArUcoSet pose estimation --------------') + + # Pose validity checking is'nt possible when only one place of the set is tracked + if len(tracked_places.keys()) == 1: + + # Get set pose from to the unique place pose + name, place = tracked_places.popitem() + F, _ = cv.Rodrigues(place.rotation) + + self._rotation, self._translation = self.__normalise_place_pose(name, place, F) + self._succeded = True + self._validity = 1 + + #print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!') + #print(f'arset rotation vector: {self._rotation[0][0]:3f} {self._rotation[1][0]:3f} {self._rotation[2][0]:3f}') + #print(f'arset translation vector: {self._translation[0]:3f} {self._translation[1]:3f} {self._translation[2]:3f}') + + # Pose validity checking processes places two by two + else: + + valid_places = [] + valid_rvecs = [] + valid_tvecs = [] + + for (A_name, A_place), (B_name, B_place) in itertools.combinations(tracked_places.items(), 2): + + #print(f'** {A_name} > {B_name}') + + # Get place rotation estimation + # Use rotation matrix instead of rotation vector + A, _ = cv.Rodrigues(A_place.rotation) + B, _ = cv.Rodrigues(B_place.rotation) + + # Rotation matrix from A place to B place + AB = B.dot(A.T) + + assert(self.__is_rotation_matrix(AB)) + + # Calculate axis-angles representation of AB rotation matrix + angle = numpy.rad2deg(numpy.arccos((numpy.trace(AB) - 1) / 2)) + expected_angle = self.__angle_cache[A_name][B_name] + + #print('angle:', angle) + #print('expected angle:', expected_angle) + + # Calculate distance between A place center and B place center + distance = numpy.linalg.norm(A_place.translation - B_place.translation) + expected_distance = self.__distance_cache[A_name][B_name] + + #print('distance: ', distance) + #print('expected distance: ', expected_distance) + + # Check angle and distance according given tolerance then normalise place pose + valid_angle = math.isclose(angle, expected_angle, abs_tol=self.angle_tolerance) + valid_distance = math.isclose(distance, expected_distance, abs_tol=self.distance_tolerance) + + if valid_angle and valid_distance: + + if A_name not in valid_places: + + # Remember this place is already validated + valid_places.append(A_name) + + rvec, tvec = self.__normalise_place_pose(A_name, A_place, A) + + # Store normalised place pose + valid_rvecs.append(rvec) + valid_tvecs.append(tvec) + + if B_name not in valid_places: + + # Remember this place is already validated + valid_places.append(B_name) + + rvec, tvec = self.__normalise_place_pose(B_name, B_place, B) + + # Store normalised place pose + valid_rvecs.append(rvec) + valid_tvecs.append(tvec) + + if len(valid_places) > 1: + + # Consider arset rotation as the mean of all valid translations + # !!! WARNING !!! This is a bad hack : processing rotations average is a very complex problem that needs to well define the distance calculation method before. + self._rotation = numpy.mean(numpy.array(valid_rvecs), axis=0) + + # Consider arset translation as the mean of all valid translations + self._translation = numpy.mean(numpy.array(valid_tvecs), axis=0) + + #print(':::::::::::::::::::::::::::::::::::::::::::::::::::') + #print(f'arset rotation vector: {self._rotation[0][0]:3f} {self._rotation[1][0]:3f} {self._rotation[2][0]:3f}') + #print(f'arset translation vector: {self._translation[0]:3f} {self._translation[1]:3f} {self._translation[2]:3f}') + + self._succeded = True + self._validity = len(valid_places) + + #print('----------------------------------------------------') + + return self._translation, self._rotation, self._succeded, self._validity + + @property + def translation(self) -> numpy.array: + """Access to set translation vector. + + .. warning:: + Setting set translation vector implies succeded status to be True and validity score to be 0.""" + + return self._translation + + @translation.setter + def translation(self, tvec): + + self._translation = tvec + self._succeded = True + self._validity = 0 + + @property + def rotation(self) -> numpy.array: + """Access to set rotation vector. + + .. warning:: + Setting set rotation vector implies succeded status to be True and validity score to be 0.""" + + return self._translation + + @rotation.setter + def rotation(self, rvec): + + self._rotation = rvec + self._succeded = True + self._validity = 0 + + @property + def succeded(self) -> bool: + """Access to set pose estimation succeded status.""" + + return self._succeded + + @property + def validity(self) -> int: + """Access to set pose estimation validity score.""" + + return self._validity + + def draw(self, frame, K, D, draw_places=True): + """Draw set axis and places.""" + + l = self.marker_size / 2 + ll = self.marker_size + + # Select color according validity score + n = 95 * self._validity if self._validity < 2 else 0 + f = 159 * self._validity if self._validity < 2 else 255 + + try: + + # Draw axis + axisPoints = numpy.float32([[ll, 0, 0], [0, ll, 0], [0, 0, ll], [0, 0, 0]]).reshape(-1, 3) + axisPoints, _ = cv.projectPoints(axisPoints, self._rotation, self._translation, K, D) + axisPoints = axisPoints.astype(int) + + cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (n,n,f), 5) # X (red) + cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 5) # Y (green) + cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 5) # Z (blue) + + # Draw places (optional) + if draw_places: + + for name, place in self.places.items(): + + if name != "top": + continue + + T = self.__translation_cache[name] + R = self.__rotation_cache[name] + + placePoints = (T + numpy.float32([R.dot([-l, -l, 0]), R.dot([l, -l, 0]), R.dot([l, l, 0]), R.dot([-l, l, 0])])).reshape(-1, 3) + placePoints, _ = cv.projectPoints(placePoints, self._rotation, self._translation, K, D) + placePoints = placePoints.astype(int) + + cv.line(frame, tuple(placePoints[0].ravel()), tuple(placePoints[1].ravel()), (f,f,f), 2) + cv.line(frame, tuple(placePoints[1].ravel()), tuple(placePoints[2].ravel()), (f,f,f), 2) + cv.line(frame, tuple(placePoints[2].ravel()), tuple(placePoints[3].ravel()), (f,f,f), 2) + cv.line(frame, tuple(placePoints[3].ravel()), tuple(placePoints[0].ravel()), (f,f,f), 2) + + except Exception as e: + + print(e) + print(self._translation) + print(self._rotation) + print(self._succeded) + print(self._validity) + print(axisPoints) |