src/robustness_analysis.py

"""

Robustness Analysis and Uncertainty Estimation for Brain Tumor Segmentation

"""

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import json
import os
from sklearn.metrics import confusion_matrix, classification_report
import pandas as pd


class RobustnessAnalyzer:
    """

    Analyze model robustness to various perturbations and corruptions

    """
    
    def __init__(self, model, input_shape=(224, 224, 3)):
        """

        Initialize robustness analyzer

        

        Args:

            model: Trained segmentation model

            input_shape: Shape of input images

        """
        self.model = model
        self.input_shape = input_shape
        self.results = {}
    
    def add_gaussian_noise(self, image, std=0.1):
        """Add Gaussian noise to image"""
        noise = np.random.normal(0, std, image.shape)
        return np.clip(image + noise, 0, 1)
    
    def add_salt_pepper_noise(self, image, salt_prob=0.01, pepper_prob=0.01):
        """Add salt and pepper noise to image"""
        noisy_image = image.copy()
        
        # Salt noise
        salt_mask = np.random.random(image.shape) < salt_prob
        noisy_image[salt_mask] = 1
        
        # Pepper noise
        pepper_mask = np.random.random(image.shape) < pepper_prob
        noisy_image[pepper_mask] = 0
        
        return noisy_image
    
    def add_gaussian_blur(self, image, kernel_size=3):
        """Apply Gaussian blur to image"""
        from scipy.ndimage import gaussian_filter
        return gaussian_filter(image, sigma=kernel_size/3)
    
    def add_motion_blur(self, image, kernel_size=5, angle=0):
        """Apply motion blur to image"""
        from scipy.ndimage import convolve
        
        # Create motion blur kernel
        kernel = np.zeros((kernel_size, kernel_size))
        center = kernel_size // 2
        kernel[center, :] = 1 / kernel_size
        
        # Rotate kernel
        from scipy.ndimage import rotate
        kernel = rotate(kernel, angle, reshape=False)
        
        # Apply convolution
        blurred = convolve(image, kernel, mode='reflect')
        return blurred
    
    def change_brightness(self, image, factor=0.5):
        """Change image brightness"""
        return np.clip(image * factor, 0, 1)
    
    def change_contrast(self, image, factor=0.5):
        """Change image contrast"""
        mean = np.mean(image)
        return np.clip((image - mean) * factor + mean, 0, 1)
    
    def rotate_image(self, image, angle=10):
        """Rotate image by angle degrees"""
        from scipy.ndimage import rotate
        return rotate(image, angle, axes=(0, 1), reshape=False, mode='reflect')
    
    def scale_image(self, image, scale_factor=0.9):
        """Scale image by factor"""
        from scipy.ndimage import zoom
        h, w = image.shape[:2]
        scaled = zoom(image, (scale_factor, scale_factor, 1), order=1)
        
        # Crop or pad to original size
        if scaled.shape[0] > h:
            start = (scaled.shape[0] - h) // 2
            scaled = scaled[start:start+h, :, :]
        elif scaled.shape[0] < h:
            pad_h = (h - scaled.shape[0]) // 2
            scaled = np.pad(scaled, ((pad_h, h - scaled.shape[0] - pad_h), (0, 0), (0, 0)), mode='constant')
        
        if scaled.shape[1] > w:
            start = (scaled.shape[1] - w) // 2
            scaled = scaled[:, start:start+w, :]
        elif scaled.shape[1] < w:
            pad_w = (w - scaled.shape[1]) // 2
            scaled = np.pad(scaled, ((0, 0), (pad_w, w - scaled.shape[1] - pad_w), (0, 0)), mode='constant')
        
        return scaled
    
    def evaluate_robustness(self, X_test, y_test, corruption_type='gaussian_noise', 

                          corruption_levels=None, metric_fn=None):
        """

        Evaluate model robustness to a specific corruption type

        

        Args:

            X_test: Test images

            y_test: Ground truth masks

            corruption_type: Type of corruption to apply

            corruption_levels: List of corruption levels to test

            metric_fn: Function to compute metric (default: Dice coefficient)

            

        Returns:

            Dictionary of robustness metrics

        """
        if corruption_levels is None:
            corruption_levels = [0.01, 0.05, 0.1, 0.2, 0.3, 0.5]
        
        if metric_fn is None:
            def metric_fn(y_true, y_pred):
                intersection = np.sum(y_true * y_pred)
                union = np.sum(y_true) + np.sum(y_pred)
                return (2. * intersection + 1e-6) / (union + 1e-6)
        
        # Get baseline performance
        baseline_preds = self.model.predict(X_test)
        baseline_score = np.mean([metric_fn(y_true, y_pred) 
                                 for y_true, y_pred in zip(y_test, baseline_preds)])
        
        # Get corruption function
        corruption_fn = getattr(self, f'add_{corruption_type}')
        
        # Evaluate at each corruption level
        scores = []
        for level in corruption_levels:
            # Apply corruption
            corrupted_X = np.array([corruption_fn(x, level) for x in X_test])
            
            # Predict on corrupted images
            preds = self.model.predict(corrupted_X)
            
            # Compute metric
            level_scores = [metric_fn(y_true, y_pred) 
                          for y_true, y_pred in zip(y_test, preds)]
            scores.append(np.mean(level_scores))
        
        # Compute robustness metrics
        results = {
            'corruption_type': corruption_type,
            'baseline_score': float(baseline_score),
            'corruption_levels': corruption_levels,
            'scores': [float(s) for s in scores],
            'mean_corruption_score': float(np.mean(scores)),
            'robustness_index': float(np.mean(scores) / baseline_score) if baseline_score > 0 else 0,
            'performance_drop': float(baseline_score - np.mean(scores))
        }
        
        self.results[corruption_type] = results
        
        return results
    
    def evaluate_all_corruptions(self, X_test, y_test, corruption_types=None, **kwargs):
        """

        Evaluate model robustness to all corruption types

        

        Args:

            X_test: Test images

            y_test: Ground truth masks

            corruption_types: List of corruption types to test

            **kwargs: Additional arguments for evaluate_robustness

            

        Returns:

            Dictionary of all robustness results

        """
        if corruption_types is None:
            corruption_types = [
                'gaussian_noise',
                'salt_pepper_noise',
                'gaussian_blur',
                'motion_blur',
                'brightness',
                'contrast',
                'rotation',
                'scaling'
            ]
        
        all_results = {}
        for corruption_type in corruption_types:
            print(f"Evaluating robustness to {corruption_type}...")
            results = self.evaluate_robustness(X_test, y_test, corruption_type, **kwargs)
            all_results[corruption_type] = results
        
        return all_results
    
    def plot_robustness_results(self, results=None, save_path=None):
        """

        Plot robustness analysis results

        

        Args:

            results: Results dictionary (if None, uses self.results)

            save_path: Path to save plot

        """
        if results is None:
            results = self.results
        
        fig, axes = plt.subplots(2, 4, figsize=(20, 10))
        axes = axes.flatten()
        
        for idx, (corruption_type, result) in enumerate(results.items()):
            if idx >= 8:
                break
                
            ax = axes[idx]
            ax.plot(result['corruption_levels'], result['scores'], 'o-', linewidth=2)
            ax.axhline(y=result['baseline_score'], color='r', linestyle='--', alpha=0.5, label='Baseline')
            ax.set_title(f"{corruption_type.replace('_', ' ').title()}\n(Robustness Index: {result['robustness_index']:.3f})")
            ax.set_xlabel('Corruption Level')
            ax.set_ylabel('Dice Score')
            ax.legend()
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def save_results(self, results, save_dir='./robustness_results'):
        """Save robustness results to files"""
        os.makedirs(save_dir, exist_ok=True)
        
        # Save as JSON
        with open(os.path.join(save_dir, 'robustness_results.json'), 'w') as f:
            json.dump(results, f, indent=2)
        
        # Save as CSV
        rows = []
        for corruption_type, result in results.items():
            for level, score in zip(result['corruption_levels'], result['scores']):
                rows.append({
                    'corruption_type': corruption_type,
                    'corruption_level': level,
                    'dice_score': score,
                    'baseline_score': result['baseline_score'],
                    'robustness_index': result['robustness_index']
                })
        
        df = pd.DataFrame(rows)
        df.to_csv(os.path.join(save_dir, 'robustness_results.csv'), index=False)
        
        # Save plots
        self.plot_robustness_results(results, os.path.join(save_dir, 'robustness_plots.png'))


class UncertaintyEstimator:
    """

    Uncertainty estimation for segmentation models using Monte Carlo Dropout

    and Deep Ensembles

    """
    
    def __init__(self, model, num_samples=50, batch_size=32):
        """

        Initialize uncertainty estimator

        

        Args:

            model: Trained segmentation model with dropout

            num_samples: Number of Monte Carlo samples

            batch_size: Batch size for predictions

        """
        self.model = model
        self.num_samples = num_samples
        self.batch_size = batch_size
    
    def enable_dropout(self):
        """Enable dropout at inference time for MC Dropout"""
        for layer in self.model.layers:
            if isinstance(layer, layers.Dropout):
                layer.trainable = True
    
    def mc_dropout_predict(self, X, num_samples=None):
        """

        Monte Carlo Dropout prediction

        

        Args:

            X: Input images

            num_samples: Number of MC samples (if None, uses self.num_samples)

            

        Returns:

            mean_prediction, uncertainty, predictions_array

        """
        if num_samples is None:
            num_samples = self.num_samples
        
        # Enable dropout
        self.enable_dropout()
        
        # Get multiple predictions with dropout enabled
        predictions = []
        for _ in range(num_samples):
            pred = self.model.predict(X, verbose=0)
            predictions.append(pred)
        
        predictions = np.array(predictions)
        
        # Compute mean and uncertainty
        mean_pred = np.mean(predictions, axis=0)
        uncertainty = np.std(predictions, axis=0)
        
        return mean_pred, uncertainty, predictions
    
    def deep_ensemble_predict(self, models, X):
        """

        Deep ensemble prediction

        

        Args:

            models: List of trained models

            X: Input images

            

        Returns:

            mean_prediction, uncertainty, predictions_array

        """
        predictions = []
        for model in models:
            pred = model.predict(X, verbose=0)
            predictions.append(pred)
        
        predictions = np.array(predictions)
        
        # Compute mean and uncertainty
        mean_pred = np.mean(predictions, axis=0)
        uncertainty = np.std(predictions, axis=0)
        
        return mean_pred, uncertainty, predictions
    
    def compute_aleatoric_uncertainty(self, model, X, num_samples=10):
        """

        Estimate aleatoric uncertainty (data uncertainty)

        

        Args:

            model: Model that outputs both prediction and uncertainty

            X: Input images

            num_samples: Number of samples for estimation

            

        Returns:

            Aleatoric uncertainty map

        """
        # For models that output uncertainty directly
        predictions = []
        for _ in range(num_samples):
            pred = model.predict(X, verbose=0)
            if isinstance(pred, list) and len(pred) > 1:
                # Assume second output is uncertainty
                predictions.append(pred[1])
            else:
                predictions.append(pred)
        
        predictions = np.array(predictions)
        aleatoric_uncertainty = np.mean(predictions, axis=0)
        
        return aleatoric_uncertainty
    
    def compute_epistemic_uncertainty(self, models, X):
        """

        Estimate epistemic uncertainty (model uncertainty) using ensemble

        

        Args:

            models: List of trained models

            X: Input images

            

        Returns:

            Epistemic uncertainty map

        """
        predictions = []
        for model in models:
            pred = model.predict(X, verbose=0)
            predictions.append(pred)
        
        predictions = np.array(predictions)
        epistemic_uncertainty = np.var(predictions, axis=0)
        
        return epistemic_uncertainty
    
    def get_confidence_intervals(self, predictions_array, confidence_level=0.95):
        """

        Compute confidence intervals from prediction samples

        

        Args:

            predictions_array: Array of predictions (num_samples, height, width, channels)

            confidence_level: Confidence level for intervals

            

        Returns:

            lower_bound, upper_bound

        """
        alpha = 1 - confidence_level
        lower_percentile = alpha / 2 * 100
        upper_percentile = (1 - alpha / 2) * 100
        
        lower_bound = np.percentile(predictions_array, lower_percentile, axis=0)
        upper_bound = np.percentile(predictions_array, upper_percentile, axis=0)
        
        return lower_bound, upper_bound
    
    def visualize_uncertainty(self, image, mean_pred, uncertainty, threshold=0.5, 

                            save_path=None, cmap='viridis'):
        """

        Visualize prediction and uncertainty

        

        Args:

            image: Input image

            mean_pred: Mean prediction

            uncertainty: Uncertainty map

            threshold: Threshold for binary prediction

            save_path: Path to save visualization

            cmap: Colormap for uncertainty

        """
        fig, axes = plt.subplots(1, 4, figsize=(16, 4))
        
        # Input image
        axes[0].imshow(image)
        axes[0].set_title('Input Image')
        axes[0].axis('off')
        
        # Mean prediction
        binary_pred = (mean_pred >= threshold).astype(int)
        axes[1].imshow(binary_pred, cmap='gray')
        axes[1].set_title(f'Binary Prediction (threshold={threshold})')
        axes[1].axis('off')
        
        # Uncertainty map
        im = axes[2].imshow(uncertainty, cmap=cmap)
        axes[2].set_title('Uncertainty Map')
        axes[2].axis('off')
        plt.colorbar(im, ax=axes[2])
        
        # Overlay uncertainty on prediction
        axes[3].imshow(binary_pred, cmap='gray', alpha=0.7)
        im = axes[3].imshow(uncertainty, cmap=cmap, alpha=0.5)
        axes[3].set_title('Prediction with Uncertainty Overlay')
        axes[3].axis('off')
        plt.colorbar(im, ax=axes[3])
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def save_uncertainty_results(self, mean_pred, uncertainty, image, save_dir='./uncertainty_results'):
        """Save uncertainty estimation results"""
        os.makedirs(save_dir, exist_ok=True)
        
        # Save predictions
        np.save(os.path.join(save_dir, 'mean_prediction.npy'), mean_pred)
        np.save(os.path.join(save_dir, 'uncertainty.npy'), uncertainty)
        
        # Save visualization
        self.visualize_uncertainty(
            image, mean_pred, uncertainty,
            save_path=os.path.join(save_dir, 'uncertainty_visualization.png')
        )
        
        # Save summary statistics
        stats = {
            'mean_uncertainty': float(np.mean(uncertainty)),
            'max_uncertainty': float(np.max(uncertainty)),
            'min_uncertainty': float(np.min(uncertainty)),
            'std_uncertainty': float(np.std(uncertainty)),
            'high_uncertainty_pixels': float(np.mean(uncertainty > 0.5)),
            'medium_uncertainty_pixels': float(np.mean((uncertainty > 0.2) & (uncertainty <= 0.5))),
            'low_uncertainty_pixels': float(np.mean(uncertainty <= 0.2))
        }
        
        with open(os.path.join(save_dir, 'uncertainty_stats.json'), 'w') as f:
            json.dump(stats, f, indent=2)


class MulticlassSegmentationModel:
    """

    Multiclass segmentation model for different tumor types

    """
    
    def __init__(self, input_shape=(224, 224, 3), num_classes=4, base_filters=64, dropout_rate=0.2):
        """

        Initialize multiclass segmentation model

        

        Args:

            input_shape: Shape of input images

            num_classes: Number of segmentation classes (including background)

            base_filters: Number of base filters

            dropout_rate: Dropout rate

        """
        self.input_shape = input_shape
        self.num_classes = num_classes
        self.base_filters = base_filters
        self.dropout_rate = dropout_rate
        self.model = None
        
        # Class names (can be customized)
        self.class_names = ['background', 'glioma', 'meningioma', 'pituitary']
    
    def build_model(self, use_attention=False):
        """

        Build multiclass U-Net model

        

        Args:

            use_attention: Whether to use attention gates

            

        Returns:

            Compiled model

        """
        inputs = layers.Input(shape=self.input_shape, name='image_input')
        
        # Normalize input
        x = layers.Rescaling(1.0 / 255)(inputs)
        
        # Encoder
        filters = self.base_filters
        skip_connections = []
        
        for i in range(4):
            # Convolutional block
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Dropout(self.dropout_rate)(x)
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            
            skip_connections.append(x)
            x = layers.MaxPooling2D(pool_size=(2, 2))(x)
            filters *= 2
        
        # Bottleneck
        x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                        kernel_initializer='he_normal')(x)
        x = layers.BatchNormalization()(x)
        x = layers.Dropout(self.dropout_rate)(x)
        x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                        kernel_initializer='he_normal')(x)
        x = layers.BatchNormalization()(x)
        
        # Decoder
        filters //= 2
        
        for i in range(4):
            # Upsampling
            x = layers.Conv2DTranspose(filters, (2, 2), strides=(2, 2), padding='same')(x)
            
            # Apply attention gate if enabled
            if use_attention:
                # Attention mechanism
                skip = skip_connections.pop()
                attention_weights = layers.Conv2D(filters, 1, padding='same', use_bias=False)(skip)
                attention_weights = layers.Activation('sigmoid')(attention_weights)
                x = layers.Concatenate()([x, layers.Multiply()([skip, attention_weights])])
            else:
                # Simple concatenation
                x = layers.Concatenate()([x, skip_connections.pop()])
            
            # Convolutional block
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Dropout(self.dropout_rate)(x)
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            
            filters //= 2
        
        # Output layer with softmax for multiclass
        outputs = layers.Conv2D(self.num_classes, (1, 1), activation='softmax', padding='same')(x)
        
        self.model = Model(inputs=[inputs], outputs=[outputs], name='multiclass_unet')
        
        return self.model
    
    def compile_model(self, learning_rate=1e-4):
        """

        Compile model with appropriate loss and metrics

        

        Args:

            learning_rate: Learning rate for optimizer

        """
        if self.model is None:
            self.build_model()
        
        # Categorical crossentropy for multiclass
        self.model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
            loss='categorical_crossentropy',
            metrics=[
                'accuracy',
                tf.keras.metrics.MeanIoU(num_classes=self.num_classes),
                self.dice_coefficient_multiclass
            ]
        )
        
        return self.model
    
    def dice_coefficient_multiclass(self, y_true, y_pred, smooth=1e-6):
        """

        Dice coefficient for multiclass segmentation

        

        Args:

            y_true: Ground truth one-hot encoded masks

            y_pred: Predicted probabilities

            smooth: Smoothing factor

            

        Returns:

            Mean Dice coefficient across classes

        """
        # Get class predictions
        y_true_classes = tf.argmax(y_true, axis=-1)
        y_pred_classes = tf.argmax(y_pred, axis=-1)
        
        # One-hot encode predictions
        y_true_onehot = tf.one_hot(y_true_classes, depth=self.num_classes)
        y_pred_onehot = tf.one_hot(y_pred_classes, depth=self.num_classes)
        
        # Compute Dice for each class
        dice_scores = []
        for i in range(self.num_classes):
            intersection = tf.reduce_sum(y_true_onehot[..., i] * y_pred_onehot[..., i])
            union = tf.reduce_sum(y_true_onehot[..., i]) + tf.reduce_sum(y_pred_onehot[..., i])
            dice = (2. * intersection + smooth) / (union + smooth)
            dice_scores.append(dice)
        
        return tf.reduce_mean(dice_scores)
    
    def prepare_multiclass_masks(self, masks, num_classes=None):
        """

        Convert integer masks to one-hot encoded masks

        

        Args:

            masks: Integer masks with values 0 to num_classes-1

            num_classes: Number of classes (if None, uses self.num_classes)

            

        Returns:

            One-hot encoded masks

        """
        if num_classes is None:
            num_classes = self.num_classes
        
        # Convert to one-hot
        one_hot = tf.one_hot(masks.astype(int), depth=num_classes)
        
        return one_hot.numpy()
    
    def train(self, X_train, y_train, X_val, y_val, epochs=100, batch_size=16, callbacks=None):
        """

        Train the multiclass model

        

        Args:

            X_train: Training images

            y_train: Training masks (integer or one-hot)

            X_val: Validation images

            y_val: Validation masks

            epochs: Number of training epochs

            batch_size: Batch size

            callbacks: List of Keras callbacks

            

        Returns:

            Training history

        """
        # Build and compile model
        self.compile_model()
        
        # Convert masks to one-hot if needed
        if len(y_train.shape) == 3 or (len(y_train.shape) == 4 and y_train.shape[-1] != self.num_classes):
            y_train = self.prepare_multiclass_masks(y_train)
            y_val = self.prepare_multiclass_masks(y_val)
        
        # Default callbacks
        if callbacks is None:
            callbacks = []
        
        callbacks.extend([
            tf.keras.callbacks.EarlyStopping(
                monitor='val_loss',
                patience=15,
                restore_best_weights=True
            ),
            tf.keras.callbacks.ReduceLROnPlateau(
                monitor='val_loss',
                factor=0.5,
                patience=5,
                min_lr=1e-7
            )
        ])
        
        # Train
        history = self.model.fit(
            X_train,
            y_train,
            validation_data=(X_val, y_val),
            epochs=epochs,
            batch_size=batch_size,
            callbacks=callbacks,
            verbose=1
        )
        
        return history
    
    def predict(self, X):
        """

        Predict segmentation masks

        

        Args:

            X: Input images

            

        Returns:

            Predicted masks (integer format)

        """
        # Get predictions
        predictions = self.model.predict(X)
        
        # Convert to integer masks
        predicted_masks = np.argmax(predictions, axis=-1)
        
        return predicted_masks
    
    def predict_proba(self, X):
        """

        Predict class probabilities

        

        Args:

            X: Input images

            

        Returns:

            Probability maps for each class

        """
        return self.model.predict(X)
    
    def evaluate_multiclass(self, X_test, y_test, class_names=None):
        """

        Evaluate multiclass segmentation with per-class metrics

        

        Args:

            X_test: Test images

            y_test: Ground truth masks (integer format)

            class_names: List of class names

            

        Returns:

            Dictionary of evaluation metrics

        """
        if class_names is None:
            class_names = self.class_names
        
        # Get predictions
        y_pred = self.predict(X_test)
        
        # Flatten for metric calculation
        y_true_flat = y_test.flatten()
        y_pred_flat = y_pred.flatten()
        
        # Compute per-class metrics
        metrics = {}
        
        # Overall metrics
        overall_accuracy = np.mean(y_true_flat == y_pred_flat)
        
        # Per-class metrics
        class_metrics = {}
        for class_idx, class_name in enumerate(class_names):
            # Binary mask for this class
            true_binary = (y_true_flat == class_idx)
            pred_binary = (y_pred_flat == class_idx)
            
            # Compute metrics
            intersection = np.sum(true_binary & pred_binary)
            union = np.sum(true_binary) + np.sum(pred_binary) - intersection
            
            if np.sum(true_binary) > 0:
                recall = intersection / np.sum(true_binary)
            else:
                recall = 0
            
            if np.sum(pred_binary) > 0:
                precision = intersection / np.sum(pred_binary)
            else:
                precision = 0
            
            if precision + recall > 0:
                f1 = 2 * precision * recall / (precision + recall)
            else:
                f1 = 0
            
            if union > 0:
                iou = intersection / union
            else:
                iou = 0
            
            dice = (2 * intersection + 1e-6) / (np.sum(true_binary) + np.sum(pred_binary) + 1e-6)
            
            class_metrics[class_name] = {
                'precision': float(precision),
                'recall': float(recall),
                'f1_score': float(f1),
                'iou': float(iou),
                'dice': float(dice),
                'support': int(np.sum(true_binary))
            }
        
        # Compute mean metrics
        mean_metrics = {}
        for metric in ['precision', 'recall', 'f1_score', 'iou', 'dice']:
            mean_metrics[f'mean_{metric}'] = float(np.mean([cm[metric] for cm in class_metrics.values()]))
        
        metrics['overall_accuracy'] = float(overall_accuracy)
        metrics['class_metrics'] = class_metrics
        metrics['mean_metrics'] = mean_metrics
        
        # Confusion matrix
        cm = confusion_matrix(y_true_flat, y_pred_flat, labels=list(range(len(class_names))))
        metrics['confusion_matrix'] = cm.tolist()
        
        return metrics
    
    def visualize_multiclass_prediction(self, image, true_mask, pred_mask, class_names=None, save_path=None):
        """

        Visualize multiclass segmentation results

        

        Args:

            image: Input image

            true_mask: Ground truth mask

            pred_mask: Predicted mask

            class_names: List of class names

            save_path: Path to save visualization

        """
        if class_names is None:
            class_names = self.class_names
        
        # Create color map for classes
        colors = plt.cm.Set3(np.linspace(0, 1, len(class_names)))
        
        fig, axes = plt.subplots(1, 4, figsize=(16, 4))
        
        # Input image
        axes[0].imshow(image)
        axes[0].set_title('Input Image')
        axes[0].axis('off')
        
        # Ground truth
        axes[1].imshow(true_mask, cmap='tab10', vmin=0, vmax=len(class_names)-1)
        axes[1].set_title('Ground Truth')
        axes[1].axis('off')
        
        # Prediction
        axes[2].imshow(pred_mask, cmap='tab10', vmin=0, vmax=len(class_names)-1)
        axes[2].set_title('Prediction')
        axes[2].axis('off')
        
        # Overlay
        axes[3].imshow(image)
        axes[3].imshow(pred_mask, cmap='tab10', vmin=0, vmax=len(class_names)-1, alpha=0.5)
        axes[3].set_title('Prediction Overlay')
        axes[3].axis('off')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()


def create_multiclass_dataset_from_binary(binary_images, binary_masks, tumor_types):
    """

    Create multiclass dataset from binary segmentation data

    

    Args:

        binary_images: List of binary images

        binary_masks: List of binary masks

        tumor_types: List of tumor type labels for each image

        

    Returns:

        Multiclass images and masks

    """
    # Create mapping from tumor type to class index
    tumor_type_to_class = {tumor_type: idx+1 for idx, tumor_type in enumerate(set(tumor_types))}
    
    # Create multiclass masks
    multiclass_masks = []
    for mask, tumor_type in zip(binary_masks, tumor_types):
        # Start with background (class 0)
        multiclass_mask = np.zeros_like(mask, dtype=np.int32)
        
        # Set tumor region to appropriate class
        class_idx = tumor_type_to_class.get(tumor_type, 0)
        multiclass_mask[mask > 0.5] = class_idx
        
        multiclass_masks.append(multiclass_mask)
    
    return np.array(binary_images), np.array(multiclass_masks)