前言

这个实战项目是在学习了Rust的基础知识后,做的一个入门实战的小工具,用来将CCPD数据集转化为yolo格式的数据。这个项目是对Rust基础入门知识的一个总结,应用到了变量的绑定、解构、符合类型、模式匹配、方法、集合类型、错误处理、模块等知识都有具体应用,还用到了log4rs、image等依赖组件。

一、开发思路

1、开发目标

开发一个命令行工具img2yolo,需要用户从命令行提供输入参数,程序内提取参数来处理数据,将CCPD 车牌数据集(原始命名标签格式)转换为 YOLO 格式标签文件。

  • 输入
    指定原始CCPD数据集的目录、数据类别、目标分类编号(Class ID)、转换后输出的目录
  • 输出
    转换后的图像文件和YOLO标签文件
2、CCPD 数据格式说明

CCPD 图像文件名中包含所有标注信息,例如:

02-90_103-258&516_428&559-428&559_258&516_270&480_420&523-0_0_6_30_30_27_31-85-21.jpg

文件名的解析结构说明:

[亮度]-[俯仰_滚转]-[bbox左上x&y_右下x&y]-[四个顶点坐标]-[车牌字符编码]-[省份编号]-[序号].jpg

程序需要的部分是:

258&516_428&559     →    bbox 左上角(258,516),右下角(428,559)

YOLO 的格式要求:

  分类编码    识别目标的中心点坐标     目标:宽   目标:高
<class_id>  <x_center> <y_center> <width> <height>
3、功能需求与逻辑
  • 根据输入参数,读取指定的CCPD数据集目录中的所有图片文件
  • 根据图片名称进行解析,获取bbox坐标
  • 根据图片的尺寸计算YOLO格式的相对坐标,即:目标所在图像的中心点坐标及其宽和高
  • 分析完成后,复制图片到新目录,并生成YOLO 标签文件(.txt,与图像同名)
  • 分析和转换过程中,记录执行日志及错误异常日志
4、模块设计
模块 功能描述
main.rs 程序主入口
args.rs 解析命令行参数
parser.rs 负责解析 CCPD 文件名中的 bbox
converter.rs 坐标转换、归一化为 YOLO 格式
writer.rs 生成 YOLO 标签文件并输出
utils.rs 读取图片尺寸、路径处理等

二、功能实现

1、项目结构
img2yolo/
├── Cargo.toml
├── log4rs.yaml
└── src/
    ├── main.rs
    ├── args.rs         
    ├── parser.rs       
    ├── converter.rs    
    ├── writer.rs       
    └── utils.rs        
2、Cargo.toml配置
[package]
name = "img2yolo"
version = "0.1.0"
edition = "2024"

[dependencies]
walkdir = "2.5.0"
image = "0.25.8"
rand = "0.10.0-rc.0"
log = "0.4.28"
log4rs = "1.4.0"

2、主程序入口 main.rs

主程序引入了其他的模块文件,并定义了程序的数据结构Bbox、YoloData、DataType
在进入主程序时,先初始化了日志目录和配置,并提取参数进行解析,根据参数获取图像文件列表,进行循环解析并创建数据集


fn main() -> Result<(), Box<dyn std::error::Error>> {

    fs::create_dir_all("logs")?;

    // 初始化 log4rs 配置, 使用log4rs方式
    log4rs::init_file("log4rs.yaml", Default::default())?;

    // 提取参数
    let mut args: Vec<String> = env::args().collect();

    //不传参数,则使用默认参数
    if args.len() < 2 {
        args = vec![
            "img2yolo".to_string(),
            "-t".parse().unwrap(),
            "train".to_string(),
            "-s".parse().unwrap(),
            "./assets/images".to_string(), // 图片目录
            "-d".parse().unwrap(),
            "./dataset".to_string(), //目标目录
        ];
    }

    // 获取解析命令参数
    let config = parse(args);

    log::info!("Starting img2yolo...");

    // 获取CCPD图片文件集合
    let mut image_files = image_files(config.src.as_str());
    if image_files.is_empty() {
        error!("No image files found in {}", config.src.as_str());
        std::process::exit(1);
    }

    // 数据集抽取比例, 0~1,可以根据比例的多少来抽取
    if config.ratio >= 0_f64 && config.ratio < 1_f64 {
        // 随机抽取数据
        let mut rng = rng();
        let sample_size = (image_files.len() as f64 * config.ratio).ceil() as usize;
        image_files.shuffle(&mut rng);
        image_files.truncate(sample_size);
    }
    log::info!("{} files found in {}", image_files.len(), config.src.as_str());

    // 创建数据集目录
    let data_type_folder = config.data_type.as_str();
    let train_path_images = PathBuf::from(config.dst.as_str())
        .join(data_type_folder).join( "images");
    let train_path_labels = PathBuf::from(config.dst.as_str())
        .join(data_type_folder).join( "labels");

    // 创建数据集
    let result = create_yolo_dataset(image_files, config.class_id, train_path_images, train_path_labels);
    if result.is_err() {
        error!("Error: {}", result.err().unwrap());
        std::process::exit(1);
    } else {
        log::info!("{} dataset created successfully", data_type_folder);
    }

    Ok(())
}
3、创建yolo数据集 create_yolo_dataset

create_yolo_dataset方法,根据图片集合、输出的图片和标签路径等参数,将输入的图片,进行解析parse_bbox_from_filename、转换bbox2yolo后,按照新的文件名,生成新的图片和标签文件

pub fn create_yolo_dataset(image_files: Vec<(String,String)>, class_id: i32, train_path_images: PathBuf, train_path_labels: PathBuf) -> Result<(), Box<dyn std::error::Error>>{
    create_dir_all(&train_path_images)?;
    create_dir_all(&train_path_labels)?;
    // 创建训练集的文件编号,默认从1开始,如果文件已经存在,则从文件编号最大值开始
    let mut i = get_max_num(&train_path_images);
    for image_file in image_files {
        // 创建新的文件名, 格式为 [类别编号]_[文件编号].jpg
        let new_filename = format!("{}_{:06}", class_id, i);

        let bbox = match parse_bbox_from_filename(&image_file.1) {
            Ok(bbox) => bbox,
            Err(e) => {
                error!("{}: Failed to parse bbox from filename '{}' (path: {}): {}",
                       new_filename, image_file.1, image_file.0, e);
                continue;
            }
        };
        // 获取图像尺寸
        let image_size = match get_image_size(&image_file.0) {
            Ok(size) => size,
            Err(e) => {
                error!("{}: Failed to get image size for file '{}' (path: {}): {}",
                       new_filename, image_file.0, image_file.0, e);
                continue;
            }
        };
        // 转换成YOLO数据
        let yolo_data = match bbox2yolo(bbox, image_size, &image_file.0) {
            Ok(data) => data,
            Err(e) => {
                error!("{}: Failed to parse YOLO data : {}", new_filename, e);
                continue;
            }
        };

        // 获取原文件的扩展名
        let src_path = Path::new(&image_file.0);
        let extension = src_path.extension().and_then(|s| s.to_str()).unwrap_or("jpg");

        // 正确构建目标文件路径
        let dst_image_path = train_path_images.join(format!("{}.{}", new_filename, extension));
        let dst_label_path = train_path_labels.join(format!("{}.txt", new_filename));

        // 复制图像文件
        if let Err(e) = copy_file(&image_file.0, dst_image_path.to_str().unwrap()) {
            error!("{}: Failed to copy image file '{}' to '{}': {}",
                   new_filename, image_file.0, dst_image_path.display(), e);
            continue;
        }
        // 保存标签文件
        if let Err(e) = write_yolo_data(yolo_data, dst_label_path.to_str().unwrap(), class_id) {
            error!("{}: Failed to save YOLO data for file '{}' to '{}': {}",
                   new_filename, image_file.1, dst_label_path.display(), e);
            continue;
        }

        info!("Processed {} ==> {}", new_filename, image_file.1);
        i += 1;
    }
    Ok(())
}
4、CCPD文件名解析 parse.rs

CCPD的文件名有其相应的结构定义规则,按照规则解析出车牌的矩形坐标位置

use crate::Bbox;

///图片命名:“025-95_113-154&383_386&473-386&473_177&454_154&383_363&402-0_0_22_27_27_33_16-37-15.jpg”
/// 1. 025:车牌区域占整个画面的比例;
/// 2. 95_113: 车牌水平和垂直角度, 水平95°, 竖直113°
/// 3. 154&383_386&473:标注框左上、右下坐标,左上(154, 383), 右下(386, 473)
/// 4. 386&473_177&454_154&383_363&402:标注框四个角点坐标,顺序为右下、左下、左上、右上
/// 5. 0_0_22_27_27_33_16:
///     车牌号码映射关系如下:
///     第一个0为省份 对应省份字典provinces中的’皖’,;
///     第二个0是该车所在地的地市一级代码,对应地市一级代码字典alphabets的’A’;
///     后5位为字母和文字, 查看车牌号ads字典,如22为Y,27为3,33为9,16为S,最终车牌号码为皖AY339S
///
pub fn parse_bbox_from_filename(filename: &str) -> Result<Bbox, Box<dyn std::error::Error>>  {
    let parts = filename.split("-").collect::<Vec<_>>();

    // 验证是否有足够的部分
    if parts.len() < 3 {
        return Err(format!("Invalid filename format: {}", filename).into());
    }

    // 386&473_177&454_154&383_363&402:标注框四个角点坐标,顺序为右下、左下、左上、右上
    // 解析矩形框坐标
    let bbox_part = parts[2].split('_').collect::<Vec<_>>();
    let left_top = bbox_part[0].split('&').collect::<Vec<_>>();
    let right_bottom = bbox_part[1].split('&').collect::<Vec<_>>();
    Ok(Bbox {
        x1: left_top[0].parse::<i32>()?,
        y1: left_top[1].parse::<i32>()?,
        x2: right_bottom[0].parse::<i32>()?,
        y2: right_bottom[1].parse::<i32>()?,
    })

}
5、矩形坐标转换为YOLO数据 converter.rs

根据解析的矩形Bbox坐标,以及图像的大小,判断数据是否正确,并返回YOLO格式的坐标数据

use crate::{Bbox, YoloData};

pub fn bbox2yolo(bbox: Bbox, image_size: (i32, i32), file_path: &str) -> Result<YoloData, Box<dyn std::error::Error>> {
    let (image_width, image_height) = image_size;

    // 验证图像尺寸不为零
    if image_width <= 0 || image_height <= 0 {
        return Err(format!("Image width and height must be positive in file: {} ({:?})", file_path, image_size).into());
    }

    // 验证边界框坐标有效性
    if bbox.x1 < 0 || bbox.y1 < 0 || bbox.x2 < 0 || bbox.y2 < 0 {
        return Err(format!("Bounding box coordinates must be non-negative in file: {} ({:?})", file_path, image_size).into());
    }

    if bbox.x1 >= image_width || bbox.x2 >= image_width ||
        bbox.y1 >= image_height || bbox.y2 >= image_height {
        return Err(format!("Bounding box coordinates exceed image dimensions {:?} in file: {} ({:?})", bbox, file_path, image_size).into());
    }

    if bbox.x1 >= bbox.x2 || bbox.y1 >= bbox.y2 {
        return Err(format!("Invalid bounding box: x1 must be less than x2 and y1 must be less than y2 in file: {} ({:?})", file_path, image_size).into());
    }

    let x_center = (((bbox.x1 + bbox.x2) as f32 / 2.0 / image_width as f32) * 1_000_000.0).round() / 1_000_000.0;
    let y_center = (((bbox.y1 + bbox.y2) as f32 / 2.0 / image_height as f32) * 1_000_000.0).round() / 1_000_000.0;
    let width = ((bbox.x2 - bbox.x1) as f32 / image_width as f32 * 1_000_000.0).round() / 1_000_000.0;
    let height = ((bbox.y2 - bbox.y1) as f32 / image_height as f32 * 1_000_000.0).round() / 1_000_000.0;

    Ok(YoloData {
        x_center,
        y_center,
        width,
        height,
    })
}
6、异常处理

在整个程序中,通过Result<T, E>这个枚举类型来处理异常,当函数处理成功时,则返回OK(T), 失败时,则返回Error(io:Error), 并通过match模式匹配来处理返回结果,例如:

// 解析YOLO数据
 let yolo_data = match bbox2yolo(bbox, image_size, &image_file.0) {
     Ok(data) => data,
     Err(e) => {
         error!("{}: Failed to parse YOLO data : {}", new_filename, e);
         continue;
     }
 };

当bbox2yolo函数处理成功,则将正常值返回赋给变量yolo_data,否则记录到错误日志中。

7、日志管理

日志管理使用的log和log4rs,日志配置在程序启动初始化

log4rs::init_file("log4rs.yaml", Default::default())?;

log4rs.yaml

appenders:
 stdout:
   kind: console
   encoder:
     pattern: "[{d}][{l}] {m}{n}"

 info_appender:
   kind: rolling_file
   path: "logs/info.log"
   append: true
   policy:
     kind: compound
     trigger:
       kind: time
       interval: 60      # 每天滚动
       modulate: true
     roller:
       kind: fixed_window
       base: 1
       count: 7
       pattern: "logs/info.{}.log"
   encoder:
     pattern: "[{d}][{l}] {m}{n}"

 error_appender:
   kind: rolling_file
   path: "logs/error.log"
   append: true
   filters:
     - kind: threshold
       level: error
   policy:
     kind: compound
     trigger:
       kind: time
       interval: 60
       modulate: true
     roller:
       kind: fixed_window
       base: 1
       count: 7
       pattern: "logs/error.{}.log"
   encoder:
     pattern: "[{d}][{l}] {m}{n}"

root:
 level: info
 appenders:
   - stdout
   - info_appender
   - error_appender

在需要记录日志的地方引入use log::{error, info};,运行后,就会按配置的规则来记录日志

三、总结

本次的Rust实践项目,通过基础入门知识构建了一个命令行的小工具,更清晰的了解到了变量的绑定、解构、模式匹配、方法、集合类型、错误处理等知识,以及cargo、编译链的使用方法,完整项目我已经推送到了GitCode中,项目地址 https://gitcode.com/szhf1980330/img2yolo.git

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐