Spaces:
Sleeping
Sleeping
File size: 6,836 Bytes
4376584 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
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)
|