Splitting draw method in two draw_axis and draw_places methods. Improving estimate_pose comment.

1 files changed, 30 insertions, 34 deletions

diff --git a/src/argaze/ArUcoMarkers/ArUcoScene.py b/src/argaze/ArUcoMarkers/ArUcoScene.py
index 552463d..a0b2a89 100644
--- a/src/argaze/ArUcoMarkers/ArUcoScene.py
+++ b/src/argaze/ArUcoMarkers/ArUcoScene.py
@@ -308,14 +308,14 @@ class ArUcoScene():
 
         return Rs, Ts
 
-    def estimate_pose(self, tracked_markers) -> Tuple[numpy.array, numpy.array, bool, int, dict]:
+    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
-            - all tracked markers considered as consistent and used to estimate the pose
+            - list of all tracked markers considered as consistent and used to estimate the pose
             - dict of identified distance or angle unconsistencies and out-of-bounds values
         """
 
@@ -493,8 +493,8 @@ class ArUcoScene():
 
         return len(self._consistent_markers)
 
-    def draw(self, frame, K, D, draw_places=True):
-        """Draw scene axis and places."""
+    def draw_axis(self, frame, K, D):
+        """Draw scene axis."""
 
         l = self.__marker_size / 2
         ll = self.__marker_size
@@ -503,39 +503,35 @@ class ArUcoScene():
         n = 95 * self.consistency if self.consistency < 2 else 0
         f = 159 * self.consistency if self.consistency < 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, numpy.array(K), numpy.array(D))
+        axisPoints = axisPoints.astype(int)
 
-            # 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, numpy.array(K), numpy.array(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)
 
-            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)
+    def draw_places(self, frame, K, D):
+        """Draw scene places."""
 
-            # Draw places (optional)
-            if draw_places:
+        l = self.__marker_size / 2
+        ll = self.__marker_size
 
-                for name, place in self.__places.items():
+        # 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
 
-                    T = self.__places[name].translation
-                    R = self.__places[name].rotation
+        for name, place in self.__places.items():
 
-                    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, numpy.array(K), numpy.array(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._consistent_markers)
-            print(axisPoints)
+            T = self.__places[name].translation
+            R = self.__places[name].rotation
+
+            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, numpy.array(K), numpy.array(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)