从TT100K到YOLOv8:Python自动化交通标志数据集转换实战

在计算机视觉领域,高质量的数据集是模型训练成功的关键前提。TT100K作为国内最具代表性的交通标志数据集之一,包含了超过10万张图像和3万多个标注实例,涵盖了中国道路上常见的200多种交通标志。然而,将TT100K的原始标注格式直接应用于YOLOv8训练并非易事——这需要经历复杂的格式转换过程。本文将深入解析如何通过Python脚本自动化完成这一转换,不仅提供完整可运行的代码,还会剖析每个关键步骤的技术细节。

1. 环境准备与数据集结构解析

在开始转换之前,我们需要先理解TT100K数据集的原始结构。下载后的数据集通常包含以下关键文件和目录:

tt100k_2021/
├── annotations_all.json  # 主标注文件
├── data/                # 原始图像存储目录
│   ├── 1.jpg
│   ├── 2.jpg
│   └── ...
└── statistics.json      # 运行脚本后生成的类别统计文件

关键依赖安装 (推荐使用Python 3.8+环境):

pip install opencv-python tqdm numpy

注意 :原始TT100K标注采用特殊的JSON格式,其中每个标注对象包含以下关键字段:

  • bbox : 采用[xmin, ymin, xmax, ymax]格式的边界框坐标
  • category : 交通标志类别名称
  • path : 对应图像文件的相对路径

2. 核心转换类设计与初始化

我们创建 TT100K2YOLO 类来封装整个转换流程,其核心结构如下:

class TT100K2YOLO:
    def __init__(self, parent_path='./tt100k_2021'):
        self.parent_path = parent_path  # 数据集根目录
        self.classes = []              # 最终保留的类别列表
        self.id_map = {}               # 类别名称到ID的映射
        
        # 创建输出目录结构
        self.output_dirs = ['train', 'val', 'test', 'labels']
        for dir_name in self.output_dirs:
            os.makedirs(os.path.join(parent_path, dir_name), exist_ok=True)

路径配置要点

  • parent_path 应指向本地数据集根目录
  • 脚本会自动创建YOLO格式所需的输出目录
  • 使用 os.path.join 确保跨平台路径兼容性

3. 类别筛选与数据统计分析

TT100K原始数据集包含200+类别,但实际应用中我们通常只需要其中部分常见类别。以下方法实现自动筛选:

def filter_classes(self, min_samples=100):
    with open(os.path.join(self.parent_path, 'annotations_all.json')) as f:
        data = json.load(f)
    
    # 统计每个类别的样本数
    class_stats = defaultdict(int)
    for img_id, img_info in data['imgs'].items():
        for obj in img_info['objects']:
            class_stats[obj['category']] += 1
    
    # 筛选满足最小样本要求的类别
    self.classes = [cls for cls, count in class_stats.items() 
                   if count >= min_samples]
    
    # 生成类别映射文件
    with open(os.path.join(self.parent_path, 'labels/classes.txt'), 'w') as f:
        f.write('\n'.join(self.classes))
    
    self.id_map = {cls: idx for idx, cls in enumerate(self.classes)}
    return self.classes

参数调整建议

  • min_samples 控制保留类别的最小样本量
  • 值设得越高,保留的类别越少但每个类别质量越高
  • 交通标志检测通常设置50-200之间

4. COCO格式中间转换实现

虽然最终目标是YOLO格式,但先转换为COCO格式可以简化过程。关键转换逻辑:

def convert_to_coco(self, phase='train', train_ratio=0.7):
    coco_data = {
        'info': {},
        'licenses': [],
        'categories': [],
        'images': [],
        'annotations': []
    }
    
    # 构建categories
    for cls_id, cls_name in enumerate(self.classes):
        coco_data['categories'].append({
            'id': cls_id,
            'name': cls_name,
            'supercategory': 'traffic_sign'
        })
    
    # 处理图像和标注
    annotation_id = 1
    all_images = [...]  # 获取所有图像列表
    random.shuffle(all_images)
    
    # 按比例划分数据集
    split_idx = int(len(all_images) * train_ratio)
    if phase == 'train':
        selected_images = all_images[:split_idx]
    elif phase == 'val':
        selected_images = all_images[split_idx:]
    
    for img_info in selected_images:
        # 添加图像信息
        coco_data['images'].append({
            'id': img_info['id'],
            'file_name': img_info['path'].split('/')[-1],
            'width': img_info['width'],
            'height': img_info['height']
        })
        
        # 转换标注
        for obj in img_info['objects']:
            if obj['category'] not in self.classes:
                continue
                
            bbox = obj['bbox']
            coco_data['annotations'].append({
                'id': annotation_id,
                'image_id': img_info['id'],
                'category_id': self.id_map[obj['category']],
                'bbox': [bbox['xmin'], bbox['ymin'], 
                        bbox['xmax']-bbox['xmin'], 
                        bbox['ymax']-bbox['ymin']],
                'area': (bbox['xmax']-bbox['xmin'])*(bbox['ymax']-bbox['ymin']),
                'iscrowd': 0
            })
            annotation_id += 1
    
    # 保存COCO格式JSON
    output_path = os.path.join(self.parent_path, f'annotations/{phase}.json')
    with open(output_path, 'w') as f:
        json.dump(coco_data, f, indent=2)

关键设计考虑

  1. 保持标注ID的连续性
  2. 正确处理边界框坐标转换
  3. 支持灵活的数据集划分比例
  4. 自动过滤不包含目标类别的图像

5. COCO到YOLO格式的最终转换

将COCO格式转换为YOLO所需的TXT格式是最后一步,也是最具技术挑战的环节:

def coco_to_yolo(self, phase='train'):
    json_path = os.path.join(self.parent_path, f'annotations/{phase}.json')
    output_dir = os.path.join(self.parent_path, 'labels', phase)
    
    os.makedirs(output_dir, exist_ok=True)
    
    with open(json_path) as f:
        coco_data = json.load(f)
    
    # 为每张图像创建对应的TXT文件
    for img in coco_data['images']:
        img_id = img['id']
        file_name = img['file_name']
        txt_name = os.path.splitext(file_name)[0] + '.txt'
        
        # 收集该图像的所有标注
        anns = [ann for ann in coco_data['annotations'] 
               if ann['image_id'] == img_id]
        
        # 转换为YOLO格式并写入文件
        with open(os.path.join(output_dir, txt_name), 'w') as f_txt:
            for ann in anns:
                # 归一化处理
                x_center = (ann['bbox'][0] + ann['bbox'][2]/2) / img['width']
                y_center = (ann['bbox'][1] + ann['bbox'][3]/2) / img['height']
                width = ann['bbox'][2] / img['width']
                height = ann['bbox'][3] / img['height']
                
                f_txt.write(f"{ann['category_id']} {x_center:.6f} {y_center:.6f} "
                          f"{width:.6f} {height:.6f}\n")
    
    print(f"转换完成!YOLO格式标注已保存到 {output_dir}")

归一化处理要点

  1. 边界框中心坐标和宽高都需要除以图像尺寸
  2. 保留6位小数确保精度
  3. 每个标注行格式: class_id x_center y_center width height
  4. 每个TXT文件对应一张图像的所有标注

6. 数据验证与常见问题排查

转换完成后,必须验证生成的数据是否符合YOLO格式要求。推荐使用以下检查方法:

def validate_yolo_labels(self, phase='train'):
    label_dir = os.path.join(self.parent_path, 'labels', phase)
    image_dir = os.path.join(self.parent_path, 'images', phase)
    
    for label_file in os.listdir(label_dir):
        if not label_file.endswith('.txt'):
            continue
            
        # 检查对应的图像文件是否存在
        img_file = os.path.splitext(label_file)[0] + '.jpg'
        if not os.path.exists(os.path.join(image_dir, img_file)):
            print(f"警告:找不到图像文件 {img_file}")
            continue
            
        # 检查标注内容格式
        with open(os.path.join(label_dir, label_file)) as f:
            lines = f.readlines()
            for line in lines:
                parts = line.strip().split()
                if len(parts) != 5:
                    print(f"格式错误:{label_file} -> {line}")
                    continue
                    
                # 验证数值范围
                try:
                    class_id = int(parts[0])
                    coords = list(map(float, parts[1:]))
                    if not (0 <= coords[0] <= 1 and 0 <= coords[1] <= 1 and
                           0 <= coords[2] <= 1 and 0 <= coords[3] <= 1):
                        print(f"坐标越界:{label_file} -> {line}")
                except ValueError:
                    print(f"数值格式错误:{label_file} -> {line}")

常见问题解决方案

问题现象 可能原因 解决方法
找不到图像文件 路径配置错误 检查 parent_path 和图像目录结构
标注文件为空 类别筛选过严 调整 min_samples 参数
坐标值大于1 归一化计算错误 检查图像尺寸读取是否正确
类别ID越界 类别映射错误 验证 classes.txt 内容

7. 完整流程封装与使用示例

将所有功能封装为完整的可执行流程:

def full_conversion_pipeline(self, min_samples=100, train_ratio=0.7):
    print("步骤1:筛选有效类别...")
    self.filter_classes(min_samples)
    
    print("步骤2:转换为COCO格式...")
    self.convert_to_coco('train', train_ratio)
    self.convert_to_coco('val', train_ratio)
    
    print("步骤3:转换为YOLO格式...")
    self.coco_to_yolo('train')
    self.coco_to_yolo('val')
    
    print("步骤4:验证结果...")
    self.validate_yolo_labels('train')
    self.validate_yolo_labels('val')
    
    print("转换流程完成!")

# 使用示例
if __name__ == '__main__':
    converter = TT100K2YOLO(parent_path='/path/to/your/tt100k_2021')
    converter.full_conversion_pipeline(min_samples=100, train_ratio=0.8)

高级技巧

  1. 添加进度条显示(使用 tqdm
  2. 支持多进程加速处理
  3. 增加日志记录功能
  4. 添加命令行参数支持

8. YOLOv8训练配置适配

转换完成后,还需要准备YOLOv8的训练配置文件。以下是典型的 data.yaml 示例:

# TT100K数据集配置
path: /path/to/your/tt100k_2021
train: images/train
val: images/val
test: images/test

# 类别信息
names:
  0: speed_limit_20
  1: speed_limit_30
  2: speed_limit_50
  # ...其他类别

训练命令示例

yolo detect train data=data.yaml model=yolov8n.pt epochs=100 imgsz=640

通过这套完整的转换流程,我们成功将TT100K数据集转化为YOLOv8可直接训练的格式。在实际项目中,根据具体需求可能还需要调整类别筛选阈值、数据集划分比例等参数。转换过程中最常遇到的路径问题可以通过严格的路径验证来避免,而标注质量问题则需要通过可视化抽查来保证。

更多推荐