image_processor.py

"""图像处理核心：加载、缩放、颜色量化"""

import numpy as np
from dataclasses import dataclass, field
from PIL import Image
from sklearn.cluster import KMeans

from constants import SYMBOLS, COLOR_NAMES_ZH


@dataclass
class ProcessResult:
    """处理结果数据对象"""
    quantized_image: Image.Image          # 量化后的小尺寸网格图
    palette: list = field(default_factory=list)  # [(rgb_tuple, symbol, count), ...]
    grid_size: tuple = (0, 0)             # (宽, 高)
    original_image: Image.Image = None    # 原图（用于PDF缩略图）
    symbol_map: dict = field(default_factory=dict)  # {rgb_tuple: symbol}
    color_names: dict = field(default_factory=dict)  # {rgb_tuple: 中文颜色名}


def approximate_color_name(rgb):
    """通过欧氏距离找到最接近的中文颜色名"""
    min_dist = float("inf")
    best_name = "未知色"
    r, g, b = rgb
    for (cr, cg, cb), name in COLOR_NAMES_ZH.items():
        dist = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2
        if dist < min_dist:
            min_dist = dist
            best_name = name
    return best_name


class ImageProcessor:
    """图像处理器 - 负责加载、缩放、量化"""

    def __init__(self):
        self._original: Image.Image | None = None
        self._resized: Image.Image | None = None
        self._current_density: int = 0

    def load_image(self, filepath: str) -> Image.Image:
        img = Image.open(filepath).convert("RGB")
        self._original = img
        self._resized = None
        self._current_density = 0
        return img

    @property
    def original(self):
        return self._original

    @staticmethod
    def calculate_grid_size(img_w: int, img_h: int, density: int) -> tuple[int, int]:
        if img_w >= img_h:
            grid_w = density
            grid_h = max(1, round(density * img_h / img_w))
        else:
            grid_h = density
            grid_w = max(1, round(density * img_w / img_h))
        return grid_w, grid_h

    def resize_to_grid(self, density: int) -> Image.Image:
        if self._original is None:
            raise ValueError("未加载图片")
        w, h = self._original.size
        grid_w, grid_h = self.calculate_grid_size(w, h, density)
        self._resized = self._original.resize((grid_w, grid_h), Image.LANCZOS)
        self._current_density = density
        return self._resized

    def resize_to_exact(self, grid_w: int, grid_h: int) -> Image.Image:
        """缩放到用户指定的精确宽高格数"""
        if self._original is None:
            raise ValueError("未加载图片")
        self._resized = self._original.resize((grid_w, grid_h), Image.LANCZOS)
        self._current_density = -1
        return self._resized

    def quantize_colors(self, n_colors: int) -> ProcessResult:
        if self._resized is None:
            raise ValueError("未缩放图片")

        img = self._resized
        pixels = np.array(img).reshape(-1, 3).astype(np.float64)

        # KMeans 聚类
        n_init = min(3, max(1, n_colors // 4))
        kmeans = KMeans(n_clusters=n_colors, n_init=n_init, max_iter=150, random_state=42)
        labels = kmeans.fit_predict(pixels)
        centers = kmeans.cluster_centers_.astype(np.uint8)

        # 统计每种颜色的像素数
        unique, counts = np.unique(labels, return_counts=True)
        color_counts = dict(zip(unique.tolist(), counts.tolist()))

        # 按数量降序排列，构建调色板
        sorted_indices = sorted(color_counts.keys(), key=lambda i: color_counts[i], reverse=True)
        symbol_map = {}
        palette = []
        color_names = {}

        for rank, idx in enumerate(sorted_indices):
            rgb = tuple(centers[idx].tolist())
            symbol = SYMBOLS[rank] if rank < len(SYMBOLS) else "?"
            count = color_counts[idx]
            symbol_map[rgb] = symbol
            palette.append((rgb, symbol, count))
            color_names[rgb] = approximate_color_name(rgb)

        # 重建量化图像
        quantized_pixels = centers[labels].reshape(img.size[1], img.size[0], 3)
        quantized_img = Image.fromarray(quantized_pixels, "RGB")

        return ProcessResult(
            quantized_image=quantized_img,
            palette=palette,
            grid_size=(img.size[0], img.size[1]),
            original_image=self._original,
            symbol_map=symbol_map,
            color_names=color_names,
        )

    def process(self, density: int, n_colors: int) -> ProcessResult:
        """完整处理流水线"""
        self.resize_to_grid(density)
        return self.quantize_colors(n_colors)

    def reprocess_colors(self, n_colors: int) -> ProcessResult:
        """仅重新量化颜色（跳过resize）"""
        if self._resized is None:
            raise ValueError("未缩放图片")
        return self.quantize_colors(n_colors)