ImagePreprocessor.py

"""
ImagePreprocessor.py
Step 1: Image preprocessing for mosaic generation.
Handles image loading, resizing, cropping, and color quantization.
"""

import cv2
import numpy as np
from PIL import Image
import os
from typing import Tuple, Optional
from sklearn.cluster import MiniBatchKMeans
import warnings

class ImagePreprocessor:
    """
    Handles image preprocessing for mosaic generation.
    Features: Smart resizing, grid-perfect cropping, fast color quantization.
    """
    
    def __init__(self, target_resolution: Tuple[int, int] = (800, 600), 
                 grid_size: Tuple[int, int] = (20, 15)):
        """
        Initialize the preprocessor.
        
        Args:
            target_resolution: Target (width, height) for processed images
            grid_size: Grid dimensions (cols, rows) for mosaic
        """
        self.target_resolution = target_resolution
        self.grid_size = grid_size
        
        # Calculate tile size for perfect grid alignment
        self.tile_width = target_resolution[0] // grid_size[0]
        self.tile_height = target_resolution[1] // grid_size[1]
        
        # Adjust target resolution to fit grid perfectly
        self.adjusted_width = self.tile_width * grid_size[0]
        self.adjusted_height = self.tile_height * grid_size[1]
        
        print(f"Target resolution: {self.adjusted_width}x{self.adjusted_height}")
        print(f"Grid size: {grid_size[0]}x{grid_size[1]}")
        print(f"Tile size: {self.tile_width}x{self.tile_height}")
    
    def load_and_preprocess_image(self, image_path: str, 
                                 apply_quantization: bool = False,
                                 n_colors: int = 16) -> Optional[np.ndarray]:
        """
        Load and preprocess image from file path.
        
        Args:
            image_path: Path to the image file
            apply_quantization: Whether to apply color quantization
            n_colors: Number of colors for quantization
            
        Returns:
            Preprocessed image as numpy array (RGB) or None if failed
        """
        try:
            # Load and convert image
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"Could not load image: {image_path}")
            
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # Resize and crop to fit grid
            processed_image = self._resize_and_crop(image)
            
            # Apply color quantization if requested
            if apply_quantization:
                processed_image = self._apply_color_quantization(processed_image, n_colors)
            
            return processed_image
            
        except Exception as e:
            print(f"Error processing {image_path}: {str(e)}")
            return None
    
    def preprocess_numpy_image(self, image: np.ndarray,
                              apply_quantization: bool = False,
                              n_colors: int = 16) -> Optional[np.ndarray]:
        """
        Preprocess numpy image array (for Gradio integration).
        
        Args:
            image: Input image as numpy array
            apply_quantization: Whether to apply color quantization
            n_colors: Number of colors for quantization
            
        Returns:
            Preprocessed image as numpy array (RGB) or None if failed
        """
        try:
            if len(image.shape) != 3 or image.shape[2] != 3:
                raise ValueError("Image must be RGB format with shape (H, W, 3)")
            
            processed_image = image.copy()
            processed_image = self._resize_and_crop(processed_image)
            
            if apply_quantization:
                processed_image = self._apply_color_quantization(processed_image, n_colors)
            
            return processed_image
            
        except Exception as e:
            print(f"Error processing numpy image: {str(e)}")
            return None
    
    def _resize_and_crop(self, image: np.ndarray) -> np.ndarray:
        """
        Resize and crop image to fit target resolution while maintaining aspect ratio.
        """
        h, w = image.shape[:2]
        target_w, target_h = self.adjusted_width, self.adjusted_height
        
        # Scale to fill target size
        scale_w = target_w / w
        scale_h = target_h / h
        scale = max(scale_w, scale_h)
        
        # Resize image
        new_w = int(w * scale)
        new_h = int(h * scale)
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
        
        # Center crop to exact target size
        start_x = (new_w - target_w) // 2
        start_y = (new_h - target_h) // 2
        cropped = resized[start_y:start_y + target_h, start_x:start_x + target_w]
        
        return cropped
    
    def _apply_color_quantization(self, image: np.ndarray, n_colors: int) -> np.ndarray:
        """
        Apply color quantization using Mini-Batch K-means for speed.
        3-4x faster than regular K-means with similar quality.
        """
        h, w, c = image.shape
        pixels = image.reshape(-1, c)
        
        # Adaptive batch size based on image size
        total_pixels = len(pixels)
        batch_size = min(max(total_pixels // 100, 1000), 10000)
        
        print(f"Applying Mini-Batch K-means quantization:")
        print(f"  Total pixels: {total_pixels:,}")
        print(f"  Batch size: {batch_size:,}")
        print(f"  Target colors: {n_colors}")
        
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=UserWarning)
            warnings.filterwarnings("ignore", message=".*ConvergenceWarning.*")
            
            kmeans = MiniBatchKMeans(
                n_clusters=n_colors, 
                batch_size=batch_size,
                random_state=42, 
                n_init=3,
                max_iter=100
            )
            labels = kmeans.fit_predict(pixels)
        
        # Replace pixels with cluster centers
        quantized_pixels = kmeans.cluster_centers_[labels]
        quantized_image = quantized_pixels.reshape(h, w, c).astype(np.uint8)
        
        return quantized_image
    
    def save_preprocessed_image(self, image: np.ndarray, output_path: str):
        """Save preprocessed image to disk."""
        try:
            pil_image = Image.fromarray(image)
            pil_image.save(output_path, quality=95)
            print(f"Saved preprocessed image to: {output_path}")
        except Exception as e:
            print(f"Error saving image: {str(e)}")

if __name__ == "__main__":
    # Test preprocessing
    preprocessor = ImagePreprocessor(
        target_resolution=(1280, 1280),
        grid_size=(32, 32)
    )
    
    test_image = "EmmaPotrait.jpg"
    
    if os.path.exists(test_image):
        processed = preprocessor.load_and_preprocess_image(
            test_image, 
            apply_quantization=True,
            n_colors=8
        )
        
        if processed is not None:
            preprocessor.save_preprocessed_image(processed, "processed_quantized.jpg")
            print("Image preprocessing completed!")
    else:
        print(f"Test image not found: {test_image}")