-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
80 lines (63 loc) · 1.98 KB
/
main.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
import cv2
import matplotlib.pyplot as plt
import numpy as np
from image import (
binarize,
find_digits,
resize_digits,
insert_into_center,
get_sample_image,
)
from train import get_model
def draw_contours(frame, contours):
for img in contours:
cv2.rectangle(
frame,
(img['x'], img['y']),
(img['x'] + img['w'], img['y'] + img['h']),
(0, 0, 0),
4
)
def preprocess(digits):
return np.vstack([digit.reshape(28*28).astype(np.float)/255
for digit in digits])
def video():
model = get_model()
cap = cv2.VideoCapture(0)
cv2.namedWindow('OCR')
last_seen = "Number: NaN"
while(True):
ret, frame = cap.read()
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
thresh = binarize(gray)
key = cv2.waitKey(1)
contours = find_digits(thresh)
draw_contours(frame, contours)
if key == ord('p'):
digits = insert_into_center(resize_digits(contours))
if digits:
X = preprocess(digits)
prediction = np.argmax(model.predict(X), axis=1)
last_seen = "Number: " + "".join(map(str, prediction))
# print last_seen
# plt.imshow(np.hstack(tuple(digits)), cmap=plt.cm.binary)
# plt.show()
cv2.putText(frame, last_seen, (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, 0)
cv2.imshow('OCR', frame)
if key == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
def static_image():
model = get_model()
frame = get_sample_image()
contours = find_digits(binarize(frame.copy()))
draw_contours(frame, contours)
digits = insert_into_center(resize_digits(contours))
X = preprocess(digits)
print np.argmax(model.predict(X), axis=1)
plt.imshow(np.hstack(tuple(digits)), cmap=plt.cm.binary)
plt.show()
if __name__ == '__main__':
# video()
static_image()