""" CTM Forensic Agent for MorphGuard ================================= Deep Forensic Reasoning Engine using Sakana AI's Continuous Thought Machine. This provides "System 2" (deliberative) analysis for detecting geometric inconsistencies in high-quality morphed images that standard models miss. Key Features: - Temporal reasoning loop (10-20 "ticks" of thinking) - Attention-based sequential image scanning - Evidence video generation showing detection reasoning - Adaptive early exit when confidence exceeds threshold """ import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import os import time import logging from typing import Dict, Any, Optional, List, Tuple from dataclasses import dataclass from PIL import Image import torchvision.transforms as transforms # Import CTM wrapper try: from external import CTMModel, CTMConfig, CTMOutput, plot_attention_video, CTM_AVAILABLE except ImportError: CTM_AVAILABLE = False CTMModel = None CTMConfig = None logger = logging.getLogger(__name__) @dataclass class ForensicResult: """Result from CTM forensic analysis""" is_morphed: bool confidence: float forensic_steps: int evidence_video_path: Optional[str] attention_regions: List[str] processing_time_ms: float tier: str = "Tier2_Forensic" def to_dict(self) -> Dict[str, Any]: return { 'is_morphed': self.is_morphed, 'confidence': self.confidence, 'forensic_steps': self.forensic_steps, 'evidence_video_path': self.evidence_video_path, 'attention_regions': self.attention_regions, 'processing_time_ms': self.processing_time_ms, 'tier': self.tier } class CTMForensicAgent: """ Deep Forensic Reasoning Agent using Continuous Thought Machine. This agent performs multi-step reasoning on face images to detect morphing artifacts that standard single-pass models miss, such as: - Geometric inconsistencies (ear symmetry, hairline logic) - Light/shadow coherence across facial features - Micro-texture continuity at blend boundaries Usage: agent = CTMForensicAgent('models/ctm_forensic.pth') result = agent.analyze(image_tensor) print(f"Is Morphed: {result.is_morphed}, Steps: {result.forensic_steps}") """ def __init__( self, checkpoint_path: Optional[str] = None, device: str = 'cuda' if torch.cuda.is_available() else 'cpu', max_thoughts: int = 16, confidence_threshold: float = 0.95, early_exit_confidence: float = 0.99, evidence_output_dir: str = 'static/evidence' ): """ Initialize CTM Forensic Agent. Args: checkpoint_path: Path to CTM model checkpoint device: Device to run inference on ('cuda' or 'cpu') max_thoughts: Maximum number of thinking steps confidence_threshold: Threshold above which fast detector result is trusted early_exit_confidence: Confidence level for early exit during thinking evidence_output_dir: Directory to save evidence videos """ self.device = device self.max_thoughts = max_thoughts self.confidence_threshold = confidence_threshold self.early_exit_confidence = early_exit_confidence self.evidence_output_dir = evidence_output_dir self.checkpoint_path = checkpoint_path # Lazy loading - model loaded on first use self._model: Optional[nn.Module] = None self._is_loaded = False # Image preprocessing (ImageNet normalization) self.transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Ensure evidence output directory exists os.makedirs(evidence_output_dir, exist_ok=True) # Facial region labels for interpretability self.region_labels = [ 'left_eye', 'right_eye', 'nose', 'mouth', 'left_ear', 'right_ear', 'forehead', 'chin', 'left_cheek', 'right_cheek', 'jawline', 'hairline', 'neck', 'background' ] logger.info(f"CTMForensicAgent initialized (lazy loading enabled)") logger.info(f" Device: {device}, Max thoughts: {max_thoughts}") logger.info(f" CTM Available: {CTM_AVAILABLE}") def load_model(self) -> bool: """Load CTM model (called lazily on first inference)""" if self._is_loaded: return True if not CTM_AVAILABLE or CTMModel is None: logger.warning("CTM not available, skipping CTM load") return False try: # Load or initialize model if self.checkpoint_path: if not os.path.exists(self.checkpoint_path): raise FileNotFoundError(f"CTM checkpoint not found at {self.checkpoint_path}") logger.info(f"Loading CTM from {self.checkpoint_path}") self._model = CTMModel.load_from_checkpoint( self.checkpoint_path, device=self.device ) else: raise ValueError("CTM checkpoint path must be provided to avoid random weights") self._model.eval() self._is_loaded = True # Warmup inference self._warmup() logger.info("CTM model loaded successfully") return True except Exception as e: logger.error(f"Failed to load CTM model: {e}") return False def _warmup(self): """Warmup model with dummy input""" try: dummy = torch.randn(1, 3, 224, 224, device=self.device) with torch.no_grad(): self._model(dummy, max_thoughts=2) logger.info("CTM warmup complete") except Exception as e: logger.warning(f"CTM warmup failed: {e}") def unload_model(self): """Unload model to free GPU memory""" if self._model is not None: del self._model self._model = None self._is_loaded = False torch.cuda.empty_cache() logger.info("CTM model unloaded") def preprocess(self, image: Any) -> torch.Tensor: """ Preprocess image for CTM inference. Args: image: PIL Image, numpy array, or file path Returns: Preprocessed tensor (1, 3, 224, 224) """ if isinstance(image, str): image = Image.open(image).convert('RGB') elif isinstance(image, np.ndarray): image = Image.fromarray(image).convert('RGB') elif isinstance(image, torch.Tensor): # Already a tensor, just ensure correct shape if image.dim() == 3: image = image.unsqueeze(0) return image.to(self.device) tensor = self.transform(image).unsqueeze(0) return tensor.to(self.device) def analyze( self, image: Any, generate_evidence: bool = True, request_id: Optional[str] = None ) -> ForensicResult: """ Perform deep forensic reasoning on a face image. This is the main entry point for CTM analysis. The model will: 1. Initialize neural state 2. Iteratively attend to different facial regions 3. Update internal representation with each "thought" 4. Exit early if highly confident 5. Generate evidence video if requested Args: image: Input image (PIL, numpy, path, or tensor) generate_evidence: Whether to generate attention replay GIF request_id: Optional ID for naming evidence files Returns: ForensicResult with prediction, confidence, and evidence """ start_time = time.time() # Ensure model is loaded if not self._is_loaded: if not self.load_model(): return ForensicResult( is_morphed=False, confidence=0.0, forensic_steps=0, evidence_video_path=None, attention_regions=[], processing_time_ms=0.0 ) # Preprocess image image_tensor = self.preprocess(image) # Keep original for visualization if isinstance(image, str): original_image = Image.open(image).convert('RGB') original_tensor = transforms.ToTensor()(original_image).unsqueeze(0) elif isinstance(image, np.ndarray): original_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float() / 255.0 else: original_tensor = image_tensor.clone() # Run CTM inference with torch.no_grad(): output, attention_history = self._model( image_tensor, max_thoughts=self.max_thoughts, early_exit_threshold=self.early_exit_confidence ) # Analyze which regions were attended to attention_regions = self._analyze_attention_regions(attention_history) # Generate evidence video evidence_path = None if generate_evidence and attention_history: evidence_filename = f"forensic_{request_id or int(time.time())}.gif" evidence_path = os.path.join(self.evidence_output_dir, evidence_filename) try: plot_attention_video(attention_history, original_tensor, evidence_path) except Exception as e: logger.warning(f"Failed to generate evidence video: {e}") evidence_path = None processing_time = (time.time() - start_time) * 1000 # Convert to ms # Create result result = ForensicResult( is_morphed=output.prediction == 1, # 1 = morphed, 0 = real confidence=output.confidence, forensic_steps=len(attention_history), evidence_video_path=evidence_path, attention_regions=attention_regions, processing_time_ms=processing_time ) logger.info(f"CTM analysis complete: morphed={result.is_morphed}, " f"confidence={result.confidence:.3f}, steps={result.forensic_steps}") return result def _analyze_attention_regions( self, attention_history: List[torch.Tensor] ) -> List[str]: """ Analyze attention maps to determine which facial regions were focused on. Maps attention peak locations to semantic facial regions. """ if not attention_history: return [] regions_attended = set() for attention in attention_history: if isinstance(attention, torch.Tensor): attn = attention.numpy() if attention.dim() == 2 else attention[0].numpy() else: attn = attention # Find peak attention location peak_idx = np.unravel_index(np.argmax(attn), attn.shape) # Map to facial region based on location in 14x14 grid # Approximate facial region mapping y, x = peak_idx region = self._map_location_to_region(x, y, attn.shape[1], attn.shape[0]) regions_attended.add(region) return list(regions_attended) def _map_location_to_region( self, x: int, y: int, width: int, height: int ) -> str: """Map grid location to facial region name""" # Normalize to 0-1 range nx = x / width ny = y / height # Simple region mapping based on typical face proportions if ny < 0.25: return 'forehead' if 0.25 < nx < 0.75 else 'hairline' elif ny < 0.45: if nx < 0.35: return 'left_eye' elif nx > 0.65: return 'right_eye' else: return 'nose' elif ny < 0.65: if nx < 0.2: return 'left_ear' elif nx > 0.8: return 'right_ear' elif nx < 0.35: return 'left_cheek' elif nx > 0.65: return 'right_cheek' else: return 'nose' elif ny < 0.8: if 0.3 < nx < 0.7: return 'mouth' else: return 'jawline' else: return 'chin' if 0.3 < nx < 0.7 else 'neck' def should_use_ctm(self, fast_confidence: float, mode: str = 'auto') -> bool: """ Determine whether to escalate to CTM analysis. Args: fast_confidence: Confidence from fast Tier 1 detector mode: Detection mode ('fast', 'deep_forensics', 'auto') Returns: True if CTM should be used """ if mode == 'deep_forensics': return True elif mode == 'fast': return False else: # auto mode return fast_confidence < self.confidence_threshold @property def is_loaded(self) -> bool: """Check if model is currently loaded""" return self._is_loaded def get_status(self) -> Dict[str, Any]: """Get agent status for monitoring""" return { 'loaded': self._is_loaded, 'device': self.device, 'max_thoughts': self.max_thoughts, 'confidence_threshold': self.confidence_threshold, 'ctm_available': CTM_AVAILABLE, 'checkpoint_path': self.checkpoint_path } # Singleton instance for lazy initialization _ctm_agent: Optional[CTMForensicAgent] = None def get_ctm_agent( checkpoint_path: Optional[str] = None, **kwargs ) -> CTMForensicAgent: """ Get or create singleton CTMForensicAgent instance. Args: checkpoint_path: Path to CTM checkpoint **kwargs: Additional arguments for CTMForensicAgent Returns: CTMForensicAgent instance """ global _ctm_agent if _ctm_agent is None: # Get default checkpoint path from environment if checkpoint_path is None: checkpoint_path = os.environ.get( 'CTM_CHECKPOINT_PATH', 'models/ctm_forensic.pth' ) # Get other config from environment max_thoughts = int(os.environ.get('CTM_MAX_THOUGHTS', 16)) confidence_threshold = float(os.environ.get('CTM_CONFIDENCE_THRESHOLD', 0.95)) _ctm_agent = CTMForensicAgent( checkpoint_path=checkpoint_path, max_thoughts=max_thoughts, confidence_threshold=confidence_threshold, **kwargs ) return _ctm_agent def reset_ctm_agent(): """Reset singleton agent (useful for testing)""" global _ctm_agent if _ctm_agent is not None: _ctm_agent.unload_model() _ctm_agent = None