src/quantization.py

from __future__ import annotations
import numpy as np
from PIL import Image
from sklearn.cluster import KMeans
from .utils import pil_to_np, np_to_pil
from .config import Config


def apply_uniform_quantization(image: Image.Image, levels: int) -> Image.Image:
    """
    Apply uniform color quantization to reduce color variations.
    
    Args:
        image: Input PIL Image
        levels: Number of quantization levels per channel
        
    Returns:
        Quantized PIL Image
    """
    img_array = pil_to_np(image)
    
    # Quantize each channel uniformly
    quantized = np.zeros_like(img_array)
    for channel in range(3):
        # Create quantization levels
        channel_data = img_array[:, :, channel]
        
        # Uniform quantization
        quantized_channel = np.round(channel_data * (levels - 1)) / (levels - 1)
        quantized_channel = np.clip(quantized_channel, 0, 1)
        quantized[:, :, channel] = quantized_channel
    
    return np_to_pil(quantized)


def apply_kmeans_quantization(image: Image.Image, k_colors: int) -> Image.Image:
    """
    Apply K-means clustering for color quantization.
    
    Args:
        image: Input PIL Image
        k_colors: Number of colors to reduce to
        
    Returns:
        Quantized PIL Image
    """
    img_array = pil_to_np(image)
    h, w, c = img_array.shape
    
    # Reshape image to list of pixels
    pixels = img_array.reshape(-1, c)
    
    # Apply K-means clustering
    kmeans = KMeans(n_clusters=k_colors, random_state=42, n_init=10)
    kmeans.fit(pixels)
    
    # Replace each pixel with its cluster center
    labels = kmeans.labels_
    quantized_pixels = kmeans.cluster_centers_[labels]
    
    # Reshape back to image
    quantized_img = quantized_pixels.reshape(h, w, c)
    
    return np_to_pil(quantized_img)


def apply_color_quantization(image: Image.Image, config: Config) -> Image.Image:
    """
    Apply color quantization based on configuration.
    
    Args:
        image: Input PIL Image
        config: Configuration object
        
    Returns:
        Quantized PIL Image
    """
    if config.use_uniform_q:
        return apply_uniform_quantization(image, config.q_levels)
    elif config.use_kmeans_q:
        return apply_kmeans_quantization(image, config.k_colors)
    else:
        # No quantization
        return image


def analyze_quantization_effect(original: Image.Image, quantized: Image.Image) -> dict:
    """
    Analyze the effect of quantization on the image.
    
    Args:
        original: Original image
        quantized: Quantized image
        
    Returns:
        Dictionary with analysis results
    """
    orig_array = pil_to_np(original)
    quant_array = pil_to_np(quantized)
    
    # Calculate differences
    diff = np.abs(orig_array - quant_array)
    
    # Calculate statistics
    mse = np.mean((orig_array - quant_array) ** 2)
    psnr = 20 * np.log10(1.0 / np.sqrt(mse)) if mse > 0 else float('inf')
    
    # Count unique colors
    orig_colors = len(np.unique(orig_array.reshape(-1, 3), axis=0))
    quant_colors = len(np.unique(quant_array.reshape(-1, 3), axis=0))
    
    return {
        'mse': float(mse),
        'psnr': float(psnr),
        'mean_difference': float(np.mean(diff)),
        'max_difference': float(np.max(diff)),
        'original_colors': orig_colors,
        'quantized_colors': quant_colors,
        'color_reduction_ratio': orig_colors / quant_colors if quant_colors > 0 else float('inf')
    }