-
Notifications
You must be signed in to change notification settings - Fork 0
/
color_range_detector.py
105 lines (73 loc) · 2.75 KB
/
color_range_detector.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
# USAGE: You need to specify a filter and "only one" image source
#
# (python) range-detector --filter RGB --image /path/to/image.png
# or
# (python) range-detector --filter HSV --webcam
import cv2
import argparse
from operator import xor
def callback(value):
pass
def setup_trackbars(range_filter):
cv2.namedWindow("Trackbars", 0)
for i in ["MIN", "MAX"]:
v = 0 if i == "MIN" else 255
for j in range_filter:
cv2.createTrackbar("%s_%s" % (j, i), "Trackbars", v, 255, callback)
def get_arguments():
ap = argparse.ArgumentParser()
ap.add_argument('-f', '--filter', required=True,
help='Range filter. RGB or HSV')
ap.add_argument('-i', '--image', required=False,
help='Path to the image')
ap.add_argument('-w', '--webcam', required=False,
help='Use webcam', action='store_true')
ap.add_argument('-p', '--preview', required=False,
help='Show a preview of the image after applying the mask',
action='store_true')
args = vars(ap.parse_args())
if not xor(bool(args['image']), bool(args['webcam'])):
ap.error("Please specify only one image source")
if not args['filter'].upper() in ['RGB', 'HSV']:
ap.error("Please speciy a correct filter.")
return args
def get_trackbar_values(range_filter):
values = []
for i in ["MIN", "MAX"]:
for j in range_filter:
v = cv2.getTrackbarPos("%s_%s" % (j, i), "Trackbars")
values.append(v)
return values
def main():
args = get_arguments()
range_filter = args['filter'].upper()
if args['image']:
image = cv2.imread(args['image'])
if range_filter == 'RGB':
frame_to_thresh = image.copy()
else:
frame_to_thresh = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
else:
camera = cv2.VideoCapture(0)
setup_trackbars(range_filter)
while True:
if args['webcam']:
ret, image = camera.read()
if not ret:
break
if range_filter == 'RGB':
frame_to_thresh = image.copy()
else:
frame_to_thresh = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
v1_min, v2_min, v3_min, v1_max, v2_max, v3_max = get_trackbar_values(range_filter)
thresh = cv2.inRange(frame_to_thresh, (v1_min, v2_min, v3_min), (v1_max, v2_max, v3_max))
if args['preview']:
preview = cv2.bitwise_and(image, image, mask=thresh)
cv2.imshow("Preview", preview)
else:
cv2.imshow("Original", image)
cv2.imshow("Thresh", thresh)
if cv2.waitKey(1) & 0xFF is ord('q'):
break
if __name__ == '__main__':
main()