CCPD数据集高效转换为YOLO格式：自动化脚本与最佳实践

1次阅读

没有评论

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

CCPD 数据集是车牌检测领域的常用数据集，但其标注格式为 VOC 风格的 XML 文件，而 YOLO 系列目标检测框架要求特定的 TXT 标注格式。这种格式差异导致直接使用 CCPD 数据集进行 YOLO 模型训练存在障碍。

主要差异点包括：

坐标表示方式：CCPD 使用绝对像素坐标，YOLO 要求归一化的相对坐标
文件结构：CCPD 每个 XML 对应一张图片，YOLO 需要将所有标注合并到单个 TXT
类别编码：YOLO 需要用数字索引替代字符串类别名

我们设计了一个 Python 自动化脚本，主要处理流程如下：

遍历指定目录下的所有 XML 标注文件
解析每个 XML 文件中的车牌位置信息
将绝对坐标转换为 YOLO 格式的归一化坐标
生成对应名称的 TXT 标注文件
提供验证机制确保转换准确性

YOLO 格式要求将边界框坐标归一化为相对图像宽高的比例值，转换公式为：

x_center = (xmin + xmax) / (2 * image_width)
y_center = (ymin + ymax) / (2 * image_height)
width = (xmax - xmin) / image_width
height = (ymax - ymin) / image_height

以下是完整的转换脚本，包含详细注释和类型注解：

import argparse
import xml.etree.ElementTree as ET
from pathlib import Path
from tqdm import tqdm
import cv2
import logging

# 配置日志记录
logging.basicConfig(filename='conversion.log', level=logging.INFO)

def parse_args() -> argparse.Namespace:
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='CCPD to YOLO format converter')
    parser.add_argument('--xml-dir', type=str, required=True, help='Directory containing CCPD XML files')
    parser.add_argument('--output-dir', type=str, required=True, help='Output directory for YOLO txt files')
    parser.add_argument('--img-dir', type=str, required=True, help='Directory containing corresponding images')
    parser.add_argument('--class-id', type=int, default=0, help='Class ID for license plates')
    return parser.parse_args()

def convert_bbox(size: tuple[int, int], box: tuple[float, float, float, float]) -> tuple[float, float, float, float]:
    """将 VOC 格式边界框转换为 YOLO 格式"""
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def process_xml(xml_file: Path, output_dir: Path, img_dir: Path, class_id: int) -> bool:
    """处理单个 XML 文件"""
    try:
        tree = ET.parse(xml_file)
        root = tree.getroot()

        # 获取图片尺寸
        size = root.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)

        # 获取边界框信息
        obj = root.find('object')
        bndbox = obj.find('bndbox')
        xmin = float(bndbox.find('xmin').text)
        xmax = float(bndbox.find('xmax').text)
        ymin = float(bndbox.find('ymin').text)
        ymax = float(bndbox.find('ymax').text)

        # 坐标转换
        yolo_box = convert_bbox((width, height), (xmin, xmax, ymin, ymax))

        # 写入 TXT 文件
        txt_path = output_dir / f"{xml_file.stem}.txt"
        with open(txt_path, 'w') as f:
            f.write(f"{class_id} {yolo_box[0]} {yolo_box[1]} {yolo_box[2]} {yolo_box[3]}")

        return True
    except Exception as e:
        logging.error(f"Error processing {xml_file}: {str(e)}")
        return False

def main():
    args = parse_args()
    xml_dir = Path(args.xml_dir)
    output_dir = Path(args.output_dir)
    img_dir = Path(args.img_dir)

    # 创建输出目录
    output_dir.mkdir(parents=True, exist_ok=True)

    # 获取所有 XML 文件
    xml_files = list(xml_dir.glob('*.xml'))

    # 处理每个文件
    success_count = 0
    for xml_file in tqdm(xml_files, desc="Processing files"):
        if process_xml(xml_file, output_dir, img_dir, args.class_id):
            success_count += 1

    print(f"Successfully processed {success_count}/{len(xml_files)} files")

if __name__ == '__main__':
    main()

为确保转换准确性，我们可以使用 OpenCV 可视化标注结果：

def visualize_yolo_annotation(img_path: Path, txt_path: Path):
    """可视化 YOLO 标注结果"""
    img = cv2.imread(str(img_path))
    h, w = img.shape[:2]

    with open(txt_path) as f:
        line = f.readline().strip()
        class_id, x, y, bw, bh = map(float, line.split())

        # 转换回像素坐标
        x1 = int((x - bw/2) * w)
        y1 = int((y - bh/2) * h)
        x2 = int((x + bw/2) * w)
        y2 = int((y + bh/2) * h)

        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.imshow('Validation', img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()