数据集准备
-
数据集下载链接下载
-
解压后的文件目录
- 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()