File size: 2,538 Bytes
c422481
 
 
 
 
388efcc
 
2ba0095
 
6bc3960
 
 
388efcc
 
c422481
388efcc
 
 
c422481
388efcc
 
 
 
 
 
 
c422481
388efcc
b80b859
 
 
 
 
 
6bc3960
 
c422481
6bc3960
 
388efcc
 
8edf546
388efcc
 
 
8edf546
388efcc
 
 
 
8edf546
388efcc
8edf546
1933832
8edf546
1933832
6bc3960
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1933832
8edf546
6bc3960
1933832
6bc3960
1933832
8edf546
1933832
8edf546
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
import numpy as np
import cv2
from sklearn.cluster import KMeans
from PIL import Image

from logic.validation import ValidationError, ensure_positive_int

IMAGE_SIZE = 400 
SAMPLE_SIZE = 50000 


def resize_img(image: Image.Image, size: int = IMAGE_SIZE) -> Image.Image:
    """Resize an image to a square of the requested size.

    Args:
        image (PIL.Image.Image): Image to resize.
        size (int, optional): Target side length in pixels. Defaults to ``IMAGE_SIZE``.

    Returns:
        PIL.Image.Image: Cropped and resized copy in RGB.

    Raises:
        ValidationError: If ``size`` is not positive.

    Example:
        >>> resized = resize_img(Image.open('input.png'), 512)
    """
    ensure_positive_int(size, "Image size", minimum=32)
    width, height = image.size
    if width != height:
        min_size = min(width, height)
        left = (width - min_size) // 2
        top = (height - min_size) // 2
        image = image.crop((left, top, left + min_size, top + min_size))
    resized_img = image.resize((size, size), Image.Resampling.LANCZOS)
    return resized_img


def color_quantize(image: Image.Image, n_colors: int = 16):
    """Reduce the number of colors with KMeans clustering.

    Args:
        image (PIL.Image.Image): Image to quantize.
        n_colors (int, optional): Number of clusters. Defaults to 16.

    Returns:
        tuple[PIL.Image.Image, np.ndarray]: Quantized image and cluster centers in BGR order.

    Raises:
        ValidationError: If ``n_colors`` is outside ``[2, 64]``.
    """
    ensure_positive_int(n_colors, "Color palette size", minimum=2, maximum=64)
    img_np = np.array(image)
    original_shape = img_np.shape
    img_flat = img_np.reshape(-1, 3)
    
    sample_size = min(SAMPLE_SIZE, img_flat.shape[0])
    indices = np.random.choice(img_flat.shape[0], sample_size, replace=False)
    img_sample = img_flat[indices]
    
    kmeans = KMeans(
        n_clusters=n_colors, 
        random_state=42, 
        init='k-means++',
        n_init=5,
        max_iter=100,
        tol=1e-3,
        algorithm='lloyd'
    )
    
    kmeans.fit(img_sample)
    labels = kmeans.predict(img_flat)
    
    quantized_flat = kmeans.cluster_centers_[labels].astype(np.uint8)
    quantized_img_np = quantized_flat.reshape(original_shape)
    
    quantized_img = Image.fromarray(quantized_img_np)
    
    color_centers = kmeans.cluster_centers_.astype(np.uint8)
    color_centers = color_centers[:, [2, 1, 0]]  # Swap R and B channels
    return quantized_img, color_centers