From 7c727f23dc73fc92f1b7472fab27c15745e54f54 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Wed, 22 Apr 2026 15:28:16 +0100 Subject: [PATCH] Restore thin-feature fidelity in edge pass --- docs/implementation-plan.md | 4 ++ scripts/render_focus_crop.py | 27 +++++----- src/repixelizer/continuous.py | 70 ++++++++++++++----------- src/repixelizer/params.py | 2 + src/repixelizer/source_reference.py | 79 ++++++++++++++++++----------- tests/test_continuous.py | 1 - tests/test_source_reference.py | 15 ++++++ 7 files changed, 123 insertions(+), 75 deletions(-) diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index f8b9f1d..17087be 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -62,6 +62,10 @@ In scope now: - make source reliability edge-aware instead of dispersion-only so high-contrast thin features are not downweighted just because they are locally messy - expand snap/refine candidate generation beyond fixed offset grids by injecting sharp per-cell exemplars, edge peaks, and gradient-guided offsets +Current implementation note: +- raw image edges are the right signal for source-detail targeting; cluster boundaries were too broad and caused thin-feature washout when reused here +- enriched edge candidates are currently gated to high-dispersion edge cells, and `snap` stays on the conservative local grid while `refine` gets the richer candidate pool + Next after that: - rerank low-confidence top lattice candidates with short real solver probes instead of relying only on the cheapest preview - rerun tuning after the edge-aware/source-guided behavior lands diff --git a/scripts/render_focus_crop.py b/scripts/render_focus_crop.py index 16c8a57..f09b0e7 100644 --- a/scripts/render_focus_crop.py +++ b/scripts/render_focus_crop.py @@ -9,9 +9,9 @@ from PIL import Image, ImageDraw from repixelizer.analysis import analyze_source from repixelizer.continuous import ( _build_candidate_positions, + _build_source_detail_reference, _discrete_refine_output, _build_source_reliability, - _cluster_boundary_map, _edge_gradient_maps, _make_patch_offsets, _make_regular_uv, @@ -119,8 +119,8 @@ def _build_focus_states( offsets_t = torch.from_numpy( _make_patch_offsets(height=height, width=width, target_height=inference.target_height, target_width=inference.target_width) ).to(device=resolved_device, dtype=torch.float32) - edge = np.maximum(analysis.edge_map, _cluster_boundary_map(analysis.cluster_map)) - edge_grad_x, edge_grad_y = _edge_gradient_maps(edge) + source_edge = analysis.edge_map.astype(np.float32) + edge_grad_x, edge_grad_y = _edge_gradient_maps(source_edge) source_lattice_reference = build_source_lattice_reference( source, target_width=inference.target_width, @@ -128,14 +128,15 @@ def _build_focus_states( phase_x=inference.phase_x, phase_y=inference.phase_y, alpha_threshold=solver_params.alpha_transparent_threshold, - edge_hint=edge, + edge_hint=source_edge, edge_grad_x_hint=edge_grad_x, edge_grad_y_hint=edge_grad_y, ) + source_detail_reference = _build_source_detail_reference(source, source_lattice_reference, solver_params) initial_patches = _sample_cell_patches(F, source_t, uv0_t, offsets_t) representative_t, _ = _representative_colors(initial_patches, solver_params) representative_t = representative_t.detach() - source_reference_t = torch.from_numpy(premultiply(source_lattice_reference.sharp_rgba)[None, ...]).to( + source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to( device=resolved_device, dtype=torch.float32, ) @@ -144,23 +145,23 @@ def _build_focus_states( dtype=torch.float32, ) source_delta_x_t = ( - torch.from_numpy(source_lattice_reference.delta_x[None, ...]).to(device=resolved_device, dtype=torch.float32) - if source_lattice_reference.delta_x is not None + torch.from_numpy((premultiply(source_detail_reference)[:, 1:, :] - premultiply(source_detail_reference)[:, :-1, :])[None, ...]).to(device=resolved_device, dtype=torch.float32) + if inference.target_width > 1 else None ) source_delta_y_t = ( - torch.from_numpy(source_lattice_reference.delta_y[None, ...]).to(device=resolved_device, dtype=torch.float32) - if source_lattice_reference.delta_y is not None + torch.from_numpy((premultiply(source_detail_reference)[1:, :, :] - premultiply(source_detail_reference)[:-1, :, :])[None, ...]).to(device=resolved_device, dtype=torch.float32) + if inference.target_height > 1 else None ) source_delta_diag_t = ( - torch.from_numpy(source_lattice_reference.delta_diag[None, ...]).to(device=resolved_device, dtype=torch.float32) - if source_lattice_reference.delta_diag is not None + torch.from_numpy((premultiply(source_detail_reference)[1:, 1:, :] - premultiply(source_detail_reference)[:-1, :-1, :])[None, ...]).to(device=resolved_device, dtype=torch.float32) + if inference.target_height > 1 and inference.target_width > 1 else None ) source_delta_anti_t = ( - torch.from_numpy(source_lattice_reference.delta_anti[None, ...]).to(device=resolved_device, dtype=torch.float32) - if source_lattice_reference.delta_anti is not None + torch.from_numpy((premultiply(source_detail_reference)[1:, :-1, :] - premultiply(source_detail_reference)[:-1, 1:, :])[None, ...]).to(device=resolved_device, dtype=torch.float32) + if inference.target_height > 1 and inference.target_width > 1 else None ) diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py index bcd4366..6cb694b 100644 --- a/src/repixelizer/continuous.py +++ b/src/repixelizer/continuous.py @@ -377,11 +377,10 @@ def _build_candidate_positions( candidates_x.append(torch.round(xs[..., None] + offset_x)) candidates_y.append(torch.round(ys[..., None] + offset_y)) - sharp_x = torch.from_numpy(source_reference.sharp_x).to(device=uv.device, dtype=uv.dtype) - sharp_y = torch.from_numpy(source_reference.sharp_y).to(device=uv.device, dtype=uv.dtype) edge_peak_x = torch.from_numpy(source_reference.edge_peak_x).to(device=uv.device, dtype=uv.dtype) edge_peak_y = torch.from_numpy(source_reference.edge_peak_y).to(device=uv.device, dtype=uv.dtype) edge_strength = torch.from_numpy(source_reference.edge_strength).to(device=uv.device, dtype=uv.dtype) + cell_dispersion = torch.from_numpy(source_reference.cell_dispersion).to(device=uv.device, dtype=uv.dtype) edge_grad_x = torch.from_numpy(source_reference.edge_grad_x).to(device=uv.device, dtype=uv.dtype) edge_grad_y = torch.from_numpy(source_reference.edge_grad_y).to(device=uv.device, dtype=uv.dtype) grad_norm = (edge_grad_x.square() + edge_grad_y.square()).sqrt() @@ -389,12 +388,12 @@ def _build_candidate_positions( unit_grad_x = torch.where(valid_grad, edge_grad_x / grad_norm.clamp_min(1e-4), torch.zeros_like(edge_grad_x)) unit_grad_y = torch.where(valid_grad, edge_grad_y / grad_norm.clamp_min(1e-4), torch.zeros_like(edge_grad_y)) - candidates_x.append(sharp_x[..., None]) - candidates_y.append(sharp_y[..., None]) - candidates_x.append(edge_peak_x[..., None]) - candidates_y.append(edge_peak_y[..., None]) + dispersion_ratio = cell_dispersion / max(float(source_reference.dispersion), 1e-4) + detail_gate = (edge_strength > solver_params.guided_candidate_edge_threshold) & (dispersion_ratio > 1.15) + edge_gate = detail_gate & valid_grad + candidates_x.append(torch.where(detail_gate, edge_peak_x, torch.round(xs))[..., None]) + candidates_y.append(torch.where(detail_gate, edge_peak_y, torch.round(ys))[..., None]) - edge_gate = (edge_strength > solver_params.guided_candidate_edge_threshold) & valid_grad guided_scales = ( -solver_params.guided_candidate_outer_scale, -solver_params.guided_candidate_inner_scale, @@ -436,6 +435,16 @@ def _reference_match_energy( return (representative_energy * representative_mix + source_energy * source_mix) / normalization +def _build_source_detail_reference(source_rgba: np.ndarray, source_reference, solver_params: SolverHyperParams) -> np.ndarray: + edge_rgba = source_rgba[source_reference.edge_peak_y, source_reference.edge_peak_x] + threshold = float(np.clip(solver_params.source_edge_detail_threshold, 0.0, 0.95)) + normalized_edge = np.clip((source_reference.edge_strength - threshold) / max(1e-4, 1.0 - threshold), 0.0, 1.0) + edge_mix = (normalized_edge * solver_params.source_edge_detail_mix)[..., None] + detail_rgba = source_reference.sharp_rgba * (1.0 - edge_mix) + edge_rgba * edge_mix + detail_rgba[..., 3] = np.maximum(source_reference.sharp_rgba[..., 3], detail_rgba[..., 3]) + return detail_rgba.astype(np.float32) + + def _edge_reliability(source_reliability, axis: str): if axis == "x": return (source_reliability[:, 1:] + source_reliability[:, :-1]) * 0.5 @@ -858,17 +867,14 @@ def _snap_output_to_source_pixels( output_height = representative.shape[0] output_width = representative.shape[1] - candidate_x, candidate_y, _, _ = _build_candidate_positions( - torch, - uv, - source_reference_np, - solver_params, - width=width, - height=height, - cell_x=cell_x, - cell_y=cell_y, - base_fraction_values=np.asarray([-0.42, -0.21, 0.0, 0.21, 0.42], dtype=np.float32), - ) + 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] representative_match = (candidate_colors - representative[..., None, :]).abs().mean(dim=-1) @@ -1369,9 +1375,10 @@ def optimize_uv_field( 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_grad_x, edge_grad_y = _edge_gradient_maps(edge) - edge_t = torch.from_numpy(edge[None, None, ...]).to(device=device, dtype=torch.float32) + source_edge = analysis.edge_map.astype(np.float32) + edge_grad_x, edge_grad_y = _edge_gradient_maps(source_edge) + guide_edge = np.maximum(source_edge, _cluster_boundary_map(analysis.cluster_map)) + edge_t = torch.from_numpy(guide_edge[None, None, ...]).to(device=device, dtype=torch.float32) uv0 = _make_regular_uv( height=height, width=width, @@ -1392,14 +1399,15 @@ def optimize_uv_field( phase_x=inference.phase_x, phase_y=inference.phase_y, alpha_threshold=solver_params.alpha_transparent_threshold, - edge_hint=edge, + edge_hint=source_edge, edge_grad_x_hint=edge_grad_x, edge_grad_y_hint=edge_grad_y, ) + source_detail_reference = _build_source_detail_reference(rgba, source_lattice_reference, solver_params) 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() - initial_source_reference_t = torch.from_numpy(premultiply(source_lattice_reference.sharp_rgba)[None, ...]).to( + initial_source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to( device=device, dtype=torch.float32, ) @@ -1408,23 +1416,23 @@ def optimize_uv_field( dtype=torch.float32, ) source_delta_x_t = ( - torch.from_numpy(source_lattice_reference.delta_x[None, ...]).to(device=device, dtype=torch.float32) - if source_lattice_reference.delta_x is not None + torch.from_numpy((premultiply(source_detail_reference)[:, 1:, :] - premultiply(source_detail_reference)[:, :-1, :])[None, ...]).to(device=device, dtype=torch.float32) + if inference.target_width > 1 else None ) source_delta_y_t = ( - torch.from_numpy(source_lattice_reference.delta_y[None, ...]).to(device=device, dtype=torch.float32) - if source_lattice_reference.delta_y is not None + torch.from_numpy((premultiply(source_detail_reference)[1:, :, :] - premultiply(source_detail_reference)[:-1, :, :])[None, ...]).to(device=device, dtype=torch.float32) + if inference.target_height > 1 else None ) source_delta_diag_t = ( - torch.from_numpy(source_lattice_reference.delta_diag[None, ...]).to(device=device, dtype=torch.float32) - if source_lattice_reference.delta_diag is not None + torch.from_numpy((premultiply(source_detail_reference)[1:, 1:, :] - premultiply(source_detail_reference)[:-1, :-1, :])[None, ...]).to(device=device, dtype=torch.float32) + if inference.target_height > 1 and inference.target_width > 1 else None ) source_delta_anti_t = ( - torch.from_numpy(source_lattice_reference.delta_anti[None, ...]).to(device=device, dtype=torch.float32) - if source_lattice_reference.delta_anti is not None + torch.from_numpy((premultiply(source_detail_reference)[1:, :-1, :] - premultiply(source_detail_reference)[:-1, 1:, :])[None, ...]).to(device=device, dtype=torch.float32) + if inference.target_height > 1 and inference.target_width > 1 else None ) snap_rgba = _snap_output_to_source_pixels( diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index 7c09859..5fc0083 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -9,6 +9,8 @@ class SolverHyperParams: source_edge_reliability_gain: float = 1.6 source_edge_reliability_floor: float = 0.25 source_edge_alpha_floor: float = 0.45 + source_edge_detail_mix: float = 0.65 + source_edge_detail_threshold: float = 0.25 boundary_probe_scale: float = 0.22 boundary_signed_weight: float = 0.55 boundary_direction_weight: float = 0.30 diff --git a/src/repixelizer/source_reference.py b/src/repixelizer/source_reference.py index 0fcdbf1..64c98af 100644 --- a/src/repixelizer/source_reference.py +++ b/src/repixelizer/source_reference.py @@ -6,6 +6,21 @@ from .io import premultiply, unpremultiply from .types import SourceLatticeReference +def _default_edge_hint(source_rgba: np.ndarray) -> np.ndarray: + rgb = source_rgba[..., :3] + alpha = source_rgba[..., 3] + lum = rgb[..., 0] * 0.2126 + rgb[..., 1] * 0.7152 + rgb[..., 2] * 0.0722 + dx = np.zeros_like(lum, dtype=np.float32) + dy = np.zeros_like(lum, dtype=np.float32) + dx[:, 1:] = np.abs(lum[:, 1:] - lum[:, :-1]) + np.abs(alpha[:, 1:] - alpha[:, :-1]) + dy[1:, :] = np.abs(lum[1:, :] - lum[:-1, :]) + np.abs(alpha[1:, :] - alpha[:-1, :]) + edge = np.sqrt(dx * dx + dy * dy, dtype=np.float32) + max_edge = float(np.max(edge)) + if max_edge > 0.0: + edge /= max_edge + return edge.astype(np.float32) + + def lattice_indices( *, height: int, @@ -48,6 +63,7 @@ def build_source_lattice_reference( ) -> SourceLatticeReference: height = source_rgba.shape[0] width = source_rgba.shape[1] + edge_hint = edge_hint.astype(np.float32) if edge_hint is not None else _default_edge_hint(source_rgba) indices = lattice_indices( height=height, width=width, @@ -61,6 +77,7 @@ def build_source_lattice_reference( flat_idx = indices.reshape(-1) flat_premul = premul.reshape(-1, premul.shape[-1]) flat_alpha = source_rgba.reshape(-1, source_rgba.shape[-1])[:, 3] + flat_edge = edge_hint.reshape(-1).astype(np.float32) counts = np.bincount(flat_idx, minlength=cell_count).astype(np.float32) safe_counts = np.maximum(counts, 1.0) @@ -103,6 +120,37 @@ def build_source_lattice_reference( flat_x = pixel_x.reshape(-1) flat_y = pixel_y.reshape(-1) exemplar_cost = pixel_diff + np.maximum(0.0, per_pixel_mean[:, 3] - flat_premul[:, 3]) * 1e-3 + + edge_score = flat_edge + flat_alpha.astype(np.float32) * 0.05 + edge_order = np.lexsort((-edge_score.astype(np.float32), flat_idx.astype(np.int64))) + edge_peak_x_flat = sharp_x_flat.copy() + edge_peak_y_flat = sharp_y_flat.copy() + edge_strength_flat = np.zeros(cell_count, dtype=np.float32) + edge_grad_x_flat = np.zeros(cell_count, dtype=np.float32) + edge_grad_y_flat = np.zeros(cell_count, dtype=np.float32) + best_edge_order = np.asarray([], dtype=np.int64) + best_edge_cells = np.asarray([], dtype=np.int64) + if edge_order.size > 0: + sorted_edge_cells = flat_idx[edge_order] + first_edge_positions = np.concatenate(([0], np.flatnonzero(np.diff(sorted_edge_cells)) + 1)) + best_edge_order = edge_order[first_edge_positions] + best_edge_cells = flat_idx[best_edge_order] + edge_peak_x_flat[best_edge_cells] = flat_x[best_edge_order] + edge_peak_y_flat[best_edge_cells] = flat_y[best_edge_order] + edge_strength_flat[best_edge_cells] = flat_edge[best_edge_order] + if edge_grad_x_hint is not None: + edge_grad_x_flat[best_edge_cells] = edge_grad_x_hint.reshape(-1)[best_edge_order].astype(np.float32) + if edge_grad_y_hint is not None: + edge_grad_y_flat[best_edge_cells] = edge_grad_y_hint.reshape(-1)[best_edge_order].astype(np.float32) + + np.maximum.at(edge_strength_flat, flat_idx, flat_edge) + if edge_grad_x_hint is not None or edge_grad_y_hint is not None: + edge_peak_index_flat = edge_peak_y_flat * width + edge_peak_x_flat + if edge_grad_x_hint is not None: + edge_grad_x_flat = edge_grad_x_hint.reshape(-1)[edge_peak_index_flat].astype(np.float32) + if edge_grad_y_hint is not None: + edge_grad_y_flat = edge_grad_y_hint.reshape(-1)[edge_peak_index_flat].astype(np.float32) + order = np.lexsort((exemplar_cost.astype(np.float32), flat_idx.astype(np.int64))) sharp_premul_flat = mean_premul_flat.copy() if order.size > 0: @@ -113,6 +161,7 @@ def build_source_lattice_reference( sharp_premul_flat[best_cells] = flat_premul[best_order] sharp_x_flat[best_cells] = flat_x[best_order] sharp_y_flat[best_cells] = flat_y[best_order] + sharp_premul = sharp_premul_flat.reshape(target_height, target_width, premul.shape[-1]) sharp_rgba = unpremultiply(sharp_premul) @@ -120,36 +169,6 @@ def build_source_lattice_reference( np.maximum.at(cell_alpha_max, flat_idx, flat_alpha) expected_pixels_per_cell = float(source_rgba.shape[0] * source_rgba.shape[1]) / float(cell_count) cell_support = (counts / max(expected_pixels_per_cell, 1.0)).reshape(target_height, target_width) - edge_peak_x_flat = sharp_x_flat.copy() - edge_peak_y_flat = sharp_y_flat.copy() - edge_strength_flat = np.zeros(cell_count, dtype=np.float32) - edge_grad_x_flat = np.zeros(cell_count, dtype=np.float32) - edge_grad_y_flat = np.zeros(cell_count, dtype=np.float32) - - if edge_hint is not None: - flat_edge = edge_hint.reshape(-1).astype(np.float32) - edge_score = flat_edge + flat_alpha.astype(np.float32) * 0.05 - edge_order = np.lexsort((-edge_score.astype(np.float32), flat_idx.astype(np.int64))) - if edge_order.size > 0: - sorted_edge_cells = flat_idx[edge_order] - first_edge_positions = np.concatenate(([0], np.flatnonzero(np.diff(sorted_edge_cells)) + 1)) - best_edge_order = edge_order[first_edge_positions] - best_edge_cells = flat_idx[best_edge_order] - edge_peak_x_flat[best_edge_cells] = flat_x[best_edge_order] - edge_peak_y_flat[best_edge_cells] = flat_y[best_edge_order] - edge_strength_flat[best_edge_cells] = flat_edge[best_edge_order] - if edge_grad_x_hint is not None: - edge_grad_x_flat[best_edge_cells] = edge_grad_x_hint.reshape(-1)[best_edge_order].astype(np.float32) - if edge_grad_y_hint is not None: - edge_grad_y_flat[best_edge_cells] = edge_grad_y_hint.reshape(-1)[best_edge_order].astype(np.float32) - - np.maximum.at(edge_strength_flat, flat_idx, flat_edge) - if edge_grad_x_hint is not None or edge_grad_y_hint is not None: - edge_peak_index_flat = edge_peak_y_flat * width + edge_peak_x_flat - if edge_grad_x_hint is not None: - edge_grad_x_flat = edge_grad_x_hint.reshape(-1)[edge_peak_index_flat].astype(np.float32) - if edge_grad_y_hint is not None: - edge_grad_y_flat = edge_grad_y_hint.reshape(-1)[edge_peak_index_flat].astype(np.float32) delta_x, delta_y, delta_diag, delta_anti = _reference_deltas(sharp_premul) return SourceLatticeReference( diff --git a/tests/test_continuous.py b/tests/test_continuous.py index 42ddd1b..dff3b08 100644 --- a/tests/test_continuous.py +++ b/tests/test_continuous.py @@ -121,7 +121,6 @@ def test_candidate_positions_include_sharp_and_edge_guided_samples() -> None: for y, x in zip(candidate_y[0, 0].cpu().tolist(), candidate_x[0, 0].cpu().tolist()) } - assert (int(reference.sharp_y[0, 0]), int(reference.sharp_x[0, 0])) in cell_candidates assert (int(reference.edge_peak_y[0, 0]), int(reference.edge_peak_x[0, 0])) in cell_candidates assert len(cell_candidates) > 2 diff --git a/tests/test_source_reference.py b/tests/test_source_reference.py index f39e81e..8068b9b 100644 --- a/tests/test_source_reference.py +++ b/tests/test_source_reference.py @@ -63,3 +63,18 @@ def test_source_lattice_reference_tracks_edge_peaks_and_gradients() -> None: assert np.all(reference.edge_strength[:, 0] == 1.0) assert np.all(reference.edge_grad_x[:, 0] == 1.0) assert np.all(reference.edge_grad_y[:, 0] == 0.0) +def test_source_lattice_reference_computes_default_edge_metadata() -> None: + source = np.zeros((6, 6, 4), dtype=np.float32) + source[..., 3] = 1.0 + source[:, 2, :3] = 1.0 + + reference = build_source_lattice_reference( + source, + target_width=2, + target_height=2, + phase_x=0.0, + phase_y=0.0, + ) + + assert np.all(reference.edge_strength[:, 0] > 0.0) + assert np.all(reference.edge_peak_x[:, 0] == 2)