这是用于文本识别的 MGP-STR 模型的 ONNX 导出版本。

运行方式如下：

import onnxruntime as ort
from transformers import MgpstrProcessor

processor = MgpstrProcessor.from_pretrained("alibaba-damo/mgp-str-base")

mpg_str_onnx = hf_hub_download(repo_id="ml6team/mgp-str-onnx", filename="mgp-str.onnx", repo_type="model")

providers = [('CUDAExecutionProvider', {"cudnn_conv_algo_search": "DEFAULT"}), 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider']

session = ort.InferenceSession(
        mpg_str_onnx, providers=providers,
)

image = Image.open("path_to_your_image").convert("RGB")

pixel_values = processor(images=image, return_tensors="pt").pixel_values
inputs = {"pixel_values": pixel_values.numpy()}

warmup = mgp_str_session.run(None, inputs)
outputs = session.run(None, inputs)

outputs = [torch.tensor(i) for i in outputs]
out_strs = processor.batch_decode(tuple(outputs))
print(out_strs["generated_text"])