captcha-engine 0.4.8

ONNX-based captcha recognition engine
Documentation
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use std::env;
use std::fs;
use std::path::Path;

const MODEL_URL: &str = "https://huggingface.co/Milang/captcha-solver/resolve/main/captcha.onnx";

fn main() {
    // Only run if the embed-model feature is enabled
    if env::var("CARGO_FEATURE_EMBED_MODEL").is_ok() {
        let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
        let out_dir = env::var("OUT_DIR").unwrap();
        let dest_path = Path::new(&out_dir).join("model.onnx");

        // 1. Try local file first (fastest, supports offline devs who have the repo)
        let local_model_path =
            Path::new(&manifest_dir).join("assets/captcha_schwarz_finetuned.onnx");

        if local_model_path.exists() {
            println!(
                "cargo:warning=Embedding local model from {}",
                local_model_path.display()
            );
            fs::copy(&local_model_path, &dest_path).expect("Failed to copy local model file");
            println!("cargo:rerun-if-changed={}", local_model_path.display());
        } else if env::var("CARGO_FEATURE_DOWNLOAD_AT_BUILD").is_ok() {
            // 2. If missing locally AND feature enabled, download it
            println!("cargo:warning=Local model not found. Downloading from HuggingFace...");

            // We use reqwest blocking client since this is a build script
            // Note: We need to handle potential errors gracefully-ish or panic clearly
            let client = reqwest::blocking::Client::new();
            let mut response = client
                .get(MODEL_URL)
                .send()
                .expect("Failed to send request to download model");

            assert!(
                response.status().is_success(),
                "Failed to download model: HTTP {}",
                response.status()
            );

            let mut file =
                fs::File::create(&dest_path).expect("Failed to create model file in OUT_DIR");
            response
                .copy_to(&mut file)
                .expect("Failed to write downloaded model to file");

            println!(
                "cargo:warning=Model downloaded successfully to {}",
                dest_path.display()
            );
        } else {
            // 3. Missing and no download feature -> Error
            println!(
                "cargo:warning=Model file not found at {}. Embedded model validation will fail.",
                local_model_path.display()
            );
            panic!(
                "Model file not found at {}. Run with --features download-at-build to download automatically, or ensure the file exists.",
                local_model_path.display()
            );
        }

        println!("cargo:rerun-if-env-changed=CARGO_FEATURE_EMBED_MODEL");
        println!("cargo:rerun-if-env-changed=CARGO_FEATURE_DOWNLOAD_AT_BUILD");
    }
}