数据集
- 数据集格式
-
如果是labelme格式的数据需要使用下面的脚本转换成coco格式
-
转换成coco格式后需要在mmpose\datasets\datasets目录下选则一个目录在其下面添加自己dataset读取文件,该文件用于训练时数据的载入,如果对数据读取无特殊要求则按以下方法编写dataset文件,如果有特殊要求则需要自己重写载入方法
-
==METAINFO==选则元信息文件,元信息文件记录了关键点的数量,连接方式颜色以及权重等重要参数
from mmpose.registry import DATASETS from ..base import BaseCocoStyleDataset @DATASETS.register_module(name='MyHandCocoDataset', force=True) class MyHandCocoDataset(BaseCocoStyleDataset): METAINFO: dict = dict(from_file='E:/Pose/mmpose-main/configs/_base_/datasets/coco_wholebody_hand.py')
- 修改完成后需要在init文件里面进行注册
-
–class_name_输入类别名,–input_输入数据集的路径,–output文件输出目录,–join_num关键点的个数,–ratio数据集划分比例
import os import sys import glob import json import shutil import argparse import numpy as np from tqdm import tqdm from sklearn.model_selection import train_test_split class Labelme2coco_keypoints(): def __init__(self, args): """ Lableme 关键点数据集转 COCO 数据集的构造函数: Args args:命令行输入的参数 - class_name 根类名字 """ self.classname_to_id = {args.class_name: 1} self.images = [] self.annotations = [] self.categories = [] self.ann_id = 0 self.img_id = 0 def read_jsonfile(self, path): with open(path, "r", encoding='utf-8') as f: return json.load(f) def save_coco_json(self, instance, save_path): json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1) def _init_categories(self): """ 初始化 COCO 的 标注类别 例如: "categories": [ { "supercategory": "hand", "id": 1, "name": "hand", "keypoints": [ "wrist", "thumb1", "thumb2", ..., ], "skeleton": [ ] } ] """ for name, id in self.classname_to_id.items(): category = {} category['supercategory'] = name category['id'] = id category['name'] = name # 21 个关键点数据 category['keypoints'] = [ "wrist", "thumb1", "thumb2", "thumb3", "thumb4", "forefinger1", "forefinger2", "forefinger3", "forefinger4", "middle_finger1", "middle_finger2", "middle_finger3", "middle_finger4", "ring_finger1", "ring_finger2", "ring_finger3", "ring_finger4", "pinky_finger1", "pinky_finger2", "pinky_finger3", ] category['skeleton'] = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [5, 9], [9, 10], [10, 11], [11, 12],[9, 13], [13, 14], [14, 15], [15, 16], [13, 17], [17, 18], [18, 19], [19, 20]] self.categories.append(category) def _image(self, obj, path): """ 解析 labelme 的 obj 对象,生成 coco 的 image 对象 生成包括:id,file_name,height,width 4个属性 示例: { "file_name": "training/rgb/00031426.jpg", "height": 224, "width": 224, "id": 31426 } """ # image = {} # image["height"] = obj["imageHeight"] # image["width"] = obj["imageWidth"] # image['id'] = self.img_id # self.img_id += 1 image = {} image["height"] = obj["imageHeight"] image["width"] = obj["imageWidth"] self.img_id = self.img_id + 1 image['id'] = self.img_id image['file_name'] = os.path.basename(path).replace(".json", ".jpg") return image def _get_box(self, points): min_x = min_y = np.inf max_x = max_y = 0 for x, y in points: min_x = min(min_x, x) min_y = min(min_y, y) max_x = max(max_x, x) max_y = max(max_y, y) return [min_x, min_y, max_x-min_x, max_y-min_y] def _get_keypoints(self, points, keypoints, num_keypoints): """ 解析 labelme 的原始数据, 生成 coco 标注的 关键点对象 例如: "keypoints": [ 67.06149888292556, # x 的值 122.5043507571318, # y 的值 1, # 相当于 Z 值,如果是2D关键点 0:不可见 1:表示可见。 82.42582269256718, 109.95672933232304, 1, ..., ], """ if points[0] == 0 and points[1] == 0: visable = 0 else: visable = 2 num_keypoints += 1 keypoints.extend([points[0], points[1], visable]) return keypoints, num_keypoints def _annotation(self, bboxes_list, keypoints_list, json_path): """ 生成coco标注 Args: bboxes_list: 矩形标注框 keypoints_list: 关键点 json_path:json文件路径 """ if len(keypoints_list) != args.join_num * len(bboxes_list): print('you loss {} keypoint(s) with file {}'.format(args.join_num * len(bboxes_list) - len(keypoints_list), json_path)) print('Please check !!!') sys.exit() i = 0 for object in bboxes_list: annotation = {} keypoints = [] num_keypoints = 0 label = object['label'] bbox = object['points'] x1, y1, x2, y2 = self._get_box(bbox) annotation['id'] = self.ann_id annotation['image_id'] = self.img_id annotation['category_id'] = int(self.classname_to_id[label]) annotation['iscrowd'] = 0 annotation['area'] = (y2 - y1) * (x2 - x1) annotation['segmentation'] = [np.asarray(bbox).flatten().tolist()] annotation['bbox'] = self._get_box(bbox) for keypoint in keypoints_list[i * args.join_num: (i + 1) * args.join_num]: point = keypoint['points'] annotation['keypoints'], num_keypoints = self._get_keypoints(point[0], keypoints, num_keypoints) annotation['num_keypoints'] = num_keypoints i += 1 self.ann_id += 1 self.annotations.append(annotation) def to_coco(self, json_path_list): """ Labelme 原始标签转换成 coco 数据集格式,生成的包括标签和图像 Args: json_path_list:原始数据集的目录 """ self._init_categories() for json_path in tqdm(json_path_list): obj = self.read_jsonfile(json_path) # 解析一个标注文件 self.images.append(self._image(obj, json_path)) # 解析图片 shapes = obj['shapes'] # 读取 labelme shape 标注 bboxes_list, keypoints_list = [], [] for shape in shapes: if shape['shape_type'] == 'rectangle': # bboxs bboxes_list.append(shape) # keypoints elif shape['shape_type'] == 'point': keypoints_list.append(shape) self._annotation(bboxes_list, keypoints_list, json_path) keypoints = {} keypoints['info'] = {'description': 'Lableme Dataset', 'version': 1.0, 'year': 2021} keypoints['license'] = ['BUAA'] keypoints['images'] = self.images keypoints['annotations'] = self.annotations keypoints['categories'] = self.categories return keypoints def init_dir(base_path): """ 初始化COCO数据集的文件夹结构; coco - annotations #标注文件路径 - train #训练数据集 - val #验证数据集 Args: base_path:数据集放置的根路径 """ if not os.path.exists(os.path.join(base_path, "coco", "images")): os.makedirs(os.path.join(base_path, "coco", "images")) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--class_name", "--n", help="class name", type=str, default='hand') parser.add_argument("--input", "--i", help="json file path (labelme)", type=str, default=r'E:\Release_Oflim_Tof_Currency\SaveImage\08_28_08_38_17\IR') parser.add_argument("--output", "--o", help="output file path (coco format)", type=str, default=r'E:\Release_Oflim_Tof_Currency\SaveImage\08_28_08_38_17') parser.add_argument("--join_num", "--j", help="number of join", type=int, default=21) parser.add_argument("--ratio", "--r", help="train and test split ratio", type=float, default=0.12) args = parser.parse_args() labelme_path = args.input saved_coco_path = args.output init_dir(saved_coco_path) # 初始化COCO数据集的文件夹结构 json_list_path = glob.glob(labelme_path + "/*.json") train_path, val_path = train_test_split(json_list_path, test_size=args.ratio) print('{} for training'.format(len(train_path)), '\n{} for testing'.format(len(val_path))) print('Start transform please wait ...') l2c_train = Labelme2coco_keypoints(args) # 构造数据集生成类 # 生成训练集 train_keypoints = l2c_train.to_coco(train_path) l2c_train.save_coco_json(train_keypoints, os.path.join(saved_coco_path, "coco", "train.json")) # 生成验证集 l2c_val = Labelme2coco_keypoints(args) val_instance = l2c_val.to_coco(val_path) l2c_val.save_coco_json(val_instance, os.path.join(saved_coco_path, "coco", "val.json")) # 拷贝 labelme 的原始图片到训练集和验证集里面 for file in train_path: shutil.copy(file.replace("json", "jpg"), os.path.join(saved_coco_path, "coco", "images")) for file in val_path: shutil.copy(file.replace("json", "jpg"), os.path.join(saved_coco_path, "coco", "images"))
配置文件
- 直接在configs目录下面选取,根据需要选则不同的配置文件,以下是我选择的配置文件
-
直接复制一份出来,然后打开将里面的==dataset_type==修改,修改成在上面注册的数据集的名称
-
data_mode修改为‘topdown’
-
data_root修改成数据集的路径
-
work_dir修改成日志与模型的保存目录
-
修改train_cfg中的max_epochs与val_interval,分别表示迭代轮数与验证间隔轮数
-
根据需要修改train_dataloader里面的batch_size,同时还需要修改里面的ann_file,指向数据集的train.json,val_dataloader同理
-
如果需要冻结训练的,则需要在model中的backbone里面添加,frozen_stages
-
把配置文件里面的val_evaluator修改成如下,使用coco评价指标
val_evaluator = [ dict(type='CocoMetric', ann_file=data_root + '/val.json'), dict(type='PCKAccuracy'), dict(type='AUC'), dict(type='NME', norm_mode='keypoint_distance', keypoint_indices=[0, 1]) ]
训练
-
方法一:直接用下面代码训练
from mmengine.config import Config from mmengine.runner import Runner if __name__ == '__main__': # 该路径为你的配置文件路径 config = Config.fromfile(r'E:\Pose\mmpose-main\data\td-hm_hrnet-w32_8xb64-210e_coco-hand-nofreeze-nopretrain-256x192-256x192.py') runner = Runner.from_cfg(config) runner.train()
-
方法二:直接使用命令行运行, 需要在mmpose目录下进行, --work-dir为模型与日志的输出目录
``
python tools/train.py
xxx/my_code/config/td-hm_mobilenetv2_8xb64-210e_coco-384x288.py
–work-dir xxx/data/train_ret/td-hm_mobilenetv2_8xb64-210e_coco-384x288``
模型测试
-
如果需要测试模型则需要在配置文件里面编写test_dataloader与test_evaluator形式与验证的写法差不多,需要注意的是test_dataloader中的参数不能与验证的一样。在命令行里面输入,–show表示显示测图片,–wait-time后面天显示的持续时间,评估完成后会在模型输出目录下输出一个json文件,里面记录了模型的评估信息,你也可以指定输出的目录
python tools/test.py \ your config \ your checkpoint \ --show \ --wait-time 2
模型导出
-
模型导出需要下载[mmdeploy](GitHub - open-mmlab/mmdeploy: OpenMMLab Model Deployment Framework)源码,源码下载好后需要在该目录下输入以下密令行对整个源码编译
python setup.py build python setup.py develop
-
使用下面的脚本导出onnx模型
from mmdeploy.apis import torch2onnx from mmdeploy.backend.sdk.export_info import export2SDK # 需要在转换pth到onnx时传入一张图片 img = 'C:/Users/86135/Desktop/h.jpg' # 保存结果路径的文件夹 work_dir = 'mmpose20240910/epoch20' # 注意这里尽量使用mmdeploy原始文档中推荐的end2end.onnx名称,后续加载onnx时,避免一些错误出现 save_file = 'det_hand_dynamic.onnx' # 使用mmdeploy源码仓库中的mmpose对应config文件(如果你使用的mmcls,那么就需要到mmcls下面找到合适的deploy config文件) deploy_cfg = 'E:\Pose\mmdeploy-main\configs\mmpose\pose-detection_onnxruntime_static.py' # 加载你训练好的模型的config配置文件(这里以“faster-rcnn_r50_fpn_2x_coco.py”举例) model_cfg = r'E:\Pose\mmpose-main\data\td-hm_hrnet-w32_8xb64-210e_coco-hand_256x192.py' # 加载使用上述模型配置文件得到的训练权重latest.pth model_checkpoint = r'E:\Pose\mmpose-main\data\work_dirs\nopretrain-nofreeze-coco-wholebody-hand-256x256\best_AUC_epoch_20.pth' # 设置device为cpu device = 'cuda:0' # 1. convert model to onnx torch2onnx(img, work_dir, save_file, deploy_cfg, model_cfg, model_checkpoint, device) # 2. extract pipeline info for sdk use (dump-info) export2SDK(deploy_cfg, model_cfg, work_dir, pth=model_checkpoint, device=device)
-
导出模型后会在你的输出目录中会有以下文件deploy.json、detail.json、pipeline.json,前面两个记录了onnx模型的信息,其中pipeline记录了输入图像的大小以及预处理信息