From b3a2d72cc1a30d2ab73a0d9e88f97b1afe59ac57 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 23 Apr 2026 04:18:24 +0100 Subject: [PATCH] Split tile-graph off shared prep spine --- README.md | 2 + docs/implementation-plan.md | 7 +- docs/tile-graph-algorithm-map.md | 68 ++++++------- src/repixelizer/analysis.py | 82 +++++++--------- src/repixelizer/continuous.py | 4 +- src/repixelizer/pipeline.py | 24 +++-- src/repixelizer/source_reference.py | 146 +++++++++++++++++----------- src/repixelizer/tile_graph.py | 109 ++++++++------------- src/repixelizer/types.py | 31 +++--- tests/test_analysis.py | 21 ++-- tests/test_continuous.py | 6 +- tests/test_pipeline.py | 84 ++++++++++++++-- tests/test_source_reference.py | 33 ++++--- tests/test_tile_graph.py | 73 +++++++------- 14 files changed, 389 insertions(+), 301 deletions(-) diff --git a/README.md b/README.md index 5853a1f..2f4ce90 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,8 @@ Recent adjacency-focused status: Current tile-graph status: - it now keeps candidates strictly local to each output coord and only uses literal source-pixel colors, so it no longer invents averaged patch colors or borrow labels from distant regions +- tile-graph prep is now a separate edge-only path: it no longer hauls continuous-only cluster analysis or the rich shared lattice reference through its core loop +- tile-graph now builds against its own lean source reference and keeps diagnostics/cache metadata outside the solver model itself - candidate seeding now starts from source-side atomic regions projected onto the output lattice, layered together with the old sharp/edge anchors instead of replacing them outright - the new deep-dive map in `docs/tile-graph-algorithm-map.md` walks the tile-graph path end to end and records the current fixed-lattice failure mode: the garbled `126x126` badge output is already bad at the tile-graph initial assignment stage, not because the fixed-lattice pipeline wrapper is corrupting it - the extraction stage now has an explicit empty-cell overlap fill pass, so output cells that actually contain opaque sampled source pixels no longer silently lose their extracted region bucket under the fixed badge lattice diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index 52e6213..c4fbd6d 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -19,11 +19,12 @@ Implemented: Verified: - local editable install in a dedicated venv -- full test suite currently passing (`66 passed`) +- full test suite currently passing (`70 passed`) - compare-mode smoke run against a real emblem image - the cleaned real badge fixture in `tests/fixtures/real/ai-badge-cleaned.png` still beats naive resize on source-lattice consistency on the continuous path - the latest full-emblem tile-graph atomic probe under `artifacts/full-emblem-tile-graph-atomic-v3-cuda/` lands at `0.0224` source-fidelity, beating the older full-emblem tile-graph CUDA baseline at `0.0283` - the new algorithm map in `docs/tile-graph-algorithm-map.md` confirms that the pinned `126x126` tile-graph badge collapse is already present in the tile-graph initial assignment; the fixed-lattice pipeline path itself is reproducing that bad state faithfully rather than introducing it, and the newer extraction coverage fix now guarantees that occupied output cells are not silently losing their source-region bucket under the corrected full-size lattice mapping +- the new architectural split keeps the fixed synthetic emblem baselines byte-identical on both engines: the pinned tile-graph output still hashes to `31e3bc2c...6ae9` and the pinned continuous output still hashes to `404748af...7903` ## Status after adjacency-first pass @@ -64,9 +65,13 @@ What landed: - the experimental `hybrid` path and its geometry-prior wiring have been removed - `tile_graph.py` no longer carries geometry-prior fields through `TileGraphModel` or tile-graph unary scoring - `tile-graph` now skips source clustering entirely and uses only edge analysis plus sharp/edge lattice references during unary scoring +- `analysis.py` now has separate prep paths: continuous gets edge plus cluster analysis, while tile-graph gets an edge-only scout report +- `source_reference.py` now has a dedicated `TileGraphSourceReference` instead of forcing tile-graph through the larger continuous/metrics reference object +- `TileGraphModel` is now solver-only state; cache metadata and build diagnostics live in a separate `TileGraphBuildStats` sidecar - `cli.py` now exposes only `continuous` and `tile-graph` as reconstruction engines - `tests/test_pipeline.py` now checks that low-confidence tile-graph runs skip rerank probes instead of rebuilding probe candidates - `tests/test_analysis.py` and `tests/test_tile_graph.py` now cover the new edge-only analysis path +- `tests/test_pipeline.py` now locks in fixed continuous and tile-graph output hashes from a pre-refactor baseline so this architectural cut cannot silently drift behavior - `docs/tile-graph-algorithm-map.md` now describes the smaller machine directly rather than documenting the removed hybrid sidecar path Current implementation note: diff --git a/docs/tile-graph-algorithm-map.md b/docs/tile-graph-algorithm-map.md index 6c7f625..0fcb4f6 100644 --- a/docs/tile-graph-algorithm-map.md +++ b/docs/tile-graph-algorithm-map.md @@ -55,16 +55,16 @@ That sentence matters because it explains the failure mode. When the lattice is flowchart TD A["source_rgba (H x W x 4)"] --> B["run_pipeline(...)"] B --> C["infer_lattice(...) or infer_fixed_lattice(...)"] - B --> D["analyze_source(...)"] + B --> D["analyze_tile_graph_source(...)"] C --> E["InferenceResult"] - D --> F["SourceAnalysis"] + D --> F["TileGraphSourceAnalysis"] E --> G["phase rerank probes (optional)"] F --> G G --> H["optimize_tile_graph(...)"] - H --> I["build_source_lattice_reference(...)"] + H --> I["build_tile_graph_source_reference(...)"] H --> J["_segment_atomic_source_regions(...)"] J --> K["_extract_source_region_tiles(...)"] - I --> L["per-output-cell references"] + I --> L["per-output-cell sharp and edge anchors"] K --> M["region buckets per output coord"] L --> N["build_tile_graph_model(...)"] M --> N @@ -110,48 +110,34 @@ Plain-language picture: This is the foreman's decision about the size and offset of the graph paper. Every later worker trusts this sheet. -### `SourceAnalysis` +### `TileGraphSourceAnalysis` - File: `src/repixelizer/types.py` - Fields: - `edge_map`: per-source-pixel edge strength - - `cluster_map`: k-means color labels - - `cluster_centers`: color centroids - - `alpha_map`: source alpha - - `cluster_preview`: diagnostic preview Important note: -- `cluster_map` is diagnostic/coarse analysis data -- it is not the current tile-graph region ownership model +- this is only the edge scout report +- tile-graph no longer drags cluster labels, alpha mirrors, or preview images through its core contract Plain-language picture: This is a scouting report about where the mural has edges and rough color families. It is binoculars, not ownership papers. -Current tile-graph note: - -- the tile-graph path now asks only for the edge scout report -- cluster labels and cluster preview are left empty on that path because the tile-graph machine does not use them - -### `SourceLatticeReference` +### `TileGraphSourceReference` - File: `src/repixelizer/types.py` -- Meaning: lattice-conditioned summary of the source under one exact `(target_width, target_height, phase_x, phase_y)` +- Meaning: tile-graph's lattice-conditioned anchor pack under one exact `(target_width, target_height, phase_x, phase_y)` - Fields: - - `mean_rgba`: average premultiplied color per inferred cell, then unpremultiplied - `sharp_rgba`: one chosen exemplar pixel per inferred cell - - `dispersion`: average per-pixel deviation from cell mean - - `lattice_indices`: source-pixel to inferred-cell assignment map - - `cell_dispersion`, `cell_counts`, `cell_support`, `cell_alpha_max` - `sharp_x`, `sharp_y` - `edge_peak_x`, `edge_peak_y` - - `edge_strength`, `edge_grad_x`, `edge_grad_y` - - `delta_x`, `delta_y`, `delta_diag`, `delta_anti` + - `edge_strength` Plain-language picture: -This is the mural after the graph paper has already been pressed onto it. Each square now gets a little dossier: its average color, one "best" pixel, one edge pixel, and a few notes about how it differs from neighbors. +This is the mural after the graph paper has already been pressed onto it, but only the notes tile-graph actually reads are kept: one sharp pixel, one edge pixel, and how strong that edge is. ### `TileGraphModel` @@ -159,14 +145,24 @@ This is the mural after the graph paper has already been pressed onto it. Each s - Meaning: the fully built discrete candidate model for one source + one fixed lattice - Fields: - `candidate_rgba`: all candidate pixel colors - - `candidate_coords`: output coord `(y, x)` each candidate belongs to - `candidate_area_ratio`: how much of a cell-sized source window the source region covered - `candidate_coverage`: clipped version of area coverage - `candidate_deltas`: expected RGBA deltas to right/down/left/up neighbors sampled one cell away in source space - `cell_candidate_offsets`, `cell_candidate_indices`: CSR-like mapping from output cells to candidate rows - `reference_sharp_rgba`, `reference_edge_rgba` - `edge_strength` - - `component_count`, `edge_density`, `average_choices` + +### `TileGraphBuildStats` + +- File: `src/repixelizer/tile_graph.py` +- Meaning: diagnostics and cache metadata carried alongside the model rather than inside it +- Fields: + - `component_count` + - `candidate_count` + - `edge_density` + - `average_choices` + - `model_device` + - `cache_hit` Plain-language picture: @@ -185,7 +181,7 @@ Main variables: - `source`: loaded RGBA source image - `fixed_dims`: either `None` or exact `(target_width, target_height)` resolved from CLI - `inference`: one `InferenceResult` -- `analysis`: one `SourceAnalysis` +- `analysis`: one `TileGraphSourceAnalysis` - `solver`: `SolverArtifacts` from continuous or tile-graph - `cleanup`: `CleanupArtifacts` - `output_rgba`: final saved output @@ -268,25 +264,21 @@ File: `src/repixelizer/analysis.py` Entry point: -- `analyze_source(rgba, seed, cluster_count=6, device=None, include_clusters=True)` +- `analyze_tile_graph_source(rgba, device=None)` Outputs: - `edge_map`: edge strength from luminance plus alpha differences -- `cluster_map`: k-means labels on opaque pixels -- `cluster_centers`: k-means centroids -- `alpha_map` -- `cluster_preview` Plain-language picture: -This is the scout walking the mural with a flashlight, marking where the paint changes abruptly and scribbling rough color families on a clipboard. +This is the scout walking the mural with a flashlight and doing only one job: marking where the paint changes abruptly. Important reality check: -- tile-graph no longer uses `cluster_map` as its actual tile ownership map +- tile-graph no longer uses coarse cluster labels as its ownership map - the tile-graph region ownership is built later by connected components over source pixels using color/alpha thresholds -- after the pruning pass, tile-graph can skip cluster assignment entirely and still run normally because it only needs `edge_map` +- after the seam split, tile-graph carries only `edge_map` through this stage ## Stage 4: Optional Phase Rerank @@ -565,8 +557,6 @@ After candidate selection, the model stores: - `candidate_rgba` - literal source-pixel colors, one row per candidate -- `candidate_coords` - - which output cell each candidate belongs to - `candidate_area_ratio` - window size match score proxy - `candidate_coverage` @@ -579,6 +569,8 @@ After candidate selection, the model stores: - source pixel at the per-cell edge peak - `edge_strength` +The diagnostic counters that used to live here now sit beside the model in `TileGraphBuildStats`, so the solver only carries solver state. + So the final discrete solver is not operating on raw tiles alone. It is operating on: diff --git a/src/repixelizer/analysis.py b/src/repixelizer/analysis.py index 244c97b..b992c0e 100644 --- a/src/repixelizer/analysis.py +++ b/src/repixelizer/analysis.py @@ -3,7 +3,7 @@ from __future__ import annotations import numpy as np from .metrics import luminance -from .types import SourceAnalysis +from .types import ContinuousSourceAnalysis, TileGraphSourceAnalysis def _require_torch(): @@ -104,23 +104,12 @@ def _kmeans_torch(torch, data_t, k: int, seed: int, iterations: int = 12): return centers_t, labels_t -def _colorize_clusters(cluster_map: np.ndarray, centers: np.ndarray) -> np.ndarray: - if centers.size == 0: - return np.zeros((*cluster_map.shape, 4), dtype=np.float32) - palette = np.clip(centers[:, :3], 0.0, 1.0) - preview = palette[np.clip(cluster_map, 0, len(palette) - 1)] - alpha = np.ones((*cluster_map.shape, 1), dtype=np.float32) - return np.concatenate([preview, alpha], axis=-1) - - -def analyze_source( +def analyze_continuous_source( rgba: np.ndarray, seed: int, cluster_count: int = 6, device: str | None = None, - *, - include_clusters: bool = True, -) -> SourceAnalysis: +) -> ContinuousSourceAnalysis: if device is not None: torch = _require_torch() resolved_device = _resolve_device(torch, device) @@ -128,45 +117,46 @@ def analyze_source( edge_map_t = _compute_edge_map_torch(torch, rgba_t) alpha_t = rgba_t[..., 3] cluster_map_t = torch.full(alpha_t.shape, -1, device=resolved_device, dtype=torch.long) - centers_t = torch.zeros((0, 4), device=resolved_device, dtype=torch.float32) - preview_t = torch.zeros((*alpha_t.shape, 4), device=resolved_device, dtype=torch.float32) - if include_clusters: - opaque_t = alpha_t > 0.05 - if bool(torch.any(opaque_t).item()): - samples_t = rgba_t[opaque_t] - k = min(cluster_count, max(2, int(np.sqrt(int(samples_t.shape[0]) // 64 + 1)))) - centers_t, labels_t = _kmeans_torch(torch, samples_t, k=k, seed=seed) - cluster_map_t[opaque_t] = labels_t - preview_palette_t = centers_t[:, :3].clamp(0.0, 1.0) if centers_t.numel() else torch.zeros((1, 3), device=resolved_device, dtype=torch.float32) - preview_indices_t = torch.clamp(cluster_map_t, 0, max(0, preview_palette_t.shape[0] - 1)) - preview_t = preview_palette_t[preview_indices_t] - preview_t = torch.cat([preview_t, alpha_t[..., None]], dim=-1) - return SourceAnalysis( + opaque_t = alpha_t > 0.05 + if bool(torch.any(opaque_t).item()): + samples_t = rgba_t[opaque_t] + k = min(cluster_count, max(2, int(np.sqrt(int(samples_t.shape[0]) // 64 + 1)))) + _centers_t, labels_t = _kmeans_torch(torch, samples_t, k=k, seed=seed) + cluster_map_t[opaque_t] = labels_t + return ContinuousSourceAnalysis( edge_map=edge_map_t.detach().cpu().numpy().astype(np.float32), cluster_map=cluster_map_t.detach().cpu().numpy().astype(np.int32), - cluster_centers=centers_t.detach().cpu().numpy().astype(np.float32), - alpha_map=alpha_t.detach().cpu().numpy().astype(np.float32), - cluster_preview=preview_t.detach().cpu().numpy().astype(np.float32), ) edge_map = _compute_edge_map(rgba) alpha = rgba[..., 3] cluster_map = np.full(alpha.shape, -1, dtype=np.int32) - centers = np.zeros((0, 4), dtype=np.float32) - preview = np.zeros((*alpha.shape, 4), dtype=np.float32) - if include_clusters: - opaque = alpha > 0.05 - if np.any(opaque): - samples = rgba[opaque] - k = min(cluster_count, max(2, int(np.sqrt(samples.shape[0] // 64 + 1)))) - centers, labels = _kmeans(samples, k=k, seed=seed) - cluster_map[opaque] = labels - preview = _colorize_clusters(np.maximum(cluster_map, 0), centers if centers.size else np.zeros((1, 4), dtype=np.float32)) - preview[..., 3] = alpha - return SourceAnalysis( + opaque = alpha > 0.05 + if np.any(opaque): + samples = rgba[opaque] + k = min(cluster_count, max(2, int(np.sqrt(samples.shape[0] // 64 + 1)))) + _centers, labels = _kmeans(samples, k=k, seed=seed) + cluster_map[opaque] = labels + return ContinuousSourceAnalysis( edge_map=edge_map, cluster_map=cluster_map, - cluster_centers=centers, - alpha_map=alpha.astype(np.float32), - cluster_preview=preview, + ) + + +def analyze_tile_graph_source( + rgba: np.ndarray, + *, + device: str | None = None, +) -> TileGraphSourceAnalysis: + if device is not None: + torch = _require_torch() + resolved_device = _resolve_device(torch, device) + rgba_t = torch.from_numpy(rgba).to(device=resolved_device, dtype=torch.float32) + edge_map_t = _compute_edge_map_torch(torch, rgba_t) + return TileGraphSourceAnalysis( + edge_map=edge_map_t.detach().cpu().numpy().astype(np.float32), + ) + + return TileGraphSourceAnalysis( + edge_map=_compute_edge_map(rgba), ) diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py index 6cb694b..1832f6d 100644 --- a/src/repixelizer/continuous.py +++ b/src/repixelizer/continuous.py @@ -6,7 +6,7 @@ from .io import premultiply, unpremultiply from .metrics import source_lattice_consistency_breakdown from .params import SolverHyperParams from .source_reference import build_source_lattice_reference -from .types import InferenceResult, SolverArtifacts, SourceAnalysis +from .types import ContinuousSourceAnalysis, InferenceResult, SolverArtifacts def _require_torch(): @@ -1358,7 +1358,7 @@ def _discrete_refine_output( def optimize_uv_field( rgba: np.ndarray, inference: InferenceResult, - analysis: SourceAnalysis, + analysis: ContinuousSourceAnalysis, steps: int, seed: int, device: str, diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index d7dcda7..f80516a 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -6,7 +6,7 @@ from typing import Any import numpy as np -from .analysis import analyze_source +from .analysis import analyze_continuous_source, analyze_tile_graph_source from .continuous import optimize_uv_field from .diagnostics import ( summarize_run, @@ -29,7 +29,7 @@ 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 +from .types import ContinuousSourceAnalysis, InferenceResult, RunResult, TileGraphSourceAnalysis def run_pipeline( @@ -79,11 +79,11 @@ def run_pipeline( device=device, ) inference_mode = "fixed" - analysis = analyze_source( + analysis = _prepare_analysis( source, seed=seed, - device=device if reconstruction_mode == "tile-graph" else None, - include_clusters=reconstruction_mode != "tile-graph", + device=device, + reconstruction_mode=reconstruction_mode, ) inference = _select_phase_candidate( source, @@ -139,8 +139,6 @@ def run_pipeline( ) write_alpha_preview(diagnostics_path / "alpha-preview.png", source, output_rgba) write_heatmap(diagnostics_path / "noise-heatmap.png", cleanup.isolated_heatmap) - if analysis.cluster_centers.size: - save_rgba(diagnostics_path / "cluster-preview.png", analysis.cluster_preview) run_json = summarize_run(result) run_json["inference"] = inference_to_json(inference) run_json["settings"] = { @@ -167,6 +165,18 @@ def run_pipeline( return result +def _prepare_analysis( + source: np.ndarray, + *, + seed: int, + device: str, + reconstruction_mode: str, +) -> ContinuousSourceAnalysis | TileGraphSourceAnalysis: + if reconstruction_mode == "tile-graph": + return analyze_tile_graph_source(source, device=device) + return analyze_continuous_source(source, seed=seed) + + def _select_phase_candidate( source: np.ndarray, inference: InferenceResult, diff --git a/src/repixelizer/source_reference.py b/src/repixelizer/source_reference.py index fcc067b..1fee6c0 100644 --- a/src/repixelizer/source_reference.py +++ b/src/repixelizer/source_reference.py @@ -3,7 +3,7 @@ from __future__ import annotations import numpy as np from .io import premultiply, unpremultiply -from .types import SourceLatticeReference +from .types import SourceLatticeReference, TileGraphSourceReference def _require_torch(): @@ -95,16 +95,6 @@ def _lattice_indices_torch( return y_idx[:, None] * target_width + x_idx[None, :] -def _reference_deltas(reference: np.ndarray) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]: - height = reference.shape[0] - width = reference.shape[1] - delta_x = reference[:, 1:, :] - reference[:, :-1, :] if width > 1 else None - delta_y = reference[1:, :, :] - reference[:-1, :, :] if height > 1 else None - delta_diag = reference[1:, 1:, :] - reference[:-1, :-1, :] if height > 1 and width > 1 else None - delta_anti = reference[1:, :-1, :] - reference[:-1, 1:, :] if height > 1 and width > 1 else None - return delta_x, delta_y, delta_diag, delta_anti - - def _unpremultiply_torch(torch, premul_t): alpha = premul_t[..., 3:4].clamp_min(1e-6) rgb = torch.where(alpha > 1e-6, premul_t[..., :3] / alpha, torch.zeros_like(premul_t[..., :3])) @@ -147,6 +137,65 @@ def build_source_lattice_reference( edge_grad_y_hint: np.ndarray | None = None, device: str | None = None, ) -> SourceLatticeReference: + payload = _build_reference_payload( + source_rgba, + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, + alpha_threshold=alpha_threshold, + edge_hint=edge_hint, + edge_grad_x_hint=edge_grad_x_hint, + edge_grad_y_hint=edge_grad_y_hint, + device=device, + ) + return SourceLatticeReference(**payload) + + +def build_tile_graph_source_reference( + source_rgba: np.ndarray, + *, + target_width: int, + target_height: int, + phase_x: float, + phase_y: float, + alpha_threshold: float = 0.05, + edge_hint: np.ndarray | None = None, + device: str | None = None, +) -> TileGraphSourceReference: + payload = _build_reference_payload( + source_rgba, + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, + alpha_threshold=alpha_threshold, + edge_hint=edge_hint, + device=device, + ) + return TileGraphSourceReference( + sharp_rgba=payload["sharp_rgba"], + sharp_x=payload["sharp_x"], + sharp_y=payload["sharp_y"], + edge_peak_x=payload["edge_peak_x"], + edge_peak_y=payload["edge_peak_y"], + edge_strength=payload["edge_strength"], + ) + + +def _build_reference_payload( + source_rgba: np.ndarray, + *, + target_width: int, + target_height: int, + phase_x: float, + phase_y: float, + alpha_threshold: float = 0.05, + edge_hint: np.ndarray | None = None, + edge_grad_x_hint: np.ndarray | None = None, + edge_grad_y_hint: np.ndarray | None = None, + device: str | None = None, +) -> dict[str, np.ndarray | float]: if device is not None: torch = _require_torch() resolved_device = _resolve_device(torch, device) @@ -261,29 +310,21 @@ def build_source_lattice_reference( expected_pixels_per_cell = float(source_rgba.shape[0] * source_rgba.shape[1]) / float(cell_count) cell_support_t = (counts_t / max(expected_pixels_per_cell, 1.0)).reshape(target_height, target_width) - sharp_premul_np = sharp_premul_t.detach().cpu().numpy().astype(np.float32) - delta_x, delta_y, delta_diag, delta_anti = _reference_deltas(sharp_premul_np) - return SourceLatticeReference( - mean_rgba=mean_rgba_t.detach().cpu().numpy().astype(np.float32), - sharp_rgba=sharp_rgba_t.detach().cpu().numpy().astype(np.float32), - dispersion=dispersion, - lattice_indices=indices_t.detach().cpu().numpy().astype(np.int32), - cell_dispersion=cell_dispersion_t.detach().cpu().numpy().astype(np.float32), - cell_counts=counts_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), - cell_support=cell_support_t.detach().cpu().numpy().astype(np.float32), - cell_alpha_max=cell_alpha_max_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), - sharp_x=sharp_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), - sharp_y=sharp_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), - edge_peak_x=edge_peak_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), - edge_peak_y=edge_peak_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), - edge_strength=edge_strength_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), - edge_grad_x=edge_grad_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), - edge_grad_y=edge_grad_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), - delta_x=delta_x.astype(np.float32) if delta_x is not None else None, - delta_y=delta_y.astype(np.float32) if delta_y is not None else None, - delta_diag=delta_diag.astype(np.float32) if delta_diag is not None else None, - delta_anti=delta_anti.astype(np.float32) if delta_anti is not None else None, - ) + return { + "mean_rgba": mean_rgba_t.detach().cpu().numpy().astype(np.float32), + "sharp_rgba": sharp_rgba_t.detach().cpu().numpy().astype(np.float32), + "dispersion": dispersion, + "cell_dispersion": cell_dispersion_t.detach().cpu().numpy().astype(np.float32), + "cell_support": cell_support_t.detach().cpu().numpy().astype(np.float32), + "cell_alpha_max": cell_alpha_max_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), + "sharp_x": sharp_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), + "sharp_y": sharp_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), + "edge_peak_x": edge_peak_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), + "edge_peak_y": edge_peak_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32), + "edge_strength": edge_strength_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), + "edge_grad_x": edge_grad_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), + "edge_grad_y": edge_grad_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32), + } height = source_rgba.shape[0] width = source_rgba.shape[1] @@ -394,25 +435,18 @@ def build_source_lattice_reference( 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) - delta_x, delta_y, delta_diag, delta_anti = _reference_deltas(sharp_premul) - return SourceLatticeReference( - mean_rgba=mean_rgba.astype(np.float32), - sharp_rgba=sharp_rgba.astype(np.float32), - dispersion=dispersion, - lattice_indices=indices.astype(np.int32), - cell_dispersion=cell_dispersion.astype(np.float32), - cell_counts=counts.reshape(target_height, target_width).astype(np.float32), - cell_support=cell_support.astype(np.float32), - cell_alpha_max=cell_alpha_max.reshape(target_height, target_width).astype(np.float32), - sharp_x=sharp_x_flat.reshape(target_height, target_width).astype(np.int32), - sharp_y=sharp_y_flat.reshape(target_height, target_width).astype(np.int32), - edge_peak_x=edge_peak_x_flat.reshape(target_height, target_width).astype(np.int32), - edge_peak_y=edge_peak_y_flat.reshape(target_height, target_width).astype(np.int32), - edge_strength=edge_strength_flat.reshape(target_height, target_width).astype(np.float32), - edge_grad_x=edge_grad_x_flat.reshape(target_height, target_width).astype(np.float32), - edge_grad_y=edge_grad_y_flat.reshape(target_height, target_width).astype(np.float32), - delta_x=delta_x.astype(np.float32) if delta_x is not None else None, - delta_y=delta_y.astype(np.float32) if delta_y is not None else None, - delta_diag=delta_diag.astype(np.float32) if delta_diag is not None else None, - delta_anti=delta_anti.astype(np.float32) if delta_anti is not None else None, - ) + return { + "mean_rgba": mean_rgba.astype(np.float32), + "sharp_rgba": sharp_rgba.astype(np.float32), + "dispersion": dispersion, + "cell_dispersion": cell_dispersion.astype(np.float32), + "cell_support": cell_support.astype(np.float32), + "cell_alpha_max": cell_alpha_max.reshape(target_height, target_width).astype(np.float32), + "sharp_x": sharp_x_flat.reshape(target_height, target_width).astype(np.int32), + "sharp_y": sharp_y_flat.reshape(target_height, target_width).astype(np.int32), + "edge_peak_x": edge_peak_x_flat.reshape(target_height, target_width).astype(np.int32), + "edge_peak_y": edge_peak_y_flat.reshape(target_height, target_width).astype(np.int32), + "edge_strength": edge_strength_flat.reshape(target_height, target_width).astype(np.float32), + "edge_grad_x": edge_grad_x_flat.reshape(target_height, target_width).astype(np.float32), + "edge_grad_y": edge_grad_y_flat.reshape(target_height, target_width).astype(np.float32), + } diff --git a/src/repixelizer/tile_graph.py b/src/repixelizer/tile_graph.py index 8a7d9b5..52480b0 100644 --- a/src/repixelizer/tile_graph.py +++ b/src/repixelizer/tile_graph.py @@ -9,8 +9,8 @@ import numpy as np from .metrics import source_lattice_consistency_breakdown from .params import SolverHyperParams -from .source_reference import build_source_lattice_reference -from .types import InferenceResult, SolverArtifacts, SourceAnalysis +from .source_reference import build_tile_graph_source_reference +from .types import InferenceResult, SolverArtifacts, TileGraphSourceAnalysis _RIGHT = 0 _DOWN = 1 @@ -54,7 +54,6 @@ def _resolve_device(torch, requested: str) -> str: @dataclass(slots=True) class TileGraphModel: candidate_rgba: np.ndarray - candidate_coords: np.ndarray candidate_area_ratio: np.ndarray candidate_coverage: np.ndarray candidate_deltas: np.ndarray @@ -63,7 +62,12 @@ class TileGraphModel: reference_sharp_rgba: np.ndarray reference_edge_rgba: np.ndarray edge_strength: np.ndarray + + +@dataclass(slots=True) +class TileGraphBuildStats: component_count: int + candidate_count: int edge_density: float average_choices: float model_device: str @@ -72,10 +76,12 @@ class TileGraphModel: _TILE_GRAPH_MODEL_CACHE_MAX_ENTRIES = 8 _TILE_GRAPH_MODEL_CACHE: OrderedDict[tuple[object, ...], TileGraphModel] = OrderedDict() +_TILE_GRAPH_BUILD_STATS_CACHE: OrderedDict[tuple[object, ...], TileGraphBuildStats] = OrderedDict() def clear_tile_graph_model_cache() -> None: _TILE_GRAPH_MODEL_CACHE.clear() + _TILE_GRAPH_BUILD_STATS_CACHE.clear() def _hash_array(array: np.ndarray) -> str: @@ -118,19 +124,28 @@ def _tile_graph_model_cache_key( ) -def _get_cached_tile_graph_model(cache_key: tuple[object, ...]) -> TileGraphModel | None: +def _get_cached_tile_graph_model(cache_key: tuple[object, ...]) -> tuple[TileGraphModel, TileGraphBuildStats] | None: cached = _TILE_GRAPH_MODEL_CACHE.get(cache_key) - if cached is None: + cached_stats = _TILE_GRAPH_BUILD_STATS_CACHE.get(cache_key) + if cached is None or cached_stats is None: return None _TILE_GRAPH_MODEL_CACHE.move_to_end(cache_key) - return cached + _TILE_GRAPH_BUILD_STATS_CACHE.move_to_end(cache_key) + return cached, cached_stats -def _store_cached_tile_graph_model(cache_key: tuple[object, ...], model: TileGraphModel) -> None: +def _store_cached_tile_graph_model( + cache_key: tuple[object, ...], + model: TileGraphModel, + stats: TileGraphBuildStats, +) -> None: _TILE_GRAPH_MODEL_CACHE[cache_key] = model _TILE_GRAPH_MODEL_CACHE.move_to_end(cache_key) + _TILE_GRAPH_BUILD_STATS_CACHE[cache_key] = stats + _TILE_GRAPH_BUILD_STATS_CACHE.move_to_end(cache_key) while len(_TILE_GRAPH_MODEL_CACHE) > _TILE_GRAPH_MODEL_CACHE_MAX_ENTRIES: - _TILE_GRAPH_MODEL_CACHE.popitem(last=False) + old_key, _old_model = _TILE_GRAPH_MODEL_CACHE.popitem(last=False) + _TILE_GRAPH_BUILD_STATS_CACHE.pop(old_key, None) def _make_regular_uv(height: int, width: int, target_height: int, target_width: int, phase_x: float, phase_y: float) -> np.ndarray: @@ -147,10 +162,6 @@ def _make_regular_uv(height: int, width: int, target_height: int, target_width: return uv.astype(np.float32) -def _rgba_luminance(rgba: np.ndarray) -> np.ndarray: - return rgba[..., 0] * 0.2126 + rgba[..., 1] * 0.7152 + rgba[..., 2] * 0.0722 - - def _component_axis_stats(member_x: np.ndarray, member_y: np.ndarray) -> tuple[float, float, float, float, float]: if member_x.size <= 1: return 1.0, 0.0, 0.0, 0.0, 0.0 @@ -871,10 +882,10 @@ def build_tile_graph_model( source_rgba: np.ndarray, *, inference: InferenceResult, - analysis: SourceAnalysis, + analysis: TileGraphSourceAnalysis, solver_params: SolverHyperParams | None = None, device: str = "cpu", -) -> TileGraphModel: +) -> tuple[TileGraphModel, TileGraphBuildStats]: torch = _require_torch() resolved_device = _resolve_device(torch, device) solver_params = solver_params or SolverHyperParams() @@ -886,8 +897,9 @@ def build_tile_graph_model( ) cached_model = _get_cached_tile_graph_model(cache_key) if cached_model is not None: - return replace(cached_model, cache_hit=True) - source_reference = build_source_lattice_reference( + model, stats = cached_model + return model, replace(stats, cache_hit=True) + source_reference = build_tile_graph_source_reference( source_rgba, target_width=inference.target_width, target_height=inference.target_height, @@ -974,7 +986,6 @@ def build_tile_graph_model( ) candidate_linear: list[int] = [] - candidate_coords: list[tuple[int, int]] = [] candidate_area_ratio: list[float] = [] candidate_coverage: list[float] = [] candidate_source_x: list[float] = [] @@ -995,7 +1006,6 @@ def build_tile_graph_model( return seen_pixels.add(pixel_index) candidate_linear.append(pixel_index) - candidate_coords.append((coord_y, coord_x)) candidate_area_ratio.append(float(max(area_ratio, 1e-3))) candidate_coverage.append(float(np.clip(coverage, 0.0, 1.0))) candidate_source_x.append(float(source_x)) @@ -1052,7 +1062,6 @@ def build_tile_graph_model( candidate_deltas_t[:, direction] = source_t[sample_y_t, sample_x_t] - center_rgba_t candidate_rgba_np = center_rgba_t.detach().cpu().numpy().astype(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 = candidate_deltas_t.detach().cpu().numpy().astype(np.float32) @@ -1068,7 +1077,6 @@ def build_tile_graph_model( ) base_model = TileGraphModel( candidate_rgba=candidate_rgba_np, - candidate_coords=candidate_coords_np, candidate_area_ratio=candidate_area_ratio_np, candidate_coverage=candidate_coverage_np, candidate_deltas=candidate_deltas_np, @@ -1077,14 +1085,17 @@ def build_tile_graph_model( reference_sharp_rgba=source_reference.sharp_rgba.astype(np.float32), reference_edge_rgba=edge_rgba_np.reshape(inference.target_height, inference.target_width, -1).astype(np.float32), edge_strength=source_reference.edge_strength.astype(np.float32), + ) + base_stats = TileGraphBuildStats( component_count=component_total, + candidate_count=int(candidate_rgba_np.shape[0]), edge_density=edge_density, average_choices=average_choices, model_device=resolved_device, cache_hit=False, ) - _store_cached_tile_graph_model(cache_key, base_model) - return replace(base_model, cache_hit=False) + _store_cached_tile_graph_model(cache_key, base_model, base_stats) + return base_model, base_stats def _cell_candidate_span(model: TileGraphModel, y: int, x: int) -> tuple[int, int]: @@ -1273,48 +1284,10 @@ def _assignment_rgba(model: TileGraphModel, selected: np.ndarray) -> np.ndarray: return model.candidate_rgba[selected] -def _pair_penalty(model: TileGraphModel, left_idx: int, right_idx: int, direction: int, solver_params: SolverHyperParams) -> float: - observed_delta = model.candidate_rgba[right_idx] - model.candidate_rgba[left_idx] - expected = model.candidate_deltas[left_idx, direction] - reverse = model.candidate_deltas[right_idx, _OPPOSITE[direction]] - delta_penalty = np.mean(np.abs(observed_delta - expected)) - reverse_penalty = np.mean(np.abs(-observed_delta - reverse)) - return float((delta_penalty + reverse_penalty) * 0.5 * solver_params.tile_graph_delta_weight) - - -def _assignment_score(selected: np.ndarray, model: TileGraphModel, unary_cost: np.ndarray, solver_params: SolverHyperParams) -> float: - height, width = selected.shape - gathered = np.zeros((height, width), dtype=np.float32) - for y in range(height): - for x in range(width): - start, end = _cell_candidate_span(model, y, x) - indices = model.cell_candidate_indices[start:end] - match = np.flatnonzero(indices == selected[y, x]) - if match.size == 0: - continue - gathered[y, x] = unary_cost[start + int(match[0])] - score = float(np.mean(gathered)) - if width > 1: - penalties = [] - for y in range(height): - for x in range(width - 1): - penalties.append(_pair_penalty(model, int(selected[y, x]), int(selected[y, x + 1]), _RIGHT, solver_params)) - if penalties: - score += float(np.mean(np.asarray(penalties, dtype=np.float32))) - if height > 1: - penalties = [] - for y in range(height - 1): - for x in range(width): - penalties.append(_pair_penalty(model, int(selected[y, x]), int(selected[y + 1, x]), _DOWN, solver_params)) - if penalties: - score += float(np.mean(np.asarray(penalties, dtype=np.float32))) - return score - - def optimize_tile_graph( rgba: np.ndarray, inference: InferenceResult, - analysis: SourceAnalysis, + analysis: TileGraphSourceAnalysis, steps: int, seed: int, device: str, @@ -1327,7 +1300,7 @@ def optimize_tile_graph( torch.manual_seed(seed) if resolved_device == "cuda": torch.cuda.manual_seed_all(seed) - model = build_tile_graph_model( + model, build_stats = build_tile_graph_model( rgba, inference=inference, analysis=analysis, @@ -1443,14 +1416,14 @@ def optimize_tile_graph( kept_initial_assignment = True diagnostics = { "mode": "tile-graph", - "tile_graph_model_device": model.model_device, + "tile_graph_model_device": build_stats.model_device, "tile_graph_solver_device": resolved_device, - "tile_graph_model_cache_hit": model.cache_hit, + "tile_graph_model_cache_hit": build_stats.cache_hit, "tile_graph_proposal_mode": "atomic-components", - "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_average_choices": model.average_choices, + "tile_graph_component_count": build_stats.component_count, + "tile_graph_candidate_count": build_stats.candidate_count, + "tile_graph_edge_density": build_stats.edge_density, + "tile_graph_average_choices": build_stats.average_choices, "tile_graph_initial_score": float(loss_history[0]), "tile_graph_final_score": float(loss_history[-1]), "tile_graph_initial_source_fidelity": float(initial_source_fidelity), diff --git a/src/repixelizer/types.py b/src/repixelizer/types.py index 518187b..0cdd18d 100644 --- a/src/repixelizer/types.py +++ b/src/repixelizer/types.py @@ -28,12 +28,17 @@ class InferenceResult: @dataclass(slots=True) -class SourceAnalysis: +class ContinuousSourceAnalysis: edge_map: np.ndarray cluster_map: np.ndarray - cluster_centers: np.ndarray - alpha_map: np.ndarray - cluster_preview: np.ndarray + + +@dataclass(slots=True) +class TileGraphSourceAnalysis: + edge_map: np.ndarray + + +SourceAnalysis = ContinuousSourceAnalysis | TileGraphSourceAnalysis @dataclass(slots=True) @@ -41,9 +46,7 @@ class SourceLatticeReference: mean_rgba: np.ndarray sharp_rgba: np.ndarray dispersion: float - lattice_indices: np.ndarray cell_dispersion: np.ndarray - cell_counts: np.ndarray cell_support: np.ndarray cell_alpha_max: np.ndarray sharp_x: np.ndarray @@ -53,10 +56,16 @@ class SourceLatticeReference: edge_strength: np.ndarray edge_grad_x: np.ndarray edge_grad_y: np.ndarray - delta_x: np.ndarray | None - delta_y: np.ndarray | None - delta_diag: np.ndarray | None - delta_anti: np.ndarray | None + + +@dataclass(slots=True) +class TileGraphSourceReference: + sharp_rgba: np.ndarray + sharp_x: np.ndarray + sharp_y: np.ndarray + edge_peak_x: np.ndarray + edge_peak_y: np.ndarray + edge_strength: np.ndarray @dataclass(slots=True) @@ -87,7 +96,7 @@ class RunResult: source_rgba: np.ndarray output_rgba: np.ndarray inference: InferenceResult - analysis: SourceAnalysis + analysis: ContinuousSourceAnalysis | TileGraphSourceAnalysis solver: SolverArtifacts cleanup: CleanupArtifacts palette_result: PaletteResult | None diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 134e08f..4482b30 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -2,33 +2,28 @@ from __future__ import annotations import numpy as np -from repixelizer.analysis import analyze_source +from repixelizer.analysis import analyze_continuous_source, analyze_tile_graph_source -def test_analyze_source_cpu_device_matches_default_path() -> None: +def test_continuous_source_analysis_cpu_device_matches_default_path() -> None: source = np.zeros((8, 8, 4), dtype=np.float32) source[..., 3] = 1.0 source[2:6, 3:5, :3] = np.asarray([1.0, 0.8, 0.2], dtype=np.float32) - default = analyze_source(source, seed=7) - accelerated = analyze_source(source, seed=7, device="cpu") + default = analyze_continuous_source(source, seed=7) + accelerated = analyze_continuous_source(source, seed=7, device="cpu") assert np.allclose(accelerated.edge_map, default.edge_map, atol=1e-5) - assert accelerated.cluster_centers.shape == default.cluster_centers.shape - assert np.allclose(accelerated.alpha_map, default.alpha_map, atol=1e-5) assert accelerated.cluster_map.shape == default.cluster_map.shape - assert accelerated.cluster_preview.shape == default.cluster_preview.shape + assert np.array_equal(accelerated.cluster_map, default.cluster_map) -def test_analyze_source_can_skip_cluster_analysis() -> None: +def test_tile_graph_source_analysis_is_edge_only() -> None: source = np.zeros((8, 8, 4), dtype=np.float32) source[..., 3] = 1.0 source[2:6, 3:5, :3] = np.asarray([1.0, 0.8, 0.2], dtype=np.float32) - analysis = analyze_source(source, seed=7, device="cpu", include_clusters=False) + analysis = analyze_tile_graph_source(source, device="cpu") assert analysis.edge_map.shape == source.shape[:2] - assert np.allclose(analysis.alpha_map, source[..., 3], atol=1e-5) - assert analysis.cluster_centers.size == 0 - assert np.all(analysis.cluster_map == -1) - assert np.count_nonzero(analysis.cluster_preview) == 0 + assert not hasattr(analysis, "cluster_map") diff --git a/tests/test_continuous.py b/tests/test_continuous.py index dff3b08..c38778b 100644 --- a/tests/test_continuous.py +++ b/tests/test_continuous.py @@ -3,7 +3,7 @@ from __future__ import annotations import numpy as np import torch -from repixelizer.analysis import analyze_source +from repixelizer.analysis import analyze_continuous_source from repixelizer.continuous import ( _build_candidate_positions, _build_source_reliability, @@ -62,7 +62,7 @@ def test_source_reliability_stays_high_for_high_contrast_edge_cells() -> None: source[..., 3] = 1.0 source[:, 1, :3] = 1.0 - analysis = analyze_source(source, seed=3) + analysis = analyze_continuous_source(source, seed=3) reference = build_source_lattice_reference( source, target_width=4, @@ -154,7 +154,7 @@ def test_optimize_uv_field_does_not_regress_from_snap_on_thin_feature_case() -> artifacts = optimize_uv_field( fake, inference=inference, - analysis=analyze_source(fake, seed=7), + analysis=analyze_continuous_source(fake, seed=7), steps=4, seed=7, device="cpu", diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 479288c..a66aafd 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,10 +1,11 @@ from __future__ import annotations +import hashlib from pathlib import Path import numpy as np -from repixelizer.analysis import analyze_source +from repixelizer.analysis import analyze_continuous_source from repixelizer.baselines import naive_resize_baseline from repixelizer.continuous import optimize_uv_field from repixelizer.metrics import coherence_breakdown, foreground_reconstruction_error, source_lattice_consistency_breakdown @@ -121,6 +122,76 @@ def test_pipeline_can_emit_initialized_output_without_optimizer_steps(tmp_path: assert result.output_rgba.shape[1] == result.inference.target_width +def test_fixed_tile_graph_pipeline_output_matches_refactor_baseline(tmp_path: Path) -> None: + source = make_emblem(24, 24) + fake = fake_pixelize( + source, + upscale=8, + phase_x=0.1, + phase_y=-0.05, + blur_radius=0.4, + warp_strength=0.15, + warp_detail=4, + seed=9, + ) + input_path = tmp_path / "tile-baseline-input.png" + output_path = tmp_path / "tile-baseline-output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, fake) + result = run_pipeline( + input_path, + output_path, + reconstruction_mode="tile-graph", + target_width=24, + target_height=24, + phase_x=0.1, + phase_y=-0.05, + enable_phase_rerank=False, + steps=0, + device="cpu", + ) + + output_hash = hashlib.sha256(np.ascontiguousarray(result.output_rgba).view(np.uint8)).hexdigest() + assert output_hash == "31e3bc2cbb802a8d11bf00fbeb73772f7fa4934e5f187888d6f0d5b356f36ae9" + + +def test_fixed_continuous_pipeline_output_matches_refactor_baseline(tmp_path: Path) -> None: + source = make_emblem(24, 24) + fake = fake_pixelize( + source, + upscale=8, + phase_x=0.1, + phase_y=-0.05, + blur_radius=0.4, + warp_strength=0.15, + warp_detail=4, + seed=9, + ) + input_path = tmp_path / "continuous-baseline-input.png" + output_path = tmp_path / "continuous-baseline-output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, fake) + result = run_pipeline( + input_path, + output_path, + reconstruction_mode="continuous", + target_width=24, + target_height=24, + phase_x=0.1, + phase_y=-0.05, + enable_phase_rerank=False, + steps=0, + device="cpu", + ) + + output_hash = hashlib.sha256(np.ascontiguousarray(result.output_rgba).view(np.uint8)).hexdigest() + assert output_hash == "404748af3c7241a6602cb91fdfed1fd7e116dd14098827e33e78c7da5e3d7903" + + def test_phase_rerank_can_override_low_confidence_inference_pick(monkeypatch) -> None: source = np.zeros((2, 2, 4), dtype=np.float32) source[0, 0] = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32) @@ -302,7 +373,7 @@ def test_pipeline_runs_tile_graph_once_without_phase_probe_reuse(tmp_path: Path, ) class DummyAnalysis: - cluster_preview = np.zeros((6, 6, 4), dtype=np.float32) + edge_map = np.zeros((6, 6), dtype=np.float32) call_count = {"tile_graph": 0} @@ -336,7 +407,7 @@ def test_pipeline_runs_tile_graph_once_without_phase_probe_reuse(tmp_path: Path, return {"score": 0.05 if target_width == 3 else 0.15} monkeypatch.setattr("repixelizer.pipeline.infer_lattice", lambda *args, **kwargs: inference) - monkeypatch.setattr("repixelizer.pipeline.analyze_source", lambda *args, **kwargs: DummyAnalysis()) + monkeypatch.setattr("repixelizer.pipeline.analyze_tile_graph_source", lambda *args, **kwargs: DummyAnalysis()) monkeypatch.setattr("repixelizer.pipeline.optimize_tile_graph", fake_optimize_tile_graph) monkeypatch.setattr("repixelizer.pipeline.source_lattice_consistency_breakdown", fake_support) monkeypatch.setattr("repixelizer.pipeline.foreground_edge_position_error", lambda *args, **kwargs: 0.1) @@ -377,7 +448,8 @@ def test_pipeline_fixed_target_and_phase_skip_search_inference(tmp_path: Path, m ) class DummyAnalysis: - cluster_preview = np.zeros((6, 6, 4), dtype=np.float32) + edge_map = np.zeros((6, 6), dtype=np.float32) + cluster_map = np.zeros((6, 6), dtype=np.int32) def fail_infer(*args, **kwargs): raise AssertionError("full lattice search should not run for fixed target/phase") @@ -398,7 +470,7 @@ def test_pipeline_fixed_target_and_phase_skip_search_inference(tmp_path: Path, m monkeypatch.setattr("repixelizer.pipeline.infer_lattice", fail_infer) monkeypatch.setattr("repixelizer.pipeline.infer_fixed_lattice", fake_infer_fixed) - monkeypatch.setattr("repixelizer.pipeline.analyze_source", lambda *args, **kwargs: DummyAnalysis()) + monkeypatch.setattr("repixelizer.pipeline.analyze_continuous_source", lambda *args, **kwargs: DummyAnalysis()) monkeypatch.setattr("repixelizer.pipeline.optimize_uv_field", fake_optimize_uv_field) result = run_pipeline( @@ -551,7 +623,7 @@ def test_badge_fixture_candidate_beats_previous_source_fidelity_baselines() -> N artifacts = optimize_uv_field( source, inference=inference, - analysis=analyze_source(source, seed=7), + analysis=analyze_continuous_source(source, seed=7), steps=2, seed=7, device="cpu", diff --git a/tests/test_source_reference.py b/tests/test_source_reference.py index 0e93e0f..cdb10db 100644 --- a/tests/test_source_reference.py +++ b/tests/test_source_reference.py @@ -2,11 +2,10 @@ from __future__ import annotations import numpy as np -from repixelizer.io import premultiply -from repixelizer.source_reference import build_source_lattice_reference +from repixelizer.source_reference import build_source_lattice_reference, build_tile_graph_source_reference -def test_source_lattice_reference_recovers_stable_cell_colors_and_deltas() -> None: +def test_source_lattice_reference_recovers_stable_cell_colors() -> None: lowres = np.asarray( [ [[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]], @@ -25,15 +24,8 @@ def test_source_lattice_reference_recovers_stable_cell_colors_and_deltas() -> No phase_y=0.0, ) - expected_premul = premultiply(lowres) assert np.allclose(reference.sharp_rgba, lowres, atol=1e-5) assert np.allclose(reference.mean_rgba[0, 1], lowres[0, 1], atol=1e-5) - assert reference.delta_x is not None - assert reference.delta_y is not None - assert reference.delta_diag is not None - assert reference.delta_anti is not None - assert np.allclose(reference.delta_x, expected_premul[:, 1:, :] - expected_premul[:, :-1, :], atol=1e-5) - assert np.allclose(reference.delta_y, expected_premul[1:, :, :] - expected_premul[:-1, :, :], atol=1e-5) assert reference.dispersion >= 0.0 @@ -63,12 +55,21 @@ 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: + + +def test_tile_graph_source_reference_matches_continuous_anchor_fields() -> None: source = np.zeros((6, 6, 4), dtype=np.float32) source[..., 3] = 1.0 source[:, 2, :3] = 1.0 - reference = build_source_lattice_reference( + full = build_source_lattice_reference( + source, + target_width=2, + target_height=2, + phase_x=0.0, + phase_y=0.0, + ) + tile = build_tile_graph_source_reference( source, target_width=2, target_height=2, @@ -76,8 +77,12 @@ def test_source_lattice_reference_computes_default_edge_metadata() -> None: phase_y=0.0, ) - assert np.all(reference.edge_strength[:, 0] > 0.0) - assert np.all(reference.edge_peak_x[:, 0] == 2) + assert np.allclose(tile.sharp_rgba, full.sharp_rgba, atol=1e-5) + assert np.array_equal(tile.sharp_x, full.sharp_x) + assert np.array_equal(tile.sharp_y, full.sharp_y) + assert np.array_equal(tile.edge_peak_x, full.edge_peak_x) + assert np.array_equal(tile.edge_peak_y, full.edge_peak_y) + assert np.allclose(tile.edge_strength, full.edge_strength, atol=1e-5) def test_source_lattice_reference_cpu_device_matches_default_path() -> None: diff --git a/tests/test_tile_graph.py b/tests/test_tile_graph.py index d188451..671ea53 100644 --- a/tests/test_tile_graph.py +++ b/tests/test_tile_graph.py @@ -5,7 +5,7 @@ from pathlib import Path import numpy as np import pytest -from repixelizer.analysis import analyze_source +from repixelizer.analysis import analyze_tile_graph_source from repixelizer.baselines import naive_resize_baseline from repixelizer.io import save_rgba from repixelizer.metrics import source_lattice_consistency_breakdown @@ -36,15 +36,20 @@ def test_tile_graph_model_extracts_candidates_and_adjacency_from_lattice_proposa top_candidates=[], ) - model = build_tile_graph_model(source, inference=inference, analysis=analyze_source(source, seed=7), device="cpu") + model, stats = build_tile_graph_model( + source, + inference=inference, + analysis=analyze_tile_graph_source(source, device="cpu"), + device="cpu", + ) assert model.candidate_rgba.shape[0] > 4 choice_counts = np.diff(model.cell_candidate_offsets) assert choice_counts.shape[0] == inference.target_width * inference.target_height assert np.all(choice_counts >= 1) - assert model.average_choices >= 1.0 - assert model.edge_density > 0.0 - assert model.model_device == "cpu" + assert stats.average_choices >= 1.0 + assert stats.edge_density > 0.0 + assert stats.model_device == "cpu" def test_tile_graph_model_works_with_edge_only_analysis() -> None: @@ -61,15 +66,15 @@ def test_tile_graph_model_works_with_edge_only_analysis() -> None: top_candidates=[], ) - model = build_tile_graph_model( + model, stats = build_tile_graph_model( source, inference=inference, - analysis=analyze_source(source, seed=5, device="cpu", include_clusters=False), + analysis=analyze_tile_graph_source(source, device="cpu"), device="cpu", ) assert model.candidate_rgba.shape[0] > 0 - assert model.model_device == "cpu" + assert stats.model_device == "cpu" def test_tile_graph_candidates_use_literal_source_pixel_colors() -> None: @@ -93,10 +98,10 @@ def test_tile_graph_candidates_use_literal_source_pixel_colors() -> None: confidence=1.0, top_candidates=[], ) - model = build_tile_graph_model( + model, _stats = build_tile_graph_model( source, inference=inference, - analysis=analyze_source(source, seed=13), + analysis=analyze_tile_graph_source(source), solver_params=SolverHyperParams(tile_graph_max_candidates=32, tile_graph_max_candidates_per_coord=2), ) @@ -106,7 +111,7 @@ def test_tile_graph_candidates_use_literal_source_pixel_colors() -> None: assert candidate_colors.issubset(source_colors | {tuple(np.zeros(4, dtype=np.float32))}) -def test_tile_graph_candidates_are_scoped_to_their_output_coord() -> None: +def test_tile_graph_candidate_buckets_are_contiguous_per_output_coord() -> None: source = np.zeros((8, 8, 4), dtype=np.float32) source[..., 3] = 1.0 source[:, 2:4] = np.asarray([1.0, 0.9, 0.2, 1.0], dtype=np.float32) @@ -120,20 +125,16 @@ def test_tile_graph_candidates_are_scoped_to_their_output_coord() -> None: top_candidates=[], ) - model = build_tile_graph_model( + model, _stats = build_tile_graph_model( source, inference=inference, - analysis=analyze_source(source, seed=11), + analysis=analyze_tile_graph_source(source), ) - for y in range(inference.target_height): - for x in range(inference.target_width): - flat = y * inference.target_width + x - start = int(model.cell_candidate_offsets[flat]) - end = int(model.cell_candidate_offsets[flat + 1]) - coords = model.candidate_coords[model.cell_candidate_indices[start:end]] - assert np.all(coords[:, 0] == y) - assert np.all(coords[:, 1] == x) + total_candidates = int(model.candidate_rgba.shape[0]) + assert int(model.cell_candidate_offsets[0]) == 0 + assert int(model.cell_candidate_offsets[-1]) == total_candidates + assert np.array_equal(model.cell_candidate_indices, np.arange(total_candidates, dtype=np.int32)) def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> None: @@ -157,7 +158,7 @@ def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> Non top_candidates=[], ) params = SolverHyperParams() - analysis = analyze_source(source, seed=7, device="cpu") + analysis = analyze_tile_graph_source(source, device="cpu") cell_h = source.shape[0] / inference.target_height cell_w = source.shape[1] / inference.target_width source_region_stride = 2 @@ -236,10 +237,10 @@ def test_tile_graph_projects_source_regions_instead_of_mixed_lattice_buckets() - top_candidates=[], ) - model = build_tile_graph_model( + model, _stats = build_tile_graph_model( source, inference=inference, - analysis=analyze_source(source, seed=19), + analysis=analyze_tile_graph_source(source), device="cpu", ) @@ -280,10 +281,10 @@ def test_tile_graph_edge_cells_include_multiple_same_cell_contour_colors() -> No confidence=1.0, top_candidates=[], ) - model = build_tile_graph_model( + model, _stats = build_tile_graph_model( source, inference=inference, - analysis=analyze_source(source, seed=17), + analysis=analyze_tile_graph_source(source), device="cpu", ) @@ -329,7 +330,7 @@ def test_tile_graph_reconstructs_synthetic_thin_feature_better_than_naive() -> N artifacts, _ = optimize_tile_graph( fake, inference=inference, - analysis=analyze_source(fake, seed=7), + analysis=analyze_tile_graph_source(fake), steps=0, seed=7, device="cpu", @@ -384,7 +385,7 @@ def test_tile_graph_keeps_shallow_connected_stroke_from_fanning_out() -> None: confidence=1.0, top_candidates=[], ), - analysis=analyze_source(source, seed=7), + analysis=analyze_tile_graph_source(source), steps=0, seed=7, device="cpu", @@ -423,7 +424,7 @@ def test_tile_graph_preserves_transparent_background_on_sparse_detail() -> None: confidence=1.0, top_candidates=[], ), - analysis=analyze_source(fake, seed=7), + analysis=analyze_tile_graph_source(fake), steps=0, seed=7, device="cpu", @@ -449,8 +450,8 @@ def test_tile_graph_cuda_matches_cpu_on_small_case() -> None: confidence=1.0, top_candidates=[], ) - cpu_analysis = analyze_source(fake, seed=7, device="cpu") - cuda_analysis = analyze_source(fake, seed=7, device="cuda") + cpu_analysis = analyze_tile_graph_source(fake, device="cpu") + cuda_analysis = analyze_tile_graph_source(fake, device="cuda") cpu_artifacts, cpu_diag = optimize_tile_graph( fake, @@ -488,23 +489,23 @@ def test_tile_graph_model_cache_reuses_fixed_lattice_build() -> None: confidence=1.0, top_candidates=[], ) - analysis = analyze_source(fake, seed=7, device="cpu") + analysis = analyze_tile_graph_source(fake, device="cpu") - first = build_tile_graph_model( + first, first_stats = build_tile_graph_model( fake, inference=inference, analysis=analysis, device="cpu", ) - second = build_tile_graph_model( + second, second_stats = build_tile_graph_model( fake, inference=inference, analysis=analysis, device="cpu", ) - assert first.cache_hit is False - assert second.cache_hit is True + assert first_stats.cache_hit is False + assert second_stats.cache_hit is True assert np.array_equal(first.candidate_rgba, second.candidate_rgba)