基于mmpose训练基础教程

数据集准备

• 数据集下载链接下载

• 解压后的文件目录
  

• json格式，将每一张图片的图片名称，图片大小，标注框的大小以及标注的点，都写在一个字典里
  面，并且添加到一个数组里面，在数组里面每一个字典只能包含一张图片的所有信息
  

编写dataset处理数据

• 需要将此文件保存在*mmpose\datasets\datasets*下的目录中，需要注意的是，所有新增的模块都需要使用注册器（Registry）进行注册，并在对应目录的 __init__.py 中进行 import，以便能够使用配置文件构建其实例，这里我将处理数据集的文件放入到了wholebody目录下，并且在==init.py==文件中注册

• METAINFO用于设置关键点的信息，如颜色，名称，连接，权重等参数

• 如果需要构建自己的数据集，则需要自己编写处理数据的类并且修改METAINFO

# Copyright (c) OpenMMLab. All rights reserved.
import json
import os.path as osp
from typing import Callable, List, Optional, Sequence, Union

import numpy as np
from mmengine.utils import check_file_exist

from mmpose.registry import DATASETS
from mmpose.datasets.datasets.base import BaseCocoStyleDataset


@DATASETS.register_module()
class TinyCocoDataset(BaseCocoStyleDataset):
    METAINFO: dict = dict(from_file='configs/_base_/datasets/coco.py')

    def _load_annotations(self) -> List[dict]:
        """Load data from annotations in MPII format."""

        check_file_exist(self.ann_file)
        with open(self.ann_file) as anno_file:
            anns = json.load(anno_file)

        data_list = []
        ann_id = 0

        for idx, ann in enumerate(anns):
            img_h, img_w = ann['image_size']

            # get bbox in shape [1, 4], formatted as xywh
            x, y, w, h = ann['bbox']
            x1 = np.clip(x, 0, img_w - 1)
            y1 = np.clip(y, 0, img_h - 1)
            x2 = np.clip(x + w, 0, img_w - 1)
            y2 = np.clip(y + h, 0, img_h - 1)

            bbox = np.array([x1, y1, x2, y2], dtype=np.float32).reshape(1, 4)

            # load keypoints in shape [1, K, 2] and keypoints_visible in [1, K]
            joints_3d = np.array(ann['keypoints']).reshape(1, -1, 3)
            num_joints = joints_3d.shape[1]
            keypoints = np.zeros((1, num_joints, 2), dtype=np.float32)
            keypoints[:, :, :2] = joints_3d[:, :, :2]
            keypoints_visible = np.minimum(1, joints_3d[:, :, 2:3])
            keypoints_visible = keypoints_visible.reshape(1, -1)

            data_info = {
                'id': ann_id,
                'img_id': int(ann['image_file'].split('.')[0]),
                'img_path': osp.join(self.data_prefix['img'], ann['image_file']),
                'bbox': bbox,
                'bbox_score': np.ones(1, dtype=np.float32),
                'keypoints': keypoints,
                'keypoints_visible': keypoints_visible,
            }

            data_list.append(data_info)
            ann_id = ann_id + 1

        return data_list, None

config配置文件

• 根据训练的需要选择对应的配置文件，配置文件在configs目录下

• 使用以下脚本生产训练需要的配置文件
  • cfg.data_root 设置数据集目录
  • cfg.work_dir 设置工作目录用于存放训练的日志
  • cfg.dataset_type 设置数据集的类型
  • cfg.train_cfg.val_interval 设置训练的日志记录，每n轮记录一次
  • cfg.train_cfg.max_epochs 设置迭代轮数
  • 可以通过设置 cfg.train_dataloader.num_batch_per_epoch 设置每个epoch的迭代轮数以便快速训练，也可以通过设置 cfg.train_dataloader.indices 设置每个epoch迭代的样本数（建议只在小批量训练时使用）

from mmengine import Config

cfg = Config.fromfile(
    './configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-256x192.py'
)

# 数据集目录
cfg.data_root = 'data/coco_tiny'
# 用于保存日志与模型
cfg.work_dir = 'work_dirs/hrnet_w32_coco_tiny_256x192'
cfg.randomness = dict(seed=0)

# set log interval
cfg.train_cfg.val_interval = 1

# 设置迭代轮次
cfg.train_cfg.max_epochs = 40

# set optimizer
cfg.optim_wrapper = dict(optimizer=dict(
    type='Adam',
    lr=5e-4,
))

# set learning rate policy
cfg.param_scheduler = [
    dict(
        type='LinearLR', begin=0, end=10, start_factor=0.001,
        by_epoch=False),  # warm-up
    dict(
        type='MultiStepLR',
        begin=0,
        end=cfg.train_cfg.max_epochs,
        milestones=[17, 35],
        gamma=0.1,
        by_epoch=True)
]


# set batch size
cfg.train_dataloader.batch_size = 16
cfg.val_dataloader.batch_size = 16
cfg.test_dataloader.batch_size = 16

"""
dataset_type需要使用注册器进行注册才能使用，此处可以根据自己的数据集的类型进行选择
"""
cfg.dataset_type = 'TinyCocoDataset'
cfg.train_dataloader.dataset.type = cfg.dataset_type
cfg.train_dataloader.dataset.ann_file = 'train.json'
cfg.train_dataloader.dataset.data_root = cfg.data_root
cfg.train_dataloader.dataset.data_prefix = dict(img='images/')


cfg.val_dataloader.dataset.type = cfg.dataset_type
cfg.val_dataloader.dataset.bbox_file = None
cfg.val_dataloader.dataset.ann_file = 'val.json'
cfg.val_dataloader.dataset.data_root = cfg.data_root
cfg.val_dataloader.dataset.data_prefix = dict(img='images/')

cfg.test_dataloader.dataset.type = cfg.dataset_type
cfg.test_dataloader.dataset.bbox_file = None
cfg.test_dataloader.dataset.ann_file = 'val.json'
cfg.test_dataloader.dataset.data_root = cfg.data_root
cfg.test_dataloader.dataset.data_prefix = dict(img='images/')

# set evaluator
cfg.val_evaluator = dict(type='PCKAccuracy')
cfg.test_evaluator = cfg.val_evaluator

cfg.default_hooks.checkpoint.save_best = 'PCK'
cfg.default_hooks.checkpoint.max_keep_ckpts = 1

print(cfg.pretty_text)

file_name = 'td-hm_hrnet-w32_8xb64-210e_coco-hand-nofreeze-nopretrain-256x192-256x192.py'
# 保存文件
cfg.dump(file_name)

• 训练
  • 如果需要冻结训练，则需要修改生成的配置文件中的 model字典，在字典中的backbone里设置frozen_stages，这里我冻结三层

from mmengine.config import Config, DictAction
from mmengine.runner import Runner

 # 载入生生的文件
config = Config.fromfile(r'E:\Pose\mmpose-main\data\td-hm_hrnet-w32_8xb64-210e_coco-hand-nofreeze-nopretrain-256x192-256x192.py')

# build the runner from config
runner = Runner.from_cfg(cfg)

# start training
runner.train()
:\Pose\mmpose-main\data\td-hm_hrnet-w32_8xb64-210e_coco-hand-nofreeze-nopretrain-256x192-256x192.py')

# build the runner from config
runner = Runner.from_cfg(cfg)

# start training
runner.train()