openpose_caffe转rknn

openpose_caffe_to_rknn.py

from rknn.api import RKNN
import cv2
import time
import numpy as np

if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN()
    
    # pre-process config
    print('--> config model')
    # 配置模型输入，用于NPU对数据输入的预处理
    # channel_mean_value='0 0 0 255'，那么模型推理时，将会对RGB数据做如下转换
    # (R - 0)/255, (G - 0)/255, (B - 0)/255。推理时，RKNN模型会自动做均值和归一化处理
    # reorder_channel=’0 1 2’用于指定是否调整RBG顺序，设置成0 1 2即按输入的RGB顺序不做调整
    # reorder_channel=’2 1 0’表示交换0和2通道，如果输入是RGB，将会被调整为BGR。如果是BGR将会被
    # 调整为BGR    
    rknn.config(channel_mean_value='0 0 0 255', reorder_channel='2 1 0')
    print('done')

    # Load tensorflow model
    print('--> Loading model')
    ret = rknn.load_caffe(model='E:\usb_test\example\yolov3\openpose_keras_18key\pose\coco\pose_deploy_linevec.prototxt', proto='caffe',
                            blobs='E:\usb_test\example\yolov3\openpose_keras_18key\pose\coco\pose_iter_440000.caffemodel')
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')

    ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
    #ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    # Export rknn model
    print('--> Export RKNN model')
    ret = rknn.export_rknn('./coco_quantization_368_654.rknn')
    if ret != 0:
        print('Export model failed!')
        exit(ret)
    print('done')

    rknn.release()

需要注意事项：由于rknn的模型是静态的，所以，在模型量化的时候，输入尺寸被固定，根据自己的图片更改caffe模型，如下：