1
+ import numpy as np
2
+ import cv2
3
+ import sys
4
+ import time
5
+ from OpenGL .GL import *
6
+ from OpenGL .GLUT import *
7
+ from OpenGL .GLU import *
8
+ import pygame
9
+ from pygame .locals import *
10
+ from objloader import OBJ
11
+
12
+ ARUCO_DICT = {
13
+ "DICT_4X4_50" : cv2 .aruco .DICT_4X4_50 ,
14
+ "DICT_4X4_100" : cv2 .aruco .DICT_4X4_100 ,
15
+ "DICT_4X4_250" : cv2 .aruco .DICT_4X4_250 ,
16
+ "DICT_4X4_1000" : cv2 .aruco .DICT_4X4_1000 ,
17
+ "DICT_5X5_50" : cv2 .aruco .DICT_5X5_50 ,
18
+ "DICT_5X5_100" : cv2 .aruco .DICT_5X5_100 ,
19
+ "DICT_5X5_250" : cv2 .aruco .DICT_5X5_250 ,
20
+ "DICT_5X5_1000" : cv2 .aruco .DICT_5X5_1000 ,
21
+ "DICT_6X6_50" : cv2 .aruco .DICT_6X6_50 ,
22
+ "DICT_6X6_100" : cv2 .aruco .DICT_6X6_100 ,
23
+ "DICT_6X6_250" : cv2 .aruco .DICT_6X6_250 ,
24
+ "DICT_6X6_1000" : cv2 .aruco .DICT_6X6_1000 ,
25
+ "DICT_7X7_50" : cv2 .aruco .DICT_7X7_50 ,
26
+ "DICT_7X7_100" : cv2 .aruco .DICT_7X7_100 ,
27
+ "DICT_7X7_250" : cv2 .aruco .DICT_7X7_250 ,
28
+ "DICT_7X7_1000" : cv2 .aruco .DICT_7X7_1000 ,
29
+ "DICT_ARUCO_ORIGINAL" : cv2 .aruco .DICT_ARUCO_ORIGINAL ,
30
+ "DICT_APRILTAG_16h5" : cv2 .aruco .DICT_APRILTAG_16h5 ,
31
+ "DICT_APRILTAG_25h9" : cv2 .aruco .DICT_APRILTAG_25h9 ,
32
+ "DICT_APRILTAG_36h10" : cv2 .aruco .DICT_APRILTAG_36h10 ,
33
+ "DICT_APRILTAG_36h11" : cv2 .aruco .DICT_APRILTAG_36h11
34
+ }
35
+
36
+ def init_gl (width , height ):
37
+ glClearColor (0.0 , 0.0 , 0.0 , 0.0 )
38
+ glClearDepth (1.0 )
39
+ glDepthFunc (GL_LESS )
40
+ glEnable (GL_DEPTH_TEST )
41
+ glShadeModel (GL_SMOOTH )
42
+ glMatrixMode (GL_PROJECTION )
43
+ glLoadIdentity ()
44
+ gluPerspective (45.0 , float (width )/ float (height ), 0.1 , 100.0 )
45
+ glMatrixMode (GL_MODELVIEW )
46
+
47
+ def draw_cube (size ):
48
+ glBegin (GL_QUADS )
49
+ glVertex3f ( size , size , - size )
50
+ glVertex3f (- size , size , - size )
51
+ glVertex3f (- size , size , size )
52
+ glVertex3f ( size , size , size )
53
+
54
+ glVertex3f ( size , - size , size )
55
+ glVertex3f (- size , - size , size )
56
+ glVertex3f (- size , - size , - size )
57
+ glVertex3f ( size , - size , - size )
58
+
59
+ glVertex3f ( size , size , size )
60
+ glVertex3f (- size , size , size )
61
+ glVertex3f (- size , - size , size )
62
+ glVertex3f ( size , - size , size )
63
+
64
+ glVertex3f ( size , - size , - size )
65
+ glVertex3f (- size , - size , - size )
66
+ glVertex3f (- size , size , - size )
67
+ glVertex3f ( size , size , - size )
68
+
69
+ glVertex3f (- size , size , size )
70
+ glVertex3f (- size , size , - size )
71
+ glVertex3f (- size , - size , - size )
72
+ glVertex3f (- size , - size , size )
73
+
74
+ glVertex3f ( size , size , - size )
75
+ glVertex3f ( size , size , size )
76
+ glVertex3f ( size , - size , size )
77
+ glVertex3f ( size , - size , - size )
78
+ glEnd ()
79
+
80
+ def pose_estimation (frame , aruco_dict_type , matrix_coefficients , distortion_coefficients ):
81
+ gray = cv2 .cvtColor (frame , cv2 .COLOR_BGR2GRAY )
82
+ cv2 .aruco_dict = cv2 .aruco .Dictionary_get (aruco_dict_type )
83
+ parameters = cv2 .aruco .DetectorParameters_create ()
84
+
85
+ corners , ids , rejected_img_points = cv2 .aruco .detectMarkers (gray , cv2 .aruco_dict , parameters = parameters )
86
+
87
+ if len (corners ) > 0 :
88
+ for i in range (0 , len (ids )):
89
+ rvec , tvec , markerPoints = cv2 .aruco .estimatePoseSingleMarkers (corners [i ], 0.02 , matrix_coefficients ,
90
+ distortion_coefficients )
91
+
92
+ cv2 .aruco .drawDetectedMarkers (frame , corners )
93
+
94
+ # Convert rotation vector to rotation matrix
95
+ rotation_matrix , _ = cv2 .Rodrigues (rvec )
96
+
97
+ # Combine rotation and translation into a single 4x4 transformation matrix
98
+ transformation_matrix = np .eye (4 )
99
+ transformation_matrix [:3 , :3 ] = rotation_matrix
100
+ transformation_matrix [:3 , 3 ] = tvec .reshape (3 )
101
+
102
+ # Set up OpenGL modelview matrix
103
+ glMatrixMode (GL_MODELVIEW )
104
+ glLoadIdentity ()
105
+ glMultMatrixf (transformation_matrix .T )
106
+
107
+ # Draw the 3D cube
108
+ glColor3f (0.0 , 1.0 , 0.0 ) # Set color to green
109
+ draw_cube (0.02 ) # Draw a cube with side length 0.02 (same as marker size)
110
+
111
+ return frame
112
+
113
+ aruco_type = "DICT_5X5_100"
114
+ arucoDict = cv2 .aruco .Dictionary_get (ARUCO_DICT [aruco_type ])
115
+ arucoParams = cv2 .aruco .DetectorParameters_create ()
116
+
117
+ intrinsic_camera = np .array (((933.15867 , 0 , 657.59 ), (0 , 933.1586 , 400.36993 ), (0 , 0 , 1 )))
118
+ distortion = np .array ((- 0.43948 , 0.18514 , 0 , 0 ))
119
+
120
+ cap = cv2 .VideoCapture (0 )
121
+ cap .set (cv2 .CAP_PROP_FRAME_WIDTH , 1280 )
122
+ cap .set (cv2 .CAP_PROP_FRAME_HEIGHT , 720 )
123
+
124
+ # Initialize Pygame and OpenGL
125
+ pygame .init ()
126
+ display = (1280 , 720 )
127
+ pygame .display .set_mode (display , DOUBLEBUF | OPENGL )
128
+ init_gl (1280 , 720 )
129
+
130
+ while True :
131
+ ret , frame = cap .read ()
132
+
133
+ if not ret :
134
+ print ("Failed to grab frame" )
135
+ break
136
+
137
+ # Clear the OpenGL buffers
138
+ glClear (GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT )
139
+
140
+ # Set up the camera matrix
141
+ glMatrixMode (GL_PROJECTION )
142
+ glLoadIdentity ()
143
+ gluPerspective (45 , (display [0 ]/ display [1 ]), 0.1 , 50.0 )
144
+ glMatrixMode (GL_MODELVIEW )
145
+ glLoadIdentity ()
146
+ gluLookAt (0 , 0 , 5 , 0 , 0 , 0 , 0 , 1 , 0 )
147
+
148
+ # Process the frame and estimate pose
149
+ output = pose_estimation (frame , ARUCO_DICT [aruco_type ], intrinsic_camera , distortion )
150
+
151
+ # Convert the OpenCV output to a Pygame surface
152
+ output = cv2 .cvtColor (output , cv2 .COLOR_BGR2RGB )
153
+ output = np .rot90 (output )
154
+ output = pygame .surfarray .make_surface (output )
155
+
156
+ # Display the Pygame surface
157
+ screen = pygame .display .get_surface ()
158
+ screen .blit (output , (0 ,0 ))
159
+
160
+ pygame .display .flip ()
161
+
162
+ for event in pygame .event .get ():
163
+ if event .type == pygame .QUIT :
164
+ pygame .quit ()
165
+ quit ()
166
+
167
+ cap .release ()
168
+ cv2 .destroyAllWindows ()
0 commit comments