From e25f0b973084b0cd187fe0c900146ed51578f472 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Wed, 22 Apr 2026 07:22:53 +0100 Subject: [PATCH] Broaden phase reranking with edge-aware metrics --- src/repixelizer/params.py | 6 ++++ src/repixelizer/pipeline.py | 66 ++++++++++++++++++++++++++++++++----- tests/test_pipeline.py | 50 +++++++++++++++++++++++++++- 3 files changed, 113 insertions(+), 9 deletions(-) diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index b211a2c..b1bd69a 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -33,6 +33,12 @@ class SolverHyperParams: alpha_representative_foreground_threshold: float = 0.60 alpha_opaque_threshold: float = 0.95 alpha_transparent_threshold: float = 0.05 + phase_rerank_support_weight: float = 0.45 + phase_rerank_edge_position_weight: float = 0.20 + phase_rerank_wobble_weight: float = 0.20 + phase_rerank_edge_concentration_weight: float = 0.10 + phase_rerank_inference_penalty_weight: float = 0.05 + phase_rerank_confidence_threshold: float = 0.12 phase_rerank_margin: float = 0.004 def to_dict(self) -> dict[str, float]: diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index c5bd413..0a3ebd9 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -19,7 +19,12 @@ from .diagnostics import ( from .discrete import cleanup_pixels from .inference import infer_lattice, inference_to_json from .io import load_rgba, nearest_resize, save_rgba -from .metrics import source_lattice_consistency_breakdown +from .metrics import ( + foreground_edge_concentration, + foreground_edge_position_error, + foreground_stroke_wobble_error, + source_lattice_consistency_breakdown, +) from .params import SolverHyperParams from .palette import load_palette, quantize_rgba, save_palette_report from .preprocess import strip_edge_background @@ -122,13 +127,11 @@ def _select_phase_candidate( solver_params: SolverHyperParams | None = None, ) -> InferenceResult: solver_params = solver_params or SolverHyperParams() - if len(inference.top_candidates) <= 1 or inference.confidence >= 0.02: + if len(inference.top_candidates) <= 1 or inference.confidence >= solver_params.phase_rerank_confidence_threshold: return inference top_score = float(inference.top_candidates[0].score) - baseline_rank: float | None = None - best_rank = float("inf") - best_candidate = inference + candidate_records: list[dict[str, float | InferenceResult]] = [] for candidate in inference.top_candidates[:8]: candidate_inference = InferenceResult( target_width=candidate.target_width, @@ -155,13 +158,60 @@ def _select_phase_candidate( phase_x=candidate.phase_x, phase_y=candidate.phase_y, ) - coherence_penalty = top_score - float(candidate.score) - rank = support["score"] + coherence_penalty * 0.5 + preview = nearest_resize(candidate_artifacts.target_rgba, width=source.shape[1], height=source.shape[0]) + candidate_records.append( + { + "inference": candidate_inference, + "support_score": support["score"], + "edge_position_error": foreground_edge_position_error(preview, source), + "stroke_wobble_error": foreground_stroke_wobble_error(preview, source), + "edge_concentration": foreground_edge_concentration(candidate_artifacts.target_rgba), + "inference_penalty": top_score - float(candidate.score), + } + ) + + if not candidate_records: + return inference + + support_penalty = _normalize_penalty(record["support_score"] for record in candidate_records) + edge_position_penalty = _normalize_penalty(record["edge_position_error"] for record in candidate_records) + wobble_penalty = _normalize_penalty(record["stroke_wobble_error"] for record in candidate_records) + edge_concentration_penalty = _normalize_penalty( + (record["edge_concentration"] for record in candidate_records), + higher_is_better=True, + ) + inference_penalty = _normalize_penalty(record["inference_penalty"] for record in candidate_records) + + baseline_rank: float | None = None + best_rank = float("inf") + best_candidate = inference + for index, record in enumerate(candidate_records): + rank = ( + solver_params.phase_rerank_support_weight * support_penalty[index] + + solver_params.phase_rerank_edge_position_weight * edge_position_penalty[index] + + solver_params.phase_rerank_wobble_weight * wobble_penalty[index] + + solver_params.phase_rerank_edge_concentration_weight * edge_concentration_penalty[index] + + solver_params.phase_rerank_inference_penalty_weight * inference_penalty[index] + ) if baseline_rank is None: baseline_rank = rank if rank < best_rank: best_rank = rank - best_candidate = candidate_inference + best_candidate = record["inference"] if baseline_rank is None or best_rank > baseline_rank - solver_params.phase_rerank_margin: return inference return best_candidate + + +def _normalize_penalty(values, *, higher_is_better: bool = False) -> list[float]: + raw = np.asarray(list(values), dtype=np.float32) + if raw.size == 0: + return [] + lo = float(np.min(raw)) + hi = float(np.max(raw)) + if hi - lo <= 1e-6: + return [0.0] * int(raw.size) + normalized = (raw - lo) / (hi - lo) + if higher_is_better: + normalized = 1.0 - normalized + return normalized.astype(np.float32).tolist() diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 842b287..017c4e4 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -123,7 +123,7 @@ def test_phase_rerank_can_override_low_confidence_inference_pick(monkeypatch) -> target_height=2, phase_x=candidate_a.phase_x, phase_y=candidate_a.phase_y, - confidence=0.0, + confidence=0.1, top_candidates=[candidate_a, candidate_b], ) outputs = { @@ -178,3 +178,51 @@ def test_phase_rerank_can_override_to_better_size_candidate(monkeypatch) -> None selected = _select_phase_candidate(source, inference, analysis=object(), seed=7, device="cpu") assert selected.target_width == candidate_b.target_width assert selected.target_height == candidate_b.target_height + + +def test_phase_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, phase_x=-0.2, phase_y=0.0, score=0.91, breakdown={}) + candidate_b = InferenceCandidate(target_width=2, target_height=2, phase_x=0.2, phase_y=0.0, score=0.909, breakdown={}) + inference = InferenceResult( + target_width=2, + target_height=2, + phase_x=candidate_a.phase_x, + phase_y=candidate_a.phase_y, + confidence=0.0, + top_candidates=[candidate_a, candidate_b], + ) + + class DummyArtifacts: + def __init__(self, rgba: np.ndarray) -> None: + self.target_rgba = rgba + + def fake_support(source_rgba, output_rgba, *, target_width, target_height, phase_x, phase_y): + 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.phase_x: np.zeros((2, 2, 4), dtype=np.float32), + candidate_b.phase_x: np.ones((2, 2, 4), dtype=np.float32), + } + + def fake_optimize_uv_field(source_rgba, inference, analysis, steps, seed, device, solver_params=None): + return DummyArtifacts(outputs[inference.phase_x]) + + monkeypatch.setattr("repixelizer.pipeline.optimize_uv_field", fake_optimize_uv_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_phase_candidate(source, inference, analysis=object(), seed=7, device="cpu") + assert selected.phase_x == candidate_b.phase_x + assert selected.phase_y == candidate_b.phase_y