Allowing to estimate pose of a marker selection only.

1 files changed, 35 insertions, 10 deletions

diff --git a/src/argaze/ArUcoMarkers/ArUcoDetector.py b/src/argaze/ArUcoMarkers/ArUcoDetector.py
index 86bcbbf..53f56ff 100644
--- a/src/argaze/ArUcoMarkers/ArUcoDetector.py
+++ b/src/argaze/ArUcoMarkers/ArUcoDetector.py
@@ -168,7 +168,7 @@ class ArUcoDetector():
         
         return self.__parameters
 
-    def detect(self, frame):
+    def detect_markers(self, frame):
         """Detect all ArUco markers into a frame.
 
         .. danger:: DON'T MIRROR FRAME
@@ -184,10 +184,13 @@ class ArUcoDetector():
         # Is there detected markers ?
         if len(self.__detected_markers_corners) > 0:
 
+            # Transform markers ids array into list
+            self.__detected_markers_ids = self.__detected_markers_ids.T[0]
+
             # Gather detected markers data and update metrics
             self.__detection_count += 1
 
-            for i, marker_id in enumerate(self.__detected_markers_ids.T[0]):
+            for i, marker_id in enumerate(self.__detected_markers_ids):
 
                 marker = ArUcoMarker.ArUcoMarker(self.dictionary, marker_id, self.marker_size)
 
@@ -202,21 +205,43 @@ class ArUcoDetector():
 
                 self.__detected_ids.append(marker_id)
 
-    def estimate_pose(self):
-        """Estimate pose of current detected markers."""
+    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_corners) > 0:
 
-            markers_rvecs, markers_tvecs, markers_points = aruco.estimatePoseSingleMarkers(self.__detected_markers_corners, self.marker_size, numpy.array(self.camera.K), numpy.array(self.camera.D)) 
+            # Is there a marker selection ?
+            if len(markers_ids) > 0:
+
+                selected_markers_corners = tuple()
+                selected_markers_ids = []
+
+                for i, marker_id in enumerate(self.__detected_markers_ids):
+
+                    if marker_id in markers_ids:
+
+                        selected_markers_corners += (self.__detected_markers_corners[i],)
+                        selected_markers_ids.append(marker_id)
+
+            # Otherwise, estimate pose of all markers
+            else:
+
+                selected_markers_corners = self.__detected_markers_corners
+                selected_markers_ids = self.__detected_markers_ids
+
+            # 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.camera.K), numpy.array(self.camera.D)) 
 
-            for i, marker_id in enumerate(self.__detected_markers_ids.T[0]):
+                for i, marker_id in enumerate(selected_markers_ids):
 
-                marker = self.__detected_markers[marker_id]
+                    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)
+                    marker.translation = markers_tvecs[i][0]
+                    marker.rotation, _ = cv.Rodrigues(markers_rvecs[i][0])
+                    marker.points = markers_points.reshape(4, 3)
 
     @property
     def detected_markers(self) -> dict[ArUcoMarkerType]: