s444420/AL-2020
1
0
Fork 2
Code Issues Pull Requests Releases Wiki Activity

recognizer.py is returnig list od codess

Browse Source
This commit is contained in:
shaaqu 2020-06-08 14:55:15 +02:00
parent d8b857bb0c
commit 3ed283eeaf
1 changed files with 36 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
coder/rocognizer.py
View File

@ -4,14 +4,22 @@ from nn_model import Net
from torchvision.transforms import transforms from torchvision.transforms import transforms
def recognizer(a_path): def recognizer(paths):
codes = []
code = [] code = []
path = a_path
transform = transforms.Compose([transforms.ToTensor(), transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)), transforms.Normalize((0.5,), (0.5,)),
]) ])
# load nn model
model = Net()
model.load_state_dict(torch.load('model.pt'))
model.eval()
for path in paths:
img = cv2.imread(path) img = cv2.imread(path)
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
@ -23,11 +31,6 @@ def recognizer(a_path):
rects = [cv2.boundingRect(ctr) for ctr in ctrs] rects = [cv2.boundingRect(ctr) for ctr in ctrs]
# load nn model
model = Net()
model.load_state_dict(torch.load('model.pt'))
model.eval()
for rect in rects: for rect in rects:
# Crop image # Crop image
crop_img = img[rect[1]:rect[1] + rect[3] + 10, rect[0]:rect[0] + rect[2] + 10, 0] crop_img = img[rect[1]:rect[1] + rect[3] + 10, rect[0]:rect[0] + rect[2] + 10, 0]
@ -44,11 +47,10 @@ def recognizer(a_path):
probab = list(ps.numpy()[0]) probab = list(ps.numpy()[0])
code.append(probab.index(max(probab))) code.append(probab.index(max(probab)))
print(code) codes.append(code)
# cv2.imshow("Code", img) # cv2.imshow("Code", img)
# cv2.waitKey() # cv2.waitKey()
return code
return codes
recognizer("55555.jpg")