* return angle and frame

Phobos97 · May 22, 2022 · ab82b11 · ab82b11
1 parent 8e2e948
commit ab82b11
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/ObjectDetection/aruco/aruco_detector.py b/ObjectDetection/aruco/aruco_detector.py
@@ -14,17 +14,17 @@ def __init__(self):
         self.camera_matrix = npzfile['mtx']
         self.dist_coeffs = npzfile['dist']
 
-    def aruco_distance(self, corners):
+    def aruco_angle(self, corners):
         rvecs, tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(corners,
                                                               self.aruco_size,
                                                               self.camera_matrix,
                                                               self.dist_coeffs)
 
         # get single marker
-        rvec = rvecs[0]
-        tvec = tvecs[0]
+        x, z = tvecs[0][0][0], tvecs[0][0][2]
+        angle = np.rad2deg(np.arctan(x/z))
 
-        return rvecs, tvecs
+        return angle
 
     def detect_markers(self, frame, return_image=False):
         # detect markers in frame
@@ -38,6 +38,8 @@ def detect_markers(self, frame, return_image=False):
             # flatten the ArUco IDs list
             ids = ids.flatten()
 
+            angle = self.aruco_angle(corners)
+
             # loop over the detected ArUCo corners
             for (markerCorner, markerID) in zip(corners, ids):
                 # extract the marker corners (which are always returned
@@ -69,3 +71,5 @@ def detect_markers(self, frame, return_image=False):
                             (topLeft[0], topLeft[1] - 15),
                             cv2.FONT_HERSHEY_SIMPLEX,
                             0.5, (0, 255, 0), 2)
+
+            return angle, frame