MNN中Python示例

MNN中Python示例

MNNConvert工具转换

/home/xxx/work/sdk/MNN/build/MNNConvert -f ONNX --modelFile ./models/pfld.onnx --MNNModel ./models/pfld.mnn --bizCode biz

python示例

import MNN
import cv2
import numpy as np
import time

def process(image_data, size):
    image_resize = cv2.resize(image_data, size).astype(float)
    input_data = np.array(image_resize)
    input_data = input_data.transpose((2, 0, 1))  # HWC --> CHW
    input_data = np.expand_dims(input_data, 0)
    return input_data


def inference(model_path, input_format, input_data):

    interpreter = MNN.Interpreter(model_path)
    session = interpreter.createSession()
    input_tensor = interpreter.getSessionInput(session)
    # 获取输入维度
    print(input_tensor.getShape())

    # NCHW : MNN.Tensor_DimensionType_Caffe
    # NC4HW : MNN.Tensor_DimensionType_Caffe_C4
    # NHWC : MNN.Tensor_DimensionType_Tensorflow
    tmp_input = MNN.Tensor(input_format, MNN.Halide_Type_Float, input_data, MNN.Tensor_DimensionType_Caffe)
    start_time = time.time()
    input_tensor.copyFrom(tmp_input)
    interpreter.runSession(session)
    infer_result = interpreter.getSessionOutput(session)
    end_time = time.time()
    print("infer time %s ms " % str((end_time - start_time) * 1000))
    # 获取输出维度
    print(infer_result.getShape())
    output_data = infer_result.getData()
    return output_data


if __name__ == "__main__":
    model_path = 'models/pfld.mnn'
    image_path = 'test_keypoint.jpg'
    resize = (112, 112)
    input_format = (1, 3, 112, 112)
    image_data = cv2.imread(image_path)
    input_data = process(image_data, resize)

    start_time = time.time()
    output_data = inference(model_path, input_format, input_data)
    end_time = time.time()
    print("handle time %s ms " % str((end_time - start_time) * 1000))
    print(output_data)
    print("Done!")

推理引擎的默认输入格式是NCHW, 若模型输入格式是NHWC的，在推理之前推理引擎自动地将数据由NHWC格式转为NCHW格式 ！！！
模型输入格式是NHWC，则input_format = (1, 112, 112, 3)
模型输入格式是NCHW，则input_format = (1, 3, 112, 112)
但是ONNX模型的推理引擎输入格式构建方式依旧为MNN.Tensor_DimensionType_Caffe)，而不是MNN.Tensor_DimensionType_Tensorflow。