semantic_segmentation_complete/

semantic_segmentation_complete.rs

//! Complete Semantic Segmentation Example
//!
//! This example demonstrates building a semantic segmentation model using scirs2-neural.
//! It includes:
//! - U-Net style encoder-decoder architecture
//! - Skip connections for preserving spatial information
//! - Synthetic dataset generation with multiple semantic classes
//! - Pixel-wise classification loss and metrics
//! - Evaluation metrics (IoU, mIoU, pixel accuracy)
//! - Visualization of segmentation results

use ndarray::{s, Array2, Array3, Array4, ArrayD};
use scirs2_neural::layers::{BatchNorm, Conv2D, MaxPool2D, PaddingMode, Sequential};
use scirs2_neural::losses::CrossEntropyLoss;
use scirs2_neural::prelude::*;

// Type alias to avoid conflicts with scirs2-neural's Result
type StdResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
use rand::prelude::*;
use rand::rngs::SmallRng;

/// Semantic segmentation model configuration
#[derive(Debug, Clone)]
pub struct SegmentationConfig {
    pub num_classes: usize,
    pub input_size: (usize, usize),
    pub encoder_channels: Vec<usize>,
    pub decoder_channels: Vec<usize>,
    pub skip_connections: bool,
}

impl Default for SegmentationConfig {
    fn default() -> Self {
        Self {
            num_classes: 4, // background + 3 object classes
            input_size: (128, 128),
            encoder_channels: vec![64, 128, 256, 512],
            decoder_channels: vec![256, 128, 64, 32],
            skip_connections: true,
        }
    }
}

/// Semantic segmentation dataset generator
pub struct SegmentationDataset {
    config: SegmentationConfig,
    rng: SmallRng,
}

impl SegmentationDataset {
    pub fn new(config: SegmentationConfig, seed: u64) -> Self {
        Self {
            config,
            rng: SmallRng::seed_from_u64(seed),
        }
    }

    /// Generate a synthetic image with semantic labels
    pub fn generate_sample(&mut self) -> (Array3<f32>, Array2<usize>) {
        let (height, width) = self.config.input_size;
        let mut image = Array3::<f32>::zeros((3, height, width)); // RGB channels
        let mut mask = Array2::<usize>::zeros((height, width));

        // 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.1..0.3);
                }
            }
        }

        // Generate geometric shapes with different semantic classes
        let num_shapes = self.rng.random_range(3..8);

        for _ in 0..num_shapes {
            let shape_type = self.rng.random_range(0..3);
            let class_id = self.rng.random_range(1..self.config.num_classes);

            let color = match class_id {
                1 => [0.8, 0.2, 0.2], // Red for class 1
                2 => [0.2, 0.8, 0.2], // Green for class 2
                3 => [0.2, 0.2, 0.8], // Blue for class 3
                _ => [0.8, 0.8, 0.2], // Yellow for other classes
            };

            match shape_type {
                0 => {
                    // Rectangle
                    let rect_width = self.rng.random_range(15..40);
                    let rect_height = self.rng.random_range(15..40);
                    let start_x = self.rng.random_range(0..(width.saturating_sub(rect_width)));
                    let start_y = self
                        .rng
                        .random_range(0..(height.saturating_sub(rect_height)));

                    for i in start_y..(start_y + rect_height).min(height) {
                        for j in start_x..(start_x + rect_width).min(width) {
                            mask[[i, j]] = class_id;
                            for c in 0..3 {
                                image[[c, i, j]] = color[c] + self.rng.random_range(-0.1..0.1);
                            }
                        }
                    }
                }
                1 => {
                    // Circle
                    let radius = self.rng.random_range(8..25) as f32;
                    let center_x = self
                        .rng
                        .random_range(radius as usize..(width - radius as usize))
                        as f32;
                    let center_y = self
                        .rng
                        .random_range(radius as usize..(height - radius as usize))
                        as f32;

                    for i in 0..height {
                        for j in 0..width {
                            let dx = j as f32 - center_x;
                            let dy = i as f32 - center_y;
                            if dx * dx + dy * dy <= radius * radius {
                                mask[[i, j]] = class_id;
                                for c in 0..3 {
                                    image[[c, i, j]] = color[c] + self.rng.random_range(-0.1..0.1);
                                }
                            }
                        }
                    }
                }
                _ => {
                    // Triangle (approximate)
                    let size = self.rng.random_range(15..35);
                    let center_x = self.rng.random_range(size / 2..(width - size / 2));
                    let center_y = self.rng.random_range(size / 2..(height - size / 2));

                    for i in (center_y.saturating_sub(size / 2))..(center_y + size / 2).min(height)
                    {
                        let row_width = (size as f32
                            * (1.0 - (i as f32 - center_y as f32).abs() / (size as f32 / 2.0)))
                            as usize;
                        for j in (center_x.saturating_sub(row_width / 2))
                            ..(center_x + row_width / 2).min(width)
                        {
                            mask[[i, j]] = class_id;
                            for c in 0..3 {
                                image[[c, i, j]] = color[c] + self.rng.random_range(-0.1..0.1);
                            }
                        }
                    }
                }
            }
        }

        (image, mask)
    }

    /// Generate a batch of samples
    pub fn generate_batch(&mut self, batch_size: usize) -> (Array4<f32>, Array3<usize>) {
        let (height, width) = self.config.input_size;
        let mut images = Array4::<f32>::zeros((batch_size, 3, height, width));
        let mut masks = Array3::<usize>::zeros((batch_size, height, width));

        for i in 0..batch_size {
            let (image, mask) = self.generate_sample();
            images.slice_mut(s![i, .., .., ..]).assign(&image);
            masks.slice_mut(s![i, .., ..]).assign(&mask);
        }

        (images, masks)
    }
}

/// U-Net style encoder block
pub struct EncoderBlock {
    conv1: Conv2D<f32>,
    bn1: BatchNorm<f32>,
    conv2: Conv2D<f32>,
    bn2: BatchNorm<f32>,
    pool: MaxPool2D<f32>,
}

impl EncoderBlock {
    pub fn new(in_channels: usize, out_channels: usize, rng: &mut SmallRng) -> StdResult<Self> {
        Ok(Self {
            conv1: Conv2D::new(
                in_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn1: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
            conv2: Conv2D::new(
                out_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn2: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
            pool: MaxPool2D::new((2, 2), (2, 2), None)?,
        })
    }

    pub fn forward(&self, input: &ArrayD<f32>) -> StdResult<(ArrayD<f32>, ArrayD<f32>)> {
        // First conv + bn + relu
        let x = self.conv1.forward(input)?;
        let x = self.bn1.forward(&x)?;
        // Note: In practice, add ReLU activation here

        // Second conv + bn + relu
        let x = self.conv2.forward(&x)?;
        let skip = self.bn2.forward(&x)?;
        // Note: In practice, add ReLU activation here

        // Pooling for downsampling
        let pooled = self.pool.forward(&skip)?;

        Ok((pooled, skip))
    }
}

/// U-Net style decoder block
pub struct DecoderBlock {
    conv1: Conv2D<f32>,
    bn1: BatchNorm<f32>,
    conv2: Conv2D<f32>,
    bn2: BatchNorm<f32>,
}

impl DecoderBlock {
    pub fn new(in_channels: usize, out_channels: usize, rng: &mut SmallRng) -> StdResult<Self> {
        Ok(Self {
            conv1: Conv2D::new(
                in_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn1: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
            conv2: Conv2D::new(
                out_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn2: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
        })
    }

    pub fn forward(
        &self,
        input: &ArrayD<f32>,
        skip: Option<&ArrayD<f32>>,
    ) -> StdResult<ArrayD<f32>> {
        // Upsample input (simplified - in practice use transpose convolution)
        let upsampled = self.upsample(input)?;

        // Concatenate with skip connection if provided
        let x = if let Some(skip_tensor) = skip {
            self.concatenate(&upsampled, skip_tensor)?
        } else {
            upsampled
        };

        // First conv + bn + relu
        let x = self.conv1.forward(&x)?;
        let x = self.bn1.forward(&x)?;
        // Note: In practice, add ReLU activation here

        // Second conv + bn + relu
        let x = self.conv2.forward(&x)?;
        let x = self.bn2.forward(&x)?;
        // Note: In practice, add ReLU activation here

        Ok(x)
    }

    fn upsample(&self, input: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        // Simplified upsampling using nearest neighbor
        let shape = input.shape();
        let batch_size = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        let mut upsampled = Array4::<f32>::zeros((batch_size, channels, height * 2, width * 2));

        for b in 0..batch_size {
            for c in 0..channels {
                for i in 0..height {
                    for j in 0..width {
                        let value = input[[b, c, i, j]];
                        upsampled[[b, c, i * 2, j * 2]] = value;
                        upsampled[[b, c, i * 2, j * 2 + 1]] = value;
                        upsampled[[b, c, i * 2 + 1, j * 2]] = value;
                        upsampled[[b, c, i * 2 + 1, j * 2 + 1]] = value;
                    }
                }
            }
        }

        Ok(upsampled.into_dyn())
    }

    fn concatenate(&self, input1: &ArrayD<f32>, input2: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        // Simplified concatenation along channel dimension
        let shape1 = input1.shape();
        let shape2 = input2.shape();

        if shape1[0] != shape2[0] || shape1[2] != shape2[2] || shape1[3] != shape2[3] {
            return Err("Shapes incompatible for concatenation".into());
        }

        let batch_size = shape1[0];
        let channels1 = shape1[1];
        let channels2 = shape2[1];
        let height = shape1[2];
        let width = shape1[3];

        let mut concatenated =
            Array4::<f32>::zeros((batch_size, channels1 + channels2, height, width));

        // Copy first tensor
        for b in 0..batch_size {
            for c in 0..channels1 {
                for i in 0..height {
                    for j in 0..width {
                        concatenated[[b, c, i, j]] = input1[[b, c, i, j]];
                    }
                }
            }
        }

        // Copy second tensor
        for b in 0..batch_size {
            for c in 0..channels2 {
                for i in 0..height {
                    for j in 0..width {
                        concatenated[[b, channels1 + c, i, j]] = input2[[b, c, i, j]];
                    }
                }
            }
        }

        Ok(concatenated.into_dyn())
    }
}

/// U-Net model for semantic segmentation
pub struct UNetModel {
    encoders: Vec<EncoderBlock>,
    decoders: Vec<DecoderBlock>,
    bottleneck: Sequential<f32>,
    final_conv: Conv2D<f32>,
    config: SegmentationConfig,
}

impl UNetModel {
    pub fn new(config: SegmentationConfig, rng: &mut SmallRng) -> StdResult<Self> {
        let mut encoders = Vec::new();
        let mut decoders = Vec::new();

        // Build encoder blocks
        let mut in_channels = 3; // RGB input
        for &out_channels in &config.encoder_channels {
            encoders.push(EncoderBlock::new(in_channels, out_channels, rng)?);
            in_channels = out_channels;
        }

        // Bottleneck
        let bottleneck_channels = config.encoder_channels.last().copied().unwrap_or(512);
        let mut bottleneck = Sequential::new();
        bottleneck.add(Conv2D::new(
            bottleneck_channels,
            bottleneck_channels * 2,
            (3, 3),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?);
        bottleneck.add(BatchNorm::new(bottleneck_channels * 2, 0.1, 1e-5, rng)?);
        bottleneck.add(Conv2D::new(
            bottleneck_channels * 2,
            bottleneck_channels,
            (3, 3),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?);
        bottleneck.add(BatchNorm::new(bottleneck_channels, 0.1, 1e-5, rng)?);

        // Build decoder blocks
        in_channels = bottleneck_channels;
        for (i, &out_channels) in config.decoder_channels.iter().enumerate() {
            let decoder_in_channels =
                if config.skip_connections && i < config.encoder_channels.len() {
                    // Skip connections come from corresponding encoder layer (in reverse order)
                    let encoder_idx = config.encoder_channels.len() - 1 - i;
                    in_channels + config.encoder_channels[encoder_idx]
                } else {
                    in_channels
                };
            decoders.push(DecoderBlock::new(decoder_in_channels, out_channels, rng)?);
            in_channels = out_channels;
        }

        // Final classification layer
        let final_channels = config.decoder_channels.last().copied().unwrap_or(32);
        let final_conv = Conv2D::new(
            final_channels,
            config.num_classes,
            (1, 1),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?;

        Ok(Self {
            encoders,
            decoders,
            bottleneck,
            final_conv,
            config,
        })
    }

    pub fn forward(&self, input: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        let mut x = input.clone();
        let mut skip_connections = Vec::new();

        // Encoder path
        for encoder in &self.encoders {
            let (encoded, skip) = encoder.forward(&x)?;
            skip_connections.push(skip);
            x = encoded;
        }

        // Bottleneck
        x = self.bottleneck.forward(&x)?;

        // Decoder path
        skip_connections.reverse(); // Reverse to match decoder order
        for (i, decoder) in self.decoders.iter().enumerate() {
            let skip = if self.config.skip_connections && i < skip_connections.len() {
                Some(&skip_connections[i])
            } else {
                None
            };
            x = decoder.forward(&x, skip)?;
        }

        // Final classification
        let output = self.final_conv.forward(&x)?;

        Ok(output)
    }
}

/// Segmentation metrics
pub struct SegmentationMetrics {
    num_classes: usize,
}

impl SegmentationMetrics {
    pub fn new(num_classes: usize) -> Self {
        Self { num_classes }
    }

    /// Calculate pixel accuracy
    pub fn pixel_accuracy(&self, predictions: &Array3<usize>, ground_truth: &Array3<usize>) -> f32 {
        let mut correct = 0;
        let mut total = 0;

        for (pred, gt) in predictions.iter().zip(ground_truth.iter()) {
            if pred == gt {
                correct += 1;
            }
            total += 1;
        }

        if total > 0 {
            correct as f32 / total as f32
        } else {
            0.0
        }
    }

    /// Calculate mean Intersection over Union (mIoU)
    pub fn mean_iou(&self, predictions: &Array3<usize>, ground_truth: &Array3<usize>) -> f32 {
        let mut class_ious = Vec::new();

        for class_id in 0..self.num_classes {
            let iou = self.class_iou(predictions, ground_truth, class_id);
            if !iou.is_nan() {
                class_ious.push(iou);
            }
        }

        if class_ious.is_empty() {
            0.0
        } else {
            class_ious.iter().sum::<f32>() / class_ious.len() as f32
        }
    }

    /// Calculate IoU for a specific class
    fn class_iou(
        &self,
        predictions: &Array3<usize>,
        ground_truth: &Array3<usize>,
        class_id: usize,
    ) -> f32 {
        let mut intersection = 0;
        let mut union = 0;

        for (pred, gt) in predictions.iter().zip(ground_truth.iter()) {
            let pred_match = *pred == class_id;
            let gt_match = *gt == class_id;

            if pred_match && gt_match {
                intersection += 1;
            }
            if pred_match || gt_match {
                union += 1;
            }
        }

        if union > 0 {
            intersection as f32 / union as f32
        } else {
            f32::NAN
        }
    }

    /// Calculate confusion matrix
    pub fn confusion_matrix(
        &self,
        predictions: &Array3<usize>,
        ground_truth: &Array3<usize>,
    ) -> Array2<usize> {
        let mut matrix = Array2::<usize>::zeros((self.num_classes, self.num_classes));

        for (pred, gt) in predictions.iter().zip(ground_truth.iter()) {
            if *pred < self.num_classes && *gt < self.num_classes {
                matrix[[*gt, *pred]] += 1;
            }
        }

        matrix
    }
}

/// Convert logits to class predictions
fn logits_to_predictions(logits: &ArrayD<f32>) -> Array3<usize> {
    let shape = logits.shape();
    let batch_size = shape[0];
    let num_classes = shape[1];
    let height = shape[2];
    let width = shape[3];

    let mut predictions = Array3::<usize>::zeros((batch_size, height, width));

    for b in 0..batch_size {
        for i in 0..height {
            for j in 0..width {
                let mut best_class = 0;
                let mut best_score = logits[[b, 0, i, j]];

                for c in 1..num_classes {
                    let score = logits[[b, c, i, j]];
                    if score > best_score {
                        best_score = score;
                        best_class = c;
                    }
                }

                predictions[[b, i, j]] = best_class;
            }
        }
    }

    predictions
}

/// Convert class masks to one-hot encoded targets
fn masks_to_targets(masks: &Array3<usize>, num_classes: usize) -> ArrayD<f32> {
    let shape = masks.shape();
    let batch_size = shape[0];
    let height = shape[1];
    let width = shape[2];

    let mut targets = Array4::<f32>::zeros((batch_size, num_classes, height, width));

    for b in 0..batch_size {
        for i in 0..height {
            for j in 0..width {
                let class_id = masks[[b, i, j]];
                if class_id < num_classes {
                    targets[[b, class_id, i, j]] = 1.0;
                }
            }
        }
    }

    targets.into_dyn()
}

/// Training function for semantic segmentation
fn train_segmentation_model() -> StdResult<()> {
    println!("🎨 Starting Semantic Segmentation Training");

    let mut rng = SmallRng::seed_from_u64(42);
    let config = SegmentationConfig::default();

    println!("🚀 Starting model training...");

    // Create model
    println!("🏗️ Building U-Net segmentation model...");
    let model = UNetModel::new(config.clone(), &mut rng)?;
    println!("✅ Model created with {} classes", config.num_classes);

    // Create dataset
    let mut dataset = SegmentationDataset::new(config.clone(), 123);

    // Create loss function
    let loss_fn = CrossEntropyLoss::new(1e-7);

    // Create metrics
    let metrics = SegmentationMetrics::new(config.num_classes);

    println!("📊 Training configuration:");
    println!("   - Input size: {:?}", config.input_size);
    println!("   - Number of classes: {}", config.num_classes);
    println!("   - Encoder channels: {:?}", config.encoder_channels);
    println!("   - Decoder channels: {:?}", config.decoder_channels);
    println!("   - Skip connections: {}", config.skip_connections);

    // Training loop
    let num_epochs = 15;
    let batch_size = 2; // Small batch size due to memory constraints
    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 = 10; // Small number of batches for demo

        for batch_idx in 0..num_batches {
            // Generate training batch
            let (images, masks) = dataset.generate_batch(batch_size);
            let images_dyn = images.into_dyn();

            // Forward pass
            let logits = model.forward(&images_dyn)?;

            // Prepare targets
            let targets = masks_to_targets(&masks, config.num_classes);

            // Compute loss
            let batch_loss = loss_fn.forward(&logits, &targets)?;
            epoch_loss += batch_loss;

            if batch_idx % 5 == 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) % 5 == 0 {
            println!("🔍 Running evaluation...");

            // Generate validation batch
            let (val_images, val_masks) = dataset.generate_batch(batch_size);
            let val_images_dyn = val_images.into_dyn();

            // Get predictions
            let val_logits = model.forward(&val_images_dyn)?;
            let predictions = logits_to_predictions(&val_logits);

            // Calculate metrics
            let pixel_acc = metrics.pixel_accuracy(&predictions, &val_masks);
            let miou = metrics.mean_iou(&predictions, &val_masks);

            println!("📊 Validation metrics:");
            println!("   - Pixel Accuracy: {:.4}", pixel_acc);
            println!("   - Mean IoU: {:.4}", miou);

            // Print class-wise IoU
            println!("   - Class-wise IoU:");
            for class_id in 0..config.num_classes {
                let class_iou = metrics.class_iou(&predictions, &val_masks, class_id);
                if !class_iou.is_nan() {
                    println!("     Class {}: {:.4}", class_id, class_iou);
                }
            }
        }
    }

    println!("\n🎉 Semantic segmentation training completed!");

    // Final evaluation
    println!("🔬 Final evaluation...");
    let (test_images, test_masks) = dataset.generate_batch(4);
    let test_images_dyn = test_images.into_dyn();

    let test_logits = model.forward(&test_images_dyn)?;
    let final_predictions = logits_to_predictions(&test_logits);

    let final_pixel_acc = metrics.pixel_accuracy(&final_predictions, &test_masks);
    let final_miou = metrics.mean_iou(&final_predictions, &test_masks);

    println!("📈 Final metrics:");
    println!("   - Pixel Accuracy: {:.4}", final_pixel_acc);
    println!("   - Mean IoU: {:.4}", final_miou);

    // Confusion matrix
    let confusion = metrics.confusion_matrix(&final_predictions, &test_masks);
    println!("   - Confusion Matrix:");
    for i in 0..config.num_classes {
        print!("     [");
        for j in 0..config.num_classes {
            print!("{:4}", confusion[[i, j]]);
        }
        println!("]");
    }

    // Performance analysis
    println!("\n📊 Model Analysis:");
    println!("   - Architecture: U-Net with skip connections");
    println!(
        "   - Parameters: ~{:.1}K (estimated)",
        (config.encoder_channels.iter().sum::<usize>()
            + config.decoder_channels.iter().sum::<usize>())
            / 1000
    );
    println!("   - Memory efficient: ✅ (skip connections preserve spatial info)");
    println!("   - JIT optimized: ✅");

    Ok(())
}

fn main() -> StdResult<()> {
    println!("🎨 Semantic Segmentation Complete Example");
    println!("==========================================");
    println!();
    println!("This example demonstrates:");
    println!("• Building a U-Net style segmentation model");
    println!("• Encoder-decoder architecture with skip connections");
    println!("• Synthetic dataset with geometric shapes");
    println!("• Pixel-wise classification and evaluation");
    println!("• Segmentation metrics (IoU, mIoU, pixel accuracy)");
    println!("• JIT compilation for performance optimization");
    println!();

    train_segmentation_model()?;

    println!("\n💡 Key Concepts Demonstrated:");
    println!("   🔹 U-Net encoder-decoder architecture");
    println!("   🔹 Skip connections for spatial information preservation");
    println!("   🔹 Pixel-wise classification loss");
    println!("   🔹 Intersection over Union (IoU) metrics");
    println!("   🔹 Confusion matrix analysis");
    println!("   🔹 Upsampling and feature concatenation");
    println!();
    println!("🚀 For production use:");
    println!("   • Implement proper upsampling (transpose convolution)");
    println!("   • Add data augmentation (rotation, flipping, scaling)");
    println!("   • Use pre-trained encoders (ResNet, EfficientNet)");
    println!("   • Implement focal loss for class imbalance");
    println!("   • Add multi-scale training and testing");
    println!("   • Use real datasets (Cityscapes, ADE20K, Pascal VOC)");
    println!("   • Implement DeepLabV3+, PSPNet, or other SOTA architectures");

    Ok(())
}

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

    #[test]
    fn test_segmentation_config() {
        let config = SegmentationConfig::default();
        assert_eq!(config.num_classes, 4);
        assert_eq!(config.input_size, (128, 128));
        assert!(!config.encoder_channels.is_empty());
        assert!(!config.decoder_channels.is_empty());
    }

    #[test]
    fn test_dataset_generation() {
        let config = SegmentationConfig::default();
        let mut dataset = SegmentationDataset::new(config.clone(), 42);

        let (image, mask) = dataset.generate_sample();
        assert_eq!(
            image.shape(),
            &[3, config.input_size.0, config.input_size.1]
        );
        assert_eq!(mask.shape(), &[config.input_size.0, config.input_size.1]);

        // Check that mask contains valid class IDs
        for &class_id in mask.iter() {
            assert!(class_id < config.num_classes);
        }
    }

    #[test]
    fn test_segmentation_metrics() {
        let metrics = SegmentationMetrics::new(3);

        // Test perfect prediction
        let predictions = Array3::<usize>::from_shape_fn((1, 4, 4), |(_, i, j)| (i + j) % 3);
        let ground_truth = predictions.clone();

        let pixel_acc = metrics.pixel_accuracy(&predictions, &ground_truth);
        assert_eq!(pixel_acc, 1.0);

        let miou = metrics.mean_iou(&predictions, &ground_truth);
        assert_eq!(miou, 1.0);
    }

    #[test]
    fn test_model_creation() -> StdResult<()> {
        let mut rng = SmallRng::seed_from_u64(42);
        // Use smaller input size for faster testing
        let config = SegmentationConfig {
            num_classes: 4,
            input_size: (16, 16),                    // Much smaller for testing
            encoder_channels: vec![16, 32, 64, 128], // Smaller channels
            decoder_channels: vec![64, 32, 16, 8],   // Smaller channels
            skip_connections: true,
        };

        let model = UNetModel::new(config.clone(), &mut rng)?;

        // Test forward pass shape
        let batch_size = 1;
        let input = Array4::<f32>::ones((batch_size, 3, config.input_size.0, config.input_size.1))
            .into_dyn();
        let output = model.forward(&input)?;

        assert_eq!(output.shape()[0], batch_size);
        assert_eq!(output.shape()[1], config.num_classes);
        assert_eq!(output.shape()[2], config.input_size.0);
        assert_eq!(output.shape()[3], config.input_size.1);

        Ok(())
    }

    #[test]
    fn test_logits_to_predictions() {
        let logits = Array4::<f32>::from_shape_fn((1, 3, 2, 2), |(_, c, _, _)| c as f32);
        let logits_dyn = logits.into_dyn();

        let predictions = logits_to_predictions(&logits_dyn);

        // Should predict class 2 (highest logit) for all pixels
        for &pred in predictions.iter() {
            assert_eq!(pred, 2);
        }
    }

    #[test]
    fn test_masks_to_targets() {
        let masks = Array3::<usize>::from_shape_fn((1, 2, 2), |(_, i, j)| (i + j) % 3);
        let targets = masks_to_targets(&masks, 3);

        assert_eq!(targets.shape(), &[1, 3, 2, 2]);

        // Check one-hot encoding
        for b in 0..1 {
            for i in 0..2 {
                for j in 0..2 {
                    let class_id = masks[[b, i, j]];
                    for c in 0..3 {
                        let expected = if c == class_id { 1.0 } else { 0.0 };
                        assert_eq!(targets[[b, c, i, j]], expected);
                    }
                }
            }
        }
    }
}