Skip to content

Commit 9eeffc6

Browse files
...
1 parent a145e2a commit 9eeffc6

File tree

1 file changed

+35
-32
lines changed

1 file changed

+35
-32
lines changed

scripts/arucoPoseEstimation.py

Lines changed: 35 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -27,44 +27,48 @@
2727
"DICT_APRILTAG_36h11": cv2.aruco.DICT_APRILTAG_36h11
2828
}
2929

30-
def draw_axis(frame, rvec, tvec, matrix_coefficients, distortion_coefficients):
31-
axis = np.float32([[0.1, 0, 0], [0, 0.1, 0], [0, 0, 0.1]]).reshape(-1, 3)
32-
imgpts, _ = cv2.projectPoints(axis, rvec, tvec, matrix_coefficients, distortion_coefficients)
33-
34-
# Ensure tvec is a 2D point
35-
origin = (int(tvec[0]), int(tvec[1]))
36-
imgpts = np.int32(imgpts).reshape(-1, 2)
37-
38-
print("Origin:", origin)
39-
print("Image Points:", imgpts)
40-
41-
if imgpts.shape[0] != 3:
42-
print("Error: imgpts does not contain 3 points. Current shape:", imgpts.shape)
43-
return
44-
45-
cv2.line(frame, origin, tuple(imgpts[0].ravel()), (255, 0, 0), 5) # X axis
46-
cv2.line(frame, origin, tuple(imgpts[1].ravel()), (0, 255, 0), 5) # Y axis
47-
cv2.line(frame, origin, tuple(imgpts[2].ravel()), (0, 0, 255), 5) # Z axis
30+
def aruco_display(corners, ids, rejected, image):
31+
if len(corners) > 0:
32+
ids = ids.flatten()
33+
for (markerCorner, markerID) in zip(corners, ids):
34+
corners = markerCorner.reshape((4, 2))
35+
(topLeft, topRight, bottomRight, bottomLeft) = corners
36+
37+
topRight = (int(topRight[0]), int(topRight[1]))
38+
bottomRight = (int(bottomRight[0]), int(bottomRight[1]))
39+
bottomLeft = (int(bottomLeft[0]), int(bottomLeft[1]))
40+
topLeft = (int(topLeft[0]), int(topLeft[1]))
41+
42+
cv2.line(image, topLeft, topRight, (0, 255, 0), 2)
43+
cv2.line(image, topRight, bottomRight, (0, 255, 0), 2)
44+
cv2.line(image, bottomRight, bottomLeft, (0, 255, 0), 2)
45+
cv2.line(image, bottomLeft, topLeft, (0, 255, 0), 2)
46+
47+
cX = int((topLeft[0] + bottomRight[0]) / 2.0)
48+
cY = int((topLeft[1] + bottomRight[1]) / 2.0)
49+
cv2.circle(image, (cX, cY), 4, (0, 0, 255), -1)
50+
51+
cv2.putText(image, str(markerID), (topLeft[0], topLeft[1] - 10), cv2.FONT_HERSHEY_SIMPLEX,
52+
0.5, (0, 255, 0), 2)
53+
print("[Inference] ArUco marker ID: {}".format(markerID))
54+
55+
return image
4856

4957
def pose_estimation(frame, aruco_dict_type, matrix_coefficients, distortion_coefficients):
5058
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
51-
aruco_dict = cv2.aruco.Dictionary_get(aruco_dict_type)
59+
cv2.aruco_dict = cv2.aruco.Dictionary_get(aruco_dict_type)
5260
parameters = cv2.aruco.DetectorParameters_create()
53-
54-
corners, ids, rejected_img_points = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=parameters)
55-
61+
62+
corners, ids, rejected_img_points = cv2.aruco.detectMarkers(gray, cv2.aruco_dict, parameters=parameters)
63+
5664
if len(corners) > 0:
5765
for i in range(0, len(ids)):
58-
rvec, tvec, markerPoints = cv2.aruco.estimatePoseSingleMarkers(corners[i], 0.02, matrix_coefficients, distortion_coefficients)
66+
rvec, tvec, markerPoints = cv2.aruco.estimatePoseSingleMarkers(corners[i], 0.02, matrix_coefficients,
67+
distortion_coefficients)
5968

6069
cv2.aruco.drawDetectedMarkers(frame, corners)
61-
62-
# Ensure rvec and tvec are in the correct shape
63-
rvec = rvec.reshape((3, 1))
64-
tvec = tvec.reshape((3, 1))
65-
66-
draw_axis(frame, rvec, tvec, matrix_coefficients, distortion_coefficients)
67-
70+
cv2.drawFrameAxes(frame, matrix_coefficients, distortion_coefficients, rvec, tvec, 0.01)
71+
6872
return frame
6973

7074
aruco_type = "DICT_5X5_100"
@@ -86,9 +90,8 @@ def pose_estimation(frame, aruco_dict_type, matrix_coefficients, distortion_coef
8690
break
8791

8892
output = pose_estimation(img, ARUCO_DICT[aruco_type], intrinsic_camera, distortion)
89-
9093
cv2.imshow('Estimated Pose', output)
91-
94+
9295
key = cv2.waitKey(1) & 0xFF
9396
if key == ord('q'):
9497
break

0 commit comments

Comments
 (0)