src/pipeline/edge_case_handler.py

"""Edge case handling for multi-character detection and ambiguous images."""

import cv2
import numpy as np
import logging
from typing import List, Dict, Any, Optional, Tuple
from PIL import Image
import torch
from transformers import pipeline

from .base import PipelineStage, CharacterAttributes

logger = logging.getLogger(__name__)

class EdgeCaseHandler(PipelineStage):
    """Handles edge cases like multi-character images and ambiguous content."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__("EdgeCaseHandler", config)
        
        # Configuration
        self.face_detection_threshold = config.get('face_detection_threshold', 0.5) if config else 0.5
        self.multi_character_threshold = config.get('multi_character_threshold', 2) if config else 2
        self.min_face_size = config.get('min_face_size', 50) if config else 50
        self.occlusion_threshold = config.get('occlusion_threshold', 0.7) if config else 0.7
        
        # Models
        self.face_cascade = None
        self.person_detector = None
        self.quality_assessor = None
        
        self._initialize_detectors()
    
    def _initialize_detectors(self):
        """Initialize detection models."""
        try:
            # OpenCV face detector
            self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
            
            # Person detection using YOLO-like approach
            try:
                self.person_detector = pipeline(
                    "object-detection",
                    model="facebook/detr-resnet-50",
                    device=0 if torch.cuda.is_available() else -1
                )
            except Exception as e:
                logger.warning(f"Could not load person detector: {e}")
                self.person_detector = None
            
            logger.info("Edge case detectors initialized")
            
        except Exception as e:
            logger.error(f"Failed to initialize detectors: {e}")
    
    def detect_multiple_characters(self, image: Image.Image) -> Dict[str, Any]:
        """Detect if image contains multiple characters."""
        try:
            # Convert PIL to OpenCV format
            cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
            
            # Face detection
            faces = self.face_cascade.detectMultiScale(
                gray,
                scaleFactor=1.1,
                minNeighbors=5,
                minSize=(self.min_face_size, self.min_face_size)
            )
            
            face_count = len(faces)
            
            # Person detection using transformer model
            person_count = 0
            person_confidence = 0.0
            
            if self.person_detector:
                try:
                    detections = self.person_detector(image)
                    persons = [d for d in detections if d['label'] == 'person' and d['score'] > self.face_detection_threshold]
                    person_count = len(persons)
                    person_confidence = np.mean([p['score'] for p in persons]) if persons else 0.0
                except Exception as e:
                    logger.warning(f"Person detection failed: {e}")
            
            # Determine if multiple characters
            is_multi_character = (
                face_count >= self.multi_character_threshold or 
                person_count >= self.multi_character_threshold
            )
            
            # Calculate confidence
            detection_confidence = max(
                face_count / self.multi_character_threshold,
                person_confidence
            )
            
            return {
                'is_multi_character': is_multi_character,
                'face_count': face_count,
                'person_count': person_count,
                'detection_confidence': min(detection_confidence, 1.0),
                'face_locations': faces.tolist() if len(faces) > 0 else [],
                'recommendation': 'skip' if is_multi_character else 'process'
            }
            
        except Exception as e:
            logger.error(f"Multi-character detection failed: {e}")
            return {
                'is_multi_character': False,
                'face_count': 0,
                'person_count': 0,
                'detection_confidence': 0.0,
                'face_locations': [],
                'recommendation': 'process',
                'error': str(e)
            }
    
    def assess_image_quality(self, image: Image.Image) -> Dict[str, Any]:
        """Assess image quality and detect potential issues."""
        try:
            # Convert to numpy array
            img_array = np.array(image)
            
            # Basic quality metrics
            quality_metrics = {
                'resolution': img_array.shape[:2],
                'aspect_ratio': img_array.shape[1] / img_array.shape[0],
                'channels': img_array.shape[2] if len(img_array.shape) > 2 else 1
            }
            
            # Convert to grayscale for analysis
            if len(img_array.shape) == 3:
                gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            else:
                gray = img_array
            
            # Blur detection (Laplacian variance)
            blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()
            is_blurry = blur_score < 100  # Threshold for blur detection
            
            # Brightness analysis
            brightness = np.mean(gray)
            is_too_dark = brightness < 50
            is_too_bright = brightness > 200
            
            # Contrast analysis
            contrast = gray.std()
            is_low_contrast = contrast < 30
            
            # Noise detection (using standard deviation in small patches)
            noise_score = self._estimate_noise(gray)
            is_noisy = noise_score > 20
            
            # Overall quality score
            quality_score = self._calculate_quality_score(
                blur_score, brightness, contrast, noise_score
            )
            
            # Determine recommendation
            issues = []
            if is_blurry:
                issues.append('blurry')
            if is_too_dark:
                issues.append('too_dark')
            if is_too_bright:
                issues.append('too_bright')
            if is_low_contrast:
                issues.append('low_contrast')
            if is_noisy:
                issues.append('noisy')
            
            recommendation = 'skip' if len(issues) >= 3 or quality_score < 0.3 else 'process'
            
            return {
                'quality_score': quality_score,
                'blur_score': blur_score,
                'brightness': brightness,
                'contrast': contrast,
                'noise_score': noise_score,
                'issues': issues,
                'is_acceptable': quality_score >= 0.3,
                'recommendation': recommendation,
                'metrics': quality_metrics
            }
            
        except Exception as e:
            logger.error(f"Quality assessment failed: {e}")
            return {
                'quality_score': 0.5,
                'issues': ['assessment_failed'],
                'is_acceptable': True,
                'recommendation': 'process',
                'error': str(e)
            }
    
    def _estimate_noise(self, gray_image: np.ndarray) -> float:
        """Estimate noise level in image."""
        try:
            # Use Laplacian to detect edges, then measure noise in non-edge regions
            laplacian = cv2.Laplacian(gray_image, cv2.CV_64F)
            edges = np.abs(laplacian) > np.percentile(np.abs(laplacian), 90)
            
            # Calculate noise in non-edge regions
            non_edge_regions = gray_image[~edges]
            if len(non_edge_regions) > 0:
                return np.std(non_edge_regions)
            else:
                return np.std(gray_image)
        except:
            return np.std(gray_image)
    
    def _calculate_quality_score(self, blur_score: float, brightness: float, 
                                contrast: float, noise_score: float) -> float:
        """Calculate overall quality score from individual metrics."""
        # Normalize scores to 0-1 range
        blur_norm = min(blur_score / 500, 1.0)  # Higher is better
        brightness_norm = 1.0 - abs(brightness - 128) / 128  # Closer to 128 is better
        contrast_norm = min(contrast / 100, 1.0)  # Higher is better
        noise_norm = max(0, 1.0 - noise_score / 50)  # Lower is better
        
        # Weighted average
        weights = [0.3, 0.2, 0.3, 0.2]  # blur, brightness, contrast, noise
        scores = [blur_norm, brightness_norm, contrast_norm, noise_norm]
        
        return sum(w * s for w, s in zip(weights, scores))
    
    def detect_occlusion(self, image: Image.Image) -> Dict[str, Any]:
        """Detect if character is significantly occluded."""
        try:
            # Convert to OpenCV format
            cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
            
            # Detect faces to estimate visible character area
            faces = self.face_cascade.detectMultiScale(gray, 1.1, 5)
            
            if len(faces) == 0:
                # No face detected - might be occluded or not a character
                return {
                    'is_occluded': True,
                    'occlusion_ratio': 1.0,
                    'visible_face_area': 0,
                    'recommendation': 'skip',
                    'reason': 'no_face_detected'
                }
            
            # Calculate largest face area
            largest_face = max(faces, key=lambda f: f[2] * f[3])
            face_area = largest_face[2] * largest_face[3]
            image_area = image.width * image.height
            
            # Estimate occlusion based on face size relative to image
            expected_face_ratio = 0.1  # Expected minimum face area ratio
            actual_face_ratio = face_area / image_area
            
            occlusion_ratio = max(0, 1 - (actual_face_ratio / expected_face_ratio))
            is_occluded = occlusion_ratio > self.occlusion_threshold
            
            return {
                'is_occluded': is_occluded,
                'occlusion_ratio': occlusion_ratio,
                'visible_face_area': face_area,
                'face_ratio': actual_face_ratio,
                'recommendation': 'skip' if is_occluded else 'process',
                'face_location': largest_face.tolist()
            }
            
        except Exception as e:
            logger.error(f"Occlusion detection failed: {e}")
            return {
                'is_occluded': False,
                'occlusion_ratio': 0.0,
                'visible_face_area': 0,
                'recommendation': 'process',
                'error': str(e)
            }
    
    def analyze_image_content(self, image: Image.Image) -> Dict[str, Any]:
        """Comprehensive analysis of image content for edge cases."""
        # Run all detection methods
        multi_char_result = self.detect_multiple_characters(image)
        quality_result = self.assess_image_quality(image)
        occlusion_result = self.detect_occlusion(image)
        
        # Combine results
        edge_cases = []
        if multi_char_result['is_multi_character']:
            edge_cases.append('multiple_characters')
        if not quality_result['is_acceptable']:
            edge_cases.append('poor_quality')
        if occlusion_result['is_occluded']:
            edge_cases.append('occluded')
        
        # Overall recommendation
        recommendations = [
            multi_char_result['recommendation'],
            quality_result['recommendation'],
            occlusion_result['recommendation']
        ]
        
        skip_count = recommendations.count('skip')
        overall_recommendation = 'skip' if skip_count >= 2 else 'process'
        
        # Confidence score for the recommendation
        confidence_scores = [
            multi_char_result.get('detection_confidence', 0.5),
            quality_result.get('quality_score', 0.5),
            1.0 - occlusion_result.get('occlusion_ratio', 0.5)
        ]
        overall_confidence = np.mean(confidence_scores)
        
        return {
            'edge_cases': edge_cases,
            'has_edge_cases': len(edge_cases) > 0,
            'recommendation': overall_recommendation,
            'confidence': overall_confidence,
            'details': {
                'multi_character': multi_char_result,
                'quality': quality_result,
                'occlusion': occlusion_result
            },
            'processing_advice': self._generate_processing_advice(edge_cases, overall_recommendation)
        }
    
    def _generate_processing_advice(self, edge_cases: List[str], recommendation: str) -> List[str]:
        """Generate advice for handling detected edge cases."""
        advice = []
        
        if 'multiple_characters' in edge_cases:
            advice.append("Image contains multiple characters - consider cropping to single character")
        
        if 'poor_quality' in edge_cases:
            advice.append("Image quality is poor - consider preprocessing or skipping")
        
        if 'occluded' in edge_cases:
            advice.append("Character appears occluded - extraction may be incomplete")
        
        if recommendation == 'skip':
            advice.append("Recommend skipping this image due to multiple edge cases")
        elif len(edge_cases) > 0:
            advice.append("Proceed with caution - edge cases detected but may be processable")
        else:
            advice.append("Image appears suitable for character extraction")
        
        return advice
    
    def process(self, input_data: Any) -> Dict[str, Any]:
        """Process image and return edge case analysis."""
        if isinstance(input_data, Image.Image):
            return self.analyze_image_content(input_data)
        elif isinstance(input_data, str):
            # Assume it's an image path
            try:
                image = Image.open(input_data)
                return self.analyze_image_content(image)
            except Exception as e:
                return {
                    'edge_cases': ['invalid_image'],
                    'has_edge_cases': True,
                    'recommendation': 'skip',
                    'confidence': 0.0,
                    'error': str(e)
                }
        else:
            raise ValueError("EdgeCaseHandler expects PIL Image or image path")
    
    def validate_input(self, input_data: Any) -> bool:
        """Validate input data."""
        return isinstance(input_data, (Image.Image, str))