27
27
"DICT_APRILTAG_36h11" : cv2 .aruco .DICT_APRILTAG_36h11
28
28
}
29
29
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
48
56
49
57
def pose_estimation (frame , aruco_dict_type , matrix_coefficients , distortion_coefficients ):
50
58
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 )
52
60
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
+
56
64
if len (corners ) > 0 :
57
65
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 )
59
68
60
69
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
+
68
72
return frame
69
73
70
74
aruco_type = "DICT_5X5_100"
@@ -86,9 +90,8 @@ def pose_estimation(frame, aruco_dict_type, matrix_coefficients, distortion_coef
86
90
break
87
91
88
92
output = pose_estimation (img , ARUCO_DICT [aruco_type ], intrinsic_camera , distortion )
89
-
90
93
cv2 .imshow ('Estimated Pose' , output )
91
-
94
+
92
95
key = cv2 .waitKey (1 ) & 0xFF
93
96
if key == ord ('q' ):
94
97
break
0 commit comments