-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
44 lines (33 loc) · 1007 Bytes
/
test.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
import cv2
import pickle
import tensorflow as tf
import numpy as np
width = 640
height = 480
threshold = 0.65
def preprocessing(img):
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = cv2.equalizeHist(img)
img = img / 255
return img
cap = cv2.VideoCapture(0)
cap.set(3, width)
cap.set(4, height)
pickle_in = open("model_trained.p", "rb")
model = pickle.load(pickle_in)
with tf.device('/gpu:0'):
success, imgOriginal = cap.read()
img = np.asarray(imgOriginal)
img = cv2.resize(img, (32, 32))
img = preprocessing(img)
img = img.reshape(1, 32, 32, 1)
# Predict
classIndex = int(model.predict_classes(img))
predictions = model.predict(img)
probVal = np.amax(predictions)
print(classIndex, probVal)
if probVal > threshold:
cv2.putText(imgOriginal, str(classIndex) + " " + str(probVal),
(50,50), cv2.FONT_HERSHEY_SCRIPT_SIMPLEX, 1, (0,0,255), 1)
cv2.imshow("Original Image", imgOriginal)
cv2.waitKey(5000)