use anyhow::{Context, Result};
use mecha10_core::model::PreprocessingConfig;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::time::Duration;
pub async fn fetch_catalog() -> Result<Vec<ModelCatalogEntry>> {
let url = mecha10_cli_core::paths::urls::mecha10_api_path("/models/catalog");
let client = reqwest::Client::builder()
.user_agent("mecha10-cli")
.timeout(Duration::from_secs(10))
.build()
.context("Failed to build HTTP client")?;
let response = client
.get(&url)
.send()
.await
.context(format!("Failed to fetch model catalog from {url}"))?
.error_for_status()
.context(format!("Model catalog request to {url} failed"))?;
response
.json::<Vec<ModelCatalogEntry>>()
.await
.context(format!("Malformed model catalog response from {url}"))
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ModelCatalogEntry {
#[schemars(example = &"yolov8n-face")]
pub name: String,
#[schemars(example = &"YOLOv8 Nano - Fast face detection")]
pub description: String,
#[schemars(example = &"object-detection")]
pub task: String,
#[schemars(example = &"deepghs/yolo-face")]
pub repo: String,
#[schemars(example = &"yolov8n-face/model.onnx")]
pub filename: String,
#[serde(default)]
#[schemars(example = &"yolo")]
pub preprocessing_preset: Option<String>,
#[serde(default)]
#[schemars(example = vec!["face".to_string()])]
pub classes: Vec<String>,
#[serde(default)]
#[schemars(example = Some(QuantizeConfig { enabled: true, method: "dynamic_int8".to_string() }))]
pub quantize: Option<QuantizeConfig>,
#[serde(default)]
#[schemars(example = Some([128_u32, 256_u32]))]
pub input_size: Option<[u32; 2]>,
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct QuantizeConfig {
#[schemars(example = true)]
pub enabled: bool,
#[schemars(example = &"dynamic_int8")]
pub method: String,
}
#[derive(Debug, Clone, Copy)]
pub enum PreprocessingPreset {
ImageNet,
Yolo,
Coco,
Zero255,
}
impl PreprocessingPreset {
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),
}
}
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,
}
}
}