Skip to content

Commit 98c1f7f

Browse files
CursedProgramsCursedPrograms
committed
...
1 parent 32609f9 commit 98c1f7f

File tree

1 file changed

+191
-149
lines changed

1 file changed

+191
-149
lines changed

‎scripts/pose_object.py

Lines changed: 191 additions & 149 deletions
Original file line numberDiff line numberDiff line change
@@ -1,172 +1,214 @@
11
import numpy as np
22
import cv2
3-
import sys
4-
import time
5-
from OpenGL.GL import *
6-
from OpenGL.GLUT import *
7-
from OpenGL.GLU import *
83
import pygame
94
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-
}
5+
from OpenGL.GL import *
6+
from OpenGL.GLU import *
357

36-
def init_gl(width, height):
37-
glClearColor(0.0, 0.0, 0.0, 0.0)
38-
glClearDepth(1.0)
39-
glDepthFunc(GL_LESS)
8+
class ObjLoader:
9+
def __init__(self, filename, swapyz=False):
10+
self.vertices = []
11+
self.normals = []
12+
self.texcoords = []
13+
self.faces = []
14+
self.gl_list = None
15+
16+
for line in open(filename, "r"):
17+
if line.startswith('#'): continue
18+
values = line.split()
19+
if not values: continue
20+
21+
if values[0] == 'v':
22+
v = list(map(float, values[1:4]))
23+
if swapyz:
24+
v = v[0], v[2], v[1]
25+
self.vertices.append(v)
26+
elif values[0] == 'vn':
27+
v = list(map(float, values[1:4]))
28+
if swapyz:
29+
v = v[0], v[2], v[1]
30+
self.normals.append(v)
31+
elif values[0] == 'vt':
32+
self.texcoords.append(list(map(float, values[1:3])))
33+
elif values[0] == 'f':
34+
face = []
35+
texcoords = []
36+
norms = []
37+
for v in values[1:]:
38+
w = v.split('/')
39+
face.append(int(w[0]))
40+
if len(w) >= 2 and len(w[1]) > 0:
41+
texcoords.append(int(w[1]))
42+
else:
43+
texcoords.append(0)
44+
if len(w) >= 3 and len(w[2]) > 0:
45+
norms.append(int(w[2]))
46+
else:
47+
norms.append(0)
48+
self.faces.append((face, norms, texcoords))
49+
50+
def create_gl_list(self):
51+
if self.gl_list is not None:
52+
return self.gl_list
53+
54+
self.gl_list = glGenLists(1)
55+
glNewList(self.gl_list, GL_COMPILE)
56+
glFrontFace(GL_CCW)
57+
for face in self.faces:
58+
vertices, normals, texture_coords = face
59+
glBegin(GL_POLYGON)
60+
for i in range(len(vertices)):
61+
if normals[i] > 0:
62+
glNormal3fv(self.normals[normals[i] - 1])
63+
glVertex3fv(self.vertices[vertices[i] - 1])
64+
glEnd()
65+
glEndList()
66+
return self.gl_list
67+
68+
def render(self):
69+
glCallList(self.create_gl_list())
70+
71+
def init_ar():
72+
pygame.init()
73+
display = (1280, 720)
74+
pygame.display.set_mode(display, DOUBLEBUF | OPENGL)
75+
4076
glEnable(GL_DEPTH_TEST)
41-
glShadeModel(GL_SMOOTH)
77+
glEnable(GL_LIGHTING)
78+
glLightfv(GL_LIGHT0, GL_POSITION, (0, 0, -2, 1))
79+
glLightfv(GL_LIGHT0, GL_AMBIENT, (0.2, 0.2, 0.2, 1))
80+
glLightfv(GL_LIGHT0, GL_DIFFUSE, (0.5, 0.5, 0.5, 1))
81+
glEnable(GL_LIGHT0)
82+
glEnable(GL_COLOR_MATERIAL)
83+
glColorMaterial(GL_FRONT_AND_BACK, GL_AMBIENT_AND_DIFFUSE)
84+
85+
return display
86+
87+
def set_projection_from_camera(intrinsic):
4288
glMatrixMode(GL_PROJECTION)
4389
glLoadIdentity()
44-
gluPerspective(45.0, float(width)/float(height), 0.1, 100.0)
45-
glMatrixMode(GL_MODELVIEW)
46-
47-
def draw_axis(img, imgpts, corner, length=0.01):
48-
origin = tuple(corner.ravel().astype(int))
49-
img = cv2.line(img, origin, tuple(imgpts[0].ravel().astype(int)), (255,0,0), 5)
50-
img = cv2.line(img, origin, tuple(imgpts[1].ravel().astype(int)), (0,255,0), 5)
51-
img = cv2.line(img, origin, tuple(imgpts[2].ravel().astype(int)), (0,0,255), 5)
52-
return img
53-
54-
def pose_estimation(frame, aruco_dict_type, matrix_coefficients, distortion_coefficients):
55-
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
56-
cv2.aruco_dict = cv2.aruco.Dictionary_get(aruco_dict_type)
57-
parameters = cv2.aruco.DetectorParameters_create()
58-
59-
corners, ids, rejected_img_points = cv2.aruco.detectMarkers(gray, cv2.aruco_dict, parameters=parameters)
60-
61-
# Debug: Print number of markers detected
62-
print(f"Number of ArUco markers detected: {len(corners)}")
63-
64-
if len(corners) > 0:
65-
for i in range(0, len(ids)):
66-
rvec, tvec, markerPoints = cv2.aruco.estimatePoseSingleMarkers(corners[i], 0.02, matrix_coefficients,
67-
distortion_coefficients)
68-
69-
# Debug: Print pose information
70-
print(f"Marker {ids[i][0]} - Rotation: {rvec}, Translation: {tvec}")
71-
72-
# Draw the axes of the marker
73-
axis_length = 0.01
74-
axis = np.float32([[axis_length,0,0], [0,axis_length,0], [0,0,-axis_length]]).reshape(-1,3)
75-
imgpts, jac = cv2.projectPoints(axis, rvec, tvec, matrix_coefficients, distortion_coefficients)
76-
frame = draw_axis(frame, imgpts, corners[i][0][0])
77-
78-
cv2.aruco.drawDetectedMarkers(frame, corners)
79-
80-
# Convert rotation vector to rotation matrix
81-
rotation_matrix, _ = cv2.Rodrigues(rvec)
82-
83-
# Combine rotation and translation into a single 4x4 transformation matrix
84-
transformation_matrix = np.eye(4)
85-
transformation_matrix[:3, :3] = rotation_matrix
86-
transformation_matrix[:3, 3] = tvec.reshape(3)
87-
88-
# Set up OpenGL modelview matrix
89-
glMatrixMode(GL_MODELVIEW)
90-
glLoadIdentity()
91-
glMultMatrixf(transformation_matrix.T)
92-
93-
# Draw the 3D object
94-
glColor3f(0.0, 1.0, 0.0) # Set color to green
95-
obj.render() # Render the loaded OBJ file
96-
97-
return frame
98-
99-
aruco_type = "DICT_5X5_100"
100-
arucoDict = cv2.aruco.Dictionary_get(ARUCO_DICT[aruco_type])
101-
arucoParams = cv2.aruco.DetectorParameters_create()
102-
103-
intrinsic_camera = np.array(((933.15867, 0, 657.59), (0, 933.1586, 400.36993), (0, 0, 1)))
104-
distortion = np.array((-0.43948, 0.18514, 0, 0))
105-
106-
cap = cv2.VideoCapture(0)
107-
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
108-
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
109-
110-
# Debug: Check if camera is opened successfully
111-
if not cap.isOpened():
112-
print("Error: Could not open camera.")
113-
sys.exit()
114-
115-
# Initialize Pygame and OpenGL
116-
pygame.init()
117-
display = (1280, 720)
118-
pygame.display.set_mode(display, DOUBLEBUF | OPENGL)
119-
init_gl(1280, 720)
90+
91+
fx = intrinsic[0,0]
92+
fy = intrinsic[1,1]
93+
fovy = 2 * np.arctan(0.5*720 / fy) * 180 / np.pi
94+
aspect = (1280 * fy) / (720 * fx)
12095

121-
# Load the OBJ file
122-
obj = Obj("objects/cube.obj", swapyz=True)
96+
gluPerspective(fovy, aspect, 0.1, 100.0)
97+
glViewport(0, 0, 1280, 720)
12398

124-
while True:
125-
ret, frame = cap.read()
99+
def set_modelview_from_camera(rvec, tvec):
100+
glMatrixMode(GL_MODELVIEW)
101+
glLoadIdentity()
102+
103+
# Adjust for the shapes returned by estimatePoseSingleMarkers
104+
rotation = rvec[0][0] # Get the 3x1 rotation vector
105+
translation = tvec[0][0] # Get the 3x1 translation vector
106+
107+
# Convert rotation vector to matrix
108+
rmtx = cv2.Rodrigues(rotation)[0]
126109

127-
if not ret:
128-
print("Failed to grab frame")
129-
break
110+
view_matrix = np.array([[rmtx[0,0], rmtx[0,1], rmtx[0,2], translation[0]],
111+
[rmtx[1,0], rmtx[1,1], rmtx[1,2], translation[1]],
112+
[rmtx[2,0], rmtx[2,1], rmtx[2,2], translation[2]],
113+
[0.0, 0.0, 0.0, 1.0]])
130114

131-
# Clear the OpenGL buffers
132-
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
115+
view_matrix = view_matrix * np.array([1, -1, -1, 1])
133116

134-
# Set up the camera matrix
117+
inverse_matrix = np.linalg.inv(view_matrix)
118+
glLoadMatrixf(inverse_matrix.T)
119+
120+
def draw_background(frame):
121+
glDisable(GL_DEPTH_TEST)
135122
glMatrixMode(GL_PROJECTION)
136123
glLoadIdentity()
137-
gluPerspective(45, (display[0]/display[1]), 0.1, 50.0)
124+
gluOrtho2D(0, 1280, 0, 720)
138125
glMatrixMode(GL_MODELVIEW)
139126
glLoadIdentity()
140-
gluLookAt(0, 0, 5, 0, 0, 0, 0, 1, 0)
141127

142-
# Process the frame and estimate pose
143-
output = pose_estimation(frame, ARUCO_DICT[aruco_type], intrinsic_camera, distortion)
128+
# Convert frame to OpenGL texture format
129+
bg_image = cv2.flip(frame, 0)
130+
bg_image = cv2.cvtColor(bg_image, cv2.COLOR_BGR2RGB)
131+
132+
glEnable(GL_TEXTURE_2D)
133+
texture_id = glGenTextures(1)
134+
glBindTexture(GL_TEXTURE_2D, texture_id)
135+
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 1280, 720, 0, GL_RGB, GL_UNSIGNED_BYTE, bg_image)
136+
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
137+
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
144138

145-
# Debug: Display the frame with ArUco markers and axes
146-
cv2.imshow('ArUco Detection', output)
139+
# Draw textured quad
140+
glBegin(GL_QUADS)
141+
glTexCoord2f(0.0, 1.0); glVertex2f(0, 0)
142+
glTexCoord2f(1.0, 1.0); glVertex2f(1280, 0)
143+
glTexCoord2f(1.0, 0.0); glVertex2f(1280, 720)
144+
glTexCoord2f(0.0, 0.0); glVertex2f(0, 720)
145+
glEnd()
147146

148-
# Convert the OpenCV output to a Pygame surface
149-
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
150-
output = np.rot90(output)
151-
output = pygame.surfarray.make_surface(output)
147+
glDeleteTextures([texture_id])
148+
glDisable(GL_TEXTURE_2D)
149+
glEnable(GL_DEPTH_TEST)
150+
151+
def main():
152+
# Initialize camera and AR
153+
cap = cv2.VideoCapture(0)
154+
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
155+
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
152156

153-
# Display the Pygame surface
154-
screen = pygame.display.get_surface()
155-
screen.blit(output, (0,0))
157+
display = init_ar()
156158

157-
pygame.display.flip()
159+
# Camera parameters
160+
camera_matrix = np.array([[933.15867, 0, 657.59],
161+
[0, 933.1586, 400.36993],
162+
[0, 0, 1]])
163+
dist_coeffs = np.array([-0.43948, 0.18514, 0, 0])
158164

159-
for event in pygame.event.get():
160-
if event.type == pygame.QUIT:
161-
pygame.quit()
162-
cap.release()
163-
cv2.destroyAllWindows()
164-
sys.exit()
165-
166-
# Break the loop if 'q' is pressed
167-
if cv2.waitKey(1) & 0xFF == ord('q'):
168-
break
169-
170-
cap.release()
171-
cv2.destroyAllWindows()
172-
pygame.quit()
165+
# Initialize ArUco detector
166+
aruco_dict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_5X5_100)
167+
parameters = cv2.aruco.DetectorParameters_create()
168+
169+
# Load 3D model
170+
obj = ObjLoader("objects/cube.obj", swapyz=True)
171+
172+
clock = pygame.time.Clock()
173+
174+
while True:
175+
for event in pygame.event.get():
176+
if event.type == pygame.QUIT:
177+
pygame.quit()
178+
cap.release()
179+
return
180+
181+
ret, frame = cap.read()
182+
if not ret:
183+
break
184+
185+
# Clear OpenGL buffers
186+
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
187+
188+
# Draw background (camera frame)
189+
draw_background(frame)
190+
191+
# Detect ArUco markers
192+
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
193+
corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=parameters)
194+
195+
if ids is not None:
196+
# Set OpenGL camera projection
197+
set_projection_from_camera(camera_matrix)
198+
199+
# Draw 3D objects for each detected marker
200+
for i in range(len(ids)):
201+
# Get pose of the marker
202+
rvec, tvec, _ = cv2.aruco.estimatePoseSingleMarkers(corners[i], 0.02, camera_matrix, dist_coeffs)
203+
204+
# Set OpenGL modelview matrix
205+
set_modelview_from_camera(rvec, tvec)
206+
207+
glColor3f(0.0, 1.0, 0.0) # Set color to green
208+
obj.render()
209+
210+
pygame.display.flip()
211+
clock.tick(60)
212+
213+
if __name__ == "__main__":
214+
main()

0 commit comments

Comments
 (0)