-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathequi_augmentation.py
199 lines (140 loc) · 6 KB
/
equi_augmentation.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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
import cv2
import numpy as np
def polar_to_3d(dphi, dtheta): #in radians
p_x = np.cos(dphi) * np.cos(dtheta);
p_y = np.cos(dphi) * np.sin(dtheta);
p_z = np.sin(dphi);
p1 = np.array([p_x,p_y,p_z])
return p1
def get_rotation_matrices(phi, theta): #in radians
## matrix to rotate. First y, then z
y_axis = np.array([0.0, 1.0, 0.0], np.float32)*phi
z_axis = np.array([0.0, 0.0, 1.0], np.float32)*theta
M1 = cv2.Rodrigues(y_axis)[0]
M2 = cv2.Rodrigues(z_axis)[0]
M = np.dot(M1,M2)
# matrix to inverse rotation
y_axis = np.array([0.0, 1.0, 0.0], np.float32)*(-phi)
z_axis = np.array([0.0, 0.0, 1.0], np.float32)*(-theta)
M1_inv = cv2.Rodrigues(z_axis)[0]
M2_inv = cv2.Rodrigues(y_axis)[0]
M_inv = np.dot(M1_inv,M2_inv)
return M, M_inv
def get8boundingpointsy_x(bounding_box):#bounding_box = (minx, miny),(maxx,maxy)
(minx, miny),(maxx,maxy) = bounding_box
middle_x = (minx+maxx)//2
middle_y = (miny+maxy)//2
t_left = (miny,minx) #top #left
t_center = (miny, middle_x) #top center
t_right = (miny, maxx) #top #right
c_left = (middle_y, minx) #left center
c_right = (middle_y, maxx) #right center
b_left = (maxy, minx) #bottom #left
b_center = (maxy, middle_x) #bottom center
b_right = (maxy, maxx) #bottom #right
return [t_left, t_center, t_right, c_left, c_right, b_left, b_center, b_right]
def project_points(points, delta, eq_cx, eq_cy, r_h, r_w, M, image):
#equi_points = equirectangular_image.copy()
eq_bound = []
for p in points:
py, px = p
py = (py/image.shape[0])*r_h
px = (px/image.shape[1])*r_w
py = py - r_h/2
px = px - r_w/2
ptheta = np.arctan(px)
pphi = np.arctan(py)
p1 = polar_to_3d(pphi, ptheta)
p1 = np.dot(p1, M)
ntheta = np.degrees(np.arctan2(p1[1], p1[0]))
nphi = np.degrees(np.arcsin(p1[2]))
x = (ntheta/delta)+eq_cx
y = (nphi/delta)+eq_cy
x = int(np.round(x))
y = int(np.round(y))
#equi_points = cv2.circle(equi_points, (x,y), 8, (0,0,255), 5)
eq_bound.append((y,x))
return eq_bound
'''
alpha = tolerance of difference of top and bottom middle points in x
'''
def get_search_ranges(eq_bound, eq_w, eq_h, phi, theta, alpha = 5):
top = min([y[0] for y in eq_bound[:3]])
bottom = max([y[0] for y in eq_bound[-3:]])
top_center = eq_bound[1][1]%(eq_w-1)
bottom_center = eq_bound[6][1]%(eq_w-1)
if abs(top_center - bottom_center) < alpha: #they are aligned
if phi < 0:
ref_left = eq_bound[0][1] #top left
ref_right = eq_bound[2][1] #top right
else:
ref_left = eq_bound[5][1] #bottom left
ref_right = eq_bound[7][1] #bottom right
if ref_right < ref_left:
part1 = np.arange(ref_left, eq_w)
part2 = np.arange(0, ref_right+1)
x_range = np.concatenate((part1,part2))
else:
x_range = np.arange(ref_left, ref_right+1)
y_range = np.arange(top, bottom+1)
else: #they are not aligned, close to the poles
x_range = np.arange(0, eq_w)
if phi > 0: #down
y_range = np.arange(top, eq_h)
else: #up
y_range = np.arange(0, bottom+1)
return x_range, y_range
def draw_points(image, points, color = (0,0,255), radius = 8, thickness = -1):
image_points = image.copy()
for i, p in enumerate(points):
image_points = cv2.circle(image_points, p, radius, color, thickness)
return image_points
'''
r_h: height of projected image, assuming that the radius
of the sphere is 1.
phi (radians): vertical angle, down is positive and up is negative
theta (radians): horizontal angle
'''
def image_projection_to_equi(equirectangular_image, image, phi = 0, theta = 0, r_h = -1, draw_intermediate = False):
eq_h, eq_w, _ = equirectangular_image.shape
eq_cx = eq_w // 2.0
eq_cy = eq_h // 2.0
delta = 180/eq_h # number of degrees per pixel
h, w, _ = image.shape
if r_h == -1:
diameter = eq_w/np.pi
r_h = h/(diameter*2)
#print(f'eq_w:{eq_w} diameter:{diameter} r_h:{r_h}')
r_w = (image.shape[1]/image.shape[0])*r_h
points = get8boundingpointsy_x(((0,0),(w-1,h-1)))#bounding_box = (minx, miny),(maxx,maxy)
M, M_inv = get_rotation_matrices(phi, theta)
points_projector = lambda ps: project_points(ps, delta, eq_cx, eq_cy, r_h, r_w, M, image)
eq_bound = points_projector(points)
x_range, y_range = get_search_ranges(eq_bound, eq_w, eq_h, phi, theta)
if draw_intermediate:
image_points = draw_points(image, points, radius = 4)
eq_points = draw_points(equirectangular_image, eq_bound, radius = 4)
eq_area = equirectangular_image.copy()
for i in y_range:
#for i in range(eq_h):
for j in x_range:
#for j in range(eq_w):
dtheta = np.radians((j - eq_cx)*delta)
dphi = np.radians((i - eq_cy)*delta)
p1 = np.dot(polar_to_3d(dphi, dtheta), M_inv)
dtheta = np.arctan2(p1[1], p1[0]);
dphi = np.arcsin(p1[2])
tanx = np.tan(dtheta) + r_w/2
tany = np.tan(dphi) + r_h/2
if tanx >= 0 and tanx <= r_w and tany>=0 and tany<=r_h:
posx = int((tanx/r_w) * image.shape[1])
posy = int((tany/r_h) * image.shape[0])
pixel = image[posy, posx]
equirectangular_image[i,j] = pixel
if draw_intermediate:
eq_area[i,j] = pixel
elif draw_intermediate:
eq_area[i,j] = [0, 0, 255]
if draw_intermediate:
return equirectangular_image, [image_points, eq_points, eq_area], points_projector, r_h
return equirectangular_image, None, points_projector, r_h