diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index d521412..f8b8a8e 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -33,6 +33,38 @@ The most provisional areas are: ## Near-term priorities +## High-value regression cases + +### AI badge emblem with baked checkerboard background + +This locally generated AI badge has become a useful stress case because it fails in a very specific, repeatable way. + +Why it is useful: +- the global target size is broadly plausible, so the failure is not just "picked the wrong lattice" +- the lower sword-tip region contains sharp single-pixel and near-single-pixel features at multiple orientations +- those features are exactly the kind of local adjacency pattern Repixelizer is supposed to preserve + +Current failure modes: +- the right side of the sword tip develops a wobbling outline instead of a smooth taper +- thin oblique outline features around the tip get collapsed into mushy transitions +- local cell assignments near the tip look inconsistent even when the overall badge size feels right + +Important note: +- the checkerboard is baked into the source image, not transparency +- when projected onto a `122x122` output grid it aliases into a visible `2-3-2` cadence +- that `2-3-2` background rhythm is expected from resampling the background pattern and should not be mistaken for proof that the inferred lattice itself is varying in size +- it is still a problem, because the background is treated as real source structure and can contaminate local scoring near the blade + +What future work should improve here: +- better weighting of thin local outline/motif preservation relative to high-contrast cell transitions +- some form of coordinated local relaxation so nearby cells can move together toward a globally better contour +- optional background suppression or de-weighting when a baked checkerboard is clearly not semantic content + +Repository fixtures: +- `tests/fixtures/real/ai-badge-cleaned.png` is the manually cleaned transparent version of the same emblem +- `tests/fixtures/real/ai-badge-cleaned.json` records why it matters +- use the cleaned fixture when isolating lattice and contour failures from background-removal failures + ### 1. Build a real benchmark corpus Add curated real-world cases, not just synthetic ones: diff --git a/src/repixelizer/cli.py b/src/repixelizer/cli.py index 61c82a2..5075674 100644 --- a/src/repixelizer/cli.py +++ b/src/repixelizer/cli.py @@ -29,6 +29,11 @@ def build_parser() -> argparse.ArgumentParser: target.add_argument("--seed", type=int, default=7, help="Random seed") target.add_argument("--steps", type=int, default=200, help="Number of optimizer steps") target.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"), help="Torch device") + target.add_argument( + "--strip-background", + action="store_true", + help="Remove light neutral edge-connected backgrounds such as fake transparency checkerboards", + ) run_parser = subparsers.add_parser("run", help="Run the optimizer.") add_shared_arguments(run_parser) @@ -120,6 +125,7 @@ def main(argv: list[str] | None = None) -> int: seed=args.seed, steps=args.steps, device=args.device, + strip_background=args.strip_background, ) return 0 if command == "benchmark": @@ -166,5 +172,6 @@ def main(argv: list[str] | None = None) -> int: seed=args.seed, steps=args.steps, device=args.device, + strip_background=args.strip_background, ) return 0 diff --git a/src/repixelizer/compare.py b/src/repixelizer/compare.py index c2b4042..689fcee 100644 --- a/src/repixelizer/compare.py +++ b/src/repixelizer/compare.py @@ -8,10 +8,12 @@ from PIL import Image, ImageDraw from .baselines import error_diffusion_baseline, naive_resize_baseline from .diagnostics import write_compare_csv, write_json -from .io import load_rgba, nearest_resize, save_rgba +from .io import nearest_resize, save_rgba from .metrics import ( coherence_breakdown, foreground_adjacency_error, + foreground_edge_concentration, + foreground_edge_position_error, foreground_motif_error, foreground_reconstruction_error, reconstruction_error, @@ -58,6 +60,7 @@ def run_compare( seed: int = 7, steps: int = 200, device: str = "auto", + strip_background: bool = False, ) -> dict[str, Any]: diagnostics_path = Path(diagnostics_dir) if diagnostics_dir else Path(output_path).with_suffix("") diagnostics_path.mkdir(parents=True, exist_ok=True) @@ -71,8 +74,9 @@ def run_compare( seed=seed, steps=steps, device=device, + strip_background=strip_background, ) - source = load_rgba(input_path) + source = result.source_rgba palette = load_palette(palette_path) if palette_path else None naive = naive_resize_baseline(source, width=result.inference.target_width, height=result.inference.target_height) diffusion = error_diffusion_baseline( @@ -101,6 +105,8 @@ def run_compare( "color_chatter": coherence["color_chatter"], "reconstruction_error": reconstruction_error(preview, source), "foreground_reconstruction_error": foreground_reconstruction_error(preview, source), + "foreground_edge_concentration": foreground_edge_concentration(image), + "foreground_edge_position_error": foreground_edge_position_error(preview, source), "foreground_adjacency_error": foreground_adjacency_error(preview, source), "foreground_motif_error": foreground_motif_error(preview, source), } diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py index 23798e5..0b1f3b3 100644 --- a/src/repixelizer/continuous.py +++ b/src/repixelizer/continuous.py @@ -180,6 +180,244 @@ def _motif_pattern_loss(representative, reference): return (motif_blocks(representative) - motif_blocks(reference)).abs().mean() +def _normalized_motif_blocks(torch, blocks): + centered = blocks - blocks.mean(dim=-2, keepdim=True) + scale = centered.abs().mean(dim=(-2, -1), keepdim=True).clamp_min(1e-4) + return centered / scale + + +def _pairwise_candidate_energy( + candidate_colors, + context_colors, + desired_delta_x, + desired_delta_y, + desired_delta_diag, + desired_delta_anti, + *, + orthogonal_weight: float, + diagonal_weight: float, +): + output_height = context_colors.shape[0] + output_width = context_colors.shape[1] + energy = candidate_colors.new_zeros(candidate_colors.shape[:3]) + + if output_width > 1: + left_context = context_colors.new_zeros(context_colors.shape) + left_context[:, 1:, :] = context_colors[:, :-1, :] + left_mask = context_colors.new_zeros((output_height, output_width, 1)) + left_mask[:, 1:, :] = 1.0 + left_desired = context_colors.new_zeros(context_colors.shape) + left_desired[:, 1:, :] = desired_delta_x + left_error = ((candidate_colors - left_context[..., None, :]) - left_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + left_error * left_mask * orthogonal_weight + + right_context = context_colors.new_zeros(context_colors.shape) + right_context[:, :-1, :] = context_colors[:, 1:, :] + right_mask = context_colors.new_zeros((output_height, output_width, 1)) + right_mask[:, :-1, :] = 1.0 + right_desired = context_colors.new_zeros(context_colors.shape) + right_desired[:, :-1, :] = desired_delta_x + right_error = ((right_context[..., None, :] - candidate_colors) - right_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + right_error * right_mask * orthogonal_weight + + if output_height > 1: + up_context = context_colors.new_zeros(context_colors.shape) + up_context[1:, :, :] = context_colors[:-1, :, :] + up_mask = context_colors.new_zeros((output_height, output_width, 1)) + up_mask[1:, :, :] = 1.0 + up_desired = context_colors.new_zeros(context_colors.shape) + up_desired[1:, :, :] = desired_delta_y + up_error = ((candidate_colors - up_context[..., None, :]) - up_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + up_error * up_mask * orthogonal_weight + + down_context = context_colors.new_zeros(context_colors.shape) + down_context[:-1, :, :] = context_colors[1:, :, :] + down_mask = context_colors.new_zeros((output_height, output_width, 1)) + down_mask[:-1, :, :] = 1.0 + down_desired = context_colors.new_zeros(context_colors.shape) + down_desired[:-1, :, :] = desired_delta_y + down_error = ((down_context[..., None, :] - candidate_colors) - down_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + down_error * down_mask * orthogonal_weight + + if output_height > 1 and output_width > 1: + diag_context = context_colors.new_zeros(context_colors.shape) + diag_context[1:, 1:, :] = context_colors[:-1, :-1, :] + diag_mask = context_colors.new_zeros((output_height, output_width, 1)) + diag_mask[1:, 1:, :] = 1.0 + diag_desired = context_colors.new_zeros(context_colors.shape) + diag_desired[1:, 1:, :] = desired_delta_diag + diag_error = ((candidate_colors - diag_context[..., None, :]) - diag_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + diag_error * diag_mask * diagonal_weight + + anti_context = context_colors.new_zeros(context_colors.shape) + anti_context[1:, :-1, :] = context_colors[:-1, 1:, :] + anti_mask = context_colors.new_zeros((output_height, output_width, 1)) + anti_mask[1:, :-1, :] = 1.0 + anti_desired = context_colors.new_zeros(context_colors.shape) + anti_desired[1:, :-1, :] = desired_delta_anti + anti_error = ((candidate_colors - anti_context[..., None, :]) - anti_desired[..., None, :]).abs().mean(dim=-1) + energy = energy + anti_error * anti_mask * diagonal_weight + + return energy + + +def _motif_candidate_energy(torch, candidate_colors, context_colors, anchor): + if candidate_colors.shape[0] < 2 or candidate_colors.shape[1] < 2: + return candidate_colors.new_zeros(candidate_colors.shape[:3]) + + energy = candidate_colors.new_zeros(candidate_colors.shape[:3]) + contributions = candidate_colors.new_zeros(candidate_colors.shape[:3]) + + def add_block_error( + y_slice, + x_slice, + top_left, + top_right, + bottom_left, + bottom_right, + anchor_top_left, + anchor_top_right, + anchor_bottom_left, + anchor_bottom_right, + ): + k = candidate_colors[y_slice, x_slice].shape[2] + + def ensure_candidate_dim(tensor): + if tensor.ndim == 4: + return tensor + return tensor[..., None, :].expand(-1, -1, k, -1) + + candidate_block = torch.stack( + [ + ensure_candidate_dim(top_left), + ensure_candidate_dim(top_right), + ensure_candidate_dim(bottom_left), + ensure_candidate_dim(bottom_right), + ], + dim=-2, + ) + anchor_block = torch.stack( + [ + anchor_top_left, + anchor_top_right, + anchor_bottom_left, + anchor_bottom_right, + ], + dim=-2, + ) + error = ( + _normalized_motif_blocks(torch, candidate_block) + - _normalized_motif_blocks(torch, anchor_block[..., None, :, :]) + ).abs().mean(dim=(-2, -1)) + energy[y_slice, x_slice] = energy[y_slice, x_slice] + error + contributions[y_slice, x_slice] = contributions[y_slice, x_slice] + 1.0 + + add_block_error( + slice(1, None), + slice(1, None), + context_colors[:-1, :-1, :], + context_colors[:-1, 1:, :], + context_colors[1:, :-1, :], + candidate_colors[1:, 1:, :, :], + anchor[:-1, :-1, :], + anchor[:-1, 1:, :], + anchor[1:, :-1, :], + anchor[1:, 1:, :], + ) + add_block_error( + slice(1, None), + slice(None, -1), + context_colors[:-1, :-1, :], + context_colors[:-1, 1:, :], + candidate_colors[1:, :-1, :, :], + context_colors[1:, 1:, :], + anchor[:-1, :-1, :], + anchor[:-1, 1:, :], + anchor[1:, :-1, :], + anchor[1:, 1:, :], + ) + add_block_error( + slice(None, -1), + slice(1, None), + context_colors[:-1, :-1, :], + candidate_colors[:-1, 1:, :, :], + context_colors[1:, :-1, :], + context_colors[1:, 1:, :], + anchor[:-1, :-1, :], + anchor[:-1, 1:, :], + anchor[1:, :-1, :], + anchor[1:, 1:, :], + ) + add_block_error( + slice(None, -1), + slice(None, -1), + candidate_colors[:-1, :-1, :, :], + context_colors[:-1, 1:, :], + context_colors[1:, :-1, :], + context_colors[1:, 1:, :], + anchor[:-1, :-1, :], + anchor[:-1, 1:, :], + anchor[1:, :-1, :], + anchor[1:, 1:, :], + ) + + return energy / contributions.clamp_min(1.0) + + +def _relax_candidate_selection( + torch, + candidate_colors, + base_energy, + anchor, + desired_delta_x, + desired_delta_y, + desired_delta_diag, + desired_delta_anti, + solver_params: SolverHyperParams, + *, + iterations: int, +): + if iterations <= 0: + return torch.argmin(base_energy, dim=-1), [] + + start_temp = max(1e-3, solver_params.relax_start_temperature) + end_temp = max(1e-3, min(start_temp, solver_params.relax_end_temperature)) + damping = float(np.clip(solver_params.relax_damping, 0.0, 0.95)) + + probs = (-base_energy / start_temp).softmax(dim=2) + loss_history: list[float] = [] + final_energy = base_energy + + for step in range(iterations): + context_colors = (candidate_colors * probs[..., None]).sum(dim=2) + final_energy = base_energy.clone() + final_energy = final_energy + _pairwise_candidate_energy( + candidate_colors, + context_colors, + desired_delta_x, + desired_delta_y, + desired_delta_diag, + desired_delta_anti, + orthogonal_weight=solver_params.refine_orthogonal_weight, + diagonal_weight=solver_params.refine_diagonal_weight, + ) + if solver_params.refine_motif_weight > 0.0: + motif_energy = _motif_candidate_energy(torch, candidate_colors, context_colors, anchor) + final_energy = final_energy + motif_energy * solver_params.refine_motif_weight + + alpha = 1.0 if iterations <= 1 else step / float(iterations - 1) + temperature = start_temp + (end_temp - start_temp) * alpha + updated_probs = (-final_energy / max(temperature, 1e-3)).softmax(dim=2) + if damping > 0.0: + probs = updated_probs * (1.0 - damping) + probs * damping + probs = probs / probs.sum(dim=2, keepdim=True).clamp_min(1e-8) + else: + probs = updated_probs + loss_history.append(float((final_energy * probs).sum(dim=2).mean().detach().cpu().item())) + + return torch.argmax(probs, dim=2), loss_history + + def _snap_output_to_source_pixels( torch, source_t, @@ -448,71 +686,41 @@ def _discrete_refine_output( desired_delta_diag = anchor_delta_diag desired_delta_anti = anchor_delta_anti - selected = torch.argmin(anchor_energy + distance_energy * 0.05, dim=-1) + passes = max(0, iterations) + relax_iterations = min(max(0, solver_params.relax_iterations), passes) if passes > 0 else 0 + selected, relax_history = _relax_candidate_selection( + torch, + candidate_colors, + anchor_energy + distance_energy * 0.05, + anchor, + desired_delta_x, + desired_delta_y, + desired_delta_diag, + desired_delta_anti, + solver_params, + iterations=relax_iterations, + ) best_selected = selected.clone() best_score = float("inf") - loss_history: list[float] = [] - passes = max(0, iterations) + loss_history: list[float] = list(relax_history) anchor_hw = anchor_t.permute(0, 2, 3, 1) for step in range(passes + 1): selected_colors = _select_colors(candidate_colors, selected) energy = base_energy.clone() - - if output_width > 1: - left_selected = torch.zeros_like(selected_colors) - left_selected[:, 1:, :] = selected_colors[:, :-1, :] - left_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) - left_mask[:, 1:, :] = 1.0 - left_desired = torch.zeros_like(selected_colors) - left_desired[:, 1:, :] = desired_delta_x - left_error = ((candidate_colors - left_selected[..., None, :]) - left_desired[..., None, :]).abs().mean(dim=-1) - energy = energy + left_error * left_mask * solver_params.refine_orthogonal_weight - - right_selected = torch.zeros_like(selected_colors) - right_selected[:, :-1, :] = selected_colors[:, 1:, :] - right_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) - right_mask[:, :-1, :] = 1.0 - right_desired = torch.zeros_like(selected_colors) - right_desired[:, :-1, :] = desired_delta_x - right_error = ((right_selected[..., None, :] - candidate_colors) - right_desired[..., None, :]).abs().mean(dim=-1) - energy = energy + right_error * right_mask * solver_params.refine_orthogonal_weight - if output_height > 1: - up_selected = torch.zeros_like(selected_colors) - up_selected[1:, :, :] = selected_colors[:-1, :, :] - up_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) - up_mask[1:, :, :] = 1.0 - up_desired = torch.zeros_like(selected_colors) - up_desired[1:, :, :] = desired_delta_y - up_error = ((candidate_colors - up_selected[..., None, :]) - up_desired[..., None, :]).abs().mean(dim=-1) - energy = energy + up_error * up_mask * solver_params.refine_orthogonal_weight - - down_selected = torch.zeros_like(selected_colors) - down_selected[:-1, :, :] = selected_colors[1:, :, :] - down_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) - down_mask[:-1, :, :] = 1.0 - down_desired = torch.zeros_like(selected_colors) - down_desired[:-1, :, :] = desired_delta_y - down_error = ((down_selected[..., None, :] - candidate_colors) - down_desired[..., None, :]).abs().mean(dim=-1) - energy = energy + down_error * down_mask * solver_params.refine_orthogonal_weight - if output_height > 1 and output_width > 1: - diag_selected = torch.zeros_like(selected_colors) - diag_selected[1:, 1:, :] = selected_colors[:-1, :-1, :] - diag_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) - diag_mask[1:, 1:, :] = 1.0 - diag_desired = torch.zeros_like(selected_colors) - diag_desired[1:, 1:, :] = desired_delta_diag - diag_error = ((candidate_colors - diag_selected[..., None, :]) - diag_desired[..., None, :]).abs().mean(dim=-1) - energy = energy + diag_error * diag_mask * solver_params.refine_diagonal_weight - - anti_selected = torch.zeros_like(selected_colors) - anti_selected[1:, :-1, :] = selected_colors[:-1, 1:, :] - anti_mask = torch.zeros((output_height, output_width, 1), device=uv.device, dtype=uv.dtype) - anti_mask[1:, :-1, :] = 1.0 - anti_desired = torch.zeros_like(selected_colors) - anti_desired[1:, :-1, :] = desired_delta_anti - anti_error = ((candidate_colors - anti_selected[..., None, :]) - anti_desired[..., None, :]).abs().mean(dim=-1) - energy = energy + anti_error * anti_mask * solver_params.refine_diagonal_weight + energy = energy + _pairwise_candidate_energy( + candidate_colors, + selected_colors, + desired_delta_x, + desired_delta_y, + desired_delta_diag, + desired_delta_anti, + orthogonal_weight=solver_params.refine_orthogonal_weight, + diagonal_weight=solver_params.refine_diagonal_weight, + ) + if solver_params.refine_motif_weight > 0.0: + motif_energy = _motif_candidate_energy(torch, candidate_colors, selected_colors, anchor) + energy = energy + motif_energy * solver_params.refine_motif_weight candidate_t = selected_colors[None, ...] score = float( diff --git a/src/repixelizer/metrics.py b/src/repixelizer/metrics.py index 3d82c85..f49d777 100644 --- a/src/repixelizer/metrics.py +++ b/src/repixelizer/metrics.py @@ -169,6 +169,50 @@ def foreground_motif_error( return float(np.mean(diff[block_mask])) +def foreground_edge_position_error( + a: np.ndarray, + b: np.ndarray, + alpha_threshold: float = 0.05, + halo: int = 1, + radius: int = 1, + distance_penalty: float = 0.35, + edge_threshold: float = 0.02, +) -> float: + edge_a = _edge_strength_map(a) + edge_b = _edge_strength_map(b) + mask = sprite_mask(a, b, alpha_threshold=alpha_threshold, halo=halo) + support_mask = mask & ((edge_a >= edge_threshold) | (edge_b >= edge_threshold)) + if not np.any(support_mask): + return 0.0 + + support_from_b = _best_local_edge_support(edge_b, radius=radius, distance_penalty=distance_penalty) + support_from_a = _best_local_edge_support(edge_a, radius=radius, distance_penalty=distance_penalty) + deficit_a = np.maximum(0.0, edge_a - support_from_b) + deficit_b = np.maximum(0.0, edge_b - support_from_a) + return float(np.mean((deficit_a[support_mask] + deficit_b[support_mask]) * 0.5)) + + +def foreground_edge_concentration( + rgba: np.ndarray, + alpha_threshold: float = 0.05, +) -> float: + premul = premultiply(rgba) + alpha = rgba[..., 3] + + delta_x = np.mean(np.abs(premul[:, 1:, :] - premul[:, :-1, :]), axis=-1) + delta_y = np.mean(np.abs(premul[1:, :, :] - premul[:-1, :, :]), axis=-1) + mask_x = (alpha[:, 1:] >= alpha_threshold) | (alpha[:, :-1] >= alpha_threshold) + mask_y = (alpha[1:, :] >= alpha_threshold) | (alpha[:-1, :] >= alpha_threshold) + + values = np.concatenate([delta_x[mask_x], delta_y[mask_y]], axis=0) + if values.size == 0: + return 1.0 + total = float(np.sum(values)) + if total <= 1e-6: + return 0.0 + return float(np.sum(values * values) / total) + + def source_lattice_consistency_breakdown( source_rgba: np.ndarray, output_rgba: np.ndarray, @@ -281,6 +325,45 @@ def _dilate_mask(mask: np.ndarray, radius: int) -> np.ndarray: return np.any(np.stack(windows, axis=0), axis=0) +def _edge_strength_map(rgba: np.ndarray) -> np.ndarray: + premul = premultiply(rgba) + right = np.zeros(rgba.shape[:2], dtype=np.float32) + down = np.zeros(rgba.shape[:2], dtype=np.float32) + right[:, :-1] = np.mean(np.abs(premul[:, 1:, :] - premul[:, :-1, :]), axis=-1) + down[:-1, :] = np.mean(np.abs(premul[1:, :, :] - premul[:-1, :, :]), axis=-1) + edge = np.maximum.reduce( + [ + right, + np.pad(right[:, :-1], ((0, 0), (1, 0))), + down, + np.pad(down[:-1, :], ((1, 0), (0, 0))), + ] + ) + return edge.astype(np.float32) + + +def _best_local_edge_support(edge: np.ndarray, *, radius: int, distance_penalty: float) -> np.ndarray: + if radius <= 0: + return edge + height, width = edge.shape + best = np.zeros_like(edge, dtype=np.float32) + for dy in range(-radius, radius + 1): + for dx in range(-radius, radius + 1): + weight = 1.0 / (1.0 + np.hypot(dx, dy) * distance_penalty) + shifted = np.zeros_like(edge, dtype=np.float32) + src_y0 = max(0, -dy) + src_y1 = min(height, height - dy) + src_x0 = max(0, -dx) + src_x1 = min(width, width - dx) + dst_y0 = max(0, dy) + dst_y1 = dst_y0 + (src_y1 - src_y0) + dst_x0 = max(0, dx) + dst_x1 = dst_x0 + (src_x1 - src_x0) + shifted[dst_y0:dst_y1, dst_x0:dst_x1] = edge[src_y0:src_y1, src_x0:src_x1] * weight + best = np.maximum(best, shifted) + return best + + def _motif_blocks(premul: np.ndarray) -> np.ndarray: blocks = np.stack( [ diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index 1f460dd..b211a2c 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -20,6 +20,11 @@ class SolverHyperParams: refine_source_delta_weight: float = 0.25 refine_orthogonal_weight: float = 0.24 refine_diagonal_weight: float = 0.08 + refine_motif_weight: float = 0.12 + relax_iterations: int = 8 + relax_start_temperature: float = 0.55 + relax_end_temperature: float = 0.10 + relax_damping: float = 0.35 structure_boundary_weight: float = 0.60 structure_anchor_adjacency_weight: float = 0.15 structure_anchor_motif_weight: float = 0.20 diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index 27f2d5f..c5bd413 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -22,6 +22,7 @@ from .io import load_rgba, nearest_resize, save_rgba from .metrics import source_lattice_consistency_breakdown from .params import SolverHyperParams from .palette import load_palette, quantize_rgba, save_palette_report +from .preprocess import strip_edge_background from .types import InferenceResult, RunResult @@ -37,10 +38,13 @@ def run_pipeline( steps: int = 200, device: str = "auto", solver_params: SolverHyperParams | None = None, + strip_background: bool = False, ) -> RunResult: started = time.perf_counter() solver_params = solver_params or SolverHyperParams() source = load_rgba(input_path) + if strip_background: + source = strip_edge_background(source) inference = infer_lattice(source, target_size=target_size, device=device) analysis = analyze_source(source, seed=seed) inference = _select_phase_candidate( @@ -80,6 +84,8 @@ def run_pipeline( if diagnostics_dir: diagnostics_path = Path(diagnostics_dir) diagnostics_path.mkdir(parents=True, exist_ok=True) + if strip_background: + save_rgba(diagnostics_path / "preprocessed-source.png", source) write_lattice_overlay(diagnostics_path / "lattice-overlay.png", source, inference) write_comparison(diagnostics_path / "comparison.png", source, output_rgba) save_rgba( @@ -97,6 +103,7 @@ def run_pipeline( "seed": seed, "steps": steps, "device": device, + "strip_background": strip_background, "solver_params": solver_params.to_dict(), } write_json(diagnostics_path / "run.json", run_json) diff --git a/src/repixelizer/preprocess.py b/src/repixelizer/preprocess.py new file mode 100644 index 0000000..b786602 --- /dev/null +++ b/src/repixelizer/preprocess.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +from collections import deque + +import numpy as np + + +def strip_edge_background( + rgba: np.ndarray, + *, + border_width: int = 8, + quantize_step: int = 8, + representative_count: int = 6, + color_threshold: float = 20.0 / 255.0, + chroma_floor: float = 6.0 / 255.0, + luminance_slack: float = 18.0 / 255.0, + enclosed_min_span: int = 24, + enclosed_max_aspect: float = 4.0, + monochrome_bucket_fraction: float = 0.7, + fringe_color_threshold: float = 28.0 / 255.0, + fringe_chroma_slack: float = 10.0 / 255.0, + fringe_luminance_slack: float = 10.0 / 255.0, +) -> np.ndarray: + image = rgba.copy() + if image.size == 0: + return image + border_rgb = _border_rgb_samples(image, border_width=border_width) + if border_rgb.size == 0: + return image + + rgb = image[..., :3] + candidate = _monochrome_candidate_mask( + rgb, + border_rgb=border_rgb, + quantize_step=quantize_step, + color_threshold=max(color_threshold, 14.0 / 255.0), + bucket_fraction=monochrome_bucket_fraction, + ) + if candidate is None: + representatives = _background_representatives( + border_rgb, + quantize_step=quantize_step, + representative_count=representative_count, + ) + border_luminance = _luminance(border_rgb) + border_chroma = border_rgb.max(axis=1) - border_rgb.min(axis=1) + luminance_min = max(0.0, float(np.percentile(border_luminance, 1.0) - luminance_slack)) + luminance_max = min(1.0, float(np.percentile(border_luminance, 99.0) + luminance_slack)) + chroma_limit = max(chroma_floor, float(np.percentile(border_chroma, 99.0) + 2.0 / 255.0)) + luminance = _luminance(rgb) + chroma = rgb.max(axis=-1) - rgb.min(axis=-1) + flat_rgb = rgb.reshape(-1, 3) + distances = np.linalg.norm( + flat_rgb[:, None, :] - representatives[None, :, :], + axis=-1, + ).min(axis=1).reshape(rgb.shape[:2]) + candidate = ( + (distances <= color_threshold) + & (chroma <= chroma_limit) + & (luminance >= luminance_min) + & (luminance <= luminance_max) + ) + relaxed_candidate = ( + (distances <= max(fringe_color_threshold, color_threshold)) + & (chroma <= chroma_limit + fringe_chroma_slack) + & (luminance >= max(0.0, luminance_min - fringe_luminance_slack)) + & (luminance <= min(1.0, luminance_max + fringe_luminance_slack)) + ) + else: + border_luminance = _luminance(border_rgb) + luminance_min = max(0.0, float(np.percentile(border_luminance, 1.0) - luminance_slack)) + luminance_max = min(1.0, float(np.percentile(border_luminance, 99.0) + luminance_slack)) + chroma = rgb.max(axis=-1) - rgb.min(axis=-1) + luminance = _luminance(rgb) + dominant_color = _dominant_border_color(border_rgb, quantize_step=quantize_step) + relaxed_distance = np.linalg.norm(rgb - dominant_color[None, None, :], axis=-1) + relaxed_candidate = ( + (relaxed_distance <= max(fringe_color_threshold, color_threshold)) + & (chroma <= max(chroma_floor + fringe_chroma_slack, 16.0 / 255.0)) + & (luminance >= max(0.0, luminance_min - fringe_luminance_slack)) + & (luminance <= min(1.0, luminance_max + fringe_luminance_slack)) + ) + + edge_connected = _edge_connected_mask(candidate) + remove_mask = edge_connected | _enclosed_background_mask( + candidate & ~edge_connected, + min_span=enclosed_min_span, + max_aspect=enclosed_max_aspect, + ) + remove_mask = _grow_from_seed(relaxed_candidate, remove_mask) + image[remove_mask, 3] = 0.0 + image[remove_mask, :3] = 0.0 + return image + + +def _border_rgb_samples(rgba: np.ndarray, *, border_width: int) -> np.ndarray: + height, width = rgba.shape[:2] + border = max(1, min(border_width, height // 2, width // 2)) + mask = np.zeros((height, width), dtype=bool) + mask[:border, :] = True + mask[-border:, :] = True + mask[:, :border] = True + mask[:, -border:] = True + opaque = rgba[..., 3] > 0.95 + return rgba[..., :3][mask & opaque] + + +def _background_representatives( + border_rgb: np.ndarray, + *, + quantize_step: int, + representative_count: int, +) -> np.ndarray: + step = max(1, quantize_step) + border8 = np.clip(np.rint(border_rgb * 255.0), 0, 255).astype(np.int16) + buckets = border8 // step + keys, inverse, counts = np.unique(buckets, axis=0, return_inverse=True, return_counts=True) + order = np.argsort(counts)[::-1] + representatives: list[np.ndarray] = [] + for bucket_index in order[: max(1, representative_count)]: + members = border_rgb[inverse == bucket_index] + representatives.append(members.mean(axis=0)) + return np.stack(representatives, axis=0).astype(np.float32) + + +def _monochrome_candidate_mask( + rgb: np.ndarray, + *, + border_rgb: np.ndarray, + quantize_step: int, + color_threshold: float, + bucket_fraction: float, +) -> np.ndarray | None: + step = max(1, quantize_step) + border8 = np.clip(np.rint(border_rgb * 255.0), 0, 255).astype(np.int16) + buckets = border8 // step + _, inverse, counts = np.unique(buckets, axis=0, return_inverse=True, return_counts=True) + dominant_index = int(np.argmax(counts)) + dominant_fraction = float(counts[dominant_index] / max(1, border8.shape[0])) + if dominant_fraction < bucket_fraction: + return None + dominant_color = border_rgb[inverse == dominant_index].mean(axis=0) + distances = np.linalg.norm(rgb - dominant_color[None, None, :], axis=-1) + return distances <= color_threshold + + +def _dominant_border_color(border_rgb: np.ndarray, *, quantize_step: int) -> np.ndarray: + step = max(1, quantize_step) + border8 = np.clip(np.rint(border_rgb * 255.0), 0, 255).astype(np.int16) + buckets = border8 // step + _, inverse, counts = np.unique(buckets, axis=0, return_inverse=True, return_counts=True) + dominant_index = int(np.argmax(counts)) + return border_rgb[inverse == dominant_index].mean(axis=0) + + +def _luminance(rgb: np.ndarray) -> np.ndarray: + return rgb[..., 0] * 0.2126 + rgb[..., 1] * 0.7152 + rgb[..., 2] * 0.0722 + + +def _edge_connected_mask(mask: np.ndarray) -> np.ndarray: + height, width = mask.shape + visited = np.zeros((height, width), dtype=bool) + queue: deque[tuple[int, int]] = deque() + for x in range(width): + if mask[0, x]: + queue.append((0, x)) + visited[0, x] = True + if mask[height - 1, x] and not visited[height - 1, x]: + queue.append((height - 1, x)) + visited[height - 1, x] = True + for y in range(height): + if mask[y, 0] and not visited[y, 0]: + queue.append((y, 0)) + visited[y, 0] = True + if mask[y, width - 1] and not visited[y, width - 1]: + queue.append((y, width - 1)) + visited[y, width - 1] = True + while queue: + y, x = queue.popleft() + for ny in range(max(0, y - 1), min(height, y + 2)): + for nx in range(max(0, x - 1), min(width, x + 2)): + if mask[ny, nx] and not visited[ny, nx]: + visited[ny, nx] = True + queue.append((ny, nx)) + return visited + + +def _enclosed_background_mask( + mask: np.ndarray, + *, + min_span: int, + max_aspect: float, +) -> np.ndarray: + height, width = mask.shape + visited = np.zeros((height, width), dtype=bool) + selected = np.zeros((height, width), dtype=bool) + for y in range(height): + for x in range(width): + if not mask[y, x] or visited[y, x]: + continue + queue: deque[tuple[int, int]] = deque([(y, x)]) + visited[y, x] = True + coords: list[tuple[int, int]] = [] + left = right = x + top = bottom = y + while queue: + cy, cx = queue.popleft() + coords.append((cy, cx)) + left = min(left, cx) + right = max(right, cx) + top = min(top, cy) + bottom = max(bottom, cy) + for ny in range(max(0, cy - 1), min(height, cy + 2)): + for nx in range(max(0, cx - 1), min(width, cx + 2)): + if mask[ny, nx] and not visited[ny, nx]: + visited[ny, nx] = True + queue.append((ny, nx)) + span_x = right - left + 1 + span_y = bottom - top + 1 + short_span = min(span_x, span_y) + long_span = max(span_x, span_y) + aspect = long_span / max(1, short_span) + if short_span < min_span or aspect > max_aspect: + continue + for cy, cx in coords: + selected[cy, cx] = True + return selected + + +def _grow_from_seed(mask: np.ndarray, seed: np.ndarray) -> np.ndarray: + height, width = mask.shape + grown = seed.copy() + queue: deque[tuple[int, int]] = deque(map(tuple, np.argwhere(seed))) + while queue: + y, x = queue.popleft() + for ny in range(max(0, y - 1), min(height, y + 2)): + for nx in range(max(0, x - 1), min(width, x + 2)): + if mask[ny, nx] and not grown[ny, nx]: + grown[ny, nx] = True + queue.append((ny, nx)) + return grown diff --git a/tests/fixtures/real/ai-badge-cleaned.json b/tests/fixtures/real/ai-badge-cleaned.json new file mode 100644 index 0000000..b4be0ce --- /dev/null +++ b/tests/fixtures/real/ai-badge-cleaned.json @@ -0,0 +1,14 @@ +{ + "id": "ai-badge-cleaned", + "kind": "real-ai-generated-emblem", + "description": "Manually cleaned transparent version of the AI badge emblem used as a real-image regression fixture.", + "source_image": "C:\\Users\\Meta\\.codex\\generated_images\\019db22f-5487-7a31-8a67-45f0c801dd7c\\ig_0b9862296469b4740169e7fbc657ac81919d4d2267673f6bae.png", + "revision": 3, + "notes": [ + "Use this version when evaluating lattice inference and thin-feature preservation without the baked checkerboard background contaminating results.", + "Primary failure region is the sword tip and nearby outline taper, plus other sharp single-pixel and near-single-pixel features at mixed orientations.", + "Keep the raw checkerboard-background version around separately when testing the automatic background stripper.", + "Revision 2 reapplies transparency as a raster mask and thresholds that mask to suppress subtle dark-background glow left by the earlier cleanup.", + "Revision 3 binarizes the repo fixture copy to strict 0/255 alpha so the regression asset stays unambiguous." + ] +} diff --git a/tests/fixtures/real/ai-badge-cleaned.png b/tests/fixtures/real/ai-badge-cleaned.png new file mode 100644 index 0000000..fe25865 Binary files /dev/null and b/tests/fixtures/real/ai-badge-cleaned.png differ diff --git a/tests/test_cli.py b/tests/test_cli.py index 7967726..9539062 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -16,3 +16,14 @@ def test_default_command_runs_without_explicit_run_subcommand(tmp_path: Path) -> exit_code = main([str(input_path), "--out", str(output_path), "--steps", "8"]) assert exit_code == 0 assert output_path.exists() + + +def test_default_command_accepts_strip_background_flag(tmp_path: Path) -> None: + source = make_emblem(16, 16) + fake = fake_pixelize(source, upscale=8, phase_x=0.2, phase_y=0.25, blur_radius=0.4) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + save_rgba(input_path, fake) + exit_code = main([str(input_path), "--out", str(output_path), "--steps", "0", "--strip-background"]) + assert exit_code == 0 + assert output_path.exists() diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 18b8519..4f4bb7d 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -5,6 +5,8 @@ import numpy as np from repixelizer.metrics import ( exact_match_ratio, foreground_adjacency_error, + foreground_edge_concentration, + foreground_edge_position_error, foreground_exact_match_ratio, foreground_motif_error, foreground_reconstruction_error, @@ -66,6 +68,35 @@ def test_texture_metrics_track_adjacency_and_local_motifs() -> None: assert foreground_motif_error(original, same_colors_wrong_structure) > 0.0 +def test_edge_position_metric_prefers_crisp_shift_to_blurry_smear() -> None: + original = np.zeros((7, 7, 4), dtype=np.float32) + original[1:6, 3] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32) + + shifted = np.zeros_like(original) + shifted[1:6, 4] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32) + + blurred = np.zeros_like(original) + blurred[1:6, 2] = np.asarray([0.4, 0.4, 0.4, 1.0], dtype=np.float32) + blurred[1:6, 3] = np.asarray([0.7, 0.7, 0.7, 1.0], dtype=np.float32) + blurred[1:6, 4] = np.asarray([0.4, 0.4, 0.4, 1.0], dtype=np.float32) + + assert foreground_edge_position_error(original, original) == 0.0 + assert foreground_edge_position_error(original, shifted) > 0.0 + assert foreground_edge_position_error(original, blurred) > foreground_edge_position_error(original, shifted) + + +def test_edge_concentration_prefers_crisp_edges_to_blurry_smear() -> None: + shifted = np.zeros((7, 7, 4), dtype=np.float32) + shifted[1:6, 4] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32) + + blurred = np.zeros_like(shifted) + blurred[1:6, 2] = np.asarray([0.4, 0.4, 0.4, 1.0], dtype=np.float32) + blurred[1:6, 3] = np.asarray([0.7, 0.7, 0.7, 1.0], dtype=np.float32) + blurred[1:6, 4] = np.asarray([0.4, 0.4, 0.4, 1.0], dtype=np.float32) + + assert foreground_edge_concentration(shifted) > foreground_edge_concentration(blurred) + + def test_source_lattice_consistency_prefers_matching_cell_structure() -> None: lowres = np.asarray( [ diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 582efe6..842b287 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -74,6 +74,30 @@ def test_pipeline_preserves_transparency_for_sprite(tmp_path: Path) -> None: assert partial.size == 0 +def test_pipeline_can_strip_checkerboard_background(tmp_path: Path) -> None: + source = make_sprite(24, 24) + checker = np.ones_like(source) + checker[..., 3] = 1.0 + for y in range(checker.shape[0]): + for x in range(checker.shape[1]): + tile = 0.97 if ((x // 4) + (y // 4)) % 2 == 0 else 0.91 + checker[y, x, :3] = tile + alpha = source[..., 3:4] + opaque_input = checker.copy() + opaque_input[..., :3] = source[..., :3] * alpha + checker[..., :3] * (1.0 - alpha) + + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, opaque_input) + result = run_pipeline(input_path, output_path, target_size=24, steps=0, strip_background=True) + assert result.source_rgba[..., 3].min() < 0.05 + assert result.output_rgba[..., 3].min() < 0.05 + assert result.output_rgba[..., 3].max() > 0.9 + + def test_pipeline_can_emit_initialized_output_without_optimizer_steps(tmp_path: Path) -> None: source = make_emblem(24, 24) fake = fake_pixelize(source, upscale=10, phase_x=0.15, phase_y=0.2, blur_radius=0.6, warp_strength=0.2, warp_detail=5) diff --git a/tests/test_preprocess.py b/tests/test_preprocess.py new file mode 100644 index 0000000..aed5b3a --- /dev/null +++ b/tests/test_preprocess.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import numpy as np + +from repixelizer.preprocess import strip_edge_background + + +def test_strip_edge_background_removes_light_checkerboard_and_preserves_subject() -> None: + rgba = np.ones((80, 80, 4), dtype=np.float32) + rgba[..., 3] = 1.0 + for y in range(80): + for x in range(80): + tile = 0.97 if ((x // 4) + (y // 4)) % 2 == 0 else 0.91 + rgba[y, x, :3] = tile + rgba[12:68, 10:70] = np.array([0.1, 0.7, 0.2, 1.0], dtype=np.float32) + for y in range(24, 56): + for x in range(26, 54): + tile = 0.97 if ((x // 4) + (y // 4)) % 2 == 0 else 0.91 + rgba[y, x, :3] = tile + rgba[20:60, 16:20] = np.array([0.98, 0.97, 0.96, 1.0], dtype=np.float32) + + stripped = strip_edge_background(rgba) + + assert stripped[0, 0, 3] == 0.0 + assert stripped[4, 4, 3] == 0.0 + assert stripped[40, 40, 3] == 0.0 + assert stripped[30, 18, 3] == 1.0 + assert np.allclose(stripped[30, 18, :3], rgba[30, 18, :3]) + assert stripped[16, 16, 3] == 1.0 + + +def test_strip_edge_background_handles_trivial_monochrome_background() -> None: + rgba = np.zeros((32, 32, 4), dtype=np.float32) + rgba[..., :3] = np.array([0.2, 0.3, 0.8], dtype=np.float32) + rgba[..., 3] = 1.0 + rgba[8:24, 8:24] = np.array([1.0, 0.6, 0.1, 1.0], dtype=np.float32) + + stripped = strip_edge_background(rgba) + + assert stripped[0, 0, 3] == 0.0 + assert stripped[16, 16, 3] == 1.0 + assert np.allclose(stripped[16, 16, :3], rgba[16, 16, :3])