NotACracker/COTR项目教程：自定义数据集完全指南

2025-07-04 02:37:53作者：管翌锬

前言

在3D目标检测领域，数据集的准备和处理是模型训练的关键环节。NotACracker/COTR项目提供了灵活的数据集自定义方案，本文将详细介绍如何根据项目需求自定义数据集。

数据集格式支持方案

方案一：转换为现有格式

对于无法直接在线读取的数据，推荐转换为KITTI格式：

数据转换器开发：编写脚本将原始数据重组并转换为KITTI格式
数据集类继承：根据数据特性继承现有数据集类进行微调
配置文件修改：调整数据路径和类别配置

优势：实现简单，可直接复用现有数据处理流程

方案二：转换为中间格式

项目支持将数据转换为pickle中间格式，每个帧数据包含以下关键信息：

image：图像元数据（路径、尺寸等）
point_cloud：点云数据信息
calib：校准参数
annos：标注信息（边界框、类别等）

这种格式的优势在于：

不强制要求原始数据组织结构
灵活性高，可适应不同数据源
便于统一处理流程

自定义数据集实现示例

数据集类实现

以下是一个完整的自定义数据集实现示例：

@DATASETS.register_module()
class MyDataset(Custom3DDataset):
    # 定义数据集类别
    CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
               'bookshelf', 'picture', 'counter', 'desk', 'curtain',
               'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
               'garbagebin')

    def __init__(self, data_root, ann_file, pipeline=None, classes=None,
                 modality=None, box_type_3d='Depth', filter_empty_gt=True,
                 test_mode=False):
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
            test_mode=test_mode)

    def get_ann_info(self, index):
        info = self.data_infos[index]
        if info['annos']['gt_num'] != 0:
            gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(np.float32)
            gt_labels_3d = info['annos']['class'].astype(np.int64)
        else:
            gt_bboxes_3d = np.zeros((0, 6), dtype=np.float32)
            gt_labels_3d = np.zeros((0,), dtype=np.int64)

        gt_bboxes_3d = DepthInstance3DBoxes(
            gt_bboxes_3d,
            box_dim=gt_bboxes_3d.shape[-1],
            with_yaw=False,
            origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)

        pts_instance_mask_path = osp.join(self.data_root, info['pts_instance_mask_path'])
        pts_semantic_mask_path = osp.join(self.data_root, info['pts_semantic_mask_path'])

        return dict(
            gt_bboxes_3d=gt_bboxes_3d,
            gt_labels_3d=gt_labels_3d,
            pts_instance_mask_path=pts_instance_mask_path,
            pts_semantic_mask_path=pts_semantic_mask_path)

配置文件调整

在配置文件中使用自定义数据集：

dataset_A_train = dict(
    type='MyDataset',
    ann_file='annotation.pkl',
    pipeline=train_pipeline
)

数据集包装器应用

项目提供了三种数据集包装器来增强数据使用：

1. 重复数据集(RepeatDataset)

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

2. 类别平衡数据集(ClassBalancedDataset)

dataset_A_train = dict(
    type='ClassBalancedDataset',
    oversample_thr=1e-3,  # 过采样阈值
    dataset=dict(
        type='Dataset_A',
        pipeline=train_pipeline
    )
)

3. 数据集拼接(ConcatDataset)

三种拼接方式：

方式一：相同类型数据集拼接

dataset_A_train = dict(
    type='Dataset_A',
    ann_file=['anno_file_1', 'anno_file_2'],
    separate_eval=False,  # 是否分开评估
    pipeline=train_pipeline
)

方式二：不同类型数据集拼接

data = dict(
    train=[dataset_A_train, dataset_B_train],
    val=dataset_A_val,
    test=dataset_A_test
)

方式三：显式使用ConcatDataset

data = dict(
    val=dict(
        type='ConcatDataset',
        datasets=[dataset_A_val, dataset_B_val],
        separate_eval=False
    )
)

数据集类别修改

可以灵活调整数据集的类别：

直接指定类别列表

classes = ('person', 'bicycle', 'car')
data = dict(
    train=dict(classes=classes),
    val=dict(classes=classes),
    test=dict(classes=classes))

从文件读取类别

classes = 'path/to/classes.txt'  # 每行一个类别名
data = dict(
    train=dict(classes=classes),
    val=dict(classes=classes),
    test=dict(classes=classes))