mirror of https://github.com/open-mmlab/mmpose
285 lines
10 KiB
Markdown
285 lines
10 KiB
Markdown
# 自定义数据集
|
||
|
||
MMPose 目前已支持了多个任务和相应的数据集。您可以在 [数据集](https://mmpose.readthedocs.io/zh_CN/latest/dataset_zoo.html) 找到它们。请按照相应的指南准备数据。
|
||
|
||
<!-- TOC -->
|
||
|
||
- [自定义数据集-将数据组织为 COCO 格式](#自定义数据集-将数据组织为-coco-格式)
|
||
- [创建自定义数据集的元信息文件](#创建自定义数据集的元信息文件)
|
||
- [创建自定义数据集类](#创建自定义数据集类)
|
||
- [创建自定义配置文件](#创建自定义配置文件)
|
||
- [数据集封装](#数据集封装)
|
||
|
||
<!-- TOC -->
|
||
|
||
## 将数据组织为 COCO 格式
|
||
|
||
最简单的使用自定义数据集的方法是将您的注释格式转换为 COCO 数据集格式。
|
||
|
||
COCO 格式的注释 JSON 文件具有以下必要键：
|
||
|
||
```python
|
||
'images': [
|
||
{
|
||
'file_name': '000000001268.jpg',
|
||
'height': 427,
|
||
'width': 640,
|
||
'id': 1268
|
||
},
|
||
...
|
||
],
|
||
'annotations': [
|
||
{
|
||
'segmentation': [[426.36,
|
||
...
|
||
424.34,
|
||
223.3]],
|
||
'keypoints': [0,0,0,
|
||
0,0,0,
|
||
0,0,0,
|
||
427,220,2,
|
||
443,222,2,
|
||
414,228,2,
|
||
449,232,2,
|
||
408,248,1,
|
||
454,261,2,
|
||
0,0,0,
|
||
0,0,0,
|
||
411,287,2,
|
||
431,287,2,
|
||
0,0,0,
|
||
458,265,2,
|
||
0,0,0,
|
||
466,300,1],
|
||
'num_keypoints': 10,
|
||
'area': 3894.5826,
|
||
'iscrowd': 0,
|
||
'image_id': 1268,
|
||
'bbox': [402.34, 205.02, 65.26, 88.45],
|
||
'category_id': 1,
|
||
'id': 215218
|
||
},
|
||
...
|
||
],
|
||
'categories': [
|
||
{'id': 1, 'name': 'person'},
|
||
]
|
||
```
|
||
|
||
JSON 标注文件中有三个关键词是必需的：
|
||
|
||
- `images`：包含所有图像信息的列表，每个图像都有一个 `file_name`、`height`、`width` 和 `id` 键。
|
||
- `annotations`：包含所有实例标注信息的列表，每个实例都有一个 `segmentation`、`keypoints`、`num_keypoints`、`area`、`iscrowd`、`image_id`、`bbox`、`category_id` 和 `id` 键。
|
||
- `categories`：包含所有类别信息的列表，每个类别都有一个 `id` 和 `name` 键。以人体姿态估计为例，`id` 为 1，`name` 为 `person`。
|
||
|
||
如果您的数据集已经是 COCO 格式的，那么您可以直接使用 `CocoDataset` 类来读取该数据集。
|
||
|
||
## 创建自定义数据集的元信息文件
|
||
|
||
对于一个新的数据集而言，您需要创建一个新的数据集元信息文件。该文件包含了数据集的基本信息，如关键点个数、排列顺序、可视化颜色、骨架连接关系等。元信息文件通常存放在 `config/_base_/datasets/` 目录下，例如：
|
||
|
||
```
|
||
config/_base_/datasets/custom.py
|
||
```
|
||
|
||
元信息文件中需要包含以下信息：
|
||
|
||
- `keypoint_info`：每个关键点的信息：
|
||
1. `name`: 关键点名称，必须是唯一的，例如 `nose`、`left_eye` 等。
|
||
2. `id`: 关键点 ID，必须是唯一的，从 0 开始。
|
||
3. `color`: 关键点可视化时的颜色，以 (\[B, G, R\]) 格式组织起来，用于可视化。
|
||
4. `type`: 关键点类型，可以是 `upper`、`lower` 或 `''`，用于数据增强 [RandomHalfBody](https://github.com/open-mmlab/mmpose/blob/b225a773d168fc2afd48cde5f76c0202d1ba2f52/mmpose/datasets/transforms/common_transforms.py#L263)。
|
||
5. `swap`: 关键点交换关系，用于水平翻转数据增强 [RandomFlip](https://github.com/open-mmlab/mmpose/blob/dev-1.x/mmpose/datasets/transforms/common_transforms.py#L94)。
|
||
- `skeleton_info`：骨架连接关系，用于可视化。
|
||
- `joint_weights`：每个关键点的权重，用于损失函数计算。
|
||
- `sigma`：标准差，用于计算 OKS 分数，详细信息请参考 [keypoints-eval](https://cocodataset.org/#keypoints-eval)。
|
||
|
||
下面是一个简化版本的元信息文件（[完整版](/configs/_base_/datasets/coco.py)）：
|
||
|
||
```python
|
||
dataset_info = dict(
|
||
dataset_name='coco',
|
||
paper_info=dict(
|
||
author='Lin, Tsung-Yi and Maire, Michael and '
|
||
'Belongie, Serge and Hays, James and '
|
||
'Perona, Pietro and Ramanan, Deva and '
|
||
r'Doll{\'a}r, Piotr and Zitnick, C Lawrence',
|
||
title='Microsoft coco: Common objects in context',
|
||
container='European conference on computer vision',
|
||
year='2014',
|
||
homepage='http://cocodataset.org/',
|
||
),
|
||
keypoint_info={
|
||
0:
|
||
dict(name='nose', id=0, color=[51, 153, 255], type='upper', swap=''),
|
||
1:
|
||
dict(
|
||
name='left_eye',
|
||
id=1,
|
||
color=[51, 153, 255],
|
||
type='upper',
|
||
swap='right_eye'),
|
||
...
|
||
16:
|
||
dict(
|
||
name='right_ankle',
|
||
id=16,
|
||
color=[255, 128, 0],
|
||
type='lower',
|
||
swap='left_ankle')
|
||
},
|
||
skeleton_info={
|
||
0:
|
||
dict(link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
|
||
...
|
||
18:
|
||
dict(
|
||
link=('right_ear', 'right_shoulder'), id=18, color=[51, 153, 255])
|
||
},
|
||
joint_weights=[
|
||
1., 1., 1., 1., 1., 1., 1., 1.2, 1.2, 1.5, 1.5, 1., 1., 1.2, 1.2, 1.5,
|
||
1.5
|
||
],
|
||
sigmas=[
|
||
0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062,
|
||
0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089
|
||
])
|
||
```
|
||
|
||
## 创建自定义数据集类
|
||
|
||
如果标注信息不是用 COCO 格式存储的，那么您需要创建一个新的数据集类。数据集类需要继承自 `BaseDataset` 类，并且需要按照以下步骤实现：
|
||
|
||
1. 在 `mmpose/datasets/datasets` 目录下找到该数据集符合的 package，如果没有符合的，则创建一个新的 package。
|
||
|
||
2. 在该 package 下创建一个新的数据集类，在对应的注册器中进行注册：
|
||
|
||
```python
|
||
from mmengine.dataset import BaseDataset
|
||
from mmpose.registry import DATASETS
|
||
|
||
@DATASETS.register_module(name='MyCustomDataset')
|
||
class MyCustomDataset(BaseDataset):
|
||
```
|
||
|
||
如果未注册，你会在运行时遇到 `KeyError: 'XXXXX is not in the dataset registry'`。
|
||
关于 `mmengine.BaseDataset` 的更多信息，请参考 [这个文档](https://mmengine.readthedocs.io/en/latest/advanced_tutorials/basedataset.html)。
|
||
|
||
3. 确保你在 package 的 `__init__.py` 中导入了该数据集类。
|
||
|
||
4. 确保你在 `mmpose/datasets/__init__.py` 中导入了该 package。
|
||
|
||
## 创建自定义配置文件
|
||
|
||
在配置文件中，你需要修改跟数据集有关的部分，例如：
|
||
|
||
```python
|
||
...
|
||
# 自定义数据集
|
||
dataset_type = 'MyCustomDataset' # or 'CocoDataset'
|
||
|
||
train_dataloader = dict(
|
||
batch_size=2,
|
||
dataset=dict(
|
||
type=dataset_type,
|
||
data_root='aaa',
|
||
# 标注文件路径为 {data_root}/{ann_file}
|
||
# 例如： aaa/annotations/xxx.json
|
||
ann_file='annotations/xxx.json',
|
||
# 图片路径为 {data_root}/{img}/
|
||
# 例如： aaa/train/c.jpg
|
||
data_prefix=dict(img='train'),
|
||
metainfo=dict(from_file='configs/_base_/datasets/custom.py'),
|
||
...),
|
||
)
|
||
|
||
val_dataloader = dict(
|
||
batch_size=2,
|
||
dataset=dict(
|
||
type=dataset_type,
|
||
data_root='aaa',
|
||
# 标注文件路径为 {data_root}/{ann_file}
|
||
# 例如： aaa/annotations/yyy.json
|
||
ann_file='annotations/yyy.json',
|
||
# 图片路径为 {data_root}/{img}/
|
||
# 例如： aaa/val/c.jpg
|
||
data_prefix=dict(img='val'),
|
||
metainfo=dict(from_file='configs/_base_/datasets/custom.py'),
|
||
...),
|
||
)
|
||
|
||
test_dataloader = dict(
|
||
batch_size=2,
|
||
dataset=dict(
|
||
type=dataset_type,
|
||
data_root='aaa',
|
||
# 标注文件路径为 {data_root}/{ann_file}
|
||
# 例如： aaa/annotations/zzz.json
|
||
ann_file='annotations/zzz.json',
|
||
# 图片路径为 {data_root}/{img}/
|
||
# 例如： aaa/test/c.jpg
|
||
data_prefix=dict(img='test'),
|
||
metainfo=dict(from_file='configs/_base_/datasets/custom.py'),
|
||
...),
|
||
)
|
||
...
|
||
```
|
||
|
||
请确保所有的路径都是正确的。
|
||
|
||
## 数据集封装
|
||
|
||
在 MMPose 中，支持使用 MMPose 实现的数据集封装和 [MMEngine](https://github.com/open-mmlab/mmengine) 实现的数据集封装。目前 [MMEngine](https://github.com/open-mmlab/mmengine) 支持以下数据集封装：
|
||
|
||
- [ConcatDataset](https://mmengine.readthedocs.io/zh_CN/latest/advanced_tutorials/basedataset.html#concatdataset)
|
||
- [RepeatDataset](https://mmengine.readthedocs.io/zh_CN/latest/advanced_tutorials/basedataset.html#repeatdataset)
|
||
|
||
### CombinedDataset
|
||
|
||
MMPose 提供了一个 [CombinedDataset](https://github.com/open-mmlab/mmpose/blob/dev-1.x/mmpose/datasets/dataset_wrappers.py#L15) 类，它可以将多个数据集封装成一个数据集。它的使用方法如下：
|
||
|
||
```python
|
||
dataset_1 = dict(
|
||
type='dataset_type_1',
|
||
data_root='aaa',
|
||
# 图片路径为 {data_root}/{img_path}/
|
||
# 例如： aaa/train/c.jpg
|
||
data_prefix=dict(img_path='train'),
|
||
# 标注文件路径为 {data_root}/{ann_file}
|
||
# 例如： aaa/annotations/train.json
|
||
ann_file='annotations/train.json',
|
||
pipeline=[
|
||
# 使用转换器将标注信息统一为需要的格式
|
||
converter_transform_1
|
||
])
|
||
|
||
dataset_2 = dict(
|
||
type='dataset_type_2',
|
||
data_root='bbb',
|
||
# 图片路径为 {data_root}/{img_path}/
|
||
# 例如： bbb/train/c.jpg
|
||
data_prefix=dict(img_path='train'),
|
||
# 标注文件路径为 {data_root}/{ann_file}
|
||
# 例如： bbb/annotations/train.json
|
||
ann_file='annotations/train.json',
|
||
pipeline=[
|
||
converter_transform_2
|
||
])
|
||
|
||
shared_pipeline = [
|
||
LoadImage(),
|
||
ParseImage(),
|
||
]
|
||
|
||
combined_dataset = dict(
|
||
type='CombinedDataset',
|
||
metainfo=dict(from_file='path/to/your/metainfo'),
|
||
datasets=[dataset_1, dataset_2],
|
||
pipeline=shared_pipeline,
|
||
)
|
||
```
|
||
|
||
- **合并数据集的元信息** 决定了标注格式，可以是子数据集的元信息，也可以是自定义的元信息。如果要自定义元信息，可以参考 [创建自定义数据集的元信息文件](#创建自定义数据集的元信息文件)。
|
||
- **KeypointConverter** 用于将不同的标注格式转换成统一的格式。比如将关键点个数不同、关键点排列顺序不同的数据集进行合并。
|
||
- 更详细的说明请前往[混合数据集训练](../user_guides/mixed_datasets.md)。
|