File size: 5,909 Bytes
f647a80
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""

ImagePreprocessor.py - FULLY OPTIMIZED

Minimal overhead, fast operations, no unnecessary prints.

"""

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


class ImagePreprocessor:
    """OPTIMIZED image preprocessing for mosaic generation."""
    
    def __init__(self, target_resolution: Tuple[int, int] = (800, 600), 

                 grid_size: Tuple[int, int] = (20, 15),

                 verbose: bool = False):
        """Initialize preprocessor - MINIMAL overhead."""
        self.target_resolution = target_resolution
        self.grid_size = grid_size
        self.verbose = verbose
        
        # Pre-compute (avoid division later)
        self.tile_width = target_resolution[0] // grid_size[0]
        self.tile_height = target_resolution[1] // grid_size[1]
        self.adjusted_width = self.tile_width * grid_size[0]
        self.adjusted_height = self.tile_height * grid_size[1]
    
    def load_and_preprocess_image(self, image_path: str, 

                                 apply_quantization: bool = False,

                                 n_colors: int = 16) -> Optional[np.ndarray]:
        """Load and preprocess - OPTIMIZED for speed."""
        try:
            # Load and convert
            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)
            processed_image = self._resize_and_crop(image)
            
            # Only quantize if explicitly requested
            if apply_quantization and n_colors < 256:
                processed_image = self._apply_color_quantization(processed_image, n_colors)
            
            return processed_image
            
        except Exception as e:
            if self.verbose:
                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 (for Gradio)."""
        try:
            if len(image.shape) != 3 or image.shape[2] != 3:
                raise ValueError("Image must be RGB (H, W, 3)")
            
            processed_image = self._resize_and_crop(image)
            
            if apply_quantization and n_colors < 256:
                processed_image = self._apply_color_quantization(processed_image, n_colors)
            
            return processed_image
            
        except Exception as e:
            if self.verbose:
                print(f"Error processing image: {str(e)}")
            return None
    
    def _resize_and_crop(self, image: np.ndarray) -> np.ndarray:
        """Resize and crop - OPTIMIZED."""
        h, w = image.shape[:2]
        target_w, target_h = self.adjusted_width, self.adjusted_height
        
        # Scale to fill
        scale = max(target_w / w, target_h / h)
        new_w = int(w * scale)
        new_h = int(h * scale)
        
        # Use INTER_LINEAR for speed (faster than INTER_AREA)
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        # Center crop
        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:
        """OPTIMIZED color quantization with aggressive sampling."""
        h, w, c = image.shape
        pixels = image.reshape(-1, c)
        total_pixels = len(pixels)
        
        # Aggressive sampling
        max_sample_size = 10000
        
        if total_pixels > max_sample_size:
            # Sample for fitting
            sample_indices = np.random.randint(0, total_pixels, size=max_sample_size)
            sampled_pixels = pixels[sample_indices].astype(np.float32)
            
            batch_size = min(max(len(sampled_pixels) // 10, 500), 2000)
            
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                
                kmeans = MiniBatchKMeans(
                    n_clusters=n_colors, batch_size=batch_size,
                    random_state=42, n_init=1, max_iter=50
                )
                
                # Fit on sample, predict on full
                kmeans.fit(sampled_pixels)
                labels = kmeans.predict(pixels.astype(np.float32))
        else:
            # Small image
            batch_size = min(max(total_pixels // 100, 500), 5000)
            
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                
                kmeans = MiniBatchKMeans(
                    n_clusters=n_colors, batch_size=batch_size,
                    random_state=42, n_init=1, max_iter=50
                )
                labels = kmeans.fit_predict(pixels.astype(np.float32))
        
        # Apply quantization
        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."""
        try:
            pil_image = Image.fromarray(image)
            pil_image.save(output_path, quality=95)
        except Exception as e:
            if self.verbose:
                print(f"Error saving image: {str(e)}")