如何高效使用BDD100K子集进行自动驾驶目标检测模型训练

1次阅读

没有评论

共计 1919 个字符，预计需要花费 5 分钟才能阅读完成。

BDD100K 作为自动驾驶领域主流数据集，包含 10 万帧高清图像，但全量训练面临三大挑战：

存储压力：解压后约 1.2TB，普通 GPU 服务器难以承载
训练效率：完整训练 YOLOv5 需 200+ 小时，迭代周期长
数据冗余：连续视频帧相似度高，存在大量无效学习

通过分析 JSON 标注中的 weather 和scene字段实现场景分层抽样：

时间维度 ：保持日间(70%)/ 黄昏(20%)/ 夜间(10%) 比例
场景维度：城市道路(60%)/ 高速公路(30%)/ 居民区(10%)
异常事件：保留所有含事故、急刹车的关键帧

统计 category 字段分布后针对性补足：

基础类别：car(35%)/person(25%)/traffic sign(20%)
长尾类别：motorcycle(8%)/traffic light(7%)/bicycle(5%)
剔除规则：删除出现次数 <50 的罕见类别

import json
from collections import defaultdict
import random

# 加载原始标注
with open('bdd100k_labels.json') as f:
    all_labels = json.load(f)

# 构建场景统计字典
scene_count = defaultdict(int)
for label in all_labels:
    scene_count[label['attributes']['weather']] += 1

# 分层抽样函数
def stratified_sample(labels, target_size=5000):
    samples = []
    weather_ratios = {'rainy': 0.1, 'snowy': 0.05, 'clear': 0.85}

    for weather, ratio in weather_ratios.items():
        subset = [x for x in labels if x['attributes']['weather'] == weather]
        samples.extend(random.sample(subset, int(target_size*ratio)))

    return samples[:target_size]

# 执行抽样
subset_labels = stratified_sample(all_labels)

关键处理嵌套 JSON 结构：

def bdd2coco(bdd_labels):
    coco = {"images": [],
        "annotations": [],
        "categories": [{"id": 1, "name": "person"},
            # ... 其他类别定义
        ]
    }

    for img_id, label in enumerate(bdd_labels):
        # 添加图像信息
        coco["images"].append({
            "id": img_id,
            "file_name": label['name'],
            "width": 1280,
            "height": 720
        })

        # 处理标注对象
        for obj in label['labels']:
            if 'box2d' not in obj: continue

            x1, y1 = obj['box2d']['x1'], obj['box2d']['y1']
            x2, y2 = obj['box2d']['x2'], obj['box2d']['y2']

            coco["annotations"].append({"id": len(coco["annotations"]),
                "image_id": img_id,
                "category_id": cat_mapping[obj['category']],
                "bbox": [x1, y1, x2-x1, y2-y1],
                "area": (x2-x1)*(y2-y1),
                "iscrowd": 0
            })

    return coco

在 RTX 3090 上的对比测试：