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模型，如下：