mecha10-models 0.6.3

//! Model catalog schema and HuggingFace preprocessor config parsing
//!
//! Types for the embedded `model_catalog.toml` catalog entries, the
//! auto-quantization config attached to a catalog entry, the built-in
//! preprocessing presets, and the schema used to parse a HuggingFace
//! `preprocessor_config.json`.

use anyhow::Result;
use mecha10_core::model::PreprocessingConfig;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

/// Model catalog entry from model_catalog.toml
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ModelCatalogEntry {
    /// Model identifier (e.g., "yolov8n-face", "mobilenet-v2")
    pub name: String,
    /// Human-readable description
    pub description: String,
    /// Task type (e.g., "object-detection", "image-classification")
    pub task: String,
    /// HuggingFace repository (e.g., "deepghs/yolo-face")
    pub repo: String,
    /// Filename within the repo (e.g., "yolov8n-face/model.onnx")
    pub filename: String,

    /// Optional preprocessing preset hint (e.g., "imagenet", "yolo", "coco")
    #[serde(default)]
    pub preprocessing_preset: Option<String>,

    /// Optional small class lists (e.g., ["face"] for face detection)
    #[serde(default)]
    pub classes: Vec<String>,

    /// Optional automatic quantization configuration
    #[serde(default)]
    pub quantize: Option<QuantizeConfig>,
}

/// Quantization configuration for automatic INT8 conversion
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct QuantizeConfig {
    /// Enable automatic quantization after download
    pub enabled: bool,
    /// Quantization method (currently only "dynamic_int8" supported)
    pub method: String,
}

/// Model catalog loaded from TOML
#[derive(Debug, Deserialize)]
pub(super) struct ModelCatalog {
    pub(super) models: Vec<ModelCatalogEntry>,
}

/// Preprocessing presets for common model families
#[derive(Debug, Clone, Copy)]
pub enum PreprocessingPreset {
    /// ImageNet normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ImageNet,
    /// YOLO normalization (mean=[0.0, 0.0, 0.0], std=[255.0, 255.0, 255.0] -> [0, 1])
    Yolo,
    /// COCO normalization (same as ImageNet)
    Coco,
    /// No normalization ([0, 255] range)
    Zero255,
}

impl PreprocessingPreset {
    /// Parse preset from string name
    pub fn from_name(name: &str) -> Result<Self> {
        match name.to_lowercase().as_str() {
            "imagenet" => Ok(Self::ImageNet),
            "yolo" => Ok(Self::Yolo),
            "coco" => Ok(Self::Coco),
            "zero255" | "0-255" => Ok(Self::Zero255),
            _ => anyhow::bail!("Unknown preprocessing preset: {}", name),
        }
    }

    /// Get preprocessing configuration for this preset
    pub fn to_config(self) -> PreprocessingConfig {
        match self {
            Self::ImageNet | Self::Coco => PreprocessingConfig {
                mean: [0.485, 0.456, 0.406],
                std: [0.229, 0.224, 0.225],
                channel_order: "RGB".to_string(),
            },
            Self::Yolo => PreprocessingConfig {
                mean: [0.0, 0.0, 0.0],
                std: [255.0, 255.0, 255.0],
                channel_order: "RGB".to_string(),
            },
            Self::Zero255 => PreprocessingConfig {
                mean: [0.0, 0.0, 0.0],
                std: [1.0, 1.0, 1.0],
                channel_order: "RGB".to_string(),
            },
        }
    }
}

/// HuggingFace preprocessor_config.json schema
#[cfg(feature = "model-download")]
#[derive(Debug, Clone, Deserialize)]
pub(super) struct HFPreprocessorConfig {
    #[serde(default)]
    pub(super) image_mean: Option<Vec<f32>>,
    #[serde(default)]
    pub(super) image_std: Option<Vec<f32>>,
    #[serde(default)]
    pub(super) size: Option<HFSize>,
    /// Some models (e.g., MobileNet) use crop_size instead of size
    #[serde(default)]
    pub(super) crop_size: Option<HFSize>,
}

#[cfg(feature = "model-download")]
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub(super) enum HFSize {
    Dict { height: u32, width: u32 },
    ShortestEdge { shortest_edge: u32 },
    Single(u32), // Square images
}

#[cfg(feature = "model-download")]
impl HFPreprocessorConfig {
    pub(super) fn to_preprocessing(&self) -> PreprocessingConfig {
        PreprocessingConfig {
            mean: [
                self.image_mean.as_ref().and_then(|v| v.first()).copied().unwrap_or(0.0),
                self.image_mean.as_ref().and_then(|v| v.get(1)).copied().unwrap_or(0.0),
                self.image_mean.as_ref().and_then(|v| v.get(2)).copied().unwrap_or(0.0),
            ],
            std: [
                self.image_std.as_ref().and_then(|v| v.first()).copied().unwrap_or(1.0),
                self.image_std.as_ref().and_then(|v| v.get(1)).copied().unwrap_or(1.0),
                self.image_std.as_ref().and_then(|v| v.get(2)).copied().unwrap_or(1.0),
            ],
            channel_order: "RGB".to_string(),
        }
    }

    pub(super) fn input_size(&self) -> Option<[u32; 2]> {
        // Check crop_size first (used by MobileNet, EfficientNet, etc.)
        // crop_size represents the actual model input dimensions
        if let Some(crop_size) = &self.crop_size {
            return match crop_size {
                HFSize::Dict { height, width } => Some([*width, *height]),
                HFSize::ShortestEdge { shortest_edge } => Some([*shortest_edge, *shortest_edge]),
                HFSize::Single(s) => Some([*s, *s]),
            };
        }

        // Fall back to size field
        match &self.size {
            Some(HFSize::Dict { height, width }) => Some([*width, *height]),
            Some(HFSize::ShortestEdge { shortest_edge }) => Some([*shortest_edge, *shortest_edge]),
            Some(HFSize::Single(s)) => Some([*s, *s]),
            None => None,
        }
    }
}