mmpose关键点训练通用流程

吗喽不爱写代码

已于 2024-09-11 14:55:50 修改

阅读量317

点赞数 10

文章标签： python 深度学习人工智能视觉检测

于 2024-09-11 14:41:26 首次发布

本文链接：https://blog.csdn.net/m0_70573069/article/details/142139001

版权

数据集

数据集格式

在这里插入图片描述

如果是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记录了输入图像的大小以及预处理信息

在这里插入图片描述

吗喽不爱写代码

关注

10
点赞
踩
10

收藏

觉得还不错? 一键收藏
0
评论
mmpose关键点训练通用流程

后面天显示的持续时间，评估完成后会在模型输出目录下输出一个json文件，里面记录了模型的评估信息，你也可以指定输出的目录。==选则元信息文件，元信息文件记录了关键点的数量，连接方式颜色以及权重等重要参数。读取文件，该文件用于训练时数据的载入，如果对数据读取无特殊要求则按以下方法编写。、pipeline.json，前面两个记录了onnx模型的信息，其中。)源码，源码下载好后需要在该目录下输入以下密令行对整个源码编译。==修改，修改成在上面注册的数据集的名称。形式与验证的写法差不多，需要注意的是。
复制链接

扫一扫