Adding estimate_markers_pose method to ArUcoDetector.

1 files changed, 36 insertions, 0 deletions

diff --git a/src/argaze/ArUcoMarkers/ArUcoDetector.py b/src/argaze/ArUcoMarkers/ArUcoDetector.py
index f178a20..01527a1 100644
--- a/src/argaze/ArUcoMarkers/ArUcoDetector.py
+++ b/src/argaze/ArUcoMarkers/ArUcoDetector.py
@@ -302,6 +302,42 @@ class ArUcoDetector():
 
         return detection_time
 
+    def estimate_markers_pose(self, markers_ids: list = []):
+        """Estimate pose of current detected markers or of given markers id list."""
+
+        # Is there detected markers ?
+        if len(self.__detected_markers) > 0:
+
+            # Select all markers by default
+            if len(markers_ids) == 0:
+
+                markers_ids = self.__detected_markers.keys()
+
+            # Prepare data for aruco.estimatePoseSingleMarkers function
+            selected_markers_corners = tuple()
+            selected_markers_ids = []
+
+            for marker_id, marker in self.__detected_markers.items():
+
+                if marker_id in markers_ids:
+
+                    selected_markers_corners += (marker.corners,)
+                    selected_markers_ids.append(marker_id)
+
+            # Estimate pose of selected markers
+            if len(selected_markers_corners) > 0:
+
+                markers_rvecs, markers_tvecs, markers_points = aruco.estimatePoseSingleMarkers(selected_markers_corners, self.marker_size, numpy.array(self.optic_parameters.K), numpy.array(self.optic_parameters.D))
+
+                for i, marker_id in enumerate(selected_markers_ids):
+
+                    marker = self.__detected_markers[marker_id]
+
+                    marker.translation = markers_tvecs[i][0]
+                    marker.rotation, _ = cv.Rodrigues(markers_rvecs[i][0])
+
+                    marker.points = markers_points.reshape(4, 3).dot(marker.rotation) + marker.translation
+
     @property
     def detected_markers(self) -> dict[ArUcoMarkerType]:
         """Access to detected markers dictionary."""