1 files changed, 47 insertions, 2 deletions

diff --git a/src/argaze/ArUcoMarker/ArUcoDetector.py b/src/argaze/ArUcoMarker/ArUcoDetector.py
index 32a7f3f..32091a4 100644
--- a/src/argaze/ArUcoMarker/ArUcoDetector.py
+++ b/src/argaze/ArUcoMarker/ArUcoDetector.py
@@ -26,7 +26,7 @@ import numpy
 from cv2 import aruco
 
 from argaze import DataFeatures
-from argaze.ArUcoMarker import ArUcoMarkerDictionary, ArUcoMarker, ArUcoOpticCalibrator
+from argaze.ArUcoMarker import ArUcoMarkerDictionary, ArUcoMarker, ArUcoOpticCalibrator, ArUcoMarkerGroup
 
 
 class DetectorParameters():
@@ -139,6 +139,11 @@ class ArUcoDetector(DataFeatures.PipelineStepObject):
 		# Init detected markers data
 		self.__detected_markers = {}
 
+		# Init pose estimation data
+		self.__pose_size = None
+		self.__pose_ids = []
+		self.__pose = None
+
 		# Init detected board data
 		self.__board = None
 		self.__board_corners_number = 0
@@ -187,6 +192,9 @@ class ArUcoDetector(DataFeatures.PipelineStepObject):
 
 		!!! danger "DON'T UNDISTORTED IMAGE"
 			Camera intrinsic parameters and distortion coefficients are used later during pose estimation.
+		
+		!!! note
+			The pose of markers will be also estimated if the pose_size attribute is not None.
 		"""
 
 		# Reset detected markers data
@@ -202,6 +210,7 @@ class ArUcoDetector(DataFeatures.PipelineStepObject):
 			detected_markers_ids = detected_markers_ids.T[0]
 
 			for i, marker_id in enumerate(detected_markers_ids):
+
 				marker = ArUcoMarker.ArUcoMarker(self.__dictionary, marker_id)
 				marker.corners = detected_markers_corners[i][0]
 
@@ -212,12 +221,48 @@ class ArUcoDetector(DataFeatures.PipelineStepObject):
 
 				self.__detected_markers[marker_id] = marker
 
+			# Estimate markers pose if required
+			if self.__pose_size is not None:
+
+				self.estimate_markers_pose(self.__pose_size, self.__pose_ids)
+
+				self.__pose = ArUcoMarkerGroup.ArUcoMarkerGroup(dictionary=self.__dictionary, places=self.__detected_markers)
+
+	@property
+	def pose_size(self) -> float:
+		"""Expected markers size in centimeters to enable pose estimation."""
+		return self.__pose_size
+
+	@pose_size.setter
+	@DataFeatures.PipelineStepAttributeSetter
+	def pose_size(self, pose_size: float):
+
+		self.__pose_size = pose_size
+
+	@property
+	def pose_ids(self) -> list:
+		"""Ids of markers to select for pose estimation (default all)."""
+		return self.__pose_ids
+
+	@pose_ids.setter
+	def pose_ids(self, pose_ids: list):
+
+		self.__pose_ids = pose_ids
+
+	@property
+	def pose(self) -> ArUcoMarkerGroup:
+		"""The estimated pose of detected markers considering their expected size."""
+		return self.__pose
+
 	def estimate_markers_pose(self, size: float, ids: list = []):
-		"""Estimate pose detected markers pose considering a marker size.
+		"""Estimate detected markers pose considering a marker size.
 
 		Parameters:
 			size: size of markers in centimeters.
 			ids: markers id list to select detected markers.
+
+		!!! warning
+			This method have to called after 'detect_markers'
 		"""
 
 		# Is there detected markers ?