forked from STOC-Machine/Optical-Flow
-
Notifications
You must be signed in to change notification settings - Fork 0
/
grid_squares.py
204 lines (164 loc) · 7.97 KB
/
grid_squares.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
200
201
202
203
import numpy as np
import cv2
from operator import attrgetter
from vector import dot, distance, cross
# Length of the grid squares in cm
SQUARE_LENGTH = 28.5
# Width of gaps between grid squares
SQUARE_GAP = 2.5
def compute_frame_squares(img):
"""
Takes an image, and returns a list of all "squares" in the image
:param img:
:return:
"""
# Create a new image with the minimum of all three channels
red, green, blue = cv2.split(img)
gray = np.minimum(np.minimum(red, green), blue)
# Run a threshold to find only white lines.
ret, out = cv2.threshold(gray, 120, 255, cv2.THRESH_BINARY) # Logitech camera
cv2.imshow('Threshold', out)
# Get contours in image in a list.
dump, contours, hierarchy = cv2.findContours(out, cv2.RETR_LIST,
cv2.CHAIN_APPROX_SIMPLE)
# Sort by area, since there are lots of contours with 0 area
contours.sort(key=cv2.contourArea, reverse=True)
squares = []
for i in range(len(contours)):
# Stop once area is too small
if cv2.contourArea(contours[i]) < 1000:
break
# Simplify contour, check if it's a square
candidate = GridSquare(contours[i])
epsilon = 0.01 * cv2.arcLength(candidate.contour, True)
candidate.corners = cv2.approxPolyDP(candidate.contour, epsilon, True)
if len(candidate.corners) == 4:
# It's a square!
squares.append(candidate)
return squares
class GridSquare:
cam_rot = None # Rotation vector of camera aligned to grid square axes.
normal = [] # Normal vector pointing out of the square
rvec = None # rvec returned by solvePNP
cam_pos = None # Equivalent to location. Used for internal testing reasons
tvec = None # tvec returned by solvePNP ransac
score = 0 # How good is this square. See compare square normals for more
corners = [] # Image coordinates of square corners
contour = None # Unsimplified image coordinates of square corners
location = [] # Square location in camera coordinates
def __init__(self, contour):
self.contour = contour
def update_pos_stats(self, cam_matrix, distortion_coefficients, object_points):
"""
Takes the square's corners found in the image, corresponding 3d
coordinates, and intrinsic camera information. Sets fields related to
extrinsic camera information: cam_rot, normal, cam_pos, location.
Note that the square might be bogus until you check its score, and
camera extrinsics are unaligned until align squares is called.
:param cam_matrix:
:param distortion_coefficients:
:param object_points:
:return: nothing
"""
temp_corners = self.corners.reshape(4, 2, 1).astype(float)
# Gets vector from camera to center of square
inliers, self.rvec, self.tvec = cv2.solvePnP(object_points, temp_corners,
cam_matrix, distortion_coefficients)
rot_matrix = cv2.Rodrigues(self.rvec)
cam_pos = np.multiply(cv2.transpose(rot_matrix[0]), -1).dot(self.tvec)
cam_to_grid_transform = np.concatenate((cv2.transpose(rot_matrix[0]), cam_pos), axis=1)
grid_to_cam_transform = np.linalg.inv(np.concatenate((cam_to_grid_transform, np.array([[0, 0, 0, 1]])), axis=0))
self.cam_rot = list(cam_to_grid_transform.dot(np.array([0, 0, 1, 0])))
self.normal = grid_to_cam_transform.dot(np.array([0, 0, 1, 0]))
self.cam_pos = cam_pos
self.location = [cam_pos[0][0],cam_pos[1][0],cam_pos[2][0]]
def align_squares(self, guess, cam_matrix, distortion_coefficients, object_points):
"""
Needed because ABCD has different camera location from BCDA. Takes a
square's corners, camera intrisics, object information, and a camera
rotation to align to. Finds the camera position that gives a rotation
vector (unit vector in cam direction using grid axes) closest to guess.
:param guess:
:param cam_matrix:
:param distortion_coefficients:
:param object_points:
:return: nothing
"""
alignment_scores = [dot(self.cam_rot, guess), 0, 0, 0]
# Loop through possible orders: ABCD, BCDA, CDAB, DABC
for rot in range(1, 4):
# Shift to the rot permutation
temp_corners = np.roll(self.corners, rot, axis=0)
# Get vector from camera to center of square
temp_corners = temp_corners.reshape(4,2,1).astype(float)
inliers, self.rvec, self.tvec = cv2.solvePnP(object_points, temp_corners,
cam_matrix, distortion_coefficients)
# Get rotation information for change to square center at origin
rot_matrix = cv2.Rodrigues(self.rvec)
cam_pos = np.multiply(cv2.transpose(rot_matrix[0]), -1).dot(self.tvec)
cam_to_grid_transform = np.concatenate((cv2.transpose(rot_matrix[0]), cam_pos), axis=1)
# Compare camera unit vector to guess
alignment_scores[rot] = dot(list(cam_to_grid_transform.dot(np.array([0, 0, 1, 0]))), guess)
# Pick the orientation that was best and recompute camera location
self.corners = np.roll(self.corners, alignment_scores.index(max(alignment_scores)), axis=0)
self.update_pos_stats(cam_matrix, distortion_coefficients, object_points)
def compare_square_normals(self, square, dim):
"""
Increments the score of two squares based on how parallel their normal
vectors are.
:param square:
:param dim: the dimensions of the image
:return: nothing
"""
# Using 1-cross product due to larger change of sin when parallel
temp_cross = cross(self.normal, square.normal)
edge = False
for point in square.corners:
if point[0][0] < 1 or point[0][0] > dim[1] - 2 or point[0][1] < 1 \
or point[0][1] > dim[0] - 2:
# Contours on the edge don't improve scores
edge = True
if not edge:
# Increment the score
self.score += 1 - abs(distance(temp_cross) / (distance(square.normal)*distance(self.normal)))
# Get Square Stats
#
def get_square_stats(img, cam_matrix, distortion_coefficients, cam_rot_guess):
"""
Returns list of squares with locations, sorted by score, aligned to
cam_rot_guess. Needs camera intrinsics. If cam_rot_guess is 0, will pick an
arbitrary square to align to.
:param img:
:param cam_matrix:
:param distortion_coefficients:
:param cam_rot_guess:
:return:
"""
# 3d grid square coordinates
object_points = np.array([[[-SQUARE_LENGTH/2, -SQUARE_LENGTH/2, 0]],
[[-SQUARE_LENGTH/2, SQUARE_LENGTH/2, 0]],
[[SQUARE_LENGTH/2, SQUARE_LENGTH/2, 0]],
[[SQUARE_LENGTH/2, -SQUARE_LENGTH/2, 0]]],
np.float32).reshape(4, 3, 1)
# Find squares in the image
squares = compute_frame_squares(img)
for square in squares:
# Get camera location for those squares
square.update_pos_stats(cam_matrix, distortion_coefficients, object_points)
# Compare each pair of squares
for i in range(len(squares)):
for j in range(len(squares)):
squares[i].compare_square_normals(squares[j], img.shape)
# Sort by score
squares.sort(key=attrgetter('score'),reverse=True)
# Filter out low scores if there are squares
if len(squares)>0:
threshold = max(.95*squares[0].score, 1.0)
squares = list(filter(lambda x: x.score > threshold, squares))
# Align squares
for square in squares:
if cam_rot_guess == 0:
# If we don't have a guess, use the highest score square
cam_rot_guess = squares[0].cam_rot
square.align_squares(cam_rot_guess, cam_matrix, distortion_coefficients, object_points)
return squares, cam_rot_guess