-
Notifications
You must be signed in to change notification settings - Fork 0
/
video.py
executable file
·374 lines (342 loc) · 15.2 KB
/
video.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
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
# Video Class with all the object recognition functions
from settings import Settings
from markers import Markers
from typing import List, Union
import cv2 as cv
import numpy as np
class Video:
# Type hinting for all the public variables
frame: np.ndarray
grayscale: np.ndarray
blur: np.ndarray
canny: np.ndarray
thrash: np.ndarray
video_output: np.ndarray
dilate: np.ndarray
erode: np.ndarray
morphologyEx: np.ndarray
blob: np.ndarray
# Type hinting for all private variables
_cap: cv.VideoCapture
def __init__(self, s: Settings):
"""
Opens camera in the in Settings object classified location
:param s: Settings object used to store thresholds and camera used data
"""
if s.use_test_image:
# If defined in the Settings object the program wil use a static test image for testing purposes
# self._cap = cv.imread(s.test_image_path)
# self.frame = cv.imread(s.test_image_path)
path = s.test_image_path + "/01.JPG"
self._cap = cv.imread(path)
self.frame = cv.imread(path)
self.__test_image_counter__ += 1
else:
self._cap = cv.VideoCapture(s.capture_device)
self._cap.set(cv.CAP_PROP_FRAME_WIDTH, s.frame_width)
self._cap.set(cv.CAP_PROP_FRAME_HEIGHT, s.frame_height)
_, self.frame = self._cap.read()
# This defines all the transformations that you can get of the images
self.grayscale = self.frame.copy()
self.blur = self.frame.copy()
self.canny = self.frame.copy()
self.thrash = self.frame.copy()
self.video_output = self.frame.copy()
self.dilate = self.frame.copy()
self.erode = self.frame.copy()
self.morphologyEx = self.frame.copy()
self.blank = self.frame.copy()
self.blob = self.frame.copy()
@staticmethod
def _get_video(s: Settings, cap: Union[np.ndarray, cv.VideoCapture]) -> np.ndarray:
"""
Updates the most recent video frame if the program is using a webcam. Otherwise will return the data in cap
:param s: Settings object used to store threshold and camera used data
:type s: Settings
:param cap: Object for VideoCapture or np array frame of loaded image
:type cap: Union[np.ndarray, cv.VideoCapture]
:return: frame image
:rtype: np.ndarray
"""
if not s.use_test_image: # Function doesn't need to update the video frame if it's a static image
_, frame = cap.read()
return frame
else:
return cap
@staticmethod
def _convert_to_grayscale(frame: np.ndarray) -> np.ndarray:
"""
Converts the image to a grayscale
:param frame: cv BGR image frame that gets converted to grayscale
:type frame: np.ndarray
:return: grayscale image
:rtype: np.ndarray
"""
grayscale = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
return grayscale
@staticmethod
def _edge_canny(s: Settings, frame: np.ndarray) -> np.ndarray:
"""
Uses the "canny_threshold_A" and "canny_threshold_B" to create an Canny view of the objects.
This function first grabs the most recent camera image an converts it to grayscale.
:param s: Settings object used to store thresholds and camera used data
:type s: Settings
:param frame: grayscale image frame for canny edge detection
:type frame: np.ndarray
:return: canny image
:rtype: np.ndarray
"""
canny = cv.Canny(frame, s.canny_threshold_A, s.canny_threshold_B)
return canny
@staticmethod
def _edge_contour(frame: np.ndarray) -> List:
"""
Detects contours found in the image
:param frame: grayscale image frame for contour edge detection
:type frame: np.ndarray
:return: List with all the found contours
"""
# hierarchy is not used
contours, hierarchy = cv.findContours(frame, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
return contours
@staticmethod
def _edge_blob(frame: np.ndarray) -> np.ndarray:
"""
TODO: not yet implemented
:param frame:
:type frame: np.ndarray
:return:
:rtype: np.ndarray
"""
detector = cv.SimpleBlobDetector()
keypoints = detector.detect(frame)
return keypoints
@staticmethod
def _edge_threshold(s: Settings, frame: np.ndarray, threshold_A=None, threshold_B=None) -> np.ndarray:
"""
Uses the "contour_threshold_A" and "contour_threshold_B" to create an threshold view of the objects.
Turns part of the image white of black based on those thresholds.
This function first grabs the most recent camera image an converts it to grayscale.
:param s: Settings object used to store thresholds and camera used data
:type s: Settings
:param frame: grayscale image for threshold detection
:type frame: np.ndarray
:param threshold_A:
:type threshold_A: Union[int, None]
:param threshold_B:
:type threshold_B: Union[int, None]
:return: threshold image
:rtype: np.ndarray
"""
if threshold_A is None:
threshold_A = s.contour_threshold_A
if threshold_B is None:
threshold_B = s.contour_threshold_B
ret, thresh = cv.threshold(frame, threshold_A, threshold_B, cv.THRESH_BINARY)
return thresh
@staticmethod
def _draw_every_contour(s: Settings, contours: Union[np.ndarray, List], width: int, height: int) -> np.ndarray:
"""
Draws every contour on an empty canvas created using _edge_contour function
:param s: Settigns object used to store threshold and camera used data
:type s: Settings
:param contours: List or np array filled with contour data
:type contours: Union[np.ndarray, List]
:param width: frame width
:type width: int
:param height: frame height
:type height: int
:return: Black and white image with with the contour lines drawn on
:rtype: np.ndarray
"""
frame = np.zeros(shape=(width, height), dtype=np.uint8)
for contour in contours:
approx = cv.approxPolyDP(contour, s.contour_poly_threshold * cv.arcLength(contour, True), True)
cv.drawContours(frame, [approx], 0, (255, 255, 255), 1)
return frame
@staticmethod
def _draw_contour(s: Settings, frame: np.ndarray, contours: List) -> np.ndarray:
"""
Function draws the found lines on the current view window.
Uses the "contour_poly_threshold" to approximate the shapes of the objects.
:param s: Settings object used to store thresholds and camera used data
:type s: Settings
:param frame: image used to draw detected contours over
:type frame: np.ndarray
:param contours: List with all the contours in the found image
:type contours: List
:return: Image with contour lines drawn on
:rtype: np.ndarray
"""
video_output = frame.copy()
for contour in contours:
approx = cv.approxPolyDP(contour, s.contour_poly_threshold * cv.arcLength(contour, True), True)
# contour_poly_threshold can be adjusted for recognition
# approximates the contours found into shapes
# cv.drawContours(self.video_output, [approx], 0, (0, 255, 0), 2)
x = approx.ravel()[0]
y = approx.ravel()[1] - 5
# if approx has 3 elements the object is a triangle
# TODO: get centre of a poligon
x_low = np.amin(approx[:, 0, 0])
x_high = np.amax(approx[:, 0, 0])
y_low = np.amin(approx[:, 0, 1])
y_high = np.amax(approx[:, 0, 1])
x_center = x_high - x_low
y_center = y_high - y_low
# if Video._check_pixel(video_output, s.triangle_1, x_center, y_center, 5):
# print("Good")
if len(approx) == 3:
cv.putText(video_output, "Triangle", (x, y), cv.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0))
cv.drawContours(video_output, [approx], 0, (0, 0, 255), 2)
if s.verbose:
print("Found triangle at:", x, ", ", y)
# if approx has 4 element the object is a square or rectangle
elif len(approx) == 4:
x, y, w, h = cv.boundingRect(approx)
cv.drawContours(video_output, [approx], 0, (0, 0, 255), 2)
aspect_ratio = float(w) / h
if 0.95 <= aspect_ratio < 1.05:
cv.putText(video_output, "Square", (x, y), cv.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0))
if s.verbose:
print("Found square at:", x, ", ", y)
else:
cv.putText(video_output, "Rectangle", (x, y), cv.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0))
if s.verbose:
print("Found rectangle at:", x, ", ", y)
else:
str = "No idea it has {} lines".format(len(approx))
cv.putText(video_output, str, (x, y), cv.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0))
cv.drawContours(video_output, [approx], 0, (0, 255, 255), 2)
return video_output
@staticmethod
def _check_pixel(frame: np.ndarray, colour: int, x: int, y: int, sigma: Union[float, int]) -> bool:
R = (colour >> 0) & 0xFF
G = (colour >> 8) & 0xFF
B = (colour >> 16) & 0xFF
if (B - sigma) < frame[x, y, 0] < (B + sigma) and (G - sigma) < frame[x, y, 1] < (G + sigma) and (R - sigma) < \
frame[x, y, 2] < (R + sigma):
return True
else:
return False
@staticmethod
def _detect_lines(frame: np.ndarray) -> np.array:
"""
Uses the HoughLineP function to detect lines.
:param frame: grayscale image for line detection
:type frame: np.ndarray
:return: Houghlinesp array with lines
:rtype: np.ndarray
"""
return cv.HoughLinesP(image=frame, rho=1, theta=np.pi / 180, threshold=50, lines=np.array([]), minLineLength=10,
maxLineGap=5)
@staticmethod
def _draw_lines(lines: np.array, width: int, height: int) -> np.ndarray:
"""
Draws the lines found using the _detect_lines() function.
:param lines: List with lines found in the image
:type lines: np.array
:param width: frame width
:type width: int
:param height: frame height
:type height: int
:return: black and white frame with the detected lines
:rtype: np.ndarray
"""
video_output = np.zeros((width, height), dtype=np.uint8)
for x1, y1, x2, y2 in lines:
cv.line(video_output, (x1, y1), (x2, y2), (255, 0, 0), 1)
return video_output
@staticmethod
def _dilate_edges(frame: np.ndarray, kernel=None, iterations=1) -> np.ndarray:
"""
Enlarges the edges found using the _edge_canny() function. This is done to remove as many double lines as
possible
:param frame: grayscale image used to dilate
:type frame: np.ndarray
:param kernel: kernel size used to dilate edges, if kernel is none size will be 3 x 3
:type kernel: Union[np.ndarray, None]
:param iterations: Non negative int for how many iterations the dilate function should be used for
:type iterations: int
:return: Dilated image
:rtype: np.ndarray
"""
if kernel is None:
kernel = np.ones((3, 3), np.uint8)
dilate = cv.dilate(frame, kernel, iterations=iterations)
return dilate
@staticmethod
def _erode_edges(frame: np.ndarray, kernel=None, iterations=1) -> np.ndarray:
"""
Erodes the edges of the detected lines.
This wil make the lines thinner and easier to detect using the detect lines functions.
:param frame: grayscale image used to erode
:type frame: np.ndarray
:param kernel: kernel size used to erode edges, if kernel is none size will be 3 x 3
:type kernel: Union[np.ndarray, None]
:return: Eroded image
:rtype: np.ndarray
"""
if kernel is None:
kernel = np.ones((3, 3), np.uint8)
erode = cv.erode(frame, kernel, iterations=iterations)
return erode
@staticmethod
def _blur_edges(s: Settings, frame: np.ndarray, ksize=None) -> np.ndarray:
"""
Blurs image using the ksize
:param s: Settigns object used to store thresholds and camera used data
:type s: Settings
:param frame: grayscale image used to blur
:type frame: np.ndarray
:param ksize: Tuple with the blur parameters if ksize=None the values are used from the Settings object
:type ksize: Union[Tuple, None]
:return: Blurred image
:rtype: np.ndarray
"""
if ksize is None:
ksize = (int(s.blur_threshold_A), int(s.blur_threshold_B))
blur = cv.blur(frame, ksize=ksize)
return blur
def detect_objects(self, s: Settings) -> None:
"""
Uses the in Settings specified method to detect the object
:param s: Settings object used to store thresholds and camera used data
:type s: Settings
:return: None
"""
m = Markers(s)
self.frame = self._get_video(s, self._cap)
marker_frame, warped_frame = m.remove_warp(s, self.frame)
self.grayscale = self._convert_to_grayscale(warped_frame)
width, height = self.frame.shape[0:2]
if s.edge_detection_type == 1: # Only threshold detection
self.thrash = self._edge_threshold(s, self.grayscale)
self.video_output = self._draw_contour(s, self.frame, self._edge_contour(self.thrash))
elif s.edge_detection_type == 2: # Threshold + Canny shape detection
gausian_blur = cv.GaussianBlur(self.grayscale, (5, 5), sigmaX=0)
self.thrash = self._edge_threshold(s, gausian_blur)
self.canny = self._edge_canny(s, self.thrash)
self.dilate = self._dilate_edges(self.canny, iterations=s.dilate_iterations)
self.erode = self._erode_edges(self.dilate, iterations=s.erode_iterations)
self.blur = self._blur_edges(s, self.erode)
contours = self._edge_contour(self.blur)
self.video_output = self._draw_contour(s, warped_frame, contours)
cv.imwrite('./output.bmp', self.video_output)
elif s.edge_detection_type == 3: # Blob + Canny shape detection
self.blob = self._edge_blob(self.grayscale)
self.canny = self._edge_canny(s, self.blob)
else: # Simple Canny method
self._edge_canny(s, self.grayscale)
def update_live_feed(self, s: Settings) -> None:
"""
Public method to update the video feed
:param s: Settings object used to store threshold and camera used data
:type s: Settings
:return: None
"""
self.frame = self._get_video(s, self._cap)
def __del__(self) -> int:
# TODO: Create deconstruct
return 0
__test_image_counter__ = 1