sbt-idp/cope2n-ai-fi/api/OCRBase/text_recognition.py

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


det_ckpt = "/models/sdsvtd/hub/wild_receipt_finetune_weights_c_lite.pth"
cls_ckpt = "satrn-lite-general-pretrain-20230106"

engine = OcrEngineForYoloX_Invoice(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(img):
    """Predict text from image

    Args:
        image_path (str): _description_

    Returns:
        list: list of words
    """
    try:
        lbboxes, lwords = engine.run_image(img)
        lWords = [Word(text=word, bndbox=bbox) for word, bbox in zip(lwords, lbboxes)]
        list_lines, _ = words_to_lines(lWords)
        return list_lines
        # return lbboxes, lwords
    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)