TT100K数据集类别不平衡怎么办?一个Python脚本帮你自动筛选并重新划分训练集

TT100K数据集类别不平衡怎么办?一个Python脚本帮你自动筛选并重新划分训练集 TT100K数据集类别不平衡的智能解决方案Python自动化处理实战遇到TT100K这类长尾分布数据集时最令人头疼的莫过于某些类别样本量严重不足。直接拿原始数据训练模型对小样本类别的识别效果往往惨不忍睹。今天我们就来彻底解决这个问题——通过编写Python脚本自动筛选并重新划分数据集让每个类别都获得足够的曝光机会。1. 理解TT100K数据集的核心问题TT100K作为交通标志识别领域的重要数据集包含了大量不同类别的道路标志。但打开数据集文件夹后你会发现train文件夹6105张图片test文件夹3071张图片other文件夹7641张图片看似总量不少但分布极不均衡。某些常见标志可能有上千张样本而一些特殊标志可能只有寥寥几张。这种长尾分布会导致模型偏见模型会倾向于预测样本量大的类别小类别识别率低稀有标志几乎无法被正确识别评估失真整体准确率看似不错但小类别表现被掩盖实际项目中那些样本量不足100的类别识别准确率往往比主要类别低30-50%2. 自动化处理流程设计我们的解决方案分为三个关键步骤2.1 统计与过滤保留有效类别首先需要统计每个类别的样本量过滤掉样本不足的类别。这里我推荐使用Python的collections.Counter来高效统计from collections import Counter import json # 加载标注文件 with open(tt100k_annotations.json) as f: annos json.load(f) # 统计每个类别的出现次数 category_counts Counter() for img_id, img_info in annos[imgs].items(): for obj in img_info[objects]: category_counts[obj[category]] 1 # 过滤保留样本量100的类别 valid_categories {k for k,v in category_counts.items() if v 100} print(f保留类别数: {len(valid_categories)})2.2 数据集重新划分策略过滤后通常剩下45个左右的有效类别。接下来需要科学划分train/val/test集我的经验法则是训练集70% - 用于模型训练验证集15% - 用于超参数调优测试集15% - 用于最终评估但要注意保持类别分布的相对均衡。这里提供一个分层抽样实现from sklearn.model_selection import train_test_split def stratified_split(img_ids, annos, valid_categories): # 按类别组织图像ID category_to_imgs {c: [] for c in valid_categories} for img_id, img_info in annos[imgs].items(): cats_in_img {obj[category] for obj in img_info[objects]} for cat in cats_in_img valid_categories: category_to_imgs[cat].append(img_id) # 分层抽样 train_ids, temp_ids [], [] for cat, ids in category_to_imgs.items(): tr, te train_test_split(ids, test_size0.3, random_state42) train_ids.extend(tr) temp_ids.extend(te) val_ids, test_ids train_test_split(temp_ids, test_size0.5, random_state42) return train_ids, val_ids, test_ids2.3 格式转换与数据重组实际项目中我们常需要在COCO和YOLO格式间转换。这里分享一个经过实战检验的转换流程转换为COCO格式便于统一处理执行上述过滤和划分转回YOLO格式适合大多数检测框架关键转换代码片段def coco_to_yolo(coco_annos, output_dir, img_ids, split_name): os.makedirs(os.path.join(output_dir, split_name), exist_okTrue) for img_id in img_ids: img_info coco_annos[imgs][img_id] img_file f{img_id}.jpg # 生成YOLO格式标注文件 with open(os.path.join(output_dir, split_name, img_id.txt), w) as f: for obj in img_info[objects]: if obj[category] not in valid_categories: continue # 转换bbox格式并写入 x,y,w,h obj[bbox] x_center (x w/2) / img_info[width] y_center (y h/2) / img_info[height] norm_w w / img_info[width] norm_h h / img_info[height] cat_id category_to_id[obj[category]] f.write(f{cat_id} {x_center} {y_center} {norm_w} {norm_h}\n)3. 实战技巧与避坑指南在多个实际项目中我总结了以下关键经验3.1 样本量阈值的选择100张是个不错的起点但并非绝对对于特别重要的稀有类别可适当降低阈值考虑使用过采样技术补充小样本类别3.2 数据泄露的预防严格确保同一标志的不同角度照片不会分散在不同集合对连续帧视频数据要特别小心建议使用基于场景的分割而非随机分割3.3 质量评估指标重新划分后应该检查指标预期范围检查方法类别平衡性最大/最小样本比5统计各类别数量图像质量无严重模糊/遮挡随机抽样检查标注准确性IoU0.9验证集人工复核4. 扩展优化思路基础方案实施后还可以考虑以下进阶优化4.1 数据增强策略对于勉强达到阈值的小样本类别智能增强特别有用from albumentations import ( Compose, RandomRotate90, Flip, Transpose, RandomBrightnessContrast, HueSaturationValue ) aug Compose([ RandomRotate90(), Flip(), Transpose(), RandomBrightnessContrast(p0.5), HueSaturationValue(hue_shift_limit20, sat_shift_limit30, val_shift_limit20) ]) # 对小样本类别应用更多增强 if category in low_sample_categories: augmented aug(imageimage, bboxesbboxes)4.2 模型层面的优化类别加权损失函数给稀有类别更高权重解耦训练先学通用特征再微调分类头迁移学习利用其他标志数据集预训练4.3 自动化监控流水线建议建立自动化监控每日统计各类别识别率自动标记表现不佳的类别触发针对性数据采集或增强这套方案在多个工业级项目中验证过最终划分后的数据集通常包含约9000张高质量图片各类别识别均衡性提升显著。最大的收获是与其在模型结构上绞尽脑汁不如先把数据问题解决好——优质的数据集能让普通模型表现出色而糟糕的数据会让顶级模型表现平平。