WTB/job.py

130 lines
4.6 KiB
Python

import numpy as np
import cv2
from imutils.object_detection import non_max_suppression
import pytesseract
# from matplotlib import pyplot as plt
from best import best
def job(folder):
folder = f'imgs/{folder}'
# tworzenie biblioteki argumentów (ta forma ułatwia rozróżnienie tych argumentów od innych zmiennych)
args = {"images": [f"{folder}/1.jpg", f"{folder}/2.jpg", f"{folder}/3.jpg"],
"east": "frozen_east_text_detection.pb",
"min_confidence": 0.5,
"width": 320, # wielokrotność 32
"height": 416 # *wielokrotność 32
}
# *uzasadnienie:
# aplikacja przeznaczona jest dla urządzeń mobilnych, dla których naturalnym ustawieniem obiektywu jest to pionowe
# (posiadające więcej px w wysokości niż szerokości). Mój telefon wykonuje zdjęcia w rodzielczości 3000x4000 px,
# zatem ich stosunek wynosi 1,(3). Domyślną wartością tutaj jest 320 px, zatem 320 * 1,3 = 416
images = []
for i in range(len(args['images'])):
images.append(cv2.imread(args['images'][i]))
# zapisywanie oryginalnych wymiarów (zdjęcia z telefonu mają takie same wymiary, więc wystarczy sprawdzenie jednego)
orig = images[0].copy()
(origH, origW) = images[0].shape[:2]
# zdefiniowanie rozmiaru czytanego obrazka (najlepiej wielokrotność liczby 32)
(newW, newH) = (args["width"], args["height"])
# obliczanie stosunku wymiarów zdjęcia do poprawnego przeskalowania (r = ratio)
rW = origW / float(newW)
rH = origH / float(newH)
# przeskalowanie zdjęć
for i in range(len(images)):
images[i] = cv2.resize(images[i], (newW, newH))
(H, W) = images[0].shape[:2]
# tworzenie blobów, które zostaną przekazane do modelu EAST w celu wykrycia tekstu
blob = cv2.dnn.blobFromImages(images, 1.0, (W, H), (123.68, 116.78, 103.94), swapRB=True, crop=False)
# załadowanie modelu EAST do wykrywania tekstu
net = cv2.dnn.readNet(args["east"])
# zgodnie z poleceniem ze źródła wyznaczamy dwie warstwy z modelu EAST
layerNames = ["feature_fusion/Conv_7/Sigmoid", "feature_fusion/concat_3"]
net.setInput(blob)
(scores, geometry) = net.forward(layerNames)
def predictions(prob_score, geo):
(numR, numC) = prob_score.shape[2:4]
boxes = []
confidence_val = []
for y in range(0, numR):
scoresData = prob_score[0, 0, y]
x0 = geo[0, 0, y]
x1 = geo[0, 1, y]
x2 = geo[0, 2, y]
x3 = geo[0, 3, y]
anglesData = geo[0, 4, y]
for i in range(0, numC):
if scoresData[i] < args["min_confidence"]:
continue
(offX, offY) = (i * 4.0, y * 4.0)
angle = anglesData[i]
cos = np.cos(angle)
sin = np.sin(angle)
h = x0[i] + x2[i]
w = x1[i] + x3[i]
endX = int(offX + (cos * x1[i]) + (sin * x2[i]))
endY = int(offY - (sin * x1[i]) + (cos * x2[i]))
startX = int(endX - w)
startY = int(endY - h)
boxes.append((startX, startY, endX, endY))
confidence_val.append(scoresData[i])
return boxes, confidence_val
(boxes, confidence_val) = predictions(scores, geometry)
boxes = non_max_suppression(np.array(boxes), probs=confidence_val)
""" Znajdowanie i rozpoznawanie tekstu """
results = []
for (startX, startY, endX, endY) in boxes:
startX = int(startX * rW)
startY = int(startY * rH)
endX = int(endX * rW)
endY = int(endY * rH)
r = orig[startY:endY, startX:endX]
configuration = "-l eng --oem 1 --psm 8"
text = pytesseract.image_to_string(r, config=configuration)
results.append(((startX, startY, endX, endY), text))
orig_image = orig.copy()
recognized = []
for ((start_X, start_Y, end_X, end_Y), text) in results:
# wyświetlanie tekstu wykrytego przez Tesseract
recognized.append(format(text))
print("{}\n".format(text))
text = "".join([x if ord(x) < 128 else "" for x in text]).strip()
cv2.rectangle(orig_image, (start_X, start_Y), (end_X, end_Y), (0, 0, 255), 2)
cv2.putText(orig_image, text, (start_X, start_Y - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
# plt.imshow(orig_image)
# plt.title('Output')
# plt.show()
info = best(recognized)
return info