""" Advanced Explainable AI (XAI) Module for MorphGuard Implements state-of-the-art interpretability techniques beyond basic gradient saliency """ import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import cv2 from typing import Dict, List, Tuple, Optional, Any, Union import logging from dataclasses import dataclass from datetime import datetime import matplotlib.pyplot as plt import seaborn as sns from captum.attr import ( IntegratedGradients, GradientShap, DeepLift, LayerConductance, LayerActivation, LayerGradientXActivation, Saliency, InputXGradient, GuidedBackprop, GuidedGradCam, Occlusion, KernelShap, Lime, ShapleyValueSampling ) from captum.attr._utils.visualization import visualize_image_attr from captum.concept import TCAV from captum.robust import PGD, FGSM import shap import lime from lime import lime_image import json logger = logging.getLogger(__name__) @dataclass class ExplanationResult: """Comprehensive explanation result with multiple XAI techniques""" method: str attribution_map: np.ndarray confidence: float interpretation_score: float feature_importance: Dict[str, float] concept_activations: Optional[Dict[str, float]] counterfactual_regions: Optional[List[Tuple[int, int, int, int]]] textual_explanation: str timestamp: datetime processing_time_ms: float @dataclass class ConceptActivation: """TCAV-based concept activation""" concept_name: str activation_strength: float statistical_significance: float directional_derivative: float class AdvancedXAIEngine: """ Advanced Explainable AI engine implementing multiple state-of-the-art techniques """ def __init__(self, model: nn.Module, device: str = 'cuda'): self.model = model self.device = device self.model.eval() # Initialize different attribution methods self.integrated_gradients = IntegratedGradients(self.model) self.gradient_shap = GradientShap(self.model) self.deep_lift = DeepLift(self.model) self.saliency = Saliency(self.model) self.input_x_gradient = InputXGradient(self.model) self.guided_backprop = GuidedBackprop(self.model) self.occlusion = Occlusion(self.model) # SHAP explainer self.shap_explainer = None # LIME explainer self.lime_explainer = lime_image.LimeImageExplainer() # Concept vectors for TCAV self.concept_vectors = {} logger.info("Advanced XAI Engine initialized with multiple attribution methods") def explain_prediction( self, image: torch.Tensor, target_class: int = 1, methods: Optional[List[str]] = None, baseline_strategy: str = 'zero' ) -> Dict[str, ExplanationResult]: """ Generate comprehensive explanations using multiple XAI techniques Args: image: Input image tensor [1, C, H, W] target_class: Target class for explanation (1 for morph) methods: List of methods to use. If None, uses all available baseline_strategy: Baseline strategy for integrated gradients Returns: Dictionary of explanation results by method name """ if methods is None: methods = [ 'integrated_gradients', 'gradient_shap', 'deep_lift', 'saliency', 'guided_backprop', 'occlusion', 'shap', 'lime' ] explanations = {} image = image.to(self.device) # Get model prediction for confidence with torch.no_grad(): prediction = self.model(image) confidence = float(prediction.sigmoid().cpu()) for method in methods: start_time = time.time() try: explanation = self._generate_explanation( image, target_class, method, baseline_strategy, confidence ) explanation.processing_time_ms = (time.time() - start_time) * 1000 explanations[method] = explanation logger.debug(f"Generated {method} explanation in {explanation.processing_time_ms:.2f}ms") except Exception as e: logger.error(f"Failed to generate {method} explanation: {e}") return explanations def _generate_explanation( self, image: torch.Tensor, target_class: int, method: str, baseline_strategy: str, confidence: float ) -> ExplanationResult: """Generate explanation using specific method""" if method == 'integrated_gradients': return self._integrated_gradients_explanation( image, target_class, baseline_strategy, confidence ) elif method == 'gradient_shap': return self._gradient_shap_explanation(image, target_class, confidence) elif method == 'deep_lift': return self._deep_lift_explanation(image, target_class, confidence) elif method == 'saliency': return self._saliency_explanation(image, target_class, confidence) elif method == 'guided_backprop': return self._guided_backprop_explanation(image, target_class, confidence) elif method == 'occlusion': return self._occlusion_explanation(image, target_class, confidence) elif method == 'shap': return self._shap_explanation(image, target_class, confidence) elif method == 'lime': return self._lime_explanation(image, target_class, confidence) else: raise ValueError(f"Unknown explanation method: {method}") def _integrated_gradients_explanation( self, image: torch.Tensor, target_class: int, baseline_strategy: str, confidence: float ) -> ExplanationResult: """Enhanced Integrated Gradients with multiple baselines""" # Generate multiple baselines baselines = self._generate_baselines(image, baseline_strategy) attributions_list = [] for baseline in baselines: attr = self.integrated_gradients.attribute( image, baselines=baseline, target=target_class, n_steps=100, internal_batch_size=1 ) attributions_list.append(attr) # Ensemble attributions final_attribution = torch.mean(torch.stack(attributions_list), dim=0) attribution_map = self._process_attribution(final_attribution) # Calculate interpretation score interpretation_score = self._calculate_interpretation_score( final_attribution, image ) # Extract feature importance feature_importance = self._extract_feature_importance(final_attribution) # Generate textual explanation textual_explanation = self._generate_textual_explanation( attribution_map, feature_importance, confidence, "Integrated Gradients" ) return ExplanationResult( method="Integrated Gradients", attribution_map=attribution_map, confidence=confidence, interpretation_score=interpretation_score, feature_importance=feature_importance, concept_activations=None, counterfactual_regions=None, textual_explanation=textual_explanation, timestamp=datetime.now() ) def _gradient_shap_explanation( self, image: torch.Tensor, target_class: int, confidence: float ) -> ExplanationResult: """Gradient SHAP explanation with noise baseline""" # Generate noise baseline noise_baseline = torch.randn_like(image) * 0.1 attribution = self.gradient_shap.attribute( image, baselines=noise_baseline, target=target_class, n_samples=50, stdevs=0.1 ) attribution_map = self._process_attribution(attribution) interpretation_score = self._calculate_interpretation_score(attribution, image) feature_importance = self._extract_feature_importance(attribution) textual_explanation = self._generate_textual_explanation( attribution_map, feature_importance, confidence, "Gradient SHAP" ) return ExplanationResult( method="Gradient SHAP", attribution_map=attribution_map, confidence=confidence, interpretation_score=interpretation_score, feature_importance=feature_importance, concept_activations=None, counterfactual_regions=None, textual_explanation=textual_explanation, timestamp=datetime.now() ) def _shap_explanation( self, image: torch.Tensor, target_class: int, confidence: float ) -> ExplanationResult: """SHAP explanation using Deep Explainer""" # Initialize SHAP explainer if not done if self.shap_explainer is None: # Create background dataset from random samples background = torch.randn(10, *image.shape[1:]).to(self.device) self.shap_explainer = shap.DeepExplainer(self.model, background) # Generate SHAP values shap_values = self.shap_explainer.shap_values(image) if isinstance(shap_values, list): shap_values = shap_values[target_class] attribution_map = self._process_attribution(torch.tensor(shap_values)) interpretation_score = np.mean(np.abs(shap_values)) feature_importance = self._extract_feature_importance(torch.tensor(shap_values)) textual_explanation = self._generate_textual_explanation( attribution_map, feature_importance, confidence, "SHAP" ) return ExplanationResult( method="SHAP", attribution_map=attribution_map, confidence=confidence, interpretation_score=float(interpretation_score), feature_importance=feature_importance, concept_activations=None, counterfactual_regions=None, textual_explanation=textual_explanation, timestamp=datetime.now() ) def _lime_explanation( self, image: torch.Tensor, target_class: int, confidence: float ) -> ExplanationResult: """LIME explanation using superpixel segmentation""" # Convert to numpy for LIME img_np = image.squeeze().cpu().numpy().transpose(1, 2, 0) if img_np.min() < 0: # Denormalize if needed img_np = (img_np + 1) / 2 img_np = (img_np * 255).astype(np.uint8) def predict_fn(images): """Prediction function for LIME""" batch = [] for img in images: # Normalize and convert to tensor img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0 img_tensor = (img_tensor - 0.5) / 0.5 # Normalize to [-1, 1] batch.append(img_tensor) batch_tensor = torch.stack(batch).to(self.device) with torch.no_grad(): predictions = self.model(batch_tensor).sigmoid().cpu().numpy() return predictions # Generate LIME explanation explanation = self.lime_explainer.explain_instance( img_np, predict_fn, top_labels=2, hide_color=0, num_samples=1000 ) # Get mask for target class temp, mask = explanation.get_image_and_mask( target_class, positive_only=False, num_features=10, hide_rest=False ) attribution_map = mask.astype(np.float32) interpretation_score = np.std(mask) # Variability as interpretation score # Feature importance from LIME segments feature_importance = {} for i, (feature_id, weight) in enumerate(explanation.local_exp[target_class]): feature_importance[f"superpixel_{feature_id}"] = float(weight) textual_explanation = self._generate_textual_explanation( attribution_map, feature_importance, confidence, "LIME" ) return ExplanationResult( method="LIME", attribution_map=attribution_map, confidence=confidence, interpretation_score=interpretation_score, feature_importance=feature_importance, concept_activations=None, counterfactual_regions=None, textual_explanation=textual_explanation, timestamp=datetime.now() ) def analyze_concepts( self, image: torch.Tensor, concept_dataset: Dict[str, torch.Tensor] ) -> Dict[str, ConceptActivation]: """ Perform Testing with Concept Activation Vectors (TCAV) analysis Args: image: Input image tensor concept_dataset: Dictionary of concept name to concept examples Returns: Dictionary of concept activations """ concept_activations = {} # This is a simplified TCAV implementation # In practice, you'd need to train linear classifiers for each concept for concept_name, concept_examples in concept_dataset.items(): try: # Extract activations from a specific layer layer_name = 'features' # Adjust based on your model architecture # Get activations for input image image_activation = self._get_layer_activation(image, layer_name) # Get activations for concept examples concept_activations_list = [] for concept_img in concept_examples: concept_activation = self._get_layer_activation(concept_img, layer_name) concept_activations_list.append(concept_activation) # Calculate concept activation strength (simplified) concept_mean = torch.mean(torch.stack(concept_activations_list), dim=0) activation_strength = torch.cosine_similarity( image_activation.flatten(), concept_mean.flatten(), dim=0 ).item() concept_activations[concept_name] = ConceptActivation( concept_name=concept_name, activation_strength=activation_strength, statistical_significance=abs(activation_strength), # Simplified directional_derivative=activation_strength ) except Exception as e: logger.error(f"Failed to analyze concept {concept_name}: {e}") return concept_activations def generate_counterfactuals( self, image: torch.Tensor, target_change: float = 0.5 ) -> List[Tuple[np.ndarray, float, List[Tuple[int, int, int, int]]]]: """ Generate counterfactual explanations by finding minimal changes Args: image: Input image tensor target_change: Desired change in prediction confidence Returns: List of (modified_image, new_confidence, changed_regions) """ counterfactuals = [] # Get original prediction with torch.no_grad(): original_pred = self.model(image).sigmoid().item() target_pred = max(0.0, min(1.0, original_pred + target_change)) # Use gradient-based optimization to find minimal changes modified_image = image.clone().requires_grad_(True) optimizer = torch.optim.Adam([modified_image], lr=0.01) for iteration in range(100): optimizer.zero_grad() prediction = self.model(modified_image).sigmoid() # Loss: difference from target + L2 regularization loss = (prediction - target_pred).pow(2) + 0.1 * (modified_image - image).pow(2).sum() loss.backward() optimizer.step() # Clamp to valid image range with torch.no_grad(): modified_image.clamp_(-1, 1) if iteration % 20 == 0: current_pred = prediction.item() if abs(current_pred - target_pred) < 0.05: break # Find changed regions diff = torch.abs(modified_image - image).squeeze() diff_np = diff.cpu().numpy() # Threshold to find significant changes threshold = np.percentile(diff_np.flatten(), 95) changed_mask = diff_np > threshold # Find bounding boxes of changed regions changed_regions = self._find_changed_regions(changed_mask) final_pred = self.model(modified_image).sigmoid().item() modified_image_np = modified_image.squeeze().cpu().numpy().transpose(1, 2, 0) counterfactuals.append((modified_image_np, final_pred, changed_regions)) return counterfactuals def _generate_baselines(self, image: torch.Tensor, strategy: str) -> List[torch.Tensor]: """Generate multiple baselines for robust attribution""" baselines = [] if strategy == 'zero': baselines.append(torch.zeros_like(image)) elif strategy == 'mean': baselines.append(torch.full_like(image, image.mean())) elif strategy == 'blur': # Gaussian blur baseline blurred = F.avg_pool2d(image, kernel_size=21, stride=1, padding=10) baselines.append(blurred) elif strategy == 'noise': baselines.append(torch.randn_like(image) * 0.1) elif strategy == 'all': # Use all baseline strategies baselines.extend(self._generate_baselines(image, 'zero')) baselines.extend(self._generate_baselines(image, 'mean')) baselines.extend(self._generate_baselines(image, 'blur')) baselines.extend(self._generate_baselines(image, 'noise')) return baselines def _process_attribution(self, attribution: torch.Tensor) -> np.ndarray: """Process attribution tensor to visualization-ready format""" attr = attribution.squeeze().cpu().numpy() # Handle multi-channel attributions if len(attr.shape) == 3: attr = np.mean(attr, axis=0) # Average across channels # Normalize attr_max = np.max(np.abs(attr)) if attr_max > 0: attr = attr / attr_max return attr def _calculate_interpretation_score( self, attribution: torch.Tensor, image: torch.Tensor ) -> float: """Calculate interpretation quality score""" # Coherence: spatial continuity of attributions attr_np = self._process_attribution(attribution) # Calculate gradient magnitude (smoothness) grad_x = np.gradient(attr_np, axis=1) grad_y = np.gradient(attr_np, axis=0) gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2) coherence = 1.0 / (1.0 + np.mean(gradient_magnitude)) # Sparsity: concentration of important regions attr_abs = np.abs(attr_np) sparsity = 1.0 - (np.count_nonzero(attr_abs > 0.1) / attr_abs.size) # Combined interpretation score interpretation_score = 0.6 * coherence + 0.4 * sparsity return float(interpretation_score) def _extract_feature_importance(self, attribution: torch.Tensor) -> Dict[str, float]: """Extract regional feature importance""" attr_np = self._process_attribution(attribution) h, w = attr_np.shape # Divide into regions and calculate importance regions = { 'top_left': attr_np[:h//2, :w//2], 'top_right': attr_np[:h//2, w//2:], 'bottom_left': attr_np[h//2:, :w//2], 'bottom_right': attr_np[h//2:, w//2:], 'center': attr_np[h//4:3*h//4, w//4:3*w//4], 'edges': np.concatenate([ attr_np[0, :], attr_np[-1, :], attr_np[:, 0], attr_np[:, -1] ]) } feature_importance = {} for region_name, region_attr in regions.items(): feature_importance[region_name] = float(np.mean(np.abs(region_attr))) return feature_importance def _generate_textual_explanation( self, attribution_map: np.ndarray, feature_importance: Dict[str, float], confidence: float, method: str ) -> str: """Generate human-readable textual explanation""" # Find most important regions sorted_regions = sorted( feature_importance.items(), key=lambda x: x[1], reverse=True ) top_regions = [region for region, _ in sorted_regions[:3]] confidence_text = "high" if confidence > 0.7 else "medium" if confidence > 0.3 else "low" prediction_text = "morphed" if confidence > 0.5 else "authentic" explanation = f"Using {method}, the model predicts this image is {prediction_text} " \ f"with {confidence_text} confidence ({confidence:.2f}). " \ f"The most influential regions are: {', '.join(top_regions)}. " # Add specific insights based on attribution patterns if feature_importance.get('edges', 0) > 0.3: explanation += "Edge artifacts suggest potential morphing. " if feature_importance.get('center', 0) > 0.5: explanation += "Central facial features show strong morphing indicators. " return explanation def _get_layer_activation(self, image: torch.Tensor, layer_name: str) -> torch.Tensor: """Extract activation from specific layer""" activations = {} def hook_fn(module, input, output): activations[layer_name] = output # Register hook (simplified - you'd need to adapt this to your model) handle = None for name, module in self.model.named_modules(): if layer_name in name: handle = module.register_forward_hook(hook_fn) break if handle is None: raise ValueError(f"Layer {layer_name} not found in model") with torch.no_grad(): _ = self.model(image) handle.remove() return activations.get(layer_name, torch.tensor([])) def _find_changed_regions(self, mask: np.ndarray) -> List[Tuple[int, int, int, int]]: """Find bounding boxes of changed regions""" # Convert to uint8 for OpenCV mask_uint8 = (mask * 255).astype(np.uint8) # Find contours contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) regions = [] for contour in contours: x, y, w, h = cv2.boundingRect(contour) if w > 10 and h > 10: # Filter small regions regions.append((x, y, x + w, y + h)) return regions def visualize_explanations( self, image: torch.Tensor, explanations: Dict[str, ExplanationResult], save_path: Optional[str] = None ) -> np.ndarray: """Create comprehensive visualization of all explanations""" num_methods = len(explanations) fig, axes = plt.subplots(2, max(4, num_methods//2), figsize=(20, 10)) axes = axes.flatten() # Original image img_np = image.squeeze().cpu().numpy().transpose(1, 2, 0) if img_np.min() < 0: img_np = (img_np + 1) / 2 axes[0].imshow(img_np) axes[0].set_title("Original Image") axes[0].axis('off') # Plot each explanation for i, (method, explanation) in enumerate(explanations.items(), 1): if i < len(axes): im = axes[i].imshow(explanation.attribution_map, cmap='RdBu_r', vmin=-1, vmax=1) axes[i].set_title(f"{method}\nScore: {explanation.interpretation_score:.3f}") axes[i].axis('off') plt.colorbar(im, ax=axes[i], fraction=0.046, pad=0.04) # Hide unused subplots for i in range(len(explanations) + 1, len(axes)): axes[i].axis('off') plt.tight_layout() if save_path: plt.savefig(save_path, dpi=300, bbox_inches='tight') # Convert to numpy array for return fig.canvas.draw() buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (3,)) plt.close(fig) return buf class ConfidenceCalibrator: """ Confidence calibration for improved model reliability """ def __init__(self): self.temperature = 1.0 self.calibration_data = [] def calibrate(self, logits: torch.Tensor, labels: torch.Tensor) -> float: """ Temperature scaling calibration Args: logits: Model logits labels: True labels Returns: Optimal temperature parameter """ # Find optimal temperature using validation set temperature = nn.Parameter(torch.ones(1) * 1.5) optimizer = torch.optim.LBFGS([temperature], lr=0.01, max_iter=50) def eval_loss(): optimizer.zero_grad() loss = F.cross_entropy(logits / temperature, labels) loss.backward() return loss optimizer.step(eval_loss) self.temperature = temperature.item() return self.temperature def apply_calibration(self, logits: torch.Tensor) -> torch.Tensor: """Apply temperature scaling to logits""" return logits / self.temperature if __name__ == "__main__": # Example usage logging.basicConfig(level=logging.INFO) # This would be your actual model model = torch.nn.Sequential( torch.nn.Conv2d(3, 64, 3), torch.nn.ReLU(), torch.nn.AdaptiveAvgPool2d(1), torch.nn.Flatten(), torch.nn.Linear(64, 1) ) xai_engine = AdvancedXAIEngine(model) # Example image image = torch.randn(1, 3, 224, 224) # Generate explanations explanations = xai_engine.explain_prediction(image) for method, explanation in explanations.items(): print(f"{method}: {explanation.textual_explanation}")