backend/services/postprocess.py

"""
Post-Processing Service

Applies subtle photographic enhancements to the generated images:
- Local contrast enhancement (CLAHE)
- Film grain simulation
- Subtle vignette effect
- Color grading adjustments
"""
import logging
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
import cv2


logger = logging.getLogger(__name__)


def apply_clahe(image: Image.Image, clip_limit: float = 2.0) -> Image.Image:
    """
    Apply Contrast Limited Adaptive Histogram Equalization (CLAHE)
    
    This enhances local contrast without over-amplifying noise.
    
    Args:
        image: Input PIL Image
        clip_limit: Threshold for contrast limiting (higher = more contrast)
        
    Returns:
        Enhanced PIL Image
    """
    # Convert to numpy array
    img_np = np.array(image)
    
    # Convert to LAB color space
    lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
    
    # Split channels
    l, a, b = cv2.split(lab)
    
    # Apply CLAHE to L channel
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
    l_clahe = clahe.apply(l)
    
    # Merge channels
    lab_clahe = cv2.merge([l_clahe, a, b])
    
    # Convert back to RGB
    rgb = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB)
    
    return Image.fromarray(rgb)


def add_film_grain(
    image: Image.Image,
    intensity: float = 0.02,
    grain_size: float = 1.0
) -> Image.Image:
    """
    Add subtle film grain for a more organic look
    
    Args:
        image: Input PIL Image
        intensity: Strength of the grain effect (0.01-0.05 recommended)
        grain_size: Size of grain particles
        
    Returns:
        Image with film grain
    """
    img_np = np.array(image).astype(np.float32) / 255.0
    
    # Generate noise
    noise = np.random.normal(0, intensity, img_np.shape)
    
    # Optional: blur noise for larger grain
    if grain_size > 1.0:
        noise = cv2.GaussianBlur(noise, (0, 0), grain_size)
    
    # Add noise to image
    noisy = img_np + noise
    noisy = np.clip(noisy, 0, 1)
    
    # Convert back to uint8
    result = (noisy * 255).astype(np.uint8)
    
    return Image.fromarray(result)


def apply_vignette(
    image: Image.Image,
    strength: float = 0.3,
    radius: float = 0.8
) -> Image.Image:
    """
    Apply a subtle vignette effect
    
    Darkens the corners and edges of the image to draw focus to the center.
    
    Args:
        image: Input PIL Image
        strength: Vignette intensity (0-1)
        radius: Radius of the unaffected center area (0-1)
        
    Returns:
        Image with vignette
    """
    width, height = image.size
    img_np = np.array(image).astype(np.float32)
    
    # Create coordinate grids
    x = np.linspace(-1, 1, width)
    y = np.linspace(-1, 1, height)
    X, Y = np.meshgrid(x, y)
    
    # Calculate distance from center
    distance = np.sqrt(X**2 + Y**2)
    
    # Create vignette mask
    vignette = 1 - np.clip((distance - radius) / (1 - radius), 0, 1) * strength
    vignette = vignette[:, :, np.newaxis]  # Add channel dimension
    
    # Apply vignette
    result = img_np * vignette
    result = np.clip(result, 0, 255).astype(np.uint8)
    
    return Image.fromarray(result)


def enhance_colors(
    image: Image.Image,
    saturation: float = 1.1,
    contrast: float = 1.05,
    brightness: float = 1.0
) -> Image.Image:
    """
    Apply subtle color grading adjustments
    
    Args:
        image: Input PIL Image
        saturation: Saturation multiplier (1.0 = no change)
        contrast: Contrast multiplier (1.0 = no change)
        brightness: Brightness multiplier (1.0 = no change)
        
    Returns:
        Color-graded image
    """
    # Adjust saturation
    if saturation != 1.0:
        enhancer = ImageEnhance.Color(image)
        image = enhancer.enhance(saturation)
    
    # Adjust contrast
    if contrast != 1.0:
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(contrast)
    
    # Adjust brightness
    if brightness != 1.0:
        enhancer = ImageEnhance.Brightness(image)
        image = enhancer.enhance(brightness)
    
    return image


def sharpen_image(image: Image.Image, strength: float = 1.0) -> Image.Image:
    """
    Apply subtle sharpening
    
    Args:
        image: Input PIL Image
        strength: Sharpening strength (0-2 recommended)
        
    Returns:
        Sharpened image
    """
    if strength <= 0:
        return image
    
    # Use UnsharpMask for better control
    from PIL import ImageFilter
    
    # Blend between original and sharpened
    sharpened = image.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
    
    if strength < 1.0:
        # Blend with original
        return Image.blend(image, sharpened, strength)
    else:
        return sharpened


def postprocess_image(
    image: Image.Image,
    apply_contrast: bool = True,
    apply_grain: bool = True,
    apply_vignette_effect: bool = True,
    apply_color_grading: bool = True,
    apply_sharpening: bool = True
) -> Image.Image:
    """
    Apply complete post-processing pipeline
    
    This function orchestrates all post-processing effects in the optimal order:
    1. Local contrast enhancement (CLAHE)
    2. Color grading
    3. Sharpening
    4. Film grain
    5. Vignette
    
    Args:
        image: Input PIL Image
        apply_contrast: Enable local contrast enhancement
        apply_grain: Enable film grain
        apply_vignette_effect: Enable vignette
        apply_color_grading: Enable color adjustments
        apply_sharpening: Enable sharpening
        
    Returns:
        Post-processed PIL Image
    """
    logger.info("Starting post-processing")
    
    try:
        # Step 1: Local contrast
        if apply_contrast:
            logger.debug("Applying CLAHE")
            image = apply_clahe(image, clip_limit=2.0)
        
        # Step 2: Color grading
        if apply_color_grading:
            logger.debug("Applying color grading")
            image = enhance_colors(
                image,
                saturation=1.08,  # Slightly more saturated
                contrast=1.03,     # Slightly more contrast
                brightness=1.0     # No brightness change
            )
        
        # Step 3: Sharpening
        if apply_sharpening:
            logger.debug("Applying sharpening")
            image = sharpen_image(image, strength=0.6)
        
        # Step 4: Film grain
        if apply_grain:
            logger.debug("Adding film grain")
            image = add_film_grain(image, intensity=0.015, grain_size=1.2)
        
        # Step 5: Vignette
        if apply_vignette_effect:
            logger.debug("Applying vignette")
            image = apply_vignette(image, strength=0.2, radius=0.85)
        
        logger.info("Post-processing completed")
        return image
        
    except Exception as e:
        logger.error(f"Error in post-processing: {e}", exc_info=True)
        logger.warning("Returning original image")
        return image


def create_comparison(
    original: Image.Image,
    processed: Image.Image,
    padding: int = 10
) -> Image.Image:
    """
    Create a side-by-side comparison image
    
    Useful for visualizing before/after results.
    
    Args:
        original: Original image
        processed: Processed image
        padding: Space between images in pixels
        
    Returns:
        Combined comparison image
    """
    # Ensure both images are the same size
    if original.size != processed.size:
        processed = processed.resize(original.size, Image.LANCZOS)
    
    width, height = original.size
    
    # Create new image with space for both
    comparison = Image.new('RGB', (width * 2 + padding, height), color='white')
    
    # Paste images
    comparison.paste(original, (0, 0))
    comparison.paste(processed, (width + padding, 0))
    
    return comparison