Restore thin-feature fidelity in edge pass

This commit is contained in:
Meta Vens 2026-04-22 15:28:16 +01:00
commit 7c727f23dc
7 changed files with 123 additions and 75 deletions

View file

@ -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

View file

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

View file

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

View file

@ -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

View file

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

View file

@ -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

View file

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