-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathfeature_extractor.py
164 lines (124 loc) · 6.44 KB
/
feature_extractor.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
import numpy as np
import cv2
from copy import copy, deepcopy
class FeatureExtractor:
def __init__(self, extractor_type="sift", params=None):
self.params = params
self.lk_params = dict(winSize=(self.params["klt"]["winSize"], self.params["klt"]["winSize"]),
maxLevel=self.params["klt"]["maxLevel"],
criteria=(eval(self.params["klt"]["criteria"]["type"]),
self.params["klt"]["criteria"]["maxCount"],
self.params["klt"]["criteria"]["epsilon"])
)
self.extractor_type = extractor_type # ["sift", "surf", "orb", "fast", "harris"]
self.extract = None
if self.extractor_type == "sift":
self.extract = self.extractSiftFeatures
self.match = self.matchSiftFeatures
elif self.extractor_type == "harris":
self.extract = self.extractHarrisCorners
self.match = None
elif self.extractor_type == "shi-tomasi":
self.extract = self.extractShiTomasiCorners
self.match = None
self.track = self.klt_tracker_masked
def extractHarrisCorners(self, image, curr_kp=[], mask_radius=7):
"""Extract Harris corners from the image with subpixel accuracy
"""
winSize = self.params["harris"]["winSize"]
zeroZone = self.params["harris"]["zeroZone"]
criteria_cc = self.params["harris"]["criteria"]
blockSize = self.params["harris"]["blockSize"]
ksize = self.params["harris"]["ksize"]
k = self.params["harris"]["k"]
thresh = self.params["harris"]["thresh"]
if len(image.shape) == 3:
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
gray = copy(image)
gray = np.float32(gray)
dst = cv2.cornerHarris(gray, blockSize, ksize, k)
dst = cv2.dilate(dst, None)
_, dst = cv2.threshold(dst, thresh * dst.max(), 255, 0)
dst = np.uint8(dst)
_, _, _, centroids = cv2.connectedComponentsWithStats(dst)
criteria = (eval(criteria_cc["type"]), criteria_cc["maxCount"], criteria_cc["epsilon"])
corners = cv2.cornerSubPix(gray, np.float32(centroids), (winSize,winSize), (zeroZone,zeroZone), criteria)
# mask out the current keypoints to avoid overlapping features
# mask = np.ones(gray.shape, dtype=np.uint8) * 255
# for i in range(len(curr_kp)):
# cv2.circle(mask, (int(curr_kp[i][0]), int(curr_kp[i][1])), mask_radius, 0, -1)
# corners_int = corners.astype(np.int32) - 1
# corners = corners[mask[corners_int[:, 1], corners_int[:, 0]] == 255]
# print(corners.shape)
return corners
def extractShiTomasiCorners(self, image, curr_kp=[], mask_radius=7):
"""Extract Shi-Tomasi corners from the image with subpixel accuracy
"""
maxCorners = self.params["shi-tomasi"]["maxCorners"]
qualityLevel = self.params["shi-tomasi"]["qualityLevel"]
minDistance = self.params["shi-tomasi"]["minDistance"]
blockSize = self.params["shi-tomasi"]["blockSize"]
if len(image.shape) == 3:
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
gray = copy(image)
# mask out the current keypoints to avoid overlapping features
mask = np.ones(gray.shape, dtype=np.uint8) * 255
for i in range(len(curr_kp)):
cv2.circle(mask, (int(curr_kp[i][0]), int(curr_kp[i][1])), mask_radius, 0, -1)
gray = np.float32(gray)
corners = cv2.goodFeaturesToTrack(gray, maxCorners=maxCorners, qualityLevel=qualityLevel, minDistance=minDistance, mask=mask, blockSize=blockSize)
corners = corners.reshape(corners.shape[0], corners.shape[2])
return corners
def extractSiftFeatures(self, image, descibe=False, curr_kp=[], mask_radius=7):
"""Extract SIFT features from the image
"""
nfeatures = self.params["sift"]["nfeatures"]
sift = cv2.SIFT_create(nfeatures)
mask = np.ones(image.shape, dtype=np.uint8) * 255
for i in range(len(curr_kp)):
cv2.circle(mask, (int(curr_kp[i][0]), int(curr_kp[i][1])), mask_radius, 0, -1)
keypoints, descriptors = sift.detectAndCompute(image, mask=mask)
keypoints = cv2.KeyPoint_convert(keypoints)
if descibe:
return keypoints, descriptors
else:
return keypoints
def matchSiftFeatures(self, descriptors1, descriptors2):
"""Match SIFT features from the image
"""
k = self.params["sift"]["k"]
dist_thresh = self.params["sift"]["dist_threshold"]
bf = cv2.BFMatcher()
matches = bf.knnMatch(descriptors1, descriptors2, k)
good = []
for m, n in matches:
if m.distance < dist_thresh * n.distance:
good.append(m)
return good
def klt_tracker(self, image1, image2, points1, max_bidrectional_error=30):
"""Track the points using KLT tracker with bidirectional error check
"""
dist_thresh = self.params["klt"]["dist_threshold"]
# max_bidrectional_error = self.params["klt"]["max_bidrectional_error"]
points1 = np.array(points1, dtype=np.float32)
points2, status, err = cv2.calcOpticalFlowPyrLK(image1, image2, points1, None, **self.lk_params)
points1r, status, err = cv2.calcOpticalFlowPyrLK(image2, image1, points2, None, **self.lk_params)
d = abs(points1 - points1r).reshape(-1, 2).max(-1)
good = d < dist_thresh
points1 = points1[good]
points2 = points2[good]
return points1, points2
def klt_tracker_masked(self, image1, image2, points1, max_bidrectional_error=30):
"""Track the points using KLT tracker with bidirectional error check
"""
dist_thresh = np.inf
# max_bidrectional_error = self.params["klt"]["max_bidrectional_error"]
points1 = np.array(points1, dtype=np.float32)
points2, status, err = cv2.calcOpticalFlowPyrLK(image1, image2, points1, None, **self.lk_params)
points1r, status, err = cv2.calcOpticalFlowPyrLK(image2, image1, points2, None, **self.lk_params)
d = abs(points1 - points1r).reshape(-1, 2).max(-1)
good = d < dist_thresh
assert len(points1) == len(points2) == len(good)
return points2, good