rustytorch/demos/image-classifier-shared/src/lib.rs

//! Shared IPC types for the Image Classifier demo
//!
//! This crate defines the data structures shared between the Rust backend
//! and the TypeScript frontend for the image classification demo.

use serde::{Deserialize, Serialize};

/// Supported model architectures for image classification
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelArchitecture {
    /// Vision Transformer Base/16
    ViTBase16,
    /// Vision Transformer Large/16
    ViTLarge16,
    /// ConvNeXt Tiny
    ConvNeXtTiny,
    /// ConvNeXt Small
    ConvNeXtSmall,
    /// ConvNeXt Base
    ConvNeXtBase,
}

impl ModelArchitecture {
    /// Get human-readable name
    pub fn display_name(&self) -> &'static str {
        match self {
            Self::ViTBase16 => "ViT-Base/16",
            Self::ViTLarge16 => "ViT-Large/16",
            Self::ConvNeXtTiny => "ConvNeXt-Tiny",
            Self::ConvNeXtSmall => "ConvNeXt-Small",
            Self::ConvNeXtBase => "ConvNeXt-Base",
        }
    }

    /// Get model parameter count (approximate)
    pub fn param_count(&self) -> usize {
        match self {
            Self::ViTBase16 => 86_000_000,
            Self::ViTLarge16 => 307_000_000,
            Self::ConvNeXtTiny => 28_000_000,
            Self::ConvNeXtSmall => 50_000_000,
            Self::ConvNeXtBase => 89_000_000,
        }
    }
}

/// Configuration for initializing the classifier
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassifierConfig {
    /// Model architecture to use
    pub architecture: ModelArchitecture,
    /// Number of classes (default: 1000 for ImageNet)
    pub num_classes: usize,
    /// Input image size (default: 224)
    pub image_size: usize,
    /// Whether to use GPU if available
    pub use_gpu: bool,
}

impl Default for ClassifierConfig {
    fn default() -> Self {
        Self {
            architecture: ModelArchitecture::ViTBase16,
            num_classes: 1000,
            image_size: 224,
            use_gpu: true,
        }
    }
}

/// A single classification prediction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Prediction {
    /// Class index (0-999 for ImageNet)
    pub class_idx: usize,
    /// Human-readable class label
    pub label: String,
    /// Confidence score (0.0 - 1.0)
    pub confidence: f32,
}

/// Result of classifying an image
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassificationResult {
    /// Top-K predictions sorted by confidence
    pub predictions: Vec<Prediction>,
    /// Inference time in milliseconds
    pub inference_time_ms: f64,
    /// Preprocessing time in milliseconds
    pub preprocess_time_ms: f64,
    /// Input image dimensions (width, height)
    pub input_dimensions: (usize, usize),
    /// Model used for inference
    pub model: String,
}

/// Status of the classifier service
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassifierStatus {
    /// Whether the model is loaded and ready
    pub initialized: bool,
    /// Current model architecture (if loaded)
    pub model: Option<String>,
    /// Compute device being used
    pub device: String,
    /// Total number of inferences performed
    pub inference_count: u64,
    /// Average inference time in milliseconds
    pub avg_inference_time_ms: f64,
}

/// Performance metrics for the classifier
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassifierMetrics {
    /// Total number of inferences
    pub total_inferences: u64,
    /// Average inference time in milliseconds
    pub avg_inference_ms: f64,
    /// Minimum inference time
    pub min_inference_ms: f64,
    /// Maximum inference time
    pub max_inference_ms: f64,
    /// Average preprocessing time
    pub avg_preprocess_ms: f64,
    /// Throughput (images per second)
    pub throughput_fps: f64,
}

/// Request to classify an image
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassifyRequest {
    /// Base64-encoded image data
    pub image_data: String,
    /// Number of top predictions to return (default: 5)
    pub top_k: Option<usize>,
}

/// List of available ImageNet-1k classes
pub const IMAGENET_CLASSES: &[&str] = &[
    "tench",
    "goldfish",
    "great white shark",
    "tiger shark",
    "hammerhead",
    "electric ray",
    "stingray",
    "cock",
    "hen",
    "ostrich",
    // ... (truncated for brevity - in production, include all 1000 classes)
    // For demo purposes, we'll load from a separate file
];

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

    #[test]
    fn test_config_default() {
        let config = ClassifierConfig::default();
        assert_eq!(config.num_classes, 1000);
        assert_eq!(config.image_size, 224);
    }

    #[test]
    fn test_architecture_display() {
        assert_eq!(ModelArchitecture::ViTBase16.display_name(), "ViT-Base/16");
    }
}