From 235526e1dce5dedcca8be9aaf4fc70c065eb44ea Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Wed, 22 Apr 2026 21:40:25 +0100 Subject: [PATCH] Rework tile graph around atomic cell candidates --- README.md | 16 +- docs/implementation-plan.md | 199 ++++++------------ src/repixelizer/params.py | 11 +- src/repixelizer/tile_graph.py | 381 +++++++++++++++++++--------------- tests/test_tile_graph.py | 9 + 5 files changed, 299 insertions(+), 317 deletions(-) diff --git a/README.md b/README.md index 229f209..7fd1b37 100644 --- a/README.md +++ b/README.md @@ -38,17 +38,15 @@ Recent adjacency-focused status: Current tile-graph status: -- it now builds literal source-pixel candidates from lattice-assigned per-cell proposals instead of walking connected components on the CPU -- it now places candidates from a per-cell local candidate set instead of treating the whole image as one reusable global tile vocabulary -- the full tile-graph path now honors `--device`: source analysis, source-lattice reference construction, proposal scoring, and the local solver all run through Torch -- it no longer invents averaged patch colors during candidate extraction; the current regression suite now checks that tile-graph candidates stay on real source colors -- each output cell now scores a small raw-pixel proposal pool built from its inferred lattice assignment, sharp exemplar, and edge peak instead of borrowing candidates from distant regions -- edge-heavy cells now widen their local candidate set with same-cell edge neighbors and strongest same-cell edge pixels, so the mode can at least consider harder contour colors instead of only a single medoid-like pick -- it already beats naive resize on the repo's synthetic thin-feature regression +- 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 +- candidate seeding now includes atomic connected color regions within each inferred lattice cell, layered together with the old sharp/edge anchors instead of replacing them outright +- the latest full-emblem probe under `artifacts/full-emblem-tile-graph-atomic-v3-cuda/` lands at `0.0224` source-fidelity, beating the earlier full-CUDA tile-graph baseline at `0.0283` +- tile-graph now falls back to its initial assignment when the propagation loop would make source-lattice fidelity worse, which is currently important for preserving sharp internal contour cells +- it still beats naive resize on the repo's synthetic thin-feature regression and now does so without letting the final propagation step blur past the initial placement +- the solver, source analysis, and shared lattice reference still honor `--device`, but the new atomic-component candidate extraction is CPU-side again and is now the main runtime bottleneck on large real fixtures - this fixes the core design mismatch that had allowed repeated distant labels to create big same-color patches and opaque black background blocks - on the current `24x24` emblem smoke case, an end-to-end `tile-graph` run dropped from about `2.57s` on CPU to `0.61s` on CUDA on this machine -- the current badge CUDA probe under `artifacts/tile-graph-cuda-pass/badge-cuda/` now completes end-to-end on the selected `126x126` lattice and lands at `0.0283` final source-fidelity -- the latest hard-edge candidate pass under `artifacts/full-emblem-tile-graph-hard-edge-v2-cuda/` improves local candidate diversity and slightly sharpens some internal contour cells, but it still regresses full-emblem source-fidelity (`0.0377` vs `0.0283` for the previous tile-graph pass), so it should be treated as an experiment rather than a settled default +- the older hard-edge-only candidate widening pass under `artifacts/full-emblem-tile-graph-hard-edge-v2-cuda/` remains a useful negative result: more edge choices alone sharpen some cells locally but still regress full-emblem source-fidelity (`0.0377`) ## Quickstart diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index cca177d..002f6dc 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -8,179 +8,109 @@ Implemented: - package metadata and editable install - CLI entrypoint with `run` and `compare` modes - lattice/phase inference with coherence scoring plus source-period prior -- source-image structural analysis, now with an optional Torch/CUDA path +- source-image structural analysis, with optional Torch/CUDA paths - PyTorch-based continuous UV-field optimization -- shared source-lattice reference built from actual inferred cell assignments, now with an optional Torch/CUDA path +- a shared source-lattice reference built from actual inferred cell assignments - source-first snap/refine scoring with explicit source-vs-representative weights - low-confidence phase reranking with soft size penalties instead of hard size-jump rejection -- GPU-friendly tile-graph proposal building from lattice-assigned raw source pixels plus per-cell edge/cluster scoring -- discrete cleanup on the output grid -- optional palette loading and quantization modes -- compare mode with naive and diffusion baselines +- an experimental `tile-graph` reconstruction path with coord-local literal source-pixel candidates - stage-aware diagnostics writing, including per-stage source-fidelity and rerank traces - synthetic test fixtures and automated tests Verified: - local editable install in a dedicated venv -- full test suite currently passing on the last completed pass (`54 passed`) +- full test suite currently passing (`55 passed`) - compare-mode smoke run against a real emblem image -- focused real-fixture probe on `tests/fixtures/real/ai-badge-cleaned.png` now beats naive resize on source-lattice consistency under the selected lattice (`0.0832` final vs `0.1304` naive with the current `126x126` / `(0.0, -0.2)` pick) -- end-to-end tile-graph CUDA smoke runs on both a `24x24` emblem case and the cleaned badge fixture; the small emblem case currently drops from about `2.57s` on CPU to `0.61s` on CUDA on this machine +- 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` ## Status after adjacency-first pass What changed: - `src/repixelizer/source_reference.py` now owns the lattice-indexed source reference used by metrics, the solver, and phase reranking -- snap and refine now prefer sharp per-cell source evidence and keep the relaxed-mode neighborhood in the final discrete choice +- snap and refine prefer sharp per-cell source evidence and can keep the relaxed-mode neighborhood in the final discrete choice - `run.json` now includes `source_fidelity` for `snap_initial`, `solver_target`, and `final_output`, plus `phase_rerank_candidates` -- `scripts/render_focus_crop.py` now prints stage-level source-fidelity when generating the guard/tip crop sheet +- `scripts/render_focus_crop.py` prints stage-level source-fidelity when generating the guard/tip crop sheet What that fixed: -- the cleaned badge regression no longer loses badly to naive resize on adjacency and motif preservation once the solver starts from the selected lattice -- the solver no longer regresses relative to its own snap stage on the added thin-feature regression case +- the cleaned badge regression no longer loses badly to naive resize once the continuous solver starts from the selected lattice +- the continuous solver no longer regresses relative to its own snap stage on the added thin-feature regression case What is still provisional: -- the cleaned badge still lands on the conservative `126x126` lattice candidate, just with a much lower confidence than before; that means candidate selection and weight tuning are improved but not "done" -- tuning has not been rerun yet against the new source-first hyperparameters, so the benchmark/tuning baseline should be treated as stale until the next sweep -- the baked-checkerboard version of the badge remains a separate background-suppression problem +- the cleaned badge still lands on the conservative `126x126` lattice candidate with low confidence, so lattice selection is improved but not "done" +- tuning has not been rerun against the newer source-first weights +- the baked-checkerboard badge background remains a separate background-suppression problem ## What is still provisional -The current implementation should be treated as a strong experimental baseline, not a final algorithm. +The current implementation should still be treated as a strong experimental baseline, not a final algorithm. The most provisional areas are: -- lattice inference on real, messy generated imagery +- 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 +- the experimental tile-graph reconstruction path, especially on real noisy inputs +- low-confidence reranking once multiple plausible lattice sizes survive the soft penalty - compare-mode metrics as a proxy for actual visual quality ## Near-term priorities ### Active pass -This pass refactors the experimental tile-graph mode into a GPU-friendly codepath. +This pass moved tile-graph back toward the original "atomic source regions" idea without throwing away the parts of the older local solver that were actually helping. What landed: -- `analysis.py` and `source_reference.py` now have Torch device paths, so the expensive per-pixel reductions can run on CPU or CUDA without changing the public CLI -- `tile_graph.py` no longer walks connected components and destructively consumes windows during model building -- instead, tile-graph now builds a small per-cell proposal pool from lattice-assigned raw source pixels, sharp exemplars, edge peaks, and cluster-aware proposal scoring -- the full tile-graph path now honors `device`, not just the final solver loop, and diagnostics now record both the model-build device and the solver device +- `tile_graph.py` now segments connected atomic color regions inside each inferred lattice cell and uses those regions to seed per-cell literal source-pixel candidates +- atomic region candidates are layered together with the older sharp/edge anchors instead of replacing them outright +- tile-graph now keeps its initial assignment whenever the propagation loop would worsen source-lattice fidelity +- diagnostics now report the atomic proposal mode plus the initial/final tile-graph source-fidelity scores and whether the initial assignment was kept Current implementation note: -- the destructive cluster walk was the main thing fighting CUDA; replacing it with lattice-assigned per-cell proposals keeps the “one cell, one local candidate set” design while turning extraction into top-k scoring and grouped reductions -- the new proposal builder is intentionally local-first: it keeps candidates tied to the inferred output coord and uses edge/cluster signals only to rank raw source pixels already assigned to that lattice cell - -Experimental tile-graph note: -- there is now a separate `tile-graph` reconstruction mode that builds literal source-pixel candidates from lattice assignments instead of cell-averaged patches or globally reusable labels -- the raw-pixel rewrite fixed an implementation mismatch where candidates had drifted into cell-averaged patch colors instead of actual source pixels -- the later per-cell-local rewrite fixed the repeated-label and opaque-black-background failures by making candidates output-coord-scoped -- the newest GPU-friendly rewrite keeps that local-candidate behavior but replaces the old component walk with per-cell top-k proposal scoring, which is fast enough to run end-to-end on CUDA -- the current badge CUDA probe under `artifacts/tile-graph-cuda-pass/badge-cuda/` completes end-to-end at the selected `126x126` lattice with `24142` candidates and `0.0283` final source-fidelity -- the latest hard-edge candidate pass widens edge-cell proposal sets with same-cell edge neighbors and strongest same-cell edge pixels, and it does reduce some medoid-like contour picks locally -- that hard-edge pass is not a clean win yet: on the full badge fixture it increases candidate count to about `30k` and average choices to about `1.92`, but regresses full-image source-fidelity to about `0.0377`, so it remains exploratory +- the new atomic-component extraction is CPU-side, so although the solver and shared lattice analysis still honor `device`, candidate extraction is once again the dominant runtime cost on large real fixtures +- the latest full-emblem atomic probe under `artifacts/full-emblem-tile-graph-atomic-v3-cuda/` lands at `0.0224` source-fidelity with `34278` candidates and about `2.16` average choices per cell +- that beats the earlier full-CUDA tile-graph baseline (`0.0283`) and materially improves on the first atomic-only attempt (`0.1571`), which chose atomic regions too eagerly and then let the solver blur them out again Next after that: +- move atomic candidate extraction closer to the user's original "consume one-cell regions and step through larger regions" design instead of only segmenting inside inferred lattice cells +- make the atomic candidate builder GPU-friendly again, because the current CPU component walk is now the main scaling bottleneck - rerank low-confidence top lattice candidates with short real solver probes instead of relying only on the cheapest preview -- improve per-cell tile-graph proposal diversity so thin contours get more than “sharp vs edge-peak” in hard cells -- separate “edge-side selection” from “more edge candidates” so we can test contour-targeted scoring without paying the current whole-image fidelity penalty -- evaluate coordinated local moves when single-cell refinement still breaks thin contours -- rerun tuning now that the tile-graph proposal and device path have changed materially -- 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 +- evaluate whether the tile-graph propagation loop should become more contour-aware or simply stay more conservative now that the initial assignment is stronger +- rerun tuning after the tile-graph objective and candidate sets stabilize +- decide whether the tile-graph path should stay an alternate solver or become a candidate generator for the continuous refine stage -## High-value regression cases +## High-value regression case ### 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. +This locally generated AI badge remains a useful stress case because it fails in a 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 +- the sword tip and guard contain 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 -- the right-hand guard wing loses short dark-light-dark adjacency motifs and collapses into a chunky gold slab -- some guard closeups look nearly identical across `snap`, `relaxed`, and `final`, which suggests the wrong local motif is often chosen before later refinement even starts +- the right side of the sword tip can still wobble instead of tapering cleanly +- thin oblique outline features can still collapse into mushy transitions +- the right-hand guard wing can still lose short dark-light-dark adjacency motifs +- on tile-graph specifically, the remaining failure is less about wrong colors entering the candidate pool and more about the propagation step preferring a smoother local arrangement than the sharper initial assignment -Current status after the adjacency-first pass: -- on the cleaned fixture, `snap`, `relaxed`, and `final` now all preserve source-lattice structure substantially better than the old baseline, with `final` currently landing at `0.0832` source-fidelity score in the reproducible probe under `artifacts/badge-final-probe/` -- that is a meaningful improvement over both the previous documented `0.1494` final score and the naive resize score under the same selected lattice -- the remaining weakness is edge-cell trust and candidate coverage, not the old snap/refine handoff collapse +Current status: +- on the cleaned fixture, the continuous path now lands at `0.0832` source-fidelity in the reproducible probe under `artifacts/badge-final-probe/`, which is materially better than the older documented `0.1494` result and better than naive resize under the same lattice +- on the full emblem, the tile-graph atomic path under `artifacts/full-emblem-tile-graph-atomic-v3-cuda/` now reaches `0.0224`, which is the best tile-graph full-emblem result so far in this repo +- the remaining weakness on tile-graph is not candidate color purity anymore; it is candidate extraction cost plus deciding how aggressive the propagation step should be once the initial assignment is already strong 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 +- that 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 What future work should improve here: - edge-aware source reliability so high-dispersion but high-contrast cells keep trusting source evidence -- richer candidate generation so snap/refine can reach sharp exemplars, cell edge peaks, and gradient-guided offsets instead of only a fixed local grid +- richer candidate generation so continuous snap/refine can reach sharp exemplars, edge peaks, and guided offsets instead of only a fixed local grid - short top-k refine probes for low-confidence lattice candidates before the final rerank decision -- a retuning sweep after the edge-aware/source-guided changes land +- an atomic candidate builder that more closely matches real source-region ownership and consumption +- a GPU-friendly reformulation of the atomic candidate extraction stage - 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 -- for the experimental tile-graph path specifically, better local proposal diversity and stronger adjacency evidence are still needed even after switching candidates back to literal source pixels -- for the experimental tile-graph path specifically, edge-cell diversity alone is not enough; the current hard-edge candidate pass shows that simply widening edge-cell choices can sharpen some contour pixels while still making the global arrangement worse -- for the experimental tile-graph path specifically, the next useful step is pruning and richer proposal families rather than more component-walk optimization, because the old CPU walk is now gone - -Current optimizer diagnosis: -- the shared source-lattice reference, source-first snap/refine scoring, and relaxed-mode handoff are now in place, so the biggest remaining losses are no longer caused by the old representative-lattice collapse -- the current weak spot is that difficult thin features often live in high-dispersion cells, which still makes the solver partially distrust the source exactly where the contrast is most informative -- candidate search is also still too centered on the regular UV neighborhood; if the right edge pixel never enters the candidate set, the solver can preserve the wrong motif very consistently -- low-confidence reranking is softer than before, but it still relies on cheap probes, so larger lattice candidates can remain underexplored even when they support thin contours better - -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 -- `scripts/render_focus_crop.py` generates reproducible docs/debug closeups from output-grid cell coordinates -- the current docs-facing guard crop uses `--cell-bbox 73 20 107 44` - -### 1. Build a real benchmark corpus - -Add curated real-world cases, not just synthetic ones: -- emblems and logos with transparency -- icons with strong outlines -- simple sprite-like figures -- examples that are known to fail naive resize-and-dither - -Each case should include: -- source image -- expected target size or plausible size range -- short notes on known failure modes - -### 2. Improve lattice inference - -The current inference is workable but still heuristic. - -Next improvements should focus on: -- better periodicity detection for noisy or partially aligned images -- stronger penalties against oversized candidate grids -- confidence calibration that reflects ambiguity more honestly - -### 3. Improve edge-cell reconstruction - -The next quality gap is no longer generic cleanup; it is preserving thin, high-contrast local structure. - -Next version should add: -- edge-aware source reliability overrides for thin outlines and hard transitions -- source-guided candidate pools that are not limited to a uniform local offset lattice -- richer diagnostics for which candidate families win on real fixtures -- later, coordinated block or contour moves when single-cell greedy updates still break lines - -### 4. Make compare mode more trustworthy - -The current metrics are useful but still internal. - -Next work should: -- expand score breakdowns -- validate whether the metrics correlate with actual human preference -- make it easier to compare runs across parameter sweeps ## Suggested next milestones @@ -191,35 +121,34 @@ Next work should: - add a repeatable batch runner - emit one summary table per corpus run -### Milestone 2: Solver tuning pass +### Milestone 2: Tile-graph extraction pass -- tune continuous-stage loss weights against the benchmark corpus after the edge-aware/candidate-guided pass lands -- test more robust thin-feature and contour preservation behavior -- reduce cases where the solver still underperforms the naive baseline or settles on the wrong low-confidence lattice +- extract atomic source regions in a way that better matches actual source texel ownership +- preserve the current sharp/edge anchors as a safety net +- port the heavy extraction work back toward a GPU-friendly formulation -### Milestone 3: Coordinated contour moves +### Milestone 3: Solver tuning pass + +- tune continuous-stage loss weights against the benchmark corpus after the next edge-aware and candidate-guided pass lands +- tune tile-graph propagation aggressiveness now that the initial-assignment fallback is in place +- reduce cases where either solver underperforms the naive baseline or settles on the wrong low-confidence lattice + +### Milestone 4: Coordinated contour moves - add contour-aware multi-cell refinement moves - protect short dark-light-dark edge motifs during block updates - reduce cases where locally correct cells still fail to form a globally smooth outline -### Milestone 4: Human-in-the-loop v2 planning - -Only after the automatic baseline is strong: -- explore optional region hints -- evaluate a minimal GUI or notebook layer -- consider multi-frame support - ## Risks -- The solver can still produce outputs that score well but feel over-smoothed. -- Real generated images may contain multiple inconsistent local lattice cues. -- Continuous optimization is flexible, but easy to overfit to the reconstruction term. -- Metrics that reward coherence can accidentally reward blandness if not balanced carefully. +- the solver can still produce outputs that score well but feel over-smoothed +- real generated images may contain multiple inconsistent local lattice cues +- metrics that reward coherence can accidentally reward blandness if not balanced carefully +- the current atomic candidate extraction can become too expensive on large fixtures unless it is reworked again ## Assumptions -- Single-image quality matters more than batch throughput in the current phase. -- Palette constraints are optional and should stay optional. -- The project is better served by improving automatic behavior first than by rushing into manual tooling. -- Python plus PyTorch remains the right stack for the current research/prototyping phase. +- single-image quality matters more than batch throughput in the current phase +- palette constraints are optional and should stay optional +- the project is better served by improving automatic behavior first than by rushing into manual tooling +- Python plus PyTorch remains the right stack for the current research/prototyping phase diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index b5f791a..8bed871 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -76,14 +76,9 @@ class SolverHyperParams: tile_graph_max_candidates: int = 768 tile_graph_max_candidates_per_coord: int = 2 tile_graph_edge_candidates_per_coord: int = 6 - 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_proposal_pool_size: int = 4 - tile_graph_proposal_edge_weight: float = 0.40 - tile_graph_proposal_distance_weight: float = 0.22 - tile_graph_proposal_cluster_weight: float = 0.12 - tile_graph_proposal_color_weight: float = 0.10 + tile_graph_component_color_threshold: float = 0.055 + tile_graph_component_alpha_threshold: float = 0.12 + tile_graph_component_min_coverage: float = 0.02 tile_graph_nonedge_sharp_weight: float = 0.85 tile_graph_nonedge_mean_weight: float = 0.15 tile_graph_edge_mean_weight: float = 0.03 diff --git a/src/repixelizer/tile_graph.py b/src/repixelizer/tile_graph.py index bd40c25..301a449 100644 --- a/src/repixelizer/tile_graph.py +++ b/src/repixelizer/tile_graph.py @@ -5,6 +5,7 @@ from typing import Any 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 @@ -81,28 +82,136 @@ def _make_regular_uv(height: int, width: int, target_height: int, target_width: return uv.astype(np.float32) -def _build_ranked_proposal_pool_torch(torch, cell_indices_t, proposal_score_t, *, cell_count: int, pool_size: int): - if pool_size <= 0: - return torch.full((cell_count, 0), -1, device=cell_indices_t.device, dtype=torch.long) - if cell_indices_t.numel() == 0: - return torch.full((cell_count, pool_size), -1, device=cell_indices_t.device, dtype=torch.long) - score_lo_t = torch.min(proposal_score_t) - score_hi_t = torch.max(proposal_score_t) - normalized_score_t = (proposal_score_t - score_lo_t) / (score_hi_t - score_lo_t).clamp_min(1e-6) - sort_key_t = cell_indices_t.to(dtype=torch.float64) * 2.0 + (1.0 - normalized_score_t.to(dtype=torch.float64)) - order_t = torch.argsort(sort_key_t) - sorted_cells_t = cell_indices_t[order_t] - positions_t = torch.arange(order_t.shape[0], device=cell_indices_t.device, dtype=torch.long) - start_flags_t = torch.ones_like(sorted_cells_t, dtype=torch.bool) - if sorted_cells_t.numel() > 1: - start_flags_t[1:] = sorted_cells_t[1:] != sorted_cells_t[:-1] - group_start_t = torch.where(start_flags_t, positions_t, torch.zeros_like(positions_t)) - group_start_t = torch.cummax(group_start_t, dim=0).values - rank_t = positions_t - group_start_t - keep_mask_t = rank_t < pool_size - pool_t = torch.full((cell_count, pool_size), -1, device=cell_indices_t.device, dtype=torch.long) - pool_t[sorted_cells_t[keep_mask_t], rank_t[keep_mask_t]] = order_t[keep_mask_t] - return pool_t +def _rgba_luminance(rgba: np.ndarray) -> np.ndarray: + return rgba[..., 0] * 0.2126 + rgba[..., 1] * 0.7152 + rgba[..., 2] * 0.0722 + + +def _segment_atomic_components_in_cell( + *, + linear_indices: np.ndarray, + flat_y: np.ndarray, + flat_x: np.ndarray, + flat_rgba: np.ndarray, + flat_premul: np.ndarray, + flat_edge: np.ndarray, + color_threshold: float, + alpha_threshold: float, +) -> list[dict[str, float | int | np.ndarray]]: + if linear_indices.size == 0: + return [] + ys = flat_y[linear_indices] + xs = flat_x[linear_indices] + premul = flat_premul[linear_indices] + rgba = flat_rgba[linear_indices] + edges = flat_edge[linear_indices] + y0 = int(np.min(ys)) + x0 = int(np.min(xs)) + local_y = ys - y0 + local_x = xs - x0 + local_index = -np.ones((int(np.max(local_y)) + 1, int(np.max(local_x)) + 1), dtype=np.int32) + local_index[local_y, local_x] = np.arange(linear_indices.shape[0], dtype=np.int32) + visited = np.zeros(linear_indices.shape[0], dtype=bool) + components: list[dict[str, float | int | np.ndarray]] = [] + + for seed in range(linear_indices.shape[0]): + if visited[seed]: + continue + visited[seed] = True + queue = [seed] + members: list[int] = [] + seed_premul = premul[seed] + seed_alpha = float(rgba[seed, 3]) + while queue: + pos = queue.pop() + members.append(pos) + cy = int(local_y[pos]) + cx = int(local_x[pos]) + for dy, dx in ((0, 1), (1, 0), (0, -1), (-1, 0)): + ny = cy + dy + nx = cx + dx + if ny < 0 or nx < 0 or ny >= local_index.shape[0] or nx >= local_index.shape[1]: + continue + neighbor = int(local_index[ny, nx]) + if neighbor < 0 or visited[neighbor]: + continue + color_diff = float(np.mean(np.abs(premul[neighbor] - seed_premul))) + alpha_diff = abs(float(rgba[neighbor, 3]) - seed_alpha) + if color_diff > color_threshold or alpha_diff > alpha_threshold: + continue + visited[neighbor] = True + queue.append(neighbor) + member_array = np.asarray(members, dtype=np.int32) + member_linear = linear_indices[member_array] + member_y = ys[member_array].astype(np.float32) + member_x = xs[member_array].astype(np.float32) + centroid_y = float(np.mean(member_y)) + centroid_x = float(np.mean(member_x)) + centroid_dist = (member_y - centroid_y) ** 2 + (member_x - centroid_x) ** 2 + centroid_pick = int(np.argmin(centroid_dist)) + edge_pick = int(np.argmax(edges[member_array])) + edge_peak = float(np.max(edges[member_array])) + rep_slot = edge_pick if edge_peak > 0.0 else centroid_pick + rep_linear = int(member_linear[rep_slot]) + rep_rgba = rgba[member_array][rep_slot].astype(np.float32) + rgba_mean = np.mean(rgba[member_array], axis=0).astype(np.float32) + components.append( + { + "rep_linear": rep_linear, + "rep_rgba": rep_rgba, + "coverage": float(member_array.size / max(1, linear_indices.size)), + "edge_peak": edge_peak, + "alpha_mean": float(np.mean(rgba[member_array, 3])), + "luminance": float(_rgba_luminance(rgba_mean[None, :])[0]), + "size": int(member_array.size), + } + ) + return components + + +def _select_atomic_cell_candidates( + *, + components: list[dict[str, float | int | np.ndarray]], + allowed_candidates: int, + min_component_coverage: float, + reference_rgba: np.ndarray, + edge_reference_rgba: np.ndarray, + prefer_edge: bool, +) -> list[dict[str, float | int | np.ndarray]]: + if not components: + return [] + picked: list[dict[str, float | int | np.ndarray]] = [] + seen_linear: set[int] = set() + + def add(component: dict[str, float | int | np.ndarray]) -> None: + rep_linear = int(component["rep_linear"]) + if rep_linear in seen_linear: + return + if float(component["coverage"]) < min_component_coverage and picked: + return + seen_linear.add(rep_linear) + picked.append(component) + + by_area = sorted(components, key=lambda comp: (-float(comp["coverage"]), -float(comp["edge_peak"]))) + by_edge = sorted(components, key=lambda comp: (-float(comp["edge_peak"]), -float(comp["coverage"]))) + by_reference = sorted( + components, + key=lambda comp: float(np.mean(np.abs(np.asarray(comp["rep_rgba"], dtype=np.float32) - reference_rgba))), + ) + by_edge_reference = sorted( + components, + key=lambda comp: float(np.mean(np.abs(np.asarray(comp["rep_rgba"], dtype=np.float32) - edge_reference_rgba))), + ) + + add(by_area[0]) + add(by_reference[0]) + if prefer_edge: + add(by_edge_reference[0]) + add(by_edge[0]) + for component in by_area: + add(component) + if len(picked) >= allowed_candidates: + break + return picked[:allowed_candidates] def build_tile_graph_model( @@ -131,18 +240,17 @@ def build_tile_graph_model( cell_w = width / max(1, inference.target_width) output_area = max(1, inference.target_height * inference.target_width) max_candidates_per_coord = max(1, int(solver_params.tile_graph_max_candidates_per_coord)) - proposal_pool_size = max(max_candidates_per_coord, int(getattr(solver_params, "tile_graph_proposal_pool_size", 4))) source_t = torch.from_numpy(source_rgba).to(device=resolved_device, dtype=torch.float32) - edge_t = torch.from_numpy(analysis.edge_map).to(device=resolved_device, dtype=torch.float32) - cluster_t = torch.from_numpy(analysis.cluster_map).to(device=resolved_device, dtype=torch.long) flat_rgba_t = source_t.reshape(-1, source_t.shape[-1]) - flat_alpha_t = flat_rgba_t[:, 3] - flat_edge_t = edge_t.reshape(-1) - flat_cluster_t = cluster_t.reshape(-1) + pixel_y_t, pixel_x_t = torch.meshgrid( + torch.arange(height, device=resolved_device, dtype=torch.long), + torch.arange(width, device=resolved_device, dtype=torch.long), + indexing="ij", + ) + flat_y_t = pixel_y_t.reshape(-1) + flat_x_t = pixel_x_t.reshape(-1) lattice_indices_t = torch.from_numpy(source_reference.lattice_indices).to(device=resolved_device, dtype=torch.long).reshape(-1) - cell_support_flat_t = torch.from_numpy(source_reference.cell_support.reshape(-1)).to(device=resolved_device, dtype=torch.float32) - cell_counts_flat_t = torch.from_numpy(source_reference.cell_counts.reshape(-1)).to(device=resolved_device, dtype=torch.float32) edge_strength_flat_t = torch.from_numpy(source_reference.edge_strength.reshape(-1)).to(device=resolved_device, dtype=torch.float32) sharp_x_flat_t = torch.from_numpy(source_reference.sharp_x.reshape(-1)).to(device=resolved_device, dtype=torch.long) sharp_y_flat_t = torch.from_numpy(source_reference.sharp_y.reshape(-1)).to(device=resolved_device, dtype=torch.long) @@ -151,101 +259,30 @@ def build_tile_graph_model( sharp_linear_t = sharp_y_flat_t * width + sharp_x_flat_t edge_linear_t = edge_peak_y_flat_t * width + edge_peak_x_flat_t - pixel_y_t, pixel_x_t = torch.meshgrid( - torch.arange(height, device=resolved_device, dtype=torch.long), - torch.arange(width, device=resolved_device, dtype=torch.long), - indexing="ij", - ) - flat_y_t = pixel_y_t.reshape(-1) - flat_x_t = pixel_x_t.reshape(-1) - sharp_rgba_flat_t = flat_rgba_t[sharp_linear_t] - edge_rgba_flat_t = flat_rgba_t[edge_linear_t] - sharp_cluster_flat_t = flat_cluster_t[sharp_linear_t] - sharp_dist_t = ( - (flat_y_t.to(dtype=torch.float32) - sharp_y_flat_t[lattice_indices_t].to(dtype=torch.float32)).abs() / max(cell_h, 1.0) - + (flat_x_t.to(dtype=torch.float32) - sharp_x_flat_t[lattice_indices_t].to(dtype=torch.float32)).abs() / max(cell_w, 1.0) - ) - edge_focus_dist_t = ( - (flat_y_t.to(dtype=torch.float32) - edge_peak_y_flat_t[lattice_indices_t].to(dtype=torch.float32)).abs() / max(cell_h, 1.0) - + (flat_x_t.to(dtype=torch.float32) - edge_peak_x_flat_t[lattice_indices_t].to(dtype=torch.float32)).abs() / max(cell_w, 1.0) - ) - same_cluster_t = ( - (flat_cluster_t >= 0) & (flat_cluster_t == sharp_cluster_flat_t[lattice_indices_t]) - ).to(dtype=torch.float32) - sharp_color_error_t = (flat_rgba_t - sharp_rgba_flat_t[lattice_indices_t]).abs().mean(dim=-1) - edge_color_error_t = (flat_rgba_t - edge_rgba_flat_t[lattice_indices_t]).abs().mean(dim=-1) - edge_cell_mask_t = edge_strength_flat_t[lattice_indices_t] >= solver_params.source_edge_detail_threshold - proposal_color_error_t = torch.where( - edge_cell_mask_t, - torch.minimum(sharp_color_error_t, edge_color_error_t), - sharp_color_error_t, - ) - edge_cell_boost_t = ( - (1.0 - edge_focus_dist_t.clamp(0.0, 1.0)) - * edge_cell_mask_t.to(dtype=torch.float32) - ) - proposal_score_t = ( - flat_alpha_t - + flat_edge_t * getattr(solver_params, "tile_graph_proposal_edge_weight", 0.40) - + same_cluster_t * getattr(solver_params, "tile_graph_proposal_cluster_weight", 0.12) - + edge_cell_boost_t * getattr(solver_params, "tile_graph_proposal_edge_weight", 0.40) * 0.5 - + cell_support_flat_t[lattice_indices_t] * 0.08 - - sharp_dist_t * getattr(solver_params, "tile_graph_proposal_distance_weight", 0.22) - - proposal_color_error_t * getattr(solver_params, "tile_graph_proposal_color_weight", 0.10) - ) - ranked_pool_t = _build_ranked_proposal_pool_torch( - torch, - lattice_indices_t, - proposal_score_t, - cell_count=output_area, - pool_size=proposal_pool_size, - ) - edge_ranked_pool_t = _build_ranked_proposal_pool_torch( - torch, - lattice_indices_t, - flat_edge_t + flat_alpha_t * 0.05, - cell_count=output_area, - pool_size=proposal_pool_size, - ) - - cluster_count = max(1, int(max(analysis.cluster_centers.shape[0], 1))) - pixel_coverage_t = torch.maximum(flat_alpha_t.clamp(0.0, 1.0), flat_alpha_t.new_zeros(flat_alpha_t.shape)) - valid_cluster_mask_t = flat_cluster_t >= 0 - if bool(valid_cluster_mask_t.any().item()): - cluster_keys_t = lattice_indices_t[valid_cluster_mask_t] * cluster_count + flat_cluster_t[valid_cluster_mask_t] - cluster_counts_t = torch.bincount(cluster_keys_t, minlength=output_area * cluster_count).to(dtype=torch.float32) - pixel_coverage_t = torch.zeros_like(flat_alpha_t, dtype=torch.float32) - pixel_coverage_t[valid_cluster_mask_t] = ( - cluster_counts_t[cluster_keys_t] / cell_counts_flat_t[lattice_indices_t[valid_cluster_mask_t]].clamp_min(1.0) - ) - pixel_coverage_t = torch.maximum(pixel_coverage_t, flat_alpha_t * 0.5) - + flat_rgba_np = source_rgba.reshape(-1, source_rgba.shape[-1]).astype(np.float32) + flat_premul_np = flat_rgba_np.copy() + flat_premul_np[:, :3] *= flat_premul_np[:, 3:4] + flat_edge_np = analysis.edge_map.reshape(-1).astype(np.float32) + flat_idx_np = source_reference.lattice_indices.reshape(-1).astype(np.int32) + counts_np = np.bincount(flat_idx_np, minlength=output_area).astype(np.int32) + offsets_np = np.zeros(output_area + 1, dtype=np.int32) + offsets_np[1:] = np.cumsum(counts_np, dtype=np.int32) + order_np = np.argsort(flat_idx_np, kind="stable") + flat_y_np, flat_x_np = np.divmod(np.arange(flat_idx_np.size, dtype=np.int64), width) + flat_y_np = flat_y_np.astype(np.int32) + flat_x_np = flat_x_np.astype(np.int32) sharp_linear_np = sharp_linear_t.detach().cpu().numpy().astype(np.int64) edge_linear_np = edge_linear_t.detach().cpu().numpy().astype(np.int64) - edge_neighbor_linear_np: list[np.ndarray] = [] - cell_ids_t = torch.arange(output_area, device=resolved_device, dtype=torch.long) - for _, dy, dx in _DIRS: - neighbor_y_t = (edge_peak_y_flat_t + dy).clamp(0, max(0, height - 1)) - neighbor_x_t = (edge_peak_x_flat_t + dx).clamp(0, max(0, width - 1)) - neighbor_linear_t = neighbor_y_t * width + neighbor_x_t - same_cell_neighbor_t = lattice_indices_t[neighbor_linear_t] == cell_ids_t - masked_neighbor_linear_t = torch.where( - same_cell_neighbor_t, - neighbor_linear_t, - torch.full_like(neighbor_linear_t, -1), - ) - edge_neighbor_linear_np.append(masked_neighbor_linear_t.detach().cpu().numpy().astype(np.int64)) - ranked_pool_np = ranked_pool_t.detach().cpu().numpy().astype(np.int64) - edge_ranked_pool_np = edge_ranked_pool_t.detach().cpu().numpy().astype(np.int64) - pixel_coverage_np = pixel_coverage_t.detach().cpu().numpy().astype(np.float32) + flat_alpha_np = flat_rgba_np[:, 3].astype(np.float32) + edge_rgba_np = flat_rgba_np[edge_linear_np] edge_strength_np = source_reference.edge_strength.reshape(-1).astype(np.float32) - flat_alpha_np = source_rgba.reshape(-1, source_rgba.shape[-1])[:, 3].astype(np.float32) candidate_linear: list[int] = [] candidate_coords: list[tuple[int, int]] = [] candidate_area_ratio: list[float] = [] candidate_coverage: list[float] = [] choice_counts_list: list[int] = [] + component_total = 0 edge_candidate_cap = max(max_candidates_per_coord, int(getattr(solver_params, "tile_graph_edge_candidates_per_coord", max_candidates_per_coord))) for flat_index in range(output_area): coord_y = flat_index // inference.target_width @@ -254,37 +291,44 @@ def build_tile_graph_model( cell_start = len(candidate_linear) edge_cell = edge_strength_np[flat_index] >= solver_params.source_edge_detail_threshold allowed_candidates = edge_candidate_cap if edge_cell else max_candidates_per_coord - proposal_pixels = [int(sharp_linear_np[flat_index])] - if edge_cell: - proposal_pixels.append(int(edge_linear_np[flat_index])) - for neighbor_linear in edge_neighbor_linear_np: - proposal_pixels.append(int(neighbor_linear[flat_index])) - proposal_pixels.extend(int(pixel_idx) for pixel_idx in edge_ranked_pool_np[flat_index].tolist() if int(pixel_idx) >= 0) - proposal_pixels.extend(int(pixel_idx) for pixel_idx in ranked_pool_np[flat_index].tolist() if int(pixel_idx) >= 0) - for pixel_index in proposal_pixels: - if pixel_index < 0 or pixel_index in seen_pixels: - continue - is_primary = len(seen_pixels) == 0 - coverage = float(pixel_coverage_np[pixel_index]) - if ( - not is_primary - and coverage < solver_params.tile_graph_window_coverage_threshold - and not edge_cell - ): - continue + + def add_candidate(pixel_index: int, coverage: float) -> None: + if pixel_index in seen_pixels or len(seen_pixels) >= allowed_candidates: + return seen_pixels.add(pixel_index) candidate_linear.append(pixel_index) candidate_coords.append((coord_y, coord_x)) candidate_area_ratio.append(1.0) - candidate_coverage.append(max(coverage, float(flat_alpha_np[pixel_index]))) - if len(seen_pixels) >= allowed_candidates: - break + candidate_coverage.append(float(np.clip(coverage, 0.0, 1.0))) + + linear_indices = order_np[offsets_np[flat_index] : offsets_np[flat_index + 1]] + components = _segment_atomic_components_in_cell( + linear_indices=linear_indices, + flat_y=flat_y_np, + flat_x=flat_x_np, + flat_rgba=flat_rgba_np, + flat_premul=flat_premul_np, + flat_edge=flat_edge_np, + color_threshold=float(getattr(solver_params, "tile_graph_component_color_threshold", 0.055)), + alpha_threshold=float(getattr(solver_params, "tile_graph_component_alpha_threshold", 0.12)), + ) + component_total += len(components) + selected_components = _select_atomic_cell_candidates( + components=components, + allowed_candidates=allowed_candidates, + min_component_coverage=float(getattr(solver_params, "tile_graph_component_min_coverage", 0.02)), + reference_rgba=source_reference.sharp_rgba.reshape(-1, 4)[flat_index], + edge_reference_rgba=edge_rgba_np[flat_index], + prefer_edge=edge_cell, + ) + add_candidate(int(sharp_linear_np[flat_index]), max(0.05, float(flat_alpha_np[int(sharp_linear_np[flat_index])]))) + if edge_cell: + add_candidate(int(edge_linear_np[flat_index]), max(0.05, float(flat_alpha_np[int(edge_linear_np[flat_index])]))) + for component in selected_components: + pixel_index = int(component["rep_linear"]) + add_candidate(pixel_index, float(component["coverage"])) if len(candidate_linear) == cell_start: - pixel_index = int(sharp_linear_np[flat_index]) - candidate_linear.append(pixel_index) - candidate_coords.append((coord_y, coord_x)) - candidate_area_ratio.append(1.0) - candidate_coverage.append(max(float(pixel_coverage_np[pixel_index]), float(flat_alpha_np[pixel_index]))) + add_candidate(int(sharp_linear_np[flat_index]), max(0.05, float(flat_alpha_np[int(sharp_linear_np[flat_index])]))) choice_counts_list.append(len(candidate_linear) - cell_start) candidate_linear_t = torch.as_tensor(candidate_linear, device=resolved_device, dtype=torch.long) @@ -309,26 +353,9 @@ def build_tile_graph_model( choice_counts = np.diff(cell_candidate_offsets) average_choices = float(np.mean(choice_counts)) if choice_counts.size else 0.0 - extracted_choice_grid = np.maximum(choice_counts.reshape(inference.target_height, inference.target_width) - 1, 0) - edge_support = 0.0 - edge_total = 0 - if inference.target_width > 1: - edge_support += float( - np.mean( - (extracted_choice_grid[:, :-1] > 0).astype(np.float32) - * (extracted_choice_grid[:, 1:] > 0).astype(np.float32) - ) - ) - edge_total += 1 - if inference.target_height > 1: - edge_support += float( - np.mean( - (extracted_choice_grid[:-1, :] > 0).astype(np.float32) - * (extracted_choice_grid[1:, :] > 0).astype(np.float32) - ) - ) - edge_total += 1 - edge_density = edge_support / max(1, edge_total) + edge_density = float( + np.mean(source_reference.edge_strength >= solver_params.source_edge_detail_threshold) + ) return TileGraphModel( candidate_rgba=candidate_rgba_np, @@ -340,9 +367,9 @@ def build_tile_graph_model( cell_candidate_indices=cell_candidate_indices_np, reference_mean_rgba=source_reference.mean_rgba.astype(np.float32), reference_sharp_rgba=source_reference.sharp_rgba.astype(np.float32), - reference_edge_rgba=edge_rgba_flat_t.reshape(inference.target_height, inference.target_width, -1).detach().cpu().numpy().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), - component_count=int(np.unique(analysis.cluster_map[analysis.alpha_map >= solver_params.alpha_transparent_threshold]).size), + component_count=component_total, edge_density=edge_density, average_choices=average_choices, model_device=resolved_device, @@ -720,17 +747,41 @@ def optimize_tile_graph( selected = selected_t.detach().cpu().numpy().astype(np.int32) initial_rgba = _assignment_rgba(model, initial_selected).astype(np.float32) target_rgba = _assignment_rgba(model, selected).astype(np.float32) + initial_source_fidelity = source_lattice_consistency_breakdown( + rgba, + initial_rgba, + target_width=inference.target_width, + target_height=inference.target_height, + phase_x=inference.phase_x, + phase_y=inference.phase_y, + )["score"] + final_source_fidelity = source_lattice_consistency_breakdown( + rgba, + target_rgba, + target_width=inference.target_width, + target_height=inference.target_height, + phase_x=inference.phase_x, + phase_y=inference.phase_y, + )["score"] + kept_initial_assignment = False + if final_source_fidelity > initial_source_fidelity + 1e-6: + target_rgba = initial_rgba.copy() + final_source_fidelity = initial_source_fidelity + kept_initial_assignment = True diagnostics = { "mode": "tile-graph", "tile_graph_model_device": model.model_device, "tile_graph_solver_device": resolved_device, - "tile_graph_proposal_mode": "lattice-topk", + "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_initial_score": float(loss_history[0]), "tile_graph_final_score": float(loss_history[-1]), + "tile_graph_initial_source_fidelity": float(initial_source_fidelity), + "tile_graph_final_source_fidelity": float(final_source_fidelity), + "tile_graph_kept_initial_assignment": kept_initial_assignment, } return ( SolverArtifacts( diff --git a/tests/test_tile_graph.py b/tests/test_tile_graph.py index 5a22459..ef70edf 100644 --- a/tests/test_tile_graph.py +++ b/tests/test_tile_graph.py @@ -187,6 +187,14 @@ def test_tile_graph_reconstructs_synthetic_thin_feature_better_than_naive() -> N phase_x=0.18, phase_y=-0.12, )["score"] + initial_tile_score = source_lattice_consistency_breakdown( + fake, + artifacts.initial_rgba, + target_width=16, + target_height=16, + phase_x=0.18, + phase_y=-0.12, + )["score"] naive_score = source_lattice_consistency_breakdown( fake, naive, @@ -196,6 +204,7 @@ def test_tile_graph_reconstructs_synthetic_thin_feature_better_than_naive() -> N phase_y=-0.12, )["score"] + assert tile_score <= initial_tile_score + 1e-6 assert tile_score <= naive_score