s444428/AL-2020
2
0
Fork 0
forked from s444420/AL-2020
Code Releases 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

70
coder/rocognizer.py
View File

@ -4,51 +4,53 @@ 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,)),
]) ])
img = cv2.imread(path)
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_gray = cv2.GaussianBlur(img_gray, (5, 5), 0)
ret, im_th = cv2.threshold(img_gray, 90, 255, cv2.THRESH_BINARY_INV)
ctrs, hier = cv2.findContours(im_th.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
rects = [cv2.boundingRect(ctr) for ctr in ctrs]
# load nn model # load nn model
model = Net() model = Net()
model.load_state_dict(torch.load('model.pt')) model.load_state_dict(torch.load('model.pt'))
model.eval() model.eval()
for rect in rects:
# Crop image
crop_img = img[rect[1]:rect[1] + rect[3] + 10, rect[0]:rect[0] + rect[2] + 10, 0]
# Resize the image
roi = cv2.resize(crop_img, (28, 28), interpolation=cv2.INTER_CUBIC)
# roi = cv2.dilate(roi, (3, 3))
# plt.imshow(roi)
# plt.show()
im = transform(roi)
im = im.view(1, 1, 28, 28)
with torch.no_grad():
logps = model(im)
ps = torch.exp(logps)
probab = list(ps.numpy()[0])
code.append(probab.index(max(probab)))
print(code) for path in paths:
# cv2.imshow("Code", img) img = cv2.imread(path)
# cv2.waitKey()
return code img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_gray = cv2.GaussianBlur(img_gray, (5, 5), 0)
ret, im_th = cv2.threshold(img_gray, 90, 255, cv2.THRESH_BINARY_INV)
ctrs, hier = cv2.findContours(im_th.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
rects = [cv2.boundingRect(ctr) for ctr in ctrs]
for rect in rects:
# Crop image
crop_img = img[rect[1]:rect[1] + rect[3] + 10, rect[0]:rect[0] + rect[2] + 10, 0]
# Resize the image
roi = cv2.resize(crop_img, (28, 28), interpolation=cv2.INTER_CUBIC)
# roi = cv2.dilate(roi, (3, 3))
# plt.imshow(roi)
# plt.show()
im = transform(roi)
im = im.view(1, 1, 28, 28)
with torch.no_grad():
logps = model(im)
ps = torch.exp(logps)
probab = list(ps.numpy()[0])
code.append(probab.index(max(probab)))
codes.append(code)
# cv2.imshow("Code", img)
# cv2.waitKey()
return codes
recognizer("55555.jpg")