CenterNet KeyPoints 关键点训练自己的数据

概述

网上搜了一圈，关于CenterNet 训练关键点数据的资料非常少，而且讲得都很模糊，没法解决实际问题，也未说明细节和要素。在踏坑许久之后，才跑通CenterNet的关键点训练，于是记录一下踏坑历程，以备后忘

环境

cuda11.0
torch1.7.1
torchvision0.8.2
numpy 1.19.2
这是我的环境版本，不是非得这个版本

数据集准备

参考我的另一篇文章
COCO KeyPoints关键点数据集准备

CenterNet 代码修改

训练代码修改

我的数据集的类别是1类，关键点是3个

新的数据集代码创建

CenterNet/src/lib/datasets/datasets 目录, coco_hp.py是原来的coco keypoints官方数据集的数据集代码，我们从这文件copy一份，命名为handKeyPoints.py

修改其中内容，先看下对比图

修改代码

num_classes = 1             #类别1num_joints = 3             #关键点个数3default_resolution = [512, 512]mean = np.array([0.636623, 0.642096, 0.649946],dtype=np.float32).reshape(1, 1, 3)                  #数据集计算出的meanstd  = np.array([0.318729, 0.316616, 0.297199],                #数据集计算出的stddtype=np.float32).reshape(1, 1, 3)flip_idx = [[1, 2]]                          #图像翻转，这个我也不懂， 就照着官方的写了一个，跟实际的关键点数对应def __init__(self, opt, split):super(HandKeyPoints, self).__init__()self.edges = [[0, 1], [1, 2]]self.acc_idxs = [1, 2, 3]#数据集文件夹 目录 data/HandCupKeyPoints ,  这个文件夹里面是annotations、test2017 、train2017三个文件夹self.data_dir = os.path.join(opt.data_dir, 'HandCupKeyPoints')                   if split == 'val':                                                                                      #这里我们的是test而不是val，所以改一下split = 'test'self.img_dir = os.path.join(self.data_dir, '{}2017'.format(split))if split == 'test':self.annot_path = os.path.join(self.data_dir, 'annotations', 'test.json')                                                       #直接指定文件名else:self.annot_path = os.path.join(self.data_dir, 'annotations', 'train.json')                                                    #直接指定文件名

这个文件就修改完毕，其他不用动

修改文件 CenterNet/src/lib/datasets/dataset_factory.py

对比如下:

我这里大小写不一致，不用管， dataset_factory 字段 key 就是刚刚创建的handKeyPoints.py 的前缀， value就是文件里 class 的类名

修改文件 CenterNet/src/lib/opts.py

第15行修改默认数据集为 handKeyPoints 数据集

第323行

   opt.flip_idx = False#dataset.flip_idxopt.heads = {'hm': opt.num_classes, 'wh': 2, 'hps': 34}          #17个点的x、y 共 34个值

修改为:

    # opt.flip_idx = Falseopt.flip_idx = dataset.flip_idxopt.heads = {'hm': opt.num_classes, 'wh': 2, 'hps': 6}     #3个点的x、y 共 6个值

第345行

 'multi_pose': {'default_resolution': [512, 512], 'num_classes': 1, 'mean': [0.408, 0.447, 0.470], 'std': [0.289, 0.274, 0.278],               #数据集的 mean std'dataset': 'coco_hp', 'num_joints': 17,                                             #关键点个数 17个'flip_idx': [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]]},                                               #flip_idx

修改为自己数据集的 mean 和std
flip_idx 修改为相应的

     'multi_pose': {'default_resolution': [512, 512], 'num_classes': 1, 'mean': [0.614, 0.612, 0.622], 'std': [0.348, 0.347, 0.329],                   #数据集的 mean std'dataset': 'handKeyPoints', 'num_joints': 3,                                            #关键点个数 3个'flip_idx': [[1, 2]]},                                                                                   #flip_idx 看情况写

到此，训练部分的代码就修改完了，可以开始训练了

编写脚本train.sh，内容

python main.py multi_pose  --arch dla_34 --dataset handKeyPoints --lr 0.25e-4  --batch_size 16 --gpus 0  --load_model ../models/multi_pose_dla_3x.pth

学习率自定义设置，
–batch_size 依据GPU显存大小，如果CUDA out of memery 就改小点
–load_model …/models/multi_pose_dla_3x.pth 这个是官方训练好的模型，

开始训练

报一堆警告，不用管，如下：

输出训练过程日志

看损失值，在验证损失不再下降的时候停止训练

训练的结果

训练的结果保存在exp目录下

模型文件如下：

这样我们就得到了训练好的模型，
我的数据集较少，不知道是由于标注质量的问题还是学习率的问题，最终验证损失loss=2.7左右时不再下降。

测试

训练代码修改

修改文件 CenterNet/src/lib/utils/debugger.py

增加handKeyPoints 数据集类别判断

第45行增加

    elif num_classes == 1 or dataset == 'handKeyPoints':               #增加自定义的handKeyPoints类别self.names = handKeyPoints_class                                         #类别名称self.names = ['p']                                                                      #类别名称self.num_class = 1                                                                    #类别数self.num_joints = 3                                                                    #关键点数量self.edges = [[0, 1], [1, 2]]                                                        #关键点连接关系self.ec = [(255, 0, 0), (0, 0, 255), (255, 0, 0)]                          #颜色self.colors_hp = [(255, 0, 255), (255, 0, 0), (0, 0, 255)]

第467行增加

handKeyPoints_class = ['hand']

修改文件 CenterNet/src/lib/detectors/multi_pose.py

第85行

 dets[:, :, :4] *= self.opt.down_ratio                                  #bbox 的4个值  dets[:, :, 5:39] *= self.opt.down_ratioemmina.mo               #5~39 是17个关键点的坐标34个值

这里修改为：

 dets[:, :, :4] *= self.opt.down_ratio                          #bbox 的4个值  dets[:, :, 5:11] *= self.opt.down_ratioemmina.mo       #5~11 是3个关键点的坐标6个值

第101行

     debugger.add_coco_bbox(bbox[:4], 0, bbox[4], img_id='multi_pose')            #bbox 的4个值  debugger.add_coco_hp(bbox[5:39], img_id='multi_pose')                                #5~39 是17个关键点的坐标34个值

修改为

     debugger.add_coco_bbox(bbox[:4], 0, bbox[4], img_id='multi_pose')            #bbox 的4个值  debugger.add_coco_hp(bbox[5:11], img_id='multi_pose')                                #5~11 是3个关键点的坐标6个值

到这里测试代码就修改好了

编写测试脚本test_keyPoints.sh, 内容如下:

python demo.py multi_pose --demo ../images/16.jpg --load_model ../models/model_best_keypoints.pth

模型用的就是上述训练好的模型

测试结果

识别出了手臂，置信度为0.9