-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathaugumented_reality_obj.py
157 lines (130 loc) · 6.24 KB
/
augumented_reality_obj.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# Author: Bharath Kumar
# Contact: [email protected]
# Reference: opencv.org
import cv2
import numpy as np
from plane_tracker import planeTracker, SelectRect
from load_obj import objLoader
import math
# Simple model of a typical house (prism over cuboid)
ar_verts = np.float32([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0],
[0, 0, 1], [0, 1, 1], [1, 1, 1], [1, 0, 1],
[0, 0.5, 2], [1, 0.5, 2]])
ar_edges = [(0, 1), (1, 2), (2, 3), (3, 0),
(4, 5), (5, 6), (6, 7), (7, 4),
(0, 4), (1, 5), (2, 6), (3, 7),
(4, 8), (5, 8), (6, 9), (7, 9), (8, 9)]
ar_faces = [(0, 1, 5, 4), (1, 2, 6, 5), (2, 3, 7, 6), (3, 0, 4, 7), (4, 5, 8), (7, 6, 9), (8, 9, 7, 4), (8, 9, 6, 5)]
color = [(0, 255, 255), (0, 255, 255), (0, 255, 255), (0, 255, 255), (0, 255, 255), (0, 255, 255), (255, 255, 255), (255, 255, 255)]
class VideoPlayer:
def __init__(self):
self.cap = cv2.VideoCapture(2)
self.frame = None
self.tracker = planeTracker()
cv2.namedWindow("PlaneTracker")
# cv2.createTrackbar('focal', 'PlaneTracker', 25, 50, self.empty)
cv2.createTrackbar('focal', 'PlaneTracker', 25, 50, self.empty)
self.rect = SelectRect("PlaneTracker", self.rect_cb)
def empty(*arg, **kw):
pass
def rect_cb(self, rect):
self.tracker.add_target(self.frame, rect)
def play(self):
# obj = objLoader("./ar_models/fox.obj", filetexture="./ar_models/texture.png", swapyz=True)
# obj = objLoader("./ar_models/fox.obj", filetexture="./ar_models/fox_texture.png", swapyz=True)
obj = objLoader("./ar_models/fox.obj", swapyz=True)
# obj = objLoader("./ar_models/rubiks.obj", filetexture="./ar_models/rubiks.png", swapyz=True)
while True:
if self.rect.tp_rect is None:
ret, frame = self.cap.read()
self.frame = frame.copy()
frame = self.frame.copy()
tracked = self.tracker.track(self.frame)
for tr in tracked:
cv2.polylines(frame, [np.int32(tr.quad)], True, (255, 255, 255), 2)
for (x, y) in np.int32(tr.p1):
cv2.circle(frame, (x, y), 2, (255, 255, 255))
# self.draw_model(frame, tr)
frame = self.render_obj(frame, obj, tr)
self.rect.draw(frame)
cv2.imshow("PlaneTracker", frame)
ret = cv2.waitKey(1)
if ret == ord('q'):
break
def draw_model(self, image, tracked):
x0, y0, x1, y1 = tracked.target.rect
quad_3d = np.float32([[x0, y0, 0], [x1, y0, 0], [x1, y1, 0], [x0, y1, 0]])
fx = 0.5 + cv2.getTrackbarPos('focal', 'PlaneTracker') / 50.0
h, w = image.shape[:2]
K = np.float64([[fx*w, 0, 0.5*(w-1)],
[0, fx*w, 0.5*(h-1)],
[0.0,0.0, 1.0]])
dist_coef = np.zeros(4)
_ret, rvec, tvec = cv2.solvePnP(quad_3d, tracked.quad, K, dist_coef)
verts = ar_verts * [(x1-x0), (y1-y0), -(x1-x0)*0.3] + (x0, y0, 0)
verts = cv2.projectPoints(verts, rvec, tvec, K, dist_coef)[0].reshape(-1, 2)
# for i, j in ar_edges:
# (x0, y0), (x1, y1) = verts[i], verts[j]
# cv2.line(image, (int(x0), int(y0)), (int(x1), int(y1)), (255, 255, 0), 2)
for ind, face in enumerate(ar_faces):
face_verts = []
for i in face:
face_verts.append(verts[i])
face_verts = np.int32(face_verts)
cv2.fillConvexPoly(image, face_verts, color[ind])
def render_obj(self, image, obj, tracked, color=True):
vertices = obj.vertices
scale_matrix = np.eye(3)*1
h, w = image.shape[:2]
fx = 0.5 + cv2.getTrackbarPos('focal', 'PlaneTracker') / 50.0
K = np.float64([[fx*w, 0, 0.5*(w-1)],
[0, fx*w, 0.5*(h-1)],
[0.0,0.0, 1.0]])
# K = np.float64([[800, 0, 320],
# [0, 800, 240],
# [0, 0, 1]])
projection = self.projection_matrix(K, tracked.H)
x0, y0, x1, y1 = tracked.target.rect
h, w = y1 - y0, x1 - x0
for face in obj.faces:
face_vertices = face[0]
points = np.array([vertices[vertex - 1] for vertex in face_vertices])
points = np.dot(points, scale_matrix)
# render model in the middle of the reference surface. To do so,
# model points must be displaced
points = np.array([[p[0] + w / 2 + x0, p[1] + h / 2 + y0, p[2]] for p in points])
# points = np.array([[p[0], p[1], p[2]] for p in points])
dst = cv2.perspectiveTransform(points.reshape(-1, 1, 3), projection)
imgpts = np.int32(dst)
if color is False:
cv2.fillConvexPoly(image, imgpts, (255, 255, 255))
else:
color = face[-1]
cv2.fillConvexPoly(image, imgpts, color)
return image
@staticmethod
def projection_matrix(camera_parameters, homography):
# A[R1 R2 T] = H
homography = homography * (-1)
rot_and_transl = np.dot(np.linalg.inv(camera_parameters), homography)
col_1 = rot_and_transl[:, 0]
col_2 = rot_and_transl[:, 1]
col_3 = rot_and_transl[:, 2]
# Normalize
l = math.sqrt(np.linalg.norm(col_1, 2) * np.linalg.norm(col_2, 2))
rot_1 = col_1 / l
rot_2 = col_2 / l
translation = col_3 / l
# compute the orthonormal basis
c = rot_1 + rot_2
p = np.cross(rot_1, rot_2)
d = np.cross(c, p)
rot_1 = np.dot(c / np.linalg.norm(c, 2) + d / np.linalg.norm(d, 2), 1 / math.sqrt(2))
rot_2 = np.dot(c / np.linalg.norm(c, 2) - d / np.linalg.norm(d, 2), 1 / math.sqrt(2))
rot_3 = np.cross(rot_1, rot_2)
# finally, compute the 3D projection matrix from the model to the current frame
projection = np.stack((rot_1, rot_2, rot_3, translation)).T
return np.dot(camera_parameters, projection)
if __name__ == "__main__":
player = VideoPlayer()
player.play()