From 3ed283eeafb762b64b373b487e27dfb8d0704816 Mon Sep 17 00:00:00 2001 From: shaaqu Date: Mon, 8 Jun 2020 14:55:15 +0200 Subject: [PATCH] recognizer.py is returnig list od codess --- coder/rocognizer.py | 70 +++++++++++++++++++++++---------------------- 1 file changed, 36 insertions(+), 34 deletions(-) diff --git a/coder/rocognizer.py b/coder/rocognizer.py index 3e84c24..fbbcda8 100644 --- a/coder/rocognizer.py +++ b/coder/rocognizer.py @@ -4,51 +4,53 @@ from nn_model import Net from torchvision.transforms import transforms -def recognizer(a_path): +def recognizer(paths): + + codes = [] code = [] - path = a_path transform = transforms.Compose([transforms.ToTensor(), 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 model = Net() model.load_state_dict(torch.load('model.pt')) 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) - # cv2.imshow("Code", img) - # cv2.waitKey() - return code + for path in paths: + 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] + + 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") -