概述

網上搜了一圈，關于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

總結

以上是生活随笔為你收集整理的CenterNet KeyPoints 关键点训练自己的数据的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。