src/pipeline/preprocessing.py

"""
Preprocessing pipeline for EL images.

This is CRITICAL for production robustness. Real-world EL images vary wildly:
- Factory cameras with different exposure settings
- Degraded modules with overall lower luminescence  
- Overexposed images from new, high-efficiency cells
- Noisy images from long-exposure captures

The preprocessing pipeline normalizes ALL inputs to look similar,
making the model's job much easier.

Pipeline:
1. Convert to grayscale (if not already)
2. CLAHE: Contrast Limited Adaptive Histogram Equalization
   - Enhances local contrast without amplifying noise
   - tileGridSize=(8,8): processes image in 8x8 blocks
   - clipLimit=2.0: prevents over-enhancement
3. Intensity normalization: scale to [0, 1]
4. Resize to consistent input size
"""

import cv2
import numpy as np
from typing import Tuple, Optional


class ELPreprocessor:
    """
    Production preprocessor for EL images.
    
    Handles: dark images, bright images, varying sizes, noise.
    Produces: consistent normalized grayscale output.
    """
    
    def __init__(
        self,
        target_size: Tuple[int, int] = (512, 512),
        clahe_clip_limit: float = 2.0,
        clahe_tile_grid: Tuple[int, int] = (8, 8),
        denoise: bool = True,
        denoise_strength: int = 7,
    ):
        """
        Args:
            target_size: (H, W) output size
            clahe_clip_limit: CLAHE contrast limit. Higher = more enhancement.
                2.0 is standard; use 3.0-4.0 for very dark images.
            clahe_tile_grid: CLAHE tile size. (8,8) is standard.
                Smaller tiles = more local contrast enhancement.
            denoise: Apply non-local means denoising
            denoise_strength: Denoising filter strength (higher = more smoothing)
        """
        self.target_size = target_size
        self.clahe = cv2.createCLAHE(
            clipLimit=clahe_clip_limit, 
            tileGridSize=clahe_tile_grid
        )
        self.denoise = denoise
        self.denoise_strength = denoise_strength
    
    def process(self, image: np.ndarray) -> np.ndarray:
        """
        Full preprocessing pipeline.
        
        Args:
            image: Input image (any format: RGB, grayscale, any size, any bit depth)
            
        Returns:
            Preprocessed grayscale image, shape (H, W), dtype float32, range [0, 1]
        """
        # Step 1: Convert to grayscale
        gray = self._to_grayscale(image)
        
        # Step 2: Denoise (before CLAHE to prevent noise amplification)
        if self.denoise:
            gray = self._denoise(gray)
        
        # Step 3: CLAHE — adaptive contrast enhancement
        # This is the most important step: makes dark images visible,
        # prevents bright images from being washed out
        enhanced = self.clahe.apply(gray)
        
        # Step 4: Intensity normalization to [0, 1]
        normalized = self._normalize_intensity(enhanced)
        
        # Step 5: Resize to target size
        resized = cv2.resize(
            normalized, 
            (self.target_size[1], self.target_size[0]),  # cv2 uses (W, H)
            interpolation=cv2.INTER_LINEAR
        )
        
        return resized.astype(np.float32)
    
    def process_for_model(self, image: np.ndarray) -> np.ndarray:
        """
        Process image and prepare for model input.
        
        Returns:
            Shape (1, H, W) float32, normalized with mean=0.5, std=0.5
        """
        processed = self.process(image)
        # Normalize to match training: (x - 0.5) / 0.5
        model_input = (processed - 0.5) / 0.5
        return model_input[np.newaxis, ...]  # Add channel dim: (1, H, W)
    
    def _to_grayscale(self, image: np.ndarray) -> np.ndarray:
        """Convert any format to 8-bit grayscale."""
        if image is None or image.size == 0:
            raise ValueError("Empty or None image received")
        
        if image.ndim == 3:
            if image.shape[2] == 4:  # RGBA
                image = cv2.cvtColor(image, cv2.COLOR_RGBA2GRAY)
            elif image.shape[2] == 3:  # RGB/BGR
                image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:
                image = image[:, :, 0]  # Take first channel
        
        # Handle 16-bit images (some industrial cameras)
        if image.dtype == np.uint16:
            image = (image / 256).astype(np.uint8)
        elif image.dtype == np.float32 or image.dtype == np.float64:
            if image.max() <= 1.0:
                image = (image * 255).astype(np.uint8)
            else:
                image = np.clip(image, 0, 255).astype(np.uint8)
        elif image.dtype != np.uint8:
            image = image.astype(np.uint8)
        
        return image
    
    def _denoise(self, gray: np.ndarray) -> np.ndarray:
        """
        Non-local means denoising.
        
        Trade-off: removes sensor noise but can blur thin cracks.
        strength=7 is conservative; increase for very noisy images.
        """
        return cv2.fastNlMeansDenoising(
            gray, 
            h=self.denoise_strength,
            templateWindowSize=7,
            searchWindowSize=21
        )
    
    def _normalize_intensity(self, image: np.ndarray) -> np.ndarray:
        """
        Percentile-based intensity normalization.
        
        Why percentile instead of min-max?
        - Hot/dead pixels don't skew the range
        - More robust for real-world images
        - 1st and 99th percentile clips extreme outliers
        """
        p_low = np.percentile(image, 1)
        p_high = np.percentile(image, 99)
        
        if p_high - p_low < 10:  # Very low contrast image
            # Fallback to full-range normalization
            p_low = image.min()
            p_high = image.max()
        
        if p_high == p_low:
            return np.zeros_like(image, dtype=np.float32)
        
        normalized = (image.astype(np.float32) - p_low) / (p_high - p_low)
        return np.clip(normalized, 0, 1)
    
    def get_image_stats(self, image: np.ndarray) -> dict:
        """
        Compute diagnostic statistics for an EL image.
        Useful for quality assessment and adaptive parameter tuning.
        """
        gray = self._to_grayscale(image)
        return {
            "mean_intensity": float(gray.mean()),
            "std_intensity": float(gray.std()),
            "min_intensity": int(gray.min()),
            "max_intensity": int(gray.max()),
            "dynamic_range": int(gray.max()) - int(gray.min()),
            "is_dark": gray.mean() < 50,
            "is_bright": gray.mean() > 200,
            "is_low_contrast": gray.std() < 20,
            "shape": gray.shape,
        }


def batch_preprocess(
    images: list, 
    preprocessor: Optional[ELPreprocessor] = None,
) -> np.ndarray:
    """
    Preprocess a batch of images for model input.
    
    Returns:
        (N, 1, H, W) float32 array ready for torch.from_numpy()
    """
    if preprocessor is None:
        preprocessor = ELPreprocessor()
    
    batch = []
    for img in images:
        processed = preprocessor.process_for_model(img)
        batch.append(processed)
    
    return np.stack(batch, axis=0)