Spaces:
Sleeping
Sleeping
| 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: | |
| """Main class for generating mosaic images from input images.""" | |
| 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: Image preprocessing - resize and crop to fit 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: Divide image into grid and analyze each cell using vectorized operations. | |
| """ | |
| 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: Replace each grid cell with corresponding tile. | |
| Optimized vectorized version. | |
| """ | |
| grid = self.config.grid | |
| tile_size = self.config.tile_size | |
| output_h, output_w = grid * tile_size, grid * tile_size | |
| # Initialize output image | |
| mosaic_array = np.zeros((output_h, output_w, 3), dtype=np.float32) | |
| # Vectorized approach - find all matches at once | |
| tile_indices = self._find_all_tile_matches_vectorized(cell_colors) | |
| # Place tiles using vectorized operations | |
| for i in range(grid): | |
| for j in range(grid): | |
| tile_idx = tile_indices[i, j] | |
| tile = self.tile_manager.tiles[tile_idx] | |
| # Place tile in output image | |
| start_h, end_h = i * tile_size, (i + 1) * tile_size | |
| start_w, end_w = j * tile_size, (j + 1) * tile_size | |
| mosaic_array[start_h:end_h, start_w:end_w] = tile | |
| return mosaic_array | |
| def generate_mosaic(self, image: Image.Image) -> Tuple[Image.Image, dict]: | |
| """ | |
| Complete mosaic generation pipeline. | |
| Returns the mosaic image and processing 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 performance for different grid sizes. | |
| """ | |
| 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: | |
| """Find all tile matches using improved vectorized operations.""" | |
| # 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) | |