YoloPose的onnx视频推理

前言

最近因为某些原因，要搞行为识别（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推理。