646 lines
25 KiB
Python
646 lines
25 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from repixelizer.baselines import naive_resize_baseline
|
|
from repixelizer.metrics import coherence_breakdown, foreground_reconstruction_error, source_lattice_consistency_breakdown
|
|
from repixelizer.params import SolverHyperParams
|
|
from repixelizer.pipeline import _run_reconstruction, _select_candidate, run_pipeline
|
|
from repixelizer.observe import PipelineCancelled
|
|
from repixelizer.phase_field import optimize_phase_field
|
|
from repixelizer.io import nearest_resize
|
|
from repixelizer.synthetic import fake_pixelize, make_emblem, make_sprite
|
|
from repixelizer.types import InferenceCandidate, InferenceResult, SolverArtifacts
|
|
|
|
|
|
def test_pipeline_writes_output_and_diagnostics(tmp_path: Path) -> None:
|
|
source = make_emblem(32, 32)
|
|
fake = fake_pixelize(source, upscale=10, blur_radius=0.5)
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
diagnostics_dir = tmp_path / "diagnostics"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, fake)
|
|
result = run_pipeline(input_path, output_path, diagnostics_dir=diagnostics_dir, steps=24)
|
|
assert output_path.exists()
|
|
assert (diagnostics_dir / "run.json").exists()
|
|
assert (diagnostics_dir / "output-preview.png").exists()
|
|
assert (diagnostics_dir / "displacement-initial_output.png").exists()
|
|
assert (diagnostics_dir / "displacement-initial_output-preview.png").exists()
|
|
import json
|
|
|
|
run_json = json.loads((diagnostics_dir / "run.json").read_text(encoding="utf-8"))
|
|
assert set(run_json["source_fidelity"].keys()) == {"initial_output", "solver_target", "final_output"}
|
|
assert "source_structure" in run_json
|
|
assert {"edge_f1", "exact_match_ratio", "score"} <= set(run_json["source_structure"].keys())
|
|
assert "candidate_rerank_candidates" in run_json
|
|
assert {"initial_output", "final_output"} <= set(run_json["optimizer_displacement"].keys())
|
|
assert result.output_rgba.shape[0] == result.inference.target_height
|
|
assert result.output_rgba.shape[1] == result.inference.target_width
|
|
|
|
|
|
def test_pipeline_beats_naive_on_roundtrip_fidelity_for_synthetic_emblem(tmp_path: Path) -> None:
|
|
source = make_emblem(32, 32)
|
|
fake = fake_pixelize(
|
|
source,
|
|
upscale=12,
|
|
blur_radius=0.65,
|
|
warp_strength=0.28,
|
|
warp_detail=6,
|
|
seed=5,
|
|
)
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, fake)
|
|
result = run_pipeline(input_path, output_path, steps=32)
|
|
naive = naive_resize_baseline(fake, width=result.inference.target_width, height=result.inference.target_height)
|
|
optimized_preview = nearest_resize(result.output_rgba, width=source.shape[1], height=source.shape[0])
|
|
naive_preview = nearest_resize(naive, width=source.shape[1], height=source.shape[0])
|
|
optimized_error = foreground_reconstruction_error(optimized_preview, source)
|
|
naive_error = foreground_reconstruction_error(naive_preview, source)
|
|
assert optimized_error <= naive_error
|
|
|
|
|
|
def test_pipeline_preserves_transparency_for_sprite(tmp_path: Path) -> None:
|
|
source = make_sprite(24, 24)
|
|
fake = fake_pixelize(source, upscale=9, blur_radius=0.45)
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, fake)
|
|
result = run_pipeline(input_path, output_path, target_size=24, steps=20)
|
|
alpha = result.output_rgba[..., 3]
|
|
assert alpha.min() < 0.05
|
|
assert alpha.max() > 0.9
|
|
partial = alpha[(alpha > 0.05) & (alpha < 0.95)]
|
|
assert partial.size == 0
|
|
|
|
|
|
def test_pipeline_can_strip_checkerboard_background(tmp_path: Path) -> None:
|
|
source = make_sprite(24, 24)
|
|
checker = np.ones_like(source)
|
|
checker[..., 3] = 1.0
|
|
for y in range(checker.shape[0]):
|
|
for x in range(checker.shape[1]):
|
|
tile = 0.97 if ((x // 4) + (y // 4)) % 2 == 0 else 0.91
|
|
checker[y, x, :3] = tile
|
|
alpha = source[..., 3:4]
|
|
opaque_input = checker.copy()
|
|
opaque_input[..., :3] = source[..., :3] * alpha + checker[..., :3] * (1.0 - alpha)
|
|
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, opaque_input)
|
|
result = run_pipeline(input_path, output_path, target_size=24, steps=0, strip_background=True)
|
|
assert result.source_rgba[..., 3].min() < 0.05
|
|
assert result.output_rgba[..., 3].min() < 0.05
|
|
assert result.output_rgba[..., 3].max() > 0.9
|
|
|
|
|
|
def test_pipeline_can_emit_initialized_output_without_optimizer_steps(tmp_path: Path) -> None:
|
|
source = make_emblem(24, 24)
|
|
fake = fake_pixelize(source, upscale=10, blur_radius=0.6, warp_strength=0.2, warp_detail=5)
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, fake)
|
|
result = run_pipeline(input_path, output_path, steps=0)
|
|
assert output_path.exists()
|
|
assert result.output_rgba.shape[0] == result.inference.target_height
|
|
assert result.output_rgba.shape[1] == result.inference.target_width
|
|
|
|
|
|
def test_fixed_phase_field_pipeline_output_matches_refactor_baseline(tmp_path: Path) -> None:
|
|
source = make_emblem(24, 24)
|
|
fake = fake_pixelize(
|
|
source,
|
|
upscale=8,
|
|
blur_radius=0.4,
|
|
warp_strength=0.15,
|
|
warp_detail=4,
|
|
seed=9,
|
|
)
|
|
input_path = tmp_path / "phase-field-baseline-input.png"
|
|
output_path = tmp_path / "phase-field-baseline-output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, fake)
|
|
result = run_pipeline(
|
|
input_path,
|
|
output_path,
|
|
target_width=24,
|
|
target_height=24,
|
|
enable_candidate_rerank=False,
|
|
steps=0,
|
|
device="cpu",
|
|
)
|
|
|
|
output_hash = hashlib.sha256(np.ascontiguousarray(result.output_rgba).view(np.uint8)).hexdigest()
|
|
assert output_hash == "6ff8f628e1d926e7a5f00d9dceea7decf079d2b070824e18ebe53676311631a0"
|
|
|
|
|
|
def test_candidate_rerank_can_override_low_confidence_inference_pick(monkeypatch) -> None:
|
|
source = np.zeros((3, 3, 4), dtype=np.float32)
|
|
source[0, 0] = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=2, target_height=2, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=3, target_height=3, score=0.909, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=2,
|
|
target_height=2,
|
|
confidence=0.1,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
outputs = {
|
|
(candidate_a.target_width, candidate_a.target_height): np.ones((2, 2, 4), dtype=np.float32),
|
|
(candidate_b.target_width, candidate_b.target_height): source.copy(),
|
|
}
|
|
|
|
class DummyArtifacts:
|
|
def __init__(self, rgba: np.ndarray) -> None:
|
|
self.target_rgba = rgba
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
assert steps == 8
|
|
return DummyArtifacts(outputs[(inference.target_width, inference.target_height)])
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
monkeypatch.setattr(
|
|
"repixelizer.pipeline.source_lattice_consistency_breakdown",
|
|
lambda source_rgba, output_rgba, *, target_width, target_height: {"score": 0.10 if target_width == 2 else 0.01},
|
|
)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_edge_position_error", lambda preview, source_rgba: 0.0)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_stroke_wobble_error", lambda preview, source_rgba: 0.0)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_edge_concentration", lambda rgba: 0.0)
|
|
|
|
selected = _select_candidate(
|
|
source,
|
|
inference,
|
|
steps=32,
|
|
seed=7,
|
|
device="cpu",
|
|
solver_params=SolverHyperParams(candidate_rerank_max_size_delta_ratio=1.0, candidate_rerank_margin=0.0),
|
|
)
|
|
assert selected.target_width == candidate_b.target_width
|
|
assert selected.target_height == candidate_b.target_height
|
|
|
|
|
|
def test_candidate_rerank_preview_respects_requested_zero_steps(monkeypatch) -> None:
|
|
source = np.zeros((2, 2, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=2, target_height=2, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=3, target_height=3, score=0.909, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=2,
|
|
target_height=2,
|
|
confidence=0.1,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
class DummyArtifacts:
|
|
def __init__(self, rgba: np.ndarray) -> None:
|
|
self.target_rgba = rgba
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
assert steps == 0
|
|
return DummyArtifacts(np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32))
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
|
|
_select_candidate(
|
|
source,
|
|
inference,
|
|
steps=0,
|
|
seed=7,
|
|
device="cpu",
|
|
solver_params=SolverHyperParams(candidate_rerank_max_size_delta_ratio=1.0),
|
|
)
|
|
|
|
|
|
def test_candidate_rerank_honors_cooperative_cancellation(monkeypatch) -> None:
|
|
source = np.zeros((2, 2, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=2, target_height=2, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=3, target_height=3, score=0.909, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=2,
|
|
target_height=2,
|
|
confidence=0.1,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
class CancelObserver:
|
|
def __call__(self, event: str, payload: dict[str, object]) -> None:
|
|
del event, payload
|
|
|
|
def check_cancelled(self) -> bool:
|
|
return True
|
|
|
|
with pytest.raises(PipelineCancelled):
|
|
_select_candidate(
|
|
source,
|
|
inference,
|
|
steps=8,
|
|
seed=7,
|
|
device="cpu",
|
|
observer=CancelObserver(),
|
|
)
|
|
|
|
|
|
def test_candidate_rerank_can_override_to_better_size_candidate(monkeypatch) -> None:
|
|
source = np.zeros((16, 16, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=16, target_height=16, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=18, target_height=18, score=0.88, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=16,
|
|
target_height=16,
|
|
confidence=0.0,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
class DummyArtifacts:
|
|
def __init__(self, rgba: np.ndarray) -> None:
|
|
self.target_rgba = rgba
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32)
|
|
return DummyArtifacts(rgba)
|
|
|
|
def fake_support(source_rgba, output_rgba, *, target_width, target_height):
|
|
if target_width == 16:
|
|
return {"score": 0.10}
|
|
return {"score": 0.01}
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
monkeypatch.setattr("repixelizer.pipeline.source_lattice_consistency_breakdown", fake_support)
|
|
|
|
selected = _select_candidate(source, inference, steps=32, seed=7, device="cpu")
|
|
assert selected.target_width == candidate_b.target_width
|
|
assert selected.target_height == candidate_b.target_height
|
|
|
|
|
|
def test_candidate_rerank_rejects_large_size_jump(monkeypatch) -> None:
|
|
source = np.zeros((16, 16, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=16, target_height=16, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=24, target_height=24, score=0.88, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=16,
|
|
target_height=16,
|
|
confidence=0.0,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
class DummyArtifacts:
|
|
def __init__(self, rgba: np.ndarray) -> None:
|
|
self.target_rgba = rgba
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32)
|
|
return DummyArtifacts(rgba)
|
|
|
|
def fake_support(source_rgba, output_rgba, *, target_width, target_height):
|
|
if target_width == 16:
|
|
return {"score": 0.10}
|
|
return {"score": 0.01}
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
monkeypatch.setattr("repixelizer.pipeline.source_lattice_consistency_breakdown", fake_support)
|
|
|
|
selected = _select_candidate(source, inference, steps=32, seed=7, device="cpu")
|
|
assert selected.target_width == candidate_a.target_width
|
|
assert selected.target_height == candidate_a.target_height
|
|
|
|
|
|
def test_select_candidate_can_skip_candidate_rerank(monkeypatch) -> None:
|
|
source = np.zeros((8, 8, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=8, target_height=8, score=0.9, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=10, target_height=10, score=0.89, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=8,
|
|
target_height=8,
|
|
confidence=0.0,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
def fail(*args, **kwargs):
|
|
raise AssertionError("candidate rerank probe should not run when disabled")
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fail)
|
|
selected = _select_candidate(
|
|
source,
|
|
inference,
|
|
steps=32,
|
|
seed=7,
|
|
device="cpu",
|
|
enable_candidate_rerank=False,
|
|
)
|
|
assert selected.target_width == candidate_a.target_width
|
|
assert selected.target_height == candidate_a.target_height
|
|
|
|
|
|
def test_pipeline_runs_phase_field_mode(tmp_path: Path, monkeypatch) -> None:
|
|
source = np.zeros((6, 6, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, source)
|
|
|
|
inference = InferenceResult(
|
|
target_width=3,
|
|
target_height=3,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
|
|
class DummyAnalysis:
|
|
edge_map = np.zeros((6, 6), dtype=np.float32)
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
del source_rgba, steps, seed, device, solver_params
|
|
rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32)
|
|
rgba[..., 1] = 1.0
|
|
rgba[..., 3] = 1.0
|
|
return SolverArtifacts(
|
|
target_rgba=rgba,
|
|
uv_field=np.zeros((inference.target_height, inference.target_width, 2), dtype=np.float32),
|
|
signal_strength=np.zeros((inference.target_height, inference.target_width), dtype=np.float32),
|
|
initial_rgba=rgba.copy(),
|
|
loss_history=[0.0],
|
|
stage_diagnostics={
|
|
"displacements": {},
|
|
"phase_field": {"mean_displacement_px": 0.0},
|
|
},
|
|
)
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.infer_autocorr_lattice", lambda *args, **kwargs: inference)
|
|
monkeypatch.setattr("repixelizer.pipeline.analyze_phase_field_source", lambda *args, **kwargs: DummyAnalysis())
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
|
|
result = run_pipeline(
|
|
input_path,
|
|
output_path,
|
|
steps=8,
|
|
device="cpu",
|
|
)
|
|
|
|
assert output_path.exists()
|
|
assert result.output_rgba.shape[:2] == (3, 3)
|
|
assert result.diagnostics["reconstruction"]["mode"] == "phase-field"
|
|
|
|
|
|
def test_pipeline_fixed_target_skips_search_inference(tmp_path: Path, monkeypatch) -> None:
|
|
source = np.zeros((6, 6, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, source)
|
|
|
|
inference = InferenceResult(
|
|
target_width=3,
|
|
target_height=3,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
|
|
class DummyAnalysis:
|
|
edge_map = np.zeros((6, 6), dtype=np.float32)
|
|
|
|
def fail_infer(*args, **kwargs):
|
|
raise AssertionError("autocorr lattice inference should not run for fixed target")
|
|
|
|
def fake_infer_fixed(*args, **kwargs):
|
|
return inference
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32)
|
|
rgba[..., 3] = 1.0
|
|
return SolverArtifacts(
|
|
target_rgba=rgba,
|
|
uv_field=np.zeros((inference.target_height, inference.target_width, 2), dtype=np.float32),
|
|
signal_strength=np.zeros((inference.target_height, inference.target_width), dtype=np.float32),
|
|
initial_rgba=rgba.copy(),
|
|
loss_history=[],
|
|
)
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.infer_autocorr_lattice", fail_infer)
|
|
monkeypatch.setattr("repixelizer.pipeline.infer_fixed_lattice", fake_infer_fixed)
|
|
monkeypatch.setattr("repixelizer.pipeline.analyze_phase_field_source", lambda *args, **kwargs: DummyAnalysis())
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
|
|
result = run_pipeline(
|
|
input_path,
|
|
output_path,
|
|
target_width=3,
|
|
target_height=3,
|
|
enable_candidate_rerank=False,
|
|
steps=0,
|
|
device="cpu",
|
|
)
|
|
|
|
assert output_path.exists()
|
|
assert result.inference.target_width == 3
|
|
assert result.inference.target_height == 3
|
|
|
|
|
|
def test_pipeline_autocorr_mode_uses_autocorr_inference(tmp_path: Path, monkeypatch) -> None:
|
|
source = np.zeros((6, 6, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "output.png"
|
|
|
|
from repixelizer.io import save_rgba
|
|
|
|
save_rgba(input_path, source)
|
|
|
|
inference = InferenceResult(
|
|
target_width=3,
|
|
target_height=3,
|
|
confidence=1.0,
|
|
top_candidates=[
|
|
InferenceCandidate(target_width=3, target_height=3, score=0.9, breakdown={})
|
|
],
|
|
)
|
|
|
|
class DummyAnalysis:
|
|
edge_map = np.zeros((6, 6), dtype=np.float32)
|
|
|
|
def fail_infer(*args, **kwargs):
|
|
raise AssertionError("alternate lattice inference should not run in autocorr mode")
|
|
|
|
def fake_infer_autocorr(*args, **kwargs):
|
|
assert kwargs["max_target_size"] == 5
|
|
return inference
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32)
|
|
rgba[..., 3] = 1.0
|
|
return SolverArtifacts(
|
|
target_rgba=rgba,
|
|
uv_field=np.zeros((inference.target_height, inference.target_width, 2), dtype=np.float32),
|
|
signal_strength=np.zeros((inference.target_height, inference.target_width), dtype=np.float32),
|
|
initial_rgba=rgba.copy(),
|
|
loss_history=[],
|
|
)
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.infer_autocorr_lattice", fail_infer)
|
|
monkeypatch.setattr("repixelizer.pipeline.infer_autocorr_lattice", fake_infer_autocorr)
|
|
monkeypatch.setattr("repixelizer.pipeline.analyze_phase_field_source", lambda *args, **kwargs: DummyAnalysis())
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
|
|
result = run_pipeline(
|
|
input_path,
|
|
output_path,
|
|
enable_candidate_rerank=False,
|
|
lattice_inference_mode="autocorr",
|
|
max_inferred_target_size=5,
|
|
steps=0,
|
|
device="cpu",
|
|
)
|
|
|
|
assert output_path.exists()
|
|
assert result.inference.target_width == 3
|
|
assert result.inference.target_height == 3
|
|
|
|
|
|
def test_candidate_rerank_can_accept_low_confidence_size_jump_with_strong_support(monkeypatch) -> None:
|
|
source = np.zeros((16, 16, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=16, target_height=16, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=21, target_height=21, score=0.89, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=16,
|
|
target_height=16,
|
|
confidence=0.0,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
class DummyArtifacts:
|
|
def __init__(self, rgba: np.ndarray) -> None:
|
|
self.target_rgba = rgba
|
|
|
|
outputs = {
|
|
candidate_a.target_width: np.zeros((candidate_a.target_height, candidate_a.target_width, 4), dtype=np.float32),
|
|
candidate_b.target_width: np.ones((candidate_b.target_height, candidate_b.target_width, 4), dtype=np.float32),
|
|
}
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
return DummyArtifacts(outputs[inference.target_width])
|
|
|
|
def fake_support(source_rgba, output_rgba, *, target_width, target_height):
|
|
return {"score": 0.22 if target_width == 16 else 0.03}
|
|
|
|
def fake_edge_position(preview, source_rgba):
|
|
return 0.07 if np.mean(preview) < 0.1 else 0.02
|
|
|
|
def fake_wobble(preview, source_rgba):
|
|
return 1.25 if np.mean(preview) < 0.1 else 0.55
|
|
|
|
def fake_concentration(rgba):
|
|
return 0.18 if np.mean(rgba) < 0.1 else 0.31
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
monkeypatch.setattr("repixelizer.pipeline.source_lattice_consistency_breakdown", fake_support)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_edge_position_error", fake_edge_position)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_stroke_wobble_error", fake_wobble)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_edge_concentration", fake_concentration)
|
|
|
|
selected = _select_candidate(source, inference, steps=32, seed=7, device="cpu")
|
|
assert selected.target_width == candidate_b.target_width
|
|
assert selected.target_height == candidate_b.target_height
|
|
|
|
|
|
def test_candidate_rerank_keeps_high_confidence_candidate_without_probe(monkeypatch) -> None:
|
|
source = np.zeros((8, 8, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=8, target_height=8, score=0.9, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=10, target_height=10, score=0.89, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=8,
|
|
target_height=8,
|
|
confidence=0.5,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
def fail(*args, **kwargs):
|
|
raise AssertionError("rerank probe should not run for high-confidence inference")
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fail)
|
|
selected = _select_candidate(source, inference, steps=32, seed=7, device="cpu")
|
|
assert selected.target_width == candidate_a.target_width
|
|
assert selected.target_height == candidate_a.target_height
|
|
|
|
|
|
def test_candidate_rerank_can_prefer_better_line_metrics(monkeypatch) -> None:
|
|
source = np.zeros((4, 4, 4), dtype=np.float32)
|
|
candidate_a = InferenceCandidate(target_width=2, target_height=2, score=0.91, breakdown={})
|
|
candidate_b = InferenceCandidate(target_width=3, target_height=3, score=0.909, breakdown={})
|
|
inference = InferenceResult(
|
|
target_width=2,
|
|
target_height=2,
|
|
confidence=0.0,
|
|
top_candidates=[candidate_a, candidate_b],
|
|
)
|
|
|
|
class DummyArtifacts:
|
|
def __init__(self, rgba: np.ndarray) -> None:
|
|
self.target_rgba = rgba
|
|
|
|
def fake_support(source_rgba, output_rgba, *, target_width, target_height):
|
|
return {"score": 0.02}
|
|
|
|
def fake_edge_position(preview, source_rgba):
|
|
return 0.05
|
|
|
|
def fake_wobble(preview, source_rgba):
|
|
return 1.2 if np.mean(preview) < 0.01 else 0.8
|
|
|
|
def fake_concentration(rgba):
|
|
return 0.15 if np.mean(rgba) < 0.01 else 0.28
|
|
|
|
outputs = {
|
|
candidate_a.target_width: np.zeros((2, 2, 4), dtype=np.float32),
|
|
candidate_b.target_width: np.ones((3, 3, 4), dtype=np.float32),
|
|
}
|
|
|
|
def fake_optimize_phase_field(source_rgba, inference, steps, seed, device, solver_params=None):
|
|
return DummyArtifacts(outputs[inference.target_width])
|
|
|
|
monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field)
|
|
monkeypatch.setattr("repixelizer.pipeline.source_lattice_consistency_breakdown", fake_support)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_edge_position_error", fake_edge_position)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_stroke_wobble_error", fake_wobble)
|
|
monkeypatch.setattr("repixelizer.pipeline.foreground_edge_concentration", fake_concentration)
|
|
|
|
selected = _select_candidate(
|
|
source,
|
|
inference,
|
|
steps=32,
|
|
seed=7,
|
|
device="cpu",
|
|
solver_params=SolverHyperParams(
|
|
candidate_rerank_size_penalty_weight=0.0,
|
|
candidate_rerank_inference_penalty_weight=0.0,
|
|
candidate_rerank_max_size_delta_ratio=1.0,
|
|
candidate_rerank_margin=0.0,
|
|
),
|
|
)
|
|
assert selected.target_width == candidate_b.target_width
|
|
assert selected.target_height == candidate_b.target_height
|