-
Notifications
You must be signed in to change notification settings - Fork 0
/
object_detection.py
158 lines (131 loc) · 5.27 KB
/
object_detection.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
import cv2 as cv
import numpy as np
import imutils
# target = cv.imread('images/soccer stars.png', cv.IMREAD_UNCHANGED)
# template = cv.imread('images/player.png', cv.IMREAD_UNCHANGED)
#
# result = cv.matchTemplate(target, template, eval(METHODS[1]))
# --------------------Single Match--------------------
# min_val, max_val, min_loc, max_loc = cv.minMaxLoc(result)
# print('Best match top left position: %s' % str(max_loc))
# print('Best match confidence: %s' % max_val)
#
# threshold = 0.8
# if max_val >= threshold:
# template_height, template_width = template.shape[0], template.shape[1]
# top_left = max_loc
#
# bottom_right = (top_left[0] + template_width, top_left[1] + template_height)
# cv.rectangle(target, top_left, bottom_right, (0, 255, 0), 2, cv.LINE_4)
#
# cv.imwrite('images/result.jpg', target)
#
# # Display the result
# cv.imshow('Template Matching Result', target)
# cv.waitKey(0)
# cv.destroyAllWindows()
# else:
# print('Template not found.')
# --------------------Multiple Match--------------------
# threshold = 0.95
# locations = np.where(result >= threshold)
# locations = list(zip(*locations[::-1]))
#
# # print(locations)
# print(len(locations))
#
# if locations:
# template_height, template_width = template.shape[0], template.shape[1]
# line_color = (0, 255, 0)
# line_type = cv.LINE_4
#
# # Loop over all the locations and draw their rectangle
# for loc in locations:
# # Determine the box positions
# top_left = loc
# bottom_right = (top_left[0] + template_width, top_left[1] + template_height)
# # Draw the box
# cv.rectangle(target, top_left, bottom_right, line_color, line_type)
#
# cv.imwrite('images/result.jpg', target)
#
# # Display the result
# cv.imshow('Matches', target)
# cv.waitKey(0)
# cv.destroyAllWindows()
# else:
# print('Needle not found.')
# Find Object Positions
class ObjectDetection:
def __init__(self, template_path, method, line_color):
self.method = eval(method)
self.imread_flag = cv.IMREAD_UNCHANGED
self.template = cv.imread(template_path, self.imread_flag)
self.template_height = self.template.shape[0]
self.template_width = self.template.shape[1]
self.group_threshold = 2
self.eps = 0.1
self.line_color = line_color
self.line_type = cv.LINE_4
self.marker_color = (255, 0, 255)
self.marker_type = cv.MARKER_CROSS
self.thickness = 2
self.marker_size = 40
def find_objects(self, target, threshold=0.7, max_results=10):
result = cv.matchTemplate(target, self.template, self.method)
locations = np.where(result >= threshold)
locations = list(zip(*locations[::-1]))
if not locations:
return np.array([], np.int32).reshape(0, 4)
rectangles = []
for loc in locations:
rect = [int(loc[0]), int(loc[1]), self.template_width, self.template_height]
rectangles.append(rect)
rectangles.append(rect)
rectangles, weights = cv.groupRectangles(rectangles, self.group_threshold, self.eps)
if len(rectangles) > max_results:
print('Warning: too many results, raise the threshold.')
rectangles = rectangles[:max_results]
return rectangles
def find_objects_rotate(self, target, threshold=0.7, max_results=10):
rectangles = []
for degrees in range(0, 360, 15):
template = imutils.rotate_bound(self.template, degrees)
result = cv.matchTemplate(target, template, self.method)
locations = np.where(result >= threshold)
locations = list(zip(*locations[::-1]))
# if not locations:
# return np.array([], np.int32).reshape(0, 4)
rects = []
for loc in locations:
rect = [int(loc[0]), int(loc[1]), self.template_width, self.template_height]
rects.append(rect)
rects.append(rect)
rects, weights = cv.groupRectangles(rects, self.group_threshold, self.eps)
if len(rects) > max_results:
print('Warning: too many results, raise the threshold.')
rects = rects[:max_results]
for rect in rects:
rectangles.append(rect)
rectangles.append(rect)
rectangles, weights = cv.groupRectangles(rectangles, self.group_threshold, self.eps)
# print("rectangles", rectangles)
return rectangles
@staticmethod
def get_click_points(rectangles):
points = []
for (x, y, w, h) in rectangles:
center_x = x + int(w / 2)
center_y = y + int(h / 2)
points.append((center_x, center_y))
return points
def draw_rectangles(self, target, rectangles):
for (x, y, w, h) in rectangles:
top_left = (x, y)
bottom_right = (x + w, y + h)
cv.rectangle(target, top_left, bottom_right, self.line_color, self.line_type, self.thickness)
return target
def draw_cross_hairs(self, target, points):
for (center_x, center_y) in points:
cv.drawMarker(target, (center_x, center_y), self.marker_color, self.marker_type, self.marker_size, self.thickness)
return target