mecha10-models 0.9.0

//! Preprocessing auto-detection and `config.json` generation
//!
//! Auto-detects normalization (mean/std) and input size for a downloaded
//! model, falling back through: HuggingFace `preprocessor_config.json` ->
//! catalog preprocessing preset -> ONNX metadata / sensible defaults. The
//! result is written out as the model's `config.json`.

use super::ModelService;
#[cfg(feature = "model-download")]
use std::path::Path;
#[cfg(feature = "model-download")]
use {
    super::catalog::{HFPreprocessorConfig, ModelCatalogEntry, PreprocessingPreset},
    anyhow::{Context, Result},
    indicatif::ProgressBar,
    mecha10_core::model::{CustomLabelsConfig, ModelConfig, PreprocessingConfig},
    tokio::fs,
};

impl ModelService {
    /// Fetch preprocessor_config.json from HuggingFace repo
    #[cfg(feature = "model-download")]
    pub(super) async fn fetch_hf_preprocessor_config(&self, repo: &str) -> Result<HFPreprocessorConfig> {
        let url = format!("https://huggingface.co/{}/raw/main/preprocessor_config.json", repo);

        let client = reqwest::Client::new();
        let response = client
            .get(&url)
            .send()
            .await
            .context(format!("Failed to fetch from {}", url))?;

        if !response.status().is_success() {
            anyhow::bail!(
                "HuggingFace preprocessor_config.json not found for {} (HTTP {})",
                repo,
                response.status()
            );
        }

        let config: HFPreprocessorConfig = response
            .json()
            .await
            .context("Failed to parse preprocessor_config.json")?;

        Ok(config)
    }

    /// Extract input size from ONNX model metadata
    ///
    /// Note: This is best-effort and may fail depending on the ONNX model format.
    /// Falls back to None if extraction fails.
    ///
    /// Only called from `auto_detect_preprocessing` below, which is itself
    /// gated on the `model-download` feature - without that cfg here too,
    /// a `model-download`-disabled build (e.g. `mecha10-schema-registry`,
    /// which depends on `mecha10-models` with no features) leaves this
    /// dead code and fails `cargo clippy -D warnings` (LAB-3569).
    #[cfg(all(feature = "model-download", feature = "vision"))]
    pub(super) fn extract_input_size_from_onnx(&self, model_path: &Path) -> Option<[u32; 2]> {
        use ort::session::Session;

        let session = Session::builder().ok()?.commit_from_file(model_path).ok()?;

        let _input = session.inputs.first()?;

        None
    }

    /// Stub when vision feature is disabled
    #[cfg(all(feature = "model-download", not(feature = "vision")))]
    pub(super) fn extract_input_size_from_onnx(&self, _model_path: &Path) -> Option<[u32; 2]> {
        None
    }

    /// Auto-detect preprocessing parameters with fallback priority:
    /// 1. HuggingFace preprocessor_config.json
    /// 2. Preprocessing preset from catalog
    /// 3. ONNX metadata (input size only, use sensible defaults for mean/std)
    #[cfg(feature = "model-download")]
    pub(super) async fn auto_detect_preprocessing(
        &self,
        entry: &ModelCatalogEntry,
        model_path: &Path,
        progress: Option<&ProgressBar>,
    ) -> Result<(PreprocessingConfig, [u32; 2])> {
        if let Some(pb) = progress {
            pb.set_message(format!("🔍 Auto-detecting preprocessing for {}", entry.name));
        }

        if let Ok(hf_config) = self.fetch_hf_preprocessor_config(&entry.repo).await {
            tracing::debug!(
                "HF config: size={:?}, crop_size={:?}",
                hf_config.size,
                hf_config.crop_size
            );

            let preprocessing = hf_config.to_preprocessing();
            let input_size = hf_config.input_size().or(entry.input_size).unwrap_or([224, 224]);

            tracing::debug!(
                "Detected preprocessing: mean={:?}, std={:?}, input_size={:?}",
                preprocessing.mean,
                preprocessing.std,
                input_size
            );

            if let Some(pb) = progress {
                pb.set_message(format!(
                    "✅ Auto-detected from HuggingFace (input_size={:?})",
                    input_size
                ));
            }

            return Ok((preprocessing, input_size));
        } else {
            tracing::debug!("Failed to fetch HuggingFace preprocessor config, falling back to preset");
        }

        if let Some(preset_name) = &entry.preprocessing_preset {
            if let Ok(preset) = PreprocessingPreset::from_name(preset_name) {
                let preprocessing = preset.to_config();

                let input_size = entry
                    .input_size
                    .or_else(|| self.extract_input_size_from_onnx(model_path))
                    .unwrap_or([224, 224]);

                if let Some(pb) = progress {
                    pb.set_message(format!(
                        "✅ Using preset '{}' (input_size={:?})",
                        preset_name, input_size
                    ));
                }

                return Ok((preprocessing, input_size));
            }
        }

        let input_size = entry
            .input_size
            .or_else(|| self.extract_input_size_from_onnx(model_path))
            .unwrap_or([224, 224]);

        let preprocessing = PreprocessingConfig {
            mean: [0.0, 0.0, 0.0],
            std: [1.0, 1.0, 1.0],
            channel_order: "RGB".to_string(),
        };

        if let Some(pb) = progress {
            pb.set_message(format!(
                "⚠️  Using fallback preprocessing (input_size={:?}). Consider editing config.json",
                input_size
            ));
        }

        Ok((preprocessing, input_size))
    }

    /// Generate and write config.json for a model with auto-detected values
    #[cfg(feature = "model-download")]
    pub(super) async fn generate_model_config(
        &self,
        entry: &ModelCatalogEntry,
        model_path: &Path,
        progress: Option<&ProgressBar>,
    ) -> Result<()> {
        let model_dir = self.catalog_model_dir(&entry.name);
        fs::create_dir_all(&model_dir)
            .await
            .context("Failed to create model directory")?;
        let config_path = model_dir.join("config.json");

        let (preprocessing, input_size) = self.auto_detect_preprocessing(entry, model_path, progress).await?;

        let num_classes = if entry.task == "object-detection" {
            entry.classes.len().max(1)
        } else {
            1000
        };

        let config = ModelConfig {
            name: entry.name.clone(),
            task: entry.task.clone(),
            repo: entry.repo.clone(),
            filename: entry.filename.clone(),
            input_size,
            preprocessing,
            num_classes,
            labels_file: "labels.txt".to_string(),
            custom_labels: CustomLabelsConfig::default(),
        };

        let json = serde_json::to_string_pretty(&config).context("Failed to serialize model config")?;

        fs::write(&config_path, json)
            .await
            .context("Failed to write model config.json")?;

        if let Some(pb) = progress {
            pb.set_message(format!("📝 Wrote config to {}", config_path.display()));
        }

        Ok(())
    }
}

#[cfg(all(test, feature = "model-download"))]
mod tests {
    use super::super::ModelService;
    use super::*;
    use std::path::PathBuf;

    fn service(models_dir: PathBuf) -> ModelService {
        ModelService::with_models_dir(models_dir).unwrap()
    }

    fn entry(name: &str, task: &str) -> ModelCatalogEntry {
        ModelCatalogEntry {
            name: name.to_string(),
            description: "test entry".to_string(),
            task: task.to_string(),
            repo: "mecha10-test-fixtures/does-not-exist".to_string(),
            filename: "model.onnx".to_string(),
            preprocessing_preset: None,
            classes: vec![],
            quantize: None,
            input_size: None,
        }
    }

    #[tokio::test]
    async fn generate_model_config_writes_to_the_namespaced_catalog_dir() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let entry = entry("mobilenet-v2", "image-classification");
        let model_path = models_dir.join("catalog").join("mobilenet-v2").join("model.onnx");

        service.generate_model_config(&entry, &model_path, None).await.unwrap();

        let expected_path = models_dir.join("catalog").join("mobilenet-v2").join("config.json");
        assert!(expected_path.exists(), "expected config.json at {:?}", expected_path);

        let content = tokio::fs::read_to_string(&expected_path).await.unwrap();
        let config: ModelConfig = serde_json::from_str(&content).unwrap();
        assert_eq!(config.name, "mobilenet-v2");

        assert!(!models_dir.join("mobilenet-v2").join("config.json").exists());
    }

    #[tokio::test]
    async fn generate_model_config_creates_the_catalog_dir_if_missing() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let entry = entry("yolov8n", "object-detection");
        let model_path = models_dir.join("catalog").join("yolov8n").join("model.onnx");

        service.generate_model_config(&entry, &model_path, None).await.unwrap();

        assert!(models_dir.join("catalog").join("yolov8n").join("config.json").exists());
    }

    #[test]
    fn extract_input_size_from_onnx_falls_back_to_none_for_a_missing_model() {
        let tmp = tempfile::tempdir().unwrap();
        let service = service(tmp.path().to_path_buf());

        let missing_model = tmp.path().join("does-not-exist.onnx");

        assert_eq!(service.extract_input_size_from_onnx(&missing_model), None);
    }
}