ultralytics-rust-core 0.1.0

//! 把官方 ultralytics/cfg（模型 YAML + default.yaml + 数据集 YAML）嵌入二进制，
//! 让 `yolo` 可执行文件像 pip 版 ultralytics 一样开箱即用。

use std::fs;
use std::path::Path;

fn main() {
    let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap();
    let cfg_dir = Path::new(&manifest).join("ultralytics/cfg");
    println!("cargo:rerun-if-changed={}", cfg_dir.display());

    let mut entries: Vec<(String, String)> = Vec::new();
    walk(&cfg_dir, &cfg_dir, &mut entries);
    entries.sort();

    let mut code = String::from(
        "// @generated by build.rs — 官方 ultralytics cfg 嵌入索引\n\
         use std::collections::HashMap;\n\
         use std::sync::OnceLock;\n\n\
         fn table() -> &'static HashMap<&'static str, &'static str> {\n\
             static TABLE: OnceLock<HashMap<&'static str, &'static str>> = OnceLock::new();\n\
             TABLE.get_or_init(|| {\n\
                 let mut m = HashMap::new();\n",
    );
    for (rel, _path) in &entries {
        code.push_str(&format!("        m.insert({rel:?}, include_str!(concat!(env!(\"CARGO_MANIFEST_DIR\"), \"/ultralytics/cfg/{rel}\")));\n", rel = rel));
    }
    code.push_str("                m\n            })\n        }\n\n");
    code.push_str(
        "/// 按相对路径取嵌入的 cfg 文本（如 \"11/yolo11.yaml\"、\"default.yaml\"）。\n\
         pub fn get(rel: &str) -> Option<&'static str> { table().get(rel).copied() }\n\n\
         /// 注册表中是否存在该相对路径。\n\
         pub fn contains(rel: &str) -> bool { table().contains_key(rel) }\n\n\
         /// 全部相对路径（稳定排序）。\n\
         pub fn keys() -> Vec<&'static str> { let mut v: Vec<&'static str> = table().keys().copied().collect(); v.sort(); v }\n",
    );

    let out = std::env::var("OUT_DIR").unwrap();
    fs::write(Path::new(&out).join("cfg_index.rs"), code).expect("write cfg_index.rs");
}

fn walk(root: &Path, dir: &Path, out: &mut Vec<(String, String)>) {
    if let Ok(rd) = fs::read_dir(dir) {
        for e in rd.flatten() {
            let p = e.path();
            if p.is_dir() {
                walk(root, &p, out);
            } else if p.extension().and_then(|x| x.to_str()) == Some("yaml") {
                let rel = p
                    .strip_prefix(root)
                    .unwrap()
                    .to_string_lossy()
                    .replace('\\', "/");
                out.push((rel, p.to_string_lossy().into_owned()));
            }
        }
    }
}