38 lines
985 B
Python
38 lines
985 B
Python
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
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:
|
||
|
|
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)
|