open-pose转rknn-爱代码爱编程
1. 前言
之前看到openpose的时候，想pt转onnx再转rknn，但是精度总是有问题，看到原来的openpose是caffe训练的，找了找caffe的模型，尝试转了rknn，没想到成功了，现在记录一下转化时候的坑！！！有坑的！！！
2. 模型和配置文件准备
可以去下载资源caffe-openpose模型,使用里面的模型，具体的代码细节参考博文中的代码！！！
3. 转换模型
(1) 修改pose_deploy_linevec.prototxt文件
特别重要！需要修改！！输入的w和h可以根据自己的需求设置，当然w和h大一点更加准确，可以自己尝试，大概在368就还不错，想要推理快一点设置288！！！
layer {
name:"image"
type:"Input"
top:"image"
input_param {
shape {
dim: 1
dim: 3
dim: 368
dim: 368
}
}
}(2) rknn_toolkit_1模型转换
from rknn.api import RKNN
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN()
# pre-process config
print('--> config model')
# 调整为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='./pose_deploy_linevec.prototxt', proto='caffe',
blobs='./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=True, dataset='./dataset.txt', pre_compile=True)
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export RKNN model')
ret = rknn.export_rknn('./pose_muti_pre.rknn')
if ret != 0:
print('Export model failed!')
exit(ret)
print('done')
rknn.release()(3)rknn_toolkit_2模型转换
from rknn.api import RKNN
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN()
# pre-process config
print('--> config model')
# 调整为BGR
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], quant_img_RGB2BGR=True, target_platform='rk3588')
print('done')
# Load tensorflow model
print('--> Loading model')
# 如有不同路径，请自行修改
ret = rknn.load_caffe(model='./pose_deploy_linevec.prototxt', blobs='./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')
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export RKNN model')
ret = rknn.export_rknn('./pose_3588_288.rknn')
if ret != 0:
print('Export model failed!')
exit(ret)
print('done')
# Ret = rknn.accuracy_analysis(inputs=['./single.jpeg'])
# if ret != 0:
# print('Accuracy analysis failed!')
# exit(ret)
rknn.release()4.模型推理
(1) rknn_toolkit_1模型推理
注意两点需要根据模型的输入尺寸进行修改！！！
1）输入尺寸
inWidth = 368
inHeight = 3682）
2）输出shape
output = output.reshape(1, 57, 46, 46)
from rknn.api import RKNN
import cv2
import time
import numpy as np
if __name__ == '__main__':
nPoints = 18
# 18个关键点信息以及对应的点对，就是互相连接的关键点
POSE_PAIRS = [ [1,0],[1,2],[1,5],[2,3],[3,4],[5,6],[6,7],[1,8],[8,9],[9,10],[1,11],[11,12],[12,13],[0,14],[0,15],[14,16],[15,17]]
# Create RKNN object
rknn = RKNN()
# 装载模型
rknn.load_rknn('./pose_deploy_linevec.rknn')
# init runtime environment
print('--> Init runtime environment')
# 初始化，指定开发板的型号、设备的id
ret = rknn.init_runtime(target='rk1808',)
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
# 输入尺寸是368x368
inWidth = 368
inHeight = 368
# 使用本地视频测试，也可以使用摄像头测试，填上对应的id
# cap = cv2.VideoCapture('1.mp4')
# hasFrame, frame = cap.read()
frame = cv2.imread('./group.jpg')
frame = cv2.resize(frame, (inWidth, inHeight), interpolation=cv2.INTER_CUBIC)
frameCopy = np.copy(frame)
frameWidth = frame.shape[1]
frameHeight = frame.shape[0]
threshold = 0.1
np.set_printoptions(threshold=np.inf)
# Inference
#print('--> Running model')
frameinput = np.transpose(frame, [2, 0, 1])
t = time.time()
[output] = rknn.inference(inputs=[frameinput], data_format="nchw")
# [output] = rknn.inference(inputs=[frameinput])
elapsed = time.time() - t
print('inference image: %.4f seconds.' % (elapsed))
np.set_printoptions(threshold=np.inf)
#print('done')
output = output.reshape(1, 57, 46, 46)
H = output.shape[2]
W = output.shape[3]
# Empty list to store the detected keypoints
points = []
for i in range(nPoints):
# confidence map of corresponding body's part.
probMap = output[0, i, :, :]
# Find global maxima of the probMap.
minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)
# Scale the point to fit on the original image
x = (frameWidth * point[0]) / W
y = (frameHeight * point[1]) / H
if prob > threshold :
cv2.circle(frame, (int(x), int(y)), 8, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
cv2.putText(frame, "{}".format(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, lineType=cv2.LINE_AA)
# Add the point to the list if the probability is greater than the threshold
points.append((int(x), int(y)))
else :
points.append(None)
# Draw Skeleton
for pair in POSE_PAIRS:
partA = pair[0]
partB = pair[1]
if points[partA] and points[partB]:
# 关键点用实心圆表示，点对用直线连接
cv2.line(frameCopy, points[partA], points[partB], (0, 255, 255), 2)
cv2.circle(frameCopy, points[partA], 8, (0, 0, 255), thickness=-1, lineType=cv2.FILLED)
Keypoint = 'Output-Keypoints'
cv2.namedWindow(Keypoint, cv2.WINDOW_NORMAL)
# cv2.imshow(Keypoint, frameCopy)
cv2.imwrite('Output-Keypoints.jpg', frameCopy)
# cv2.imwrite('Output-Skeleton.jpg', frame)
# print("Total time taken : {:.3f}".format(time.time() - t))
rknn.release()
(2) rknn_toolkit_2模型推理使用的是lite版本
同样修改输入和输出的形状！！！
from rknnlite.api import RKNNLite
import cv2
import time
import numpy as np
if __name__ == '__main__':
nPoints = 18
# 18个关键点信息以及对应的点对，就是互相连接的关键点
POSE_PAIRS = [ [1,0],[1,2],[1,5],[2,3],[3,4],[5,6],[6,7],[1,8],[8,9],[9,10],[1,11],[11,12],[12,13],[0,14],[0,15],[14,16],[15,17]]
# Create RKNN object
rknn_lite = RKNNLite()
# 装载模型
rknn_lite.load_rknn('./pose_deploy_linevec.rknn')
# init runtime environment
print('--> Init runtime environment')
# 初始化，指定开发板的型号、设备的id
ret = rknn_lite.init_runtime(target='rk1808',)
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
# 输入尺寸是368x368
inWidth = 368
inHeight = 368
# 使用本地视频测试，也可以使用摄像头测试，填上对应的id
# cap = cv2.VideoCapture('1.mp4')
# hasFrame, frame = cap.read()
frame = cv2.imread('./group.jpg')
frame = cv2.resize(frame, (inWidth, inHeight), interpolation=cv2.INTER_CUBIC)
frameCopy = np.copy(frame)
frameWidth = frame.shape[1]
frameHeight = frame.shape[0]
threshold = 0.1
np.set_printoptions(threshold=np.inf)
# Inference
#print('--> Running model')
frameinput = np.transpose(frame, [2, 0, 1])
t = time.time()
# [output] = rknn_lite.inference(inputs=[frameinput], data_format="nchw")
[output] = rknn_lite.inference(inputs=[frameinput])
elapsed = time.time() - t
print('inference image: %.4f seconds.' % (elapsed))
np.set_printoptions(threshold=np.inf)
#print('done')
output = output.reshape(1, 57, 46, 46)
H = output.shape[2]
W = output.shape[3]
# Empty list to store the detected keypoints
points = []
for i in range(nPoints):
# confidence map of corresponding body's part.
probMap = output[0, i, :, :]
# Find global maxima of the probMap.
minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)
# Scale the point to fit on the original image
x = (frameWidth * point[0]) / W
y = (frameHeight * point[1]) / H
if prob > threshold :
cv2.circle(frame, (int(x), int(y)), 8, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
cv2.putText(frame, "{}".format(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, lineType=cv2.LINE_AA)
# Add the point to the list if the probability is greater than the threshold
points.append((int(x), int(y)))
else :
points.append(None)
# Draw Skeleton
for pair in POSE_PAIRS:
partA = pair[0]
partB = pair[1]
if points[partA] and points[partB]:
# 关键点用实心圆表示，点对用直线连接
cv2.line(frameCopy, points[partA], points[partB], (0, 255, 255), 2)
cv2.circle(frameCopy, points[partA], 8, (0, 0, 255), thickness=-1, lineType=cv2.FILLED)
Keypoint = 'Output-Keypoints'
cv2.namedWindow(Keypoint, cv2.WINDOW_NORMAL)
# cv2.imshow(Keypoint, frameCopy)
cv2.imwrite('Output-Keypoints.jpg', frameCopy)
# cv2.imwrite('Output-Skeleton.jpg', frame)
# print("Total time taken : {:.3f}".format(time.time() - t))
rknn_lite.release()
5.检测结果展示
这个是单人的姿态检测！！！多人不适用呢
