mecha10_core/

model.rs

//! Model configuration utilities for AI nodes
//!
//! This module provides shared utilities for loading and managing ONNX model configurations.
//! It's used by both the CLI (for model management) and AI nodes (for runtime inference).
//!
//! # Design
//!
//! Model configurations are stored in `models/<name>/config.json` and include:
//! - Model metadata (name, task, repo, filename)
//! - Preprocessing parameters (input_size, mean, std, channel_order)
//! - Class labels configuration
//!
//! # Usage
//!
//! ```no_run
//! use mecha10_core::model::{load_model_config, load_labels};
//!
//! // Load complete model configuration
//! let config = load_model_config("mobilenet-v2")?;
//!
//! // Load class labels
//! let labels = load_labels("mobilenet-v2")?;
//! ```

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

/// Model configuration stored in models/<name>/config.json
///
/// This is the complete configuration for a model, including metadata,
/// preprocessing parameters, and label information.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
    /// Model name
    pub name: String,
    /// Task type (e.g., "image-classification", "object-detection")
    pub task: String,
    /// HuggingFace repository
    pub repo: String,
    /// Filename within the repo
    pub filename: String,
    /// Input size [width, height]
    pub input_size: [u32; 2],
    /// Preprocessing configuration
    pub preprocessing: PreprocessingConfig,
    /// Number of output classes
    pub num_classes: usize,
    /// Path to labels file (relative to model directory)
    pub labels_file: String,
    /// Custom labels configuration (for fine-tuning)
    #[serde(default)]
    pub custom_labels: CustomLabelsConfig,
}

/// Preprocessing configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PreprocessingConfig {
    /// Mean values for normalization [R, G, B]
    pub mean: [f32; 3],
    /// Standard deviation for normalization [R, G, B]
    pub std: [f32; 3],
    /// Channel order (RGB or BGR)
    pub channel_order: String,
}

/// Custom labels for fine-tuned models
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CustomLabelsConfig {
    /// Whether custom labels are enabled
    #[serde(default)]
    pub enabled: bool,
    /// Additional classes for fine-tuned models
    #[serde(default)]
    pub additional_classes: Vec<String>,
}

/// Load model configuration from disk
///
/// Loads the complete model configuration from `models/<name>/config.json`.
///
/// # Arguments
///
/// * `model_name` - Name of the model (e.g., "mobilenet-v2")
///
/// # Returns
///
/// The complete `ModelConfig` with all metadata and preprocessing parameters.
///
/// # Errors
///
/// Returns an error if:
/// - The config file doesn't exist
/// - The config file is not valid JSON
/// - The config file is missing required fields
///
/// # Example
///
/// ```no_run
/// use mecha10_core::model::load_model_config;
///
/// let config = load_model_config("mobilenet-v2")?;
/// println!("Model: {} ({})", config.name, config.task);
/// println!("Input size: {:?}", config.input_size);
/// println!("Preprocessing: mean={:?}, std={:?}", config.preprocessing.mean, config.preprocessing.std);
/// # Ok::<(), anyhow::Error>(())
/// ```
pub fn load_model_config(model_name: &str) -> Result<ModelConfig> {
    let config_path = PathBuf::from("models").join(model_name).join("config.json");

    if !config_path.exists() {
        anyhow::bail!(
            "Model config not found for '{}' at {}. Run: mecha10 models pull {}",
            model_name,
            config_path.display(),
            model_name
        );
    }

    let content = std::fs::read_to_string(&config_path)
        .context(format!("Failed to read model config: {}", config_path.display()))?;

    let config: ModelConfig = serde_json::from_str(&content)
        .context(format!("Failed to parse model config.json: {}", config_path.display()))?;

    Ok(config)
}

/// Load class labels from model directory
///
/// Loads class labels from `models/<name>/labels.txt` with support for custom labels.
/// If custom labels are enabled in the model config, they will be appended to the base labels.
///
/// # Arguments
///
/// * `model_name` - Name of the model (e.g., "mobilenet-v2")
///
/// # Returns
///
/// A vector of class label strings. The index in the vector corresponds to the class ID.
///
/// # Errors
///
/// Returns an error if:
/// - The labels file doesn't exist
/// - The labels file cannot be read
/// - The model config cannot be loaded (when checking for custom labels)
///
/// # Example
///
/// ```no_run
/// use mecha10_core::model::load_labels;
///
/// let labels = load_labels("mobilenet-v2")?;
/// println!("Loaded {} class labels", labels.len());
/// println!("First label: {}", labels[0]);
/// # Ok::<(), anyhow::Error>(())
/// ```
pub fn load_labels(model_name: &str) -> Result<Vec<String>> {
    // Load from models/<name>/labels.txt
    let labels_path = PathBuf::from("models").join(model_name).join("labels.txt");

    if !labels_path.exists() {
        anyhow::bail!(
            "Labels file not found for model '{}' at {}. Run: mecha10 models pull {}",
            model_name,
            labels_path.display(),
            model_name
        );
    }

    tracing::info!("📋 Loading labels from: {}", labels_path.display());

    let content = std::fs::read_to_string(&labels_path)
        .context(format!("Failed to read labels file: {}", labels_path.display()))?;

    let mut labels: Vec<String> = content.lines().map(|s| s.trim().to_string()).collect();

    // Load model config to check for custom labels
    let config_path = PathBuf::from("models").join(model_name).join("config.json");
    if config_path.exists() {
        if let Ok(config_content) = std::fs::read_to_string(&config_path) {
            if let Ok(config) = serde_json::from_str::<ModelConfig>(&config_content) {
                if config.custom_labels.enabled {
                    // Append custom labels
                    let num_custom = config.custom_labels.additional_classes.len();
                    labels.extend(config.custom_labels.additional_classes);
                    tracing::info!("✅ Added {} custom labels (total: {} labels)", num_custom, labels.len());
                }
            }
        }
    }

    tracing::info!("✅ Loaded {} labels", labels.len());
    Ok(labels)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model_config_serialization() {
        let config = ModelConfig {
            name: "test-model".to_string(),
            task: "image-classification".to_string(),
            repo: "test/repo".to_string(),
            filename: "model.onnx".to_string(),
            input_size: [224, 224],
            preprocessing: PreprocessingConfig {
                mean: [0.485, 0.456, 0.406],
                std: [0.229, 0.224, 0.225],
                channel_order: "RGB".to_string(),
            },
            num_classes: 1000,
            labels_file: "labels.txt".to_string(),
            custom_labels: CustomLabelsConfig::default(),
        };

        let json = serde_json::to_string(&config).unwrap();
        let deserialized: ModelConfig = serde_json::from_str(&json).unwrap();

        assert_eq!(config.name, deserialized.name);
        assert_eq!(config.input_size, deserialized.input_size);
        assert_eq!(config.preprocessing.mean, deserialized.preprocessing.mean);
    }
}