一、背景
工作项目需要用到关键点检测，yolo系列在工业视觉领域应用相对广泛，技术积累相对成熟（翻译：出bug方便搜），在检测速度和准确率方面也表现优异，毫无疑问成为首选。作为彼时yolo家族的最新款，本人也想尝尝鲜。
原始的yolov7-pose是针对人体17个关键点的检测模型，因此需要对其改造以适应项目需求（项目的使用场景中只需要3个关键点），开始漫长的改造和调试之路。
源码：GitHub - WongKinYiu/yolov7 at pose
二、具体工作
1. 准备数据集
数据标注工具使用的是labelme，生成标注文件，由于labelme无法直接生成yolo格式的标签，因此需要将其生成的coco格式的json文件转换为yolo格式的txt文件。
yolo格式的关键点标签文件内容如下：
前五个为目标类别和标注框中心点和长宽，根据图像宽高统一进行归一化，后面即为关键点坐标和关键点可见度
，
和
同样进行了归一化操作，
，0表示该点不存在，1表示该点不可见（遮挡），2表示该点可见。
转换代码网上很多，不多做介绍，主要参考[1]。
数据集生成之后，在data/coco_kpts.yaml中修改自己的数据集路径及类别信息。
2. 代码修改
数据集准备好之后就可以进行模型训练了，如果需要预训练权重可以从作者github上下载。由于官方提供的是针对人体关键点的检测代码，且并没有提供一个统一的关键点数量修改入口，在数据集加载、损失函数计算、画图等文件里默认都是1个类别和17个关键点，所以训练之前需要自己修改相关代码文件，将其中与关键点数量相关的部分全都修改为自己需要的关键点数量。
（1）cfg/yolov7-w6-pose.yaml
模型结构和anchor的配置文件，将里面的nc和nkpt修改为自己需要检测的类别数和关键点数量。
（2）model/yolo.py
Detect()对应的是yolo的head部分，处理模型最终的预测值，将其转换为我们最终需要的预测坐标和类别信息等。此脚本中作者在yolov5的Detect()中加入了对关键点预测信息的处理方法，并新增了IDetect()和IKeypoint()，三者之间大同小异，用哪个都可以，以下用IKeypoint()为例进行修改。
Line265-266，坐标种类划分
# 将6改为5+nc
x_det = x[i][..., :5 + self.nc]
x_kpt = x[i][..., 5 + self.nc:]Line283-284，关键点坐标处理
# 17改为nkpt
x_kpt[..., 0::3] = (x_kpt[..., ::3] * 2. - 0.5 + kpt_grid_x.repeat(1,1,1,1,self.nkpt)) * self.stride[i] # xy
x_kpt[..., 1::3] = (x_kpt[..., 1::3] * 2. - 0.5 + kpt_grid_y.repeat(1,1,1,1,self.nkpt)) * self.stride[i] # xy此部分修改仅在模型推理的后处理阶段有效，训练过程中并不会用到。
（3）utils/datasets.py
在create_dataloader()，LoadImagesAndLabels()，random_perspective()中添加参数nkpt
def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix='', tidl_load=False,
kpt_label=False, nkpt=0):class LoadImagesAndLabels(Dataset): # for training/testing
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
cache_images=False, single_cls=False, stride=32, pad=0.0, prefix='', square=False, tidl_load=False,
kpt_label=True, nkpt=0):def random_perspective(img, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0,
border=(0, 0), kpt_label=False, nkpt=0):处理标签数据，在LoadImagesAndLabels()中添加成员变量nkpt，同时修改操作关键点翻转的flip_index
self.nkpt = nkpt
"""
源文件中flip_index = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]，分别表示：
[0鼻子，1左眼，2右眼，3左耳，4右耳，5左键，6右肩，7左肘，8右肘，9左腕，10右腕，11左胯，12右胯，13左膝，14右膝，15左踝，16右踝]
除0之外，其它成对关键点两两交换，实现左右翻转效果，可以根据自己的关键点要求进行修改，比如标注三个关键点[0, 1, 2]，前两个关键点可以进行翻转，则flip_index=[1, 0, 2]
"""
self.flip_index = [1, 0, 2]LoadImagesAndLabels()的cache_labels()方法
Line 496
# 56改为5+nkpt*3
assert l.shape[1] == 5 + self.nkpt * 3, 'labels require 14 columns each'Line 500,505,513,517，39改为5+nkpt*2
Line 592后添加
img, labels = random_perspective(img, labels,
degrees=hyp['degrees'],
translate=hyp['translate'],
scale=hyp['scale'],
shear=hyp['shear'],
perspective=hyp['perspective'],
kpt_label=self.kpt_label,
nkpt=self.nkpt) # nkpt传参Line 986-987，random_perspective()，17改为nkpt
Line 989-997，34改为2*nkpt
（4）utils/loss.py
ComputeLoss添加参数nkpt和nc
class ComputeLoss:
# Compute losses
def __init__(self, model, autobalance=False, kpt_label=False, nkpt=17, nc=1):
super(ComputeLoss, self).__init__()
self.kpt_label = kpt_label
self.nkpt = nkpt
self.nc = ncLine 119，修改sigmas，长度=nkpt，修改原则参考[2]，简单点说sigmas中每个元素值（OKS中的）对应一个关键点的惩罚，值越小惩罚力度越大，相应模型对这些关键点的学习关注度越高，具体值的设定沿用yolo-pose[5]，论文中说是从coco验证集中得出的经验值。
原始代码中sigmas=[.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]/10.0，所以修改也尽量参考这些值。
def __call__(self, p, targets): # predictions, targets, model
device = targets.device
lcls, lbox, lobj, lkpt, lkptv = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
# 修改sigmas
sigmas = torch.tensor([.79, .79, .25], device=device) / 10.0Line 187，build_targets()，41改为7+2*nkpt
Line 202，40改为6+2*nkpt，19改为2+nkpt
（5）utils/general.py
Line 493,590，56改为3*nkpt+5
non_max_suppression()
non_max_suppression_export()
Line 594，57改为6+3*nkpt
Line 506,540-541，6改为5+nc[3]
（6）utils/common.py
网络组件，根据[2]中作者的解释，如果使用的模型是基于较新版本的yolov5， 模型组件有所变化，SPP替换为SPPF，因此需要在common.py中添加SPPF的实现。但yolov7-pose中作者提供的yolov7-w6-pose.yaml文件中没有使用SPPF模块，所以加不加都可以。
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
super().__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))（7）utils/plot.py
Line 97，skeleton，绘制骨架结构，哪两个关键点之间需要连接就添加一组，比如我检测的三个关键点，1和2需要连接，3不需要连接，可以这样写：
skeleton = [[1, 2], [3, 3]]Line 101，pose_limb_color，绘制连线颜色（猜的）；
Line 102，pose_kpt_color，绘制关键点颜色（猜的）
这两处根据自己的关键点数量修改，从palette调色板中选颜色，不改也可以。
Line 219，40改为2*nkpt+6，参考[2]热评
（8）train.py
Line 201,211，create_dataloader()传参nkpt
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
world_size=opt.world_size, workers=opt.workers,
image_weights=opt.image_weights,quad=opt.quad, prefix=colorstr('train: '), kpt_label=kpt_label,
nkpt=nkpt)Line 369,441，test()传参nkpt
results, _, _ = test.test(opt.data,
batch_size=batch_size * 2,
imgsz=imgsz_test,
conf_thres=0.001,
iou_thres=0.7,
model=attempt_load(m, device).half(),
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=True,
plots=False,
is_coco=is_coco,
kpt_label=kpt_label,
nkpt=nkpt)（9）test.py
Line 22，test()添加参数nkpt
def test(data,
weights=None,
batch_size=32,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
save_json_kpt=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_txt_tidl=False, # for auto-labelling
save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences
plots=True,
wandb_logger=None,
compute_loss=None,
half_precision=True,
is_coco=False,
opt=None,
tidl_load=False,
dump_img=False,
kpt_label=False,
flip_test=False,
nkpt=17):Line 80，flip_index，与datasets.py中LoadImagesAndLabels()的flip_index相同
Line 99，create_dataloader()传参nkpt
dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True,
prefix=colorstr(f'{task}: '), tidl_load=tidl_load, kpt_label=kpt_label, nkpt=nkpt)[0]补充（2024-08-07）
Line 266,269，plot_images()传参nkpt
plot_images(img, targets, paths, f, names, kpt_label=kpt_label, orig_shape=shapes[si], nkpt=nkpt)plot_images(img, output_to_target(out), paths, f, names, kpt_label=kpt_label, steps=3, orig_shape=shapes[si], nkpt=nkpt)至此，完成所有需要修改的位置，可以训练自己的数据集了。
三、问题
本人在初次使用过程中遇到过各种bug，除去一些低级错误和忘掉的错误，总结一下印象比较深刻的一些问题。
1. 某个关键点预测偏差大
训练完成后，发现模型的预测结果中三个关键点里总有一个点预测飞了（时隔太久图已经删了），效果类似于这样
试过是否模型没有收敛？标签生成有问题？那个参数设置的不对？模型后处理有问题？修改关键点数量的时候那个地方没改对？最终在这里找到解决方法（大佬牛逼！）
https://github.com/TexasInstruments/edgeai-yolov5/issues/47
修改loss.py中OKS的指数部分的计算方法，论文中OKS的指数的分母是，不知道为啥代码里就变成了
还有一种修改方式，改为
，但这种方式没有试过，不知道是否有效
补充（2024-09-06）实测该方法对我预测的飞点问题无效
2. 模型导出
如果需要将模型部署到边缘设备， 比如部署到CPU/NPU/TPU上进行加速，都要把模型转换到特定格式，转换之前要先把权重文件导出为onnx模型，同时修改yolo.py的Detect()部分。
def forward(self, x):
# x = x.copy() # for profiling
z = [] # inference output
self.training |= self.export
for i in range(self.nl):
if self.nkpt is None or self.nkpt==0:
x[i] = self.m[i](x[i])
else :
x[i] = torch.cat((self.m[i](x[i]), self.m_kpt[i](x[i])), axis=1)
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
# 注释掉此行
# x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
x_det = x[i][..., :6]
x_kpt = x[i][..., 6:]
if not self.training: # inference
if self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
kpt_grid_x = self.grid[i][..., 0:1]
kpt_grid_y = self.grid[i][..., 1:2]
if self.nkpt == 0:
y = x[i].sigmoid()
else:
y = x_det.sigmoid()
if self.inplace:
# 此处的y改为x_det
xy = (x_det[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (x_det[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
if self.nkpt != 0:
# 此处的17改为关键点数量
x_kpt[..., 0::3] = (x_kpt[..., ::3] * 2. - 0.5 + kpt_grid_x.repeat(1,1,1,1,self.nkpt)) * self.stride[i] # xy
x_kpt[..., 1::3] = (x_kpt[..., 1::3] * 2. - 0.5 + kpt_grid_y.repeat(1,1,1,1,self.nkpt)) * self.stride[i] # xy
#x_kpt[..., 0::3] = ((x_kpt[..., 0::3].tanh() * 2.) ** 3 * self.anchor_grid[i][:,0].repeat(self.nkpt,1).permute(1,0).view(1, self.na, 1, 1, self.nkpt)) + kpt_grid_x.repeat(1,1,1,1,17) * self.stride[i] # xy
#x_kpt[..., 1::3] = ((x_kpt[..., 1::3].tanh() * 2.) ** 3 * self.anchor_grid[i][:,0].repeat(self.nkpt,1).permute(1,0).view(1, self.na, 1, 1, self.nkpt)) + kpt_grid_y.repeat(1,1,1,1,17) * self.stride[i] # xy
x_kpt[..., 2::3] = x_kpt[..., 2::3].sigmoid()
# 此处的y改为x_det
y = torch.cat((xy, wh, x_det[..., 4:], x_kpt), dim = -1)
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
if self.nkpt != 0:
y[..., 6:] = (y[..., 6:] * 2. - 0.5 + self.grid[i].repeat((1,1,1,1,self.nkpt))) * self.stride[i] # xy
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
return x if self.training else (torch.cat(z, 1), x)
最终输出形状，参考[4]
四、总结
第一次使用关键点检测模型并进行魔改，说实话前期还是费了不少功夫的，在没有人带的情况下只能自己慢慢摸索，最终总算是走通了，也用到了项目里，过程当中也学到了不少东西，总感觉不记下来时间长了就忘记了，比如调过的bug很多已经忘了，绝不止前面两个
。所以时隔一年多，又从头捋了一遍，把这个过程尽量详细的记录下来，就当复习了。
yolov7-pose在yolo-pose的基础上改了模型结构，层数更深，多了一个特征图尺度。经过测试，在预训练权重上进行训练比从零训练准确率要高一些，yolov7-pose比yolo-pose的准确率要高，看来增加模型深度和特征图尺度确实有利于模型学习。如果想进一步提升准确率，之前尝试过加入CBAM模块，准确率有微小提升（大概1%），模型参数会增加，训练会变慢，这个根据实际需要来吧。
本人技术菜鸡，写文章纯属学习记录，大佬勿喷。
参考资料
[1] Yolov7-pose 训练body+foot关键点-CSDN博客
[2] https://zhuanlan.zhihu.com/p/603799078
[3] 基于yoloV7-pose添加任意关键点 ＋ 多类别分类网络修改_yolov7s-pose分类代码-CSDN博客
[4] GitHub - nanmi/yolov7-pose: pose detection base on yolov7
[5] https://arxiv.org/abs/2204.06806
2万+
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
