Retruning validity score instead of simple validity state. Adding set_pose method. Drawing cube according its validity score.

1 files changed, 69 insertions, 54 deletions

diff --git a/src/argaze/ArUcoMarkers/ArUcoCube.py b/src/argaze/ArUcoMarkers/ArUcoCube.py
index a1090fe..e2983df 100644
--- a/src/argaze/ArUcoMarkers/ArUcoCube.py
+++ b/src/argaze/ArUcoMarkers/ArUcoCube.py
@@ -80,7 +80,7 @@ class ArUcoCube():
             self.__translation = numpy.zeros(3)
             self.__rotation = numpy.zeros(3)
             self.__succeded = False
-            self.__validated = False
+            self.__validity = 0
 
             # Process markers ids to speed up further calculations
             self.__identifier_cache = {}
@@ -205,7 +205,12 @@ class ArUcoCube():
         self.__translation = numpy.zeros(3)
         self.__rotation = numpy.zeros(3)
         self.__succeded = False
-        self.__validated = False
+        self.__validity = 0
+
+        # Don't try to estimate pose if there is no tracked markers
+        if len(tracked_markers) == 0:
+
+            return self.get_pose()
 
         # Look for faces related to tracked markers
         tracked_faces = {}
@@ -229,6 +234,7 @@ class ArUcoCube():
 
             self.__rotation, self.__translation = self.__normalise_face_pose(name,face, F)
             self.__succeded = True
+            self.__validity = 1
 
             #print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
             #print(f'arcube rotation vector: {self.__rotation[0][0]:3f} {self.__rotation[1][0]:3f} {self.__rotation[2][0]:3f}')
@@ -312,7 +318,7 @@ class ArUcoCube():
                 #print(f'arcube translation vector: {self.__translation[0]:3f} {self.__translation[1]:3f} {self.__translation[2]:3f}')
 
                 self.__succeded = True
-                self.__validated = True
+                self.__validity = len(valid_faces)
 
         #print('----------------------------------------------------')
 
@@ -320,67 +326,76 @@ class ArUcoCube():
 
     def get_pose(self):
 
-        return self.__translation, self.__rotation, self.__succeded, self.__validated
-
-    def offset_pose(self, tvec = numpy.zeros(3), rvec = numpy.zeros(3)):
-
-        # Offset cube translation
-        self.__translation += tvec
+        return self.__translation, self.__rotation, self.__succeded, self.__validity
 
-        # Offset cube rotation
-        C, _ = cv.Rodrigues(self.__rotation)
+    def set_pose(self, tvec = numpy.array([]), rvec = numpy.array([])):
 
-        R_offset = self.__make_rotation_matrix(*rvec)
-        C_offset = C.dot(R_offset.T)  
+        if tvec.size == 3:
+            self.__translation = tvec
 
-        self.__rotation, _ = cv.Rodrigues(C_offset)
+        if rvec.size == 3:
+            self.__rotation = rvec
 
-        # Pose is no more validated
         self.__succeded = True
-        self.__validated = False
+        self.__validity = 0
 
     def draw(self, frame, K, D):
 
         l = self.edge_size / 2
         ll = self.edge_size
 
-        # 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, K, D)
-        axisPoints = axisPoints.astype(int)
-
-        frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (0,0,255), 5) # X (red)
-        frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (0,255,0), 5) # Y (green)
-        frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (255,0,0), 5) # Z (blue)
-
-        # Draw left face
-        leftPoints = numpy.float32([[-l, l, l], [-l, -l, l], [-l, -l, -l], [-l, l, -l]]).reshape(-1, 3)
-        leftPoints, _ = cv.projectPoints(leftPoints, self.__rotation, self.__translation, K, (0, 0, 0, 0))
-        leftPoints = leftPoints.astype(int)
-        
-        frame = cv.line(frame, tuple(leftPoints[0].ravel()), tuple(leftPoints[1].ravel()), (0,0,255), 2)
-        frame = cv.line(frame, tuple(leftPoints[1].ravel()), tuple(leftPoints[2].ravel()), (0,0,255), 2)
-        frame = cv.line(frame, tuple(leftPoints[2].ravel()), tuple(leftPoints[3].ravel()), (0,0,255), 2)
-        frame = cv.line(frame, tuple(leftPoints[3].ravel()), tuple(leftPoints[0].ravel()), (0,0,255), 2)
-
-        # Draw top face
-        topPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, l, -l], [l, l, -l]]).reshape(-1, 3)
-        topPoints, _ = cv.projectPoints(topPoints, self.__rotation, self.__translation, K, (0, 0, 0, 0))
-        topPoints = topPoints.astype(int)
-
-        frame = cv.line(frame, tuple(topPoints[0].ravel()), tuple(topPoints[1].ravel()), (0,255,0), 2)
-        frame = cv.line(frame, tuple(topPoints[1].ravel()), tuple(topPoints[2].ravel()), (0,255,0), 2)
-        frame = cv.line(frame, tuple(topPoints[2].ravel()), tuple(topPoints[3].ravel()), (0,255,0), 2)
-        frame = cv.line(frame, tuple(topPoints[3].ravel()), tuple(topPoints[0].ravel()), (0,255,0), 2)
-
-        # Draw front face
-        frontPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, -l, l], [l, -l, l]]).reshape(-1, 3)
-        frontPoints, _ = cv.projectPoints(frontPoints, self.__rotation, self.__translation, K, (0, 0, 0, 0))
-        frontPoints = frontPoints.astype(int)
-        
-        frame = cv.line(frame, tuple(frontPoints[0].ravel()), tuple(frontPoints[1].ravel()), (255,0,0), 2)
-        frame = cv.line(frame, tuple(frontPoints[1].ravel()), tuple(frontPoints[2].ravel()), (255,0,0), 2)
-        frame = cv.line(frame, tuple(frontPoints[2].ravel()), tuple(frontPoints[3].ravel()), (255,0,0), 2)
-        frame = cv.line(frame, tuple(frontPoints[3].ravel()), tuple(frontPoints[0].ravel()), (255,0,0), 2)
+        # Select color according validity score
+        n = 95 * self.__validity if self.__validity < 2 else 0
+        f = 159 * self.__validity if self.__validity < 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, K, D)
+            axisPoints = axisPoints.astype(int)
+
+            frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (n,n,f), 5) # X (red)
+            frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 5) # Y (green)
+            frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 5) # Z (blue)
+
+            # Draw left face
+            leftPoints = numpy.float32([[-l, l, l], [-l, -l, l], [-l, -l, -l], [-l, l, -l]]).reshape(-1, 3)
+            leftPoints, _ = cv.projectPoints(leftPoints, self.__rotation, self.__translation, K, D)
+            leftPoints = leftPoints.astype(int)
+            
+            frame = cv.line(frame, tuple(leftPoints[0].ravel()), tuple(leftPoints[1].ravel()), (n,n,f), 2)
+            frame = cv.line(frame, tuple(leftPoints[1].ravel()), tuple(leftPoints[2].ravel()), (n,n,f), 2)
+            frame = cv.line(frame, tuple(leftPoints[2].ravel()), tuple(leftPoints[3].ravel()), (n,n,f), 2)
+            frame = cv.line(frame, tuple(leftPoints[3].ravel()), tuple(leftPoints[0].ravel()), (n,n,f), 2)
+
+            # Draw top face
+            topPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, l, -l], [l, l, -l]]).reshape(-1, 3)
+            topPoints, _ = cv.projectPoints(topPoints, self.__rotation, self.__translation, K, D)
+            topPoints = topPoints.astype(int)
+
+            frame = cv.line(frame, tuple(topPoints[0].ravel()), tuple(topPoints[1].ravel()), (n,f,n), 2)
+            frame = cv.line(frame, tuple(topPoints[1].ravel()), tuple(topPoints[2].ravel()), (n,f,n), 2)
+            frame = cv.line(frame, tuple(topPoints[2].ravel()), tuple(topPoints[3].ravel()), (n,f,n), 2)
+            frame = cv.line(frame, tuple(topPoints[3].ravel()), tuple(topPoints[0].ravel()), (n,f,n), 2)
+
+            # Draw front face
+            frontPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, -l, l], [l, -l, l]]).reshape(-1, 3)
+            frontPoints, _ = cv.projectPoints(frontPoints, self.__rotation, self.__translation, K, D)
+            frontPoints = frontPoints.astype(int)
+            
+            frame = cv.line(frame, tuple(frontPoints[0].ravel()), tuple(frontPoints[1].ravel()), (f,n,n), 2)
+            frame = cv.line(frame, tuple(frontPoints[1].ravel()), tuple(frontPoints[2].ravel()), (f,n,n), 2)
+            frame = cv.line(frame, tuple(frontPoints[2].ravel()), tuple(frontPoints[3].ravel()), (f,n,n), 2)
+            frame = cv.line(frame, tuple(frontPoints[3].ravel()), tuple(frontPoints[0].ravel()), (f,n,n), 2)
+
+        except Exception as e:
+
+            print(e)
+            print(self.__translation)
+            print(self.__rotation)
+            print(self.__succeded)
+            print(self.__validity)
+            print(axisPoints)
 
         return frame