Hugging Peers

ocr_model

☆ Save
by lsmodel ocr-recognition other
149 downloads 0 likes on ModelScope updated

Model card

简介

本仓库用于托管从OCR项目转换而来的ONNX模型。

SDK下载

#安装ModelScope
pip install modelscope



#SDK模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('lsmodel/ocr_model')

Git下载

#Git模型下载
git clone https://www.modelscope.cn/lsmodel/ocr_model.git

期望模型使用方式以及适用范围

本模型主要用于给输入图片输出图中文字内容，具体地，模型输出内容以字符串形式输出。用户可以自行尝试各种输入图片。

模型局限性以及可能的偏差

  • 模型是在中英文数据集上训练的，在其他语言的数据上有可能产生一定偏差，请用户自行评测后决定如何使用。
  • 当前版本在python3.7的CPU环境和单GPU环境测试通过，其他环境下可用性待测试。

ONNX模型使用

import torch
import numpy as np
import torch.nn.functional as F
import onnxruntime as rt
import cv2

def keepratio_resize(img):
    cur_ratio = img.shape[1] / float(img.shape[0])
    mask_height = 32
    mask_width = 804
    if cur_ratio > float(mask_width) / mask_height:
        cur_target_height = mask_height
        cur_target_width = mask_width
    else:
        cur_target_height = mask_height
        cur_target_width = int(mask_height * cur_ratio)
    img = cv2.resize(img, (cur_target_width, cur_target_height))
    mask = np.zeros([mask_height, mask_width, 3]).astype(np.uint8)
    mask[:img.shape[0], :img.shape[1], :] = img
    img = mask
    return img

img = cv2.imread('ocr_recognition.jpg') # 请在替换本地测试图片路径
img = keepratio_resize(img)
img = torch.FloatTensor(img)
chunk_img = []
for i in range(3):
    left = (300 - 48) * i
    chunk_img.append(img[:, left:left + 300, :])
merge_img = torch.cat(chunk_img, 0)
data = merge_img.view(3, 32, 300, 3) / 255.
data = data.permute(0, 3, 1, 2).cuda() 
input_data = data.cpu().numpy()

# inference
sess = rt.InferenceSession('model.onnx')
input_name = sess.get_inputs()[0].name
output_name= sess.get_outputs()[0].name
res = sess.run([output_name], {input_name: input_data})
outprobs = F.softmax(torch.tensor(res[0]), dim=-1)
preds = torch.argmax(outprobs, -1)

# load dict and CTC decode
# vocab.txt可从模型主页下载
labelMapping = dict()
with open('vocab.txt', 'r', encoding='utf-8') as f:
    lines = f.readlines()
    cnt = 2
    for line in lines:
        line = line.strip('\n')
        labelMapping[cnt] = line
        cnt += 1

batchSize, length = preds.shape
final_str_list = []
for i in range(batchSize):
    pred_idx = preds[i].cpu().data.tolist()
    last_p = 0
    str_pred = []
    for p in pred_idx:
        if p != last_p and p != 0:
            str_pred.append(labelMapping[p])
        last_p = p
    final_str = ''.join(str_pred)
    final_str_list.append(final_str)

print(final_str_list)