from __future__ import annotations import numpy as np from PIL import Image from typing import List, Tuple import time from scipy.spatial.distance import cdist from .utils import pil_to_np, np_to_pil, resize_and_crop_to_grid, cell_means from .config import Config, Implementation, MatchSpace from .tiles import TileManager from .quantization import apply_color_quantization class MosaicGenerator: """Generate photomosaic images using configuration-driven stages.""" def __init__(self, config: Config): self.config = config self.tile_manager = TileManager(config) self.processing_time = {} def preprocess_image(self, image: Image.Image) -> Image.Image: """ Step 1: Resize/crop the source image to align with the configured grid. Args: image (Image.Image): Input photo supplied by the user. Returns: Image.Image: Processed RGB image whose dimensions are divisible by `grid`. """ # Resize and crop to ensure grid compatibility processed_img = resize_and_crop_to_grid( image, self.config.out_w, self.config.out_h, self.config.grid ) # Apply color quantization if enabled if self.config.use_uniform_q or self.config.use_kmeans_q: processed_img = apply_color_quantization(processed_img, self.config) return processed_img def analyze_grid_cells(self, image: Image.Image) -> np.ndarray: """ Step 2: Compute representative colors for every grid cell. Args: image (Image.Image): Preprocessed image from `preprocess_image`. Returns: np.ndarray: Array of shape (grid, grid, 3) containing mean RGB values. """ img_array = pil_to_np(image) # Always use vectorized operations for better performance cell_colors = cell_means(img_array, self.config.grid) return cell_colors def map_tiles_to_grid(self, cell_colors: np.ndarray) -> np.ndarray: """ Step 3: Assemble the mosaic by mapping each cell color to a tile. Args: cell_colors (np.ndarray): Array produced by `analyze_grid_cells`. Returns: np.ndarray: Final mosaic pixels as a float32 array in [0, 1]. """ grid = self.config.grid tile_size = self.config.tile_size output_h, output_w = grid * tile_size, grid * tile_size # Vectorized approach - find all matches at once tile_indices = self._find_all_tile_matches_vectorized(cell_colors) # Stack tile bank once and gather the selected tiles in bulk tile_bank = np.stack(self.tile_manager.tiles, axis=0).astype(np.float32, copy=False) selected_tiles = tile_bank[tile_indices] # (grid, grid, tile_size, tile_size, 3) mosaic_array = ( selected_tiles.transpose(0, 2, 1, 3, 4) .reshape(output_h, output_w, 3) .copy() ) return mosaic_array def generate_mosaic(self, image: Image.Image) -> Tuple[Image.Image, dict]: """ Execute preprocessing, grid analysis, and tile mapping in sequence. Args: image (Image.Image): Input RGB image. Returns: Tuple[Image.Image, dict]: (mosaic image, timing statistics). """ start_time = time.time() # Step 1: Preprocessing preprocess_start = time.time() processed_img = self.preprocess_image(image) self.processing_time['preprocessing'] = time.time() - preprocess_start # Step 2: Grid analysis analysis_start = time.time() cell_colors = self.analyze_grid_cells(processed_img) self.processing_time['grid_analysis'] = time.time() - analysis_start # Step 3: Tile mapping mapping_start = time.time() mosaic_array = self.map_tiles_to_grid(cell_colors) self.processing_time['tile_mapping'] = time.time() - mapping_start # Convert to PIL Image mosaic_img = np_to_pil(mosaic_array) total_time = time.time() - start_time self.processing_time['total'] = total_time # Prepare statistics stats = { 'grid_size': self.config.grid, 'tile_size': self.config.tile_size, 'output_resolution': f"{mosaic_img.width}x{mosaic_img.height}", 'processing_time': self.processing_time.copy(), 'implementation': self.config.impl.value, 'match_space': self.config.match_space.value } return mosaic_img, stats def benchmark_grid_sizes(self, image: Image.Image, grid_sizes: List[int]) -> dict: """ Benchmark mosaic generation for multiple grid sizes. Args: image (Image.Image): Input image. grid_sizes (List[int]): Grid sizes (NxN) to evaluate. Returns: dict: Mapping of grid size to timing and mosaic metadata. """ results = {} original_grid = self.config.grid for grid_size in grid_sizes: self.config.grid = grid_size # Update output dimensions to maintain aspect ratio self.config.out_w = (image.width // grid_size) * grid_size self.config.out_h = (image.height // grid_size) * grid_size # Time the generation start_time = time.time() mosaic_img, stats = self.generate_mosaic(image) total_time = time.time() - start_time results[grid_size] = { 'processing_time': total_time, 'output_resolution': f"{mosaic_img.width}x{mosaic_img.height}", 'total_tiles': grid_size * grid_size } # Restore original grid size self.config.grid = original_grid return results def _find_all_tile_matches_vectorized(self, cell_colors: np.ndarray) -> np.ndarray: """ Return the best tile index for every grid cell using NumPy distance matrices. Args: cell_colors (np.ndarray): Array of cell mean colors (grid, grid, 3). Returns: np.ndarray: Tile indices shaped like the grid. """ # Ensure tiles are loaded self.tile_manager._ensure_tiles_loaded() if not self.tile_manager.tiles: return np.zeros(cell_colors.shape[:2], dtype=int) grid_h, grid_w = cell_colors.shape[:2] cell_colors_reshaped = cell_colors.reshape(-1, 3) if self.config.match_space == MatchSpace.LAB: cell_colors_lab = np.array([self.tile_manager._rgb_to_lab(color) for color in cell_colors_reshaped]) # (N,3) tile_colors_array = np.array(self.tile_manager.tile_colors_lab) # (M,3) distances = self.tile_manager._calculate_perceptual_distance(cell_colors_lab, tile_colors_array) # (N,M) else: tile_colors_array = np.array(self.tile_manager.tile_colors) # (M,3) distances = self.tile_manager._calculate_rgb_distance(cell_colors_reshaped, tile_colors_array) # (N,M) # Add small randomness per candidate to avoid ties noise_factor = 0.01 distances = distances * (1 + noise_factor * np.random.random(distances.shape)) # Find best tile per cell (argmin over tiles axis) best_indices = np.argmin(distances, axis=1) # Reshape back to grid return best_indices.reshape(grid_h, grid_w)