-
Notifications
You must be signed in to change notification settings - Fork 5
/
main.py
54 lines (38 loc) · 1.48 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
import numpy as np
import cv2
import torch
from model import Net
cap = cv2.VideoCapture(0)
cap.set(3, 700)
cap.set(4, 480)
modelo = torch.load('model_trained.pt')
modelo.eval()
signs = {'0': 'A', '1': 'B', '2': 'C', '3': 'D', '4': 'E', '5': 'F', '6': 'G', '7': 'H', '8': 'I',
'10': 'K', '11': 'L', '12': 'M', '13': 'N', '14': 'O', '15': 'P', '16': 'Q', '17': 'R',
'18': 'S', '19': 'T', '20': 'U', '21': 'V', '22': 'W', '23': 'X', '24': 'Y' }
while True:
ret, frame = cap.read()
# Lugar de la imagen donde se toma la muestra
img = frame[20:250, 20:250]
res = cv2.resize(img, dsize=(28, 28), interpolation = cv2.INTER_CUBIC)
res = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY)
res1 = np.reshape(res, (1, 1, 28, 28)) / 255
res1 = torch.from_numpy(res1)
res1 = res1.type(torch.FloatTensor)
out = modelo(res1)
# Probabilidades
probs, label = torch.topk(out, 25)
probs = torch.nn.functional.softmax(probs, 1)
pred = out.max(1, keepdim=True)[1]
if float(probs[0,0]) < 0.4:
texto_mostrar = 'Sign not detected'
else:
texto_mostrar = signs[str(int(pred))] + ': ' + '{:.2f}'.format(float(probs[0,0])) + '%'
font = cv2.FONT_HERSHEY_SIMPLEX
frame = cv2.putText(frame, texto_mostrar, (60,285), font, 1, (255,0,0), 2, cv2.LINE_AA)
frame = cv2.rectangle(frame, (20, 20), (250, 250), (0, 255, 0), 3)
cv2.imshow('Cam', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()