Fix EasyOCR model path to use script-relative directory
This commit is contained in:
52
text_recognition.py
Normal file
52
text_recognition.py
Normal file
@ -0,0 +1,52 @@
|
|||||||
|
import easyocr
|
||||||
|
import cv2
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
class DetectText():
|
||||||
|
def __init__(self, logger, threshold=0.0, allowlist=None, text_args=None):
|
||||||
|
# Set EasyOCR model storage directory (default: models/easyocr relative to this script)
|
||||||
|
default_model_dir = str(Path(__file__).parent / "models" / "easyocr")
|
||||||
|
model_storage_directory = os.environ.get('EASYOCR_MODEL_DIR', default_model_dir)
|
||||||
|
|
||||||
|
# This needs to run only once to load the model into memory
|
||||||
|
self.reader = easyocr.Reader(['en'], model_storage_directory=model_storage_directory)
|
||||||
|
self.threshold = threshold
|
||||||
|
self.logger = logger
|
||||||
|
self.allowlist = allowlist
|
||||||
|
self.text_args = text_args.split(',') if text_args else []
|
||||||
|
|
||||||
|
def detect(self, img): # expects CV2 image
|
||||||
|
|
||||||
|
if 'threshold' in self.text_args:
|
||||||
|
ret, img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
|
||||||
|
|
||||||
|
if 'blur' in self.text_args:
|
||||||
|
img = cv2.blur(img, (5, 5))
|
||||||
|
|
||||||
|
if 'grayscale' in self.text_args:
|
||||||
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||||
|
|
||||||
|
if 'mag2' in self.text_args:
|
||||||
|
mag_ratio = 2.0
|
||||||
|
else:
|
||||||
|
mag_ratio = 1.0
|
||||||
|
|
||||||
|
output = []
|
||||||
|
boxes = []
|
||||||
|
# run OCR
|
||||||
|
results = self.reader.readtext(img, allowlist=self.allowlist, mag_ratio=mag_ratio)
|
||||||
|
|
||||||
|
for res in results:
|
||||||
|
top_left = (int(res[0][0][0]), int(res[0][0][1]))
|
||||||
|
bottom_right = (int(res[0][2][0]), int(res[0][2][1]))
|
||||||
|
|
||||||
|
text = res[1]
|
||||||
|
confidence = res[2]
|
||||||
|
|
||||||
|
if confidence >= self.threshold:
|
||||||
|
output.append((text, confidence))
|
||||||
|
boxes.append([top_left, bottom_right])
|
||||||
|
|
||||||
|
return output, boxes
|
||||||
Reference in New Issue
Block a user