MMsegmentation项目自定义数据集开发指南

MMsegmentation项目自定义数据集开发指南

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation

前言

在语义分割任务中，使用自定义数据集是常见需求。本文将详细介绍如何在MMsegmentation框架中添加和管理自定义数据集，包括数据集类实现、数据组织方式以及高级数据集混合技术。

一、实现自定义数据集类

1. 创建数据集类

在MMsegmentation中，自定义数据集需要继承BaseSegDataset基类，并通过装饰器注册到数据注册表中：

from mmseg.registry import DATASETS
from .basesegdataset import BaseSegDataset

@DATASETS.register_module()
class ExampleDataset(BaseSegDataset):
    # 定义类别信息和颜色调色板
    METAINFO = dict(
        classes=('road', 'sidewalk', 'building'),
        palette=[[128,64,128], [244,35,232], [70,70,70]])
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

关键点说明：

• METAINFO必须包含classes和palette两个字段
• classes定义所有语义类别名称
• palette为每个类别指定可视化颜色(RGB格式)

2. 注册数据集类

创建完数据集类后，需要在__init__.py中导入：

from .example import ExampleDataset

3. 创建配置文件

在配置文件中指定数据集类型和路径：

dataset_type = 'ExampleDataset'
data_root = 'data/custom_dataset/'

二、数据组织规范

1. 标准目录结构

推荐按以下结构组织数据：

data/custom_dataset/
├── img_dir
│   ├── train
│   │   ├── 0001.jpg
│   │   ├── 0002.jpg
│   ├── val
├── ann_dir
│   ├── train
│   │   ├── 0001.png
│   │   ├── 0002.png
│   ├── val

注意：

• 图像和标注文件应具有相同文件名前缀
• 标注文件应为单通道PNG图像
• 像素值范围[0, num_classes-1]，表示类别索引

2. 标注文件处理技巧

对于小型数据集，可以使用Pillow库的调色板模式(P模式)存储标注：

from PIL import Image
import numpy as np

# 将类别索引转换为调色板图像
palette = np.array([[128,64,128], [244,35,232], [70,70,70]])
label = np.array([[0,1], [2,0]])  # 示例标签
img = Image.fromarray(label.astype(np.uint8), mode='P')
img.putpalette(palette.flatten())
img.save('label.png')

三、高级数据集操作

1. 数据集重复

当数据量不足时，可使用RepeatDataset进行数据重复：

train_dataset = dict(
    type='RepeatDataset',
    times=10,  # 重复10次
    dataset=dict(
        type='CustomDataset',
        pipeline=train_pipeline
    )
)

2. 数据集拼接

合并多个数据集进行联合训练：

concatenate_dataset = dict(
    type='ConcatDataset',
    datasets=[dataset_A_cfg, dataset_B_cfg]
)

3. 多图混合增强

实现Mosaic等增强策略：

train_dataset = dict(
    type='MultiImageMixDataset',
    dataset=dict(
        type='CustomDataset',
        pipeline=[load_pipeline]
    ),
    pipeline=[
        dict(type='RandomMosaic', prob=1),
        dict(type='Resize', img_scale=(1024, 512)),
        dict(type='PackSegInputs')
    ]
)

四、最佳实践建议

1. 数据预处理：对于非常规格式数据，建议先转换为标准格式再使用
2. 类别平衡：在自定义数据集中注意各类别样本数量的平衡
3. 验证集划分：如果没有官方测试集标注，可将验证集作为测试集
4. 性能优化：大数据集建议使用DefaultSampler提高加载效率

通过以上方法，您可以灵活地将各种格式的语义分割数据集成到MMsegmentation框架中，充分利用其强大的训练和评估功能。

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation