CBGSDataset类-带类别平衡采样的数据集封装器

python 复制代码

class CBGSDataset(object):
    """带类别平衡采样的数据集封装器，实现论文《Class-balanced Grouping and Sampling for Point Cloud 3D Object Detection》
    (https://arxiv.org/abs/1908.09492) 提出的方法。

    通过类别平衡采样策略平衡不同类别场景的数量。

    参数：
        dataset (:obj:`CustomDataset`): 需要进行类别平衡采样的原始数据集。
    """

    def __init__(self, dataset):
        self.dataset = dataset
        self.CLASSES = dataset.CLASSES  # 获取原始数据集的类别列表
        self.cat2id = {name: i for i, name in enumerate(self.CLASSES)}  # 构建类别名称到ID的映射
        self.sample_indices = self._get_sample_indices()  # 计算平衡采样后的索引
        
        # 如果原始数据集有flag属性（如训练/验证标记），则继承该属性
        if hasattr(self.dataset, 'flag'):
            self.flag = np.array(
                [self.dataset.flag[ind] for ind in self.sample_indices],
                dtype=np.uint8)

    def _get_sample_indices(self):
        """通过类别平衡策略生成采样索引列表。

        返回：
            list[int]: 平衡采样后的样本索引列表
        """
        # 初始化字典：记录每个类别对应的所有样本索引
        class_sample_idxs = {cat_id: [] for cat_id in self.cat2id.values()}
        
        # 遍历数据集，统计每个类别的样本索引
        for idx in range(len(self.dataset)):
            sample_cat_ids = self.dataset.get_cat_ids(idx)  # 获取当前样本包含的类别ID
            for cat_id in sample_cat_ids:
                class_sample_idxs[cat_id].append(idx)
        
        # 计算总样本数（考虑一个样本可能属于多个类别）
        duplicated_samples = sum([len(v) for _, v in class_sample_idxs.items()])
        
        # 计算当前每个类别的分布比例
        class_distribution = {
            k: len(v) / duplicated_samples
            for k, v in class_sample_idxs.items()
        }

        sample_indices = []
        frac = 1.0 / len(self.CLASSES)  # 目标分布：均匀分布
        
        # 计算每个类别的采样比率
        ratios = [frac / v for v in class_distribution.values()]
        
        # 按比率进行随机采样
        for cls_inds, ratio in zip(list(class_sample_idxs.values()), ratios):
            sample_indices += np.random.choice(
                cls_inds, 
                int(len(cls_inds) * ratio),  # 按比率调整采样数量
                replace=True  # 允许重复采样（用于过采样）
            ).tolist()
            
        return sample_indices

核心功能说明

1. 问题背景

在3D点云目标检测任务中（如自动驾驶场景），数据通常存在严重的类别不平衡问题（例如"汽车"样本远多于"行人"）。
直接训练会导致模型对高频类别过拟合，低频类别检测效果差。

2. 解决方案

过采样(Oversampling)：对稀有类别（如行人）的样本进行重复采样。
欠采样(Undersampling)：对高频类别（如汽车）的样本进行随机丢弃。
最终使每个类别的样本贡献度相等。

3. 算法关键步骤

统计原始分布
- 遍历数据集，记录每个类别出现的所有样本索引。
计算平衡比率
- 目标分布：若共有N个类别，则每个类别占比应为1/N。
- 对每个类别计算采样比率：ratio = (目标比例) / (当前比例)
执行重采样
- 使用np.random.choice按比率随机选择样本，允许重复（replace=True）。

使用示例

假设原始数据分布：

汽车：1000个样本
行人：100个样本
自行车：50个样本

经过CBGSDataset处理后：

每个类别的目标比例：33.3%（3个类别）
重采样后每个类别约383个样本（通过过采样/欠采样实现）

与FastBEV的关系

在BEV（鸟瞰图）感知任务中，类别平衡能显著提升小物体检测效果（如行人、自行车）。
可配合FastBEV的多相机特征融合模块使用，改善3D检测性能。