From d2e8ccac4d7ded068042c1b139dfc8b78fd032a5 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Tue, 21 Apr 2026 21:56:25 +0100 Subject: [PATCH] Add corpus benchmarks and lattice refinement --- .gitignore | 17 + README.md | 158 +++----- examples/corpus/README.md | 31 ++ examples/corpus/originals/.gitkeep | 1 + src/repixelizer/benchmark.py | 286 ++++++++++++++ src/repixelizer/cli.py | 105 ++++- src/repixelizer/compare.py | 13 +- src/repixelizer/continuous.py | 610 ++++++++++++++++++++++++++--- src/repixelizer/corpus.py | 500 +++++++++++++++++++++++ src/repixelizer/diagnostics.py | 16 +- src/repixelizer/discrete.py | 6 +- src/repixelizer/inference.py | 321 +++++++++++++-- src/repixelizer/metrics.py | 229 ++++++++++- src/repixelizer/params.py | 34 ++ src/repixelizer/pipeline.py | 81 +++- src/repixelizer/synthetic.py | 178 ++++++++- src/repixelizer/tuning.py | 275 +++++++++++++ tests/test_benchmark.py | 86 ++++ tests/test_corpus.py | 59 +++ tests/test_inference.py | 11 +- tests/test_metrics.py | 98 +++++ tests/test_pipeline.py | 69 +++- tests/test_tuning.py | 32 ++ 23 files changed, 2997 insertions(+), 219 deletions(-) create mode 100644 examples/corpus/README.md create mode 100644 examples/corpus/originals/.gitkeep create mode 100644 src/repixelizer/benchmark.py create mode 100644 src/repixelizer/corpus.py create mode 100644 src/repixelizer/params.py create mode 100644 src/repixelizer/tuning.py create mode 100644 tests/test_benchmark.py create mode 100644 tests/test_corpus.py create mode 100644 tests/test_metrics.py create mode 100644 tests/test_tuning.py diff --git a/.gitignore b/.gitignore index a80289e..7e507c5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,9 +1,26 @@ .venv/ __pycache__/ *.pyc +*.pyo +*.pyd .pytest_cache/ +.mypy_cache/ +.coverage +htmlcov/ build/ dist/ *.egg-info/ + +.vscode/ +.idea/ +Thumbs.db +.DS_Store + artifacts/ diagnostics/ + +examples/corpus/source-sheets/ +examples/corpus/prepare-corpus.json +examples/corpus/ATTRIBUTION.md +examples/corpus/originals/* +!examples/corpus/originals/.gitkeep diff --git a/README.md b/README.md index 1737a20..2237362 100644 --- a/README.md +++ b/README.md @@ -2,35 +2,19 @@ Repixelizer is a standalone Python CLI for turning "fake pixel art" into true grid-aligned pixel art. -The core idea is simple: generated images often look like pixel art without actually obeying a single clean pixel grid. Repixelizer treats that as an optimization problem instead of a resize problem. It estimates the implied lattice, builds a real target grid, optimizes how output pixels sample the source image, then cleans the result up on the discrete grid. +The project targets images that look locally pixelated but do not actually sit on one clean lattice. Instead of treating that as a resize problem, Repixelizer treats it as a lattice-inference and local-structure-preservation problem: infer the implied grid, choose real output cells, and snap them back onto a coherent pixel lattice while preserving the source's local adjacency patterns. -This project is meant to stand on its own. It should be understandable without any StreamPixels-specific context. +## Current Status -## What it does +This is already a serious experimental baseline, not a toy prototype: -Repixelizer currently focuses on single-image inputs such as: -- icons -- emblems -- logos -- simple sprites +- CUDA-backed inference and solver paths are wired up +- round-trip corpus benchmarking is built in +- soft and crisp corruption profiles are supported +- diagnostics and baseline comparisons are written automatically +- a black-box tuning harness exists for longer offline parameter sweeps -The current pipeline is: - -1. infer the source image's implied fake-pixel lattice -2. choose a target output resolution, unless the user overrides it -3. initialize a UV/sample field over the target grid -4. optimize that field in premultiplied RGBA space with PyTorch -5. project onto a true pixel grid -6. run discrete cleanup to reduce noise and isolated-pixel garbage -7. optionally quantize into a palette-constrained result -8. write diagnostics and baseline comparisons - -## Project docs - -- Product and technical spec: [docs/spec.md]() -- Living implementation plan: [docs/implementation-plan.md]() - -If you are moving this project into a fresh workspace, those two docs are the main handoff. +The current pipeline is especially strong on benchmarked sprite and emblem inputs where local fake-pixel texture is meaningful but the global grid is inconsistent. ## Quickstart @@ -43,7 +27,7 @@ Run the optimizer: ```powershell repixelize input.png --out output.png -repixelize input.png --out output.png --target-size 120 --diagnostics-dir diagnostics +repixelize input.png --out output.png --diagnostics-dir diagnostics --device auto ``` Run the optimizer plus baselines: @@ -52,106 +36,78 @@ Run the optimizer plus baselines: repixelize compare input.png --out output.png --diagnostics-dir diagnostics ``` -## CLI contract +## Corpus And Benchmarks -Primary command shape: +Prepare a local Creative Commons corpus: ```powershell -repixelize input.png --out output.png -repixelize compare input.png --out output.png +repixelize prepare-corpus --corpus-dir examples/corpus ``` -Important flags: +Run the round-trip benchmark: -- `--target-size `: override inferred target max dimension -- `--palette `: palette file in `.gpl`, `.txt`, or `.json` -- `--palette-mode off|fit|strict`: palette behavior -- `--diagnostics-dir `: write visual and JSON diagnostics -- `--seed `: deterministic optimizer seed -- `--steps `: number of optimization steps -- `--device cpu|cuda`: PyTorch device +```powershell +repixelize benchmark --corpus-dir examples/corpus --out-dir artifacts/benchmark +``` -### Palette modes +Useful benchmark flags: -- `off`: keep the output unconstrained -- `fit`: adapt the result to a derived or supplied palette -- `strict`: stay strictly inside the supplied palette +- `--profile soft` or `--profile crisp` +- `--case ` repeated for a focused slice +- `--limit ` for the first `n` matching cases +- `--infer-size` to test automatic size inference +- `--keep-existing` if you explicitly want to preserve an old output directory -Palette constraints are optional by design. Many fake-pixel-art inputs are not born from a coherent palette, so palette enforcement is useful but not assumed. +By default the benchmark clears its output directory before each run so the artifacts folder stays readable. -## Outputs +## Tuning -Main output: -- final RGBA PNG +Repixelizer includes a black-box tuning loop for longer offline searches over solver weights: -Optional diagnostics: -- `run.json`: chosen size, phase, settings, timings, and score breakdown -- `lattice-overlay.png`: inferred lattice preview over the source image -- `comparison.png`: source, optimized output, and nearest-neighbor preview -- `alpha-preview.png`: alpha behavior before and after -- `noise-heatmap.png`: isolated-pixel and local-noise hotspots -- `cluster-preview.png`: coarse material/color clustering +```powershell +repixelize tune --corpus-dir examples/corpus --out-dir artifacts/tuning --profile soft --limit 8 +``` -Additional compare-mode outputs: -- `compare.json` -- `compare.csv` -- `compare-sheet.png` +This is intentionally a search-based tuner, not gradient descent. The objective depends on discrete argmins, thresholding, and benchmark comparisons, so it is better treated as a repeatable black-box optimization problem. -## Repo layout +## Repository Hygiene + +This repo is set up to be safe for local experimentation: + +- generated outputs under `artifacts/` are ignored +- local corpus payloads under `examples/corpus/originals/` are ignored +- archived raw sprite sheets under `examples/corpus/source-sheets/` are ignored +- generated corpus metadata files are ignored + +That keeps benchmark assets, attribution exports, and tuning runs from polluting upstream history while still keeping the code, docs, and corpus layout tracked. + +## Repo Layout - `src/repixelizer`: core package -- `tests`: synthetic and pipeline-focused tests -- `examples`: reserved for future real example inputs -- `artifacts`: local smoke-test outputs, not core source +- `tests`: focused regression tests +- `docs/spec.md`: product and technical spec +- `docs/implementation-plan.md`: working roadmap +- `examples/corpus/README.md`: local corpus layout and attribution workflow Core modules: -- `inference.py`: target-size and phase inference -- `analysis.py`: edge, alpha, and coarse clustering analysis -- `continuous.py`: UV-field optimization -- `discrete.py`: local cleanup after projection -- `palette.py`: palette loading and optional quantization -- `compare.py`: baseline comparisons and contact-sheet output -- `synthetic.py`: benchmark generators for tests -## Current status - -The project is bootstrapped and runnable: -- editable install works -- PyTorch-based optimization is wired in -- synthetic tests are in place -- compare mode writes diagnostics and baseline metrics - -What exists today is a serious experimental baseline, not a finished research result. It is strong enough to iterate on algorithmically, but it is not yet proven across a broad real-world image corpus. - -## Current limitations - -- single-image only -- no multi-frame sprite-sheet consistency yet -- no GUI yet -- no user-painted masks or region hints yet -- lattice inference is heuristic, not learned -- continuous optimization is real, but still early and intentionally conservative +- `inference.py`: target size and phase inference +- `continuous.py`: source-aware lattice snapping and local discrete refinement +- `metrics.py`: fidelity, adjacency, and motif metrics +- `benchmark.py`: corpus benchmark runner +- `tuning.py`: black-box hyperparameter search harness +- `synthetic.py`: facsimile generation for tests and benchmarks ## Validation -The current test suite checks: -- resolution inference on synthetic emblem inputs -- palette loading and strict palette enforcement -- pipeline output and diagnostics writing -- compare-mode artifact generation -- a basic default CLI execution path - -Run it with: +Run the focused test suite with: ```powershell .venv\Scripts\python -m pytest -q ``` -## Near-term direction +## Notes -The next useful steps are: -- improve inference on real-world generated inputs -- make the optimizer outperform resize-and-dither baselines more consistently -- expand the curated benchmark corpus -- add smarter discrete cleanup rules for outlines, highlights, and material bands -- optionally add manual guidance in a later version, once the automatic baseline is stable +- The local corpus is optional; the codebase is still usable without it. +- The benchmark and tuning commands are designed for iterative local work, not for shipping generated artifacts in git. +- Palette enforcement is optional and can be left off for unconstrained recovery runs. diff --git a/examples/corpus/README.md b/examples/corpus/README.md new file mode 100644 index 0000000..22cf021 --- /dev/null +++ b/examples/corpus/README.md @@ -0,0 +1,31 @@ +# Corpus Layout + +Put Creative Commons source pixel art in [originals](). +If you import sprite sheets or collage sheets first, run `repixelize prepare-corpus --corpus-dir examples/corpus` to archive the raw sheets into `source-sheets/` and rewrite `originals/` into one-sprite-per-file inputs for benchmarking. +That step also refreshes [ATTRIBUTION.md]() from the processed sidecars. + +Recommended layout: + +- `examples/corpus/originals/.png` +- `examples/corpus/originals/.json` + +The JSON sidecar is optional, but useful for attribution and notes. Suggested fields: + +```json +{ + "title": "Sprite Name", + "author": "Artist Name", + "license": "CC-BY 4.0", + "source_url": "https://example.com/original", + "notes": "Strong outline, metallic highlights." +} +``` + +Run the round-trip benchmark with: + +```powershell +repixelize prepare-corpus --corpus-dir examples/corpus +repixelize benchmark --corpus-dir examples/corpus --out-dir artifacts/benchmark +``` + +By default the benchmark locks the target size to the original image dimensions so we can measure recovery quality cleanly. Add `--infer-size` if you want to evaluate automatic size inference too. diff --git a/examples/corpus/originals/.gitkeep b/examples/corpus/originals/.gitkeep new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/examples/corpus/originals/.gitkeep @@ -0,0 +1 @@ + diff --git a/src/repixelizer/benchmark.py b/src/repixelizer/benchmark.py new file mode 100644 index 0000000..c7ebcd8 --- /dev/null +++ b/src/repixelizer/benchmark.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import Any +import json + +import numpy as np + +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, + exact_match_ratio, + foreground_adjacency_error, + foreground_coverage, + foreground_exact_match_ratio, + foreground_motif_error, + foreground_reconstruction_error, + reconstruction_error, +) +from .params import SolverHyperParams +from .pipeline import run_pipeline +from .synthetic import fake_pixelize + +SUPPORTED_EXTENSIONS = {".png", ".bmp", ".gif", ".webp"} + + +def run_roundtrip_benchmark( + corpus_dir: str | Path, + out_dir: str | Path, + *, + variants: int = 3, + profiles: list[str] | None = None, + seed: int = 7, + steps: int = 200, + device: str = "auto", + infer_size: bool = False, + include_cases: list[str] | None = None, + limit_cases: int | None = None, + keep_existing: bool = False, + solver_params: SolverHyperParams | None = None, +) -> dict[str, Any]: + corpus_path = Path(corpus_dir) + originals_dir = corpus_path / "originals" + all_original_paths = sorted( + path for path in originals_dir.rglob("*") if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS + ) + original_paths = [ + path + for path in all_original_paths + if _case_matches(path.relative_to(originals_dir).with_suffix("").as_posix(), path, include_cases or []) + ] + if limit_cases is not None: + original_paths = original_paths[: max(0, limit_cases)] + if not original_paths: + raise RuntimeError(f"No source pixel-art originals found in {originals_dir}.") + selected_profiles = profiles or ["soft", "crisp"] + + out_path = Path(out_dir) + if out_path.exists() and not keep_existing: + shutil.rmtree(out_path) + out_path.mkdir(parents=True, exist_ok=True) + rows: list[dict[str, Any]] = [] + for case_index, original_path in enumerate(original_paths): + original = load_rgba(original_path) + metadata = _load_metadata(original_path) + case_id = original_path.relative_to(originals_dir).with_suffix("").as_posix() + case_dir = out_path / "cases" / Path(case_id) + case_dir.mkdir(parents=True, exist_ok=True) + save_rgba(case_dir / "original.png", original) + for profile_index, profile in enumerate(selected_profiles): + for variant_index in range(variants): + variant_seed = seed + case_index * 1000 + profile_index * 100 + variant_index + settings = _variant_settings(variant_seed, profile=profile) + variant_dir = case_dir / f"profile-{profile}" / f"variant-{variant_index + 1:02d}" + diagnostics_dir = variant_dir / "diagnostics" + diagnostics_dir.mkdir(parents=True, exist_ok=True) + fake = fake_pixelize(original, seed=variant_seed, **settings) + input_path = variant_dir / "input.png" + output_path = variant_dir / "optimized.png" + save_rgba(input_path, fake) + + locked_target_size = max(original.shape[0], original.shape[1]) if not infer_size else None + result = run_pipeline( + input_path, + output_path, + target_size=locked_target_size, + diagnostics_dir=diagnostics_dir, + seed=variant_seed, + steps=steps, + device=device, + solver_params=solver_params, + ) + + naive = naive_resize_baseline(fake, width=original.shape[1], height=original.shape[0]) + diffusion = error_diffusion_baseline(fake, width=original.shape[1], height=original.shape[0]) + save_rgba(variant_dir / "naive.png", naive) + save_rgba(variant_dir / "diffusion.png", diffusion) + _save_preview(variant_dir / "optimized-preview.png", result.output_rgba, original) + _save_preview(variant_dir / "naive-preview.png", naive, original) + _save_preview(variant_dir / "diffusion-preview.png", diffusion, original) + + optimized_preview = nearest_resize(result.output_rgba, width=original.shape[1], height=original.shape[0]) + optimized_error = foreground_reconstruction_error(optimized_preview, original) + naive_error = foreground_reconstruction_error(naive, original) + diffusion_error = foreground_reconstruction_error(diffusion, original) + optimized_adjacency = foreground_adjacency_error(optimized_preview, original) + naive_adjacency = foreground_adjacency_error(naive, original) + diffusion_adjacency = foreground_adjacency_error(diffusion, original) + optimized_motif = foreground_motif_error(optimized_preview, original) + naive_motif = foreground_motif_error(naive, original) + diffusion_motif = foreground_motif_error(diffusion, original) + row = { + "case_id": case_id, + "profile": profile, + "variant": variant_index + 1, + "source_file": str(original_path), + "title": metadata.get("title", ""), + "author": metadata.get("author", ""), + "license": metadata.get("license", ""), + "source_url": metadata.get("source_url", ""), + "original_width": int(original.shape[1]), + "original_height": int(original.shape[0]), + "input_width": int(fake.shape[1]), + "input_height": int(fake.shape[0]), + "upscale": settings["upscale"], + "phase_x": settings["phase_x"], + "phase_y": settings["phase_y"], + "blur_radius": settings["blur_radius"], + "warp_strength": settings["warp_strength"], + "warp_detail": settings["warp_detail"], + "warp_sample_mode": settings["warp_sample_mode"], + "primary_metric": "foreground_premultiplied_mae", + "reference_foreground_coverage": foreground_coverage(original, original), + "target_size_locked": not infer_size, + "inferred_width": result.inference.target_width, + "inferred_height": result.inference.target_height, + "inference_confidence": result.inference.confidence, + "optimized_error_to_original": optimized_error, + "naive_error_to_original": naive_error, + "diffusion_error_to_original": diffusion_error, + "optimized_adjacency_error_to_original": optimized_adjacency, + "naive_adjacency_error_to_original": naive_adjacency, + "diffusion_adjacency_error_to_original": diffusion_adjacency, + "optimized_motif_error_to_original": optimized_motif, + "naive_motif_error_to_original": naive_motif, + "diffusion_motif_error_to_original": diffusion_motif, + "optimized_canvas_error_to_original": reconstruction_error(optimized_preview, original), + "naive_canvas_error_to_original": reconstruction_error(naive, original), + "diffusion_canvas_error_to_original": reconstruction_error(diffusion, original), + "optimized_exact_match": foreground_exact_match_ratio(optimized_preview, original), + "naive_exact_match": foreground_exact_match_ratio(naive, original), + "diffusion_exact_match": foreground_exact_match_ratio(diffusion, original), + "optimized_canvas_exact_match": exact_match_ratio(optimized_preview, original), + "naive_canvas_exact_match": exact_match_ratio(naive, original), + "diffusion_canvas_exact_match": exact_match_ratio(diffusion, original), + "optimized_coherence": coherence_breakdown(result.output_rgba)["coherence_score"], + "naive_coherence": coherence_breakdown(naive)["coherence_score"], + "diffusion_coherence": coherence_breakdown(diffusion)["coherence_score"], + "optimized_beats_naive_error": optimized_error <= naive_error, + "optimized_beats_diffusion_error": optimized_error <= diffusion_error, + "optimized_beats_naive_adjacency": optimized_adjacency <= naive_adjacency, + "optimized_beats_diffusion_adjacency": optimized_adjacency <= diffusion_adjacency, + "optimized_beats_naive_motif": optimized_motif <= naive_motif, + "optimized_beats_diffusion_motif": optimized_motif <= diffusion_motif, + } + rows.append(row) + + summary = { + "case_count": len(original_paths), + "row_count": len(rows), + "variants_per_case": variants, + "profiles": selected_profiles, + "target_size_locked": not infer_size, + "primary_metric": "foreground_premultiplied_mae", + "solver_params": solver_params.to_dict() if solver_params is not None else SolverHyperParams().to_dict(), + "cases": _summarize_cases(rows), + "rows": rows, + } + write_compare_csv(out_path / "benchmark.csv", rows) + write_json(out_path / "benchmark.json", summary) + return summary + + +def _load_metadata(original_path: Path) -> dict[str, Any]: + metadata_path = original_path.with_suffix(".json") + if not metadata_path.exists(): + return {} + return json.loads(metadata_path.read_text(encoding="utf-8")) + + +def _case_matches(case_id: str, original_path: Path, include_cases: list[str]) -> bool: + if not include_cases: + return True + normalized = {token.lower() for token in include_cases} + aliases = { + case_id, + case_id.lower(), + original_path.stem, + original_path.stem.lower(), + original_path.name, + original_path.name.lower(), + } + return any(token in aliases for token in normalized) + + +def _variant_settings(seed: int, profile: str) -> dict[str, Any]: + rng = np.random.default_rng(seed) + if profile == "crisp": + return { + "upscale": int(rng.integers(8, 15)), + "phase_x": float(rng.uniform(-0.45, 0.45)), + "phase_y": float(rng.uniform(-0.45, 0.45)), + "blur_radius": 0.0, + "warp_strength": float(rng.uniform(0.14, 0.46)), + "warp_detail": int(rng.integers(5, 9)), + "warp_sample_mode": "nearest", + } + if profile == "soft": + return { + "upscale": int(rng.integers(8, 15)), + "phase_x": float(rng.uniform(-0.45, 0.45)), + "phase_y": float(rng.uniform(-0.45, 0.45)), + "blur_radius": float(rng.uniform(0.45, 1.15)), + "warp_strength": float(rng.uniform(0.18, 0.55)), + "warp_detail": int(rng.integers(4, 8)), + "warp_sample_mode": "bilinear", + } + raise ValueError(f"Unsupported benchmark profile: {profile}") + + +def _save_preview(path: Path, rgba: np.ndarray, original: np.ndarray) -> None: + preview = nearest_resize(rgba, width=original.shape[1] * 8, height=original.shape[0] * 8) + save_rgba(path, preview) + + +def _summarize_cases(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[tuple[str, str], list[dict[str, Any]]] = {} + for row in rows: + grouped.setdefault((str(row["case_id"]), str(row["profile"])), []).append(row) + summaries = [] + for (case_id, profile), case_rows in sorted(grouped.items()): + optimized = [float(row["optimized_error_to_original"]) for row in case_rows] + naive = [float(row["naive_error_to_original"]) for row in case_rows] + diffusion = [float(row["diffusion_error_to_original"]) for row in case_rows] + optimized_adjacency = [float(row["optimized_adjacency_error_to_original"]) for row in case_rows] + naive_adjacency = [float(row["naive_adjacency_error_to_original"]) for row in case_rows] + diffusion_adjacency = [float(row["diffusion_adjacency_error_to_original"]) for row in case_rows] + optimized_motif = [float(row["optimized_motif_error_to_original"]) for row in case_rows] + naive_motif = [float(row["naive_motif_error_to_original"]) for row in case_rows] + diffusion_motif = [float(row["diffusion_motif_error_to_original"]) for row in case_rows] + summaries.append( + { + "case_id": case_id, + "profile": profile, + "variants": len(case_rows), + "optimized_error_mean": float(np.mean(optimized)), + "naive_error_mean": float(np.mean(naive)), + "diffusion_error_mean": float(np.mean(diffusion)), + "optimized_adjacency_error_mean": float(np.mean(optimized_adjacency)), + "naive_adjacency_error_mean": float(np.mean(naive_adjacency)), + "diffusion_adjacency_error_mean": float(np.mean(diffusion_adjacency)), + "optimized_motif_error_mean": float(np.mean(optimized_motif)), + "naive_motif_error_mean": float(np.mean(naive_motif)), + "diffusion_motif_error_mean": float(np.mean(diffusion_motif)), + "optimized_beats_naive_rate": float(np.mean([row["optimized_beats_naive_error"] for row in case_rows])), + "optimized_beats_diffusion_rate": float( + np.mean([row["optimized_beats_diffusion_error"] for row in case_rows]) + ), + "optimized_beats_naive_adjacency_rate": float( + np.mean([row["optimized_beats_naive_adjacency"] for row in case_rows]) + ), + "optimized_beats_diffusion_adjacency_rate": float( + np.mean([row["optimized_beats_diffusion_adjacency"] for row in case_rows]) + ), + "optimized_beats_naive_motif_rate": float( + np.mean([row["optimized_beats_naive_motif"] for row in case_rows]) + ), + "optimized_beats_diffusion_motif_rate": float( + np.mean([row["optimized_beats_diffusion_motif"] for row in case_rows]) + ), + } + ) + return summaries diff --git a/src/repixelizer/cli.py b/src/repixelizer/cli.py index bd21ad7..61c82a2 100644 --- a/src/repixelizer/cli.py +++ b/src/repixelizer/cli.py @@ -3,8 +3,11 @@ from __future__ import annotations import argparse import sys +from .benchmark import run_roundtrip_benchmark from .compare import run_compare +from .corpus import prepare_corpus from .pipeline import run_pipeline +from .tuning import tune_solver_hyperparams def build_parser() -> argparse.ArgumentParser: @@ -25,18 +28,83 @@ def build_parser() -> argparse.ArgumentParser: 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") + target.add_argument("--device", default="auto", choices=("auto", "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) + benchmark_parser = subparsers.add_parser("benchmark", help="Run the round-trip corpus benchmark.") + benchmark_parser.add_argument("--corpus-dir", default="examples/corpus", help="Corpus root containing originals/") + benchmark_parser.add_argument("--out-dir", default="artifacts/benchmark", help="Directory for benchmark outputs") + benchmark_parser.add_argument("--variants", type=int, default=3, help="Facsimile variants per original") + benchmark_parser.add_argument( + "--profile", + action="append", + choices=("soft", "crisp"), + default=None, + help="Corruption profile to include; may be repeated. Defaults to both soft and crisp.", + ) + benchmark_parser.add_argument("--seed", type=int, default=7, help="Random seed") + benchmark_parser.add_argument("--steps", type=int, default=200, help="Number of optimizer steps") + benchmark_parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"), help="Torch device") + benchmark_parser.add_argument( + "--case", + action="append", + default=None, + help="Restrict benchmark to a specific case id or sprite basename; may be repeated", + ) + benchmark_parser.add_argument("--limit", type=int, default=None, help="Only benchmark the first N matching cases") + benchmark_parser.add_argument( + "--infer-size", + action="store_true", + help="Let the pipeline infer output size instead of locking it to the original ground-truth size", + ) + benchmark_parser.add_argument( + "--keep-existing", + action="store_true", + help="Keep any existing files in the benchmark output directory instead of clearing it first", + ) + tune_parser = subparsers.add_parser("tune", help="Run a black-box hyperparameter search on a benchmark slice.") + tune_parser.add_argument("--corpus-dir", default="examples/corpus", help="Corpus root containing originals/") + tune_parser.add_argument("--out-dir", default="artifacts/tuning", help="Directory for tuning outputs") + tune_parser.add_argument("--trials", type=int, default=8, help="Number of parameter sets to evaluate") + tune_parser.add_argument("--variants", type=int, default=1, help="Facsimile variants per original") + tune_parser.add_argument( + "--profile", + action="append", + choices=("soft", "crisp"), + default=None, + help="Corruption profile to include; may be repeated. Defaults to soft only for tuning.", + ) + tune_parser.add_argument("--seed", type=int, default=7, help="Random seed") + tune_parser.add_argument("--steps", type=int, default=48, help="Number of optimizer steps") + tune_parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"), help="Torch device") + tune_parser.add_argument( + "--case", + action="append", + default=None, + help="Restrict tuning to a specific case id or sprite basename; may be repeated", + ) + tune_parser.add_argument( + "--limit", + type=int, + default=None, + help="Only tune on the first N matching cases. Defaults to 8 when no cases are specified.", + ) + tune_parser.add_argument( + "--infer-size", + action="store_true", + help="Let the pipeline infer output size instead of locking it to the original ground-truth size", + ) + prepare_parser = subparsers.add_parser("prepare-corpus", help="Normalize imported corpus sheets into single sprites.") + prepare_parser.add_argument("--corpus-dir", default="examples/corpus", help="Corpus root containing originals/") 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"}: + if not raw_argv or raw_argv[0] not in {"run", "compare", "benchmark", "tune", "prepare-corpus", "-h", "--help"}: raw_argv = ["run", *raw_argv] parser = build_parser() args = parser.parse_args(raw_argv) @@ -54,6 +122,39 @@ def main(argv: list[str] | None = None) -> int: device=args.device, ) return 0 + if command == "benchmark": + run_roundtrip_benchmark( + args.corpus_dir, + args.out_dir, + variants=args.variants, + profiles=args.profile, + seed=args.seed, + steps=args.steps, + device=args.device, + infer_size=args.infer_size, + include_cases=args.case, + limit_cases=args.limit, + keep_existing=args.keep_existing, + ) + return 0 + if command == "tune": + tune_solver_hyperparams( + args.corpus_dir, + args.out_dir, + trials=args.trials, + variants=args.variants, + profiles=args.profile, + seed=args.seed, + steps=args.steps, + device=args.device, + infer_size=args.infer_size, + include_cases=args.case, + limit_cases=args.limit, + ) + return 0 + if command == "prepare-corpus": + prepare_corpus(args.corpus_dir) + return 0 run_pipeline( args.input, diff --git a/src/repixelizer/compare.py b/src/repixelizer/compare.py index fc291bb..c2b4042 100644 --- a/src/repixelizer/compare.py +++ b/src/repixelizer/compare.py @@ -9,7 +9,13 @@ 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 .metrics import ( + coherence_breakdown, + foreground_adjacency_error, + foreground_motif_error, + foreground_reconstruction_error, + reconstruction_error, +) from .palette import load_palette from .pipeline import run_pipeline @@ -51,7 +57,7 @@ def run_compare( diagnostics_dir: str | Path | None = None, seed: int = 7, steps: int = 200, - device: str = "cpu", + device: str = "auto", ) -> 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) @@ -94,6 +100,9 @@ def run_compare( "isolated_penalty": coherence["isolated_penalty"], "color_chatter": coherence["color_chatter"], "reconstruction_error": reconstruction_error(preview, source), + "foreground_reconstruction_error": foreground_reconstruction_error(preview, source), + "foreground_adjacency_error": foreground_adjacency_error(preview, source), + "foreground_motif_error": foreground_motif_error(preview, source), } ) write_compare_csv(diagnostics_path / "compare.csv", rows) diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py index a07f07a..23798e5 100644 --- a/src/repixelizer/continuous.py +++ b/src/repixelizer/continuous.py @@ -1,10 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass - import numpy as np from .io import premultiply, unpremultiply +from .params import SolverHyperParams from .types import InferenceResult, SolverArtifacts, SourceAnalysis @@ -19,6 +18,17 @@ def _require_torch(): return torch, F +def _resolve_device(torch, requested: str) -> str: + if requested == "auto": + return "cuda" if torch.cuda.is_available() else "cpu" + if requested == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested, but this PyTorch build does not have a usable CUDA device. " + "Install a CUDA-enabled PyTorch build or use --device cpu." + ) + return requested + + 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 @@ -42,6 +52,505 @@ def _cluster_boundary_map(cluster_map: np.ndarray) -> np.ndarray: return boundary +def _make_patch_offsets(height: int, width: int, target_height: int, target_width: int) -> np.ndarray: + cell_x = width / max(1, target_width) + cell_y = height / max(1, target_height) + fractions = np.asarray([-0.36, -0.12, 0.12, 0.36], dtype=np.float32) + offsets: list[tuple[float, float]] = [] + for fy in fractions: + for fx in fractions: + offsets.append( + ( + (fx * cell_x / max(1.0, width - 1)) * 2.0, + (fy * cell_y / max(1.0, height - 1)) * 2.0, + ) + ) + return np.asarray(offsets, dtype=np.float32) + + +def _sample_source(F, source_t, grid, *, mode: str = "bilinear"): + sampled = F.grid_sample(source_t, grid, align_corners=True, mode=mode, padding_mode="border") + return sampled.permute(0, 2, 3, 1) + + +def _sample_cell_patches(F, source_t, uv, offsets_t): + patch_grid = uv[:, :, :, None, :] + offsets_t[None, None, None, :, :] + patch_grid = patch_grid.clamp(-1.0, 1.0) + height = patch_grid.shape[1] + width = patch_grid.shape[2] + samples = patch_grid.shape[3] + flattened = patch_grid.reshape(1, height, width * samples, 2) + sampled = _sample_source(F, source_t, flattened, mode="bilinear") + return sampled.reshape(1, height, width, samples, sampled.shape[-1]) + + +def _representative_colors(patches, solver_params: SolverHyperParams): + patch_mean = patches.mean(dim=3, keepdim=True) + distances = (patches - patch_mean).abs().mean(dim=-1) + weights = distances.mul(-solver_params.representative_softmax_scale).softmax(dim=3) + representative = (patches * weights[..., None]).sum(dim=3) + coherence = (patches - representative[:, :, :, None, :]).abs().mean() + return representative, coherence + + +def _source_boundary_deltas(F, source_t, uv, axis: str, probe_scale: float = 0.22): + if axis == "x": + if uv.shape[2] < 2: + return None + delta = uv[:, :, 1:, :] - uv[:, :, :-1, :] + midpoint = (uv[:, :, 1:, :] + uv[:, :, :-1, :]) * 0.5 + else: + if uv.shape[1] < 2: + return None + delta = uv[:, 1:, :, :] - uv[:, :-1, :, :] + midpoint = (uv[:, 1:, :, :] + uv[:, :-1, :, :]) * 0.5 + left = (midpoint - delta * probe_scale).clamp(-1.0, 1.0) + right = (midpoint + delta * probe_scale).clamp(-1.0, 1.0) + source_left = _sample_source(F, source_t, left, mode="bilinear") + source_right = _sample_source(F, source_t, right, mode="bilinear") + return source_right - source_left + + +def _boundary_pattern_loss(F, source_t, uv, representative, axis: str, solver_params: SolverHyperParams): + source_delta = _source_boundary_deltas( + F, + source_t, + uv, + axis=axis, + probe_scale=solver_params.boundary_probe_scale, + ) + if source_delta is None: + return representative.new_tensor(0.0) + if axis == "x": + rep_delta = representative[:, :, 1:, :] - representative[:, :, :-1, :] + else: + rep_delta = representative[:, 1:, :, :] - representative[:, :-1, :, :] + magnitude_loss = ( + rep_delta.abs().mean(dim=-1) - source_delta.abs().mean(dim=-1) + ).abs().mean() + signed_loss = (rep_delta - source_delta).abs().mean() + rep_norm = rep_delta / rep_delta.abs().mean(dim=-1, keepdim=True).clamp_min(1e-4) + source_norm = source_delta / source_delta.abs().mean(dim=-1, keepdim=True).clamp_min(1e-4) + direction_loss = (rep_norm - source_norm).abs().mean() + return ( + signed_loss * solver_params.boundary_signed_weight + + direction_loss * solver_params.boundary_direction_weight + + magnitude_loss * solver_params.boundary_magnitude_weight + ) + + +def _adjacency_pattern_loss(representative, reference): + losses = [] + if representative.shape[2] > 1: + rep_dx = representative[:, :, 1:, :] - representative[:, :, :-1, :] + ref_dx = reference[:, :, 1:, :] - reference[:, :, :-1, :] + rep_dx_norm = rep_dx / rep_dx.abs().mean(dim=-1, keepdim=True).clamp_min(1e-4) + ref_dx_norm = ref_dx / ref_dx.abs().mean(dim=-1, keepdim=True).clamp_min(1e-4) + losses.append((rep_dx - ref_dx).abs().mean() * 0.65 + (rep_dx_norm - ref_dx_norm).abs().mean() * 0.35) + if representative.shape[1] > 1: + rep_dy = representative[:, 1:, :, :] - representative[:, :-1, :, :] + ref_dy = reference[:, 1:, :, :] - reference[:, :-1, :, :] + rep_dy_norm = rep_dy / rep_dy.abs().mean(dim=-1, keepdim=True).clamp_min(1e-4) + ref_dy_norm = ref_dy / ref_dy.abs().mean(dim=-1, keepdim=True).clamp_min(1e-4) + losses.append((rep_dy - ref_dy).abs().mean() * 0.65 + (rep_dy_norm - ref_dy_norm).abs().mean() * 0.35) + if not losses: + return representative.new_tensor(0.0) + return sum(losses) / len(losses) + + +def _motif_pattern_loss(representative, reference): + if representative.shape[1] < 2 or representative.shape[2] < 2: + return _adjacency_pattern_loss(representative, reference) + import torch + + def motif_blocks(tensor): + blocks = torch.stack( + [ + tensor[:, :-1, :-1, :], + tensor[:, :-1, 1:, :], + tensor[:, 1:, :-1, :], + tensor[:, 1:, 1:, :], + ], + dim=-2, + ) + centered = blocks - blocks.mean(dim=-2, keepdim=True) + scale = centered.abs().mean(dim=(-2, -1), keepdim=True).clamp_min(1e-4) + return centered / scale + + return (motif_blocks(representative) - motif_blocks(reference)).abs().mean() + + +def _snap_output_to_source_pixels( + torch, + source_t, + uv_t, + representative_t, + solver_params: SolverHyperParams, + *, + cell_x: float, + cell_y: float, +) -> np.ndarray: + source_hw = source_t[0].permute(1, 2, 0) + height = source_hw.shape[0] + width = source_hw.shape[1] + uv = uv_t[0] + representative = representative_t[0] + output_height = representative.shape[0] + output_width = representative.shape[1] + + fractions = torch.tensor([-0.42, -0.21, 0.0, 0.21, 0.42], device=uv.device, dtype=uv.dtype) + offset_y, offset_x = torch.meshgrid(fractions, fractions, indexing="ij") + offset_x = (offset_x.reshape(-1) * cell_x).to(dtype=uv.dtype) + offset_y = (offset_y.reshape(-1) * cell_y).to(dtype=uv.dtype) + + xs = (uv[..., 0] + 1.0) * 0.5 * max(1.0, float(width - 1)) + ys = (uv[..., 1] + 1.0) * 0.5 * max(1.0, float(height - 1)) + candidate_x = torch.round(xs[..., None] + offset_x).clamp(0, max(0, width - 1)).to(dtype=torch.long) + candidate_y = torch.round(ys[..., None] + offset_y).clamp(0, max(0, height - 1)).to(dtype=torch.long) + candidate_colors = source_hw[candidate_y, candidate_x] + + base_energy = (candidate_colors - representative[..., None, :]).abs().mean(dim=-1) * solver_params.snap_base_match_weight + representative_delta_x = representative[:, 1:, :] - representative[:, :-1, :] if output_width > 1 else None + representative_delta_y = representative[1:, :, :] - representative[:-1, :, :] if output_height > 1 else None + representative_delta_diag = ( + representative[1:, 1:, :] - representative[:-1, :-1, :] if output_height > 1 and output_width > 1 else None + ) + representative_delta_anti = ( + representative[1:, :-1, :] - representative[:-1, 1:, :] if output_height > 1 and output_width > 1 else None + ) + + selected = torch.argmin(base_energy, dim=-1) + for _ in range(4): + selected_colors = torch.gather( + candidate_colors, + dim=2, + index=selected[..., None, None].expand(output_height, output_width, 1, candidate_colors.shape[-1]), + ).squeeze(2) + energy = base_energy.clone() + if output_width > 1: + left_selected = torch.zeros_like(selected_colors) + left_selected[:, 1:, :] = selected_colors[:, :-1, :] + left_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + left_mask[:, 1:, :] = 1.0 + left_desired = torch.zeros_like(selected_colors) + left_desired[:, 1:, :] = representative_delta_x + left_error = ((candidate_colors - left_selected[..., None, :]) - left_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + left_error * left_mask * solver_params.snap_neighbor_weight + + right_selected = torch.zeros_like(selected_colors) + right_selected[:, :-1, :] = selected_colors[:, 1:, :] + right_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + right_mask[:, :-1, :] = 1.0 + right_desired = torch.zeros_like(selected_colors) + right_desired[:, :-1, :] = representative_delta_x + right_error = ((right_selected[..., None, :] - candidate_colors) - right_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + right_error * right_mask * solver_params.snap_neighbor_weight + if output_height > 1: + up_selected = torch.zeros_like(selected_colors) + up_selected[1:, :, :] = selected_colors[:-1, :, :] + up_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + up_mask[1:, :, :] = 1.0 + up_desired = torch.zeros_like(selected_colors) + up_desired[1:, :, :] = representative_delta_y + up_error = ((candidate_colors - up_selected[..., None, :]) - up_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + up_error * up_mask * solver_params.snap_neighbor_weight + + down_selected = torch.zeros_like(selected_colors) + down_selected[:-1, :, :] = selected_colors[1:, :, :] + down_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + down_mask[:-1, :, :] = 1.0 + down_desired = torch.zeros_like(selected_colors) + down_desired[:-1, :, :] = representative_delta_y + down_error = ((down_selected[..., None, :] - candidate_colors) - down_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + down_error * down_mask * solver_params.snap_neighbor_weight + if output_height > 1 and output_width > 1: + diag_selected = torch.zeros_like(selected_colors) + diag_selected[1:, 1:, :] = selected_colors[:-1, :-1, :] + diag_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + diag_mask[1:, 1:, :] = 1.0 + diag_desired = torch.zeros_like(selected_colors) + diag_desired[1:, 1:, :] = representative_delta_diag + diag_error = ((candidate_colors - diag_selected[..., None, :]) - diag_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + diag_error * diag_mask * solver_params.snap_diagonal_weight + + anti_selected = torch.zeros_like(selected_colors) + anti_selected[1:, :-1, :] = selected_colors[:-1, 1:, :] + anti_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + anti_mask[1:, :-1, :] = 1.0 + anti_desired = torch.zeros_like(selected_colors) + anti_desired[1:, :-1, :] = representative_delta_anti + anti_error = ((candidate_colors - anti_selected[..., None, :]) - anti_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + anti_error * anti_mask * solver_params.snap_diagonal_weight + selected = torch.argmin(energy, dim=-1) + + hardened_selected, use_opaque, use_transparent = _harden_binary_alpha_selection( + torch, + candidate_colors, + selected, + base_energy, + representative[..., 3], + solver_params=solver_params, + ) + final_colors = _select_colors(candidate_colors, hardened_selected) + return _finalize_output_rgba(final_colors, use_opaque, use_transparent) + + +def _select_colors(candidate_colors, selected): + output_height = candidate_colors.shape[0] + output_width = candidate_colors.shape[1] + return candidate_colors.gather( + dim=2, + index=selected[..., None, None].expand(output_height, output_width, 1, candidate_colors.shape[-1]), + ).squeeze(2) + + +def _harden_binary_alpha_selection( + torch, + candidate_colors, + selected, + ranking_energy, + representative_alpha, + solver_params: SolverHyperParams, +): + selected_colors = _select_colors(candidate_colors, selected) + selected_alpha = selected_colors[..., 3] + partial_mask = ( + selected_alpha > solver_params.alpha_transparent_threshold + ) & ( + selected_alpha < solver_params.alpha_opaque_threshold + ) + prefer_foreground = (selected_alpha >= solver_params.alpha_foreground_threshold) | ( + representative_alpha >= solver_params.alpha_representative_foreground_threshold + ) + + opaque_mask = candidate_colors[..., 3] >= solver_params.alpha_opaque_threshold + transparent_mask = candidate_colors[..., 3] <= solver_params.alpha_transparent_threshold + large_penalty = ranking_energy.new_full(ranking_energy.shape, 1e6) + opaque_scores = torch.where(opaque_mask, ranking_energy, ranking_energy + large_penalty) + transparent_scores = torch.where(transparent_mask, ranking_energy, ranking_energy + large_penalty) + + opaque_idx = torch.argmin(opaque_scores, dim=2) + transparent_idx = torch.argmin(transparent_scores, dim=2) + opaque_available = torch.any(opaque_mask, dim=2) + transparent_available = torch.any(transparent_mask, dim=2) + + hardened = selected.clone() + use_opaque = partial_mask & prefer_foreground & opaque_available + use_transparent = partial_mask & (~prefer_foreground) & transparent_available + hardened = torch.where(use_opaque, opaque_idx, hardened) + hardened = torch.where(use_transparent, transparent_idx, hardened) + return hardened, use_opaque, use_transparent + + +def _finalize_output_rgba(best_colors, force_opaque_mask, force_transparent_mask) -> np.ndarray: + rgba = unpremultiply(best_colors.detach().cpu().numpy().clip(0.0, 1.0)) + if force_opaque_mask is not None: + opaque = force_opaque_mask.detach().cpu().numpy() + rgba[..., 3][opaque] = 1.0 + if force_transparent_mask is not None: + transparent = force_transparent_mask.detach().cpu().numpy() + rgba[..., 3][transparent] = 0.0 + rgba[..., :3][transparent] = 0.0 + return rgba + + +def _structure_score( + F, + source_t, + uv_t, + candidate_t, + representative_t, + anchor_t, + solver_params: SolverHyperParams, +): + boundary = _boundary_pattern_loss(F, source_t, uv_t, candidate_t, "x", solver_params) + _boundary_pattern_loss( + F, + source_t, + uv_t, + candidate_t, + "y", + solver_params, + ) + anchor_adj = _adjacency_pattern_loss(candidate_t, anchor_t) + anchor_motif = _motif_pattern_loss(candidate_t, anchor_t) + representative_match = (candidate_t - representative_t).abs().mean() + return ( + boundary * solver_params.structure_boundary_weight + + anchor_adj * solver_params.structure_anchor_adjacency_weight + + anchor_motif * solver_params.structure_anchor_motif_weight + + representative_match * solver_params.structure_representative_weight + ) + + +def _discrete_refine_output( + torch, + F, + source_t, + uv_t, + representative_t, + anchor_t, + source_delta_x_t, + source_delta_y_t, + solver_params: SolverHyperParams, + *, + cell_x: float, + cell_y: float, + iterations: int, +) -> tuple[np.ndarray, list[float]]: + source_hw = source_t[0].permute(1, 2, 0) + height = source_hw.shape[0] + width = source_hw.shape[1] + uv = uv_t[0] + representative = representative_t[0] + anchor = anchor_t[0].permute(1, 2, 0) + output_height = representative.shape[0] + output_width = representative.shape[1] + + fractions = torch.tensor([-0.42, -0.21, 0.0, 0.21, 0.42], device=uv.device, dtype=uv.dtype) + offset_y, offset_x = torch.meshgrid(fractions, fractions, indexing="ij") + offset_x = (offset_x.reshape(-1) * cell_x).to(dtype=uv.dtype) + offset_y = (offset_y.reshape(-1) * cell_y).to(dtype=uv.dtype) + + xs = (uv[..., 0] + 1.0) * 0.5 * max(1.0, float(width - 1)) + ys = (uv[..., 1] + 1.0) * 0.5 * max(1.0, float(height - 1)) + candidate_x = torch.round(xs[..., None] + offset_x).clamp(0, max(0, width - 1)).to(dtype=torch.long) + candidate_y = torch.round(ys[..., None] + offset_y).clamp(0, max(0, height - 1)).to(dtype=torch.long) + candidate_colors = source_hw[candidate_y, candidate_x] + anchor_energy = (candidate_colors - anchor[..., None, :]).abs().mean(dim=-1) + rep_energy = (candidate_colors - representative[..., None, :]).abs().mean(dim=-1) + alpha_energy = (candidate_colors[..., 3] - anchor[..., None, 3]).abs() + distance_energy = ((offset_x / max(cell_x, 1e-4)) ** 2 + (offset_y / max(cell_y, 1e-4)) ** 2).reshape(1, 1, -1) + base_energy = ( + anchor_energy * solver_params.refine_anchor_weight + + rep_energy * solver_params.refine_representative_weight + + alpha_energy * solver_params.refine_alpha_weight + + distance_energy * solver_params.refine_distance_weight + ) + + anchor_delta_x = anchor[:, 1:, :] - anchor[:, :-1, :] if output_width > 1 else None + anchor_delta_y = anchor[1:, :, :] - anchor[:-1, :, :] if output_height > 1 else None + anchor_delta_diag = anchor[1:, 1:, :] - anchor[:-1, :-1, :] if output_height > 1 and output_width > 1 else None + anchor_delta_anti = anchor[1:, :-1, :] - anchor[:-1, 1:, :] if output_height > 1 and output_width > 1 else None + source_delta_x = source_delta_x_t[0] if source_delta_x_t is not None else None + source_delta_y = source_delta_y_t[0] if source_delta_y_t is not None else None + + def desired_delta(anchor_delta, source_delta): + if anchor_delta is None: + return None + if source_delta is None: + return anchor_delta + source_weight = solver_params.refine_source_delta_weight + anchor_weight = 1.0 - source_weight + return anchor_delta * anchor_weight + source_delta * source_weight + + desired_delta_x = desired_delta(anchor_delta_x, source_delta_x) + desired_delta_y = desired_delta(anchor_delta_y, source_delta_y) + desired_delta_diag = anchor_delta_diag + desired_delta_anti = anchor_delta_anti + + selected = torch.argmin(anchor_energy + distance_energy * 0.05, dim=-1) + best_selected = selected.clone() + best_score = float("inf") + loss_history: list[float] = [] + passes = max(0, iterations) + anchor_hw = anchor_t.permute(0, 2, 3, 1) + + for step in range(passes + 1): + selected_colors = _select_colors(candidate_colors, selected) + energy = base_energy.clone() + + if output_width > 1: + left_selected = torch.zeros_like(selected_colors) + left_selected[:, 1:, :] = selected_colors[:, :-1, :] + left_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + left_mask[:, 1:, :] = 1.0 + left_desired = torch.zeros_like(selected_colors) + left_desired[:, 1:, :] = desired_delta_x + left_error = ((candidate_colors - left_selected[..., None, :]) - left_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + left_error * left_mask * solver_params.refine_orthogonal_weight + + right_selected = torch.zeros_like(selected_colors) + right_selected[:, :-1, :] = selected_colors[:, 1:, :] + right_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + right_mask[:, :-1, :] = 1.0 + right_desired = torch.zeros_like(selected_colors) + right_desired[:, :-1, :] = desired_delta_x + right_error = ((right_selected[..., None, :] - candidate_colors) - right_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + right_error * right_mask * solver_params.refine_orthogonal_weight + if output_height > 1: + up_selected = torch.zeros_like(selected_colors) + up_selected[1:, :, :] = selected_colors[:-1, :, :] + up_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + up_mask[1:, :, :] = 1.0 + up_desired = torch.zeros_like(selected_colors) + up_desired[1:, :, :] = desired_delta_y + up_error = ((candidate_colors - up_selected[..., None, :]) - up_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + up_error * up_mask * solver_params.refine_orthogonal_weight + + down_selected = torch.zeros_like(selected_colors) + down_selected[:-1, :, :] = selected_colors[1:, :, :] + down_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + down_mask[:-1, :, :] = 1.0 + down_desired = torch.zeros_like(selected_colors) + down_desired[:-1, :, :] = desired_delta_y + down_error = ((down_selected[..., None, :] - candidate_colors) - down_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + down_error * down_mask * solver_params.refine_orthogonal_weight + if output_height > 1 and output_width > 1: + diag_selected = torch.zeros_like(selected_colors) + diag_selected[1:, 1:, :] = selected_colors[:-1, :-1, :] + diag_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + diag_mask[1:, 1:, :] = 1.0 + diag_desired = torch.zeros_like(selected_colors) + diag_desired[1:, 1:, :] = desired_delta_diag + diag_error = ((candidate_colors - diag_selected[..., None, :]) - diag_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + diag_error * diag_mask * solver_params.refine_diagonal_weight + + anti_selected = torch.zeros_like(selected_colors) + anti_selected[1:, :-1, :] = selected_colors[:-1, 1:, :] + anti_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) + anti_mask[1:, :-1, :] = 1.0 + anti_desired = torch.zeros_like(selected_colors) + anti_desired[1:, :-1, :] = desired_delta_anti + anti_error = ((candidate_colors - anti_selected[..., None, :]) - anti_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + anti_error * anti_mask * solver_params.refine_diagonal_weight + + candidate_t = selected_colors[None, ...] + score = float( + _structure_score( + F, + source_t, + uv_t, + candidate_t, + representative_t, + anchor_hw, + solver_params, + ).detach().cpu().item() + ) + loss_history.append(score) + if score < best_score: + best_score = score + best_selected = selected.clone() + + if step >= passes: + break + next_selected = torch.argmin(energy, dim=-1) + if torch.equal(next_selected, selected): + selected = next_selected + break + selected = next_selected + + hardened_selected, use_opaque, use_transparent = _harden_binary_alpha_selection( + torch, + candidate_colors, + best_selected, + base_energy, + representative[..., 3], + solver_params=solver_params, + ) + best_colors = _select_colors(candidate_colors, hardened_selected) + return _finalize_output_rgba(best_colors, use_opaque, use_transparent), loss_history + + def optimize_uv_field( rgba: np.ndarray, inference: InferenceResult, @@ -49,11 +558,18 @@ def optimize_uv_field( steps: int, seed: int, device: str, + solver_params: SolverHyperParams | None = None, ) -> SolverArtifacts: torch, F = _require_torch() + device = _resolve_device(torch, device) + solver_params = solver_params or SolverHyperParams() torch.manual_seed(seed) + if device == "cuda": + torch.cuda.manual_seed_all(seed) source = premultiply(rgba) height, width = source.shape[:2] + cell_x = width / max(1, inference.target_width) + cell_y = height / max(1, inference.target_height) 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) @@ -66,59 +582,51 @@ def optimize_uv_field( 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())) - + offsets_t = torch.from_numpy( + _make_patch_offsets(height=height, width=width, target_height=inference.target_height, target_width=inference.target_width) + ).to(device=device, dtype=torch.float32) + guide_small = F.interpolate(edge_t, size=(inference.target_height, inference.target_width), mode="bilinear", align_corners=True) + initial_patches = _sample_cell_patches(F, source_t, uv0_t, offsets_t) + initial_representative_t, _ = _representative_colors(initial_patches, solver_params) + initial_representative_t = initial_representative_t.detach() + source_delta_x_t = _source_boundary_deltas(F, source_t, uv0_t, axis="x") + source_delta_y_t = _source_boundary_deltas(F, source_t, uv0_t, axis="y") + snap_rgba = _snap_output_to_source_pixels( + torch, + source_t, + uv0_t, + initial_representative_t, + solver_params, + cell_x=cell_x, + cell_y=cell_y, + ) + snap_t = torch.from_numpy(premultiply(snap_rgba).transpose(2, 0, 1)[None, ...]).to( + device=device, + dtype=torch.float32, + ) 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() + target_rgba, loss_history = _discrete_refine_output( + torch, + F, + source_t, + uv0_t, + initial_representative_t, + snap_t, + source_delta_x_t, + source_delta_y_t, + solver_params, + cell_x=cell_x, + cell_y=cell_y, + iterations=steps, + ) + initial_rgba = snap_rgba + uv_np = uv0_t.detach().cpu().numpy()[0] guidance = guide_small[0, 0].detach().cpu().numpy() return SolverArtifacts( - target_rgba=unpremultiply(target_rgba), - uv_field=uv.detach().cpu().numpy()[0], + target_rgba=target_rgba, + uv_field=uv_np, guidance_strength=guidance.astype(np.float32), - initial_rgba=unpremultiply(initial_rgba_np), + initial_rgba=initial_rgba, loss_history=loss_history, ) diff --git a/src/repixelizer/corpus.py b/src/repixelizer/corpus.py new file mode 100644 index 0000000..591c83e --- /dev/null +++ b/src/repixelizer/corpus.py @@ -0,0 +1,500 @@ +from __future__ import annotations + +from collections import Counter, deque +import json +from pathlib import Path +import re +import shutil +from typing import Any + +import numpy as np + +from .io import load_rgba, save_rgba + +FRAME_SHEETS: dict[str, int] = { + "mage-1-85x94.png": 0, + "mage-2-122x110.png": 0, + "mage-3-87x110.png": 0, + "minion-45x66.png": 0, + "andromalius-57x88.png": 0, + "disciple-45x51.png": 0, + "gnu-120x100.png": 0, + "shadow-80x70.png": -1, +} + +MULTI_CHARACTER_SHEETS: tuple[str, ...] = ( + "9RPGenemies.PNG", + "DAGRONS5.png", + "more rpg enemies.PNG", + "threeformsPJ2.png", +) + +MIN_COMPONENT_PIXELS = 100 + +SOURCE_ATTRIBUTION: dict[str, dict[str, Any]] = { + "9RPGenemies.PNG": { + "source_title": "More RPG enemies!", + "source_authors": ["Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/more-rpg-enemies", + "licenses": ["CC-BY 3.0", "CC-BY-SA 3.0", "OGA-BY 3.0"], + "attribution_text": "Stephen Challener (Redshrike), hosted by OpenGameArt.org", + "attribution_instructions": "Credit Stephen Challener (Redshrike) and include a link back to OpenGameArt.org.", + "attribution_confidence": "exact", + }, + "DAGRONS5.png": { + "source_title": "RPG Enemies: 11 Dragons", + "source_authors": ["Stephen Challener (Redshrike)", "MrBeast", "Surt", "Blarumyrran", "Sharm", "Zabin"], + "source_url": "https://opengameart.org/content/rpg-enemies-11-dragons", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Credit to: Stephen \"Redshrike\" Challener, MrBeast, Surt, Blarumyrran, Sharm, Zabin", + "attribution_instructions": "Credit the listed contributors and include a link to the source page.", + "attribution_confidence": "exact", + }, + "more rpg enemies.PNG": { + "source_title": "6 More RPG Enemies", + "source_authors": ["Stephen Challener (Redshrike)", "Blarumyrran", "LordNeo"], + "source_url": "https://opengameart.org/content/6-more-rpg-enemies", + "licenses": ["CC-BY 3.0", "OGA-BY 3.0"], + "attribution_text": "Stephen Challener (Redshrike), Blarumyrran and LordNeo, hosted by OpenGameArt.org", + "attribution_instructions": "Credit Stephen Challener (Redshrike), Blarumyrran, and LordNeo, and include a link back to OpenGameArt.org.", + "attribution_confidence": "exact", + }, + "threeformsPJ2.png": { + "source_title": "3-form RPG boss: \"Harlequin Epicycle.\"", + "source_authors": ["Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/3-form-rpg-boss-harlequin-epicycle", + "licenses": ["CC-BY 3.0", "GPL 3.0", "GPL 2.0", "OGA-BY 3.0"], + "attribution_text": "Stephen Challener (Redshrike), hosted by OpenGameArt.org", + "attribution_instructions": "Credit Stephen Challener (Redshrike) and include a link back to OpenGameArt.org.", + "attribution_confidence": "exact", + }, + "mage-1-85x94.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Balmer; attribute Stephen Challener (Redshrike) for the original sprites and link back to OpenGameArt.org", + "attribution_instructions": "Best-effort recovery: credit Balmer and Stephen Challener (Redshrike), and include a link to the OpenGameArt source page.", + "attribution_confidence": "inferred", + }, + "mage-2-122x110.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Balmer; attribute Stephen Challener (Redshrike) for the original sprites and link back to OpenGameArt.org", + "attribution_instructions": "Best-effort recovery: credit Balmer and Stephen Challener (Redshrike), and include a link to the OpenGameArt source page.", + "attribution_confidence": "inferred", + }, + "mage-3-87x110.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Balmer; attribute Stephen Challener (Redshrike) for the original sprites and link back to OpenGameArt.org", + "attribution_instructions": "Best-effort recovery: credit Balmer and Stephen Challener (Redshrike), and include a link to the OpenGameArt source page.", + "attribution_confidence": "inferred", + }, + "minion-45x66.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Balmer; attribute Stephen Challener (Redshrike) for the original sprites and link back to OpenGameArt.org", + "attribution_instructions": "Best-effort recovery: credit Balmer and Stephen Challener (Redshrike), and include a link to the OpenGameArt source page.", + "attribution_confidence": "inferred", + }, + "andromalius-57x88.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Balmer; attribute Stephen Challener (Redshrike) for the original sprites and link back to OpenGameArt.org", + "attribution_instructions": "Best-effort recovery: credit Balmer and Stephen Challener (Redshrike), and include a link to the OpenGameArt source page.", + "attribution_confidence": "inferred", + }, + "disciple-45x51.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY 3.0"], + "attribution_text": "Balmer; attribute Stephen Challener (Redshrike) for the original sprites and link back to OpenGameArt.org", + "attribution_instructions": "Best-effort recovery: credit Balmer and Stephen Challener (Redshrike), and include a link to the OpenGameArt source page.", + "attribution_confidence": "inferred", + }, + "gnu-120x100.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)", "Ben \"Cookiez\" Potter"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY-SA 3.0", "GPL 3.0"], + "attribution_text": "Balmer; also credit Stephen Challener (Redshrike) and Ben \"Cookiez\" Potter, with links back to OpenGameArt.org", + "attribution_instructions": "Recovered from the Ars Notoria sheet plus the original Gnu Mage provenance. Use CC-BY-SA 3.0 or GPL 3.0 terms.", + "attribution_confidence": "inferred", + "original_source_title": "LPC in battle RPG sprites", + "original_source_url": "https://opengameart.org/content/lpc-in-battle-rpg-sprites", + "license_note": "Redshrike later clarified in the Ars Notoria page comments that the Gnu Man should inherit CC-BY-SA terms from its original source.", + }, + "shadow-80x70.png": { + "source_title": "Bosses and monsters spritesheets (Ars Notoria)", + "source_authors": ["Balmer", "Stephen Challener (Redshrike)", "Clint Bellanger"], + "source_url": "https://opengameart.org/content/bosses-and-monsters-spritesheets-ars-notoria", + "licenses": ["CC-BY-SA 3.0"], + "attribution_text": "Balmer; also credit Stephen Challener (Redshrike) and Clint Bellanger, with links back to OpenGameArt.org", + "attribution_instructions": "Recovered from the Ars Notoria sheet plus the original shadow source provenance. Use CC-BY-SA 3.0 terms.", + "attribution_confidence": "inferred", + "original_source_title": "First Person Dungeon Crawl Enemies Remixed", + "original_source_url": "https://opengameart.org/content/first-person-dungeon-crawl-enemies-remixed", + "license_note": "Redshrike later clarified in the Ars Notoria page comments that the Shadow Soul should inherit CC-BY-SA terms from its original source.", + }, +} + + +def prepare_corpus( + corpus_dir: str | Path, + *, + archive_dir: str | Path | None = None, +) -> dict[str, Any]: + corpus_path = Path(corpus_dir) + originals_dir = corpus_path / "originals" + source_sheets_dir = Path(archive_dir) if archive_dir is not None else corpus_path / "source-sheets" + originals_dir.mkdir(parents=True, exist_ok=True) + source_sheets_dir.mkdir(parents=True, exist_ok=True) + + outputs: list[dict[str, Any]] = [] + for source_name, frame_index in FRAME_SHEETS.items(): + source_path = _resolve_source_path(source_name, originals_dir, source_sheets_dir) + if source_path is None: + continue + rgba = load_rgba(source_path) + frame_width, frame_height = _parse_frame_size(source_path.name) + frame = _extract_frame(rgba, frame_width=frame_width, frame_height=frame_height, frame_index=frame_index) + output_name = f"{_strip_size_suffix(source_path.stem)}.png" + output_path = originals_dir / output_name + save_rgba(output_path, frame) + sidecar_path = output_path.with_suffix(".json") + sidecar_payload = _merge_sidecar( + sidecar_path, + { + "source_sheet": source_path.name, + "kind": "animation-frame", + "frame_index": _normalize_frame_index(frame_index, rgba.shape[1] // frame_width, rgba.shape[0] // frame_height), + "frame_width": frame_width, + "frame_height": frame_height, + **_source_attribution(source_path.name), + }, + ) + sidecar_path.write_text( + json.dumps(sidecar_payload, indent=2), + encoding="utf-8", + ) + _archive_source_sheet(source_path, originals_dir, source_sheets_dir) + outputs.append({"source": source_path.name, "output": output_name, "kind": "animation-frame"}) + + for source_name in MULTI_CHARACTER_SHEETS: + source_path = _resolve_source_path(source_name, originals_dir, source_sheets_dir) + if source_path is None: + continue + rgba = load_rgba(source_path) + sprites = _split_multi_character_sheet(rgba) + stem = _sanitize_name(source_path.stem) + for index, sprite in enumerate(sprites, start=1): + output_name = f"{stem}-{index:02d}.png" + output_path = originals_dir / output_name + save_rgba(output_path, sprite["rgba"]) + sidecar_path = output_path.with_suffix(".json") + sidecar_payload = _merge_sidecar( + sidecar_path, + { + "source_sheet": source_path.name, + "kind": "multi-character-crop", + "component_index": index, + "bbox": { + "left": sprite["bbox"][0], + "top": sprite["bbox"][1], + "right": sprite["bbox"][2], + "bottom": sprite["bbox"][3], + }, + **_source_attribution(source_path.name), + }, + ) + sidecar_path.write_text( + json.dumps(sidecar_payload, indent=2), + encoding="utf-8", + ) + outputs.append({"source": source_path.name, "output": output_name, "kind": "multi-character-crop"}) + _archive_source_sheet(source_path, originals_dir, source_sheets_dir) + + summary = { + "corpus_dir": str(corpus_path), + "archive_dir": str(source_sheets_dir), + "outputs": outputs, + "output_count": len(outputs), + } + (corpus_path / "prepare-corpus.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + write_attribution_markdown(corpus_path) + return summary + + +def write_attribution_markdown(corpus_dir: str | Path) -> Path: + corpus_path = Path(corpus_dir) + originals_dir = corpus_path / "originals" + entries = _load_processed_entries(originals_dir) + attribution_path = corpus_path / "ATTRIBUTION.md" + attribution_path.write_text(_render_attribution_markdown(entries), encoding="utf-8") + return attribution_path + + +def _resolve_source_path(source_name: str, originals_dir: Path, source_sheets_dir: Path) -> Path | None: + original_path = originals_dir / source_name + if original_path.exists(): + return original_path + archived_path = source_sheets_dir / source_name + if archived_path.exists(): + return archived_path + return None + + +def _source_attribution(source_name: str) -> dict[str, Any]: + return dict(SOURCE_ATTRIBUTION.get(source_name, {})) + + +def _merge_sidecar(sidecar_path: Path, generated: dict[str, Any]) -> dict[str, Any]: + existing: dict[str, Any] = {} + if sidecar_path.exists(): + existing = json.loads(sidecar_path.read_text(encoding="utf-8")) + merged = dict(existing) + merged.update(generated) + return merged + + +def _load_processed_entries(originals_dir: Path) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + for sidecar_path in sorted(originals_dir.glob("*.json")): + payload = json.loads(sidecar_path.read_text(encoding="utf-8")) + if "source_sheet" not in payload: + continue + payload["derived_file"] = f"{sidecar_path.stem}.png" + entries.append(payload) + return entries + + +def _render_attribution_markdown(entries: list[dict[str, Any]]) -> str: + grouped: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + for entry in entries: + key = ( + entry.get("source_title", ""), + entry.get("source_url", ""), + tuple(entry.get("source_authors", [])), + tuple(entry.get("licenses", [])), + entry.get("attribution_confidence", ""), + entry.get("attribution_text", ""), + entry.get("attribution_instructions", ""), + entry.get("original_source_title", ""), + entry.get("original_source_url", ""), + entry.get("license_note", ""), + ) + grouped.setdefault(key, []).append(entry) + + lines = [ + "# Corpus Attribution", + "", + "Generated by `repixelize prepare-corpus`.", + "", + "This file summarizes attribution metadata for processed sprites in `examples/corpus/originals`.", + "Entries marked `inferred` are best-effort recoveries and should be reviewed before distribution.", + "", + ] + sort_keys = sorted( + grouped, + key=lambda key: ( + str(key[0]), + str(key[1]), + min(str(entry["source_sheet"]) for entry in grouped[key]), + ), + ) + for key in sort_keys: + source_entries = sorted(grouped[key], key=lambda item: (str(item["source_sheet"]), str(item["derived_file"]))) + sample = source_entries[0] + title = sample.get("source_title", str(sample["source_sheet"])) + lines.append(f"## {title}") + lines.append("") + source_sheets = sorted({str(entry["source_sheet"]) for entry in source_entries}) + if len(source_sheets) == 1: + lines.append(f"- Source sheet: `{source_sheets[0]}`") + else: + lines.append(f"- Source sheets: {', '.join(f'`{sheet}`' for sheet in source_sheets)}") + if sample.get("source_url"): + lines.append(f"- Source URL: {sample['source_url']}") + if sample.get("source_authors"): + lines.append(f"- Authors: {', '.join(sample['source_authors'])}") + if sample.get("licenses"): + lines.append(f"- Licenses: {', '.join(sample['licenses'])}") + if sample.get("attribution_confidence"): + lines.append(f"- Attribution confidence: {sample['attribution_confidence']}") + if sample.get("attribution_text"): + lines.append(f"- Attribution text: {sample['attribution_text']}") + if sample.get("attribution_instructions"): + lines.append(f"- Notes: {sample['attribution_instructions']}") + if sample.get("original_source_title"): + original = sample["original_source_title"] + if sample.get("original_source_url"): + original = f"{original} ({sample['original_source_url']})" + lines.append(f"- Original upstream source: {original}") + if sample.get("license_note"): + lines.append(f"- License note: {sample['license_note']}") + lines.append("- Derived files:") + for entry in source_entries: + lines.append(f" - `{entry['derived_file']}`") + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + +def _archive_source_sheet(source_path: Path, originals_dir: Path, source_sheets_dir: Path) -> None: + if source_path.parent != originals_dir: + return + destination = source_sheets_dir / source_path.name + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists(): + destination.unlink() + shutil.move(str(source_path), str(destination)) + + +def _parse_frame_size(name: str) -> tuple[int, int]: + match = re.search(r"-(\d+)x(\d+)\.[^.]+$", name, flags=re.IGNORECASE) + if match is None: + raise RuntimeError(f"Could not parse frame size from {name}.") + return int(match.group(1)), int(match.group(2)) + + +def _strip_size_suffix(stem: str) -> str: + return re.sub(r"-\d+x\d+$", "", stem) + + +def _sanitize_name(value: str) -> str: + text = re.sub(r"[^a-zA-Z0-9]+", "-", value.lower()).strip("-") + return re.sub(r"-{2,}", "-", text) + + +def _normalize_frame_index(frame_index: int, columns: int, rows: int) -> int: + total = columns * rows + return total + frame_index if frame_index < 0 else frame_index + + +def _extract_frame(rgba: np.ndarray, *, frame_width: int, frame_height: int, frame_index: int) -> np.ndarray: + height, width = rgba.shape[:2] + if width % frame_width != 0 or height % frame_height != 0: + raise RuntimeError("Sheet dimensions do not divide evenly into animation frames.") + columns = width // frame_width + rows = height // frame_height + index = _normalize_frame_index(frame_index, columns, rows) + if index < 0 or index >= columns * rows: + raise RuntimeError("Frame index is out of range for the sprite sheet.") + row = index // columns + column = index % columns + cell = rgba[row * frame_height : (row + 1) * frame_height, column * frame_width : (column + 1) * frame_width].copy() + return _remove_edge_connected_background(cell) + + +def _split_multi_character_sheet(rgba: np.ndarray) -> list[dict[str, Any]]: + transparent = _remove_edge_connected_background(rgba) + alpha = transparent[..., 3] > 0.0 + components = _connected_components(alpha) + sprites: list[dict[str, Any]] = [] + for component in components: + if component["pixels"] < MIN_COMPONENT_PIXELS: + continue + left, top, right, bottom = component["bbox"] + crop = _crop_with_padding(transparent, left, top, right, bottom, padding=1) + sprites.append({"rgba": crop, "bbox": [left, top, right, bottom]}) + sprites.sort(key=lambda item: (item["bbox"][1], item["bbox"][0])) + return sprites + + +def _crop_with_padding(rgba: np.ndarray, left: int, top: int, right: int, bottom: int, *, padding: int) -> np.ndarray: + height, width = rgba.shape[:2] + x0 = max(0, left - padding) + y0 = max(0, top - padding) + x1 = min(width, right + padding) + y1 = min(height, bottom + padding) + return rgba[y0:y1, x0:x1].copy() + + +def _remove_edge_connected_background(rgba: np.ndarray) -> np.ndarray: + image = rgba.copy() + bg_rgb = _background_color(image) + bg_match = np.all(np.abs(image[..., :3] - bg_rgb[None, None, :]) < 0.5 / 255.0, axis=-1) + edge_connected = _edge_connected_mask(bg_match) + image[edge_connected, 3] = 0.0 + image[edge_connected, :3] = 0.0 + return image + + +def _background_color(rgba: np.ndarray) -> np.ndarray: + height, width = rgba.shape[:2] + corners = [ + tuple(np.rint(rgba[0, 0, :3] * 255.0).astype(int)), + tuple(np.rint(rgba[0, width - 1, :3] * 255.0).astype(int)), + tuple(np.rint(rgba[height - 1, 0, :3] * 255.0).astype(int)), + tuple(np.rint(rgba[height - 1, width - 1, :3] * 255.0).astype(int)), + ] + bg = Counter(corners).most_common(1)[0][0] + return np.asarray(bg, dtype=np.float32) / 255.0 + + +def _edge_connected_mask(mask: np.ndarray) -> np.ndarray: + height, width = mask.shape + visited = np.zeros((height, width), dtype=bool) + queue: deque[tuple[int, int]] = deque() + for x in range(width): + if mask[0, x]: + queue.append((0, x)) + visited[0, x] = True + if mask[height - 1, x] and not visited[height - 1, x]: + queue.append((height - 1, x)) + visited[height - 1, x] = True + for y in range(height): + if mask[y, 0] and not visited[y, 0]: + queue.append((y, 0)) + visited[y, 0] = True + if mask[y, width - 1] and not visited[y, width - 1]: + queue.append((y, width - 1)) + visited[y, width - 1] = True + while queue: + y, x = queue.popleft() + for ny in range(max(0, y - 1), min(height, y + 2)): + for nx in range(max(0, x - 1), min(width, x + 2)): + if mask[ny, nx] and not visited[ny, nx]: + visited[ny, nx] = True + queue.append((ny, nx)) + return visited + + +def _connected_components(mask: np.ndarray) -> list[dict[str, Any]]: + height, width = mask.shape + visited = np.zeros((height, width), dtype=bool) + components: list[dict[str, Any]] = [] + for y in range(height): + for x in range(width): + if not mask[y, x] or visited[y, x]: + continue + queue: deque[tuple[int, int]] = deque([(y, x)]) + visited[y, x] = True + pixels = 0 + left = x + right = x + 1 + top = y + bottom = y + 1 + while queue: + cy, cx = queue.popleft() + pixels += 1 + left = min(left, cx) + right = max(right, cx + 1) + top = min(top, cy) + bottom = max(bottom, cy + 1) + for ny in range(max(0, cy - 1), min(height, cy + 2)): + for nx in range(max(0, cx - 1), min(width, cx + 2)): + if mask[ny, nx] and not visited[ny, nx]: + visited[ny, nx] = True + queue.append((ny, nx)) + components.append({"pixels": pixels, "bbox": [left, top, right, bottom]}) + return components diff --git a/src/repixelizer/diagnostics.py b/src/repixelizer/diagnostics.py index a721d82..aa5cf9d 100644 --- a/src/repixelizer/diagnostics.py +++ b/src/repixelizer/diagnostics.py @@ -9,7 +9,7 @@ import numpy as np from PIL import Image, ImageDraw from .io import nearest_resize, save_rgba -from .metrics import coherence_breakdown, reconstruction_error +from .metrics import coherence_breakdown, foreground_reconstruction_error, reconstruction_error from .types import InferenceResult, RunResult @@ -55,7 +55,7 @@ def write_comparison(path: str | Path, source_rgba: np.ndarray, output_rgba: np. 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"] + labels = ["Source Facsimile", "Output @ Source Size", "Output Grid"] canvas = Image.new("RGBA", (width * len(tiles), height + 24), (24, 24, 24, 255)) for index, tile in enumerate(tiles): panel = _make_checker((width, height)) @@ -113,5 +113,17 @@ def summarize_run(result: RunResult) -> dict[str, Any]: "confidence": result.inference.confidence, "coherence": coherence_breakdown(result.output_rgba), "source_preview_reconstruction_error": reconstruction_error(source_preview, result.source_rgba), + "source_preview_foreground_error": foreground_reconstruction_error(source_preview, result.source_rgba), + "output_colors_from_source_ratio": _source_color_ratio(result.source_rgba, result.output_rgba), "loss_history": result.solver.loss_history, } + + +def _source_color_ratio(source_rgba: np.ndarray, output_rgba: np.ndarray) -> float: + source8 = _to_uint8(source_rgba).reshape(-1, 4) + output8 = _to_uint8(output_rgba).reshape(-1, 4) + source_colors = {tuple(px.tolist()) for px in source8} + output_colors = {tuple(px.tolist()) for px in output8} + if not output_colors: + return 1.0 + return float(sum(color in source_colors for color in output_colors) / len(output_colors)) diff --git a/src/repixelizer/discrete.py b/src/repixelizer/discrete.py index 439f773..02a7f97 100644 --- a/src/repixelizer/discrete.py +++ b/src/repixelizer/discrete.py @@ -24,10 +24,12 @@ def _local_energy(patch: np.ndarray) -> float: return isolation + alpha_kink * 0.25 -def cleanup_pixels(rgba: np.ndarray, source_guidance: np.ndarray, iterations: int = 2) -> CleanupArtifacts: +def cleanup_pixels(rgba: np.ndarray, source_guidance: np.ndarray, iterations: int = 0) -> CleanupArtifacts: result = rgba.copy() height, width = result.shape[:2] heatmap = np.zeros((height, width), dtype=np.float32) + if iterations <= 0: + return CleanupArtifacts(cleaned_rgba=result, isolated_heatmap=heatmap) for _ in range(iterations): updated = result.copy() for y in range(1, height - 1): @@ -42,8 +44,6 @@ def cleanup_pixels(rgba: np.ndarray, source_guidance: np.ndarray, iterations: in 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 diff --git a/src/repixelizer/inference.py b/src/repixelizer/inference.py index 8c206b1..d095a71 100644 --- a/src/repixelizer/inference.py +++ b/src/repixelizer/inference.py @@ -1,25 +1,31 @@ from __future__ import annotations import numpy as np -from PIL import Image -from .metrics import coherence_breakdown +from .io import premultiply 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 _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 lattice inference. Install project dependencies first." + ) from exc + return torch, F -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 _resolve_device(torch, requested: str) -> str: + if requested == "auto": + return "cuda" if torch.cuda.is_available() else "cpu" + if requested == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested for lattice inference, but this PyTorch build does not have a usable CUDA device. " + "Install a CUDA-enabled PyTorch build or use --device cpu." + ) + return requested def _candidate_dims(width: int, height: int, target_size: int | None) -> list[tuple[int, int]]: @@ -56,6 +62,88 @@ def _estimate_cell_size(profile: np.ndarray) -> float: return float(best_lag) +def _estimate_spacing_cell_size(deltas: np.ndarray, axis_length: int) -> tuple[float | None, float]: + intervals, weights = _collect_change_intervals(deltas) + if intervals.size < 6: + return None, 0.0 + max_candidate = min(64, max(4, axis_length // 2)) + candidates = np.arange(2, max_candidate + 1, dtype=np.float32) + scores = np.asarray([_spacing_score(intervals, weights, cell) for cell in candidates], dtype=np.float32) + best_index = int(np.argmax(scores)) + best_cell = float(candidates[best_index]) + best_score = float(scores[best_index]) + multiples = np.maximum(1.0, np.rint(intervals / best_cell)) + residual = np.abs(intervals / best_cell - multiples) + keep = residual < 0.2 + if np.any(keep): + refined_weights = weights[keep] / np.sqrt(multiples[keep]) + refined = float(np.average(intervals[keep] / multiples[keep], weights=refined_weights)) + else: + refined = best_cell + coverage = min(1.0, intervals.size / 24.0) + confidence = np.clip((best_score - 0.35) / 0.45, 0.0, 1.0) * coverage + return refined, float(confidence) + + +def _collect_change_intervals(deltas: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + positive = deltas[deltas > 1e-4] + if positive.size < 8: + return np.zeros(0, dtype=np.float32), np.zeros(0, dtype=np.float32) + threshold = float(np.clip(np.quantile(positive, 0.72) * 0.55, 0.018, 0.16)) + intervals: list[np.ndarray] = [] + weights: list[np.ndarray] = [] + for line in deltas: + centers, strengths = _boundary_runs(line, threshold) + if centers.size < 2: + continue + line_intervals = np.diff(centers) + if line_intervals.size == 0: + continue + keep = line_intervals >= 2.0 + if not np.any(keep): + continue + line_weights = ((strengths[:-1] + strengths[1:]) * 0.5)[keep] + intervals.append(line_intervals[keep].astype(np.float32)) + weights.append(line_weights.astype(np.float32)) + if not intervals: + return np.zeros(0, dtype=np.float32), np.zeros(0, dtype=np.float32) + return np.concatenate(intervals), np.concatenate(weights) + + +def _boundary_runs(line: np.ndarray, threshold: float) -> tuple[np.ndarray, np.ndarray]: + indices = np.flatnonzero(line >= threshold) + if indices.size == 0: + return np.zeros(0, dtype=np.float32), np.zeros(0, dtype=np.float32) + splits = np.where(np.diff(indices) > 1)[0] + starts = np.concatenate([indices[:1], indices[splits + 1]]) + ends = np.concatenate([indices[splits], indices[-1:]]) + centers = (starts.astype(np.float32) + ends.astype(np.float32) + 1.0) * 0.5 + strengths = np.asarray([float(np.max(line[start : end + 1])) for start, end in zip(starts, ends)], dtype=np.float32) + return centers, strengths + + +def _spacing_score(intervals: np.ndarray, weights: np.ndarray, cell_size: float) -> float: + multiples = np.maximum(1.0, np.rint(intervals / cell_size)) + residual = np.abs(intervals / cell_size - multiples) + closeness = np.exp(-((residual / 0.18) ** 2)) + interval_weights = weights / np.sqrt(multiples) + return float(np.sum(interval_weights * closeness) / (np.sum(interval_weights) + 1e-6)) + + +def _estimate_lattice_spacing(rgba: np.ndarray) -> tuple[tuple[float | None, float], tuple[float | None, float]]: + premult_rgba = premultiply(rgba) + delta_x = np.linalg.norm(premult_rgba[:, 1:] - premult_rgba[:, :-1], axis=-1) + delta_y = np.linalg.norm(premult_rgba[1:, :] - premult_rgba[:-1, :], axis=-1) + return _estimate_spacing_cell_size(delta_x, premult_rgba.shape[1]), _estimate_spacing_cell_size(delta_y, premult_rgba.shape[0]) + + +def _blend_prior(spacing: float | None, spacing_confidence: float, autocorr: float) -> float: + if spacing is None: + return float(autocorr) + weight = float(np.clip(spacing_confidence, 0.0, 0.85)) + return float(spacing * weight + autocorr * (1.0 - weight)) + + 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 @@ -65,38 +153,193 @@ def _estimate_lattice_prior(rgba: np.ndarray) -> tuple[float, float]: 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) + spacing_x, spacing_y = _estimate_lattice_spacing(rgba) + autocorr_x = _estimate_cell_size(profile_x) + autocorr_y = _estimate_cell_size(profile_y) + return _blend_prior(spacing_x[0], spacing_x[1], autocorr_x), _blend_prior(spacing_y[0], spacing_y[1], autocorr_y) -def infer_lattice(rgba: np.ndarray, target_size: int | None = None) -> InferenceResult: +def _build_phase_grid( + torch, + *, + width: int, + height: int, + target_width: int, + target_height: int, + phase_xs, + phase_ys, + device: str, +): + batch = phase_xs.shape[0] + cell_x = width / target_width + cell_y = height / target_height + base_x = (torch.arange(target_width, device=device, dtype=torch.float32) + 0.5) * cell_x - 0.5 + base_y = (torch.arange(target_height, device=device, dtype=torch.float32) + 0.5) * cell_y - 0.5 + xs = base_x[None, None, :].expand(batch, target_height, target_width) + phase_xs[:, None, None] * cell_x + ys = base_y[None, :, None].expand(batch, target_height, target_width) + phase_ys[:, None, None] * cell_y + xs = xs.clamp(0.0, max(0.0, width - 1)) + ys = ys.clamp(0.0, max(0.0, height - 1)) + grid_x = (xs / max(1.0, width - 1)) * 2.0 - 1.0 + grid_y = (ys / max(1.0, height - 1)) * 2.0 - 1.0 + return torch.stack([grid_x, grid_y], dim=-1) + + +def _unpremultiply_batch(torch, premult_batch): + alpha = premult_batch[:, 3:4].clamp_min(1e-6) + rgb = torch.where(alpha > 1e-6, premult_batch[:, :3] / alpha, torch.zeros_like(premult_batch[:, :3])) + return torch.cat([rgb, premult_batch[:, 3:4]], dim=1).clamp(0.0, 1.0) + + +def _coherence_breakdown_torch(torch, F, rgba_batch): + batch = rgba_batch.shape[0] + color_rgba = rgba_batch.permute(0, 2, 3, 1) + alpha = rgba_batch[:, 3] + + cluster_sims = [] + for dy, dx in ((0, 1), (1, 0)): + shifted = torch.roll(color_rgba, shifts=(dy, dx), dims=(1, 2)) + dist = torch.linalg.vector_norm(color_rgba - shifted, dim=-1) + cluster_sims.append(torch.exp(-dist * 8.0)) + cluster_continuity = torch.stack(cluster_sims, dim=0).mean(dim=(0, 2, 3)) + + alpha_crispness = 1.0 - (4.0 * alpha * (1.0 - alpha)).mean(dim=(1, 2)) + + neighbors = [] + for dy, dx in ((0, 1), (1, 0), (0, -1), (-1, 0)): + shifted = torch.roll(color_rgba, shifts=(dy, dx), dims=(1, 2)) + dist = torch.linalg.vector_norm(color_rgba - shifted, dim=-1) + neighbors.append(dist > 0.15) + isolated_penalty = (torch.stack(neighbors, dim=0).sum(dim=0) >= 3).float().mean(dim=(1, 2)) + + mask = alpha >= 0.5 + has_mask = mask.flatten(1).any(dim=1) + if mask.shape[1] > 1 and mask.shape[2] > 1: + a = mask[:, :-1, :-1].to(dtype=torch.int32) + b = mask[:, :-1, 1:].to(dtype=torch.int32) + c = mask[:, 1:, :-1].to(dtype=torch.int32) + d = mask[:, 1:, 1:].to(dtype=torch.int32) + block_sum = a + b + c + d + corner_blocks = ((block_sum == 1) | (block_sum == 3)).sum(dim=(1, 2)).to(dtype=torch.float32) + else: + corner_blocks = torch.zeros(batch, device=rgba_batch.device, dtype=torch.float32) + edge_pixels = (mask ^ torch.roll(mask, shifts=1, dims=1)) | (mask ^ torch.roll(mask, shifts=1, dims=2)) + edge_count = torch.count_nonzero(edge_pixels, dim=(1, 2)).to(dtype=torch.float32).clamp_min(1.0) + outline_straightness = torch.where( + has_mask, + 1.0 - corner_blocks / edge_count, + torch.ones_like(edge_count), + ) + + color_rgb = rgba_batch[:, :3] + padded = F.pad(color_rgb, (1, 1, 1, 1), mode="replicate") + patches = padded.unfold(2, 3, 1).unfold(3, 3, 1).contiguous() + patches = patches.view(batch, 3, rgba_batch.shape[2], rgba_batch.shape[3], 9) + median = patches.median(dim=-1).values + color_chatter = torch.linalg.vector_norm(color_rgb - median, dim=1).mean(dim=(1, 2)) + + coherence_score = ( + 0.25 * cluster_continuity + + 0.20 * alpha_crispness + + 0.20 * outline_straightness + + 0.20 * (1.0 - isolated_penalty) + + 0.15 * (1.0 - torch.minimum(torch.ones_like(color_chatter), color_chatter)) + ) + return { + "cluster_continuity": cluster_continuity, + "alpha_crispness": alpha_crispness, + "outline_straightness": outline_straightness, + "isolated_penalty": isolated_penalty, + "color_chatter": color_chatter, + "coherence_score": coherence_score, + } + + +def _score_phase_group( + rgba: np.ndarray, + *, + target_width: int, + target_height: int, + prior_cell_x: float, + prior_cell_y: float, + phase_values: np.ndarray, + device: str, +): + torch, F = _require_torch() + height, width = rgba.shape[:2] + premult_rgba = premultiply(rgba) + source_t = torch.from_numpy(premult_rgba.transpose(2, 0, 1)[None, ...]).to(device=device, dtype=torch.float32) + phase_xs, phase_ys = np.meshgrid(phase_values, phase_values, indexing="xy") + phase_x_t = torch.from_numpy(phase_xs.reshape(-1).astype(np.float32)).to(device=device) + phase_y_t = torch.from_numpy(phase_ys.reshape(-1).astype(np.float32)).to(device=device) + batch = phase_x_t.shape[0] + grid = _build_phase_grid( + torch, + width=width, + height=height, + target_width=target_width, + target_height=target_height, + phase_xs=phase_x_t, + phase_ys=phase_y_t, + device=device, + ) + source_batch = source_t.expand(batch, -1, -1, -1) + sampled = F.grid_sample(source_batch, grid, align_corners=True, mode="bilinear", padding_mode="border") + rgba_sample = _unpremultiply_batch(torch, sampled) + breakdown = _coherence_breakdown_torch(torch, F, rgba_sample) + cell_x = width / target_width + cell_y = height / target_height + 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) + score = breakdown["coherence_score"] * 0.55 + size_prior * 0.45 + score_np = score.detach().cpu().numpy() + + candidates: list[InferenceCandidate] = [] + for index, phase_x in enumerate(phase_x_t.detach().cpu().numpy()): + phase_y = float(phase_y_t[index].detach().cpu().item()) + candidate_breakdown = { + "cluster_continuity": float(breakdown["cluster_continuity"][index].detach().cpu().item()), + "alpha_crispness": float(breakdown["alpha_crispness"][index].detach().cpu().item()), + "outline_straightness": float(breakdown["outline_straightness"][index].detach().cpu().item()), + "isolated_penalty": float(breakdown["isolated_penalty"][index].detach().cpu().item()), + "color_chatter": float(breakdown["color_chatter"][index].detach().cpu().item()), + "coherence_score": float(breakdown["coherence_score"][index].detach().cpu().item()), + "size_prior": size_prior, + } + candidates.append( + InferenceCandidate( + target_width=target_width, + target_height=target_height, + phase_x=float(phase_x), + phase_y=phase_y, + score=float(score_np[index]), + breakdown=candidate_breakdown, + ) + ) + return candidates + + +def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str = "auto") -> InferenceResult: + torch, _ = _require_torch() + resolved_device = _resolve_device(torch, device) height, width = rgba.shape[:2] - image = _to_pil(rgba) prior_cell_x, prior_cell_y = _estimate_lattice_prior(rgba) + phase_values = np.linspace(-0.4, 0.4, num=5, dtype=np.float32) + 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.extend( + _score_phase_group( + rgba, + target_width=target_width, + target_height=target_height, + prior_cell_x=prior_cell_x, + prior_cell_y=prior_cell_y, + phase_values=phase_values, + device=resolved_device, + ) + ) + candidates.sort(key=lambda item: item.score, reverse=True) best = candidates[0] second = candidates[1] if len(candidates) > 1 else best diff --git a/src/repixelizer/metrics.py b/src/repixelizer/metrics.py index 2c04561..a212a79 100644 --- a/src/repixelizer/metrics.py +++ b/src/repixelizer/metrics.py @@ -2,6 +2,8 @@ from __future__ import annotations import numpy as np +from .io import premultiply + def luminance(rgba: np.ndarray) -> np.ndarray: return rgba[..., 0] * 0.2126 + rgba[..., 1] * 0.7152 + rgba[..., 2] * 0.0722 @@ -65,7 +67,137 @@ def outline_straightness(rgba: np.ndarray) -> float: def reconstruction_error(a: np.ndarray, b: np.ndarray) -> float: - return float(np.mean(np.abs(a - b))) + return float(np.mean(np.abs(premultiply(a) - premultiply(b)))) + + +def sprite_mask(a: np.ndarray, b: np.ndarray, alpha_threshold: float = 0.05, halo: int = 1) -> np.ndarray: + mask = np.maximum(a[..., 3], b[..., 3]) >= alpha_threshold + if np.any(mask) and halo > 0: + mask = _dilate_mask(mask, radius=halo) + return mask + + +def foreground_coverage(a: np.ndarray, b: np.ndarray, alpha_threshold: float = 0.05, halo: int = 1) -> float: + return float(np.mean(sprite_mask(a, b, alpha_threshold=alpha_threshold, halo=halo))) + + +def foreground_reconstruction_error( + a: np.ndarray, + b: np.ndarray, + alpha_threshold: float = 0.05, + halo: int = 1, +) -> float: + mask = sprite_mask(a, b, alpha_threshold=alpha_threshold, halo=halo) + if not np.any(mask): + return reconstruction_error(a, b) + diff = np.mean(np.abs(premultiply(a) - premultiply(b)), axis=-1) + return float(np.mean(diff[mask])) + + +def foreground_adjacency_error( + a: np.ndarray, + b: np.ndarray, + alpha_threshold: float = 0.05, + halo: int = 1, +) -> float: + premul_a = premultiply(a) + premul_b = premultiply(b) + mask = sprite_mask(a, b, alpha_threshold=alpha_threshold, halo=halo) + diffs: list[np.ndarray] = [] + masks: list[np.ndarray] = [] + + for axis in (0, 1): + delta_a = np.diff(premul_a, axis=axis) + delta_b = np.diff(premul_b, axis=axis) + if axis == 0: + edge_mask = mask[1:, :] | mask[:-1, :] + else: + edge_mask = mask[:, 1:] | mask[:, :-1] + diffs.append(np.mean(np.abs(delta_a - delta_b), axis=-1)) + masks.append(edge_mask) + + weights = np.concatenate([edge_mask.reshape(-1) for edge_mask in masks], axis=0) + values = np.concatenate([diff.reshape(-1) for diff in diffs], axis=0) + if not np.any(weights): + return 0.0 + return float(np.mean(values[weights])) + + +def foreground_motif_error( + a: np.ndarray, + b: np.ndarray, + alpha_threshold: float = 0.05, + halo: int = 1, +) -> float: + premul_a = premultiply(a) + premul_b = premultiply(b) + mask = sprite_mask(a, b, alpha_threshold=alpha_threshold, halo=halo) + if premul_a.shape[0] < 2 or premul_a.shape[1] < 2: + return foreground_adjacency_error(a, b, alpha_threshold=alpha_threshold, halo=halo) + + blocks_a = _motif_blocks(premul_a) + blocks_b = _motif_blocks(premul_b) + block_mask = mask[:-1, :-1] | mask[:-1, 1:] | mask[1:, :-1] | mask[1:, 1:] + if not np.any(block_mask): + return 0.0 + diff = np.mean(np.abs(blocks_a - blocks_b), axis=(-2, -1)) + return float(np.mean(diff[block_mask])) + + +def source_lattice_consistency_breakdown( + source_rgba: np.ndarray, + output_rgba: np.ndarray, + *, + target_width: int, + target_height: int, + phase_x: float, + phase_y: float, + alpha_threshold: float = 0.05, +) -> dict[str, float]: + lattice_source, dispersion = lattice_source_rgba( + source_rgba, + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, + alpha_threshold=alpha_threshold, + ) + cell_error = foreground_reconstruction_error(output_rgba, lattice_source, alpha_threshold=alpha_threshold, halo=0) + adjacency_error = foreground_adjacency_error(output_rgba, lattice_source, alpha_threshold=alpha_threshold, halo=0) + motif_error = foreground_motif_error(output_rgba, lattice_source, alpha_threshold=alpha_threshold, halo=0) + score = ( + cell_error * 0.34 + + adjacency_error * 0.24 + + motif_error * 0.24 + + dispersion * 0.18 + ) + return { + "cell_error": cell_error, + "adjacency_error": adjacency_error, + "motif_error": motif_error, + "cell_dispersion": dispersion, + "score": score, + } + + +def exact_match_ratio(a: np.ndarray, b: np.ndarray) -> float: + a8 = np.clip(np.rint(premultiply(a) * 255.0), 0, 255).astype(np.uint8) + b8 = np.clip(np.rint(premultiply(b) * 255.0), 0, 255).astype(np.uint8) + return float(np.mean(np.all(a8 == b8, axis=-1))) + + +def foreground_exact_match_ratio( + a: np.ndarray, + b: np.ndarray, + alpha_threshold: float = 0.05, + halo: int = 1, +) -> float: + mask = sprite_mask(a, b, alpha_threshold=alpha_threshold, halo=halo) + if not np.any(mask): + return 1.0 + a8 = np.clip(np.rint(premultiply(a) * 255.0), 0, 255).astype(np.uint8) + b8 = np.clip(np.rint(premultiply(b) * 255.0), 0, 255).astype(np.uint8) + return float(np.mean(np.all(a8[mask] == b8[mask], axis=-1))) def coherence_breakdown(rgba: np.ndarray) -> dict[str, float]: @@ -84,3 +216,98 @@ def coherence_breakdown(rgba: np.ndarray) -> dict[str, float]: + 0.15 * (1.0 - min(1.0, breakdown["color_chatter"])) ) return breakdown + + +def _dilate_mask(mask: np.ndarray, radius: int) -> np.ndarray: + if radius <= 0: + return mask + height, width = mask.shape + padded = np.pad(mask, radius, mode="constant", constant_values=False) + windows = [] + for dy in range(radius * 2 + 1): + for dx in range(radius * 2 + 1): + windows.append(padded[dy : dy + height, dx : dx + width]) + return np.any(np.stack(windows, axis=0), axis=0) + + +def _motif_blocks(premul: np.ndarray) -> np.ndarray: + blocks = np.stack( + [ + premul[:-1, :-1, :], + premul[:-1, 1:, :], + premul[1:, :-1, :], + premul[1:, 1:, :], + ], + axis=-2, + ) + centered = blocks - np.mean(blocks, axis=-2, keepdims=True) + scale = np.maximum(1e-4, np.mean(np.abs(centered), axis=(-2, -1), keepdims=True)) + return centered / scale + + +def lattice_source_rgba( + source_rgba: np.ndarray, + *, + target_width: int, + target_height: int, + phase_x: float, + phase_y: float, + alpha_threshold: float = 0.05, +) -> tuple[np.ndarray, float]: + indices = _lattice_indices( + height=source_rgba.shape[0], + width=source_rgba.shape[1], + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, + ) + cell_count = max(1, target_width * target_height) + premul = premultiply(source_rgba) + flat_idx = indices.reshape(-1) + flat_premul = premul.reshape(-1, premul.shape[-1]) + counts = np.bincount(flat_idx, minlength=cell_count).astype(np.float32) + safe_counts = np.maximum(counts, 1.0) + channel_sums = [ + np.bincount(flat_idx, weights=flat_premul[:, channel], minlength=cell_count).astype(np.float32) + for channel in range(flat_premul.shape[-1]) + ] + mean_premul = np.stack(channel_sums, axis=-1) / safe_counts[:, None] + mean_premul = mean_premul.reshape(target_height, target_width, premul.shape[-1]) + mean_rgba = _unpremultiply_array(mean_premul) + + per_pixel_mean = mean_premul.reshape(-1, premul.shape[-1])[flat_idx] + pixel_diff = np.mean(np.abs(flat_premul - per_pixel_mean), axis=-1) + source_alpha = source_rgba.reshape(-1, source_rgba.shape[-1])[:, 3] + lattice_alpha = mean_rgba.reshape(-1, mean_rgba.shape[-1])[flat_idx, 3] + support_mask = np.maximum(source_alpha, lattice_alpha) >= alpha_threshold + if np.any(support_mask): + dispersion = float(np.mean(pixel_diff[support_mask])) + else: + dispersion = float(np.mean(pixel_diff)) if pixel_diff.size else 0.0 + return mean_rgba.astype(np.float32), dispersion + + +def _lattice_indices( + *, + height: int, + width: int, + target_width: int, + target_height: int, + phase_x: float, + phase_y: float, +) -> np.ndarray: + cell_x = width / max(1, target_width) + cell_y = height / max(1, target_height) + xs = (np.arange(width, dtype=np.float32) + 0.5) / cell_x - phase_x + ys = (np.arange(height, dtype=np.float32) + 0.5) / cell_y - phase_y + x_idx = np.clip(np.floor(xs).astype(np.int32), 0, max(0, target_width - 1)) + y_idx = np.clip(np.floor(ys).astype(np.int32), 0, max(0, target_height - 1)) + return y_idx[:, None] * target_width + x_idx[None, :] + + +def _unpremultiply_array(premul: np.ndarray, eps: float = 1e-6) -> np.ndarray: + out = premul.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) diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py new file mode 100644 index 0000000..1f460dd --- /dev/null +++ b/src/repixelizer/params.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass + + +@dataclass(slots=True) +class SolverHyperParams: + representative_softmax_scale: float = 18.0 + boundary_probe_scale: float = 0.22 + boundary_signed_weight: float = 0.55 + boundary_direction_weight: float = 0.30 + boundary_magnitude_weight: float = 0.15 + snap_base_match_weight: float = 0.65 + snap_neighbor_weight: float = 0.35 + snap_diagonal_weight: float = 0.14 + refine_anchor_weight: float = 0.60 + refine_representative_weight: float = 0.20 + refine_alpha_weight: float = 0.10 + refine_distance_weight: float = 0.10 + refine_source_delta_weight: float = 0.25 + refine_orthogonal_weight: float = 0.24 + refine_diagonal_weight: float = 0.08 + structure_boundary_weight: float = 0.60 + structure_anchor_adjacency_weight: float = 0.15 + structure_anchor_motif_weight: float = 0.20 + structure_representative_weight: float = 0.05 + alpha_foreground_threshold: float = 0.50 + alpha_representative_foreground_threshold: float = 0.60 + alpha_opaque_threshold: float = 0.95 + alpha_transparent_threshold: float = 0.05 + phase_rerank_margin: float = 0.004 + + def to_dict(self) -> dict[str, float]: + return asdict(self) diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index 5909e9f..27f2d5f 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -18,9 +18,11 @@ from .diagnostics import ( ) from .discrete import cleanup_pixels from .inference import infer_lattice, inference_to_json -from .io import load_rgba, save_rgba +from .io import load_rgba, nearest_resize, save_rgba +from .metrics import source_lattice_consistency_breakdown +from .params import SolverHyperParams from .palette import load_palette, quantize_rgba, save_palette_report -from .types import RunResult +from .types import InferenceResult, RunResult def run_pipeline( @@ -33,12 +35,22 @@ def run_pipeline( diagnostics_dir: str | Path | None = None, seed: int = 7, steps: int = 200, - device: str = "cpu", + device: str = "auto", + solver_params: SolverHyperParams | None = None, ) -> RunResult: started = time.perf_counter() + solver_params = solver_params or SolverHyperParams() source = load_rgba(input_path) - inference = infer_lattice(source, target_size=target_size) + inference = infer_lattice(source, target_size=target_size, device=device) analysis = analyze_source(source, seed=seed) + inference = _select_phase_candidate( + source, + inference, + analysis=analysis, + seed=seed, + device=device, + solver_params=solver_params, + ) solver = optimize_uv_field( source, inference=inference, @@ -46,6 +58,7 @@ def run_pipeline( steps=steps, seed=seed, device=device, + solver_params=solver_params, ) cleanup = cleanup_pixels(solver.target_rgba, source_guidance=solver.guidance_strength) palette = load_palette(palette_path) if palette_path else None @@ -69,6 +82,10 @@ def run_pipeline( 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) + save_rgba( + diagnostics_path / "output-preview.png", + nearest_resize(output_rgba, width=output_rgba.shape[1] * 8, height=output_rgba.shape[0] * 8), + ) 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) @@ -80,8 +97,64 @@ def run_pipeline( "seed": seed, "steps": steps, "device": device, + "solver_params": solver_params.to_dict(), } 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 + + +def _select_phase_candidate( + source: np.ndarray, + inference: InferenceResult, + *, + analysis, + seed: int, + device: str, + solver_params: SolverHyperParams | None = None, +) -> InferenceResult: + solver_params = solver_params or SolverHyperParams() + if len(inference.top_candidates) <= 1 or inference.confidence >= 0.02: + return inference + + top_score = float(inference.top_candidates[0].score) + baseline_rank: float | None = None + best_rank = float("inf") + best_candidate = inference + for candidate in inference.top_candidates[:8]: + candidate_inference = InferenceResult( + target_width=candidate.target_width, + target_height=candidate.target_height, + phase_x=candidate.phase_x, + phase_y=candidate.phase_y, + confidence=inference.confidence, + top_candidates=inference.top_candidates, + ) + candidate_artifacts = optimize_uv_field( + source, + inference=candidate_inference, + analysis=analysis, + steps=0, + seed=seed, + device=device, + solver_params=solver_params, + ) + support = source_lattice_consistency_breakdown( + source, + candidate_artifacts.target_rgba, + target_width=candidate.target_width, + target_height=candidate.target_height, + phase_x=candidate.phase_x, + phase_y=candidate.phase_y, + ) + coherence_penalty = top_score - float(candidate.score) + rank = support["score"] + coherence_penalty * 0.5 + if baseline_rank is None: + baseline_rank = rank + if rank < best_rank: + best_rank = rank + best_candidate = candidate_inference + if baseline_rank is None or best_rank > baseline_rank - solver_params.phase_rerank_margin: + return inference + return best_candidate diff --git a/src/repixelizer/synthetic.py b/src/repixelizer/synthetic.py index 7dcacd1..01a12fb 100644 --- a/src/repixelizer/synthetic.py +++ b/src/repixelizer/synthetic.py @@ -1,10 +1,10 @@ from __future__ import annotations -from pathlib import Path - import numpy as np from PIL import Image, ImageDraw, ImageFilter +from .io import premultiply, unpremultiply + def make_emblem(width: int = 32, height: int = 32) -> np.ndarray: image = Image.new("RGBA", (width, height), (0, 0, 0, 0)) @@ -34,13 +34,177 @@ def fake_pixelize( phase_x: float = 0.22, phase_y: float = 0.31, blur_radius: float = 0.7, + warp_strength: float = 0.0, + warp_detail: int = 4, + warp_sample_mode: str = "bilinear", + seed: int = 0, ) -> 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)) + target_rgba = np.asarray(target, dtype=np.float32) / 255.0 + height_px, width_px = target_rgba.shape[:2] + grid_x, grid_y = np.meshgrid(np.arange(width_px, dtype=np.float32), np.arange(height_px, dtype=np.float32)) + sample_x = grid_x - phase_x * float(upscale) + sample_y = grid_y - phase_y * float(upscale) + if warp_strength > 0.0: + rng = np.random.default_rng(seed) + strength_px = float(upscale) * warp_strength + warp_x, warp_y = _warp_fields(width_px, height_px, upscale=upscale, detail=warp_detail, rng=rng) + sample_x -= warp_x * strength_px + sample_y -= warp_y * strength_px + warped = _sample_rgba(target_rgba, sample_x, sample_y, mode=warp_sample_mode) + if blur_radius <= 0.0: + return warped + blurred = Image.fromarray(np.clip(np.rint(warped * 255.0), 0, 255).astype(np.uint8), mode="RGBA").filter( + ImageFilter.GaussianBlur(radius=blur_radius) + ) return np.asarray(blurred, dtype=np.float32) / 255.0 + + +def _warp_fields(width: int, height: int, *, upscale: int, detail: int, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]: + smooth_tile = max(float(upscale * (detail + 2)), 2.0) + block_tile = max(float(upscale * max(1.5, 7.5 - detail)), 2.0) + patch_tile = max(float(upscale * max(1.0, 5.5 - detail * 0.5)), 2.0) + stripe_tile = max(float(upscale * max(1.0, 6.5 - detail)), 2.0) + smooth_blur = max(0.0, upscale * 0.2) + block_blur = max(0.0, upscale * 0.32) + patch_blur = max(0.0, upscale * 0.12) + stripe_blur = max(0.0, upscale * 0.16) + warp_x = ( + _noise_field(width, height, tile_size=smooth_tile, rng=rng, blur_radius=smooth_blur) * 0.15 + + _noise_field( + width, + height, + tile_size=block_tile, + rng=rng, + resample=Image.Resampling.NEAREST, + blur_radius=block_blur, + ) + * 0.45 + + _noise_field( + width, + height, + tile_size=patch_tile, + rng=rng, + resample=Image.Resampling.NEAREST, + blur_radius=patch_blur, + ) + * 0.25 + + _stripe_field( + width, + height, + tile_size=stripe_tile, + rng=rng, + axis="x", + blur_radius=stripe_blur, + ) + * 0.15 + ) + warp_y = ( + _noise_field(width, height, tile_size=smooth_tile, rng=rng, blur_radius=smooth_blur) * 0.15 + + _noise_field( + width, + height, + tile_size=block_tile, + rng=rng, + resample=Image.Resampling.NEAREST, + blur_radius=block_blur, + ) + * 0.45 + + _noise_field( + width, + height, + tile_size=patch_tile, + rng=rng, + resample=Image.Resampling.NEAREST, + blur_radius=patch_blur, + ) + * 0.25 + + _stripe_field( + width, + height, + tile_size=stripe_tile, + rng=rng, + axis="y", + blur_radius=stripe_blur, + ) + * 0.15 + ) + return np.clip(warp_x, -1.25, 1.25), np.clip(warp_y, -1.25, 1.25) + + +def _noise_field( + width: int, + height: int, + *, + tile_size: float, + rng: np.random.Generator, + resample: int = Image.Resampling.BILINEAR, + blur_radius: float = 0.0, +) -> np.ndarray: + coarse_width = max(2, int(np.ceil(width / max(tile_size, 1.0))) + 1) + coarse_height = max(2, int(np.ceil(height / max(tile_size, 1.0))) + 1) + noise = (rng.random((coarse_height, coarse_width)) * 255.0).astype(np.uint8) + image = Image.fromarray(noise, mode="L").resize((width, height), resample=resample) + if blur_radius > 0.0: + image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + return np.asarray(image, dtype=np.float32) / 255.0 * 2.0 - 1.0 + + +def _stripe_field( + width: int, + height: int, + *, + tile_size: float, + rng: np.random.Generator, + axis: str, + blur_radius: float, +) -> np.ndarray: + if axis == "x": + coarse_height = max(2, int(np.ceil(height / max(tile_size, 1.0))) + 1) + noise = (rng.random((coarse_height, 1)) * 255.0).astype(np.uint8) + image = Image.fromarray(noise, mode="L").resize((1, height), resample=Image.Resampling.NEAREST) + if blur_radius > 0.0: + image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + values = np.asarray(image, dtype=np.float32) / 255.0 * 2.0 - 1.0 + return np.repeat(values, width, axis=1) + coarse_width = max(2, int(np.ceil(width / max(tile_size, 1.0))) + 1) + noise = (rng.random((1, coarse_width)) * 255.0).astype(np.uint8) + image = Image.fromarray(noise, mode="L").resize((width, 1), resample=Image.Resampling.NEAREST) + if blur_radius > 0.0: + image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + values = np.asarray(image, dtype=np.float32) / 255.0 * 2.0 - 1.0 + return np.repeat(values, height, axis=0) + + +def _sample_rgba(rgba: np.ndarray, sample_x: np.ndarray, sample_y: np.ndarray, *, mode: str) -> np.ndarray: + if mode == "nearest": + return _nearest_sample_rgba(rgba, sample_x, sample_y) + if mode == "bilinear": + return _bilinear_sample_rgba(rgba, sample_x, sample_y) + raise ValueError(f"Unsupported warp sample mode: {mode}") + + +def _nearest_sample_rgba(rgba: np.ndarray, sample_x: np.ndarray, sample_y: np.ndarray) -> np.ndarray: + height, width = rgba.shape[:2] + x = np.clip(np.rint(sample_x), 0, max(0, width - 1)).astype(np.int32) + y = np.clip(np.rint(sample_y), 0, max(0, height - 1)).astype(np.int32) + return rgba[y, x] + + +def _bilinear_sample_rgba(rgba: np.ndarray, sample_x: np.ndarray, sample_y: np.ndarray) -> np.ndarray: + premult = premultiply(rgba) + height, width = premult.shape[:2] + clamped_x = np.clip(sample_x, 0.0, max(0.0, width - 1)) + clamped_y = np.clip(sample_y, 0.0, max(0.0, height - 1)) + x0 = np.floor(clamped_x).astype(np.int32) + y0 = np.floor(clamped_y).astype(np.int32) + x1 = np.clip(x0 + 1, 0, max(0, width - 1)) + y1 = np.clip(y0 + 1, 0, max(0, height - 1)) + wx = (clamped_x - x0).astype(np.float32)[..., None] + wy = (clamped_y - y0).astype(np.float32)[..., None] + top = premult[y0, x0] * (1.0 - wx) + premult[y0, x1] * wx + bottom = premult[y1, x0] * (1.0 - wx) + premult[y1, x1] * wx + sampled = top * (1.0 - wy) + bottom * wy + return unpremultiply(np.clip(sampled, 0.0, 1.0)) diff --git a/src/repixelizer/tuning.py b/src/repixelizer/tuning.py new file mode 100644 index 0000000..e96a168 --- /dev/null +++ b/src/repixelizer/tuning.py @@ -0,0 +1,275 @@ +from __future__ import annotations + +import shutil +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np + +from .benchmark import run_roundtrip_benchmark +from .diagnostics import write_json +from .params import SolverHyperParams + + +def _safe_ratio(numerator: float, denominator: float) -> float: + return float(numerator / max(denominator, 1e-6)) + + +def _row_score(row: dict[str, Any]) -> float: + error_ratio = _safe_ratio(float(row["optimized_error_to_original"]), float(row["naive_error_to_original"])) + adjacency_ratio = _safe_ratio( + float(row["optimized_adjacency_error_to_original"]), + float(row["naive_adjacency_error_to_original"]), + ) + motif_ratio = _safe_ratio(float(row["optimized_motif_error_to_original"]), float(row["naive_motif_error_to_original"])) + return error_ratio * 0.50 + adjacency_ratio * 0.25 + motif_ratio * 0.25 + + +def _score_summary(summary: dict[str, Any]) -> dict[str, float]: + row_scores = np.asarray([_row_score(row) for row in summary["rows"]], dtype=np.float64) + mean_row_score = float(np.mean(row_scores)) + worst_row_score = float(np.max(row_scores)) + score = mean_row_score * 0.70 + worst_row_score * 0.30 + return { + "score": score, + "mean_row_score": mean_row_score, + "worst_row_score": worst_row_score, + } + + +def _normalize_group(values: np.ndarray) -> np.ndarray: + values = np.clip(values.astype(np.float64), 1e-4, None) + return values / np.sum(values) + + +def _mutate_group(rng: np.random.Generator, values: list[float], scale: float) -> np.ndarray: + base = np.asarray(values, dtype=np.float64) + noise = np.exp(rng.normal(0.0, 0.28 * scale, size=base.shape)) + return _normalize_group(base * noise) + + +def _mutate_positive( + rng: np.random.Generator, + value: float, + *, + scale: float, + low: float, + high: float, +) -> float: + mutated = value * float(np.exp(rng.normal(0.0, 0.22 * scale))) + return float(np.clip(mutated, low, high)) + + +def _mutate_linear( + rng: np.random.Generator, + value: float, + *, + scale: float, + sigma: float, + low: float, + high: float, +) -> float: + mutated = value + float(rng.normal(0.0, sigma * scale)) + return float(np.clip(mutated, low, high)) + + +def _mutate_params(base: SolverHyperParams, rng: np.random.Generator, scale: float) -> SolverHyperParams: + refine_group = _mutate_group( + rng, + [ + base.refine_anchor_weight, + base.refine_representative_weight, + base.refine_alpha_weight, + base.refine_distance_weight, + ], + scale, + ) + structure_group = _mutate_group( + rng, + [ + base.structure_boundary_weight, + base.structure_anchor_adjacency_weight, + base.structure_anchor_motif_weight, + base.structure_representative_weight, + ], + scale, + ) + return replace( + base, + representative_softmax_scale=_mutate_positive( + rng, + base.representative_softmax_scale, + scale=scale, + low=8.0, + high=36.0, + ), + boundary_probe_scale=_mutate_linear( + rng, + base.boundary_probe_scale, + scale=scale, + sigma=0.03, + low=0.14, + high=0.32, + ), + refine_anchor_weight=float(refine_group[0]), + refine_representative_weight=float(refine_group[1]), + refine_alpha_weight=float(refine_group[2]), + refine_distance_weight=float(refine_group[3]), + refine_source_delta_weight=_mutate_linear( + rng, + base.refine_source_delta_weight, + scale=scale, + sigma=0.08, + low=0.05, + high=0.45, + ), + refine_orthogonal_weight=_mutate_linear( + rng, + base.refine_orthogonal_weight, + scale=scale, + sigma=0.05, + low=0.10, + high=0.40, + ), + refine_diagonal_weight=_mutate_linear( + rng, + base.refine_diagonal_weight, + scale=scale, + sigma=0.025, + low=0.02, + high=0.18, + ), + structure_boundary_weight=float(structure_group[0]), + structure_anchor_adjacency_weight=float(structure_group[1]), + structure_anchor_motif_weight=float(structure_group[2]), + structure_representative_weight=float(structure_group[3]), + phase_rerank_margin=_mutate_linear( + rng, + base.phase_rerank_margin, + scale=scale, + sigma=0.0015, + low=0.0005, + high=0.010, + ), + ) + + +def tune_solver_hyperparams( + corpus_dir: str | Path, + out_dir: str | Path, + *, + trials: int = 8, + variants: int = 1, + profiles: list[str] | None = None, + seed: int = 7, + steps: int = 48, + device: str = "auto", + infer_size: bool = False, + include_cases: list[str] | None = None, + limit_cases: int | None = None, +) -> dict[str, Any]: + if trials < 1: + raise ValueError("trials must be at least 1") + + selected_profiles = profiles or ["soft"] + effective_limit = limit_cases + if effective_limit is None and not include_cases: + effective_limit = 8 + + out_path = Path(out_dir) + if out_path.exists(): + shutil.rmtree(out_path) + out_path.mkdir(parents=True, exist_ok=True) + scratch_dir = out_path / "_scratch" + best_dir = out_path / "best-benchmark" + rng = np.random.default_rng(seed) + + best_params = SolverHyperParams() + best_score = float("inf") + best_summary: dict[str, Any] | None = None + trials_payload: list[dict[str, Any]] = [] + + for trial_index in range(trials): + if trial_index == 0: + candidate = best_params + origin = "default" + else: + anneal = max(0.35, 1.0 - (trial_index / max(2, trials))) + candidate = _mutate_params(best_params, rng, scale=anneal) + origin = "mutated-best" + + summary = run_roundtrip_benchmark( + corpus_dir, + scratch_dir, + variants=variants, + profiles=selected_profiles, + seed=seed, + steps=steps, + device=device, + infer_size=infer_size, + include_cases=include_cases, + limit_cases=effective_limit, + keep_existing=False, + solver_params=candidate, + ) + score_info = _score_summary(summary) + improved = score_info["score"] < best_score + trials_payload.append( + { + "trial": trial_index + 1, + "origin": origin, + "score": score_info["score"], + "mean_row_score": score_info["mean_row_score"], + "worst_row_score": score_info["worst_row_score"], + "params": candidate.to_dict(), + "improved_best": improved, + } + ) + if improved: + best_score = score_info["score"] + best_params = candidate + best_summary = summary + + final_summary = run_roundtrip_benchmark( + corpus_dir, + best_dir, + variants=variants, + profiles=selected_profiles, + seed=seed, + steps=steps, + device=device, + infer_size=infer_size, + include_cases=include_cases, + limit_cases=effective_limit, + keep_existing=False, + solver_params=best_params, + ) + final_score = _score_summary(final_summary) + + if scratch_dir.exists(): + shutil.rmtree(scratch_dir) + + payload = { + "trials_requested": trials, + "trials_completed": len(trials_payload), + "variants": variants, + "profiles": selected_profiles, + "seed": seed, + "steps": steps, + "device": device, + "infer_size": infer_size, + "include_cases": include_cases or [], + "limit_cases_effective": effective_limit, + "objective": "0.5 * error_ratio_to_naive + 0.25 * adjacency_ratio_to_naive + 0.25 * motif_ratio_to_naive; lower is better", + "trial_results": trials_payload, + "best_score": final_score["score"], + "best_mean_row_score": final_score["mean_row_score"], + "best_worst_row_score": final_score["worst_row_score"], + "best_params": best_params.to_dict(), + "best_benchmark_dir": str(best_dir), + "best_benchmark": final_summary, + "initial_benchmark": best_summary, + } + write_json(out_path / "tuning-results.json", payload) + return payload diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py new file mode 100644 index 0000000..8d54a01 --- /dev/null +++ b/tests/test_benchmark.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from repixelizer.benchmark import run_roundtrip_benchmark +from repixelizer.io import save_rgba +from repixelizer.synthetic import make_emblem + + +def test_roundtrip_benchmark_writes_summary_and_case_outputs(tmp_path: Path) -> None: + corpus_dir = tmp_path / "corpus" + originals_dir = corpus_dir / "originals" + originals_dir.mkdir(parents=True) + save_rgba(originals_dir / "crest.png", make_emblem(16, 16)) + (originals_dir / "crest.json").write_text( + json.dumps( + { + "title": "Crest", + "author": "Test Artist", + "license": "CC-BY 4.0", + "source_url": "https://example.com/crest", + } + ), + encoding="utf-8", + ) + + out_dir = tmp_path / "benchmark-out" + payload = run_roundtrip_benchmark(corpus_dir, out_dir, variants=1, steps=4, seed=3) + + assert payload["case_count"] == 1 + assert payload["row_count"] == 2 + assert payload["profiles"] == ["soft", "crisp"] + assert payload["primary_metric"] == "foreground_premultiplied_mae" + assert (out_dir / "benchmark.csv").exists() + assert (out_dir / "benchmark.json").exists() + assert (out_dir / "cases" / "crest" / "original.png").exists() + assert (out_dir / "cases" / "crest" / "profile-soft" / "variant-01" / "input.png").exists() + assert (out_dir / "cases" / "crest" / "profile-crisp" / "variant-01" / "optimized.png").exists() + assert "reference_foreground_coverage" in payload["rows"][0] + assert "optimized_canvas_error_to_original" in payload["rows"][0] + assert "optimized_adjacency_error_to_original" in payload["rows"][0] + assert "optimized_motif_error_to_original" in payload["rows"][0] + assert "profile" in payload["rows"][0] + assert "warp_sample_mode" in payload["rows"][0] + + +def test_roundtrip_benchmark_can_filter_cases(tmp_path: Path) -> None: + corpus_dir = tmp_path / "corpus" + originals_dir = corpus_dir / "originals" + originals_dir.mkdir(parents=True) + save_rgba(originals_dir / "crest-a.png", make_emblem(16, 16)) + save_rgba(originals_dir / "crest-b.png", make_emblem(18, 18)) + + out_dir = tmp_path / "benchmark-out" + payload = run_roundtrip_benchmark( + corpus_dir, + out_dir, + variants=1, + profiles=["crisp"], + steps=2, + seed=3, + include_cases=["crest-b"], + ) + + assert payload["case_count"] == 1 + assert payload["row_count"] == 1 + assert payload["rows"][0]["case_id"] == "crest-b" + assert payload["rows"][0]["profile"] == "crisp" + + +def test_roundtrip_benchmark_clears_existing_output_dir_by_default(tmp_path: Path) -> None: + corpus_dir = tmp_path / "corpus" + originals_dir = corpus_dir / "originals" + originals_dir.mkdir(parents=True) + save_rgba(originals_dir / "crest.png", make_emblem(16, 16)) + + out_dir = tmp_path / "benchmark-out" + stale_file = out_dir / "stale.txt" + stale_file.parent.mkdir(parents=True) + stale_file.write_text("old", encoding="utf-8") + + run_roundtrip_benchmark(corpus_dir, out_dir, variants=1, profiles=["soft"], steps=2, seed=3) + + assert not stale_file.exists() + assert (out_dir / "benchmark.json").exists() diff --git a/tests/test_corpus.py b/tests/test_corpus.py new file mode 100644 index 0000000..de3a6fd --- /dev/null +++ b/tests/test_corpus.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np + +from repixelizer.corpus import prepare_corpus +from repixelizer.io import save_rgba + + +def test_prepare_corpus_extracts_frames_splits_sheets_and_archives_sources(tmp_path: Path) -> None: + corpus_dir = tmp_path / "corpus" + originals_dir = corpus_dir / "originals" + originals_dir.mkdir(parents=True) + + mage = np.zeros((188, 340, 4), dtype=np.float32) + mage[..., 1] = 1.0 + mage[20:60, 20:50] = np.array([1.0, 0.0, 0.0, 1.0], dtype=np.float32) + mage[110:160, 270:320] = np.array([0.0, 0.0, 1.0, 1.0], dtype=np.float32) + save_rgba(originals_dir / "mage-1-85x94.png", mage) + + shadow = np.zeros((350, 320, 4), dtype=np.float32) + shadow[..., :3] = np.array([0.0, 1.0 / 255.0, 1.0 / 255.0], dtype=np.float32) + shadow[..., 3] = 1.0 + shadow[290:340, 250:300] = np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32) + save_rgba(originals_dir / "shadow-80x70.png", shadow) + + multi = np.ones((110, 317, 4), dtype=np.float32) + multi[..., 3] = 1.0 + multi[8:102, 0:79] = np.array([1.0, 0.0, 0.0, 1.0], dtype=np.float32) + multi[0:103, 88:210] = np.array([0.0, 1.0, 0.0, 1.0], dtype=np.float32) + multi[0:110, 231:317] = np.array([0.0, 0.0, 1.0, 1.0], dtype=np.float32) + save_rgba(originals_dir / "threeformsPJ2.png", multi) + + payload = prepare_corpus(corpus_dir) + + assert payload["output_count"] == 5 + assert (originals_dir / "mage-1.png").exists() + assert (originals_dir / "shadow.png").exists() + assert (originals_dir / "threeformspj2-01.png").exists() + assert (originals_dir / "threeformspj2-02.png").exists() + assert (originals_dir / "threeformspj2-03.png").exists() + assert (corpus_dir / "source-sheets" / "mage-1-85x94.png").exists() + assert (corpus_dir / "source-sheets" / "shadow-80x70.png").exists() + assert (corpus_dir / "source-sheets" / "threeformsPJ2.png").exists() + attribution_text = (corpus_dir / "ATTRIBUTION.md").read_text(encoding="utf-8") + mage_sidecar = json.loads((originals_dir / "mage-1.json").read_text(encoding="utf-8")) + shadow_sidecar = json.loads((originals_dir / "shadow.json").read_text(encoding="utf-8")) + boss_sidecar = json.loads((originals_dir / "threeformspj2-01.json").read_text(encoding="utf-8")) + assert mage_sidecar["source_title"] == "Bosses and monsters spritesheets (Ars Notoria)" + assert mage_sidecar["attribution_confidence"] == "inferred" + assert shadow_sidecar["licenses"] == ["CC-BY-SA 3.0"] + assert shadow_sidecar["original_source_title"] == "First Person Dungeon Crawl Enemies Remixed" + assert boss_sidecar["source_title"] == '3-form RPG boss: "Harlequin Epicycle."' + assert "## Bosses and monsters spritesheets (Ars Notoria)" in attribution_text + assert "`mage-1.png`" in attribution_text + assert "`shadow.png`" in attribution_text + assert "## 3-form RPG boss: \"Harlequin Epicycle.\"" in attribution_text diff --git a/tests/test_inference.py b/tests/test_inference.py index 10069af..0ad1c88 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -6,7 +6,16 @@ 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) + fake = fake_pixelize( + source, + upscale=12, + phase_x=0.2, + phase_y=0.35, + blur_radius=0.75, + warp_strength=0.28, + warp_detail=6, + seed=5, + ) result = infer_lattice(fake) assert result.target_width in range(28, 37) assert result.target_height in range(28, 37) diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..18b8519 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import numpy as np + +from repixelizer.metrics import ( + exact_match_ratio, + foreground_adjacency_error, + foreground_exact_match_ratio, + foreground_motif_error, + foreground_reconstruction_error, + source_lattice_consistency_breakdown, + reconstruction_error, +) + + +def test_reconstruction_error_ignores_hidden_rgb_in_transparent_pixels() -> None: + transparent_a = np.zeros((4, 4, 4), dtype=np.float32) + transparent_b = transparent_a.copy() + transparent_b[..., 0] = 1.0 + + assert reconstruction_error(transparent_a, transparent_b) == 0.0 + assert exact_match_ratio(transparent_a, transparent_b) == 1.0 + + +def test_foreground_metrics_penalize_visible_sprite_errors_and_leaks() -> None: + original = np.zeros((4, 4, 4), dtype=np.float32) + original[1, 1] = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32) + + hidden_rgb = original.copy() + hidden_rgb[0, 0, :3] = np.asarray([0.2, 0.8, 0.4], dtype=np.float32) + assert foreground_reconstruction_error(original, hidden_rgb) == 0.0 + assert foreground_exact_match_ratio(original, hidden_rgb) == 1.0 + + wrong_foreground = original.copy() + wrong_foreground[1, 1] = np.asarray([0.0, 1.0, 0.0, 1.0], dtype=np.float32) + assert foreground_reconstruction_error(original, wrong_foreground) > 0.0 + assert foreground_exact_match_ratio(original, wrong_foreground) < 1.0 + + leaked_foreground = original.copy() + leaked_foreground[0, 0] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32) + assert foreground_reconstruction_error(original, leaked_foreground) > 0.0 + + +def test_texture_metrics_track_adjacency_and_local_motifs() -> None: + original = np.zeros((4, 4, 4), dtype=np.float32) + original[1:3, 1:3] = np.asarray( + [ + [[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]], + [[0.0, 0.0, 1.0, 1.0], [1.0, 0.0, 0.0, 1.0]], + ], + dtype=np.float32, + ) + + same_colors_wrong_structure = original.copy() + same_colors_wrong_structure[1:3, 1:3] = np.asarray( + [ + [[1.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 1.0]], + [[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0]], + ], + dtype=np.float32, + ) + + assert foreground_adjacency_error(original, original) == 0.0 + assert foreground_motif_error(original, original) == 0.0 + assert foreground_adjacency_error(original, same_colors_wrong_structure) > 0.0 + assert foreground_motif_error(original, same_colors_wrong_structure) > 0.0 + + +def test_source_lattice_consistency_prefers_matching_cell_structure() -> None: + lowres = np.asarray( + [ + [[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]], + [[0.0, 1.0, 0.0, 1.0], [1.0, 1.0, 0.0, 1.0]], + ], + dtype=np.float32, + ) + source = np.repeat(np.repeat(lowres, 4, axis=0), 4, axis=1) + good = lowres.copy() + bad = lowres[:, ::-1].copy() + + good_score = source_lattice_consistency_breakdown( + source, + good, + target_width=2, + target_height=2, + phase_x=0.0, + phase_y=0.0, + )["score"] + bad_score = source_lattice_consistency_breakdown( + source, + bad, + target_width=2, + target_height=2, + phase_x=0.0, + phase_y=0.0, + )["score"] + + assert good_score < bad_score diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 3344e9c..9309a6b 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -2,10 +2,14 @@ from __future__ import annotations from pathlib import Path +import numpy as np + from repixelizer.baselines import naive_resize_baseline -from repixelizer.metrics import coherence_breakdown -from repixelizer.pipeline import run_pipeline +from repixelizer.metrics import coherence_breakdown, foreground_reconstruction_error +from repixelizer.pipeline import _select_phase_candidate, run_pipeline +from repixelizer.io import nearest_resize from repixelizer.synthetic import fake_pixelize, make_emblem, make_sprite +from repixelizer.types import InferenceCandidate, InferenceResult def test_pipeline_writes_output_and_diagnostics(tmp_path: Path) -> None: @@ -21,11 +25,12 @@ def test_pipeline_writes_output_and_diagnostics(tmp_path: Path) -> None: result = run_pipeline(input_path, output_path, diagnostics_dir=diagnostics_dir, steps=24) assert output_path.exists() assert (diagnostics_dir / "run.json").exists() + assert (diagnostics_dir / "output-preview.png").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: +def test_pipeline_beats_naive_on_roundtrip_fidelity_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" @@ -36,9 +41,11 @@ def test_pipeline_beats_naive_on_coherence_for_synthetic_emblem(tmp_path: Path) 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 + optimized_preview = nearest_resize(result.output_rgba, width=source.shape[1], height=source.shape[0]) + naive_preview = nearest_resize(naive, width=source.shape[1], height=source.shape[0]) + optimized_error = foreground_reconstruction_error(optimized_preview, source) + naive_error = foreground_reconstruction_error(naive_preview, source) + assert optimized_error <= naive_error def test_pipeline_preserves_transparency_for_sprite(tmp_path: Path) -> None: @@ -54,3 +61,53 @@ def test_pipeline_preserves_transparency_for_sprite(tmp_path: Path) -> None: alpha = result.output_rgba[..., 3] assert alpha.min() < 0.05 assert alpha.max() > 0.9 + partial = alpha[(alpha > 0.05) & (alpha < 0.95)] + assert partial.size == 0 + + +def test_pipeline_can_emit_initialized_output_without_optimizer_steps(tmp_path: Path) -> None: + source = make_emblem(24, 24) + fake = fake_pixelize(source, upscale=10, phase_x=0.15, phase_y=0.2, blur_radius=0.6, warp_strength=0.2, warp_detail=5) + 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=0) + assert output_path.exists() + assert result.output_rgba.shape[0] == result.inference.target_height + assert result.output_rgba.shape[1] == result.inference.target_width + + +def test_phase_rerank_can_override_low_confidence_inference_pick(monkeypatch) -> None: + source = np.zeros((2, 2, 4), dtype=np.float32) + source[0, 0] = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32) + candidate_a = InferenceCandidate(target_width=2, target_height=2, phase_x=-0.4, phase_y=0.2, score=0.91, breakdown={}) + candidate_b = InferenceCandidate(target_width=2, target_height=2, phase_x=0.4, phase_y=0.2, score=0.909, breakdown={}) + inference = InferenceResult( + target_width=2, + target_height=2, + phase_x=candidate_a.phase_x, + phase_y=candidate_a.phase_y, + confidence=0.0, + top_candidates=[candidate_a, candidate_b], + ) + outputs = { + (candidate_a.phase_x, candidate_a.phase_y): np.ones_like(source), + (candidate_b.phase_x, candidate_b.phase_y): source.copy(), + } + + class DummyArtifacts: + def __init__(self, rgba: np.ndarray) -> None: + self.target_rgba = rgba + + def fake_optimize_uv_field(source_rgba, inference, analysis, steps, seed, device, solver_params=None): + assert steps == 0 + return DummyArtifacts(outputs[(inference.phase_x, inference.phase_y)]) + + monkeypatch.setattr("repixelizer.pipeline.optimize_uv_field", fake_optimize_uv_field) + + selected = _select_phase_candidate(source, inference, analysis=object(), seed=7, device="cpu") + assert selected.phase_x == candidate_b.phase_x + assert selected.phase_y == candidate_b.phase_y diff --git a/tests/test_tuning.py b/tests/test_tuning.py new file mode 100644 index 0000000..1b0c59d --- /dev/null +++ b/tests/test_tuning.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path + +from repixelizer.io import save_rgba +from repixelizer.synthetic import make_emblem +from repixelizer.tuning import tune_solver_hyperparams + + +def test_tuning_writes_results_and_best_benchmark(tmp_path: Path) -> None: + corpus_dir = tmp_path / "corpus" + originals_dir = corpus_dir / "originals" + originals_dir.mkdir(parents=True) + save_rgba(originals_dir / "crest.png", make_emblem(16, 16)) + + out_dir = tmp_path / "tuning-out" + payload = tune_solver_hyperparams( + corpus_dir, + out_dir, + trials=2, + variants=1, + profiles=["soft"], + steps=2, + seed=3, + device="cpu", + ) + + assert payload["trials_completed"] == 2 + assert "best_params" in payload + assert (out_dir / "tuning-results.json").exists() + assert (out_dir / "best-benchmark" / "benchmark.json").exists() + assert not (out_dir / "_scratch").exists()