mecha10-models 4.2.6

//! Model catalog schema, static compiled-in list, and HuggingFace preprocessor config parsing
//!
//! Types for the public model catalog's 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`.
//!
//! The catalog itself used to be a TOML file (`model_catalog.toml`) compiled directly into
//! this CLI binary via `include_str!`, then later a live `GET /api/models/catalog` call to
//! `mecha10-api` (`mecha10-industries/mecha10-api`'s `models::registry::catalog` module). Both
//! of `mecha10-api`'s public model routes (`/api/models/catalog` and
//! `/api/downloads/model/{repo}/{filename}`) have since been removed - the catalog is nothing
//! but a list of public HuggingFace repos, so there's no reason to route either the listing or
//! the actual model bytes through `mecha10-api` at all. [`catalog_entries`] is now the single
//! source of truth: a static, compiled-in `Vec`, mirrored from the same public models every
//! entry already names on HuggingFace Hub (`download::pull_from_repo` fetches the artifact
//! itself directly from `https://huggingface.co/{repo}/resolve/main/{filename}`). No network
//! call is needed to list or resolve a catalog model by name anymore.

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

fn entry(
    name: &str,
    description: &str,
    task: &str,
    repo: &str,
    filename: &str,
    preprocessing_preset: Option<&str>,
    classes: &[&str],
) -> ModelCatalogEntry {
    ModelCatalogEntry {
        name: name.to_string(),
        description: description.to_string(),
        task: task.to_string(),
        repo: repo.to_string(),
        filename: filename.to_string(),
        preprocessing_preset: preprocessing_preset.map(|p| p.to_string()),
        classes: classes.iter().map(|c| c.to_string()).collect(),
        quantize: None,
        input_size: None,
    }
}

/// The public model catalog - a static, compiled-in list of public HuggingFace repos. This is
/// the sole source of truth (see the module doc comment); no network call, caching, or
/// fallback logic is involved in producing it.
pub fn catalog_entries() -> Vec<ModelCatalogEntry> {
    vec![
        entry(
            "yolov8n-face",
            "YOLOv8 Nano - Fast face detection",
            "object-detection",
            "deepghs/yolo-face",
            "yolov8n-face/model.onnx",
            Some("yolo"),
            &["face"],
        ),
        entry(
            "yolov8m-face",
            "YOLOv8 Medium - Better accuracy face detection",
            "object-detection",
            "deepghs/yolo-face",
            "yolov8m-face/model.onnx",
            Some("yolo"),
            &["face"],
        ),
        entry(
            "yolov8n",
            "YOLOv8 Nano - Fast object detection (80 COCO classes)",
            "object-detection",
            "deepghs/yolos",
            "yolov8n/model.onnx",
            Some("yolo"),
            &[],
        ),
        entry(
            "yolov8s",
            "YOLOv8 Small - Balanced speed/accuracy (80 COCO classes)",
            "object-detection",
            "deepghs/yolos",
            "yolov8s/model.onnx",
            Some("yolo"),
            &[],
        ),
        entry(
            "yolov8m",
            "YOLOv8 Medium - Better accuracy (80 COCO classes)",
            "object-detection",
            "deepghs/yolos",
            "yolov8m/model.onnx",
            Some("yolo"),
            &[],
        ),
        entry(
            "yolov8l",
            "YOLOv8 Large - High accuracy (80 COCO classes)",
            "object-detection",
            "deepghs/yolos",
            "yolov8l/model.onnx",
            Some("yolo"),
            &[],
        ),
        entry(
            "yolov8x",
            "YOLOv8 Extra-Large - Best accuracy (80 COCO classes)",
            "object-detection",
            "deepghs/yolos",
            "yolov8x/model.onnx",
            Some("yolo"),
            &[],
        ),
        entry(
            "mobilenet-v2",
            "MobileNetV2 - Fast image classification (1000 ImageNet classes)",
            "image-classification",
            "onnx-community/mobilenet_v2_0.75_160",
            "onnx/model.onnx",
            Some("imagenet"),
            &[],
        ),
        entry(
            "resnet18",
            "ResNet-18 - Image classification (1000 ImageNet classes)",
            "image-classification",
            "onnx-community/resnet-18",
            "onnx/model.onnx",
            Some("imagenet"),
            &[],
        ),
        ModelCatalogEntry {
            // OSNet's export fixes a non-square 256x128 (h, w) input, unlike the generic
            // ImageNet 224x224 square default `extract_input_size_from_onnx` falls back to
            // when it can't introspect the ONNX graph - which, for this model, is always
            // (that extraction is still a stub). Without this override, DeepSORT's reference
            // embedding extraction sends the model a [1,3,224,224] tensor and the ONNX
            // runtime rejects it outright.
            input_size: Some([128, 256]),
            ..entry(
                "osnet_x0_25",
                "OSNet x0.25 - Lightweight person re-identification features for DeepSORT (512-dim, MSMT17)",
                "feature-extraction",
                "anriha/osnet_x0_25_msmt17",
                "osnet_x0_25_msmt17.onnx",
                Some("imagenet"),
                &[],
            )
        },
    ]
}

/// Model catalog entry, as returned by [`catalog_entries`]
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ModelCatalogEntry {
    /// Model identifier (e.g., "yolov8n-face", "mobilenet-v2")
    #[schemars(example = &"yolov8n-face")]
    pub name: String,
    /// Human-readable description
    #[schemars(example = &"YOLOv8 Nano - Fast face detection")]
    pub description: String,
    /// Task type (e.g., "object-detection", "image-classification")
    #[schemars(example = &"object-detection")]
    pub task: String,
    /// HuggingFace repository (e.g., "deepghs/yolo-face")
    #[schemars(example = &"deepghs/yolo-face")]
    pub repo: String,
    /// Filename within the repo (e.g., "yolov8n-face/model.onnx")
    #[schemars(example = &"yolov8n-face/model.onnx")]
    pub filename: String,

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

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

    /// Optional automatic quantization configuration
    #[serde(default)]
    #[schemars(example = Some(QuantizeConfig { enabled: true, method: "dynamic_int8".to_string() }))]
    pub quantize: Option<QuantizeConfig>,

    /// Explicit `[width, height]` override for the model's input tensor size - takes
    /// precedence over auto-detection (HuggingFace `preprocessor_config.json`, then ONNX
    /// graph introspection) in `preprocessing::auto_detect_preprocessing`. Most entries leave
    /// this `None`; it exists for models where auto-detection is unreliable, e.g. a non-square
    /// re-id backbone that would otherwise silently fall back to the generic 224x224 default.
    #[serde(default)]
    #[schemars(example = Some([128_u32, 256_u32]))]
    pub input_size: Option<[u32; 2]>,
}

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

#[derive(Debug, Clone, Copy)]
pub enum PreprocessingPreset {
    ImageNet,
    Yolo,
    Coco,
    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(),
            },
        }
    }
}

#[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>,
    #[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),
}

#[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]> {
        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]),
            };
        }

        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,
        }
    }
}

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

    #[test]
    fn catalog_entries_are_non_empty_and_uniquely_named() {
        let entries = catalog_entries();
        assert!(!entries.is_empty());

        let mut names: Vec<&str> = entries.iter().map(|e| e.name.as_str()).collect();
        names.sort();
        names.dedup();
        assert_eq!(names.len(), entries.len(), "catalog entry names must be unique");
    }

    #[test]
    fn yolov8n_has_no_classes_override_meaning_default_coco() {
        let entries = catalog_entries();
        let yolov8n = entries.iter().find(|e| e.name == "yolov8n").expect("yolov8n entry");
        assert!(yolov8n.classes.is_empty());
    }

    #[test]
    fn face_models_carry_an_explicit_small_class_list() {
        let entries = catalog_entries();
        let face_model = entries
            .iter()
            .find(|e| e.name == "yolov8n-face")
            .expect("yolov8n-face entry");
        assert_eq!(face_model.classes, vec!["face".to_string()]);
    }

    #[test]
    fn osnet_carries_an_explicit_non_square_input_size_override() {
        let entries = catalog_entries();
        let osnet = entries
            .iter()
            .find(|e| e.name == "osnet_x0_25")
            .expect("osnet_x0_25 entry");
        assert_eq!(osnet.input_size, Some([128, 256]));
    }

    #[test]
    fn every_entry_resolves_to_a_valid_preprocessing_preset() {
        let entries = catalog_entries();
        for entry in &entries {
            let preset = entry
                .preprocessing_preset
                .as_deref()
                .unwrap_or_else(|| panic!("entry '{}' has no preprocessing_preset", entry.name));
            PreprocessingPreset::from_name(preset).unwrap_or_else(|_| {
                panic!(
                    "entry '{}' has an invalid preprocessing preset '{}'",
                    entry.name, preset
                )
            });
        }
    }
}