Add corpus benchmarks and lattice refinement
This commit is contained in:
Meta Vens
parent
634ca7f002
commit
d2e8ccac4d
23 changed files with 2995 additions and 217 deletions
286
src/repixelizer/benchmark.py
Normal file
286
src/repixelizer/benchmark.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
500
src/repixelizer/corpus.py
Normal file
500
src/repixelizer/corpus.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
34
src/repixelizer/params.py
Normal file
34
src/repixelizer/params.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
275
src/repixelizer/tuning.py
Normal file
275
src/repixelizer/tuning.py
Normal file
|
|
@ -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
|
||||
Loading…
Reference in a new issue