40 lines
1.0 KiB
Python
Executable File
40 lines
1.0 KiB
Python
Executable File
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)
|
|
|
|
|
|
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:
|
|
print(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)
|