123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
from paddleocr import PaddleOCR
|
|
import easyocr
|
|
from typing import Optional
|
|
from rapidocr_onnxruntime import RapidOCR
|
|
import langid, sys,os
|
|
from utils import contains_lang, standardize_lang
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
from logging_config import logger
|
|
# PaddleOCR
|
|
# Paddleocr supports Chinese, English, French, German, Korean and Japanese.
|
|
# You can set the parameter `lang` as `ch`, `en`, `fr`, `german`, `korean`, `japan`
|
|
# to switch the language model in order.
|
|
# need to run only once to download and load model into memory
|
|
|
|
default_languages = ['en', 'ch', 'ja', 'ko']
|
|
|
|
|
|
def _paddle_init(paddle_lang, use_angle_cls=False, use_GPU=True, **kwargs):
|
|
return PaddleOCR(use_angle_cls=use_angle_cls, lang=paddle_lang, use_GPU=use_GPU, **kwargs)
|
|
|
|
|
|
def _paddle_ocr(ocr, image) -> list:
|
|
### return a list containing the bounding box, text and confidence of the detected text
|
|
result = ocr.ocr(image, cls=False)[0]
|
|
if not isinstance(result, list):
|
|
return []
|
|
result = [ (pos, text[0], text[1]) for pos, text in result]
|
|
return result
|
|
|
|
# EasyOCR has support for many languages
|
|
|
|
def _easy_init(easy_languages: list, use_GPU=True, **kwargs):
|
|
return easyocr.Reader(easy_languages, gpu=use_GPU, **kwargs)
|
|
|
|
def _easy_ocr(ocr,image) -> list:
|
|
detected_texts = ocr.readtext(image)
|
|
return detected_texts
|
|
|
|
# RapidOCR mostly for mandarin and some other asian languages
|
|
# default only supports chinese and english
|
|
def _rapid_init(use_GPU=True, **kwargs):
|
|
return RapidOCR(use_gpu=use_GPU, **kwargs)
|
|
|
|
def _rapid_ocr(ocr, image) -> list:
|
|
return ocr(image)[0]
|
|
|
|
### Initialize the OCR model
|
|
def init_OCR(model='paddle', easy_languages: Optional[list] = ['ch_sim','en'], paddle_lang: Optional[str] = 'ch_sim', use_GPU=True):
|
|
if model == 'paddle':
|
|
paddle_lang = standardize_lang(paddle_lang)['paddleocr_lang']
|
|
return _paddle_init(paddle_lang=paddle_lang, use_GPU=use_GPU)
|
|
elif model == 'easy':
|
|
langs = []
|
|
for lang in easy_languages:
|
|
langs.append(standardize_lang(lang)['easyocr_lang'])
|
|
return _easy_init(easy_languages=langs, use_GPU=use_GPU)
|
|
elif model == 'rapid':
|
|
return _rapid_init(use_GPU=use_GPU)
|
|
|
|
### Perform OCR on the image
|
|
def _identify(ocr, image) -> list:
|
|
if isinstance(ocr, PaddleOCR):
|
|
return _paddle_ocr(ocr, image)
|
|
elif isinstance(ocr, easyocr.Reader):
|
|
return _easy_ocr(ocr, image)
|
|
elif isinstance(ocr, RapidOCR):
|
|
return _rapid_ocr(ocr, image)
|
|
else:
|
|
raise ValueError("Invalid OCR model. Please initialise the OCR model first with init() and pass it as an argument to _identify().")
|
|
|
|
|
|
### Filter out the results that are not in the source language. Slower but for a wider range of languages
|
|
# not working but also not very reliable so don't worry about it
|
|
def _id_filtered(ocr, image, lang) -> list:
|
|
lang = standardize_lang(lang)['id_model_lang']
|
|
result = _identify(ocr, image)
|
|
### Parallelise since langid is slow
|
|
def classify_text(entry):
|
|
return entry if langid.classify(entry[1])[0] == lang else None
|
|
with ThreadPoolExecutor() as executor:
|
|
results_no_eng = list(filter(None, executor.map(classify_text, result)))
|
|
return results_no_eng
|
|
|
|
|
|
# ch_sim, ch_tra, ja, ko, en input
|
|
def _id_lang(ocr, image, lang) -> list:
|
|
result = _identify(ocr, image)
|
|
lang = standardize_lang(lang)['id_model_lang']
|
|
# try:
|
|
logger.info(f"Filtering out phrases not in {lang}.")
|
|
filtered = [entry for entry in result if contains_lang(entry[1], lang)]
|
|
# except:
|
|
# logger.error(f"Selected language not part of default: {default_languages}.")
|
|
# raise ValueError(f"Selected language not part of default: {default_languages}.")
|
|
return filtered
|
|
|
|
def id_keep_source_lang(ocr, image, lang) -> list:
|
|
try:
|
|
return _id_lang(ocr, image, lang)
|
|
except ValueError:
|
|
try:
|
|
return _id_filtered(ocr, image, lang)
|
|
except Exception as e:
|
|
print(f'Probably an issue with the _id_filtered function. {e}')
|
|
sys.exit(1)
|
|
|
|
|
|
|
|
def get_words(ocr_output) -> list:
|
|
return [entry[1] for entry in ocr_output]
|
|
|
|
def get_positions(ocr_output) -> list:
|
|
return [entry[0] for entry in ocr_output]
|
|
|
|
def get_confidences(ocr_output) -> list:
|
|
return [entry[2] for entry in ocr_output]
|
|
|
|
|
|
if __name__ == '__main__':
|
|
OCR_languages = ['ch_sim','en']
|
|
reader = easyocr.Reader(OCR_languages, gpu=True)
|
|
|