diff --git a/src/repixelizer/inference.py b/src/repixelizer/inference.py index f60170e..2a47538 100644 --- a/src/repixelizer/inference.py +++ b/src/repixelizer/inference.py @@ -61,6 +61,52 @@ def _candidate_dims( return dims +def _strong_spacing_size_window( + hinted_sizes: list[int], + *, + spacing_x: tuple[float | None, float], + spacing_y: tuple[float | None, float], + prior_reliability: float, +) -> tuple[int, int] | None: + if not hinted_sizes: + return None + axis_confidences = [float(confidence) for _spacing, confidence in (spacing_x, spacing_y) if confidence > 1e-6] + if not axis_confidences: + return None + center = int(round(float(np.median(np.asarray(hinted_sizes, dtype=np.float32))))) + spread = max(abs(int(size) - center) for size in hinted_sizes) + mean_confidence = float(np.mean(axis_confidences)) + max_confidence = float(max(axis_confidences)) + if prior_reliability < 0.45 or max_confidence < 0.55 or spread > 2: + return None + radius = 0 if spread == 0 and mean_confidence >= 0.72 and prior_reliability >= 0.68 else 1 + return center, radius + + +def _resolve_candidate_dims_from_spacing( + width: int, + height: int, + target_size: int | None, + *, + hinted_sizes: list[int], + spacing_x: tuple[float | None, float], + spacing_y: tuple[float | None, float], + prior_reliability: float, +) -> list[tuple[int, int]]: + dims = _candidate_dims(width, height, target_size, hinted_sizes=hinted_sizes) + size_window = _strong_spacing_size_window( + hinted_sizes, + spacing_x=spacing_x, + spacing_y=spacing_y, + prior_reliability=prior_reliability, + ) + if size_window is None: + return dims + center, radius = size_window + narrowed = [dim for dim in dims if abs(max(dim) - center) <= radius] + return narrowed or dims + + def _estimate_cell_size(profile: np.ndarray) -> float: centered = profile.astype(np.float32) - float(np.mean(profile)) max_lag = min(64, max(4, profile.shape[0] // 3)) @@ -192,7 +238,12 @@ def _combine_axis_priors(axis_priors: list[tuple[float, float]]) -> tuple[float, return shared_prior, reliability -def _estimate_lattice_prior_details(rgba: np.ndarray) -> tuple[float, float, float]: +def _estimate_lattice_prior_details( + rgba: np.ndarray, + *, + spacing_x: tuple[float | None, float] | None = None, + spacing_y: tuple[float | None, float] | None = None, +) -> tuple[float, float, float]: alpha = rgba[..., 3] luminance = rgba[..., 0] * 0.2126 + rgba[..., 1] * 0.7152 + rgba[..., 2] * 0.0722 dx = np.zeros_like(luminance) @@ -201,7 +252,10 @@ def _estimate_lattice_prior_details(rgba: np.ndarray) -> tuple[float, float, flo 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) - spacing_x, spacing_y = _estimate_lattice_spacing(rgba) + if spacing_x is None or spacing_y is None: + inferred_spacing_x, inferred_spacing_y = _estimate_lattice_spacing(rgba) + spacing_x = inferred_spacing_x if spacing_x is None else spacing_x + spacing_y = inferred_spacing_y if spacing_y is None else spacing_y autocorr_x = _estimate_cell_size(profile_x) autocorr_y = _estimate_cell_size(profile_y) axis_x = _axis_prior_from_estimates(spacing_x[0], spacing_x[1], autocorr_x) @@ -483,9 +537,21 @@ def infer_lattice( height, width = rgba.shape[:2] spacing_x, spacing_y = _estimate_lattice_spacing(rgba) hinted_sizes = _hint_target_sizes_from_spacing(width, height, spacing_x, spacing_y) - prior_cell_x, prior_cell_y, prior_reliability = _estimate_lattice_prior_details(rgba) + prior_cell_x, prior_cell_y, prior_reliability = _estimate_lattice_prior_details( + rgba, + spacing_x=spacing_x, + spacing_y=spacing_y, + ) phase_values = np.linspace(-0.4, 0.4, num=5, dtype=np.float32) - candidate_dims = _candidate_dims(width, height, target_size, hinted_sizes=hinted_sizes) + candidate_dims = _resolve_candidate_dims_from_spacing( + width, + height, + target_size, + hinted_sizes=hinted_sizes, + spacing_x=spacing_x, + spacing_y=spacing_y, + prior_reliability=prior_reliability, + ) phase_sample_count = int(phase_values.size * phase_values.size) emit_observer( diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index 5255b17..de774e7 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -439,7 +439,9 @@ def _select_phase_candidate_with_reconstruction( phase_y=float(candidate.phase_y), total_steps=preview_steps, final_loss=( - None if not candidate_artifacts.loss_history else float(candidate_artifacts.loss_history[-1]) + None + if not getattr(candidate_artifacts, "loss_history", None) + else float(candidate_artifacts.loss_history[-1]) ), ) support = source_lattice_consistency_breakdown( diff --git a/tests/test_gui.py b/tests/test_gui.py index 239e728..fcdf75e 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -100,7 +100,7 @@ def test_infer_lattice_emits_search_progress_events(monkeypatch) -> None: monkeypatch.setattr(inference_module, "_require_torch", lambda: (_FakeTorch(), object())) monkeypatch.setattr(inference_module, "_estimate_lattice_spacing", lambda rgba: ((None, 0.0), (None, 0.0))) monkeypatch.setattr(inference_module, "_hint_target_sizes_from_spacing", lambda width, height, spacing_x, spacing_y: []) - monkeypatch.setattr(inference_module, "_estimate_lattice_prior_details", lambda rgba: (4.0, 4.0, 0.5)) + monkeypatch.setattr(inference_module, "_estimate_lattice_prior_details", lambda rgba, **kwargs: (4.0, 4.0, 0.5)) monkeypatch.setattr( inference_module, "_candidate_dims", diff --git a/tests/test_inference.py b/tests/test_inference.py index 6934605..2ca99cc 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -5,6 +5,7 @@ from repixelizer.inference import ( _candidate_dims, _combine_axis_priors, _hint_target_sizes_from_spacing, + _resolve_candidate_dims_from_spacing, _top_candidates_by_size, infer_fixed_lattice, infer_lattice, @@ -37,6 +38,35 @@ def test_spacing_hints_can_add_dense_candidates_around_true_size() -> None: assert 128 in sizes +def test_strong_spacing_signal_can_collapse_candidate_search_to_single_size() -> None: + dims = _resolve_candidate_dims_from_spacing( + 1024, + 1024, + None, + hinted_sizes=[128], + spacing_x=(8.0, 0.82), + spacing_y=(8.02, 0.78), + prior_reliability=0.74, + ) + sizes = {width for width, _ in dims} + assert sizes == {128} + + +def test_weak_spacing_signal_keeps_broad_candidate_search() -> None: + dims = _resolve_candidate_dims_from_spacing( + 1024, + 1024, + None, + hinted_sizes=[128], + spacing_x=(8.0, 0.18), + spacing_y=(8.1, 0.12), + prior_reliability=0.24, + ) + sizes = {width for width, _ in dims} + assert 128 in sizes + assert len(sizes) > 10 + + def test_top_candidates_are_diversified_by_size() -> None: candidates = [ InferenceCandidate(target_width=113, target_height=113, phase_x=0.0, phase_y=0.0, score=0.90, breakdown={}),