/
cokecap.py
92 lines (72 loc) · 1.98 KB
/
cokecap.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
import cv2
import threading
import sys
from train import *
from constants import *
train = False
lock = threading.Lock()
timer = None
stop_timer = False
backup = False
# Start overlay variables
ocr_boxes = []
files = []
# End overlay variables
for arg in sys.argv:
if arg == 'train':
train = True
break
#if arg == 'backup':
# backup = True
frame = None
ocr = CokeOCR(train)
#get all of the prepopulated files
i=1
while i<=TOTAL_NUM:
files.append('iPhone'+str(i)+'.jpg')
i = i+1
#for filess in files:
# print filess
def draw_overlays():
#print boxes
for box in ocr_boxes:
[x, y, w, h] = box
cv2.rectangle(frame, (TARGET_RECTANGLE_X + x, TARGET_RECTANGLE_Y + y), (TARGET_RECTANGLE_X + x + w, TARGET_RECTANGLE_Y + y + h), (0, 255, 0), 2)
def rect_from_center(image, (X, Y), width, height, (B, G, R)):
cv2.rectangle(image, (X - width / 2, Y - height / 2), (X + width / 2, Y + height / 2), (B, G, R), 2)
cv2.circle(image, (X, Y), TARGET_CIRCLE_RADIUS, (B, G, R), 2)
def take_picture():
lock.acquire()
global ocr_boxes
ocr_boxes = ocr.ocr(frame)
print "take picture!"
if not stop_timer:
timer = threading.Timer(3, take_picture)
timer.start()
lock.release()
cv2.namedWindow("VideoWindow")
webcam = cv2.VideoCapture(0)
if webcam.isOpened(): # try to get the first frame
rval, frame = webcam.read()
else:
rval = False
if not train:
timer = threading.Timer(3, take_picture)
timer.start()
while rval:
lock.acquire()
draw_overlays()
rect_from_center(frame, (CENTER_X, CENTER_Y), TARGET_RECTANGLE_WIDTH, TARGET_RECTANGLE_HEIGHT, (0, 0, 200))
cv2.imshow("VideoWindow", frame)
rval, frame = webcam.read()
lock.release()
key = cv2.waitKey(20)
if key == 27: # exit on ESC
if train:
ocr.publish()
break
if train and key == 32:
lock.acquire()
ocr.analyze(frame)
lock.release()
stop_timer = True