Adding ArUco axis feature to estimate pose using orthogonal axis defined by ArUco markers.

2 files changed, 204 insertions, 164 deletions

diff --git a/src/argaze/ArScene.py b/src/argaze/ArScene.py
index f1cba12..3702c6b 100644
--- a/src/argaze/ArScene.py
+++ b/src/argaze/ArScene.py
@@ -57,6 +57,15 @@ class ArScene():
 	aoi_scene: AOI3DScene.AOI3DScene = field(init=False, default_factory=AOI3DScene.AOI3DScene)
 	"""AOI 3D scene ..."""
 
+	angle_tolerance: float
+	"""Angle error tolerance allowed to validate marker pose in degree."""
+
+	distance_tolerance: float
+	"""Distance error tolerance allowed to validate marker pose in centimeter."""
+
+	aruco_axis: dict
+	"""Dictionary of orthogonal axis where each axis is defined by list of 3 markers identifier (first is origin)."""
+
 	def __init__(self, **kwargs):
 
 		self.aruco_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(kwargs.pop('aruco_dictionary'))
@@ -67,15 +76,15 @@ class ArScene():
 
 		self.aruco_tracker = ArUcoTracker.ArUcoTracker(self.aruco_dictionary, self.aruco_marker_size, self.aruco_camera, **kwargs.pop('aruco_tracker'))
 
-		# Check aruco_scene.places value type
-		aruco_scene_places_value = kwargs['aruco_scene']['places']
+		# Check aruco_scene value type
+		aruco_scene_value = kwargs.pop('aruco_scene')
 
 		# str: relative path to .obj file
-		if type(aruco_scene_places_value) == str:
+		if type(aruco_scene_value) == str:
 
-			kwargs['aruco_scene']['places'] = os.path.join(self.__current_directory, aruco_scene_places_value)
+			aruco_scene_value = os.path.join(self.__current_directory, aruco_scene_value)
 
-		self.aruco_scene = ArUcoScene.ArUcoScene(self.aruco_dictionary, self.aruco_marker_size, **kwargs.pop('aruco_scene'))
+		self.aruco_scene = ArUcoScene.ArUcoScene(self.aruco_dictionary, self.aruco_marker_size, aruco_scene_value)
 
 		# Check aoi_scene value type
 		aoi_scene_value = kwargs.pop('aoi_scene')
@@ -90,8 +99,32 @@ class ArScene():
 		else:
 			self.aoi_scene = AOI3DScene.AOI3DScene(aoi_scene_value)
 
+		# TODO : allow to define AOI using marker corner position
+		'''
+    	"aruco_aoi": {
+        	"Screen": [[4, 3], [5, 2], [6, 1], [7, 0]] # [[MARKER_ID, MARKER_CORNER_ID], ...]
+    	}
+		'''
+
+		# Init axis markers
+		self.aruco_axis = {}
+
+		# Update all attributes from arguments
 		self.__dict__.update(kwargs)
 
+		# Convert axis markers identifier into ARUCO_DICT_NAME#ID string
+		aruco_axis_string = {}
+		for axis_name, markers_id in self.aruco_axis.items():
+
+			# Estimate pose from axis markers
+			aruco_axis_names = []
+			for marker_id in markers_id:
+				aruco_axis_names.append(f'{self.aruco_dictionary.name}#{marker_id}')
+
+			aruco_axis_string[axis_name] = aruco_axis_names
+
+		self.aruco_axis = aruco_axis_string
+
 	@classmethod
 	def from_json(self, json_filepath: str) -> ArSceneType:
 		"""Load ArGaze project from .json file."""
@@ -115,28 +148,67 @@ class ArScene():
 
 		return output
 
-	def estimate_pose(self, frame) -> Tuple[numpy.array, numpy.array, list]:
+	def estimate_pose(self, frame) -> Tuple[numpy.array, numpy.array, dict]:
 		"""Estimate scene pose from ArUco markers into frame.
 
 		* **Returns:** 
-            - scene translation vector
-            - scene rotation matrix
-            - list of all tracked markers considered as consistent and used to estimate the pose
-        """
+			- scene translation vector
+			- scene rotation matrix
+			- dict of markers used to estimate the pose
+		"""
 
 		self.aruco_tracker.track(frame)
 
-		# When no marker is detected, no pose estimation can't be done
+		# Pose estimation fails when no marker is detected
 		if len(self.aruco_tracker.tracked_markers) == 0:
 			
 			raise PoseEstimationFailed('No marker detected')
 
-		# Estimate scene pose from tracked markers
-		tvec, rmat, success, consistent_markers, unconsistencies = self.aruco_scene.estimate_pose(self.aruco_tracker.tracked_markers)
+		scene_markers, _ = self.aruco_scene.filter_markers(self.aruco_tracker.tracked_markers)
+
+		# Pose estimation fails when no marker belongs to the scene
+		if len(scene_markers) == 0:
+
+			raise PoseEstimationFailed('No marker belongs to the scene')
+
+		# Estimate scene pose from unique marker transformations
+		elif len(scene_markers) == 1:
+
+			tvec, rmat = self.aruco_scene.estimate_pose_from_any_markers(scene_markers)
+
+			return tvec, rmat, scene_markers
+
+		# Try to estimate scene pose from 3 markers defining an orthogonal axis
+		elif len(scene_markers) >= 3 and len(self.aruco_axis) > 0:
+
+			for axis_name, markers_names in self.aruco_axis.items():
+
+				try:
+
+					axis_markers = []
+					for name in markers_names:
+						axis_markers.append((name, scene_markers[name]))
+
+					tvec, rmat = self.aruco_scene.estimate_pose_from_axis_markers(axis_markers)
+
+					return tvec, rmat, axis_markers
+
+				except:
+					pass
+
+			raise PoseEstimationFailed('No marker axis')
+
+		# Otherwise, check markers consistency
+		consistent_markers, unconsistent_markers, unconsistencies = self.aruco_scene.check_markers_consistency(scene_markers, self.angle_tolerance, self.distance_tolerance)
+
+		# Pose estimation fails when no marker passes consistency checking
+		if len(consistent_markers) == 0:
 
-		if not success:
 			raise PoseEstimationFailed('Unconsistent marker poses', unconsistencies)
 
+		# Otherwise, estimate scene pose from all markers transformations
+		tvec, rmat = self.aruco_scene.estimate_pose_from_any_markers(consistent_markers)
+
 		return tvec, rmat, consistent_markers
 
 	def project(self, tvec, rmat, visual_hfov=0) -> AOI2DSceneType:
diff --git a/src/argaze/ArUcoMarkers/ArUcoScene.py b/src/argaze/ArUcoMarkers/ArUcoScene.py
index 06d63d3..3990904 100644
--- a/src/argaze/ArUcoMarkers/ArUcoScene.py
+++ b/src/argaze/ArUcoMarkers/ArUcoScene.py
@@ -38,21 +38,12 @@ class Place():
 class ArUcoScene():
     """Define abstract class to handle group of ArUco markers as one unique spatial entity and estimate its pose."""
 
-    angle_tolerance: float
-    """Angle error tolerance allowed to validate place pose in degree."""
-
-    distance_tolerance: float
-    """Distance error tolerance allowed to validate place pose in centimeter."""
-
-    def __init__(self, dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary, marker_size: float, angle_tolerance: float, distance_tolerance: float, places: dict | str = None):
+    def __init__(self, dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary, marker_size: float, places: dict | str = None):
         """Define scene attributes."""
 
         self.__dictionary = dictionary
         self.__marker_size = marker_size
 
-        self.angle_tolerance = angle_tolerance
-        self.distance_tolerance = distance_tolerance
-
         # NEVER USE {} as default function argument
         self.places = places
 
@@ -82,8 +73,6 @@ class ArUcoScene():
         # Init pose data
         self._translation = numpy.zeros(3)
         self._rotation = numpy.zeros(3)
-        self._succeded = False
-        self._consistent_markers = 0
 
         # Process markers ids to speed up further calculations
         self.__identifier_cache = {}
@@ -296,166 +285,164 @@ class ArUcoScene():
         except IOError:
             raise IOError(f'File not found: {obj_filepath}')
 
-    def __normalise_marker_pose(self, place, R, T):
+    def filter_markers(self, tracked_markers) -> Tuple[dict, dict]:
+        """Sort markers belonging to the scene from a given tracked markers list (cf ArUcoTracker.track()).
 
-        # Transform marker rotation into scene rotation matrix
-        Rs = R.dot(place.rotation.T)
+        * **Returns:**
+            - dict of markers belonging to this scene
+            - dict of remaining markers not belonging to this scene
+        """
 
-        # Transform marker translation into scene translation vector
-        #Ts = T + place.rotation.dot(R.dot(-place.translation))
-        #Ts = T - place.rotation.dot(R.dot(place.translation))
-        #Ts = T - place.translation.dot(place.rotation.T)
-        Ts = T - place.translation.dot(place.rotation).dot(R.T)
+        scene_markers = {}
+        remaining_markers = {}
 
-        return Rs, Ts
+        for (marker_id, marker) in tracked_markers.items():
 
-    def estimate_pose(self, tracked_markers) -> Tuple[numpy.array, numpy.array, bool, list, dict]:
-        """Estimate scene pose from tracked markers (cf ArUcoTracker.track()) and validate its consistency according expected scene places.
-        
-        * **Returns:** 
-            - scene translation vector
-            - scene rotation matrix
-            - scene pose estimation success status
-            - list of all tracked markers considered as consistent and used to estimate the pose
+            try:
+                name = self.__identifier_cache[marker_id]
+                scene_markers[name] = marker
+
+            except KeyError:
+                
+                remaining_markers[marker_id] = marker
+
+        return scene_markers, remaining_markers
+
+    def check_markers_consistency(self, scene_markers, angle_tolerance: float, distance_tolerance: float) -> Tuple[dict, dict]:
+        """Evaluate if given markers configuration match related places configuration.
+
+        * **Returns:**
+            - dict of consistent markers
+            - dict of unconsistent markers
             - dict of identified distance or angle unconsistencies and out-of-bounds values
         """
 
-        # Init pose data
-        self._translation = T0
-        self._rotation = R0
-        self._succeded = False
-        self._consistent_markers = {}
-        self._unconsistencies = {}
+        consistent_markers = {}
+        unconsistencies = {}
 
-        # Don't try to estimate pose if there is no tracked markers
-        if len(tracked_markers) == 0:
+        for (A_name, A_marker), (B_name, B_marker) in itertools.combinations(scene_markers.items(), 2):
 
-            return self._translation, self._rotation, self._succeded, self._consistent_markers, self._unconsistencies
+            # Rotation matrix from A marker to B marker
+            AB = B_marker.rotation.dot(A_marker.rotation.T)
 
-        # Look for tracked markers which belong to the scene and store them by name
-        scene_markers = {}
-        for (marker_id, marker) in tracked_markers.items():
+            # 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]
 
-            try:
-                name = self.__identifier_cache[marker_id]
-                scene_markers[name] = marker
+            # Calculate distance between A marker center and B marker center
+            distance = numpy.linalg.norm(A_marker.translation - B_marker.translation)
+            expected_distance = self.__distance_cache[A_name][B_name]
 
-            except KeyError:
-                continue
+            # Check angle and distance according given tolerance then normalise marker pose
+            consistent_angle = math.isclose(angle, expected_angle, abs_tol=angle_tolerance)
+            consistent_distance = math.isclose(distance, expected_distance, abs_tol=distance_tolerance)
 
-        # Pose consistency checking is'nt possible when only one marker is tracked
-        if len(scene_markers) == 1:
+            if consistent_angle and consistent_distance:
 
-            # Get scene pose from the unique marker
-            name, marker = scene_markers.popitem()
+                if A_name not in consistent_markers.keys():
 
-            self._rotation, self._translation = self.__normalise_marker_pose(self.__places[name], marker.rotation, marker.translation)
-            self._succeded = True
-            self._consistent_markers[name] = marker
+                    # Remember this marker is already validated
+                    consistent_markers[A_name] = A_marker
 
-        # Pose consistency checking processes markers two by two
-        else:
+                if B_name not in consistent_markers.keys():
 
-            consistent_markers = []
-            consistent_rmats = []
-            consistent_tvecs = []
+                    # Remember this marker is already validated
+                    consistent_markers[B_name] = B_marker
 
-            for (A_name, A_marker), (B_name, B_marker) in itertools.combinations(scene_markers.items(), 2):
+            else:
 
-                # Rotation matrix from A marker to B marker
-                AB = B_marker.rotation.dot(A_marker.rotation.T)
+                if not consistent_angle:
+                    unconsistencies[f'{A_name}/{B_name} angle'] = angle
+                
+                if not consistent_distance:
+                    unconsistencies[f'{A_name}/{B_name} distance'] = distance
 
-                # 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]
+        # Gather unconsistent markers
+        unconsistent_markers = {}
 
-                # Calculate distance between A marker center and B marker center
-                distance = numpy.linalg.norm(A_marker.translation - B_marker.translation)
-                expected_distance = self.__distance_cache[A_name][B_name]
+        for name, marker in scene_markers.items():
 
-                # Check angle and distance according given tolerance then normalise marker pose
-                consistent_angle = math.isclose(angle, expected_angle, abs_tol=self.angle_tolerance)
-                consistent_distance = math.isclose(distance, expected_distance, abs_tol=self.distance_tolerance)
+            if name not in consistent_markers:
 
-                if consistent_angle and consistent_distance:
+                unconsistent_markers[name] = marker
 
-                    if A_name not in consistent_markers:
+        return consistent_markers, unconsistent_markers, unconsistencies
 
-                        # Remember this marker is already validated
-                        consistent_markers.append(A_name)
+    def estimate_pose_from_axis_markers(self, axis_markers) -> Tuple[numpy.array, numpy.array]:
+        """Calculate tranformations (rotation and translation) from a list of 3 (name, marker) tuples defining an orthogonal axis."""
 
-                        rmat, tvec = self.__normalise_marker_pose(self.__places[A_name], A_marker.rotation, A_marker.translation)
+        O_name, O_marker = axis_markers[0]
+        A_name, A_marker = axis_markers[1]
+        B_name, B_marker = axis_markers[2]
 
-                        # Store normalised marker pose
-                        consistent_rmats.append(rmat)
-                        consistent_tvecs.append(tvec)
-                        
-                    if B_name not in consistent_markers:
+        O_place = self.__places[O_name]
+        A_place = self.__places[A_name]
+        B_place = self.__places[B_name]
 
-                        # Remember this marker is already validated
-                        consistent_markers.append(B_name)
+        # Place axis
+        OA = A_place.translation - O_place.translation
+        OA = OA / numpy.linalg.norm(OA)
 
-                        rmat, tvec = self.__normalise_marker_pose(self.__places[B_name], B_marker.rotation, B_marker.translation)
+        OB = B_place.translation - O_place.translation
+        OB = OB / numpy.linalg.norm(OB)
 
-                        # Store normalised marker pose
-                        consistent_rmats.append(rmat)
-                        consistent_tvecs.append(tvec)
+        # Detect and correct bad place axis orientation
+        X_sign = numpy.sign(OA)[0]
+        Y_sign = numpy.sign(OB)[1]
 
-                else:
+        P = numpy.array([OA*X_sign, OB*Y_sign, numpy.cross(OA*X_sign, OB*Y_sign)])
 
-                    if not consistent_angle:
-                        self._unconsistencies[f'{A_name}/{B_name} angle'] = angle
-                    
-                    if not consistent_distance:
-                        self._unconsistencies[f'{A_name}/{B_name} distance'] = distance
+        OA = A_marker.translation - O_marker.translation
+        OA = OA / numpy.linalg.norm(OA)
 
-            if len(consistent_markers) > 1:
+        OB = B_marker.translation - O_marker.translation
+        OB = OB / numpy.linalg.norm(OB)
 
-                # Consider ArUcoScene rotation as the mean of all consistent 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(consistent_rmats), axis=0)
+        # Detect and correct bad place axis orientation
+        X_sign = numpy.sign(OA)[0]
+        Y_sign = -numpy.sign(OB)[1]
 
-                # Consider ArUcoScene translation as the mean of all consistent translations
-                self._translation = numpy.mean(numpy.array(consistent_tvecs), axis=0)
+        M = numpy.array([OA*X_sign, OB*Y_sign, numpy.cross(OA*X_sign, OB*Y_sign)])
 
-                self._succeded = True
-                self._consistent_markers = consistent_markers
+        # Then estimate ArUcoScene rotation
+        self._rotation = P.dot(M.T)
 
-            else:
+        # Consider ArUcoScene translation as the translation of the marker at axis origin
+        self._translation = O_marker.translation - O_place.translation.dot(O_place.rotation).dot(O_marker.rotation.T)
+
+        return self._translation, self._rotation
 
-                unconsistent_rmats = []
-                unconsistent_tvecs = []
+    def estimate_pose_from_any_markers(self, consistent_markers) -> Tuple[numpy.array, numpy.array]:
+        """Calculate transformations (rotation and translation) that move each marker to its related place."""
 
-                # Gather unconsistent pose estimations
-                for name, marker in scene_markers.items():
+        rotations = []
+        translations = []
 
-                    if name not in consistent_markers:
+        for (name, marker) in consistent_markers.items():
 
-                        rmat, tvec = self.__normalise_marker_pose(self.__places[name], marker.rotation, marker.translation)
+            place = self.__places[name]
 
-                        unconsistent_rmats = [rmat]
-                        unconsistent_tvecs = [tvec]
+            # Rotation matrix that transform marker to related place
+            R = marker.rotation.dot(place.rotation.T)
 
-                # Consider ArUcoScene rotation as the mean of all unconsistent 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(unconsistent_rmats), axis=0)
+            # Translation vector that transform marker to related place
+            T = marker.translation - place.translation.dot(place.rotation).dot(marker.rotation.T)
 
-                # Consider ArUcoScene translation as the mean of all unconsistent translations
-                self._translation = numpy.mean(numpy.array(unconsistent_tvecs), axis=0)
+            rotations.append(R)
+            translations.append(T)
 
-                self._succeded = False
-                self._consistent_markers = {}
+        # Consider ArUcoScene rotation as the mean of all marker rotations
+        # !!! 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(rotations), axis=0)
 
-        #print('----------------------------------------------------')
+        # Consider ArUcoScene translation as the mean of all marker translations
+        self._translation = numpy.mean(numpy.array(translations), axis=0)
 
-        return self._translation, self._rotation, self._succeded, self._consistent_markers, self._unconsistencies
+        return self._translation, self._rotation
 
     @property
     def translation(self) -> numpy.array:
-        """Access to scene translation vector.
-
-        .. warning::
-           Setting scene translation vector implies succeded status to be True and consistency score to be 0."""
+        """Access to scene translation vector."""
 
         return self._translation
 
@@ -463,15 +450,10 @@ class ArUcoScene():
     def translation(self, tvec):
 
         self._translation = tvec
-        self._succeded = True
-        self._consistent_markers = 0
 
     @property
     def rotation(self) -> numpy.array:
-        """Access to scene rotation matrix.
-
-        .. warning::
-           Setting scene rotation matrix implies succeded status to be True and consistency score to be 0."""
+        """Access to scene rotation matrix."""
 
         return self._translation
 
@@ -479,30 +461,16 @@ class ArUcoScene():
     def rotation(self, rmat):
 
         self._rotation = rmat
-        self._succeded = True
-        self._consistent_markers = 0
-
-    @property
-    def succeded(self) -> bool:
-        """Access to scene pose estimation succeded status."""
-
-        return self._succeded
-
-    @property
-    def consistency(self) -> int:
-        """Access to scene pose estimation consistency score."""
-
-        return len(self._consistent_markers)
 
-    def draw_axis(self, frame, K, D):
-        """Draw scene axis."""
+    def draw_axis(self, frame, K, D, consistency=2):
+        """Draw scene axis according a consistency score."""
 
         l = self.__marker_size / 2
         ll = self.__marker_size
 
         # Select color according consistency score
-        n = 95 * self.consistency if self.consistency < 2 else 0
-        f = 159 * self.consistency if self.consistency < 2 else 255
+        n = 95 * consistency if consistency < 2 else 0
+        f = 159 * consistency if consistency < 2 else 255
 
         # Draw axis
         axisPoints = numpy.float32([[ll, 0, 0], [0, ll, 0], [0, 0, ll], [0, 0, 0]]).reshape(-1, 3)
@@ -513,15 +481,15 @@ class ArUcoScene():
         cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 6) # Y (green)
         cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 6) # Z (blue)
 
-    def draw_places(self, frame, K, D):
-        """Draw scene places and their axis."""
+    def draw_places(self, frame, K, D, consistency=2):
+        """Draw scene places and their axis according a consistency score."""
 
         l = self.__marker_size / 2
         ll = self.__marker_size
 
         # Select color according consistency score
-        n = 95 * self.consistency if self.consistency < 2 else 0
-        f = 159 * self.consistency if self.consistency < 2 else 255
+        n = 95 * consistency if consistency < 2 else 0
+        f = 159 * consistency if consistency < 2 else 255
 
         for name, place in self.__places.items():