sbt-idp/cope2n-ai-fi/common/ocr.py

from common.utils.ocr_yolox import OcrEngineForYoloX_ID_Driving
from common.utils.word_formation import Word, words_to_lines

det_ckpt = "yolox-s-general-text-pretrain-20221226"
cls_ckpt = "satrn-lite-general-pretrain-20230106"

engine = OcrEngineForYoloX_ID_Driving(det_ckpt, cls_ckpt)
import logging
import logging.config
from utils.logging.logging import LOGGER_CONFIG
# Load the logging configuration
logging.config.dictConfig(LOGGER_CONFIG)
# Get the logger
logger = logging.getLogger(__name__)


def ocr_predict(image):
    """Predict text from image

    Args:
        image_path (str): _description_

    Returns:
        list: list of words
    """
    try:
        lbboxes, lwords = engine.run_image(image)
        lWords = [Word(text=word, bndbox=bbox) for word, bbox in zip(lwords, lbboxes)]
        list_lines, _ = words_to_lines(lWords)
        return list_lines
    except AssertionError as e:
        logger.info(e)
        list_lines = []
        return list_lines


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=str, required=True)
    args = parser.parse_args()

    list_lines = ocr_predict(args.image)