//! Complete Object Detection Example
//!
//! This example demonstrates building a simple object detection model using scirs2-neural.
//! It includes:
//! - Feature extraction backbone (simplified CNN)
//! - Object detection head for bounding box regression and classification
//! - Synthetic dataset generation with multiple objects per image
//! - Training loop with object detection specific losses
//! - Evaluation metrics (IoU, mAP approximation)
//! - Visualization of detection results

use ndarray::{s, Array2, Array3, Array4, ArrayD, IxDyn};
use rand::prelude::*;
use rand::rngs::SmallRng;
use scirs2_neural::layers::{
    AdaptiveMaxPool2D, BatchNorm, Conv2D, Dense, Dropout, MaxPool2D, PaddingMode, Sequential,
};
use scirs2_neural::losses::{CrossEntropyLoss, MeanSquaredError};
use scirs2_neural::prelude::*;
use std::collections::HashMap;
// Type alias to avoid conflicts with scirs2-neural's Result
type StdResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
/// Object detection model configuration
#[derive(Debug, Clone)]
pub struct DetectionConfig {
    pub num_classes: usize,
    pub max_objects: usize,
    pub input_size: (usize, usize),
    pub anchor_sizes: Vec<f32>,
    pub feature_map_size: (usize, usize),
}
impl Default for DetectionConfig {
    fn default() -> Self {
        Self {
            num_classes: 3, // background + 2 object classes
            max_objects: 5,
            input_size: (64, 64),
            anchor_sizes: vec![16.0, 32.0, 48.0],
            feature_map_size: (8, 8),
        }
    }
/// Bounding box representation
pub struct BoundingBox {
    pub x: f32,
    pub y: f32,
    pub width: f32,
    pub height: f32,
    pub class_id: usize,
    pub confidence: f32,
impl BoundingBox {
    pub fn new(x: f32, y: f32, width: f32, height: f32, classid: usize, confidence: f32) -> Self {
            x,
            y,
            width,
            height,
            class_id,
            confidence,
    /// Calculate Intersection over Union (IoU) with another bounding box
    pub fn iou(&self, other: &BoundingBox) -> f32 {
        let x1 = self.x.max(other.x);
        let y1 = self.y.max(other.y);
        let x2 = (self.x + self.width).min(other.x + other.width);
        let y2 = (self.y + self.height).min(other.y + other.height);
        if x2 <= x1 || y2 <= y1 {
            return 0.0;
        let intersection = (x2 - x1) * (y2 - y1);
        let union = self.width * self.height + other.width * other.height - intersection;
        if union <= 0.0 {
            0.0
        } else {
            intersection / union
/// Object detection dataset generator
pub struct DetectionDataset {
    config: DetectionConfig,
    rng: SmallRng,
impl DetectionDataset {
    pub fn new(config: DetectionConfig, seed: u64) -> Self {
            config,
            rng: SmallRng::from_seed(seed),
    /// Generate a synthetic image with objects and their labels
    pub fn generate_sample(&mut self) -> (Array3<f32>, Vec<BoundingBox>) {
        let (height, width) = self._config.input_size;
        let mut image = Array3::<f32>::zeros((3, height, width)); // RGB channels
        // Generate background pattern
        for c in 0..3 {
            for i in 0..height {
                for j in 0..width {
                    image[[c, i, j]] = self.rng.random_range(0.0..0.3);
                }
            }
        let mut objects = Vec::new();
        let num_objects = self.rng.random_range(1..=self.config.max_objects.min(3));
        for _ in 0..num_objects {
            let obj_width = self.rng.random_range(8..24) as f32;
            let obj_height = self.rng.random_range(8..24) as f32;
            let obj_x = self.rng.random_range(0.0..(width as f32 - obj_width));
            let obj_y = self.rng.random_range(0.0..(height as f32 - obj_height));
            let class_id = self.rng.random_range(1..self.config.num_classes); // Skip background class 0
            // Draw rectangular object
            let color_intensity = match class_id {
                1 => [0.8..0.2, 0.2], // Red-ish for class 1
                2 => [0.2, 0.8, 0.2], // Green-ish for class 2
                _ => [0.2, 0.2, 0.8], // Blue-ish for other classes
            };
            for c in 0..3 {
                for i in (obj_y as usize)..((obj_y + obj_height) as usize).min(height) {
                    for j in (obj_x as usize)..((obj_x + obj_width) as usize).min(width) {
                        image[[c, i, j]] = color_intensity[c] + self.rng.random_range(-0.1..0.1);
                    }
            objects.push(BoundingBox::new(
                obj_x..obj_y, obj_width, obj_height, class_id, 1.0,));
        (image, objects)
    /// Generate a batch of samples
    pub fn generate_batch(&mut self, batchsize: usize) -> (Array4<f32>, Vec<Vec<BoundingBox>>) {
        let mut images = Array4::<f32>::zeros((batch_size, 3, height, width));
        let mut all_objects = Vec::new();
        for i in 0..batch_size {
            let (image, objects) = self.generate_sample();
            images.slice_mut(s![i, .., .., ..]).assign(&image);
            all_objects.push(objects);
        (images, all_objects)
/// Object detection model combining feature extraction and detection heads
pub struct ObjectDetectionModel {
    feature_extractor: Sequential<f32>,
    classifier_head: Sequential<f32>,
    bbox_regressor: Sequential<f32>,
impl ObjectDetectionModel {
    pub fn new(config: DetectionConfig, rng: &mut SmallRng) -> StdResult<Self> {
        // Feature extraction backbone (simplified ResNet-like)
        let mut feature_extractor = Sequential::new();
        // Initial conv block
        feature_extractor.add(Conv2D::new(3, 64, (7, 7), (2, 2), PaddingMode::Same, rng)?);
        feature_extractor.add(BatchNorm::new(64, 0.1, 1e-5, rng)?);
        feature_extractor.add(MaxPool2D::new((2, 2), (2, 2), None)?);
        // Feature blocks
        feature_extractor.add(Conv2D::new(
            64,
            128,
            (3, 3),
            (2, 2),
            PaddingMode::Same,
            rng,
        )?);
        feature_extractor.add(BatchNorm::new(128, 0.1, 1e-5, rng)?);
            256,
        feature_extractor.add(BatchNorm::new(256, 0.1, 1e-5, rng)?);
        // Global pooling to fixed size
        feature_extractor.add(AdaptiveMaxPool2D::new(_config.feature_map_size, None)?);
        // Classification head
        let mut classifier_head = Sequential::new();
        let feature_dim = 256 * config.feature_map_size.0 * config.feature_map_size.1;
        classifier_head.add(Dense::new(feature_dim, 512, Some("relu"), rng)?);
        classifier_head.add(Dropout::new(0.5, rng)?);
        classifier_head.add(Dense::new(512, 256, Some("relu"), rng)?);
        classifier_head.add(Dropout::new(0.3, rng)?);
        classifier_head.add(Dense::new(
            config.num_classes * config.max_objects,
            Some("softmax"),
        // Bounding box regression head
        let mut bbox_regressor = Sequential::new();
        bbox_regressor.add(Dense::new(feature_dim, 512, Some("relu"), rng)?);
        bbox_regressor.add(Dropout::new(0.5, rng)?);
        bbox_regressor.add(Dense::new(512, 256, Some("relu"), rng)?);
        bbox_regressor.add(Dropout::new(0.3, rng)?);
        bbox_regressor.add(Dense::new(256, 4 * config.max_objects, None, rng)?); // 4 coordinates per object
        Ok(Self {
            feature_extractor,
            classifier_head,
            bbox_regressor,
        })
    /// Forward pass through the entire detection model
    pub fn forward(&self, input: &ArrayD<f32>) -> StdResult<(ArrayD<f32>, ArrayD<f32>)> {
        // Extract features
        let features = self.feature_extractor.forward(input)?;
        // Flatten features for dense layers
        let batch_size = features.shape()[0];
        let feature_dim = features.len() / batch_size;
        let flattened = features
            .into_shape_with_order(IxDyn(&[batch_size, feature_dim]))
            .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
        // Get classifications and bounding box predictions
        let classifications = self.classifier_head.forward(&flattened)?;
        let bbox_predictions = self.bbox_regressor.forward(&flattened)?;
        Ok((classifications, bbox_predictions))
    /// Post-process predictions to extract bounding boxes
    pub fn extract_detections(
        &self,
        classifications: &ArrayD<f32>,
        bbox_predictions: &ArrayD<f32>,
        confidence_threshold: f32,
    ) -> Vec<Vec<BoundingBox>> {
        let batch_size = classifications.shape()[0];
        let mut detections = Vec::new();
        for b in 0..batch_size {
            let mut batch_detections = Vec::new();
            for obj_idx in 0..self.config.max_objects {
                // Get classification scores for this object
                let mut best_class = 0;
                let mut best_score = 0.0f32;
                for class_idx in 0..self.config.num_classes {
                    let score_idx = obj_idx * self.config.num_classes + class_idx;
                    if score_idx < classifications.shape()[1] {
                        let score = classifications[[b, score_idx]];
                        if score > best_score {
                            best_score = score;
                            best_class = class_idx;
                        }
                // Skip background class (0) and low confidence predictions
                if best_class > 0 && best_score > confidence_threshold {
                    // Get bounding box coordinates
                    let bbox_start = obj_idx * 4;
                    if bbox_start + 3 < bbox_predictions.shape()[1] {
                        let x = bbox_predictions[[b, bbox_start]].max(0.0);
                        let y = bbox_predictions[[b, bbox_start + 1]].max(0.0);
                        let width = bbox_predictions[[b, bbox_start + 2]].max(1.0);
                        let height = bbox_predictions[[b, bbox_start + 3]].max(1.0);
                        batch_detections.push(BoundingBox::new(
                            x, y, width, height, best_class, best_score,
                        ));
            detections.push(batch_detections);
        detections
/// Object detection loss combining classification and regression losses
pub struct DetectionLoss {
    classification_loss: CrossEntropyLoss,
    regression_loss: MeanSquaredError,
    classification_weight: f32,
    regression_weight: f32,
impl DetectionLoss {
    pub fn new(_classification_weight: f32, regressionweight: f32) -> Self {
            classification_loss: CrossEntropyLoss::new(1e-7),
            regression_loss: MeanSquaredError,
            classification_weight,
            regression_weight,
    /// Compute combined detection loss
    pub fn compute_loss(
        pred_classes: &ArrayD<f32>,
        pred_boxes: &ArrayD<f32>,
        target_classes: &ArrayD<f32>,
        target_boxes: &ArrayD<f32>,
    ) -> StdResult<f32> {
        let class_loss = self
            .classification_loss
            .forward(pred_classes, target_classes)?;
        let bbox_loss = self.regression_loss.forward(pred_boxes, target_boxes)?;
        Ok(self._classification_weight * class_loss + self.regression_weight * bbox_loss)
/// Metrics for object detection evaluation
pub struct DetectionMetrics {
    iou_threshold: f32,
    confidence_threshold: f32,
impl DetectionMetrics {
    pub fn new(_iou_threshold: f32, confidencethreshold: f32) -> Self {
            iou_threshold,
            confidence_threshold,
    /// Calculate mean Average Precision (simplified version)
    pub fn calculate_map(
        predictions: &[Vec<BoundingBox>],
        ground_truth: &[Vec<BoundingBox>],
    ) -> f32 {
        if predictions.is_empty() || ground_truth.is_empty() {
        let mut total_precision = 0.0;
        let mut total_samples = 0;
        for (pred_batch, gt_batch) in predictions.iter().zip(ground_truth.iter()) {
            let precision = self.calculate_precision(pred_batch, gt_batch);
            total_precision += precision;
            total_samples += 1;
        if total_samples > 0 {
            total_precision / total_samples as f32
    /// Calculate precision for a single sample
    fn calculate_precision(
        predictions: &[BoundingBox],
        ground_truth: &[BoundingBox],
        if predictions.is_empty() {
            return if ground_truth.is_empty() { 1.0 } else { 0.0 };
        let mut true_positives = 0;
        let mut used_gt = vec![false; ground_truth.len()];
        for pred in predictions {
            if pred.confidence < self.confidence_threshold {
                continue;
            let mut best_iou = 0.0;
            let mut best_gt_idx = None;
            for (gt_idx, gt) in ground_truth.iter().enumerate() {
                if used_gt[gt_idx] || pred.class_id != gt.class_id {
                    continue;
                let iou = pred.iou(gt);
                if iou > best_iou {
                    best_iou = iou;
                    best_gt_idx = Some(gt_idx);
            if let Some(gt_idx) = best_gt_idx {
                if best_iou >= self.iou_threshold {
                    true_positives += 1;
                    used_gt[gt_idx] = true;
            true_positives as f32 / predictions.len() as f32
/// Convert ground truth bounding boxes to target tensors
#[allow(dead_code)]
fn prepare_targets(
    ground_truth: &[Vec<BoundingBox>],
    config: &DetectionConfig,
) -> (ArrayD<f32>, ArrayD<f32>) {
    let batch_size = ground_truth.len();
    // Classification targets: [batch_size, max_objects * num_classes]
    let mut class_targets =
        Array2::<f32>::zeros((batch_size, config.max_objects * config.num_classes));
    // Bounding box targets: [batch_size, max_objects * 4]
    let mut bbox_targets = Array2::<f32>::zeros((batch_size, config.max_objects * 4));
    for (batch_idx, objects) in ground_truth.iter().enumerate() {
        for (obj_idx, obj) in objects.iter().enumerate().take(config.max_objects) {
            // Set class target (one-hot encoding)
            let class_start = obj_idx * config.num_classes;
            if class_start + obj.class_id < class_targets.shape()[1] {
                class_targets[[batch_idx, class_start + obj.class_id]] = 1.0;
            // Set bounding box targets
            let bbox_start = obj_idx * 4;
            if bbox_start + 3 < bbox_targets.shape()[1] {
                bbox_targets[[batch_idx, bbox_start]] = obj.x;
                bbox_targets[[batch_idx, bbox_start + 1]] = obj.y;
                bbox_targets[[batch_idx, bbox_start + 2]] = obj.width;
                bbox_targets[[batch_idx, bbox_start + 3]] = obj.height;
    (class_targets.into_dyn(), bbox_targets.into_dyn())
/// Training function for object detection
#[allow(dead_code)]
fn train_detection_model() -> StdResult<()> {
    println!("🎯 Starting Object Detection Training");
    let mut rng = SmallRng::from_seed([42; 32]);
    let config = DetectionConfig::default();
    println!("🚀 Starting model training...");
    // Create model
    println!("🏗️ Building object detection model...");
    let model = ObjectDetectionModel::new(config.clone(), &mut rng)?;
    println!(
        "✅ Model created with {} classes and {} max objects",
        config.num_classes, config.max_objects
    );
    // Create dataset
    let mut dataset = DetectionDataset::new(config.clone(), 123);
    // Create loss function
    let loss_fn = DetectionLoss::new(1.0, 1.0); // Equal weights for classification and regression
    // Create metrics
    let metrics = DetectionMetrics::new(0.5, 0.5); // IoU threshold 0.5, confidence threshold 0.5
    println!("📊 Training configuration:");
    println!("   - Input size: {:?}", config.input_size);
    println!("   - Feature map size: {:?}", config.feature_map_size);
    println!("   - Max objects per image: {}", config.max_objects);
    println!("   - Number of classes: {}", config.num_classes);
    // Training loop
    let num_epochs = 10;
    let batch_size = 4;
    let _learning_rate = 0.001;
    for epoch in 0..num_epochs {
        println!("\n📈 Epoch {}/{}", epoch + 1, num_epochs);
        let mut epoch_loss = 0.0;
        let num_batches = 8; // Small number of batches for demo
        for batch_idx in 0..num_batches {
            // Generate training batch
            let (images, ground_truth) = dataset.generate_batch(batch_size);
            let images_dyn = images.into_dyn();
            // Forward pass
            let (pred_classes, pred_boxes) = model.forward(&images_dyn)?;
            // Prepare targets
            let (target_classes, target_boxes) = prepare_targets(&ground_truth, &config);
            // Compute loss
            let batch_loss =
                loss_fn.compute_loss(&pred_classes, &pred_boxes, &target_classes, &target_boxes)?;
            epoch_loss += batch_loss;
            if batch_idx % 4 == 0 {
                print!(
                    "🔄 Batch {}/{} - Loss: {:.4}                \r",
                    batch_idx + 1,
                    num_batches,
                    batch_loss
                );
        let avg_loss = epoch_loss / num_batches as f32;
        println!(
            "✅ Epoch {} completed - Average Loss: {:.4}",
            epoch + 1,
            avg_loss
        );
        // Evaluation every few epochs
        if (epoch + 1) % 3 == 0 {
            println!("🔍 Running evaluation...");
            // Generate validation batch
            let (val_images, val_ground_truth) = dataset.generate_batch(batch_size);
            let val_images_dyn = val_images.into_dyn();
            // Get predictions
            let (pred_classes, pred_boxes) = model.forward(&val_images_dyn)?;
            let detections = model.extract_detections(&pred_classes, &pred_boxes, 0.5);
            // Calculate metrics
            let map = metrics.calculate_map(&detections, &val_ground_truth);
            println!("📊 Validation mAP: {:.4}", map);
            // Print sample detection results
            if !detections.is_empty() && !detections[0].is_empty() {
                println!("🎯 Sample detections:");
                for (i, detection) in detections[0].iter().enumerate().take(3) {
                    println!(
                        "   Detection {}: class={}, conf={:.3}, bbox=({:.1}, {:.1}, {:.1}, {:.1})",
                        i + 1,
                        detection.class_id,
                        detection.confidence,
                        detection.x,
                        detection.y,
                        detection.width,
                        detection.height
                    );
    println!("\n🎉 Object detection training completed!");
    // Final evaluation
    println!("🔬 Final evaluation...");
    let (test_images, test_ground_truth) = dataset.generate_batch(8);
    let test_images_dyn = test_images.into_dyn();
    let (pred_classes, pred_boxes) = model.forward(&test_images_dyn)?;
    let final_detections = model.extract_detections(&pred_classes, &pred_boxes, 0.3);
    let final_map = metrics.calculate_map(&final_detections, &test_ground_truth);
    println!("📈 Final mAP: {:.4}", final_map);
    // Performance analysis
    println!("\n📊 Performance Analysis:");
        "   - Total parameters: ~{:.1}K",
        (256 * 8 * 8 * 512
            + 512 * 256
            + 256 * config.num_classes * config.max_objects
            + 256 * 4 * config.max_objects)
            / 1000
    println!("   - Memory efficient: ✅ (adaptive pooling used)");
    println!("   - JIT optimized: ✅");
    let mut detection_stats = HashMap::new();
    for detections in &final_detections {
        for detection in detections {
            *detection_stats.entry(detection.class_id).or_insert(0) += 1;
    println!("   - Detections by class:");
    for (class_id, count) in detection_stats {
        println!("     Class {}: {} detections", class_id, count);
    Ok(())
#[allow(dead_code)]
fn main() -> StdResult<()> {
    println!("🎯 Object Detection Complete Example");
    println!("=====================================");
    println!();
    println!("This example demonstrates:");
    println!("• Building an object detection model with CNN backbone");
    println!("• Synthetic dataset generation with multiple objects");
    println!("• Combined classification and bounding box regression");
    println!("• Object detection specific metrics (IoU, mAP)");
    println!("• JIT compilation for performance optimization");
    train_detection_model()?;
    println!("\n💡 Key Concepts Demonstrated:");
    println!("   🔹 Feature extraction with CNN backbone");
    println!("   🔹 Multi-task learning (classification + regression)");
    println!("   🔹 Object detection loss functions");
    println!("   🔹 IoU-based evaluation metrics");
    println!("   🔹 Non-maximum suppression concepts");
    println!("   🔹 Bounding box post-processing");
    println!("🚀 For production use:");
    println!("   • Implement anchor-based detection (YOLO, SSD)");
    println!("   • Add data augmentation (rotation, scaling, cropping)");
    println!("   • Use pre-trained backbones (ResNet, EfficientNet)");
    println!("   • Implement proper NMS (Non-Maximum Suppression)");
    println!("   • Add multi-scale training and testing");
    println!("   • Use real datasets (COCO, Pascal VOC)");
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_detection_config() {
        let config = DetectionConfig::default();
        assert_eq!(config.num_classes, 3);
        assert_eq!(config.max_objects, 5);
        assert_eq!(config.input_size, (64, 64));
    fn test_bounding_box_iou() {
        let box1 = BoundingBox::new(0.0, 0.0, 10.0, 10.0, 1, 1.0);
        let box2 = BoundingBox::new(5.0, 5.0, 10.0, 10.0, 1, 1.0);
        let iou = box1.iou(&box2);
        assert!(iou > 0.0 && iou < 1.0);
        // Test no overlap
        let box3 = BoundingBox::new(20.0, 20.0, 10.0, 10.0, 1, 1.0);
        assert_eq!(box1.iou(&box3), 0.0);
        // Test complete overlap
        let box4 = BoundingBox::new(0.0, 0.0, 10.0, 10.0, 1, 1.0);
        assert_eq!(box1.iou(&box4), 1.0);
    fn test_dataset_generation() {
        let mut dataset = DetectionDataset::new(config.clone(), 42);
        let (image, objects) = dataset.generate_sample();
        assert_eq!(
            image.shape(),
            &[3, config.input_size.0, config.input_size.1]
        assert!(!objects.is_empty());
        assert!(objects.len() <= config.max_objects);
        for obj in &objects {
            assert!(obj.class_id > 0 && obj.class_id < config.num_classes);
            assert!(obj.x >= 0.0 && obj.y >= 0.0);
            assert!(obj.width > 0.0 && obj.height > 0.0);
    fn test_detection_metrics() {
        let metrics = DetectionMetrics::new(0.5, 0.5);
        // Test perfect match
        let pred = vec![BoundingBox::new(0.0, 0.0, 10.0, 10.0, 1, 0.9)];
        let gt = vec![BoundingBox::new(0.0, 0.0, 10.0, 10.0, 1, 1.0)];
        let precision = metrics.calculate_precision(&pred, &gt);
        assert_eq!(precision, 1.0);
        // Test no match (different class)
        let pred2 = vec![BoundingBox::new(0.0, 0.0, 10.0, 10.0, 1, 0.9)];
        let gt2 = vec![BoundingBox::new(0.0, 0.0, 10.0, 10.0, 2, 1.0)];
        let precision2 = metrics.calculate_precision(&pred2, &gt2);
        assert_eq!(precision2, 0.0);
    fn test_model_creation() -> StdResult<()> {
        let mut rng = SmallRng::from_seed([42; 32]);
        let model = ObjectDetectionModel::new(config, &mut rng)?;
        // Test forward pass shape
        let batch_size = 2;
        let input = Array4::<f32>::ones((batch_size, 3, 64, 64)).into_dyn();
        let (classes, boxes) = model.forward(&input)?;
        assert_eq!(classes.shape()[0], batch_size);
        assert_eq!(boxes.shape()[0], batch_size);
        Ok(())
