-
Notifications
You must be signed in to change notification settings - Fork 0
/
video.py
128 lines (101 loc) · 4.05 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
from sys import exit as Die
try:
import sys
import cv2
from colordetection import ColorDetector
except ImportError as err:
Die(err)
class Webcam:
def __init__(self):
self.cam = cv2.VideoCapture(0)
self.stickers = self.get_sticker_coordinates('main')
self.current_stickers = self.get_sticker_coordinates('current')
self.preview_stickers = self.get_sticker_coordinates('preview')
def get_sticker_coordinates(self, name):
stickers = {
'main': [
[200, 120], [300, 120], [400, 120],
[200, 220], [300, 220], [400, 220],
[200, 320], [300, 320], [400, 320]
],
'current': [
[20, 20], [54, 20], [88, 20],
[20, 54], [54, 54], [88, 54],
[20, 88], [54, 88], [88, 88]
],
'preview': [
[20, 130], [54, 130], [88, 130],
[20, 164], [54, 164], [88, 164],
[20, 198], [54, 198], [88, 198]
]
}
return stickers[name]
def draw_main_stickers(self, frame):
"""Draws the 9 stickers in the frame."""
for x,y in self.stickers:
cv2.rectangle(frame, (x,y), (x+30, y+30), (255,255,255), 2)
def draw_current_stickers(self, frame, state):
"""Draws the 9 current stickers in the frame."""
for index,(x,y) in enumerate(self.current_stickers):
cv2.rectangle(frame, (x,y), (x+32, y+32), ColorDetector.name_to_rgb(state[index]), -1)
def draw_preview_stickers(self, frame, state):
"""Draws the 9 preview stickers in the frame."""
for index,(x,y) in enumerate(self.preview_stickers):
cv2.rectangle(frame, (x,y), (x+32, y+32), ColorDetector.name_to_rgb(state[index]), -1)
def color_to_notation(self, color):
notation = {
#'green' : 'F',
#'white' : 'U',
#'blue' : 'B',
#'red' : 'R',
#'orange' : 'L',
#'yellow' : 'D',
'green' : 'B',
'white' : 'D',
'blue' : 'F',
'red' : 'R',
'orange' : 'L',
'yellow' : 'U'
}
return notation[color]
def scan(self):
sides = {}
preview = ['white','white','white',
'white','white','white',
'white','white','white']
state = [0,0,0,
0,0,0,
0,0,0]
while True:
_, frame = self.cam.read()
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
key = cv2.waitKey(10) & 0xff
# init certain stickers.
self.draw_main_stickers(frame)
self.draw_preview_stickers(frame, preview)
for index,(x,y) in enumerate(self.stickers):
roi = hsv[y:y+32, x:x+32]
avg_hsv = ColorDetector.average_hsv(roi)
color_name = ColorDetector.get_color_name(avg_hsv)
state[index] = color_name
# update when space bar is pressed.
if key == 32:
preview = list(state)
self.draw_preview_stickers(frame, state)
face = self.color_to_notation(state[4])
notation = [self.color_to_notation(color) for color in state]
sides[face] = notation
# show the new stickers
self.draw_current_stickers(frame, state)
# append amount of scanned sides
text = 'scanned sides: {}/6'.format(len(sides))
cv2.putText(frame, text, (20, 460), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255,255,255), 1, cv2.LINE_AA)
# quit on escape.
if key == 27:
break
# show result
cv2.imshow("default", frame)
self.cam.release()
cv2.destroyAllWindows()
return sides if len(sides) == 6 else False
webcam = Webcam()