mecha10_models/services/model_service/
catalog.rs

1//! Model catalog schema and HuggingFace preprocessor config parsing
2//!
3//! Types for the embedded `model_catalog.toml` catalog entries, the
4//! auto-quantization config attached to a catalog entry, the built-in
5//! preprocessing presets, and the schema used to parse a HuggingFace
6//! `preprocessor_config.json`.
7
8use anyhow::Result;
9use mecha10_core::model::PreprocessingConfig;
10use schemars::JsonSchema;
11use serde::{Deserialize, Serialize};
12
13/// Model catalog entry from model_catalog.toml
14#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
15pub struct ModelCatalogEntry {
16    /// Model identifier (e.g., "yolov8n-face", "mobilenet-v2")
17    pub name: String,
18    /// Human-readable description
19    pub description: String,
20    /// Task type (e.g., "object-detection", "image-classification")
21    pub task: String,
22    /// HuggingFace repository (e.g., "deepghs/yolo-face")
23    pub repo: String,
24    /// Filename within the repo (e.g., "yolov8n-face/model.onnx")
25    pub filename: String,
26
27    /// Optional preprocessing preset hint (e.g., "imagenet", "yolo", "coco")
28    #[serde(default)]
29    pub preprocessing_preset: Option<String>,
30
31    /// Optional small class lists (e.g., ["face"] for face detection)
32    #[serde(default)]
33    pub classes: Vec<String>,
34
35    /// Optional automatic quantization configuration
36    #[serde(default)]
37    pub quantize: Option<QuantizeConfig>,
38}
39
40/// Quantization configuration for automatic INT8 conversion
41#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
42pub struct QuantizeConfig {
43    /// Enable automatic quantization after download
44    pub enabled: bool,
45    /// Quantization method (currently only "dynamic_int8" supported)
46    pub method: String,
47}
48
49/// Model catalog loaded from TOML
50#[derive(Debug, Deserialize)]
51pub(super) struct ModelCatalog {
52    pub(super) models: Vec<ModelCatalogEntry>,
53}
54
55/// Preprocessing presets for common model families
56#[derive(Debug, Clone, Copy)]
57pub enum PreprocessingPreset {
58    /// ImageNet normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
59    ImageNet,
60    /// YOLO normalization (mean=[0.0, 0.0, 0.0], std=[255.0, 255.0, 255.0] -> [0, 1])
61    Yolo,
62    /// COCO normalization (same as ImageNet)
63    Coco,
64    /// No normalization ([0, 255] range)
65    Zero255,
66}
67
68impl PreprocessingPreset {
69    /// Parse preset from string name
70    pub fn from_name(name: &str) -> Result<Self> {
71        match name.to_lowercase().as_str() {
72            "imagenet" => Ok(Self::ImageNet),
73            "yolo" => Ok(Self::Yolo),
74            "coco" => Ok(Self::Coco),
75            "zero255" | "0-255" => Ok(Self::Zero255),
76            _ => anyhow::bail!("Unknown preprocessing preset: {}", name),
77        }
78    }
79
80    /// Get preprocessing configuration for this preset
81    pub fn to_config(self) -> PreprocessingConfig {
82        match self {
83            Self::ImageNet | Self::Coco => PreprocessingConfig {
84                mean: [0.485, 0.456, 0.406],
85                std: [0.229, 0.224, 0.225],
86                channel_order: "RGB".to_string(),
87            },
88            Self::Yolo => PreprocessingConfig {
89                mean: [0.0, 0.0, 0.0],
90                std: [255.0, 255.0, 255.0],
91                channel_order: "RGB".to_string(),
92            },
93            Self::Zero255 => PreprocessingConfig {
94                mean: [0.0, 0.0, 0.0],
95                std: [1.0, 1.0, 1.0],
96                channel_order: "RGB".to_string(),
97            },
98        }
99    }
100}
101
102/// HuggingFace preprocessor_config.json schema
103#[cfg(feature = "model-download")]
104#[derive(Debug, Clone, Deserialize)]
105pub(super) struct HFPreprocessorConfig {
106    #[serde(default)]
107    pub(super) image_mean: Option<Vec<f32>>,
108    #[serde(default)]
109    pub(super) image_std: Option<Vec<f32>>,
110    #[serde(default)]
111    pub(super) size: Option<HFSize>,
112    /// Some models (e.g., MobileNet) use crop_size instead of size
113    #[serde(default)]
114    pub(super) crop_size: Option<HFSize>,
115}
116
117#[cfg(feature = "model-download")]
118#[derive(Debug, Clone, Deserialize)]
119#[serde(untagged)]
120pub(super) enum HFSize {
121    Dict { height: u32, width: u32 },
122    ShortestEdge { shortest_edge: u32 },
123    Single(u32), // Square images
124}
125
126#[cfg(feature = "model-download")]
127impl HFPreprocessorConfig {
128    pub(super) fn to_preprocessing(&self) -> PreprocessingConfig {
129        PreprocessingConfig {
130            mean: [
131                self.image_mean.as_ref().and_then(|v| v.first()).copied().unwrap_or(0.0),
132                self.image_mean.as_ref().and_then(|v| v.get(1)).copied().unwrap_or(0.0),
133                self.image_mean.as_ref().and_then(|v| v.get(2)).copied().unwrap_or(0.0),
134            ],
135            std: [
136                self.image_std.as_ref().and_then(|v| v.first()).copied().unwrap_or(1.0),
137                self.image_std.as_ref().and_then(|v| v.get(1)).copied().unwrap_or(1.0),
138                self.image_std.as_ref().and_then(|v| v.get(2)).copied().unwrap_or(1.0),
139            ],
140            channel_order: "RGB".to_string(),
141        }
142    }
143
144    pub(super) fn input_size(&self) -> Option<[u32; 2]> {
145        // Check crop_size first (used by MobileNet, EfficientNet, etc.)
146        // crop_size represents the actual model input dimensions
147        if let Some(crop_size) = &self.crop_size {
148            return match crop_size {
149                HFSize::Dict { height, width } => Some([*width, *height]),
150                HFSize::ShortestEdge { shortest_edge } => Some([*shortest_edge, *shortest_edge]),
151                HFSize::Single(s) => Some([*s, *s]),
152            };
153        }
154
155        // Fall back to size field
156        match &self.size {
157            Some(HFSize::Dict { height, width }) => Some([*width, *height]),
158            Some(HFSize::ShortestEdge { shortest_edge }) => Some([*shortest_edge, *shortest_edge]),
159            Some(HFSize::Single(s)) => Some([*s, *s]),
160            None => None,
161        }
162    }
163}
mecha10_models/services/model_service/catalog.rs

mecha10_models/services/model_service/
catalog.rs
