-
Notifications
You must be signed in to change notification settings - Fork 0
/
Filter.py
70 lines (49 loc) · 1.81 KB
/
Filter.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
import numpy as np
from collections import Counter
import cv2
# Set the link between method name and function
def select_method(method_name):
if method_name == 'mean':
return cv2.mean
elif method_name == 'hsv':
return hsv_dominant
elif method_name == 'rgb':
return bgr_dominant
elif method_name == 'kmean':
return kmean_dominant
# Return the dominant using the kmean method
def kmean_dominant(frame):
return generate_kmean_image(frame, 8, True)
# Return the dominant using hsv histogram
def hsv_dominant(frame):
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
hist = cv2.calcHist([hsv], [0, 1, 2], None, [180, 256, 256], [0, 180, 0, 256, 0, 256])
# Get max position
i, j, k = np.unravel_index(hist.argmax(), hist.shape)
color = np.uint8([[[i, j, k]]])
return cv2.cvtColor(color, cv2.COLOR_HSV2BGR)[0][0]
# Return the dominant using bgr histogram
def bgr_dominant(frame):
hist = cv2.calcHist([frame], [0, 1, 2], None, [256, 256, 256], [0, 255, 0, 255, 0, 255])
# Get max position
i, j, k = np.unravel_index(hist.argmax(), hist.shape)
color = np.uint8([[[i, j, k]]])
return color[0][0]
def most_common(lst):
data = Counter(lst)
return data.most_common(1)[0][0]
# Return the kmean image or the dominant color
def generate_kmean_image(frame, k, dominant=False, display=False):
z = frame.reshape((-1, 3))
z = np.float32(z)
# K-Mean
criteria = (cv2.TermCriteria_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
ret, label, center = cv2.kmeans(z, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
center = np.uint8(center)
res = center[label.flatten()]
res2 = res.reshape(frame.shape)
if display:
cv2.imshow('res2', res2)
if dominant:
return center[most_common(label[0])]
return res2