前言
最近因为某些原因，要搞行为识别（action recognition），刚好看了YoloPose这篇文章，但是网上看到的中文资料不多，基本都是一篇文章复制粘贴，所以在这里抛砖引玉，写一篇从头开始接触的笔记。
现在源码前端时间已经在github上开源了：YoloPose
但是一开始我直接以git clone的方式下载源码，发现只能够下载到Yolov5的源码，里面没有任何和YoloPose相关的文件，我很疑惑，但是看线上仓库里文件确实不太一样，于是我直接下载压缩包，终于得到了新加的文件：
- test.py：应该是用来训练YoloPose的代码
- onnx_inference / yolo_pose_onnx_inference.py：用于onnx推理的代码
- detect.py：用于pytorch推理的代码（这个之后讲）
这篇博文主要做了以下工作：
- YoloPose的onnx推理代码添加video视频推理函数
- 解决非模型标准尺寸输入，造成的检测框和keypoint坐标偏移
先跑一下demo
想要尝尝鲜跑一下demo的话
- 安装依赖
pip install -r requirements.txt
- 下载模型文件：model（这是德州仪器处理过的轻量化onnx模型，仓库里还有很多别pytorch模型可供下载）
- 进入onnx_inference文件内，运行命令：
python yolo_pose_onnx_inference.py --img-path ./sample_ips.txt --model-path yolov5s6_pose_640_ti_lite_54p9_82p2.onnx
这样就能在onnx_inference/sample_ops_onnxrt文件夹内看到生成的结果图片了，还有储存在检测框信息的txt文档了。
如果想要跑自己的照片的话，要确保自己的照片resize到 640 x 640 的大小，还得是png格式，然后将照片的路径写到 sample_ips.txt 这个文本里。其实我也不确定前面是不是一定要这样，但是我这里不这样做的话，感觉画出来的关键点和检测框都会偏移。
感觉在后处理的时候，没有根据缩放对坐标点进行重新映射。
目前为止，至少在2022.07.22日的时候，官方只给出了图片推理的代码，没有给出视频推理的代码。因为项目要求，我觉得还是得搞一下。
模型的输入输出
因为要利用它来实现视频推理，网络内部姑且不管，把它当成一个黑箱，但是推理的输入输出还是要搞清楚的。（如果想看YoloPose原理的其实可以移步了）
输入
def read_img(img_file, img_mean=127.5, img_scale=1/127.5):
img = cv2.imread(img_file)[:,:,::-1] # bgr->rgb，对通道进行翻转
img = cv2.resize(img, (640, 640), interpolation=cv2.INTER_LINEAR)
img = (img - img_mean) * img_scale # 像素归一化
img = np.asarray(img, dtype=np.float32) # 转成float32位的数组
img = np.expand_dims(img, 0) # 在axis=0处加入一个维度，img.shape=(1,640,640,3)
img = img.transpose(0,3,1,2) # 调整维度顺序：（batchsize，channel，height，width）
return img
可以看出来，就是正常推理前图片预处理过程。
我们知道model_inference(model_path, input)就是将预处理过的图片送进网络推理的函数，所以也不用管它，只要知道输入就是一个模型的地址model_path和一个四维度的图片数据input就行了。
输出
我们需要关注的是推理之后得到的输出信息。
output通过打印信息，大概知道是一个list，现在的情况list内只有一个array，因为我们的batchsize=1（预处理阶段新加的维度），就是说我们每次送进网络进行推理的照片只有一张，如果我们一次一起送很多张照片进网络的话：list的结果应该是[array_1, array_2, ..., array_n]，每一个array都代表着一张照片的信息。
因为batchsize=1，直接取output[0]来进行后续处理。output[0]讲道理就是numpy.nadarray类型的数据，里面包含了单张图片的全部信息。就拿github仓库里那两只人陪狗玩的测试图片来举例吧。这个array的维度将会是(2, 57)，就是output[0].shape = (2, 57)，因为，图里检测到2个人的骨架关键点，而描述一个人的骨架各种信息，需要57个数据去描述。下面我会对这个57个元素的含义一一进行讲解。
[0:4]：检测框的box（4）[4]：检测框置信率conf（1）[5]：检测目标的标签（1）[6:]：人的关键点有17个，每个关键点都是(x, y, conf)组成的（51）
这样就刚刚好将全部元素都分配完成了。看来onnx的预训练模型已经包含了nms的模块。（当然，你也可以将官方提供的pytorch模型转成onnx模型，但是使用export.py转的时候，需要带上包含nms模块的选项，不知道为啥，我想转成带nms模快的onnx模型老是报错，之后有时间处理一下。）
搞清楚输入和输出之后，就可以利用它们来写我们的视频推理代码了。
YoloPose视频推理
我这部分代码是参考PP-human写的
输入视频
视频推理也分两种
- 摄像头输入
- 视频文件输入
第一步肯定是输入视频和准备输出文件夹：
if camera_id != -1:
video_name = 'output.mp4' # 先给个输出视频名字
capture = cv2.VideoCapture(camera_id)
else:
video_name = os.path.split(video_file)[-1] # 直接原文件名输出
capture = cv2.VideoCaprure(video_file)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
out_path = os.path.join(output_dir, video_name)
从这里知道，参数列表里一定要有camera_id摄像头设备编号和video_file视频输入路径，和视频输出路径output_dir。
获取视频信息
包括长宽，fps（每秒几帧），总帧数
width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(capture.get(cv2.CAP_PROP_FPS))
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
输出视频配置
利用原视频的配置，来对要保存的视频进行配置：
fourcc = cv2.VideoWriter_fourcc(* 'mp4v') # 规定输出视频格式
writer= cv2.VideoWriter(out_path, fourcc, fps, (width heiht))
逐帧处理
接下来就要开始正式对视频进行处理，逐帧送进网络里进行推理了。
index = 0 # 用于记录帧数
while(1):
success, frame = capture.read() # 一帧帧的读取，帧为bgr格式的
if not success: break
index += 1
print('detect frame: %d' % (index))
# 将frame输入网络中，在这之前，我们要对frame进行一些图象预处理
# 和read_img的操作一样，直接抄过来就好了
# 将img_mean和img_scale都直接用实数表示了
img = frame[:,:,::-1] # bgr->rgb，对通道进行翻转
# 图片被线性拉伸或者压缩了，后面要进行坐标的重新映射
img = cv2.resize(img, (640, 640), interpolation=cv2.INTER_LINEAR)
img = (img - mean) * scale # 像素归一化
img = np.asarray(img, dtype=np.float32) # 转成float32位的数组
img = np.expand_dims(img, 0) # 在axis=0处加入一个维度，img.shape=(1,640,640,3)
frame = img.transpose(0,3,1,2) # 调整维度顺序：（batchsize，channel，height，width）
# 因为也是一张张的传进去推理的，所以batchsize=1
output = model_inference(model_path, frame) # 这意味着，参数里也要有model_path
# 判断该帧是否有目标，没有目标的话，就直接原样输出该帧
if output[0].shape[0] == 0:
writer.write(frame)
continue
# 有目标的话就对图片进行后处理
det_bboxes = output[0][:, 0:4]
det_scores = output[0][4]
det_labels = output[0][5]
kpts = output[0][:, 6:]
for idx in range(len(det_bboxes)):
det_bbox = det_bboxes[idx] # 选取一个目标的检测框
kpt = kpts[idx] # 选取其对应的关键点数据
if det_scores[idx] > score_threshold: # 参数列表里还得有score_threshold
color_map = _CLASS_COLOR_MAP[int(det_labels[idx])]
# 这里要对坐标进行转换，已知视频的height和width，变换后为640，可知：
# 真实坐标truth = now * original_size / 640
# 画识别框
start_point = (det_bbox[0]*width / 640, det_bbox[1]*height / 640)
end_point = (det_bbox[2]*width / 640, det_bbox[3]*height / 640)
frame = cv2.rectangle(frame, start_point, end_point, color_map[::-1],2)
# 写点文字到图上去
cv2.putText(frame, "id:{}".format(int(det_labels[idx])),
(int(det_bbox[0]*width / 640 + 5),int(det_bbox[1]*width / 640 + 15)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color_map[::-1], 2)
cv2.putText(frame, "score:{:2.1f}".format(int(det_labels[idx])),
(int(det_bbox[0]*width / 640 + 5),int(det_bbox[1]*width / 640 + 30)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color_map[::-1], 2)
# 关键点需要重新定位，kpt是一个51个元素的array，（x, y, conf）排列，处理感觉好麻烦
# 感觉实时的话，这里真会浪费很多时间，终于知道为啥要求640 x 640的视频了
plot_skeleton_kpts(frame, kpt)
writer.write(frame)
if camera_id != -1:
cv2.imshow('Mask Detection', frame)
if cv2.waitKey(1) & 0xFF == ord('q'): break
writer.release()
大致上是写完了，但是我还没有测试过，明天再测试吧。
总结一下要传入的参数有：score_threshold，model_path，camera_id，video_file，output_dir，mean，scale感觉也不是很多的样子。
测试已经完成了，能够正常的输出实时的视频。还顺便修了一下非(640x640)视频输入，关键点和识别框标注偏移的问题。
完整的函数代码我放在下面：
完整代码
def video_inference(model_path, video_file, camera_id, output_dir, score_threshold=0.3, mean, scale):
video_name = 'output.mp4' # 先给个输出视频名字
if camera_id != -1:
capture = cv2.VideoCapture(camera_id)
else:
video_name = os.path.split(video_file)[-1] # 直接原文件名输出
capture = cv2.VideoCapture(video_file)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
out_path = os.path.join(output_dir, video_name)
width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(capture.get(cv2.CAP_PROP_FPS))
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
fourcc = cv2.VideoWriter_fourcc(* 'mp4v') # 规定输出视频格式
writer= cv2.VideoWriter(out_path, fourcc, fps, (width, height))
print('fps', fps)
start = time.time()
index = 0 # 用于记录帧数
while(1):
success, frame = capture.read() # 一帧帧的读取，帧为bgr格式的
if not success: break
index += 1
frame1 = frame
# 将frame输入网络中，在这之前，我们要对frame进行一些图象预处理
# 和read_img的操作一样，直接抄过来就好了
# 将img_mean和img_scale都直接用实数表示了
img = frame1[:,:,::-1] # bgr->rgb，对通道进行翻转
# 图片被线性拉伸或者压缩了，后面要进行坐标的重新映射
img = cv2.resize(img, (640, 640), interpolation=cv2.INTER_LINEAR)
img = (img - mean) * scale # 像素归一化
img = np.asarray(img, dtype=np.float32) # 转成float32位的数组
img = np.expand_dims(img, 0) # 在axis=0处加入一个维度，img.shape=(1,640,640,3)
frame1 = img.transpose(0,3,1,2) # 调整维度顺序：（batchsize，channel，height，width）
# 因为也是一张张的传进去推理的，所以batchsize=1
output = model_inference(model_path, frame1) # 这意味着，参数里也要有model_path
# 判断该帧是否有目标，没有目标的话，就直接原样输出该帧
if output[0].shape[0] == 0:
writer.write(frame)
continue
print('detect frame:{}, {} targets'.format(index, output[0].shape[0]))
# 有目标的话就对图片进行后处理
det_bboxes = output[0][:, 0:4]
det_scores = output[0][:, 4]
det_labels = output[0][:, 5]
kpts = output[0][:, 6:]
for idx in range(len(det_bboxes)):
det_bbox = det_bboxes[idx] # 选取一个目标的检测框
kpt = kpts[idx] # 选取其对应的关键点数据
for i in range(len(kpt)):
if i % 3 == 0:
kpt[i] = (kpt[i] * width) / 640
if i % 3 == 1:
kpt[i] = (kpt[i] * height) / 640
if det_scores[idx] > score_threshold: # 参数列表里还得有score_threshold
color_map = _CLASS_COLOR_MAP[int(det_labels[idx])]
# 这里要对坐标进行转换，已知视频的height和width，变换后为640，可知：
# 真实坐标truth = now * original_size / 640
# 画识别框
start_point = (round(det_bbox[0]*width / 640), round(det_bbox[1]*height / 640))
end_point = (round(det_bbox[2]*width / 640), round(det_bbox[3]*height / 640))
frame = cv2.rectangle(frame, start_point, end_point, color_map[::-1],2)
# 写点文字到图上去
cv2.putText(frame, "id:{}".format(int(det_labels[idx])),
(int(det_bbox[0]*width / 640 + 5),int(det_bbox[1]*width / 640 + 15)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color_map[::-1], 2)
cv2.putText(frame, "class:{:2.1f}".format(int(det_labels[idx])),
(int(det_bbox[0]*width / 640 + 5),int(det_bbox[1]*width / 640 + 30)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color_map[::-1], 2)
# 关键点需要重新定位，kpt是一个51个元素的array，（x, y, conf）排列，处理感觉好麻烦
# 感觉实时的话，这里真会浪费很多时间，终于知道为啥要求640 x 640的视频了
plot_skeleton_kpts(frame, kpt)
writer.write(frame)
if camera_id != -1:
cv2.namedWindow('Mask Detection', 0)
cv2.imshow('Mask Detection', frame)
if cv2.waitKey(1) & 0xFF == ord('q'): break
writer.release()
print("times:", time.time() - start)
直接在main()函数里面调用就行了。输入视频文件和摄像头都行。
如果大家有什么不了解的地方，尽管提问，一起进步，也希望大家多多拷打我。
大家随意转载，但是如果能把我这链接带上最好了！
下一篇会写一下YOLOpose的pytorch推理。