YOLOv7-Pose尝鲜！基于YOLOv7的关键点模型测评

↑ 点击 蓝字  关注极市平台

作者丨陈er

来源丨GiantPandaCV

编辑丨极市平台

极市导读

 

本文主要围绕以YOLO为基线的关键点检测器，介绍了该框架的演变，并提供了在onnxruntime的推理框架下实现YOLOv7-Pose的具体代码和相应解释。 >>极市七夕粉丝福利活动：搞科研的日子是364天，但七夕只有一天！

前言

目前人体姿态估计总体分为Top-down和Bottom-up两种，与目标检测不同，无论是基于热力图或是基于检测器处理的关键点检测算法，都较为依赖计算资源，推理耗时略长，今年出现了以YOLO为基线的关键点检测器。玩过目标检测的童鞋都知道YOLO以及各种变种目前算是工业落地较多的一类检测器，其简单的设计思想，长期活跃的社区生态，使其始终占据着较高的话题度。

演变

在ECCV 2022和CVPRW 2022会议上，YoLo-Pose和KaPao（下称为yolo-like-pose）都基于流行的YOLO目标检测框架提出一种新颖的无热力图的方法，类似于很久以前谷歌使用回归计算关键点的思想，yolo-like-pose一不使用检测器进行二阶处理，二部使用热力图拼接，虽然是一种暴力回归关键点的检测算法，但在处理速度上具有一定优势。

KaPao

去年11月，滑铁卢大学率先提出了 KaPao：Rethinking Keypoint Representations: Modeling Keypoints and Poses as Objects for Multi-Person Human Pose Estimation，基于YOLOv5进行关键点检测，该文章目前已被ECCV 2022接收，该算法所取得的性能如下：

论文链接： https://arxiv.org/abs/2111.08557 

代码链接： https://github.com/wmcnally/kapao

yolov5-pose

今年4月，yolo-pose也挂在了arvix，在论文中，通过调研发现 HeatMap 的方式普遍使用L1 Loss。然而，L1损失并不一定适合获得最佳的OKS。且由于HeatMap是概率图，因此在基于纯HeatMap的方法中不可能使用OKS作为loss，只有当回归到关键点位置时，OKS才能被用作损失函数。因此，yolo-pose使用oks loss作为关键点的损失

详细代码可见： https://github.com/TexasInstruments/edgeai-yolov5/blob/yolo-pose/utils/loss.py

if self.kpt_label:
                    #Direct kpt prediction
                    pkpt_x = ps[:, 6::3] * 2. - 0.5
                    pkpt_y = ps[:, 7::3] * 2. - 0.5
                    pkpt_score = ps[:, 8::3]
                    #mask
                    kpt_mask = (tkpt[i][:, 0::2] != 0)
                    lkptv += self.BCEcls(pkpt_score, kpt_mask.float()) 
                    #l2 distance based loss
                    #lkpt += (((pkpt-tkpt[i])*kpt_mask)**2).mean()  #Try to make this loss based on distance instead of ordinary difference
                    #oks based loss
                    d = (pkpt_x-tkpt[i][:,0::2])**2 + (pkpt_y-tkpt[i][:,1::2])**2
                    s = torch.prod(tbox[i][:,-2:], dim=1, keepdim=True)
                    kpt_loss_factor = (torch.sum(kpt_mask != 0) + torch.sum(kpt_mask == 0))/torch.sum(kpt_mask != 0)
                    lkpt += kpt_loss_factor*((1 - torch.exp(-d/(s*(4*sigmas**2)+1e-9)))*kpt_mask).mean()

相关性能如下：

yolov7-pose

上个星期，YOLOv7的作者也放出了关于人体关键点检测的模型，该模型基于YOLOv7-w6。

目前作者提供了.pt文件和推理测试的脚本，有兴趣的童靴可以去看看，本文的重点更偏向于对yolov7-pose.pt进行onnx文件的抽取和推理。

yolov7-pose + onnxruntime

首先下载好官方的预训练模型，使用提供的脚本进行推理：

% weigths = torch.load('weights/yolov7-w6-pose.pt')  
% image = cv2.imread('sample/pose.jpeg')  
!python pose.py   

一、yolov7-w6 VS yolov7-w6-pose

1. 首先看下yolov7-w6使用的检测头

• 表示一共有四组不同尺度的检测头，分别为15×15,30×30,60×60,120×120，对应输出的节点为114,115,116,117
• nc对应coco的80个类别
• no表示

2. 再看看yolov7-w6-pose使用的检测头：

上述重复的地方不累述，讲几个点：

• 代表person一个类别
• nkpt表示人体的17个关键点

二、修改export脚本

如果直接使用export脚本进行onnx的抽取一定报错，在上一节我们已经看到pose.pt模型使用的检测头为IKeypoint，那么脚本需要进行相应更改：在export.py的这个位置插入：

    # 原代码:
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
        if isinstance(m, models.common.Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
     model.model[-1].export = not opt.grid  # set Detect() layer grid export
                
    # 修改代码:
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
        if isinstance(m, models.common.Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        elif isinstance(m, models.yolo.IKeypoint):
            m.forward = m.forward_keypoint  # assign forward (optional)
            # 此处切换检测头
    model.model[-1].export = not opt.grid  # set Detect() layer grid export

forward_keypoint在原始的yolov7 repo源码中有，作者已经封装好，但估计是还没打算开放使用。使用以下命令进行抽取：

python export.py --weights 'weights/yolov7-w6-pose.pt' --img-size 960 --simplify True  

抽取后的onnx检测头：

三、onnxruntime推理

onnxruntime推理代码：

import onnxruntime  
import matplotlib.pyplot as plt  
import torch  
import cv2  
from torchvision import transforms  
import numpy as np  
from utils.datasets import letterbox  
from utils.general import non_max_suppression_kpt  
from utils.plots import output_to_keypoint, plot_skeleton_kpts  
  
device = torch.device("cpu")  
  
image = cv2.imread('sample/pose.jpeg')  
image = letterbox(image, 960, stride=64, auto=True)[0]  
image_ = image.copy()  
image = transforms.ToTensor()(image)  
image = torch.tensor(np.array([image.numpy()]))  
  
print(image.shape)  
sess = onnxruntime.InferenceSession('weights/yolov7-w6-pose.onnx')  
out = sess.run(['output'], {'images': image.numpy()})[0]  
out = torch.from_numpy(out)  
  
output = non_max_suppression_kpt(out, 0.25, 0.65, nc=1, nkpt=17, kpt_label=True)  
output = output_to_keypoint(output)  
nimg = image[0].permute(1, 2, 0) * 255  
nimg = nimg.cpu().numpy().astype(np.uint8)  
nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)  
for idx in range(output.shape[0]):  
    plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)  
  
# matplotlib inline  
plt.figure(figsize=(8, 8))  
plt.axis('off')  
plt.imshow(nimg)  
plt.show()  
plt.savefig("tmp") 

推理效果几乎无损，但耗时会缩短一倍左右，另外有几个点：

• image = letterbox(image, 960, stride=64, auto=True)[0] 中stride指的是最大步长，yolov7-w6和yolov5s下采样多了一步，导致在8，16，32的基础上多了64的下采样步长
• output = non_max_suppression_kpt(out, 0.25, 0.65, nc=1, nkpt=17, kpt_label=True) ，nc 和 kpt_label 等信息在netron打印模型文件时可以看到
• 所得到的onnx相比原半精度模型大了将近三倍，后续排查原因
• yolov7-w6-pose极度吃显存，推理一张960×960的图像，需要2-4G的显存，训练更难以想象

公众号后台回复“ECCV2022”获取论文分类资源下载～

‍

‍△点击卡片关注极市平台，获取 最新CV干货‍

‍

极市干货

算法项目： CV工业项目落地实战 ｜ 目标检测算法上新！（年均分成5万）

实操教程 ：Pytorch - 弹性训练原理分析｜《CUDA C 编程指南》导读

极视角动态： 极视角作为重点项目入选「2022青岛十大资本青睐企业」榜单！ ｜ 极视角发布EQP激励计划，招募优质算法团队展开多维度生态合作！ ｜极市AI校园大使招募

“

点击阅读原文进入CV社区

收获更多技术干货