diff --git a/README.md b/README.md index 963d744..8ac4852 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,14 @@ Recent adjacency-focused status: - the cleaned AI badge regression now beats naive resize on the repo's source-lattice consistency metric under the selected lattice - the snap-to-refine handoff is now source-first rather than dominated by the softened representative lattice - low-confidence phase reranking now uses a soft size penalty instead of a hard size-jump reject +- an experimental `tile-graph` reconstruction mode now exists behind `--reconstruction-mode tile-graph` + +Current tile-graph status: + +- it builds source-driven cell candidates from connected source clusters plus lattice-aligned fallback coverage +- it learns directional adjacency preferences between candidates and places them with a soft discrete propagation loop +- it already beats naive resize on the repo's synthetic thin-feature regression +- it is not yet good enough to replace the default continuous path on the cleaned AI badge stress case ## Quickstart @@ -47,6 +55,7 @@ Run the optimizer: ```powershell repixelize input.png --out output.png repixelize input.png --out output.png --diagnostics-dir diagnostics --device auto +repixelize input.png --out output.png --reconstruction-mode tile-graph --diagnostics-dir diagnostics --device cpu ``` Run the optimizer plus baselines: diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index 17087be..0084c4b 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -49,6 +49,7 @@ The current implementation should be treated as a strong experimental baseline, The most provisional areas are: - lattice inference on real, messy generated imagery - edge-cell reliability and candidate coverage in the continuous stage +- the new experimental tile-graph reconstruction path, especially on real noisy inputs - low-confidence reranking once multiple plausible lattice sizes survive the new soft penalty - compare-mode metrics as a proxy for actual visual quality @@ -66,10 +67,15 @@ 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 +Experimental tile-graph note: +- there is now a separate `tile-graph` reconstruction mode that extracts cell candidates from connected source clusters, fills gaps with lattice-aligned fallback candidates, learns candidate adjacency, and places tiles with a soft propagation loop +- that path is already promising on synthetic thin-feature cases, but it currently overcommits to larger lattice candidates and underperforms the default continuous path on the cleaned badge fixture + 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 - evaluate coordinated local moves when single-cell refinement still breaks thin contours +- decide whether the tile-graph path should stay a diagnostics-only experiment, become a refine-stage candidate generator, or mature into a full alternate solver ## High-value regression cases diff --git a/src/repixelizer/cli.py b/src/repixelizer/cli.py index d83c7a1..34409fb 100644 --- a/src/repixelizer/cli.py +++ b/src/repixelizer/cli.py @@ -29,6 +29,12 @@ 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( + "--reconstruction-mode", + default="continuous", + choices=("continuous", "tile-graph"), + help="Reconstruction engine to run after lattice estimation", + ) target.add_argument( "--strip-background", action="store_true", @@ -125,6 +131,7 @@ def main(argv: list[str] | None = None) -> int: seed=args.seed, steps=args.steps, device=args.device, + reconstruction_mode=args.reconstruction_mode, strip_background=args.strip_background, ) return 0 @@ -172,6 +179,7 @@ def main(argv: list[str] | None = None) -> int: seed=args.seed, steps=args.steps, device=args.device, + reconstruction_mode=args.reconstruction_mode, strip_background=args.strip_background, ) return 0 diff --git a/src/repixelizer/compare.py b/src/repixelizer/compare.py index 8d61221..ed3eccb 100644 --- a/src/repixelizer/compare.py +++ b/src/repixelizer/compare.py @@ -64,6 +64,7 @@ def run_compare( steps: int = 200, device: str = "auto", strip_background: bool = False, + reconstruction_mode: str = "continuous", ) -> 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) @@ -78,6 +79,7 @@ def run_compare( steps=steps, device=device, strip_background=strip_background, + reconstruction_mode=reconstruction_mode, ) source = result.source_rgba palette = load_palette(palette_path) if palette_path else None diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index 5fc0083..8c593ee 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -73,6 +73,15 @@ class SolverHyperParams: phase_rerank_confidence_threshold: float = 0.12 phase_rerank_max_size_delta_ratio: float = 0.40 phase_rerank_margin: float = 0.004 + tile_graph_max_candidates: int = 160 + tile_graph_component_min_area_ratio: float = 0.10 + tile_graph_large_component_ratio: float = 1.35 + tile_graph_window_coverage_threshold: float = 0.16 + tile_graph_coordinate_weight: float = 0.08 + tile_graph_area_weight: float = 0.03 + tile_graph_count_weight: float = 0.55 + tile_graph_delta_weight: float = 0.45 + tile_graph_iterations: int = 10 def to_dict(self) -> dict[str, float | int]: return asdict(self) diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index a2d8140..471386e 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -28,6 +28,7 @@ from .metrics import ( from .params import SolverHyperParams from .palette import load_palette, quantize_rgba, save_palette_report from .preprocess import strip_edge_background +from .tile_graph import optimize_tile_graph from .types import InferenceResult, RunResult @@ -44,6 +45,7 @@ def run_pipeline( device: str = "auto", solver_params: SolverHyperParams | None = None, strip_background: bool = False, + reconstruction_mode: str = "continuous", ) -> RunResult: started = time.perf_counter() solver_params = solver_params or SolverHyperParams() @@ -59,8 +61,9 @@ def run_pipeline( seed=seed, device=device, solver_params=solver_params, + reconstruction_mode=reconstruction_mode, ) - solver = optimize_uv_field( + solver, reconstruction_diagnostics = _run_reconstruction( source, inference=inference, analysis=analysis, @@ -68,14 +71,19 @@ def run_pipeline( seed=seed, device=device, solver_params=solver_params, + reconstruction_mode=reconstruction_mode, ) cleanup = cleanup_pixels(solver.target_rgba, source_guidance=solver.guidance_strength) palette = load_palette(palette_path) if palette_path else None palette_result = quantize_rgba(cleanup.cleaned_rgba, mode=palette_mode, palette=palette) output_rgba = palette_result.rgba if palette_result else cleanup.cleaned_rgba + Path(output_path).parent.mkdir(parents=True, exist_ok=True) save_rgba(output_path, output_rgba) - diagnostics: dict[str, Any] = {"elapsed_seconds": time.perf_counter() - started} + diagnostics: dict[str, Any] = { + "elapsed_seconds": time.perf_counter() - started, + "reconstruction": reconstruction_diagnostics, + } result = RunResult( source_rgba=source, output_rgba=output_rgba, @@ -109,8 +117,11 @@ def run_pipeline( "steps": steps, "device": device, "strip_background": strip_background, + "reconstruction_mode": reconstruction_mode, "solver_params": solver_params.to_dict(), } + if diagnostics.get("reconstruction"): + run_json["reconstruction"] = diagnostics["reconstruction"] write_json(diagnostics_path / "run.json", run_json) if palette_result is not None: save_palette_report(diagnostics_path / "palette-report.json", palette_result.palette) @@ -125,6 +136,7 @@ def _select_phase_candidate( seed: int, device: str, solver_params: SolverHyperParams | None = None, + reconstruction_mode: str = "continuous", ) -> InferenceResult: solver_params = solver_params or SolverHyperParams() if len(inference.top_candidates) <= 1 or inference.confidence >= solver_params.phase_rerank_confidence_threshold: @@ -141,7 +153,7 @@ def _select_phase_candidate( confidence=inference.confidence, top_candidates=inference.top_candidates, ) - candidate_artifacts = optimize_uv_field( + candidate_artifacts, _ = _run_reconstruction( source, inference=candidate_inference, analysis=analysis, @@ -149,6 +161,7 @@ def _select_phase_candidate( seed=seed, device=device, solver_params=solver_params, + reconstruction_mode=reconstruction_mode, ) support = source_lattice_consistency_breakdown( source, @@ -272,3 +285,36 @@ def _size_delta_ratio(a: InferenceResult, b: InferenceResult) -> float: width_ratio = abs(b.target_width - a.target_width) / max(1.0, float(a.target_width)) height_ratio = abs(b.target_height - a.target_height) / max(1.0, float(a.target_height)) return max(width_ratio, height_ratio) + + +def _run_reconstruction( + source: np.ndarray, + *, + inference: InferenceResult, + analysis, + steps: int, + seed: int, + device: str, + solver_params: SolverHyperParams, + reconstruction_mode: str, +): + if reconstruction_mode == "tile-graph": + return optimize_tile_graph( + source, + inference=inference, + analysis=analysis, + steps=steps, + seed=seed, + device=device, + solver_params=solver_params, + ) + solver = optimize_uv_field( + source, + inference=inference, + analysis=analysis, + steps=steps, + seed=seed, + device=device, + solver_params=solver_params, + ) + return solver, {"mode": "continuous"} diff --git a/src/repixelizer/tile_graph.py b/src/repixelizer/tile_graph.py new file mode 100644 index 0000000..045155e --- /dev/null +++ b/src/repixelizer/tile_graph.py @@ -0,0 +1,515 @@ +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .io import premultiply, unpremultiply +from .params import SolverHyperParams +from .source_reference import build_source_lattice_reference +from .types import InferenceResult, SolverArtifacts, SourceAnalysis + +_RIGHT = 0 +_DOWN = 1 +_LEFT = 2 +_UP = 3 +_DIRS = ( + (_RIGHT, 0, 1), + (_DOWN, 1, 0), + (_LEFT, 0, -1), + (_UP, -1, 0), +) +_OPPOSITE = { + _RIGHT: _LEFT, + _DOWN: _UP, + _LEFT: _RIGHT, + _UP: _DOWN, +} + + +@dataclass(slots=True) +class TileGraphModel: + candidate_rgba: np.ndarray + candidate_coords: np.ndarray + candidate_area_ratio: np.ndarray + candidate_coverage: np.ndarray + pair_penalty: np.ndarray + reference_mean_rgba: np.ndarray + reference_sharp_rgba: np.ndarray + component_count: int + edge_density: float + + +@dataclass(slots=True) +class _Component: + label: int + coords: np.ndarray + bbox: tuple[int, int, int, int] + mask: np.ndarray + integral: np.ndarray + area: int + centroid_y: float + centroid_x: float + + +def _make_regular_uv(height: int, width: int, target_height: int, target_width: int, phase_x: float, phase_y: float) -> np.ndarray: + cell_x = width / target_width + cell_y = height / target_height + xs = (np.arange(target_width, dtype=np.float32) + 0.5 + phase_x) * cell_x - 0.5 + ys = (np.arange(target_height, dtype=np.float32) + 0.5 + phase_y) * cell_y - 0.5 + xs = np.clip(xs, 0.0, max(0.0, width - 1)) + ys = np.clip(ys, 0.0, max(0.0, height - 1)) + grid_x, grid_y = np.meshgrid(xs, ys) + uv = np.stack([grid_x, grid_y], axis=-1) + uv[..., 0] = (uv[..., 0] / max(1.0, width - 1)) * 2.0 - 1.0 + uv[..., 1] = (uv[..., 1] / max(1.0, height - 1)) * 2.0 - 1.0 + return uv.astype(np.float32) + + +def _component_integral(mask: np.ndarray) -> np.ndarray: + integral = np.zeros((mask.shape[0] + 1, mask.shape[1] + 1), dtype=np.int32) + integral[1:, 1:] = np.cumsum(np.cumsum(mask.astype(np.int32), axis=0), axis=1) + return integral + + +def _integral_sum(integral: np.ndarray, y0: int, x0: int, y1: int, x1: int) -> int: + return int(integral[y1, x1] - integral[y0, x1] - integral[y1, x0] + integral[y0, x0]) + + +def _window_bounds(center_y: float, center_x: float, cell_h: float, cell_w: float, height: int, width: int) -> tuple[int, int, int, int]: + y0 = int(np.floor(center_y - cell_h * 0.5)) + x0 = int(np.floor(center_x - cell_w * 0.5)) + y1 = int(np.ceil(center_y + cell_h * 0.5)) + x1 = int(np.ceil(center_x + cell_w * 0.5)) + y0 = max(0, min(height - 1, y0)) + x0 = max(0, min(width - 1, x0)) + y1 = max(y0 + 1, min(height, y1)) + x1 = max(x0 + 1, min(width, x1)) + return y0, x0, y1, x1 + + +def _source_to_cell_coord( + center_y: float, + center_x: float, + *, + cell_h: float, + cell_w: float, + target_height: int, + target_width: int, + phase_y: float, + phase_x: float, +) -> tuple[int, int]: + x_idx = int(np.floor((center_x + 0.5) / max(cell_w, 1e-4) - phase_x)) + y_idx = int(np.floor((center_y + 0.5) / max(cell_h, 1e-4) - phase_y)) + x_idx = max(0, min(target_width - 1, x_idx)) + y_idx = max(0, min(target_height - 1, y_idx)) + return y_idx, x_idx + + +def _window_mean_premul(source_premul: np.ndarray, center_y: float, center_x: float, cell_h: float, cell_w: float) -> np.ndarray: + y0, x0, y1, x1 = _window_bounds(center_y, center_x, cell_h, cell_w, source_premul.shape[0], source_premul.shape[1]) + window = source_premul[y0:y1, x0:x1] + if window.size == 0: + return np.zeros(4, dtype=np.float32) + return window.reshape(-1, 4).mean(axis=0).astype(np.float32) + + +def _window_component_coverage(component: _Component, center_y: float, center_x: float, cell_h: float, cell_w: float) -> float: + y0, x0, y1, x1 = _window_bounds(center_y, center_x, cell_h, cell_w, component.mask.shape[0] + component.bbox[0], component.mask.shape[1] + component.bbox[1]) + bbox_y0, bbox_x0, _, _ = component.bbox + local_y0 = max(0, y0 - bbox_y0) + local_x0 = max(0, x0 - bbox_x0) + local_y1 = min(component.mask.shape[0], y1 - bbox_y0) + local_x1 = min(component.mask.shape[1], x1 - bbox_x0) + if local_y1 <= local_y0 or local_x1 <= local_x0: + return 0.0 + covered = _integral_sum(component.integral, local_y0, local_x0, local_y1, local_x1) + total = max(1, (local_y1 - local_y0) * (local_x1 - local_x0)) + return float(covered / total) + + +def _nearest_component_pixel(component: _Component, center_y: float, center_x: float) -> tuple[float, float]: + deltas = component.coords.astype(np.float32) - np.asarray([center_y, center_x], dtype=np.float32)[None, :] + best = int(np.argmin(np.sum(deltas * deltas, axis=1))) + return float(component.coords[best, 0]), float(component.coords[best, 1]) + + +def _best_stepped_center( + component: _Component, + center_y: float, + center_x: float, + *, + step_y: float, + step_x: float, + cell_h: float, + cell_w: float, +) -> tuple[float, float, float] | None: + offsets_y = np.asarray([0.0, -step_y * 0.35, step_y * 0.35], dtype=np.float32) + offsets_x = np.asarray([0.0, -step_x * 0.35, step_x * 0.35], dtype=np.float32) + best: tuple[float, float, float] | None = None + for dy in offsets_y: + for dx in offsets_x: + probe_y, probe_x = _nearest_component_pixel(component, center_y + float(dy), center_x + float(dx)) + coverage = _window_component_coverage(component, probe_y, probe_x, cell_h, cell_w) + if best is None or coverage > best[2]: + best = (probe_y, probe_x, coverage) + return best + + +def _extract_components(cluster_map: np.ndarray, alpha_map: np.ndarray, alpha_threshold: float) -> list[_Component]: + valid = alpha_map >= alpha_threshold + visited = np.zeros(valid.shape, dtype=bool) + components: list[_Component] = [] + height, width = valid.shape + for y in range(height): + for x in range(width): + if not valid[y, x] or visited[y, x]: + continue + label = int(cluster_map[y, x]) + queue = deque([(y, x)]) + visited[y, x] = True + coords: list[tuple[int, int]] = [] + min_y = max_y = y + min_x = max_x = x + while queue: + cy, cx = queue.popleft() + coords.append((cy, cx)) + if cy < min_y: + min_y = cy + if cy > max_y: + max_y = cy + if cx < min_x: + min_x = cx + if cx > max_x: + max_x = cx + for dy, dx in ((0, 1), (1, 0), (0, -1), (-1, 0)): + ny = cy + dy + nx = cx + dx + if ny < 0 or ny >= height or nx < 0 or nx >= width: + continue + if visited[ny, nx] or not valid[ny, nx] or int(cluster_map[ny, nx]) != label: + continue + visited[ny, nx] = True + queue.append((ny, nx)) + coords_array = np.asarray(coords, dtype=np.int32) + mask = np.zeros((max_y - min_y + 1, max_x - min_x + 1), dtype=bool) + local = coords_array - np.asarray([min_y, min_x], dtype=np.int32)[None, :] + mask[local[:, 0], local[:, 1]] = True + components.append( + _Component( + label=label, + coords=coords_array, + bbox=(min_y, min_x, max_y + 1, max_x + 1), + mask=mask, + integral=_component_integral(mask), + area=int(coords_array.shape[0]), + centroid_y=float(np.mean(coords_array[:, 0])), + centroid_x=float(np.mean(coords_array[:, 1])), + ) + ) + return components + + +def build_tile_graph_model( + source_rgba: np.ndarray, + *, + inference: InferenceResult, + analysis: SourceAnalysis, + solver_params: SolverHyperParams | None = None, +) -> TileGraphModel: + solver_params = solver_params or SolverHyperParams() + source_reference = build_source_lattice_reference( + source_rgba, + target_width=inference.target_width, + target_height=inference.target_height, + phase_x=inference.phase_x, + phase_y=inference.phase_y, + alpha_threshold=solver_params.alpha_transparent_threshold, + edge_hint=analysis.edge_map, + ) + reference_premul = premultiply(source_reference.sharp_rgba) + source_premul = premultiply(source_rgba) + height, width = source_rgba.shape[:2] + cell_h = height / max(1, inference.target_height) + cell_w = width / max(1, inference.target_width) + target_area = cell_h * cell_w + components = _extract_components(analysis.cluster_map, analysis.alpha_map, solver_params.alpha_transparent_threshold) + components.sort(key=lambda component: (abs(component.area / max(target_area, 1e-4) - 1.0), -component.area)) + + candidate_rgba: list[np.ndarray] = [np.zeros(4, dtype=np.float32)] + candidate_coords: list[tuple[int, int]] = [(0, 0)] + candidate_area_ratio: list[float] = [0.0] + candidate_coverage: list[float] = [1.0] + candidate_deltas: list[np.ndarray] = [np.zeros((4, 4), dtype=np.float32)] + + used_coords: set[tuple[int, int]] = set() + step_y = max(1.0, cell_h) + step_x = max(1.0, cell_w) + for component in components: + area_ratio = float(component.area / max(target_area, 1e-4)) + if area_ratio < solver_params.tile_graph_component_min_area_ratio: + continue + per_component_limit = max(1, min(24, int(np.ceil(max(1.0, area_ratio))))) + seed_y, seed_x = _nearest_component_pixel(component, component.centroid_y, component.centroid_x) + queue = deque([(seed_y, seed_x)]) + local_seen: set[tuple[int, int]] = set() + added = 0 + while queue and added < per_component_limit and len(candidate_rgba) < solver_params.tile_graph_max_candidates: + center_y, center_x = queue.popleft() + coord = _source_to_cell_coord( + center_y, + center_x, + cell_h=cell_h, + cell_w=cell_w, + target_height=inference.target_height, + target_width=inference.target_width, + phase_y=inference.phase_y, + phase_x=inference.phase_x, + ) + if coord in local_seen or coord in used_coords: + continue + local_seen.add(coord) + coverage = _window_component_coverage(component, center_y, center_x, cell_h, cell_w) + if coverage < solver_params.tile_graph_window_coverage_threshold and added > 0: + continue + center_premul = _window_mean_premul(source_premul, center_y, center_x, cell_h, cell_w) + delta_features = np.zeros((4, 4), dtype=np.float32) + for direction, dy, dx in _DIRS: + neighbor = _window_mean_premul( + source_premul, + center_y + dy * step_y, + center_x + dx * step_x, + cell_h, + cell_w, + ) + delta_features[direction] = neighbor - center_premul + candidate_rgba.append(unpremultiply(center_premul[None, None, :])[0, 0].astype(np.float32)) + candidate_coords.append(coord) + candidate_area_ratio.append(area_ratio) + candidate_coverage.append(coverage) + candidate_deltas.append(delta_features) + used_coords.add(coord) + added += 1 + if area_ratio < solver_params.tile_graph_large_component_ratio: + continue + for _, dy, dx in _DIRS: + best = _best_stepped_center( + component, + center_y + dy * step_y, + center_x + dx * step_x, + step_y=step_y, + step_x=step_x, + cell_h=cell_h, + cell_w=cell_w, + ) + if best is None: + continue + probe_y, probe_x, probe_coverage = best + if probe_coverage < solver_params.tile_graph_window_coverage_threshold: + continue + probe_coord = _source_to_cell_coord( + probe_y, + probe_x, + cell_h=cell_h, + cell_w=cell_w, + target_height=inference.target_height, + target_width=inference.target_width, + phase_y=inference.phase_y, + phase_x=inference.phase_x, + ) + if probe_coord not in local_seen and probe_coord not in used_coords: + queue.append((probe_y, probe_x)) + + def add_reference_candidate(coord_y: int, coord_x: int) -> None: + if len(candidate_rgba) >= solver_params.tile_graph_max_candidates: + return + if (coord_y, coord_x) in used_coords: + return + delta_features = np.zeros((4, 4), dtype=np.float32) + center_premul = reference_premul[coord_y, coord_x] + for direction, dy, dx in _DIRS: + neighbor_y = min(max(coord_y + dy, 0), inference.target_height - 1) + neighbor_x = min(max(coord_x + dx, 0), inference.target_width - 1) + delta_features[direction] = reference_premul[neighbor_y, neighbor_x] - center_premul + candidate_rgba.append(source_reference.sharp_rgba[coord_y, coord_x].astype(np.float32)) + candidate_coords.append((coord_y, coord_x)) + candidate_area_ratio.append(1.0) + candidate_coverage.append(1.0) + candidate_deltas.append(delta_features) + used_coords.add((coord_y, coord_x)) + + if len(candidate_rgba) < solver_params.tile_graph_max_candidates: + remaining = solver_params.tile_graph_max_candidates - len(candidate_rgba) + stride = max(1, int(np.sqrt((inference.target_height * inference.target_width) / max(1, remaining)))) + for coord_y in range(0, inference.target_height, stride): + for coord_x in range(0, inference.target_width, stride): + add_reference_candidate(coord_y, coord_x) + if len(candidate_rgba) >= solver_params.tile_graph_max_candidates: + break + if len(candidate_rgba) >= solver_params.tile_graph_max_candidates: + break + if len(candidate_rgba) < solver_params.tile_graph_max_candidates: + priority = ( + source_reference.edge_strength + + np.clip(source_reference.cell_dispersion / max(source_reference.dispersion, 1e-4), 0.0, 4.0) * 0.35 + ) + for flat_index in np.argsort(priority.reshape(-1))[::-1]: + coord_y = int(flat_index // inference.target_width) + coord_x = int(flat_index % inference.target_width) + add_reference_candidate(coord_y, coord_x) + if len(candidate_rgba) >= solver_params.tile_graph_max_candidates: + break + + candidate_rgba_np = np.asarray(candidate_rgba, dtype=np.float32) + candidate_coords_np = np.asarray(candidate_coords, dtype=np.int32) + candidate_area_ratio_np = np.asarray(candidate_area_ratio, dtype=np.float32) + candidate_coverage_np = np.asarray(candidate_coverage, dtype=np.float32) + candidate_deltas_np = np.asarray(candidate_deltas, dtype=np.float32) + + candidate_count = candidate_rgba_np.shape[0] + adjacency_counts = np.zeros((4, candidate_count, candidate_count), dtype=np.float32) + coord_to_indices: dict[tuple[int, int], list[int]] = {} + for index, coord in enumerate(candidate_coords_np.tolist()): + coord_to_indices.setdefault((int(coord[0]), int(coord[1])), []).append(index) + for index, (coord_y, coord_x) in enumerate(candidate_coords_np.tolist()): + for direction, dy, dx in _DIRS: + for neighbor_index in coord_to_indices.get((int(coord_y + dy), int(coord_x + dx)), []): + adjacency_counts[direction, index, neighbor_index] += 1.0 + + pair_penalty = np.zeros((4, candidate_count, candidate_count), dtype=np.float32) + observed_delta = candidate_rgba_np[None, :, :] - candidate_rgba_np[:, None, :] + for direction, _, _ in _DIRS: + row_sum = adjacency_counts[direction].sum(axis=1, keepdims=True) + probability = (adjacency_counts[direction] + 1.0) / (row_sum + float(candidate_count)) + count_penalty = -np.log(probability) + count_penalty = count_penalty / max(float(np.max(count_penalty)), 1e-4) + expected = candidate_deltas_np[:, direction, :][:, None, :] + reverse = candidate_deltas_np[:, _OPPOSITE[direction], :][None, :, :] + delta_penalty = np.mean(np.abs(observed_delta - expected), axis=-1) + reverse_penalty = np.mean(np.abs(-observed_delta - reverse), axis=-1) + delta_penalty = (delta_penalty + reverse_penalty) * 0.5 + delta_penalty = delta_penalty / max(float(np.max(delta_penalty)), 1e-4) + pair_penalty[direction] = ( + count_penalty * solver_params.tile_graph_count_weight + + delta_penalty * solver_params.tile_graph_delta_weight + ).astype(np.float32) + + edge_density = float(np.mean(adjacency_counts > 0.0)) if candidate_count > 0 else 0.0 + return TileGraphModel( + candidate_rgba=candidate_rgba_np, + candidate_coords=candidate_coords_np, + candidate_area_ratio=candidate_area_ratio_np, + candidate_coverage=candidate_coverage_np, + pair_penalty=pair_penalty, + reference_mean_rgba=source_reference.mean_rgba.astype(np.float32), + reference_sharp_rgba=source_reference.sharp_rgba.astype(np.float32), + component_count=len(components), + edge_density=edge_density, + ) + + +def _tile_graph_unary_cost(model: TileGraphModel, solver_params: SolverHyperParams) -> np.ndarray: + height, width = model.reference_sharp_rgba.shape[:2] + color_error = ( + np.mean(np.abs(model.reference_sharp_rgba[..., None, :] - model.candidate_rgba[None, None, :, :]), axis=-1) * 0.6 + + np.mean(np.abs(model.reference_mean_rgba[..., None, :] - model.candidate_rgba[None, None, :, :]), axis=-1) * 0.4 + ) + grid_y, grid_x = np.indices((height, width), dtype=np.float32) + coord_error = ( + np.abs(grid_y[..., None] - model.candidate_coords[None, None, :, 0].astype(np.float32)) + + np.abs(grid_x[..., None] - model.candidate_coords[None, None, :, 1].astype(np.float32)) + ) / max(1.0, float(height + width)) + area_error = np.abs(np.log(np.clip(model.candidate_area_ratio, 1e-4, None) + 1e-4))[None, None, :] + return ( + color_error + + coord_error * solver_params.tile_graph_coordinate_weight + + area_error * solver_params.tile_graph_area_weight + ).astype(np.float32) + + +def _assignment_rgba(model: TileGraphModel, selected: np.ndarray) -> np.ndarray: + return model.candidate_rgba[selected] + + +def _assignment_score(selected: np.ndarray, unary_cost: np.ndarray, pair_penalty: np.ndarray) -> float: + gathered = np.take_along_axis(unary_cost, selected[..., None], axis=2)[..., 0] + score = float(np.mean(gathered)) + if selected.shape[1] > 1: + score += float(np.mean(pair_penalty[_RIGHT, selected[:, :-1], selected[:, 1:]])) + if selected.shape[0] > 1: + score += float(np.mean(pair_penalty[_DOWN, selected[:-1, :], selected[1:, :]])) + return score + + +def optimize_tile_graph( + rgba: np.ndarray, + inference: InferenceResult, + analysis: SourceAnalysis, + steps: int, + seed: int, + device: str, + solver_params: SolverHyperParams | None = None, +) -> tuple[SolverArtifacts, dict[str, Any]]: + del steps, seed, device + solver_params = solver_params or SolverHyperParams() + model = build_tile_graph_model( + rgba, + inference=inference, + analysis=analysis, + solver_params=solver_params, + ) + unary_cost = _tile_graph_unary_cost(model, solver_params) + selected = np.argmin(unary_cost, axis=2).astype(np.int32) + initial_selected = selected.copy() + loss_history = [_assignment_score(selected, unary_cost, model.pair_penalty)] + for _ in range(max(0, int(solver_params.tile_graph_iterations))): + changed = 0 + for parity in (0, 1): + for y in range(selected.shape[0]): + for x in range((y + parity) % 2, selected.shape[1], 2): + cost = unary_cost[y, x].copy() + if x > 0: + cost += model.pair_penalty[_RIGHT, selected[y, x - 1], :] + if x + 1 < selected.shape[1]: + cost += model.pair_penalty[_RIGHT, :, selected[y, x + 1]] + if y > 0: + cost += model.pair_penalty[_DOWN, selected[y - 1, x], :] + if y + 1 < selected.shape[0]: + cost += model.pair_penalty[_DOWN, :, selected[y + 1, x]] + best = int(np.argmin(cost)) + if best != int(selected[y, x]): + selected[y, x] = best + changed += 1 + loss_history.append(_assignment_score(selected, unary_cost, model.pair_penalty)) + if changed == 0: + break + + initial_rgba = _assignment_rgba(model, initial_selected).astype(np.float32) + target_rgba = _assignment_rgba(model, selected).astype(np.float32) + diagnostics = { + "mode": "tile-graph", + "tile_graph_component_count": model.component_count, + "tile_graph_candidate_count": int(model.candidate_rgba.shape[0]), + "tile_graph_edge_density": model.edge_density, + "tile_graph_initial_score": float(loss_history[0]), + "tile_graph_final_score": float(loss_history[-1]), + } + return ( + SolverArtifacts( + target_rgba=target_rgba, + uv_field=_make_regular_uv( + height=rgba.shape[0], + width=rgba.shape[1], + target_height=inference.target_height, + target_width=inference.target_width, + phase_x=inference.phase_x, + phase_y=inference.phase_y, + ), + guidance_strength=np.zeros((inference.target_height, inference.target_width), dtype=np.float32), + initial_rgba=initial_rgba, + loss_history=[float(value) for value in loss_history], + ), + diagnostics, + ) diff --git a/tests/test_tile_graph.py b/tests/test_tile_graph.py new file mode 100644 index 0000000..3800868 --- /dev/null +++ b/tests/test_tile_graph.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from repixelizer.analysis import analyze_source +from repixelizer.baselines import naive_resize_baseline +from repixelizer.io import save_rgba +from repixelizer.metrics import source_lattice_consistency_breakdown +from repixelizer.pipeline import run_pipeline +from repixelizer.synthetic import fake_pixelize, make_emblem +from repixelizer.tile_graph import build_tile_graph_model, optimize_tile_graph +from repixelizer.types import InferenceResult + + +def test_tile_graph_model_extracts_candidates_and_adjacency_from_component_walk() -> None: + lowres = np.zeros((4, 4, 4), dtype=np.float32) + lowres[..., 3] = 1.0 + lowres[:, 1] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32) + source = np.repeat(np.repeat(lowres, 8, axis=0), 8, axis=1) + inference = InferenceResult( + target_width=4, + target_height=4, + phase_x=0.0, + phase_y=0.0, + confidence=1.0, + top_candidates=[], + ) + + model = build_tile_graph_model( + source, + inference=inference, + analysis=analyze_source(source, seed=7), + ) + + assert model.candidate_rgba.shape[0] > 4 + assert model.pair_penalty.shape == (4, model.candidate_rgba.shape[0], model.candidate_rgba.shape[0]) + assert model.edge_density > 0.0 + + +def test_tile_graph_reconstructs_synthetic_thin_feature_better_than_naive() -> None: + source = np.zeros((16, 16, 4), dtype=np.float32) + source[4:12, 8] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32) + source[4, 7:10] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32) + source[11, 7] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32) + source[10, 6] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32) + source[9, 5] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32) + + fake = fake_pixelize( + source, + upscale=9, + phase_x=0.18, + phase_y=-0.12, + blur_radius=0.45, + warp_strength=0.2, + warp_detail=5, + seed=11, + ) + inference = InferenceResult( + target_width=16, + target_height=16, + phase_x=0.18, + phase_y=-0.12, + confidence=1.0, + top_candidates=[], + ) + artifacts, _ = optimize_tile_graph( + fake, + inference=inference, + analysis=analyze_source(fake, seed=7), + steps=0, + seed=7, + device="cpu", + ) + naive = naive_resize_baseline(fake, width=16, height=16) + + tile_score = source_lattice_consistency_breakdown( + fake, + artifacts.target_rgba, + target_width=16, + target_height=16, + phase_x=0.18, + phase_y=-0.12, + )["score"] + naive_score = source_lattice_consistency_breakdown( + fake, + naive, + target_width=16, + target_height=16, + phase_x=0.18, + phase_y=-0.12, + )["score"] + + assert tile_score <= naive_score + + +def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(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) + input_path = tmp_path / "input.png" + output_path = tmp_path / "nested" / "output.png" + diagnostics_dir = tmp_path / "diagnostics" + + save_rgba(input_path, fake) + result = run_pipeline( + input_path, + output_path, + diagnostics_dir=diagnostics_dir, + steps=24, + reconstruction_mode="tile-graph", + ) + + import json + + run_json = json.loads((diagnostics_dir / "run.json").read_text(encoding="utf-8")) + assert output_path.exists() + assert run_json["settings"]["reconstruction_mode"] == "tile-graph" + assert run_json["reconstruction"]["mode"] == "tile-graph" + assert run_json["reconstruction"]["tile_graph_candidate_count"] > 0 + assert set(run_json["source_fidelity"].keys()) == {"snap_initial", "solver_target", "final_output"} + assert result.diagnostics["reconstruction"]["mode"] == "tile-graph"