From aceefe46b2ea5b198d987e5dadc9ada09ececf7d Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Tue, 21 Apr 2026 16:10:55 +0100 Subject: [PATCH] Bootstrap the repixelizer CLI project --- .gitignore | 9 +++ README.md | 62 ++++++++++++++++ pyproject.toml | 37 +++++++++ src/repixelizer/__init__.py | 5 ++ src/repixelizer/__main__.py | 5 ++ src/repixelizer/analysis.py | 75 +++++++++++++++++++ src/repixelizer/baselines.py | 47 ++++++++++++ src/repixelizer/cli.py | 69 +++++++++++++++++ src/repixelizer/compare.py | 102 +++++++++++++++++++++++++ src/repixelizer/continuous.py | 124 +++++++++++++++++++++++++++++++ src/repixelizer/diagnostics.py | 117 +++++++++++++++++++++++++++++ src/repixelizer/discrete.py | 60 +++++++++++++++ src/repixelizer/inference.py | 132 +++++++++++++++++++++++++++++++++ src/repixelizer/io.py | 47 ++++++++++++ src/repixelizer/metrics.py | 86 +++++++++++++++++++++ src/repixelizer/palette.py | 79 ++++++++++++++++++++ src/repixelizer/pipeline.py | 87 ++++++++++++++++++++++ src/repixelizer/synthetic.py | 46 ++++++++++++ src/repixelizer/types.py | 71 ++++++++++++++++++ tests/test_cli.py | 18 +++++ tests/test_compare.py | 22 ++++++ tests/test_inference.py | 13 ++++ tests/test_palette.py | 29 ++++++++ tests/test_pipeline.py | 56 ++++++++++++++ 24 files changed, 1398 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 src/repixelizer/__init__.py create mode 100644 src/repixelizer/__main__.py create mode 100644 src/repixelizer/analysis.py create mode 100644 src/repixelizer/baselines.py create mode 100644 src/repixelizer/cli.py create mode 100644 src/repixelizer/compare.py create mode 100644 src/repixelizer/continuous.py create mode 100644 src/repixelizer/diagnostics.py create mode 100644 src/repixelizer/discrete.py create mode 100644 src/repixelizer/inference.py create mode 100644 src/repixelizer/io.py create mode 100644 src/repixelizer/metrics.py create mode 100644 src/repixelizer/palette.py create mode 100644 src/repixelizer/pipeline.py create mode 100644 src/repixelizer/synthetic.py create mode 100644 src/repixelizer/types.py create mode 100644 tests/test_cli.py create mode 100644 tests/test_compare.py create mode 100644 tests/test_inference.py create mode 100644 tests/test_palette.py create mode 100644 tests/test_pipeline.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a80289e --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +.venv/ +__pycache__/ +*.pyc +.pytest_cache/ +build/ +dist/ +*.egg-info/ +artifacts/ +diagnostics/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..04bea47 --- /dev/null +++ b/README.md @@ -0,0 +1,62 @@ +# Repixelizer + +Repixelizer is a standalone Python CLI that turns "fake pixel art" into true grid-aligned pixel art for icons and sprites. + +It does not rely on a single resize pass. Instead it: + +1. infers the source image's implied pixel lattice, +2. initializes a target grid at that resolution, +3. runs a continuous UV optimization stage in premultiplied RGBA space, +4. projects the result onto a real pixel grid with discrete cleanup, +5. optionally constrains the output to a palette, and +6. writes diagnostics and baseline comparisons. + +## Installation + +```powershell +python -m venv .venv +.venv\Scripts\python -m pip install -e .[dev] +``` + +## Usage + +```powershell +repixelize input.png --out output.png +repixelize input.png --out output.png --target-size 120 --diagnostics-dir diagnostics +repixelize compare input.png --out output.png --diagnostics-dir diagnostics +``` + +### Palette modes + +- `off`: keep the full-color RGBA output. +- `fit`: derive or adapt to a useful palette and write palette diagnostics. +- `strict`: require a supplied palette file and stay inside it exactly. + +Palette files can be `.gpl`, `.txt`, or `.json`. + +## Diagnostics + +When `--diagnostics-dir` is supplied, Repixelizer writes: + +- `run.json`: chosen size, phase, score breakdown, timings, and settings +- `lattice-overlay.png`: inferred lattice drawn over the source +- `comparison.png`: source, optimized output, and nearest-neighbor preview +- `alpha-preview.png`: source/result alpha +- `noise-heatmap.png`: local noise and isolated-pixel hotspots +- `cluster-preview.png`: coarse material/color clustering +- compare mode extras: + - `compare.json` + - `compare.csv` + - `compare-sheet.png` + +## Architecture + +- `repixelizer.inference`: lattice and target-size inference +- `repixelizer.continuous`: PyTorch UV-field optimization +- `repixelizer.discrete`: local-search cleanup on the final grid +- `repixelizer.baselines`: naive and error-diffusion baselines +- `repixelizer.diagnostics`: machine-readable and visual debug outputs + +## Current scope + +Version `0.1.0` is single-image only. It is aimed at icons, emblems, and simple sprites first, with automatic resolution inference and no manual masks. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..1311fae --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "repixelizer" +version = "0.1.0" +description = "Optimization-based repixelizer for converting fake pixel art into true grid-aligned sprite and icon art." +readme = "README.md" +requires-python = ">=3.12" +license = { text = "MIT" } +authors = [ + { name = "OpenAI Codex" } +] +dependencies = [ + "numpy>=2.1.0", + "Pillow>=10.4.0", + "torch>=2.6.0" +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.3.0" +] + +[project.scripts] +repixelize = "repixelizer.cli:main" + +[tool.setuptools] +package-dir = { "" = "src" } + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +pythonpath = ["src"] +testpaths = ["tests"] diff --git a/src/repixelizer/__init__.py b/src/repixelizer/__init__.py new file mode 100644 index 0000000..5c9548e --- /dev/null +++ b/src/repixelizer/__init__.py @@ -0,0 +1,5 @@ +"""Optimization-based repixelizer for fake pixel art.""" + +from .pipeline import run_pipeline + +__all__ = ["run_pipeline"] diff --git a/src/repixelizer/__main__.py b/src/repixelizer/__main__.py new file mode 100644 index 0000000..a049ad7 --- /dev/null +++ b/src/repixelizer/__main__.py @@ -0,0 +1,5 @@ +from .cli import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/repixelizer/analysis.py b/src/repixelizer/analysis.py new file mode 100644 index 0000000..de9def7 --- /dev/null +++ b/src/repixelizer/analysis.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import numpy as np + +from .metrics import luminance +from .types import SourceAnalysis + + +def _compute_edge_map(rgba: np.ndarray) -> np.ndarray: + lum = luminance(rgba) + alpha = rgba[..., 3] + dx = np.zeros_like(lum) + dy = np.zeros_like(lum) + dx[:, 1:] = np.abs(lum[:, 1:] - lum[:, :-1]) + np.abs(alpha[:, 1:] - alpha[:, :-1]) + dy[1:, :] = np.abs(lum[1:, :] - lum[:-1, :]) + np.abs(alpha[1:, :] - alpha[:-1, :]) + edge = np.sqrt(dx * dx + dy * dy) + max_edge = float(np.max(edge)) + if max_edge > 0: + edge /= max_edge + return edge.astype(np.float32) + + +def _kmeans(data: np.ndarray, k: int, seed: int, iterations: int = 12) -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(seed) + if data.shape[0] <= k: + centers = data.copy() + labels = np.arange(data.shape[0], dtype=np.int32) + return centers, labels + picks = rng.choice(data.shape[0], size=k, replace=False) + centers = data[picks].copy() + for _ in range(iterations): + distances = np.linalg.norm(data[:, None, :] - centers[None, :, :], axis=-1) + labels = np.argmin(distances, axis=1) + new_centers = [] + for idx in range(k): + members = data[labels == idx] + if members.size == 0: + new_centers.append(centers[idx]) + else: + new_centers.append(np.mean(members, axis=0)) + centers = np.asarray(new_centers, dtype=np.float32) + distances = np.linalg.norm(data[:, None, :] - centers[None, :, :], axis=-1) + labels = np.argmin(distances, axis=1) + return centers.astype(np.float32), labels.astype(np.int32) + + +def _colorize_clusters(cluster_map: np.ndarray, centers: np.ndarray) -> np.ndarray: + if centers.size == 0: + return np.zeros((*cluster_map.shape, 4), dtype=np.float32) + palette = np.clip(centers[:, :3], 0.0, 1.0) + preview = palette[np.clip(cluster_map, 0, len(palette) - 1)] + alpha = np.ones((*cluster_map.shape, 1), dtype=np.float32) + return np.concatenate([preview, alpha], axis=-1) + + +def analyze_source(rgba: np.ndarray, seed: int, cluster_count: int = 6) -> SourceAnalysis: + edge_map = _compute_edge_map(rgba) + alpha = rgba[..., 3] + opaque = alpha > 0.05 + cluster_map = np.full(alpha.shape, -1, dtype=np.int32) + centers = np.zeros((0, 4), dtype=np.float32) + if np.any(opaque): + samples = rgba[opaque] + k = min(cluster_count, max(2, int(np.sqrt(samples.shape[0] // 64 + 1)))) + centers, labels = _kmeans(samples, k=k, seed=seed) + cluster_map[opaque] = labels + preview = _colorize_clusters(np.maximum(cluster_map, 0), centers if centers.size else np.zeros((1, 4), dtype=np.float32)) + preview[..., 3] = alpha + return SourceAnalysis( + edge_map=edge_map, + cluster_map=cluster_map, + cluster_centers=centers, + alpha_map=alpha.astype(np.float32), + cluster_preview=preview, + ) diff --git a/src/repixelizer/baselines.py b/src/repixelizer/baselines.py new file mode 100644 index 0000000..a26f67a --- /dev/null +++ b/src/repixelizer/baselines.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import numpy as np + +from .io import box_resize +from .palette import derive_palette + + +def naive_resize_baseline(rgba: np.ndarray, width: int, height: int) -> np.ndarray: + return box_resize(rgba, width=width, height=height) + + +def _nearest_palette_color(pixel: np.ndarray, palette: np.ndarray) -> np.ndarray: + distances = np.sum((palette - pixel[None, :]) ** 2, axis=1) + return palette[int(np.argmin(distances))] + + +def error_diffusion_baseline( + rgba: np.ndarray, + width: int, + height: int, + palette: list[tuple[int, int, int]] | None = None, +) -> np.ndarray: + resized = box_resize(rgba, width=width, height=height).copy() + palette_array = np.asarray(palette or derive_palette(resized), dtype=np.float32) / 255.0 + rgb = resized[..., :3].copy() + alpha = resized[..., 3:4].copy() + height_px, width_px = rgb.shape[:2] + for y in range(height_px): + serpentine = y % 2 == 1 + x_range = range(width_px - 1, -1, -1) if serpentine else range(width_px) + for x in x_range: + old = rgb[y, x].copy() + new = _nearest_palette_color(old, palette_array) + rgb[y, x] = new + error = old - new + neighbors = ( + [(-1, 1, 7 / 16), (0, 1, 5 / 16), (-1, 0, 3 / 16), (1, 0, 1 / 16)] + if serpentine + else [(1, 0, 7 / 16), (0, 1, 5 / 16), (-1, 1, 3 / 16), (1, 1, 1 / 16)] + ) + for dx, dy, weight in neighbors: + nx = x + dx + ny = y + dy + if 0 <= nx < width_px and 0 <= ny < height_px: + rgb[ny, nx] = np.clip(rgb[ny, nx] + error * weight, 0.0, 1.0) + return np.concatenate([rgb, alpha], axis=-1) diff --git a/src/repixelizer/cli.py b/src/repixelizer/cli.py new file mode 100644 index 0000000..bd21ad7 --- /dev/null +++ b/src/repixelizer/cli.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import argparse +import sys + +from .compare import run_compare +from .pipeline import run_pipeline + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="repixelize", description="Optimization-based repixelizer for fake pixel art.") + subparsers = parser.add_subparsers(dest="command") + + def add_shared_arguments(target: argparse.ArgumentParser) -> None: + target.add_argument("input", help="Source image path") + target.add_argument("--out", required=True, help="Output image path") + target.add_argument("--target-size", type=int, default=None, help="Override the inferred target max dimension") + target.add_argument("--palette", default=None, help="Optional palette file (.gpl, .txt, .json)") + target.add_argument( + "--palette-mode", + choices=("off", "fit", "strict"), + default="off", + help="Palette handling mode", + ) + target.add_argument("--diagnostics-dir", default=None, help="Directory for debug artifacts") + target.add_argument("--seed", type=int, default=7, help="Random seed") + target.add_argument("--steps", type=int, default=200, help="Number of optimizer steps") + target.add_argument("--device", default="cpu", choices=("cpu", "cuda"), help="Torch device") + + run_parser = subparsers.add_parser("run", help="Run the optimizer.") + add_shared_arguments(run_parser) + compare_parser = subparsers.add_parser("compare", help="Run the optimizer and comparison baselines.") + add_shared_arguments(compare_parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + raw_argv = list(sys.argv[1:] if argv is None else argv) + if not raw_argv or raw_argv[0] not in {"run", "compare", "-h", "--help"}: + raw_argv = ["run", *raw_argv] + parser = build_parser() + args = parser.parse_args(raw_argv) + command = args.command or "run" + if command == "compare": + run_compare( + args.input, + args.out, + target_size=args.target_size, + palette_path=args.palette, + palette_mode=args.palette_mode, + diagnostics_dir=args.diagnostics_dir, + seed=args.seed, + steps=args.steps, + device=args.device, + ) + return 0 + + run_pipeline( + args.input, + args.out, + target_size=args.target_size, + palette_path=args.palette, + palette_mode=args.palette_mode, + diagnostics_dir=args.diagnostics_dir, + seed=args.seed, + steps=args.steps, + device=args.device, + ) + return 0 diff --git a/src/repixelizer/compare.py b/src/repixelizer/compare.py new file mode 100644 index 0000000..fc291bb --- /dev/null +++ b/src/repixelizer/compare.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +from .baselines import error_diffusion_baseline, naive_resize_baseline +from .diagnostics import write_compare_csv, write_json +from .io import load_rgba, nearest_resize, save_rgba +from .metrics import coherence_breakdown, reconstruction_error +from .palette import load_palette +from .pipeline import run_pipeline + + +def _to_uint8(rgba: np.ndarray) -> np.ndarray: + return np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8) + + +def _contact_sheet( + source: np.ndarray, + optimized: np.ndarray, + naive: np.ndarray, + diffusion: np.ndarray, + path: str | Path, +) -> None: + height, width = source.shape[:2] + panels = [ + ("Source", source), + ("Optimized", nearest_resize(optimized, width=width, height=height)), + ("Naive", nearest_resize(naive, width=width, height=height)), + ("Diffusion", nearest_resize(diffusion, width=width, height=height)), + ] + canvas = Image.new("RGBA", (width * len(panels), height + 24), (20, 20, 20, 255)) + draw = ImageDraw.Draw(canvas) + for index, (label, rgba) in enumerate(panels): + image = Image.fromarray(_to_uint8(rgba), mode="RGBA") + canvas.alpha_composite(image, (index * width, 24)) + draw.text((index * width + 8, 4), label, fill=(255, 255, 255, 255)) + canvas.save(path) + + +def run_compare( + input_path: str | Path, + output_path: str | Path, + *, + target_size: int | None = None, + palette_path: str | Path | None = None, + palette_mode: str = "off", + diagnostics_dir: str | Path | None = None, + seed: int = 7, + steps: int = 200, + device: str = "cpu", +) -> dict[str, Any]: + diagnostics_path = Path(diagnostics_dir) if diagnostics_dir else Path(output_path).with_suffix("") + diagnostics_path.mkdir(parents=True, exist_ok=True) + result = run_pipeline( + input_path, + output_path, + target_size=target_size, + palette_path=palette_path, + palette_mode=palette_mode, + diagnostics_dir=diagnostics_path, + seed=seed, + steps=steps, + device=device, + ) + source = load_rgba(input_path) + palette = load_palette(palette_path) if palette_path else None + naive = naive_resize_baseline(source, width=result.inference.target_width, height=result.inference.target_height) + diffusion = error_diffusion_baseline( + source, + width=result.inference.target_width, + height=result.inference.target_height, + palette=palette, + ) + _contact_sheet(source, result.output_rgba, naive, diffusion, diagnostics_path / "compare-sheet.png") + rows = [] + for name, image in ( + ("optimized", result.output_rgba), + ("naive", naive), + ("diffusion", diffusion), + ): + preview = nearest_resize(image, width=source.shape[1], height=source.shape[0]) + coherence = coherence_breakdown(image) + rows.append( + { + "name": name, + "coherence_score": coherence["coherence_score"], + "cluster_continuity": coherence["cluster_continuity"], + "alpha_crispness": coherence["alpha_crispness"], + "outline_straightness": coherence["outline_straightness"], + "isolated_penalty": coherence["isolated_penalty"], + "color_chatter": coherence["color_chatter"], + "reconstruction_error": reconstruction_error(preview, source), + } + ) + write_compare_csv(diagnostics_path / "compare.csv", rows) + payload = {"rows": rows} + write_json(diagnostics_path / "compare.json", payload) + return payload diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py new file mode 100644 index 0000000..a07f07a --- /dev/null +++ b/src/repixelizer/continuous.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .io import premultiply, unpremultiply +from .types import InferenceResult, SolverArtifacts, SourceAnalysis + + +def _require_torch(): + try: + import torch + import torch.nn.functional as F + except ImportError as exc: # pragma: no cover - exercised only when torch missing + raise RuntimeError( + "PyTorch is required for the continuous optimization stage. Install project dependencies first." + ) from exc + return torch, F + + +def _make_regular_uv(height: int, width: int, target_height: int, target_width: int, phase_x: float, phase_y: float) -> np.ndarray: + cell_x = width / target_width + cell_y = height / target_height + xs = (np.arange(target_width, dtype=np.float32) + 0.5 + phase_x) * cell_x - 0.5 + ys = (np.arange(target_height, dtype=np.float32) + 0.5 + phase_y) * cell_y - 0.5 + xs = np.clip(xs, 0.0, max(0.0, width - 1)) + ys = np.clip(ys, 0.0, max(0.0, height - 1)) + grid_x, grid_y = np.meshgrid(xs, ys) + uv = np.stack([grid_x, grid_y], axis=-1) + uv[..., 0] = (uv[..., 0] / max(1.0, width - 1)) * 2.0 - 1.0 + uv[..., 1] = (uv[..., 1] / max(1.0, height - 1)) * 2.0 - 1.0 + return uv.astype(np.float32) + + +def _cluster_boundary_map(cluster_map: np.ndarray) -> np.ndarray: + if cluster_map.size == 0: + return np.zeros_like(cluster_map, dtype=np.float32) + boundary = np.zeros(cluster_map.shape, dtype=np.float32) + boundary[:, 1:] = np.maximum(boundary[:, 1:], (cluster_map[:, 1:] != cluster_map[:, :-1]).astype(np.float32)) + boundary[1:, :] = np.maximum(boundary[1:, :], (cluster_map[1:, :] != cluster_map[:-1, :]).astype(np.float32)) + return boundary + + +def optimize_uv_field( + rgba: np.ndarray, + inference: InferenceResult, + analysis: SourceAnalysis, + steps: int, + seed: int, + device: str, +) -> SolverArtifacts: + torch, F = _require_torch() + torch.manual_seed(seed) + source = premultiply(rgba) + height, width = source.shape[:2] + source_t = torch.from_numpy(source.transpose(2, 0, 1)[None, ...]).to(device=device, dtype=torch.float32) + edge = np.maximum(analysis.edge_map, _cluster_boundary_map(analysis.cluster_map)) + edge_t = torch.from_numpy(edge[None, None, ...]).to(device=device, dtype=torch.float32) + uv0 = _make_regular_uv( + height=height, + width=width, + target_height=inference.target_height, + target_width=inference.target_width, + phase_x=inference.phase_x, + phase_y=inference.phase_y, + ) + uv0_t = torch.from_numpy(uv0[None, ...]).to(device=device, dtype=torch.float32) + uv = torch.nn.Parameter(uv0_t.clone()) + optimizer = torch.optim.Adam([uv], lr=0.02) + loss_history: list[float] = [] + for _ in range(steps): + optimizer.zero_grad(set_to_none=True) + sampled = F.grid_sample(source_t, uv, align_corners=True, mode="bilinear", padding_mode="border") + reconstruction = F.interpolate(sampled, size=(height, width), mode="nearest") + recon_loss = torch.mean(torch.abs(reconstruction - source_t)) + + dx = uv[:, :, 1:, :] - uv[:, :, :-1, :] + dy = uv[:, 1:, :, :] - uv[:, :-1, :, :] + dx0 = uv0_t[:, :, 1:, :] - uv0_t[:, :, :-1, :] + dy0 = uv0_t[:, 1:, :, :] - uv0_t[:, :-1, :, :] + + guide_small = F.interpolate(edge_t, size=(inference.target_height, inference.target_width), mode="bilinear", align_corners=True) + guide_x = 1.0 - guide_small[:, :, :, 1:] + guide_y = 1.0 - guide_small[:, :, 1:, :] + + smooth_loss = torch.mean((dx - dx0) ** 2 * guide_x.permute(0, 2, 3, 1)) + torch.mean( + (dy - dy0) ** 2 * guide_y.permute(0, 2, 3, 1) + ) + + target_rgba = sampled.permute(0, 2, 3, 1) + neighbor_mean = ( + torch.roll(target_rgba, 1, dims=1) + + torch.roll(target_rgba, -1, dims=1) + + torch.roll(target_rgba, 1, dims=2) + + torch.roll(target_rgba, -1, dims=2) + ) / 4.0 + speckle = torch.relu(torch.linalg.norm(target_rgba - neighbor_mean, dim=-1) - 0.08) + anti_speckle = torch.mean(speckle * (1.0 - guide_small[:, 0])) + + alpha = target_rgba[..., 3] + alpha_crisp = torch.mean(alpha * (1.0 - alpha)) + + loss = recon_loss + smooth_loss * 0.35 + anti_speckle * 0.25 + alpha_crisp * 0.05 + loss.backward() + optimizer.step() + with torch.no_grad(): + uv.clamp_(-1.0, 1.0) + loss_history.append(float(loss.detach().cpu().item())) + + with torch.no_grad(): + sampled = F.grid_sample(source_t, uv, align_corners=True, mode="bilinear", padding_mode="border") + target_rgba = sampled[0].permute(1, 2, 0).detach().cpu().numpy() + initial_rgba = F.grid_sample(source_t, uv0_t, align_corners=True, mode="bilinear", padding_mode="border")[0].permute(1, 2, 0) + initial_rgba_np = initial_rgba.detach().cpu().numpy() + guidance = guide_small[0, 0].detach().cpu().numpy() + + return SolverArtifacts( + target_rgba=unpremultiply(target_rgba), + uv_field=uv.detach().cpu().numpy()[0], + guidance_strength=guidance.astype(np.float32), + initial_rgba=unpremultiply(initial_rgba_np), + loss_history=loss_history, + ) diff --git a/src/repixelizer/diagnostics.py b/src/repixelizer/diagnostics.py new file mode 100644 index 0000000..a721d82 --- /dev/null +++ b/src/repixelizer/diagnostics.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import csv +import json +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +from .io import nearest_resize, save_rgba +from .metrics import coherence_breakdown, reconstruction_error +from .types import InferenceResult, RunResult + + +def _to_uint8(rgba: np.ndarray) -> np.ndarray: + return np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8) + + +def write_lattice_overlay(path: str | Path, source_rgba: np.ndarray, inference: InferenceResult) -> None: + image = Image.fromarray(_to_uint8(source_rgba), mode="RGBA").convert("RGBA") + draw = ImageDraw.Draw(image) + height, width = source_rgba.shape[:2] + cell_x = width / inference.target_width + cell_y = height / inference.target_height + phase_x = inference.phase_x * cell_x + phase_y = inference.phase_y * cell_y + for ix in range(inference.target_width + 1): + x = ix * cell_x + phase_x + draw.line((x, 0, x, height), fill=(255, 255, 255, 96), width=1) + for iy in range(inference.target_height + 1): + y = iy * cell_y + phase_y + draw.line((0, y, width, y), fill=(255, 255, 255, 96), width=1) + image.save(path) + + +def _make_checker(size: tuple[int, int], tile: int = 8) -> Image.Image: + width, height = size + checker = Image.new("RGBA", size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(checker) + a = (245, 245, 245, 255) + b = (220, 220, 220, 255) + for y in range(0, height, tile): + for x in range(0, width, tile): + color = a if ((x // tile) + (y // tile)) % 2 == 0 else b + draw.rectangle((x, y, x + tile - 1, y + tile - 1), fill=color) + return checker + + +def write_comparison(path: str | Path, source_rgba: np.ndarray, output_rgba: np.ndarray) -> None: + height, width = source_rgba.shape[:2] + preview = nearest_resize(output_rgba, width=width, height=height) + tiles = [ + Image.fromarray(_to_uint8(source_rgba), mode="RGBA"), + Image.fromarray(_to_uint8(preview), mode="RGBA"), + Image.fromarray(_to_uint8(output_rgba), mode="RGBA").resize((width, height), resample=Image.Resampling.NEAREST), + ] + labels = ["Source", "NN Preview", "Target"] + canvas = Image.new("RGBA", (width * len(tiles), height + 24), (24, 24, 24, 255)) + for index, tile in enumerate(tiles): + panel = _make_checker((width, height)) + panel.alpha_composite(tile, (0, 0)) + canvas.alpha_composite(panel, (index * width, 24)) + draw = ImageDraw.Draw(canvas) + for index, label in enumerate(labels): + draw.text((index * width + 8, 4), label, fill=(255, 255, 255, 255)) + canvas.save(path) + + +def write_alpha_preview(path: str | Path, source_rgba: np.ndarray, output_rgba: np.ndarray) -> None: + source_alpha = np.repeat(source_rgba[..., 3:4], 4, axis=-1) + source_alpha[..., 3] = 1.0 + output_alpha = np.repeat(output_rgba[..., 3:4], 4, axis=-1) + output_alpha[..., 3] = 1.0 + write_comparison(path, source_alpha, output_alpha) + + +def write_heatmap(path: str | Path, heatmap: np.ndarray) -> None: + max_value = float(np.max(heatmap)) + scaled = heatmap / max_value if max_value > 0 else heatmap + rgba = np.zeros((*scaled.shape, 4), dtype=np.float32) + rgba[..., 0] = scaled + rgba[..., 1] = np.sqrt(scaled) * 0.5 + rgba[..., 3] = np.clip(scaled * 1.5, 0.0, 1.0) + save_rgba(path, rgba) + + +def write_json(path: str | Path, payload: dict[str, Any]) -> None: + Path(path).write_text(json.dumps(payload, indent=2), encoding="utf-8") + + +def write_compare_csv(path: str | Path, rows: list[dict[str, Any]]) -> None: + if not rows: + return + keys = list(rows[0].keys()) + with Path(path).open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=keys) + writer.writeheader() + writer.writerows(rows) + + +def summarize_run(result: RunResult) -> dict[str, Any]: + source_preview = nearest_resize( + result.output_rgba, + width=result.source_rgba.shape[1], + height=result.source_rgba.shape[0], + ) + return { + "target_width": result.inference.target_width, + "target_height": result.inference.target_height, + "phase_x": result.inference.phase_x, + "phase_y": result.inference.phase_y, + "confidence": result.inference.confidence, + "coherence": coherence_breakdown(result.output_rgba), + "source_preview_reconstruction_error": reconstruction_error(source_preview, result.source_rgba), + "loss_history": result.solver.loss_history, + } diff --git a/src/repixelizer/discrete.py b/src/repixelizer/discrete.py new file mode 100644 index 0000000..439f773 --- /dev/null +++ b/src/repixelizer/discrete.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import numpy as np + +from .metrics import isolated_pixel_rate, luminance +from .types import CleanupArtifacts + + +def _local_energy(patch: np.ndarray) -> float: + center = patch[1, 1] + neighbors = np.stack( + [ + patch[1, 0], + patch[1, 2], + patch[0, 1], + patch[2, 1], + ], + axis=0, + ) + dist = np.linalg.norm(neighbors - center[None, :], axis=-1) + isolation = float(np.mean(dist)) + alpha = patch[..., 3] + alpha_kink = float(np.mean(alpha * (1.0 - alpha))) + return isolation + alpha_kink * 0.25 + + +def cleanup_pixels(rgba: np.ndarray, source_guidance: np.ndarray, iterations: int = 2) -> CleanupArtifacts: + result = rgba.copy() + height, width = result.shape[:2] + heatmap = np.zeros((height, width), dtype=np.float32) + for _ in range(iterations): + updated = result.copy() + for y in range(1, height - 1): + for x in range(1, width - 1): + if source_guidance[y, x] > 0.55: + continue + patch = result[y - 1 : y + 2, x - 1 : x + 2] + current_energy = _local_energy(patch) + neighbors = [ + result[y, x], + result[y - 1, x], + result[y + 1, x], + result[y, x - 1], + result[y, x + 1], + np.mean(patch.reshape(-1, 4), axis=0), + np.median(patch.reshape(-1, 4), axis=0), + ] + best = result[y, x] + best_energy = current_energy + for candidate in neighbors[1:]: + patch_candidate = patch.copy() + patch_candidate[1, 1] = candidate + energy = _local_energy(patch_candidate) + if energy + 1e-6 < best_energy: + best_energy = energy + best = candidate + updated[y, x] = best + heatmap[y, x] = max(heatmap[y, x], current_energy - best_energy) + result = np.clip(updated, 0.0, 1.0) + return CleanupArtifacts(cleaned_rgba=result, isolated_heatmap=heatmap) diff --git a/src/repixelizer/inference.py b/src/repixelizer/inference.py new file mode 100644 index 0000000..8c206b1 --- /dev/null +++ b/src/repixelizer/inference.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import numpy as np +from PIL import Image + +from .metrics import coherence_breakdown +from .types import InferenceCandidate, InferenceResult + + +def _to_pil(rgba: np.ndarray) -> Image.Image: + return Image.fromarray(np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8), mode="RGBA") + + +def _apply_phase(image: Image.Image, offset_x: float, offset_y: float) -> Image.Image: + rgba = np.asarray(image, dtype=np.uint8) + pad_x = int(np.ceil(abs(offset_x))) + 2 + pad_y = int(np.ceil(abs(offset_y))) + 2 + padded = np.pad(rgba, ((pad_y, pad_y), (pad_x, pad_x), (0, 0)), mode="edge") + start_x = pad_x + int(round(offset_x)) + start_y = pad_y + int(round(offset_y)) + cropped = padded[start_y : start_y + rgba.shape[0], start_x : start_x + rgba.shape[1]] + return Image.fromarray(cropped, mode="RGBA") + + +def _candidate_dims(width: int, height: int, target_size: int | None) -> list[tuple[int, int]]: + if target_size is not None: + if width >= height: + return [(target_size, max(1, round(height * target_size / width)))] + return [(max(1, round(width * target_size / height)), target_size)] + max_dim = max(width, height) + min_size = max(12, int(round(max_dim / 48))) + max_size = min(256, max(24, int(round(max_dim / 2.5)))) + step = 2 if max_size <= 96 else 4 + dims: list[tuple[int, int]] = [] + for size in range(min_size, max_size + 1, step): + if width >= height: + dims.append((size, max(1, round(height * size / width)))) + else: + dims.append((max(1, round(width * size / height)), size)) + return dims + + +def _estimate_cell_size(profile: np.ndarray) -> float: + centered = profile.astype(np.float32) - float(np.mean(profile)) + max_lag = min(64, max(4, profile.shape[0] // 3)) + best_lag = 1 + best_score = -float("inf") + for lag in range(2, max_lag + 1): + left = centered[:-lag] + right = centered[lag:] + denom = float(np.linalg.norm(left) * np.linalg.norm(right)) + 1e-6 + score = float(np.dot(left, right) / denom) + if score > best_score: + best_score = score + best_lag = lag + return float(best_lag) + + +def _estimate_lattice_prior(rgba: np.ndarray) -> tuple[float, float]: + alpha = rgba[..., 3] + luminance = rgba[..., 0] * 0.2126 + rgba[..., 1] * 0.7152 + rgba[..., 2] * 0.0722 + dx = np.zeros_like(luminance) + dy = np.zeros_like(luminance) + dx[:, 1:] = np.abs(luminance[:, 1:] - luminance[:, :-1]) + np.abs(alpha[:, 1:] - alpha[:, :-1]) + dy[1:, :] = np.abs(luminance[1:, :] - luminance[:-1, :]) + np.abs(alpha[1:, :] - alpha[:-1, :]) + profile_x = (dx.mean(axis=0) + alpha.mean(axis=0) * 0.1).astype(np.float32) + profile_y = (dy.mean(axis=1) + alpha.mean(axis=1) * 0.1).astype(np.float32) + return _estimate_cell_size(profile_x), _estimate_cell_size(profile_y) + + +def infer_lattice(rgba: np.ndarray, target_size: int | None = None) -> InferenceResult: + height, width = rgba.shape[:2] + image = _to_pil(rgba) + prior_cell_x, prior_cell_y = _estimate_lattice_prior(rgba) + candidates: list[InferenceCandidate] = [] + for target_width, target_height in _candidate_dims(width, height, target_size): + cell_x = width / target_width + cell_y = height / target_height + for phase_x in np.linspace(-0.4, 0.4, num=5): + for phase_y in np.linspace(-0.4, 0.4, num=5): + shifted = _apply_phase(image, phase_x * cell_x, phase_y * cell_y) + sample = shifted.resize((target_width, target_height), resample=Image.Resampling.BOX) + rgba_sample = np.asarray(sample, dtype=np.float32) / 255.0 + breakdown = coherence_breakdown(rgba_sample) + prior_score_x = np.exp(-abs(np.log((cell_x + 1e-6) / (prior_cell_x + 1e-6))) * 0.8) + prior_score_y = np.exp(-abs(np.log((cell_y + 1e-6) / (prior_cell_y + 1e-6))) * 0.8) + size_prior = float((prior_score_x + prior_score_y) * 0.5) + breakdown["size_prior"] = size_prior + score = breakdown["coherence_score"] * 0.65 + size_prior * 0.35 + candidates.append( + InferenceCandidate( + target_width=target_width, + target_height=target_height, + phase_x=float(phase_x), + phase_y=float(phase_y), + score=float(score), + breakdown=breakdown, + ) + ) + candidates.sort(key=lambda item: item.score, reverse=True) + best = candidates[0] + second = candidates[1] if len(candidates) > 1 else best + confidence = max(0.0, best.score - second.score) + return InferenceResult( + target_width=best.target_width, + target_height=best.target_height, + phase_x=best.phase_x, + phase_y=best.phase_y, + confidence=confidence, + top_candidates=candidates[:8], + ) + + +def inference_to_json(result: InferenceResult) -> dict[str, object]: + return { + "target_width": result.target_width, + "target_height": result.target_height, + "phase_x": result.phase_x, + "phase_y": result.phase_y, + "confidence": result.confidence, + "top_candidates": [ + { + "target_width": candidate.target_width, + "target_height": candidate.target_height, + "phase_x": candidate.phase_x, + "phase_y": candidate.phase_y, + "score": candidate.score, + "breakdown": candidate.breakdown, + } + for candidate in result.top_candidates + ], + } diff --git a/src/repixelizer/io.py b/src/repixelizer/io.py new file mode 100644 index 0000000..19752d8 --- /dev/null +++ b/src/repixelizer/io.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from PIL import Image + + +def load_rgba(path: str | Path) -> np.ndarray: + image = Image.open(path).convert("RGBA") + return np.asarray(image, dtype=np.float32) / 255.0 + + +def save_rgba(path: str | Path, rgba: np.ndarray) -> None: + clipped = np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8) + Image.fromarray(clipped, mode="RGBA").save(path) + + +def premultiply(rgba: np.ndarray) -> np.ndarray: + premult = rgba.copy() + premult[..., :3] *= premult[..., 3:4] + return premult + + +def unpremultiply(rgba: np.ndarray, eps: float = 1e-6) -> np.ndarray: + out = rgba.copy() + alpha = np.maximum(out[..., 3:4], eps) + out[..., :3] = np.where(out[..., 3:4] > eps, out[..., :3] / alpha, 0.0) + return np.clip(out, 0.0, 1.0) + + +def nearest_resize(rgba: np.ndarray, width: int, height: int) -> np.ndarray: + image = Image.fromarray(np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8), mode="RGBA") + resized = image.resize((width, height), resample=Image.Resampling.NEAREST) + return np.asarray(resized, dtype=np.float32) / 255.0 + + +def box_resize(rgba: np.ndarray, width: int, height: int) -> np.ndarray: + image = Image.fromarray(np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8), mode="RGBA") + resized = image.resize((width, height), resample=Image.Resampling.BOX) + return np.asarray(resized, dtype=np.float32) / 255.0 + + +def bilinear_resize(rgba: np.ndarray, width: int, height: int) -> np.ndarray: + image = Image.fromarray(np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8), mode="RGBA") + resized = image.resize((width, height), resample=Image.Resampling.BILINEAR) + return np.asarray(resized, dtype=np.float32) / 255.0 diff --git a/src/repixelizer/metrics.py b/src/repixelizer/metrics.py new file mode 100644 index 0000000..2c04561 --- /dev/null +++ b/src/repixelizer/metrics.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import numpy as np + + +def luminance(rgba: np.ndarray) -> np.ndarray: + return rgba[..., 0] * 0.2126 + rgba[..., 1] * 0.7152 + rgba[..., 2] * 0.0722 + + +def alpha_crispness(rgba: np.ndarray) -> float: + alpha = rgba[..., 3] + return float(1.0 - np.mean(4.0 * alpha * (1.0 - alpha))) + + +def isolated_pixel_rate(rgba: np.ndarray, threshold: float = 0.15) -> float: + color = rgba[..., :4] + diffs = [] + neighbors = [] + for dy, dx in ((0, 1), (1, 0), (0, -1), (-1, 0)): + shifted = np.roll(color, shift=(dy, dx), axis=(0, 1)) + dist = np.linalg.norm(color - shifted, axis=-1) + diffs.append(dist) + neighbors.append(dist > threshold) + stacked = np.stack(neighbors, axis=0) + isolated = np.sum(stacked, axis=0) >= 3 + return float(np.mean(isolated)) + + +def cluster_continuity(rgba: np.ndarray) -> float: + color = rgba[..., :4] + sims = [] + for dy, dx in ((0, 1), (1, 0)): + shifted = np.roll(color, shift=(dy, dx), axis=(0, 1)) + dist = np.linalg.norm(color - shifted, axis=-1) + sims.append(np.exp(-dist * 8.0)) + return float(np.mean(np.stack(sims, axis=0))) + + +def color_chatter(rgba: np.ndarray) -> float: + color = rgba[..., :3] + padded = np.pad(color, ((1, 1), (1, 1), (0, 0)), mode="edge") + windows = [] + for y in range(3): + for x in range(3): + windows.append(padded[y : y + color.shape[0], x : x + color.shape[1]]) + stack = np.stack(windows, axis=0) + med = np.median(stack, axis=0) + dev = np.linalg.norm(color - med, axis=-1) + return float(np.mean(dev)) + + +def outline_straightness(rgba: np.ndarray) -> float: + mask = rgba[..., 3] >= 0.5 + if not np.any(mask): + return 1.0 + a = mask[:-1, :-1] + b = mask[:-1, 1:] + c = mask[1:, :-1] + d = mask[1:, 1:] + block_sum = a.astype(np.int8) + b.astype(np.int8) + c.astype(np.int8) + d.astype(np.int8) + corner_blocks = np.logical_or(block_sum == 1, block_sum == 3) + edge_pixels = np.logical_xor(mask, np.roll(mask, 1, axis=0)) | np.logical_xor(mask, np.roll(mask, 1, axis=1)) + edge_count = np.maximum(1, np.count_nonzero(edge_pixels)) + return float(1.0 - np.count_nonzero(corner_blocks) / edge_count) + + +def reconstruction_error(a: np.ndarray, b: np.ndarray) -> float: + return float(np.mean(np.abs(a - b))) + + +def coherence_breakdown(rgba: np.ndarray) -> dict[str, float]: + breakdown = { + "cluster_continuity": cluster_continuity(rgba), + "alpha_crispness": alpha_crispness(rgba), + "outline_straightness": outline_straightness(rgba), + "isolated_penalty": isolated_pixel_rate(rgba), + "color_chatter": color_chatter(rgba), + } + breakdown["coherence_score"] = ( + 0.25 * breakdown["cluster_continuity"] + + 0.20 * breakdown["alpha_crispness"] + + 0.20 * breakdown["outline_straightness"] + + 0.20 * (1.0 - breakdown["isolated_penalty"]) + + 0.15 * (1.0 - min(1.0, breakdown["color_chatter"])) + ) + return breakdown diff --git a/src/repixelizer/palette.py b/src/repixelizer/palette.py new file mode 100644 index 0000000..93c86b5 --- /dev/null +++ b/src/repixelizer/palette.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +from PIL import Image + +from .types import PaletteResult + + +def _parse_hex_color(value: str) -> tuple[int, int, int] | None: + token = value.strip().lstrip("#") + if len(token) != 6: + return None + try: + return int(token[0:2], 16), int(token[2:4], 16), int(token[4:6], 16) + except ValueError: + return None + + +def load_palette(path: str | Path) -> list[tuple[int, int, int]]: + palette_path = Path(path) + if palette_path.suffix.lower() == ".json": + data = json.loads(palette_path.read_text(encoding="utf-8")) + return [tuple(color) for color in data["palette"]] + colors: list[tuple[int, int, int]] = [] + for raw_line in palette_path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or line.startswith("GIMP Palette") or line.startswith("Name:") or line.startswith("Columns:"): + continue + if line.startswith("Channels:"): + continue + hex_color = _parse_hex_color(line) + if hex_color is not None: + colors.append(hex_color) + continue + parts = line.split() + if len(parts) >= 3 and all(part.isdigit() for part in parts[:3]): + colors.append((int(parts[0]), int(parts[1]), int(parts[2]))) + if not colors: + raise ValueError(f"No palette colors found in {palette_path}") + return colors + + +def save_palette_report(path: str | Path, palette: list[tuple[int, int, int]]) -> None: + report = {"palette": palette} + Path(path).write_text(json.dumps(report, indent=2), encoding="utf-8") + + +def derive_palette(rgba: np.ndarray, max_colors: int = 32) -> list[tuple[int, int, int]]: + image = Image.fromarray(np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8), mode="RGBA") + opaque_rgb = np.asarray(image)[..., :3] + unique = np.unique(opaque_rgb.reshape(-1, 3), axis=0) + colors = int(np.clip(unique.shape[0], 8, max_colors)) + quantized = image.convert("RGB").quantize(colors=colors, method=Image.Quantize.MEDIANCUT) + palette = quantized.getpalette()[: colors * 3] + return [tuple(palette[index : index + 3]) for index in range(0, len(palette), 3)] + + +def quantize_rgba( + rgba: np.ndarray, + mode: str, + palette: list[tuple[int, int, int]] | None = None, +) -> PaletteResult | None: + if mode == "off": + return None + if mode == "strict" and not palette: + raise ValueError("Strict palette mode requires a palette file.") + use_palette = palette or derive_palette(rgba) + rgb = np.clip(np.rint(rgba[..., :3] * 255.0), 0, 255).astype(np.int16) + alpha = rgba[..., 3:4] + palette_array = np.asarray(use_palette, dtype=np.int16) + flat = rgb.reshape(-1, 3) + distances = np.sum((flat[:, None, :] - palette_array[None, :, :]) ** 2, axis=-1) + indices = np.argmin(distances, axis=1) + quantized_rgb = palette_array[indices].reshape(rgb.shape).astype(np.float32) / 255.0 + quantized_rgba = np.concatenate([quantized_rgb, alpha], axis=-1) + return PaletteResult(rgba=quantized_rgba, palette=use_palette, indexed_rgba=None, indexed_png_path=None) diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py new file mode 100644 index 0000000..5909e9f --- /dev/null +++ b/src/repixelizer/pipeline.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import time +from pathlib import Path +from typing import Any + +import numpy as np + +from .analysis import analyze_source +from .continuous import optimize_uv_field +from .diagnostics import ( + summarize_run, + write_alpha_preview, + write_comparison, + write_heatmap, + write_json, + write_lattice_overlay, +) +from .discrete import cleanup_pixels +from .inference import infer_lattice, inference_to_json +from .io import load_rgba, save_rgba +from .palette import load_palette, quantize_rgba, save_palette_report +from .types import RunResult + + +def run_pipeline( + input_path: str | Path, + output_path: str | Path, + *, + target_size: int | None = None, + palette_path: str | Path | None = None, + palette_mode: str = "off", + diagnostics_dir: str | Path | None = None, + seed: int = 7, + steps: int = 200, + device: str = "cpu", +) -> RunResult: + started = time.perf_counter() + source = load_rgba(input_path) + inference = infer_lattice(source, target_size=target_size) + analysis = analyze_source(source, seed=seed) + solver = optimize_uv_field( + source, + inference=inference, + analysis=analysis, + steps=steps, + seed=seed, + device=device, + ) + cleanup = cleanup_pixels(solver.target_rgba, source_guidance=solver.guidance_strength) + palette = load_palette(palette_path) if palette_path else None + palette_result = quantize_rgba(cleanup.cleaned_rgba, mode=palette_mode, palette=palette) + output_rgba = palette_result.rgba if palette_result else cleanup.cleaned_rgba + save_rgba(output_path, output_rgba) + + diagnostics: dict[str, Any] = {"elapsed_seconds": time.perf_counter() - started} + result = RunResult( + source_rgba=source, + output_rgba=output_rgba, + inference=inference, + analysis=analysis, + solver=solver, + cleanup=cleanup, + palette_result=palette_result, + diagnostics=diagnostics, + ) + if diagnostics_dir: + diagnostics_path = Path(diagnostics_dir) + diagnostics_path.mkdir(parents=True, exist_ok=True) + write_lattice_overlay(diagnostics_path / "lattice-overlay.png", source, inference) + write_comparison(diagnostics_path / "comparison.png", source, output_rgba) + write_alpha_preview(diagnostics_path / "alpha-preview.png", source, output_rgba) + write_heatmap(diagnostics_path / "noise-heatmap.png", cleanup.isolated_heatmap) + save_rgba(diagnostics_path / "cluster-preview.png", analysis.cluster_preview) + run_json = summarize_run(result) + run_json["inference"] = inference_to_json(inference) + run_json["settings"] = { + "target_size": target_size, + "palette_mode": palette_mode, + "seed": seed, + "steps": steps, + "device": device, + } + write_json(diagnostics_path / "run.json", run_json) + if palette_result is not None: + save_palette_report(diagnostics_path / "palette-report.json", palette_result.palette) + return result diff --git a/src/repixelizer/synthetic.py b/src/repixelizer/synthetic.py new file mode 100644 index 0000000..7dcacd1 --- /dev/null +++ b/src/repixelizer/synthetic.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter + + +def make_emblem(width: int = 32, height: int = 32) -> np.ndarray: + image = Image.new("RGBA", (width, height), (0, 0, 0, 0)) + draw = ImageDraw.Draw(image) + draw.rounded_rectangle((3, 6, 27, 28), radius=8, fill=(194, 24, 24, 255), outline=(255, 220, 96, 255), width=2) + draw.polygon([(16, 3), (20, 9), (16, 29), (12, 9)], fill=(185, 220, 255, 255), outline=(40, 52, 80, 255)) + draw.ellipse((4, 4, 15, 15), fill=(24, 188, 120, 255)) + draw.ellipse((8, 7, 12, 11), fill=(240, 255, 255, 255)) + return np.asarray(image, dtype=np.float32) / 255.0 + + +def make_sprite(width: int = 24, height: int = 24) -> np.ndarray: + image = Image.new("RGBA", (width, height), (0, 0, 0, 0)) + draw = ImageDraw.Draw(image) + draw.rectangle((8, 4, 15, 18), fill=(48, 88, 220, 255), outline=(244, 248, 255, 255)) + draw.rectangle((9, 2, 14, 7), fill=(255, 210, 162, 255), outline=(72, 32, 24, 255)) + draw.line((10, 18, 8, 22), fill=(42, 42, 42, 255), width=1) + draw.line((13, 18, 15, 22), fill=(42, 42, 42, 255), width=1) + draw.line((8, 10, 4, 14), fill=(255, 210, 162, 255), width=1) + draw.line((15, 10, 19, 6), fill=(255, 210, 162, 255), width=1) + return np.asarray(image, dtype=np.float32) / 255.0 + + +def fake_pixelize( + rgba: np.ndarray, + upscale: int = 14, + phase_x: float = 0.22, + phase_y: float = 0.31, + blur_radius: float = 0.7, +) -> np.ndarray: + source = Image.fromarray(np.clip(np.rint(rgba * 255.0), 0, 255).astype(np.uint8), mode="RGBA") + width, height = source.size + target = source.resize((width * upscale, height * upscale), resample=Image.Resampling.NEAREST) + shift_x = int(round(upscale * phase_x)) + shift_y = int(round(upscale * phase_y)) + shifted = Image.new("RGBA", target.size, (0, 0, 0, 0)) + shifted.alpha_composite(target, (shift_x, shift_y)) + blurred = shifted.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + return np.asarray(blurred, dtype=np.float32) / 255.0 diff --git a/src/repixelizer/types.py b/src/repixelizer/types.py new file mode 100644 index 0000000..72dfbdd --- /dev/null +++ b/src/repixelizer/types.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np + + +@dataclass(slots=True) +class InferenceCandidate: + target_width: int + target_height: int + phase_x: float + phase_y: float + score: float + breakdown: dict[str, float] + + +@dataclass(slots=True) +class InferenceResult: + target_width: int + target_height: int + phase_x: float + phase_y: float + confidence: float + top_candidates: list[InferenceCandidate] = field(default_factory=list) + + +@dataclass(slots=True) +class SourceAnalysis: + edge_map: np.ndarray + cluster_map: np.ndarray + cluster_centers: np.ndarray + alpha_map: np.ndarray + cluster_preview: np.ndarray + + +@dataclass(slots=True) +class SolverArtifacts: + target_rgba: np.ndarray + uv_field: np.ndarray + guidance_strength: np.ndarray + initial_rgba: np.ndarray + loss_history: list[float] + + +@dataclass(slots=True) +class CleanupArtifacts: + cleaned_rgba: np.ndarray + isolated_heatmap: np.ndarray + + +@dataclass(slots=True) +class PaletteResult: + rgba: np.ndarray + palette: list[tuple[int, int, int]] + indexed_rgba: np.ndarray | None = None + indexed_png_path: Path | None = None + + +@dataclass(slots=True) +class RunResult: + source_rgba: np.ndarray + output_rgba: np.ndarray + inference: InferenceResult + analysis: SourceAnalysis + solver: SolverArtifacts + cleanup: CleanupArtifacts + palette_result: PaletteResult | None + diagnostics: dict[str, Any] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..7967726 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from pathlib import Path + +from repixelizer.cli import main +from repixelizer.io import save_rgba +from repixelizer.synthetic import fake_pixelize, make_emblem + + +def test_default_command_runs_without_explicit_run_subcommand(tmp_path: Path) -> None: + source = make_emblem(16, 16) + fake = fake_pixelize(source, upscale=8, phase_x=0.2, phase_y=0.25, blur_radius=0.4) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + save_rgba(input_path, fake) + exit_code = main([str(input_path), "--out", str(output_path), "--steps", "8"]) + assert exit_code == 0 + assert output_path.exists() diff --git a/tests/test_compare.py b/tests/test_compare.py new file mode 100644 index 0000000..a32414f --- /dev/null +++ b/tests/test_compare.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from pathlib import Path + +from repixelizer.compare import run_compare +from repixelizer.io import save_rgba +from repixelizer.synthetic import fake_pixelize, make_emblem + + +def test_compare_writes_comparison_artifacts(tmp_path: Path) -> None: + source = make_emblem(32, 32) + fake = fake_pixelize(source, upscale=10, phase_x=0.2, phase_y=0.3, blur_radius=0.5) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + diagnostics_dir = tmp_path / "diagnostics" + save_rgba(input_path, fake) + + payload = run_compare(input_path, output_path, diagnostics_dir=diagnostics_dir, steps=20) + assert output_path.exists() + assert (diagnostics_dir / "compare-sheet.png").exists() + assert (diagnostics_dir / "compare.csv").exists() + assert len(payload["rows"]) == 3 diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..10069af --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from repixelizer.inference import infer_lattice +from repixelizer.synthetic import fake_pixelize, make_emblem + + +def test_infer_lattice_recovers_emblem_scale() -> None: + source = make_emblem(32, 32) + fake = fake_pixelize(source, upscale=12, phase_x=0.2, phase_y=0.35, blur_radius=0.6) + result = infer_lattice(fake) + assert result.target_width in range(28, 37) + assert result.target_height in range(28, 37) + assert result.confidence >= 0.0 diff --git a/tests/test_palette.py b/tests/test_palette.py new file mode 100644 index 0000000..688c2cf --- /dev/null +++ b/tests/test_palette.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from pathlib import Path + +from repixelizer.palette import load_palette, quantize_rgba +from repixelizer.synthetic import make_emblem + + +def test_load_gpl_palette(tmp_path: Path) -> None: + palette_path = tmp_path / "palette.gpl" + palette_path.write_text( + "GIMP Palette\nName: Test\nColumns: 2\n255 0 0 Red\n0 255 0 Green\n", + encoding="utf-8", + ) + palette = load_palette(palette_path) + assert palette == [(255, 0, 0), (0, 255, 0)] + + +def test_strict_quantization_stays_in_palette() -> None: + rgba = make_emblem(8, 8) + palette = [(255, 0, 0), (0, 255, 0), (255, 255, 255)] + result = quantize_rgba(rgba, mode="strict", palette=palette) + assert result is not None + unique = { + tuple((pixel[:3] * 255.0).round().astype(int)) + for pixel in result.rgba.reshape(-1, 4) + if pixel[3] > 0.05 + } + assert unique.issubset(set(palette)) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..3344e9c --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from pathlib import Path + +from repixelizer.baselines import naive_resize_baseline +from repixelizer.metrics import coherence_breakdown +from repixelizer.pipeline import run_pipeline +from repixelizer.synthetic import fake_pixelize, make_emblem, make_sprite + + +def test_pipeline_writes_output_and_diagnostics(tmp_path: Path) -> None: + source = make_emblem(32, 32) + fake = fake_pixelize(source, upscale=10, phase_x=0.15, phase_y=0.25, blur_radius=0.5) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + diagnostics_dir = tmp_path / "diagnostics" + + from repixelizer.io import save_rgba + + save_rgba(input_path, fake) + result = run_pipeline(input_path, output_path, diagnostics_dir=diagnostics_dir, steps=24) + assert output_path.exists() + assert (diagnostics_dir / "run.json").exists() + assert result.output_rgba.shape[0] == result.inference.target_height + assert result.output_rgba.shape[1] == result.inference.target_width + + +def test_pipeline_beats_naive_on_coherence_for_synthetic_emblem(tmp_path: Path) -> None: + source = make_emblem(32, 32) + fake = fake_pixelize(source, upscale=12, phase_x=0.2, phase_y=0.35, blur_radius=0.65) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, fake) + result = run_pipeline(input_path, output_path, steps=32) + naive = naive_resize_baseline(fake, width=result.inference.target_width, height=result.inference.target_height) + optimized_score = coherence_breakdown(result.output_rgba)["coherence_score"] + naive_score = coherence_breakdown(naive)["coherence_score"] + assert optimized_score >= naive_score + + +def test_pipeline_preserves_transparency_for_sprite(tmp_path: Path) -> None: + source = make_sprite(24, 24) + fake = fake_pixelize(source, upscale=9, phase_x=0.1, phase_y=0.22, blur_radius=0.45) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, fake) + result = run_pipeline(input_path, output_path, steps=20) + alpha = result.output_rgba[..., 3] + assert alpha.min() < 0.05 + assert alpha.max() > 0.9