Numpy 目标检测数据集 按各类别的边界框数量分割

Posted 2022-09-30 荷碧·TongZJ

在对目标检测数据集进行分割的时候，如果只是简单地将图像数据随机地按比例分割，难免需要面对以下问题：

• 如果数据集中某一类别的边界框均聚集在训练集，神经网络在对该类别的学习可能过拟合，在验证集上的表现可信度较低
• 如果数据集中某一类别的边界框均聚集在验证集，神经网络在训练集上无法充分学习到该类别的特征，同时在验证集上的表现也将较差

如果要将目标检测的数据集按照 4:1 的比例分割，那么最优的解决方案显然不是暴力地按图像数量分割，而是按各类别的边界框数量分割，使得：

• 类别 1 的边界框在训练集中的数量 : 在验证集中的数量 ≈ 4:1
• 类别 2 的边界框在训练集中的数量 : 在验证集中的数量 ≈ 4:1
• ……

所以我提出了这样一个分割目标检测数据集的算法，定义以下变量：

	训练集统计：训练集中各个类别的边界框总数，初始全为 0
	验证集统计：验证集中各个类别的边界框总数，初始全为 0
	图像统计：第 i 张图像各个类别的边界框总数

例如，对于下面这个数据集，第 i 行第 j 列的数据表示了第 i 张图像中第 j 个类别的边界框总数：

	0	1	2
train_0.jpg	1	9	9
train_1.jpg	3	5	8
train_2.jpg	9	3	5
……	…	…	…
train_99.jpg	8	0	6

初始状态下，；给定比例为 4:1，定义某状态的 SSE 损失值为：

打乱数据集的图像顺序后，依次取出数据集中的图像 id、图像统计 ，将该图像加入训练集 / 验证集，则衍生 2 个子状态：

• 加入训练集：训练集统计变为 ，验证集统计仍为 ，对应损失值为 
• 加入验证集：训练集统计仍为 ，验证集统计变为 ，对应损失值为 

执行较小的损失值所对应的动作后，更新 ，最后可使得数据集的分割结果逼近我们想要的“按各类别的边界框数量分割”

加上存储图像 id 的操作，该算法的大致流程如下：

下面这个函数的主要要求 label_count 和 scale 两个参数：

• label_count 是图像边界框的统计信息，为 pandas 的 DataFrame 类型；其形式与上文示例中的表格一致，列索引为类别，行索引为图像 id (建议设为图像的文件名)
• scale 表示某类别边界框分布在训练集中的比例，当希望上文中提到的比例为 n:1 时，则该参数的值为 

这个函数的用法已经编写在代码最下方的 if 语句下，兄弟姐妹们别忘了一键三连喔

import logging

import numpy as np
import pandas as pd
from tqdm import tqdm

logging.basicConfig(format='%(message)s', level=logging.INFO)
LOGGER = logging.getLogger(__name__)


def bbox_split(label_count: pd.DataFrame, scale: float,
               verbose=False, seed=0):
    ''' 对目标检测数据集按各类别的边界框数量分割
        label_count: Dataframe[classes, img_id]
        scale: 某类别边界框分布在训练集中的比例
        return: 训练集id列表, 验证集id列表'''
    exp_radio = scale / (1 - scale)
    # 打乱图像的次序
    idx = list(label_count.index)
    np.random.seed(seed)
    np.random.shuffle(idx)
    label_count = label_count.loc[idx]
    # 记录训练集、验证集当前的边界框数量
    data_count = np.zeros([2, len(label_count.columns)])
    # 防止后续运算中分母出现 0
    data_count[1] += 1e-4
    # 存放训练集、验证集数据的 id
    data_pool = [[] for _ in range(2)]
    pbar = tqdm(label_count.iterrows(), total=len(idx))
    for i, cnt in pbar:
        cnt = np.array(cnt)
        # 计算期望比例 与 执行动作后比例的 SSE 损失
        loss = np.zeros(2)
        for j, next_sit in enumerate([data_count.copy() for _ in range(2)]):
            next_sit[j] += cnt
            loss[j] = np.square(next_sit[0] / next_sit[1] - exp_radio).sum()
        # 根据损失值选择下一步动作
        choose = loss.argmin()
        data_count[choose] += cnt
        data_pool[choose].append(i)
        # 输出当前的训练集各类别比例的上下限
        cur_scale = data_count[0] / data_count.sum(axis=0) - scale
        pbar.set_description(f'BBox Scale Error ∈ [cur_scale.min():.3f, cur_scale.max():.3f]')
    if verbose:
        LOGGER.info(f'Train Set Length: len(data_pool[0])\\t\\tEval Set Length: len(data_pool[1])')
    return data_pool


if __name__ == '__main__':
    # 5000 张图像各个类别边界框数目统计结果
    example = (np.random.random([5000, 5]) * 10).astype(np.int32)
    example = pd.DataFrame(example, index=[f'train_i.jpg' for i in range(example.shape[0])])
    print(example)

    train_id, eval_id = bbox_split(example, 0.8, verbose=True)