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