rustytorch/demos/rtx-image-classifier-demo/src/classifier.rs

//! Image classifier implementation using rtx-vision models

use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use base64::Engine;
use image::GenericImageView;
use rtx_tensor::{Device, Tensor};
use rtx_vision::{ConvNeXt, ConvNeXtConfig, ViT, ViTConfig};

use crate::error::{ClassifierError, Result};
use crate::imagenet_labels::get_label;
use image_classifier_shared::{
    ClassificationResult, ClassifierConfig, ClassifierMetrics, ClassifierStatus, ModelArchitecture,
    Prediction,
};

/// Image classifier using Vision Transformers and `ConvNeXt` models
pub struct ImageClassifier {
    /// Current model (boxed for type erasure)
    model: RwLock<Option<ClassifierModel>>,
    /// Configuration
    config: ClassifierConfig,
    /// Device being used
    device: Device,
    /// Inference statistics
    stats: InferenceStats,
}

/// Type-erased model wrapper
enum ClassifierModel {
    ViT(ViT),
    ConvNeXt(ConvNeXt),
}

impl ClassifierModel {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        match self {
            Self::ViT(model) => model.forward(input).map_err(ClassifierError::from),
            Self::ConvNeXt(model) => model.forward(input).map_err(ClassifierError::from),
        }
    }

    fn name(&self) -> &'static str {
        match self {
            Self::ViT(_) => "ViT",
            Self::ConvNeXt(_) => "ConvNeXt",
        }
    }
}

/// Inference statistics tracker
struct InferenceStats {
    total_inferences: AtomicU64,
    total_time_ms: RwLock<f64>,
    min_time_ms: RwLock<f64>,
    max_time_ms: RwLock<f64>,
    total_preprocess_ms: RwLock<f64>,
}

impl InferenceStats {
    fn new() -> Self {
        Self {
            total_inferences: AtomicU64::new(0),
            total_time_ms: RwLock::new(0.0),
            min_time_ms: RwLock::new(f64::MAX),
            max_time_ms: RwLock::new(0.0),
            total_preprocess_ms: RwLock::new(0.0),
        }
    }

    fn record(&self, inference_ms: f64, preprocess_ms: f64) {
        self.total_inferences.fetch_add(1, Ordering::Relaxed);

        let mut total = self.total_time_ms.write().unwrap();
        *total += inference_ms;

        let mut min = self.min_time_ms.write().unwrap();
        if inference_ms < *min {
            *min = inference_ms;
        }

        let mut max = self.max_time_ms.write().unwrap();
        if inference_ms > *max {
            *max = inference_ms;
        }

        let mut preprocess = self.total_preprocess_ms.write().unwrap();
        *preprocess += preprocess_ms;
    }

    fn get_metrics(&self) -> ClassifierMetrics {
        let count = self.total_inferences.load(Ordering::Relaxed);
        let total_ms = *self.total_time_ms.read().unwrap();
        let min_ms = *self.min_time_ms.read().unwrap();
        let max_ms = *self.max_time_ms.read().unwrap();
        let preprocess_ms = *self.total_preprocess_ms.read().unwrap();

        let avg_inference = if count > 0 {
            total_ms / count as f64
        } else {
            0.0
        };
        let avg_preprocess = if count > 0 {
            preprocess_ms / count as f64
        } else {
            0.0
        };
        let throughput = if avg_inference > 0.0 {
            1000.0 / avg_inference
        } else {
            0.0
        };

        ClassifierMetrics {
            total_inferences: count,
            avg_inference_ms: avg_inference,
            min_inference_ms: if min_ms == f64::MAX { 0.0 } else { min_ms },
            max_inference_ms: max_ms,
            avg_preprocess_ms: avg_preprocess,
            throughput_fps: throughput,
        }
    }
}

impl ImageClassifier {
    /// Create a new image classifier with the given configuration
    pub fn new(config: ClassifierConfig) -> Result<Self> {
        // Select device based on configuration and availability
        let device = Self::select_device(config.use_gpu);

        Ok(Self {
            model: RwLock::new(None),
            config,
            device,
            stats: InferenceStats::new(),
        })
    }

    /// Select the best available device
    fn select_device(use_gpu: bool) -> Device {
        if use_gpu {
            // Try to get the best GPU device
            if let Ok(device) = Device::try_default()
                && device.is_gpu()
            {
                return device;
            }
        }

        // Fall back to CPU
        Device::cpu()
    }

    /// Initialize the model
    pub fn initialize(&self) -> Result<()> {
        let model = self.create_model()?;
        *self.model.write().unwrap() = Some(model);
        tracing::info!(
            "Initialized {} model on {:?}",
            self.config.architecture.display_name(),
            self.device_name()
        );
        Ok(())
    }

    /// Create the model based on configuration
    fn create_model(&self) -> Result<ClassifierModel> {
        match self.config.architecture {
            ModelArchitecture::ViTBase16 => {
                let mut config = ViTConfig::base_16();
                config.num_classes = self.config.num_classes;
                config.image_size = self.config.image_size;
                let model = ViT::new(config, &self.device)?;
                Ok(ClassifierModel::ViT(model))
            }
            ModelArchitecture::ViTLarge16 => {
                let mut config = ViTConfig::large_16();
                config.num_classes = self.config.num_classes;
                config.image_size = self.config.image_size;
                let model = ViT::new(config, &self.device)?;
                Ok(ClassifierModel::ViT(model))
            }
            ModelArchitecture::ConvNeXtTiny => {
                let mut config = ConvNeXtConfig::tiny();
                config.num_classes = self.config.num_classes;
                let model = ConvNeXt::new(config, &self.device)?;
                Ok(ClassifierModel::ConvNeXt(model))
            }
            ModelArchitecture::ConvNeXtSmall => {
                let mut config = ConvNeXtConfig::small();
                config.num_classes = self.config.num_classes;
                let model = ConvNeXt::new(config, &self.device)?;
                Ok(ClassifierModel::ConvNeXt(model))
            }
            ModelArchitecture::ConvNeXtBase => {
                let mut config = ConvNeXtConfig::base();
                config.num_classes = self.config.num_classes;
                let model = ConvNeXt::new(config, &self.device)?;
                Ok(ClassifierModel::ConvNeXt(model))
            }
        }
    }

    /// Check if the classifier is initialized
    pub fn is_initialized(&self) -> bool {
        self.model.read().unwrap().is_some()
    }

    /// Get device name
    pub fn device_name(&self) -> String {
        self.device.to_string()
    }

    /// Classify an image from base64-encoded data
    pub fn classify_base64(&self, image_data: &str, top_k: usize) -> Result<ClassificationResult> {
        // Decode base64
        let bytes = base64::engine::general_purpose::STANDARD.decode(image_data)?;

        self.classify_bytes(&bytes, top_k)
    }

    /// Classify an image from raw bytes
    pub fn classify_bytes(&self, image_bytes: &[u8], top_k: usize) -> Result<ClassificationResult> {
        let model_guard = self.model.read().unwrap();
        let model = model_guard
            .as_ref()
            .ok_or(ClassifierError::NotInitialized)?;

        // Preprocess image
        let preprocess_start = Instant::now();
        let (input_tensor, original_dims) = self.preprocess_image(image_bytes)?;
        let preprocess_time = preprocess_start.elapsed().as_secs_f64() * 1000.0;

        // Run inference
        let inference_start = Instant::now();
        let logits = model.forward(&input_tensor)?;
        let inference_time = inference_start.elapsed().as_secs_f64() * 1000.0;

        // Post-process to get predictions
        let predictions = self.postprocess_logits(&logits, top_k)?;

        // Record stats
        self.stats.record(inference_time, preprocess_time);

        Ok(ClassificationResult {
            predictions,
            inference_time_ms: inference_time,
            preprocess_time_ms: preprocess_time,
            input_dimensions: original_dims,
            model: model.name().to_string(),
        })
    }

    /// Preprocess image bytes into a tensor
    fn preprocess_image(&self, bytes: &[u8]) -> Result<(Tensor, (usize, usize))> {
        // Load image
        let img = image::load_from_memory(bytes)?;
        let (orig_width, orig_height) = img.dimensions();

        // Resize to model input size
        let resized = img.resize_exact(
            self.config.image_size as u32,
            self.config.image_size as u32,
            image::imageops::FilterType::Lanczos3,
        );

        // Convert to RGB and normalize
        let rgb = resized.to_rgb8();
        let (width, height) = rgb.dimensions();

        // Convert to CHW format with ImageNet normalization
        // mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]
        let mean = [0.485, 0.456, 0.406];
        let std = [0.229, 0.224, 0.225];

        let mut data = vec![0.0f32; 3 * (width as usize) * (height as usize)];

        for (x, y, pixel) in rgb.enumerate_pixels() {
            let idx_base = (y as usize) * (width as usize) + (x as usize);
            for c in 0..3 {
                let value = f32::from(pixel[c]) / 255.0;
                let normalized = (value - mean[c]) / std[c];
                let idx = c * (width as usize) * (height as usize) + idx_base;
                data[idx] = normalized;
            }
        }

        // Create tensor [1, C, H, W]
        let tensor = Tensor::from_data(
            data,
            vec![1, 3, height as usize, width as usize],
            &self.device,
        )?;

        Ok((tensor, (orig_width as usize, orig_height as usize)))
    }

    /// Post-process logits to get top-K predictions
    fn postprocess_logits(&self, logits: &Tensor, top_k: usize) -> Result<Vec<Prediction>> {
        // Get logits as flat vector
        let logits_data = logits.to_vec()?;

        // Apply softmax
        let max_logit = logits_data
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);
        let exp_sum: f32 = logits_data.iter().map(|&x| (x - max_logit).exp()).sum();
        let probs: Vec<f32> = logits_data
            .iter()
            .map(|&x| (x - max_logit).exp() / exp_sum)
            .collect();

        // Get top-K indices
        let mut indexed: Vec<(usize, f32)> = probs.into_iter().enumerate().collect();
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        let predictions = indexed
            .into_iter()
            .take(top_k)
            .map(|(idx, conf)| Prediction {
                class_idx: idx,
                label: get_label(idx).to_string(),
                confidence: conf,
            })
            .collect();

        Ok(predictions)
    }

    /// Get current status
    pub fn status(&self) -> ClassifierStatus {
        let metrics = self.stats.get_metrics();
        let model_guard = self.model.read().unwrap();

        ClassifierStatus {
            initialized: model_guard.is_some(),
            model: model_guard
                .as_ref()
                .map(|m| format!("{} ({})", self.config.architecture.display_name(), m.name())),
            device: self.device_name(),
            inference_count: metrics.total_inferences,
            avg_inference_time_ms: metrics.avg_inference_ms,
        }
    }

    /// Get performance metrics
    pub fn metrics(&self) -> ClassifierMetrics {
        self.stats.get_metrics()
    }

    /// Reset the classifier
    pub fn reset(&self) {
        *self.model.write().unwrap() = None;
    }
}

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

    #[test]
    fn test_classifier_creation() {
        let config = ClassifierConfig::default();
        let classifier = ImageClassifier::new(config).unwrap();
        assert!(!classifier.is_initialized());
    }

    #[test]
    fn test_device_selection() {
        let device = ImageClassifier::select_device(false);
        matches!(device, Device::Cpu(_));
    }
}