-
Notifications
You must be signed in to change notification settings - Fork 0
/
boat_finder.py
50 lines (43 loc) · 2.04 KB
/
boat_finder.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
import numpy as np
import cv2
from hull_filter import hull_filter_lower, hull_filter_upper
# HSV thresholds (red wraps around from 179 to 0)
color_lower = np.array([174,140,50])
color_upper = np.array([180,255,255])
color2_lower = np.array([0,140,50])
color2_upper = np.array([15,255,255])
# Kernel size for erode and dilate
kernel = np.ones((5,5), np.uint8)
def parse_image(image_frame):
# Convert to HSV color space
hsv_frame = cv2.cvtColor(image_frame, cv2.COLOR_BGR2HSV)
# Split out the HSV values
hue, sat, val = cv2.split(hsv_frame)
# Mask off the desired color in the HSV color space
hsv_mask = cv2.inRange(hsv_frame, color_lower, color_upper)
hsv2_mask = cv2.inRange(hsv_frame, color2_lower, color2_upper)
# Combine the two masks b/c red wraps around from 255 to 0
hsv_mask = cv2.add(hsv_mask, hsv2_mask)
#Erode and Dilate to get rid of image noise
hsv_mask_filt = cv2.erode(hsv_mask, kernel, iterations=1)
hsv_mask_filt = cv2.dilate(hsv_mask_filt, kernel, iterations=2)
# Find contours in the masked image to find possible buoy locations
contours = cv2.findContours(hsv_mask_filt.copy(), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)[-2]
# Only search if there are contours available
buoys = []
if len(contours) > 0:
#Show the filtered out area
#cv2.rectangle(image_frame, hull_filter_upper, hull_filter_lower, (0,255,255), thickness=2)
# find the largest contour, and compute its minimum enclosing circle, and centroid
for c in contours:
((x_loc,y_loc), radius) = cv2.minEnclosingCircle(c)
M = cv2.moments(c)
#center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
# draw a circle on the largest thing it found
cv2.circle(image_frame, (int(x_loc),int(y_loc)), int(radius), (0,0,255), 2)
# Filter out detections of the hull
if (x_loc > hull_filter_lower[0]) and (x_loc < hull_filter_upper[0]) and (y_loc > hull_filter_lower[1]) and (y_loc < hull_filter_upper[1]):
continue
buoys.append(((x_loc,y_loc),radius))
return buoys