MMagic项目数据集准备指南：从基础到自定义-优快云博客

MMagic项目数据集准备指南：从基础到自定义

mmagic OpenMMLab Multimodal Advanced, Generative, and Intelligent Creation Toolbox. Unlock the magic 🪄: Generative-AI (AIGC), easy-to-use APIs, awsome model zoo, diffusion models, for text-to-image generation, image/video restoration/enhancement, etc. 项目地址: https://gitcode.com/gh_mirrors/mm/mmagic

前言

在计算机视觉任务中，高质量的数据集是模型训练的基础。MMagic作为一个强大的多媒体内容生成与编辑工具库，提供了多种数据集处理方式，支持从图像超分辨率到视频插帧等多种任务。本文将详细介绍MMagic支持的数据集格式，并指导如何为特定任务准备自定义数据集。

数据集基础架构

MMagic中的所有数据集都继承自BaseDataset基类，其核心工作流程包含两个关键步骤：

数据列表加载：通过load_data_list方法加载数据信息（如文件路径）
数据预处理：在__getitem__方法中调用prepare_data进行数据预处理

这种设计使得数据集处理流程清晰且易于扩展。

核心数据集类型详解

1. 基础图像数据集(BasicImageDataset)

适用场景：图像超分辨率、修复、无条件图像生成等低层视觉任务

特点：

支持可选标注文件
灵活的文件名模板配置
多种数据组织结构支持

典型配置示例：

dataset = BasicImageDataset(
    ann_file='',  # 可选标注文件
    metainfo=dict(dataset_type='div2k', task_name='sisr'),
    data_root='data/DIV2K',
    data_prefix=dict(gt='DIV2K_train_HR', img='DIV2K_train_LR_bicubic/X4'),
    filename_tmpl=dict(img='{}_x4', gt='{}'),
    pipeline=[]
)

2. 基础帧数据集(BasicFramesDataset)

适用场景：视频超分辨率、视频帧插值等时序相关任务

特点：

支持多帧输入配置
可指定加载特定帧序列
适应不同视频数据集结构

典型配置示例：

dataset = BasicFramesDataset(
    ann_file='tri_trainlist.txt',
    metainfo=dict(dataset_type='vimeo90k', task_name='vfi'),
    data_root='data/vimeo-triplet',
    data_prefix=dict(img='sequences', gt='sequences'),
    pipeline=[],
    depth=2,
    load_frames_list=dict(img=['img1.png', 'img3.png'], gt=['img2.png'])
)

3. 条件生成数据集(BasicConditionalDataset)

适用场景：SAGAN、BigGAN等条件生成模型

支持三种标注格式：

行读取格式：每行包含图像路径和类别索引
字典格式：JSON文件，键为图像路径，值为标签
文件夹结构：按类别组织文件夹，无需标注文件

典型配置示例：

dataset = dict(
    type='BasicConditionalDataset',
    data_root='./data/eg3d/shapenet-car',
    ann_file='annotation.json',
    pipeline=train_pipeline
)

专用数据集类型

1. 图像合成数据集(AdobeComp1kDataset)

特点：

支持在线合成与离线加载两种模式
处理alpha通道、前景和背景图像
适用于图像合成任务

2. 动态缩放数据集(GrowScaleImgDataset)

特点：

支持训练过程中动态切换数据分辨率
用于PGGAN、StyleGANv1等渐进式增长模型
自动匹配当前训练阶段所需分辨率

3. 单图像GAN数据集(SinGANDataset)

特点：

专为SinGAN设计
构建图像金字塔结构
不按索引迭代，返回固定预处理图像

4. 配对图像数据集(PairedImageDataset)

特点：

用于Pix2Pix等需要精确配对的模型
图像文件包含配对图像的拼接
简化配对数据加载流程

5. 非配对图像数据集(UnpairedImageDataset)

特点：

用于CycleGAN等不需要精确配对的模型
训练时随机采样图像对
测试时按顺序配对

自定义数据集开发指南

当现有数据集格式无法满足需求时，可以按照以下步骤创建自定义数据集：

选择合适的基类：根据任务特点继承BasicImageDataset或BasicFramesDataset
实现核心方法：
- load_annotations：加载并解析数据标注
- 根据需要重写__getitem__方法
注册数据集：使用@DATASETS.register_module()装饰器

示例代码框架：

from mmagic.registry import DATASETS

@DATASETS.register_module()
class CustomVideoDataset(BasicFramesDataset):
    """自定义视频数据集说明文档"""
    
    def __init__(self, ann_file, metainfo, data_root, data_prefix, pipeline, test_mode=False):
        super().__init__(ann_file, metainfo, data_root, data_prefix, pipeline, test_mode)
        self.data_infos = self.load_annotations()
    
    def load_annotations(self):
        """实现自定义标注加载逻辑"""
        data_infos = []
        # 自定义解析逻辑
        return data_infos

数据集重复技巧

当数据量不足时，可以使用RepeatDataset包装器扩展数据集：

dataset_A_train = dict(
    type='RepeatDataset',
    times=N,  # 重复次数
    dataset=dict(  # 原始数据集配置
        type='Dataset_A',
        pipeline=train_pipeline
    )
)