diff --git a/ObjectDetection/aruco/aruco_detector.py b/ObjectDetection/aruco/aruco_detector.py index 9673cd1..3780467 100644 --- 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