src/pipeline/preprocessor.py

"""Preprocessing pipeline for style normalization and occlusion handling."""

import cv2
import numpy as np
import logging
import time
from typing import Dict, Any, Optional, Tuple, List, Union
from PIL import Image, ImageEnhance, ImageFilter
import torch
import torchvision.transforms as transforms
from sklearn.cluster import KMeans

from .base import PipelineStage

logger = logging.getLogger(__name__)

class StyleNormalizer:
    """Normalizes different art styles for consistent processing."""
    
    def __init__(self, config: Dict[str, Any]):
        self.target_size = config.get('target_size', (512, 512))
        self.normalize_brightness = config.get('normalize_brightness', True)
        self.normalize_contrast = config.get('normalize_contrast', True)
        self.normalize_saturation = config.get('normalize_saturation', True)
        
        # Style detection thresholds
        self.anime_threshold = config.get('anime_threshold', 0.7)
        self.realistic_threshold = config.get('realistic_threshold', 0.6)
    
    def detect_art_style(self, image: Image.Image) -> Dict[str, Any]:
        """Detect the art style of the image."""
        try:
            # Convert to numpy for analysis
            img_array = np.array(image)
            
            # Color analysis
            color_variance = np.var(img_array, axis=(0, 1))
            avg_color_variance = np.mean(color_variance)
            
            # Edge analysis
            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            edge_density = np.sum(edges > 0) / edges.size
            
            # Saturation analysis
            hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
            saturation = hsv[:, :, 1]
            avg_saturation = np.mean(saturation)
            
            # Style classification heuristics
            anime_score = 0.0
            realistic_score = 0.0
            
            # High saturation + low color variance = anime-like
            if avg_saturation > 100 and avg_color_variance < 2000:
                anime_score += 0.4
            
            # Sharp edges = anime-like
            if edge_density > 0.1:
                anime_score += 0.3
            
            # High color variance + moderate saturation = realistic
            if avg_color_variance > 3000 and 50 < avg_saturation < 150:
                realistic_score += 0.5
            
            # Determine primary style
            if anime_score > self.anime_threshold:
                style = 'anime'
                confidence = anime_score
            elif realistic_score > self.realistic_threshold:
                style = 'realistic'
                confidence = realistic_score
            else:
                style = 'mixed'
                confidence = max(anime_score, realistic_score)
            
            return {
                'style': style,
                'confidence': confidence,
                'metrics': {
                    'color_variance': avg_color_variance,
                    'edge_density': edge_density,
                    'saturation': avg_saturation,
                    'anime_score': anime_score,
                    'realistic_score': realistic_score
                }
            }
            
        except Exception as e:
            logger.error(f"Style detection failed: {e}")
            return {
                'style': 'unknown',
                'confidence': 0.0,
                'error': str(e)
            }
    
    def normalize_image(self, image: Image.Image, style_info: Dict[str, Any]) -> Image.Image:
        """Normalize image based on detected style."""
        try:
            normalized = image.copy()
            
            # Resize to target size
            normalized = normalized.resize(self.target_size, Image.Resampling.LANCZOS)
            
            # Style-specific normalization
            style = style_info.get('style', 'unknown')
            
            if style == 'anime':
                normalized = self._normalize_anime_style(normalized)
            elif style == 'realistic':
                normalized = self._normalize_realistic_style(normalized)
            else:
                normalized = self._normalize_generic_style(normalized)
            
            return normalized
            
        except Exception as e:
            logger.error(f"Image normalization failed: {e}")
            return image.resize(self.target_size, Image.Resampling.LANCZOS)
    
    def _normalize_anime_style(self, image: Image.Image) -> Image.Image:
        """Normalize anime-style images."""
        # Anime images often have high saturation and sharp edges
        # Slightly reduce saturation for better CLIP processing
        if self.normalize_saturation:
            enhancer = ImageEnhance.Color(image)
            image = enhancer.enhance(0.9)  # Reduce saturation slightly
        
        # Enhance contrast for better feature detection
        if self.normalize_contrast:
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(1.1)
        
        return image
    
    def _normalize_realistic_style(self, image: Image.Image) -> Image.Image:
        """Normalize realistic-style images."""
        # Realistic images may need brightness and contrast adjustment
        if self.normalize_brightness:
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(1.05)  # Slight brightness boost
        
        if self.normalize_contrast:
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(1.15)  # Enhance contrast
        
        # Slight sharpening for better feature detection
        image = image.filter(ImageFilter.UnsharpMask(radius=1, percent=110, threshold=3))
        
        return image
    
    def _normalize_generic_style(self, image: Image.Image) -> Image.Image:
        """Generic normalization for unknown styles."""
        # Conservative normalization
        if self.normalize_brightness:
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(1.02)
        
        if self.normalize_contrast:
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(1.05)
        
        return image

class OcclusionHandler:
    """Handles occluded or partially visible characters."""
    
    def __init__(self, config: Dict[str, Any]):
        self.min_visible_ratio = config.get('min_visible_ratio', 0.3)
        self.inpainting_enabled = config.get('inpainting_enabled', False)
        
        # Face detection for occlusion analysis
        self.face_cascade = cv2.CascadeClassifier(
            cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
        )
    
    def detect_occlusion_regions(self, image: Image.Image) -> Dict[str, Any]:
        """Detect occluded regions in the image."""
        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
            faces = self.face_cascade.detectMultiScale(gray, 1.1, 5)
            
            if len(faces) == 0:
                return {
                    'has_occlusion': True,
                    'occlusion_type': 'no_face_detected',
                    'visible_ratio': 0.0,
                    'recommendation': 'skip'
                }
            
            # Analyze largest face
            largest_face = max(faces, key=lambda f: f[2] * f[3])
            x, y, w, h = largest_face
            
            # Extract face region
            face_region = gray[y:y+h, x:x+w]
            
            # Detect occlusion using edge density and uniformity
            edges = cv2.Canny(face_region, 50, 150)
            edge_density = np.sum(edges > 0) / edges.size
            
            # Check for uniform regions (potential occlusion)
            uniform_threshold = 10
            uniform_regions = np.std(face_region) < uniform_threshold
            
            # Estimate visible ratio
            visible_ratio = edge_density * (1.0 if not uniform_regions else 0.5)
            
            has_occlusion = visible_ratio < self.min_visible_ratio
            
            return {
                'has_occlusion': has_occlusion,
                'occlusion_type': 'partial' if has_occlusion else 'none',
                'visible_ratio': visible_ratio,
                'face_location': largest_face.tolist(),
                'edge_density': edge_density,
                'recommendation': 'process_with_caution' if has_occlusion else 'process'
            }
            
        except Exception as e:
            logger.error(f"Occlusion detection failed: {e}")
            return {
                'has_occlusion': False,
                'occlusion_type': 'unknown',
                'visible_ratio': 1.0,
                'recommendation': 'process',
                'error': str(e)
            }
    
    def enhance_occluded_image(self, image: Image.Image, 
                              occlusion_info: Dict[str, Any]) -> Image.Image:
        """Enhance occluded images for better attribute extraction."""
        try:
            enhanced = image.copy()
            
            if occlusion_info.get('has_occlusion', False):
                # Apply enhancement based on occlusion type
                occlusion_type = occlusion_info.get('occlusion_type', 'unknown')
                
                if occlusion_type == 'partial':
                    # Enhance contrast and sharpness for partially occluded images
                    enhancer = ImageEnhance.Contrast(enhanced)
                    enhanced = enhancer.enhance(1.3)
                    
                    enhancer = ImageEnhance.Sharpness(enhanced)
                    enhanced = enhancer.enhance(1.2)
                    
                    # Apply unsharp mask
                    enhanced = enhanced.filter(ImageFilter.UnsharpMask(radius=2, percent=150, threshold=3))
                
                elif occlusion_type == 'no_face_detected':
                    # Try to enhance overall image quality
                    enhancer = ImageEnhance.Brightness(enhanced)
                    enhanced = enhancer.enhance(1.1)
                    
                    enhancer = ImageEnhance.Contrast(enhanced)
                    enhanced = enhancer.enhance(1.2)
            
            return enhanced
            
        except Exception as e:
            logger.error(f"Image enhancement failed: {e}")
            return image

class ImagePreprocessor(PipelineStage):
    """Comprehensive image preprocessing for improved attribute extraction."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__("ImagePreprocessor", config)
        
        # Configuration
        if config:
            self.enable_style_normalization = config.get('enable_style_normalization', True)
            self.enable_occlusion_handling = config.get('enable_occlusion_handling', True)
            self.enable_quality_enhancement = config.get('enable_quality_enhancement', True)
            self.target_size = config.get('target_size', (512, 512))
        else:
            self.enable_style_normalization = True
            self.enable_occlusion_handling = True
            self.enable_quality_enhancement = True
            self.target_size = (512, 512)
        
        # Components
        self.style_normalizer = StyleNormalizer(config or {})
        self.occlusion_handler = OcclusionHandler(config or {})
        
        # Preprocessing statistics
        self.stats = {
            'processed': 0,
            'style_normalized': 0,
            'occlusion_handled': 0,
            'quality_enhanced': 0,
            'skipped': 0
        }
    
    def preprocess_image(self, image: Image.Image) -> Dict[str, Any]:
        """Comprehensive image preprocessing."""
        try:
            start_time = time.time()
            processed_image = image.copy()
            preprocessing_info = {
                'original_size': image.size,
                'steps_applied': [],
                'style_info': {},
                'occlusion_info': {},
                'quality_info': {}
            }
            
            # Step 1: Style detection and normalization
            if self.enable_style_normalization:
                style_info = self.style_normalizer.detect_art_style(processed_image)
                processed_image = self.style_normalizer.normalize_image(processed_image, style_info)
                preprocessing_info['style_info'] = style_info
                preprocessing_info['steps_applied'].append('style_normalization')
                self.stats['style_normalized'] += 1
            
            # Step 2: Occlusion detection and handling
            if self.enable_occlusion_handling:
                occlusion_info = self.occlusion_handler.detect_occlusion_regions(processed_image)
                
                if occlusion_info.get('has_occlusion', False):
                    processed_image = self.occlusion_handler.enhance_occluded_image(
                        processed_image, occlusion_info
                    )
                    preprocessing_info['steps_applied'].append('occlusion_handling')
                    self.stats['occlusion_handled'] += 1
                
                preprocessing_info['occlusion_info'] = occlusion_info
            
            # Step 3: Quality enhancement
            if self.enable_quality_enhancement:
                quality_info = self._assess_and_enhance_quality(processed_image)
                if quality_info.get('enhanced', False):
                    processed_image = quality_info['enhanced_image']
                    preprocessing_info['steps_applied'].append('quality_enhancement')
                    self.stats['quality_enhanced'] += 1
                
                preprocessing_info['quality_info'] = quality_info
            
            # Final validation
            should_skip = self._should_skip_image(preprocessing_info)
            if should_skip:
                self.stats['skipped'] += 1
                preprocessing_info['recommendation'] = 'skip'
                preprocessing_info['skip_reason'] = should_skip
            else:
                preprocessing_info['recommendation'] = 'process'
            
            processing_time = time.time() - start_time
            preprocessing_info['processing_time'] = processing_time
            
            self.stats['processed'] += 1
            
            return {
                'processed_image': processed_image,
                'preprocessing_info': preprocessing_info,
                'should_skip': should_skip is not False
            }
            
        except Exception as e:
            logger.error(f"Preprocessing failed: {e}")
            return {
                'processed_image': image,
                'preprocessing_info': {'error': str(e)},
                'should_skip': False
            }
    
    def _assess_and_enhance_quality(self, image: Image.Image) -> Dict[str, Any]:
        """Assess and enhance image quality."""
        try:
            # Convert to numpy for analysis
            img_array = np.array(image)
            
            # Quality metrics
            brightness = np.mean(img_array)
            contrast = np.std(img_array)
            
            # Blur detection
            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()
            
            quality_issues = []
            enhanced_image = image.copy()
            
            # Brightness correction
            if brightness < 80:
                quality_issues.append('too_dark')
                enhancer = ImageEnhance.Brightness(enhanced_image)
                enhanced_image = enhancer.enhance(1.2)
            elif brightness > 180:
                quality_issues.append('too_bright')
                enhancer = ImageEnhance.Brightness(enhanced_image)
                enhanced_image = enhancer.enhance(0.9)
            
            # Contrast enhancement
            if contrast < 40:
                quality_issues.append('low_contrast')
                enhancer = ImageEnhance.Contrast(enhanced_image)
                enhanced_image = enhancer.enhance(1.3)
            
            # Blur correction
            if blur_score < 100:
                quality_issues.append('blurry')
                enhanced_image = enhanced_image.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=3))
            
            return {
                'quality_issues': quality_issues,
                'enhanced': len(quality_issues) > 0,
                'enhanced_image': enhanced_image if quality_issues else image,
                'metrics': {
                    'brightness': brightness,
                    'contrast': contrast,
                    'blur_score': blur_score
                }
            }
            
        except Exception as e:
            logger.error(f"Quality assessment failed: {e}")
            return {
                'quality_issues': [],
                'enhanced': False,
                'enhanced_image': image,
                'error': str(e)
            }
    
    def _should_skip_image(self, preprocessing_info: Dict[str, Any]) -> Union[str, bool]:
        """Determine if image should be skipped based on preprocessing results."""
        # Check occlusion
        occlusion_info = preprocessing_info.get('occlusion_info', {})
        if occlusion_info.get('recommendation') == 'skip':
            return 'severe_occlusion'
        
        # Check style confidence
        style_info = preprocessing_info.get('style_info', {})
        if style_info.get('confidence', 1.0) < 0.2:
            return 'unrecognizable_style'
        
        # Check quality issues
        quality_info = preprocessing_info.get('quality_info', {})
        quality_issues = quality_info.get('quality_issues', [])
        if len(quality_issues) >= 3:
            return 'poor_quality'
        
        return False
    
    def batch_preprocess(self, images: List[Image.Image]) -> List[Dict[str, Any]]:
        """Preprocess a batch of images efficiently."""
        results = []
        
        for i, image in enumerate(images):
            try:
                result = self.preprocess_image(image)
                results.append(result)
            except Exception as e:
                logger.error(f"Failed to preprocess image {i}: {e}")
                results.append({
                    'processed_image': image,
                    'preprocessing_info': {'error': str(e)},
                    'should_skip': True
                })
        
        return results
    
    def get_preprocessing_stats(self) -> Dict[str, Any]:
        """Get preprocessing statistics."""
        total = self.stats['processed']
        
        return {
            'total_processed': total,
            'style_normalized': self.stats['style_normalized'],
            'occlusion_handled': self.stats['occlusion_handled'],
            'quality_enhanced': self.stats['quality_enhanced'],
            'skipped': self.stats['skipped'],
            'skip_rate': self.stats['skipped'] / total if total > 0 else 0,
            'enhancement_rate': (self.stats['style_normalized'] + self.stats['quality_enhanced']) / total if total > 0 else 0
        }
    
    def process(self, input_data: Any) -> Any:
        """Process image preprocessing."""
        if isinstance(input_data, Image.Image):
            return self.preprocess_image(input_data)
        elif isinstance(input_data, list) and all(isinstance(img, Image.Image) for img in input_data):
            return self.batch_preprocess(input_data)
        elif isinstance(input_data, dict) and input_data.get('operation') == 'stats':
            return self.get_preprocessing_stats()
        else:
            raise ValueError("ImagePreprocessor expects PIL Image, list of images, or stats operation")
    
    def validate_input(self, input_data: Any) -> bool:
        """Validate input data."""
        return (
            isinstance(input_data, Image.Image) or
            (isinstance(input_data, list) and all(isinstance(img, Image.Image) for img in input_data)) or
            (isinstance(input_data, dict) and 'operation' in input_data)
        )