GameCult/repixelizer
Watch
9
0
Fork
You've already forked repixelizer
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Rework tile graph around atomic cell candidates

Browse source
This commit is contained in:
Meta Vens 2026-04-22 21:40:25 +01:00
commit 235526e1dc
5 changed files with 299 additions and 317 deletions
Download patch file Download diff file Expand all files Collapse all files

16
README.md
View file

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

199
docs/implementation-plan.md
View file

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

11
src/repixelizer/params.py
View file

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

381
src/repixelizer/tile_graph.py
View file

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

9
tests/test_tile_graph.py
View file

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