Python- ocr识别模型（MNN模型）预测

import MNN
import cv2
import numpy as np
import time

def normlize_with_pad(img, width, height, ratio):
    h, w = img.shape[:2]  # 长边缩放为min_side
    
    if h / w > height / width * ratio:
        val = int(img[-1,-1])
        try:
            img = cv2.resize(img, (int(ratio * height * w // h), height))
            img = cv2.copyMakeBorder(img, 0, 0, 0, (width - int(ratio * height * w // h)),
                                     cv2.BORDER_CONSTANT, value=[0, 0, 0])
        except Exception as e:
            print('error image shape {} {}'.format(h, w))
    img = cv2.resize(img, (width, height))
    return img

def process(image_data, size):
    image_resize = normlize_with_pad(image_data, size[1],size[0], 1)
    input_data = np.array(image_resize)
    # input_data = np.ascontiguousarray(input_data)
    input_data = input_data.astype(np.float32)
    input_data = input_data / 255
    input_data = np.expand_dims(input_data, 0)
    input_data = np.expand_dims(input_data, 0)
    return input_data

def decode_out(str_index,logit, characters):
    char_list = []
    char_logit = []
    for i in range(len(str_index)):
        if str_index[i] != 0 and (not (i > 0 and str_index[i - 1] == str_index[i])):
            char_list.append(characters[str_index[i]-1])
            char_logit.append(logit[i].numpy())
    # char_l=1
    # for charl in char_logit:
    #     char_l*=charl
    # # print(char_l)
    # if not type(char_l)==int:
    #     char_l=char_l.numpy()
    #     if char_l.ndim>0:
    #         char_l=char_l[0]
    # print(char_logit)
    char_l=np.mean(char_logit)
    return ''.join(char_list),char_l




if __name__ == "__main__":
    import torch
    model_path = 'densenet_rnn.mnn'
    # image_path = '61c785d7180e455aa6a7f892a44b733f_0_1713158555.jpg'
    image_path='OCRAExtended/0.jpg'
    resize = (32, 320)
    # characters= '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz[!"#$%&()*+.,/:;<=>?@\\^-_`{|}~]'
    characters= '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz/:-'
    # characters= '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'

    t1 = time.time()
    # (1) load model
    net = MNN.nn.load_module_from_file(model_path,  ["images"], ["outputs"])
    # net = MNN.nn.load_module_from_file(model_path,  ["Input:0"], ["model/swin_tiny_patch4_window7_224/out/truediv:0"])
    
    # preprocess
    print(image_path)
    image_data = cv2.imread(image_path,0)
    input_data = process(image_data, resize)
  
    # (2) 构建一个Var类型的占位符来保存numpy，placeholder(shape, format, dtype)
    # print(input_data.shape) #(1, 1, 320, 32)
    input_var = MNN.expr.placeholder(input_data.shape, MNN.expr.NCHW)
    input_var.write(input_data)
    # (3) cv2 read shape is NHWC, Module's need is NC4HW4, convert it
    input_var = MNN.expr.convert(input_var, MNN.expr.NC4HW4)
    # print(input_var.shape)
    # (4) inference
    output_var = net.forward(input_var)
    # print(output_var.shape) #[1, 160, 37]
    # (5) the output from net may be NC4HW4, turn to linear layout
    # output_var = MNN.expr.convert(output_var, MNN.expr.NCHW)
    # print(output_var.shape)
    output_var = output_var.read()
    output_var = torch.tensor(output_var)
    # print(output_var.shape) #(80, 1, 66) #torch.Size([1, 160, 77])
    logit, preds = output_var.max(2)
    logit = torch.exp(logit)
    preds = preds.transpose(1, 0).contiguous().view(-1)
    # print(preds)
    lab2str,char_logit = decode_out(preds,logit,characters )
    # lab2str,char_logit = decode_out(preds,logit[0],characters )
    print(lab2str,char_logit)
    t2 = time.time()



# ./MNNConvert -f ONNX --modelFile "kang-slim.onnx" --MNNModel "kang-slim.mnn" --bizCode MNN

参考文章：Ubuntu18.04上MNN编译与使用（Python版）_mnn使用python cpu推理 demo-CSDN博客