mecha10-models 0.9.0

//! Embedding resolved class-name labels into a model's own ONNX `metadata_props`
//! (shells out to an embedded Python + `onnx` script)
//!
//! `object-detector` reads its detectable-class list exclusively from a loaded
//! model's own ONNX `metadata_props` (a `class_names` custom-metadata key holding
//! a JSON array of strings - see `ObjectDetectorNode::parse_class_names_metadata`
//! in `packages/nodes/object-detector/src/lib.rs`), never from a hardcoded list or
//! a sidecar file. `labels.rs` already resolves a catalog model's labels from a
//! legitimate source (the HuggingFace repo's own `labels.json`, an inline catalog
//! `classes` list, or the ImageNet labels file) and writes them to `labels.txt` -
//! this module takes that already-resolved list and embeds it into the `.onnx`
//! file itself, so `object-detector` actually sees it.
//!
//! No pure-Rust crate for editing ONNX `metadata_props` is present anywhere in the
//! workspace (`ort`, the workspace's only ONNX dependency, is inference-only), so
//! this mirrors `quantize.rs`'s existing pattern of shelling out to an embedded
//! Python script - here using the `onnx` package's `metadata_props` API rather
//! than `onnxruntime`'s quantization tools.

use super::ModelService;
use anyhow::Result;
use indicatif::ProgressBar;
use std::path::Path;
use tokio::fs;

impl ModelService {
    /// Embed `class_names` into `model_path`'s ONNX `metadata_props`, in place.
    ///
    /// Only ever called with a non-empty `class_names` - the caller (`pull`) skips
    /// this entirely when label resolution didn't find anything, mirroring the
    /// sidecar `labels.txt` resolution's own "found nothing, move on" behavior
    /// rather than writing an empty/placeholder `class_names` array.
    ///
    /// Embedding is best-effort: if Python or the `onnx` package isn't available,
    /// this logs a warning and returns `Ok(())` rather than failing the whole
    /// `pull` - the model and its `labels.txt` sidecar are still usable by
    /// `image-classifier`, and failing to embed only degrades `object-detector`
    /// (which will fall back to zero detections, per PR #897) rather than the
    /// download itself.
    #[cfg(feature = "model-download")]
    pub(super) async fn embed_class_names_metadata(
        &self,
        model_path: &Path,
        class_names: &[String],
        progress: Option<&ProgressBar>,
    ) -> Result<()> {
        if class_names.is_empty() {
            return Ok(());
        }

        let python = match self.find_python() {
            Ok(python) => python,
            Err(e) => {
                tracing::warn!(
                    "Skipping embedding \"class_names\" metadata into {}: {e}. \
                     object-detector will report zero detections for this model until it's re-pulled \
                     with Python 3 available.",
                    model_path.display()
                );
                return Ok(());
            }
        };

        if let Some(pb) = progress {
            pb.set_message("Embedding class names into model metadata...");
        }

        let script = include_str!("../../../scripts/embed_class_names.py");
        let script_path = std::env::temp_dir().join("mecha10_embed_class_names.py");
        fs::write(&script_path, script).await?;

        let class_names_json = serde_json::to_string(class_names)?;
        let class_names_path = model_path.with_file_name("class_names.json");
        fs::write(&class_names_path, &class_names_json).await?;

        let output_result = tokio::process::Command::new(&python)
            .arg(&script_path)
            .arg(model_path)
            .arg(&class_names_path)
            .arg(model_path)
            .output()
            .await?;

        let _ = fs::remove_file(&script_path).await;
        let _ = fs::remove_file(&class_names_path).await;

        if !output_result.status.success() {
            let stderr = String::from_utf8_lossy(&output_result.stderr);
            tracing::warn!(
                "Skipping embedding \"class_names\" metadata into {}: {}\n\nTip: install with 'pip install onnx' \
                 to let object-detector see this model's classes.",
                model_path.display(),
                stderr.trim()
            );
            if let Some(pb) = progress {
                pb.set_message("⚠️  Skipped embedding class names (see logs)");
            }
            return Ok(());
        }

        if let Some(pb) = progress {
            pb.set_message("✅ Embedded class names into model metadata");
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::super::ModelService;
    use std::path::PathBuf;

    /// The tiny fixture model at `tests/fixtures/tiny.onnx` (a single `Identity`
    /// node, no `metadata_props`) - small enough to commit, just big enough to be
    /// a valid `.onnx` file `onnx.load`/`onnx.save` round-trips cleanly.
    fn fixture_model_bytes() -> &'static [u8] {
        include_bytes!("../../../tests/fixtures/tiny.onnx")
    }

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

    /// Skip (rather than fail) when the environment has no usable `python3`/`onnx`,
    /// the same "best effort" tool availability `embed_class_names_metadata` itself
    /// tolerates.
    ///
    /// CI images are expected to carry both, but a bare host might not (this repo's
    /// own onnx runtime doesn't even link on the bare host - see `launch-container`).
    async fn skip_if_onnx_unavailable(service: &ModelService) -> bool {
        let Ok(python) = service.find_python() else {
            return true;
        };

        let output = tokio::process::Command::new(&python)
            .arg("-c")
            .arg("import onnx")
            .output()
            .await;

        !matches!(output, Ok(out) if out.status.success())
    }

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

        if skip_if_onnx_unavailable(&service).await {
            return;
        }

        let model_path = tmp.path().join("model.onnx");
        tokio::fs::write(&model_path, fixture_model_bytes()).await.unwrap();

        let class_names = vec!["person".to_string(), "bicycle".to_string(), "car".to_string()];

        service
            .embed_class_names_metadata(&model_path, &class_names, None)
            .await
            .unwrap();

        let python = service.find_python().unwrap();
        let output = tokio::process::Command::new(&python)
            .arg("-c")
            .arg(format!(
                "import onnx, json, sys; m = onnx.load({model_path:?}); \
                 props = {{p.key: p.value for p in m.metadata_props}}; \
                 sys.stdout.write(props.get('class_names', ''))"
            ))
            .output()
            .await
            .unwrap();

        assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr));
        let raw = String::from_utf8_lossy(&output.stdout);
        let parsed: Vec<String> = serde_json::from_str(&raw).expect("class_names metadata must be a JSON string array");
        assert_eq!(parsed, class_names);
    }

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

        let model_path = tmp.path().join("model.onnx");
        tokio::fs::write(&model_path, fixture_model_bytes()).await.unwrap();
        let before = tokio::fs::read(&model_path).await.unwrap();

        service
            .embed_class_names_metadata(&model_path, &[], None)
            .await
            .unwrap();

        let after = tokio::fs::read(&model_path).await.unwrap();
        assert_eq!(
            before, after,
            "an empty class list must not write a placeholder class_names key"
        );
    }
}