Broaden phase reranking with edge-aware metrics

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]:

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()

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