simple_mosaic.py

# simple_mosaic.py
from pathlib import Path
from typing import List, Tuple, Iterator
import numpy as np
from PIL import Image

class SimpleMosaicImage:
    def __init__(self, path: str):
        self.path = path
        self.img = Image.open(path).convert("RGB")
        self.width, self.height = self.img.size
        print(f"[INFO] Loaded: {path} | size={self.width}x{self.height}")

    def resize(self, longest_side: int = 512) -> "SimpleMosaicImage":
        w, h = self.width, self.height
        scale = longest_side / max(w, h)
        if scale < 1.0:
            new_w = int(w * scale)
            new_h = int(h * scale)
            self.img = self.img.resize((new_w, new_h), Image.BICUBIC)
            self.width, self.height = new_w, new_h
            print(f"[INFO] Resized to {new_w}x{new_h}")
        return self

    def crop_to_grid(self, grid_size: int = 32) -> "SimpleMosaicImage":
        new_w = (self.width // grid_size) * grid_size
        new_h = (self.height // grid_size) * grid_size
        if new_w != self.width or new_h != self.height:
            self.img = self.img.crop((0, 0, new_w, new_h))
            self.width, self.height = new_w, new_h
            print(f"[INFO] Cropped to {new_w}x{new_h} for grid {grid_size}")
        return self

    def _as_array(self):
        return np.asarray(self.img, dtype=np.uint8)

    def iter_cells(self, grid_size: int):
        for y in range(0, self.height, grid_size):
            for x in range(0, self.width, grid_size):
                yield (x, y, grid_size, grid_size)

    @staticmethod
    def _cell_mean(arr, x, y, w, h):
        block = arr[y:y+h, x:x+w, :]
        mean = block.mean(axis=(0,1))
        return tuple(int(round(v)) for v in mean)

    @staticmethod
    def _nearest_color(target, palette):
        tr, tg, tb = target
        pal = np.array(palette, dtype=np.int16)
        diff = pal - np.array([tr, tg, tb], dtype=np.int16)
        dist2 = np.sum(diff*diff, axis=1)
        idx = int(np.argmin(dist2))
        return tuple(int(v) for v in pal[idx])

    def mosaic_average_color(self, grid_size: int = 32):
        arr = self._as_array()
        out = np.empty_like(arr)
        for (x, y, w, h) in self.iter_cells(grid_size):
            color = self._cell_mean(arr, x, y, w, h)
            out[y:y+h, x:x+w, :] = color
        return Image.fromarray(out, mode="RGB")

    def mosaic_with_palette(self, grid_size: int, palette: List[Tuple[int,int,int]]):
        arr = self._as_array()
        out = np.empty_like(arr)
        for (x, y, w, h) in self.iter_cells(grid_size):
            avg = self._cell_mean(arr, x, y, w, h)
            color = self._nearest_color(avg, palette)
            out[y:y+h, x:x+w, :] = color
        return Image.fromarray(out, mode="RGB")

    def save(self, image: Image.Image, out_path: str) -> None:
        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        image.save(out_path)
        print(f"[INFO] Saved: {out_path}")


PALETTE_16 = [
    (0,0,0), (255,255,255), (255,0,0), (0,255,0), (0,0,255),
    (255,255,0), (255,0,255), (0,255,255),
    (128,128,128), (128,0,0), (0,128,0), (0,0,128),
    (128,128,0), (128,0,128), (0,128,128), (200,200,200)
]

loader = SimpleMosaicImage("./samples/IMG_5090.jpg")
loader.resize(longest_side=512).crop_to_grid(grid_size=2)

mosaic = loader.mosaic_average_color(grid_size=2)
loader.save(mosaic, "./out/mosaic_avg_32.png")