preprocess.py

"""
Naari Studio VTON - Preprocessing Module
=========================================
Shared image preprocessing utilities for all engines.

v3.1.0 - Enhanced with background artifact fixes

Author: AnahataSri (Naari Studio)
License: MIT
"""

import logging
from typing import Tuple, Optional
from PIL import Image, ImageStat, ImageFilter
import numpy as np

from config import MAX_IMAGE_SIZE, DIFFUSION_SIZE

logger = logging.getLogger(__name__)


# =============================================================================
# IMAGE VALIDATION
# =============================================================================

def validate_image(image: Image.Image, name: str = "Image") -> Tuple[bool, str]:
    """
    Validate an input image.
    
    Args:
        image: PIL Image to validate
        name: Name for error messages
    
    Returns:
        Tuple of (is_valid, error_message)
    """
    if image is None:
        return False, f"❌ {name} is required"
    
    # Check if it's a valid PIL Image
    try:
        width, height = image.size
    except:
        return False, f"❌ {name} is not a valid image"
    
    # Check minimum size
    if width < 64 or height < 64:
        return False, f"❌ {name} is too small (minimum 64x64)"
    
    # Check maximum size
    if width > 4096 or height > 4096:
        return False, f"❌ {name} is too large (maximum 4096x4096)"
    
    return True, ""


def validate_inputs(
    person_image: Image.Image,
    garment_image: Image.Image
) -> Tuple[bool, str]:
    """
    Validate both person and garment images.
    
    Args:
        person_image: PIL Image of person
        garment_image: PIL Image of garment
    
    Returns:
        Tuple of (is_valid, error_message)
    """
    is_valid, msg = validate_image(person_image, "Person image")
    if not is_valid:
        return False, msg
    
    is_valid, msg = validate_image(garment_image, "Garment image")
    if not is_valid:
        return False, msg
    
    return True, ""


# =============================================================================
# BACKGROUND ARTIFACT FIX
# =============================================================================

def detect_dark_background(image: Image.Image, threshold: int = 60) -> bool:
    """
    Detect if image has a predominantly dark background.
    
    Dark backgrounds can be misinterpreted as hair by the model.
    
    Args:
        image: PIL Image to analyze
        threshold: Brightness threshold (0-255). Below this is considered dark.
    
    Returns:
        True if background appears dark
    """
    # Sample edges to check background
    width, height = image.size
    
    # Sample strips from edges (10% of width/height)
    edge_width = max(10, width // 10)
    edge_height = max(10, height // 10)
    
    # Get edge regions
    top = image.crop((0, 0, width, edge_height))
    bottom = image.crop((0, height - edge_height, width, height))
    left = image.crop((0, 0, edge_width, height))
    right = image.crop((width - edge_width, 0, width, height))
    
    # Calculate average brightness of edges
    edge_brightnesses = []
    for edge in [top, bottom, left, right]:
        stat = ImageStat.Stat(edge.convert('L'))
        edge_brightnesses.append(stat.mean[0])
    
    avg_edge_brightness = sum(edge_brightnesses) / len(edge_brightnesses)
    
    logger.info(f"Edge brightness: {avg_edge_brightness:.1f}")
    
    return avg_edge_brightness < threshold


def replace_dark_background(
    image: Image.Image,
    target_color: Tuple[int, int, int] = (230, 230, 230),
    edge_threshold: int = 60
) -> Image.Image:
    """
    Replace dark background with a neutral light color.
    
    This prevents the model from misinterpreting dark backgrounds as hair.
    
    Args:
        image: PIL Image to process
        target_color: RGB color to use for background (default: light gray)
        edge_threshold: Brightness threshold for edge detection
    
    Returns:
        Image with replaced background
    """
    # Convert to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Create a simple edge-based mask
    # This is a lightweight approach that doesn't require heavy models
    img_array = np.array(image)
    
    # Convert to grayscale for brightness detection
    gray = np.array(image.convert('L'))
    
    # Create mask where pixels are darker than threshold
    dark_mask = gray < edge_threshold
    
    # Dilate edges slightly to ensure full coverage
    from scipy import ndimage
    dark_mask = ndimage.binary_dilation(dark_mask, iterations=2)
    
    # Only replace pixels that are likely background (edges + dark)
    width, height = image.size
    edge_margin = min(width, height) // 8
    
    # Create edge mask
    edge_mask = np.zeros_like(dark_mask)
    edge_mask[:edge_margin, :] = True  # Top
    edge_mask[-edge_margin:, :] = True  # Bottom
    edge_mask[:, :edge_margin] = True  # Left
    edge_mask[:, -edge_margin:] = True  # Right
    
    # Combine: only replace dark pixels at edges
    replace_mask = dark_mask & edge_mask
    
    # Replace background
    result = img_array.copy()
    result[replace_mask] = target_color
    
    logger.info(f"Replaced {replace_mask.sum()} background pixels")
    
    return Image.fromarray(result)


# =============================================================================
# IMAGE NORMALIZATION
# =============================================================================

def normalize_image(
    image: Image.Image,
    max_size: int = None,
    target_mode: str = "RGB",
    fix_background: bool = True
) -> Image.Image:
    """
    Normalize an image for processing.
    
    - Converts to target color mode
    - Resizes if larger than max_size while maintaining aspect ratio
    - Optionally fixes dark background artifacts
    
    Args:
        image: PIL Image to normalize
        max_size: Maximum dimension (uses config default if None)
        target_mode: Target color mode (default RGB)
        fix_background: Whether to fix dark background artifacts
    
    Returns:
        Normalized PIL Image
    """
    max_size = max_size or MAX_IMAGE_SIZE
    
    # Convert color mode
    if image.mode != target_mode:
        image = image.convert(target_mode)
    
    # Fix dark background if enabled
    if fix_background:
        if detect_dark_background(image):
            logger.info("Dark background detected, applying fix")
            try:
                image = replace_dark_background(image)
            except Exception as e:
                logger.warning(f"Background fix failed: {e}, continuing without fix")
    
    # Resize if too large
    width, height = image.size
    if max(width, height) > max_size:
        ratio = max_size / max(width, height)
        new_width = int(width * ratio)
        new_height = int(height * ratio)
        image = image.resize((new_width, new_height), Image.LANCZOS)
        logger.info(f"Resized image from {width}x{height} to {new_width}x{new_height}")
    
    return image


def resize_for_diffusion(
    image: Image.Image,
    target_size: int = None
) -> Image.Image:
    """
    Resize image for diffusion models.
    
    Ensures dimensions are multiples of 8 for VAE compatibility.
    
    Args:
        image: PIL Image to resize
        target_size: Target max dimension (uses config default if None)
    
    Returns:
        Resized PIL Image with dimensions divisible by 8
    """
    target_size = target_size or DIFFUSION_SIZE
    
    width, height = image.size
    
    # Calculate new size maintaining aspect ratio
    if max(width, height) > target_size:
        ratio = target_size / max(width, height)
        width = int(width * ratio)
        height = int(height * ratio)
    
    # Round to nearest multiple of 8
    width = (width // 8) * 8
    height = (height // 8) * 8
    
    # Ensure minimum size
    width = max(width, 64)
    height = max(height, 64)
    
    return image.resize((width, height), Image.LANCZOS)


# =============================================================================
# INPUT PREPROCESSING
# =============================================================================

def preprocess_inputs(
    person_image: Image.Image,
    garment_image: Image.Image,
    fix_background: bool = True
) -> Tuple[Image.Image, Image.Image, str]:
    """
    Preprocess person and garment images for try-on.
    
    This is the main preprocessing function that:
    1. Validates both images
    2. Normalizes color modes
    3. Fixes dark background artifacts (person image only)
    4. Resizes to appropriate dimensions
    
    Args:
        person_image: PIL Image of person
        garment_image: PIL Image of garment
        fix_background: Whether to apply background artifact fixes
    
    Returns:
        Tuple of (processed_person, processed_garment, status_message)
        If validation fails, returns (None, None, error_message)
    """
    # Validate
    is_valid, error_msg = validate_inputs(person_image, garment_image)
    if not is_valid:
        return None, None, error_msg
    
    # Normalize person image with background fix
    person_processed = normalize_image(person_image, fix_background=fix_background)
    
    # Normalize garment image without background fix (not needed)
    garment_processed = normalize_image(garment_image, fix_background=False)
    
    logger.info(f"Preprocessed: person={person_processed.size}, garment={garment_processed.size}")
    
    return person_processed, garment_processed, "✅ Images preprocessed"


# =============================================================================
# NUMPY CONVERSION UTILITIES
# =============================================================================

def pil_to_numpy(image: Image.Image) -> np.ndarray:
    """Convert PIL Image to numpy array."""
    return np.array(image)


def numpy_to_pil(array: np.ndarray) -> Image.Image:
    """Convert numpy array to PIL Image."""
    return Image.fromarray(array.astype(np.uint8))


def ensure_pil(image) -> Optional[Image.Image]:
    """
    Ensure input is a PIL Image.
    
    Handles numpy arrays and existing PIL Images.
    
    Args:
        image: Input (numpy array or PIL Image)
    
    Returns:
        PIL Image or None if conversion fails
    """
    if image is None:
        return None
    
    if isinstance(image, Image.Image):
        return image
    
    if isinstance(image, np.ndarray):
        return numpy_to_pil(image)
    
    logger.warning(f"Unknown image type: {type(image)}")
    return None