From ca68154f181eb97892a47982ad1db5a641dbd004 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 23 Apr 2026 05:20:14 +0100 Subject: [PATCH] Replace tile stepping with overlap reduction --- README.md | 10 +- docs/implementation-plan.md | 22 +- docs/tile-graph-algorithm-map.md | 461 ++++++++------------ src/repixelizer/params.py | 3 - src/repixelizer/tile_graph.py | 708 ++++++++----------------------- tests/test_pipeline.py | 2 +- tests/test_tile_graph.py | 11 +- 7 files changed, 390 insertions(+), 827 deletions(-) diff --git a/README.md b/README.md index e1225db..f02a777 100644 --- a/README.md +++ b/README.md @@ -44,25 +44,25 @@ Current tile-graph status: - candidate selection is now source-owned: extracted tiles are ranked by their own area, coverage, and edge signal instead of by resemblance to `sharp_rgba` / `edge_rgba` - occupied foreground cells no longer get papered over with injected `sharp_pixel` or `edge_pixel` fallbacks; if extraction cannot produce a real tile there, the build fails instead of lying - the adjacency term is now learned from extracted tiles themselves, not from sampled one-cell-away RGBA deltas +- tile extraction no longer walks components with seed queues and one-cell windows; it now directly reduces `(component_id, output_cell)` overlaps into candidate shards - the deep-dive map in `docs/tile-graph-algorithm-map.md` now reflects the living machine directly instead of documenting the dead portrait/stride/delta path - the extraction stage still has an explicit empty-cell overlap fill pass, so output cells that actually contain opaque source pixels do not silently lose their extracted region bucket - 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 - source-region connected-components now have a device-side Torch path, so the expensive labeling step no longer depends on Python flood fill -- profiling the selected badge candidate shows the main remaining bottleneck is model construction, not the solver loop: `build_tile_graph_model(...)` takes about `142.3s`, and about `131.2s` of that is `_extract_source_region_tiles(...)` +- profiling the older selected badge candidate showed the main bottleneck was model construction rather than the solver loop: `build_tile_graph_model(...)` took about `142.3s`, and about `131.2s` of that was `_extract_source_region_tiles(...)` - after the first pruning pass, tile-graph no longer spends iteration time on pipeline rerank probes at all; it now runs the chosen or pinned lattice directly - tile-graph now also skips cluster analysis entirely and no longer uses sharp/edge portrait anchors in unary scoring; the path keeps only edge scouting plus direct per-cell source summaries -- the remaining large-fixture bottleneck is the per-component one-cell window cutting pass inside source-region extraction, which is still Python/NumPy-heavy after the GPU CCL stage +- after the reduce-by-key rewrite, the large-fixture bottleneck moved: extraction is much cheaper, and the main cold-build cost is now connected-component labeling / label compression rather than per-component window cutting - the pipeline now has a direct-control path for iteration: `--target-width` / `--target-height` plus optional `--phase-x` / `--phase-y` let you run an exact lattice without paying the full lattice search, and `--skip-phase-rerank` lets you keep the pipeline from second-guessing that choice - repeated fixed-lattice `tile-graph` runs in the same Python process now reuse the expensive model build; on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, the first CUDA run took about `10.2s` and the second cached rerun took about `2.1s` -- elongated source regions now get a stroke-aware slicing pass that follows their principal axis instead of only marching with cardinal queue steps, and the repo now has a shallow-stroke regression that checks for fan-out on that kind of component -- that stroke-aware slicer is a real synthetic improvement but not a real-badge win yet: the latest badge probe under `artifacts/badge-tile-graph-stroke-v2-cuda/` still keeps the initial assignment and slightly regresses source-fidelity (`0.1832` vs `0.1800`) +- the latest reduce-by-key pinned badge probe under `artifacts/tile-graph-reducebykey-v1-badge-126/` finishes in about `16.1s` cold and `6.0s` warm on CUDA, versus the previous full-resolution cut at about `96.9s` cold and `5.4s` warm +- that speed pass is a real throughput win but a small quality regression: the same pinned badge case moves from `0.1814` to `0.2036` source-fidelity - after the first pruning pass, tile-graph no longer participates in pipeline phase-rerank probes and no longer carries hybrid geometry priors through its unary cost; the path is now one lattice-conditioned candidate generator plus one local discrete solver - 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 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`) -- on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, the latest cut lands at `0.1814` source-fidelity instead of the post-refactor collapse around `0.5055`, which is a real recovery even though it is still far from the continuous path ## Quickstart diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index 03351ce..5a81eec 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -19,7 +19,7 @@ Implemented: Verified: - local editable install in a dedicated venv -- full test suite currently passing before the current algorithm-cut pass (`70 passed`) +- full test suite currently passing before the current speed/refactor pass (`69 passed` once the new grouped-extraction path landed) - compare-mode smoke run against a real emblem image - the cleaned real badge fixture in `tests/fixtures/real/ai-badge-cleaned.png` still beats naive resize on source-lattice consistency on the continuous path - the latest full-emblem tile-graph atomic probe under `artifacts/full-emblem-tile-graph-atomic-v3-cuda/` lands at `0.0224` source-fidelity, beating the older full-emblem tile-graph CUDA baseline at `0.0283` @@ -58,7 +58,7 @@ The most provisional areas are: ### Active pass -This pass is the first pruning pass after the new algorithm map: strip away the optimizer-era sidecars that are not part of the core tile-graph machine. +This pass is the first serious speed pass after the algorithm cleanup: replace the expensive per-component window-walk extractor with a grouped `(component_id, output_cell)` overlap reducer and then measure the tradeoff honestly. What landed: - `tile-graph` no longer participates in pipeline phase rerank probes; rerank is now a continuous-only wrapper @@ -73,24 +73,28 @@ What landed: - `tests/test_analysis.py` and `tests/test_tile_graph.py` now cover the new edge-only analysis path - `tests/test_pipeline.py` now locks in fixed continuous and tile-graph output hashes from a pre-refactor baseline so this architectural cut cannot silently drift behavior - `docs/tile-graph-algorithm-map.md` now describes the smaller machine directly rather than documenting the removed hybrid sidecar path +- the old per-component queue/seed/window extractor is gone; `_extract_source_region_tiles(...)` now directly reduces `(component_id, output_cell)` overlaps into candidate shards +- the old stroke-aware slicing machinery and its tuning knobs are gone along with that extractor +- same-component continuation now defines tile-graph adjacency directly, replacing the previous search-heavy neighbor heuristic inside extraction Current implementation note: -- the earlier profiling result still stands: on the searched badge path, the solver loop is not the problem; the heavy costs are `infer_lattice(...)`, continuous-mode phase rerank probes, and `_extract_source_region_tiles(...)` +- the earlier profiling result still stands conceptually: on the searched badge path, the solver loop is not the problem - the direct-control path is specifically for dodging those costs during iteration when we already know which lattice we want to inspect - tile-graph now skips pipeline rerank entirely, so iteration no longer rebuilds multiple probe candidates before the final tile-graph solve -- on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, a fixed-lattice CUDA `tile-graph` run now took about `10.2s` on the first same-process run and about `2.1s` on the cached rerun, with the same output and a reported `tile_graph_model_cache_hit` on the second pass +- on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, the new grouped-extraction CUDA path now takes about `16.1s` on the first same-process run and about `6.0s` on the cached rerun, with `tile_graph_model_cache_hit=True` on the second pass - the current deep-dive diagnosis is that the fixed-lattice garbling is not a wrapper bug and not mainly a parity-solver bug; it is born in tile cutting, candidate truncation, local source support, and the first unary argmin assignment - the extraction coverage bug turned out to be real but secondary: the new `_extract_source_region_tiles(...)` overlap fill pass now guarantees that any output cell containing opaque sampled source pixels gets at least one extracted region bucket, but the pinned `126x126` badge output remains numerically unchanged (`0.4998626`), which narrows the remaining failure to candidate ranking and lattice-conditioned reference usage rather than empty region buckets - 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 - the latest core algorithm cut removes `source_region_stride`, removes portrait-based candidate ranking, removes sharp/edge foreground fallback injection, and replaces sampled-delta pairwise scoring with adjacency learned from extracted tiles - that cut also removes the dedicated `TileGraphSourceReference` entirely: tile-graph now grounds itself in full-resolution extracted tiles plus direct per-cell source summaries (`cell_mean_rgba`, `cell_alpha_mean`, `cell_edge_strength`) -- on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, the new cut recovers the fixed-lattice run from roughly `0.5055` source-fidelity down to `0.1814`; that is still weak, but it is no longer total algorithmic collapse +- on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, the portrait-free cut first recovered the fixed-lattice run from roughly `0.5055` source-fidelity down to `0.1814`; the new grouped-extraction speed pass regresses a bit to `0.2036` while making the same case roughly `6x` faster cold on CUDA +- a fresh CPU profile of the grouped extractor shows the bottleneck has moved: `build_tile_graph_model(...)` now takes about `61.0s` instead of `145.8s`, and `_extract_source_region_tiles(...)` no longer dominates; the new hot spot is `_segment_atomic_source_regions(...)` at about `59.5s` Next after that: - split tile-graph iteration into a cached model-build phase and a near-free solve phase so weight tuning does not keep paying the `~131s` source-region cutting bill - extend the current fixed-lattice path so repeated CLI runs can reuse cached model artifacts across processes instead of only within one Python process -- port or reformulate `_extract_source_region_tiles(...)` so it no longer spends most of the badge runtime in Python/NumPy window cutting +- replace or accelerate `_segment_atomic_source_regions(...)`, because that is now the cold-build bottleneck rather than extraction - decide whether the continuous path should stay the only global search/rerank engine or whether tile-graph eventually needs its own lighter lattice chooser - 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 @@ -118,8 +122,8 @@ Current status: - 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 - on the cleaned badge, the new CCL-backed tile-graph run under `artifacts/badge-tile-graph-ccl-cuda/` still lands at a weak `0.1800`, so the current CCL work should be treated as an infrastructure/performance step rather than a visual-quality fix -- the new stroke-aware badge run under `artifacts/badge-tile-graph-stroke-v2-cuda/` still does not solve the guard-contour problem and slightly regresses to `0.1832`, so the next stroke pass should be judged on whether it changes the source-side path model more fundamentally -- older searched-path profiling showed that iteration was dominated by preprocessing work rather than the tile-graph parity solver: about `44.1s` in inference, `223.9s` in low-confidence rerank probes, and `142.3s` in the final selected-candidate build; the pruned tile-graph path no longer pays those rerank probes, but `_extract_source_region_tiles(...)` is still the heavy remaining step +- the new grouped-extraction badge run under `artifacts/tile-graph-reducebykey-v1-badge-126/` is the first speed-focused cut that actually changes the preprocessing profile in the right direction: cold CUDA time drops to about `16.1s`, warm cached time is about `6.0s`, and the remaining quality problem is contour precision rather than total collapse +- the latest CPU profile shows that preprocessing is still the story, but the story changed: extraction is no longer the monster, and connected-component labeling / label compression is now the thing to attack next Important note: - the checkerboard is baked into the source image, not transparency @@ -131,7 +135,7 @@ What future work should improve here: - 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 - an atomic candidate builder that more closely matches real source-region ownership and consumption -- a GPU-friendly reformulation of the atomic candidate extraction stage +- a faster labeler or grouped label-reduction path so tile-graph can keep the new simpler overlap model without paying for label propagation so heavily - some form of coordinated local relaxation so nearby cells can move together toward a globally better contour ## Suggested next milestones diff --git a/docs/tile-graph-algorithm-map.md b/docs/tile-graph-algorithm-map.md index 33b539d..aef903e 100644 --- a/docs/tile-graph-algorithm-map.md +++ b/docs/tile-graph-algorithm-map.md @@ -2,47 +2,43 @@ ## Why this map exists -The tile-graph path is trying to do one simple, stubborn thing: +The tile-graph path is trying to do one stubbornly honest thing: -Take a fake-pixel mural, cut it into honest one-cell tiles that really belong to the source image, learn which tiles naturally sit beside which, and then lay those tiles back down on a clean output lattice. +Take a fake-pixel mural, split it into real opaque paint islands, slice those islands by the chosen output lattice, learn which slices continue into neighboring cells, and then reassemble the image using only those real slices. -That is the dream. +When the output looks wrong, this map is here to answer: -When the output looks wrong, this map is here to answer one question: - -Where, exactly, did the machine start lying? +Where did the machine start lying? ## One-sentence machine The current tile-graph machine is: -`source image -> edge scout -> atomic opaque regions -> one-cell tile cuts -> per-cell candidate buckets -> learned tile adjacency -> local discrete assignment` +`source image -> edge scout -> connected components -> component-cell overlap reduction -> per-cell candidate buckets -> learned same-component adjacency -> local discrete assignment` -Or in more visual language: +Or, in pictures: - the source image is the mural -- connected components are the paint islands -- extracted tiles are the stones we chip out of that mural -- candidate buckets are the bins beside each slot in the final mosaic -- the solver is the mason choosing which chipped stone goes in each slot while trying to keep neighboring stones compatible +- connected components are paint islands +- the chosen lattice is a sheet of graph paper laid over the mural +- the overlap reducer is a cookie cutter that stamps each island against each graph-paper square it actually touches +- the solver is the mason choosing which stamped shard goes in each final slot ## What was cut away -The machine used to carry several ideas that did not belong: +These ideas are gone: - subsampling the mural before extraction with `source_region_stride` -- ranking region candidates by resemblance to a lattice portrait (`sharp_rgba`, `edge_rgba`) -- papering over missing foreground tiles by injecting `sharp_pixel` or `edge_pixel` -- pretending a learned tile graph existed while actually using sampled color deltas from one cell away in source space -- building and hauling a dedicated tile-graph source-reference object through the core loop +- per-component seed walking and window stepping +- stroke-specific PCA slicing +- ranking candidates by resemblance to a lattice portrait (`sharp_rgba`, `edge_rgba`) +- injecting `sharp_pixel` / `edge_pixel` when extraction failed +- sampled one-cell-away RGBA deltas pretending to be a graph +- a dedicated tile-graph source-reference object -Those are gone. +That matters because all of those were ways for the machine to stop listening to the actual source-owned tiles. -That matters because they were smuggling continuous-solver assumptions back into a path that is supposed to be source-owned and discrete. - -## The living machine - -### Stage 0: Pipeline chooses a lattice +## Stage 0: The pipeline picks a ruler Function: @@ -57,16 +53,15 @@ Inputs: - `phase_x` - `phase_y` -What it means: +Meaning: -The pipeline chooses the grid we are going to commit to. This is not the tile-graph solver yet. This is just choosing the size of the chessboard before we decide which pieces go on it. +Before anything else, the system chooses the graph paper. -Important truth: +If the graph paper is wrong, every later step is cutting the mural against the wrong ruler. -- if the lattice is wrong, every later step is forced to cut the mural against the wrong ruler -- but once the lattice is pinned, the tile-graph path now runs directly on that lattice without phase-rerank probes or hybrid sidecars +Once the lattice is pinned, the tile-graph path now runs directly on that lattice. It no longer burns time on phase-rerank probes or hybrid sidecars. -### Stage 1: Edge scout +## Stage 1: Edge scout Function: @@ -76,71 +71,66 @@ Output: - `TileGraphSourceAnalysis.edge_map` -What it means: +Meaning: -This stage is the scout walking over the mural with a lantern, marking where the sharp cliffs are. It is not deciding colors. It is only pointing out where the paint changes abruptly. +This is a scout walking over the mural with a lantern and chalk, marking where the paint changes sharply. + +It is not deciding colors. It is only measuring local cliff faces. Important variable: - `edge_map[y, x]` - - a per-source-pixel edge strength in `[0, 1]` Why it exists: -- region extraction needs to know where a component has its sharpest internal detail -- candidate budgeting uses edge strength to give difficult cells more room +- connected components need a detail hint +- candidate budgeting still uses edge-heavy cells as "hard mode" -### Stage 2: Source pixels are projected onto the output lattice +## Stage 2: Every source pixel is assigned to an output cell Function: -- `build_tile_graph_model(...)` in `src/repixelizer/tile_graph.py` +- `build_tile_graph_model(...)` Key variables: - `cell_w` - `cell_h` -- `projected_coord_x` -- `projected_coord_y` - `projected_flat_index` - -What it means: - -Every source pixel is asked: - -"If this final grid is the one we believe in, which output cell do you live under?" - -This is like dropping a transparent graph-paper sheet over the mural and writing the output cell index on every source pixel. - -Derived summaries: - -- `cell_counts_flat` - - how many source pixels fall into each output cell - `cell_mean_rgba_flat` - - the average source color already living under that cell -- `cell_alpha_support_flat` - - the strongest alpha seen in that cell - `cell_alpha_mean_flat` - - the average alpha in that cell +- `cell_alpha_support_flat` - `cell_edge_strength_flat` - - the strongest edge signal seen in that cell -These summaries are not a tile portrait from a separate source-reference system. They are direct summaries of the actual source pixels already assigned to that output cell under the chosen lattice. +Meaning: -### Stage 3: Atomic opaque regions +Imagine dropping transparent graph paper over the mural and writing the output cell index on every source pixel. -Function: +That gives us direct per-cell source summaries: + +- average color already under the slot +- average alpha already under the slot +- strongest alpha under the slot +- strongest edge under the slot + +These are not a separate portrait of what the slot "should" be. They are literally summaries of the pixels already living under that slot. + +## Stage 3: Connected components + +Functions: - `_segment_atomic_source_regions(...)` -- CPU path: `_segment_atomic_source_regions_cpu(...)` +- `_segment_atomic_source_regions_cpu(...)` -Inputs: +Output: -- full-resolution `source_rgba` -- full-resolution `edge_map` -- alpha/color join thresholds +- `AtomicRegionLabeling` + - `pixel_linear` + - `component_ids` + - `component_sizes` + - `component_count` -What it means: +Meaning: Now the mural is split into paint islands. @@ -150,168 +140,106 @@ Pixels join the same island only if: - their premultiplied colors are similar enough - their alpha is similar enough -This is not a "cluster by vibes" step. It is literal connected-component labeling over the full-resolution mural. +This is literal connected-component labeling over the full-resolution source. Important consequence: -- every opaque source pixel belongs to some component -- if an occupied output cell later ends up with no real extracted tile, that is not normal; it is a bug or a failure in cutting/projection logic +- every opaque source pixel belongs to exactly one component -### Stage 4: Cut one-cell tiles out of each component +## Stage 4: Reduce component-cell overlaps Function: - `_extract_source_region_tiles(...)` -Inputs: +This is the new heart of the machine. -- `components` -- `flat_rgba` -- `flat_edge` -- `flat_x` -- `flat_y` -- lattice geometry: `cell_w`, `cell_h`, `phase_x`, `phase_y` +Instead of walking each component with seed queues and cell-sized windows, the reducer does something much simpler: -What it means: +1. take every opaque labeled source pixel +2. project it onto the chosen output lattice +3. form a compound key: `(component_id, output_cell)` +4. sort by that key +5. reduce each run into one candidate shard -This is the quarry. +This is the same basic mental pattern as `reduce_by_key` / `segment_reduce`: sort once, then aggregate consecutive runs. -For each paint island, we try to chip out cell-sized stones. - -There are two modes: - -- ordinary components: - - start from the centroid and edge peak - - accept a one-cell window if it contains enough of the component - - march outward one cell at a time -- elongated stroke-like components: - - estimate the component’s principal axis - - seed along that axis - - cut tiles in bands so a long stroke is not immediately turned into a rectangular blob - -Important variables inside each extracted tile: +What each reduced shard stores: - `rep_linear` - - the literal source pixel index chosen to represent the tile - `rep_rgba` - - the tile color - `area_ratio` - - how much component area the tile owns relative to one output cell - `coverage` - - `area_ratio` clipped to `[0, 1]` - `edge_peak` - - strongest edge value inside that tile footprint -- `source_center_x`, `source_center_y` - - where this tile lives in source space +- `source_center_x` +- `source_center_y` - `coord_x`, `coord_y`, `flat_index` - - which output cell this tile projects into - `component_id` - - which paint island it came from Natural-language picture: -The component is a strip of stained glass. We are cutting out little panes, each about one output cell wide, and tagging where each pane came from. +Instead of hand-carving the island with a pocketknife, we stamp every island against the graph paper and keep one shard per `(island, cell)` overlap. -### Stage 5: Fill empty occupied cells honestly +## Stage 5: Candidate buckets Still inside: - `_extract_source_region_tiles(...)` +- `_select_source_region_candidates(...)` -What it means: +Meaning: -After the first pass of tile cutting, some output cells may still have no tile in their bucket even though component pixels really do overlap them. +Each output cell now gets a bucket of legal shards. -To fix that, the builder does an overlap-based empty-cell fill: +The reducer first creates every real `(component, cell)` overlap shard. -- for each component, look at every output cell it physically overlaps -- if that bucket is still empty, add one overlap tile there - -This matters because: - -- a foreground cell is allowed to have multiple candidates -- a foreground cell is not allowed to have zero real extracted candidates - -Background-only cells are different: - -- they have no opaque source ownership -- they get a transparent candidate later in model building - -### Stage 6: Learn adjacency from extracted tiles - -Still inside: - -- `_extract_source_region_tiles(...)` - -What it means: - -Each extracted tile now looks one cell away in all four directions and asks: - -"Among the tiles that actually landed in that neighboring output slot, which one looks like my real neighbor back in source space?" - -For each direction, we search the neighboring bucket and choose the best observed neighbor by: - -- source-center displacement closeness -- same-component tie preference -- area fit -- coverage -- edge strength - -Stored per tile: - -- `neighbor_rgba[4, 4]` -- `neighbor_mask[4]` - -This is the first honest adjacency graph the tile-graph path has had. - -It is not: - -- sampled one-cell-away colors from the source image -- a lattice portrait -- a guessed delta field - -It is: - -- "this extracted stone most naturally sat beside that extracted stone" - -### Stage 7: Build per-cell candidate buckets - -Function: - -- `build_tile_graph_model(...)` -- helper: `_select_source_region_candidates(...)` - -What it means: - -Each output cell gets a bucket of legal stones. - -Candidate selection now keeps only tile-owned signals: +Then bucket pruning keeps the best local candidates using only tile-owned evidence: - area fit - coverage - edge peak -It does not use: - -- `sharp_rgba` similarity -- `edge_rgba` similarity -- injected sharp/edge fallback pixels - Important rule: -- if a cell has real source-owned region tiles, those are the truth -- if a cell has no opaque source support, it gets a transparent candidate -- if a cell has opaque source support but no extracted region candidate, the build fails loudly +- if a cell has opaque source support, it must have at least one real extracted shard +- if pruning would empty a supported cell, the best overlap shard for that cell is kept anyway +- background or near-background cells can also get a transparent candidate -That last rule is deliberate. The old machine would quietly make something up. The new one would rather stop the line than lie. +What this prevents: -Additional background option: +- silent empty foreground buckets +- fake sharp/edge fallback pixels -- if `cell_alpha_mean_flat < 0.98`, a transparent candidate is also added +## Stage 6: Learn adjacency from the reduced shards -This lets the solver choose "leave this slot empty" in partially occupied or fringe cells instead of being forced to place an opaque stone just because a tiny sliver of source detail touched the slot. +Still inside: -### Stage 8: The solver model +- `_extract_source_region_tiles(...)` + +Meaning: + +Each shard now asks a simpler question than before: + +"Does my same component continue into the cell to the right, down, left, or up?" + +If a same-component overlap shard exists in the neighboring cell, that shard's representative color becomes the expected neighbor color for that direction. + +Stored per shard: + +- `neighbor_rgba[4, 4]` +- `neighbor_mask[4]` + +This is an intentionally narrower graph than the earlier heuristic: + +- it only learns continuity that the same connected paint island actually exhibits +- it does not try to smooth across boundaries between different components + +That is a carefully chosen bias: + +- continuity within a real island is trustworthy +- continuity across a contour is where fake smoothing usually starts + +## Stage 7: Solver model Dataclass: @@ -331,45 +259,41 @@ Current essential fields: - `cell_alpha_mean` - `cell_edge_strength` -What it means: +Meaning: -This is the parts tray beside the machine. +This is the tray of legal shards plus the local evidence under each output slot. -- the first block describes the legal stones -- the second block says which stones belong to which output slot -- the last block is the local source evidence already living under each slot - -### Stage 9: Unary cost +## Stage 8: Unary cost Function: - `_tile_graph_unary_cost_torch(...)` -Current unary terms: +Terms: - `color_error` - - how far the candidate color is from the actual mean source color already under that output cell + - candidate versus actual mean color already under the slot - `area_error` - - how far the candidate footprint is from one cell + - how close the shard is to one full cell - `alpha_error` - - how far the candidate alpha is from the actual mean alpha under that cell + - candidate alpha versus actual mean alpha under the slot - `coverage_error` - - how incomplete the candidate’s source ownership is + - how incomplete the overlap is - `edge_error` - - how far the candidate’s edge peak is from the strongest edge already under that cell + - candidate edge peak versus strongest edge under the slot -Natural-language picture: +Meaning: -This is the bouncer at each slot saying: +The unary cost is the bouncer at each slot asking: -"Does this stone even look like it belongs in the patch of mural that sits under this square?" +"Does this shard honestly fit the patch of mural under this square?" Important difference from the old machine: - this is local source support -- not a separate lattice portrait pretending to know the answer in advance +- not a pre-baked portrait from another reference object -### Stage 10: Pairwise adjacency cost +## Stage 9: Pairwise cost Functions: @@ -379,24 +303,18 @@ Functions: - `_pair_penalty_option_down_torch(...)` - `_pair_penalty_option_up_torch(...)` -What it means: +Meaning: Now the mason cares about seams. -If tile A says it naturally expects tile B to the right, then choosing a very different right neighbor should cost something. +If a shard says "my component really continues to the right as color X," then placing a very different right neighbor costs something. -The current pairwise cost compares: +This is narrower than the old fake graph, but more honest: -- the actually chosen neighbor color -- against the neighbor color learned from extracted tiles +- it rewards real within-component continuation +- it stops encouraging smoothing across unrelated components -This is now a real observed compatibility signal. - -It is no longer: - -- "the source image one cell away from this center happened to look like X, so maybe that is the right delta" - -### Stage 11: Checkerboard parity refinement +## Stage 10: Parity refinement Function: @@ -404,119 +322,106 @@ Function: What it does: -1. build the model +1. build candidate buckets 2. compute unary cost 3. choose an initial argmin candidate per cell -4. alternate parity updates over the grid -5. keep the initial assignment if the refinement step makes source-lattice fidelity worse +4. alternate parity updates across the grid +5. if refinement worsens source-fidelity, keep the initial assignment -Natural-language picture: +Meaning: -The mason lays down a first draft, then alternates black squares and white squares of the chessboard, trying to improve local seam compatibility without moving everything at once. +The mason lays an initial mosaic, then alternates black and white squares of the chessboard trying to improve local seams without moving everything at once. Important truth: -- this is still a local greedy/discrete smoother -- it is not a global optimal graph solver -- if the initial candidate buckets are wrong, parity updates will not save the image +- this is still a local discrete optimizer +- it is not globally optimal +- if the buckets are wrong, refinement will not save the image ## Data dictionary -### `flat_rgba_np` +### `AtomicRegionLabeling.pixel_linear` -- full-resolution source pixels flattened into `[pixel, rgba]` +- flattened source-pixel indices for every opaque labeled pixel + +### `AtomicRegionLabeling.component_ids` + +- component id for each `pixel_linear` ### `projected_flat_index` -- for each source pixel, which output cell it belongs to under the chosen lattice +- output-cell index for each source pixel under the chosen lattice -### `components` +### `compound_key` -- connected opaque paint islands in the source +- `(component_id * output_area) + output_cell` + +This is the reducer's key. Every identical key means "these pixels belong to the same component-cell overlap shard." ### `region_buckets[flat_index]` -- all extracted real tiles currently legal for one output cell - -### `candidate_rgba` - -- literal colors the solver is allowed to place +- all legal source-owned shards currently available for one output cell ### `candidate_neighbor_rgba` -- per-candidate remembered neighbor colors from real extracted-tile adjacency +- expected neighboring shard colors learned from same-component continuation -### `cell_mean_rgba` - -- actual average source color already living under each output cell - -### `cell_alpha_mean` - -- actual average source alpha already living under each output cell - -### `cell_edge_strength` - -- strongest edge signal already living under each output cell - -## What this machine is trying to maximize +## What this machine is optimizing for In plain language: -- choose only real, source-owned stones -- prefer stones that genuinely fit the patch of mural beneath the slot -- prefer neighboring stones that were actually observed beside each other -- never silently hallucinate a foreground tile just because the lattice wanted one +- choose only real source-owned shards +- prefer shards that genuinely fit the local patch of mural +- preserve continuation within real paint islands +- refuse to invent a foreground shard when extraction did not produce one -## What would make the machine lie now +## What can still go wrong -If the output is still wrong after this pass, the likely failure points are no longer the old ones. +The likely remaining lies are now: -The remaining likely lies are: - -- region cutting still slices a component into the wrong one-cell panes -- candidate truncation still throws away the only good tile for a difficult slot -- learned adjacency still chooses the wrong observed neighbor when multiple neighbors are plausible -- the parity solver is still too local and can settle into a bad arrangement even with legal tiles -- the lattice itself is wrong, which means the whole quarry is being cut against the wrong ruler +- the lattice is wrong +- component labeling merged or split the wrong source regions +- the direct overlap shard for a difficult contour cell is too coarse +- candidate truncation threw away the only good shard in a crowded cell +- same-component adjacency is too narrow to express a more complex local motif +- the parity solver settles into a mediocre local arrangement ## Current reality check -Pinned cleaned badge run: +Pinned cleaned badge run after the reduce-by-key extraction rewrite: - target: `126x126` - phase: `(0.0, -0.2)` - mode: `tile-graph` - phase rerank: off -Current result after the latest cut: +Measured result: -- initial source-fidelity: `0.1814` -- final source-fidelity: `0.1814` -- the solver kept the initial assignment because refinement did not improve it +- initial source-fidelity: `0.2036` +- final source-fidelity: `0.2036` +- the solver kept the initial assignment -That is still not good enough, but it is dramatically more honest than the earlier architectural-collapse run around `0.5055`. - -In other words: - -- the machine is no longer babbling nonsense -- it is still not carving the right stones often enough +That is a small quality regression from the previous `0.1814` cut, but it is still far better than the earlier fixed-lattice collapse around `0.5055`. ## Big-picture judgment -Does the current machine make sense? +Does the current machine make more sense than the old one? -Yes, much more than before. +Yes. -The data now flows in one believable direction: +The dataflow is now much cleaner: 1. observe the mural -2. find real paint islands -3. cut one-cell stones out of those islands -4. learn which stones really sat beside which -5. choose among those stones for each output slot +2. label the paint islands +3. stamp those islands against the chosen lattice +4. keep the resulting overlap shards +5. learn only the continuation those shards actually exhibit +6. assemble the output from those shards -That is a coherent mental model. +That is the right shape. -What is still missing is not conceptual honesty. +The remaining problem is no longer "why is the machine inventing stories?" -What is still missing is quality in the stone cutting and compatibility choice. +It is: + +"Are these overlap shards expressive enough for the finest contours, and can we compute the labels fast enough?" diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index 48a99b3..31dface 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -80,9 +80,6 @@ class SolverHyperParams: tile_graph_component_alpha_threshold: float = 0.12 tile_graph_source_region_min_area_ratio: float = 0.06 tile_graph_source_region_window_coverage: float = 0.12 - tile_graph_stroke_linearity_threshold: float = 0.72 - tile_graph_stroke_step_scale: float = 0.95 - tile_graph_stroke_minor_limit_scale: float = 0.55 tile_graph_area_weight: float = 0.03 tile_graph_alpha_weight: float = 0.25 tile_graph_coverage_weight: float = 0.05 diff --git a/src/repixelizer/tile_graph.py b/src/repixelizer/tile_graph.py index e174daa..b0accc7 100644 --- a/src/repixelizer/tile_graph.py +++ b/src/repixelizer/tile_graph.py @@ -1,7 +1,7 @@ from __future__ import annotations import hashlib -from collections import OrderedDict, deque +from collections import OrderedDict from dataclasses import dataclass, replace from typing import Any @@ -75,6 +75,14 @@ class TileGraphBuildStats: cache_hit: bool = False +@dataclass(slots=True) +class AtomicRegionLabeling: + pixel_linear: np.ndarray + component_ids: np.ndarray + component_sizes: np.ndarray + component_count: int + + _TILE_GRAPH_MODEL_CACHE_MAX_ENTRIES = 8 _TILE_GRAPH_MODEL_CACHE: OrderedDict[tuple[object, ...], TileGraphModel] = OrderedDict() _TILE_GRAPH_BUILD_STATS_CACHE: OrderedDict[tuple[object, ...], TileGraphBuildStats] = OrderedDict() @@ -106,9 +114,6 @@ def _tile_graph_model_cache_key( "tile_graph_component_alpha_threshold", "tile_graph_source_region_min_area_ratio", "tile_graph_source_region_window_coverage", - "tile_graph_stroke_linearity_threshold", - "tile_graph_stroke_step_scale", - "tile_graph_stroke_minor_limit_scale", "tile_graph_area_weight", "tile_graph_alpha_weight", "tile_graph_coverage_weight", @@ -167,51 +172,28 @@ def _make_regular_uv(height: int, width: int, target_height: int, target_width: return uv.astype(np.float32) -def _component_axis_stats(member_x: np.ndarray, member_y: np.ndarray) -> tuple[float, float, float, float, float]: - if member_x.size <= 1: - return 1.0, 0.0, 0.0, 0.0, 0.0 - centered_x = member_x.astype(np.float32) - float(np.mean(member_x)) - centered_y = member_y.astype(np.float32) - float(np.mean(member_y)) - covariance = np.asarray( - [ - [float(np.mean(centered_x * centered_x)), float(np.mean(centered_x * centered_y))], - [float(np.mean(centered_x * centered_y)), float(np.mean(centered_y * centered_y))], - ], - dtype=np.float32, - ) - eigenvalues, eigenvectors = np.linalg.eigh(covariance) - order = np.argsort(eigenvalues) - minor_value = float(max(0.0, eigenvalues[order[0]])) - major_value = float(max(0.0, eigenvalues[order[1]])) - major_vector = eigenvectors[:, order[1]].astype(np.float32) - norm = float(np.linalg.norm(major_vector)) - if norm <= 1e-6: - return 1.0, 0.0, 0.0, 0.0, 0.0 - major_vector /= norm - axis_x = float(major_vector[0]) - axis_y = float(major_vector[1]) - projections = centered_x * axis_x + centered_y * axis_y - perpendicular = centered_x * (-axis_y) + centered_y * axis_x - major_span = float(np.max(projections) - np.min(projections)) if projections.size else 0.0 - minor_span = float(np.max(perpendicular) - np.min(perpendicular)) if perpendicular.size else 0.0 - linearity = float(np.clip(1.0 - (minor_value / max(major_value, 1e-6)), 0.0, 1.0)) - return axis_x, axis_y, linearity, major_span, minor_span - - -def _project_source_point_to_output_coord( +def _project_pixels_to_output_cells( *, - x: float, - y: float, + pixel_x: np.ndarray, + pixel_y: np.ndarray, cell_w: float, cell_h: float, target_width: int, target_height: int, phase_x: float, phase_y: float, -) -> tuple[int, int, int]: - coord_x = int(np.clip(np.floor((x + 0.5) / max(cell_w, 1e-6) - phase_x), 0, max(0, target_width - 1))) - coord_y = int(np.clip(np.floor((y + 0.5) / max(cell_h, 1e-6) - phase_y), 0, max(0, target_height - 1))) - return coord_y, coord_x, coord_y * target_width + coord_x +) -> np.ndarray: + coord_x = np.clip( + np.floor((pixel_x.astype(np.float32) + 0.5) / max(cell_w, 1e-6) - phase_x).astype(np.int32), + 0, + max(0, target_width - 1), + ) + coord_y = np.clip( + np.floor((pixel_y.astype(np.float32) + 0.5) / max(cell_h, 1e-6) - phase_y).astype(np.int32), + 0, + max(0, target_height - 1), + ) + return coord_y * target_width + coord_x def _segment_atomic_source_regions_cpu( @@ -241,7 +223,10 @@ def _segment_atomic_source_regions_cpu( pixel_count = width * height visited = ~opaque.reshape(-1) - components: list[dict[str, float | int | np.ndarray]] = [] + pixel_linear_parts: list[np.ndarray] = [] + component_ids_parts: list[np.ndarray] = [] + component_sizes: list[int] = [] + component_id = 0 for seed in range(pixel_count): if visited[seed]: @@ -275,32 +260,23 @@ def _segment_atomic_source_regions_cpu( visited[neighbor] = True queue.append(neighbor) member_array = np.asarray(members, dtype=np.int32) - member_y = (member_array // width).astype(np.float32) - member_x = (member_array % width).astype(np.float32) - centroid_y = float(np.mean(member_y)) - centroid_x = float(np.mean(member_x)) - axis_x, axis_y, linearity, major_span, minor_span = _component_axis_stats(member_x, member_y) - centroid_dist = (member_y - centroid_y) ** 2 + (member_x - centroid_x) ** 2 - centroid_pick = int(np.argmin(centroid_dist)) - edge_pick = int(np.argmax(flat_edge[member_array])) - edge_peak = float(np.max(flat_edge[member_array])) - components.append( - { - "member_linear": member_array, - "edge_peak_linear": int(member_array[edge_pick]), - "edge_peak": edge_peak, - "centroid_linear": int(member_array[centroid_pick]), - "centroid_x": centroid_x, - "centroid_y": centroid_y, - "axis_x": axis_x, - "axis_y": axis_y, - "linearity": linearity, - "major_span": major_span, - "minor_span": minor_span, - "size": int(member_array.size), - } + pixel_linear_parts.append(member_array) + component_ids_parts.append(np.full(member_array.shape[0], component_id, dtype=np.int32)) + component_sizes.append(int(member_array.size)) + component_id += 1 + if not pixel_linear_parts: + return AtomicRegionLabeling( + pixel_linear=np.zeros((0,), dtype=np.int32), + component_ids=np.zeros((0,), dtype=np.int32), + component_sizes=np.zeros((0,), dtype=np.int32), + component_count=0, ) - return components + return AtomicRegionLabeling( + pixel_linear=np.concatenate(pixel_linear_parts).astype(np.int32, copy=False), + component_ids=np.concatenate(component_ids_parts).astype(np.int32, copy=False), + component_sizes=np.asarray(component_sizes, dtype=np.int32), + component_count=component_id, + ) def _segment_atomic_source_regions( @@ -311,12 +287,17 @@ def _segment_atomic_source_regions( color_threshold: float, alpha_threshold: float, device: str, -) -> list[dict[str, float | int | np.ndarray]]: +) -> AtomicRegionLabeling: torch = _require_torch() resolved_device = _resolve_device(torch, device) height, width = source_rgba.shape[:2] if height <= 0 or width <= 0: - return [] + return AtomicRegionLabeling( + pixel_linear=np.zeros((0,), dtype=np.int32), + component_ids=np.zeros((0,), dtype=np.int32), + component_sizes=np.zeros((0,), dtype=np.int32), + component_count=0, + ) if resolved_device == "cpu" and (height * width) <= 4096: flat_rgba = source_rgba.reshape(-1, source_rgba.shape[-1]).astype(np.float32) @@ -338,7 +319,12 @@ def _segment_atomic_source_regions( alpha_map = source_t[..., 3] opaque = alpha_map >= alpha_floor if not bool(opaque.any().item()): - return [] + return AtomicRegionLabeling( + pixel_linear=np.zeros((0,), dtype=np.int32), + component_ids=np.zeros((0,), dtype=np.int32), + component_sizes=np.zeros((0,), dtype=np.int32), + component_count=0, + ) premul = source_t.clone() premul[..., :3] *= premul[..., 3:4] @@ -390,142 +376,30 @@ def _segment_atomic_source_regions( valid_linear_t = torch.nonzero(labels.reshape(-1) < sentinel, as_tuple=False).flatten() if valid_linear_t.numel() <= 0: - return [] + return AtomicRegionLabeling( + pixel_linear=np.zeros((0,), dtype=np.int32), + component_ids=np.zeros((0,), dtype=np.int32), + component_sizes=np.zeros((0,), dtype=np.int32), + component_count=0, + ) component_labels_t = labels.reshape(-1)[valid_linear_t] _, inverse_t = torch.unique(component_labels_t, sorted=True, return_inverse=True) valid_linear = valid_linear_t.detach().cpu().numpy().astype(np.int32) inverse = inverse_t.detach().cpu().numpy().astype(np.int32) - flat_edge = edge_t.reshape(-1).detach().cpu().numpy().astype(np.float32) - - order = np.argsort(inverse, kind="stable") - sorted_linear = valid_linear[order] - sorted_inverse = inverse[order] - boundaries = np.flatnonzero( - np.concatenate(( - np.asarray([True], dtype=bool), - sorted_inverse[1:] != sorted_inverse[:-1], - )) + component_count = int(inverse.max()) + 1 if inverse.size else 0 + component_sizes = np.bincount(inverse, minlength=component_count).astype(np.int32) + return AtomicRegionLabeling( + pixel_linear=valid_linear, + component_ids=inverse, + component_sizes=component_sizes, + component_count=component_count, ) - ends = np.concatenate((boundaries[1:], np.asarray([sorted_inverse.shape[0]], dtype=np.int64))) - - components: list[dict[str, float | int | np.ndarray]] = [] - for start, end in zip(boundaries.tolist(), ends.tolist(), strict=False): - member_array = sorted_linear[start:end].astype(np.int32, copy=False) - member_y = (member_array // width).astype(np.float32) - member_x = (member_array % width).astype(np.float32) - centroid_y = float(np.mean(member_y)) - centroid_x = float(np.mean(member_x)) - axis_x, axis_y, linearity, major_span, minor_span = _component_axis_stats(member_x, member_y) - centroid_dist = (member_y - centroid_y) ** 2 + (member_x - centroid_x) ** 2 - centroid_pick = int(np.argmin(centroid_dist)) - edge_pick = int(np.argmax(flat_edge[member_array])) - edge_peak = float(np.max(flat_edge[member_array])) - components.append( - { - "member_linear": member_array, - "edge_peak_linear": int(member_array[edge_pick]), - "edge_peak": edge_peak, - "centroid_linear": int(member_array[centroid_pick]), - "centroid_x": centroid_x, - "centroid_y": centroid_y, - "axis_x": axis_x, - "axis_y": axis_y, - "linearity": linearity, - "major_span": major_span, - "minor_span": minor_span, - "size": int(member_array.size), - } - ) - return components - - -def _stroke_window_mask( - *, - member_x: np.ndarray, - member_y: np.ndarray, - remaining: np.ndarray, - center_x: float, - center_y: float, - axis_x: float, - axis_y: float, - major_half: float, - minor_half: float, -) -> np.ndarray: - delta_x = member_x - float(center_x) - delta_y = member_y - float(center_y) - major_delta = delta_x * axis_x + delta_y * axis_y - minor_delta = delta_x * (-axis_y) + delta_y * axis_x - return remaining & (np.abs(major_delta) <= major_half) & (np.abs(minor_delta) <= minor_half) - - -def _stroke_seed_centers( - *, - component: dict[str, float | int | np.ndarray], - member_x: np.ndarray, - member_y: np.ndarray, - flat_x: np.ndarray, - flat_y: np.ndarray, - cell_w: float, - cell_h: float, - step_scale: float, -) -> list[tuple[float, float]]: - axis_x = float(component["axis_x"]) - axis_y = float(component["axis_y"]) - centroid_x = float(component["centroid_x"]) - centroid_y = float(component["centroid_y"]) - projection = (member_x - centroid_x) * axis_x + (member_y - centroid_y) * axis_y - if projection.size <= 0: - return [(centroid_y, centroid_x)] - step = max( - 1.0, - (abs(axis_x) * cell_w + abs(axis_y) * cell_h) * float(step_scale), - np.hypot(axis_x * cell_w, axis_y * cell_h) * 0.9, - ) - if step <= 1e-6: - return [(centroid_y, centroid_x)] - low = float(np.min(projection)) - high = float(np.max(projection)) - edge_peak_linear = int(component["edge_peak_linear"]) - anchor_x = float(flat_x[edge_peak_linear]) - anchor_y = float(flat_y[edge_peak_linear]) - anchor_projection = (anchor_x - centroid_x) * axis_x + (anchor_y - centroid_y) * axis_y - if not np.isfinite(anchor_projection): - anchor_projection = 0.0 - - seed_projections: list[float] = [anchor_projection] - cursor = anchor_projection - step - while cursor >= low - 0.5 * step: - seed_projections.append(cursor) - cursor -= step - cursor = anchor_projection + step - while cursor <= high + 0.5 * step: - seed_projections.append(cursor) - cursor += step - if len(seed_projections) == 1 and abs(high - low) > 0.6 * step: - center_projection = 0.5 * (low + high) - seed_projections.extend([center_projection - 0.5 * step, center_projection + 0.5 * step]) - - seen: set[tuple[int, int]] = set() - centers: list[tuple[float, float]] = [] - perpendicular = (member_x - centroid_x) * (-axis_y) + (member_y - centroid_y) * axis_x - for proj_center in sorted(seed_projections): - score = np.abs(projection - proj_center) + np.abs(perpendicular) * 0.35 - pick = int(np.argmin(score)) - center = (float(member_y[pick]), float(member_x[pick])) - key = (int(round(center[0] * 2.0)), int(round(center[1] * 2.0))) - if key in seen: - continue - seen.add(key) - centers.append(center) - if not centers: - centers.append((centroid_y, centroid_x)) - return centers def _extract_source_region_tiles( *, - components: list[dict[str, float | int | np.ndarray]], + labeling: AtomicRegionLabeling, flat_rgba: np.ndarray, flat_edge: np.ndarray, flat_x: np.ndarray, @@ -538,345 +412,134 @@ def _extract_source_region_tiles( phase_y: float, min_region_area_ratio: float, min_window_coverage: float, - stroke_linearity_threshold: float, - stroke_step_scale: float, - stroke_minor_limit_scale: float, ) -> list[list[dict[str, float | int | np.ndarray]]]: output_area = max(1, target_width * target_height) cell_area = max(cell_w * cell_h, 1e-6) - half_w = cell_w * 0.5 - half_h = cell_h * 0.5 buckets: list[list[dict[str, float | int | np.ndarray]]] = [[] for _ in range(output_area)] - min_region_pixels = max(1, int(np.ceil(cell_area * min_region_area_ratio))) - ordered_components = sorted( - components, - key=lambda comp: (abs(float(comp["size"]) - cell_area), -float(comp["edge_peak"])), - ) - primary_tiles: list[dict[str, float | int | np.ndarray]] = [] + if labeling.pixel_linear.size <= 0 or labeling.component_count <= 0: + return buckets - def overlap_tiles_for_component( - component: dict[str, float | int | np.ndarray], - *, - component_id: int, - ) -> list[dict[str, float | int | np.ndarray]]: - member_linear = np.asarray(component["member_linear"], dtype=np.int32) - if member_linear.size <= 0: - return [] - member_x = flat_x[member_linear].astype(np.float32) - member_y = flat_y[member_linear].astype(np.float32) - member_coord_x = np.clip( - np.floor((member_x + 0.5) / max(cell_w, 1e-6) - phase_x).astype(np.int32), - 0, - max(0, target_width - 1), + pixel_linear = labeling.pixel_linear.astype(np.int32, copy=False) + component_ids = labeling.component_ids.astype(np.int32, copy=False) + component_sizes = labeling.component_sizes.astype(np.int32, copy=False) + pixel_x = flat_x[pixel_linear].astype(np.int32, copy=False) + pixel_y = flat_y[pixel_linear].astype(np.int32, copy=False) + projected_flat_index = _project_pixels_to_output_cells( + pixel_x=pixel_x, + pixel_y=pixel_y, + cell_w=cell_w, + cell_h=cell_h, + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, + ).astype(np.int32, copy=False) + compound_key = component_ids.astype(np.int64) * output_area + projected_flat_index.astype(np.int64) + order = np.argsort(compound_key, kind="stable") + sorted_key = compound_key[order] + sorted_linear = pixel_linear[order] + sorted_component = component_ids[order] + sorted_cell = projected_flat_index[order] + sorted_x = pixel_x[order].astype(np.float32) + sorted_y = pixel_y[order].astype(np.float32) + sorted_edge = flat_edge[sorted_linear].astype(np.float32) + starts = np.concatenate( + ( + np.asarray([0], dtype=np.int64), + np.flatnonzero(sorted_key[1:] != sorted_key[:-1]) + 1, ) - member_coord_y = np.clip( - np.floor((member_y + 0.5) / max(cell_h, 1e-6) - phase_y).astype(np.int32), - 0, - max(0, target_height - 1), - ) - member_flat_index = member_coord_y * target_width + member_coord_x - order = np.argsort(member_flat_index, kind="stable") - sorted_flat_index = member_flat_index[order] - starts = np.concatenate( - ( - np.asarray([0], dtype=np.int64), - np.flatnonzero(sorted_flat_index[1:] != sorted_flat_index[:-1]) + 1, + ) + ends = np.concatenate((starts[1:], np.asarray([sorted_key.shape[0]], dtype=np.int64))) + counts = (ends - starts).astype(np.int32) + area_ratio = counts.astype(np.float32) / float(cell_area) + coverage = np.clip(area_ratio, 0.0, 1.0).astype(np.float32) + sum_x = np.add.reduceat(sorted_x, starts).astype(np.float32) + sum_y = np.add.reduceat(sorted_y, starts).astype(np.float32) + center_x = (sum_x / np.maximum(1, counts)).astype(np.float32) + center_y = (sum_y / np.maximum(1, counts)).astype(np.float32) + edge_peak = np.maximum.reduceat(sorted_edge, starts).astype(np.float32) + segment_component = sorted_component[starts].astype(np.int32, copy=False) + segment_cell = sorted_cell[starts].astype(np.int32, copy=False) + component_area_ratio = component_sizes[segment_component].astype(np.float32) / float(cell_area) + keep = (component_area_ratio >= float(min_region_area_ratio)) & (area_ratio >= float(min_window_coverage)) + + rep_linear = np.empty(starts.shape[0], dtype=np.int32) + for segment_index, (start, end) in enumerate(zip(starts.tolist(), ends.tolist(), strict=False)): + segment_linear = sorted_linear[start:end] + if segment_linear.size <= 0: + rep_linear[segment_index] = 0 + continue + if edge_peak[segment_index] > 0.0: + local_edge = sorted_edge[start:end] + rep_linear[segment_index] = int(segment_linear[int(np.argmax(local_edge))]) + else: + rep_linear[segment_index] = int(segment_linear[0]) + + kept_by_cell: dict[int, list[int]] = {} + all_by_cell: dict[int, list[int]] = {} + for segment_index, flat_index in enumerate(segment_cell.tolist()): + all_by_cell.setdefault(flat_index, []).append(segment_index) + if keep[segment_index]: + kept_by_cell.setdefault(flat_index, []).append(segment_index) + + for flat_index, segment_indices in all_by_cell.items(): + use_segments = kept_by_cell.get(flat_index) + if not use_segments: + best_segment = min( + segment_indices, + key=lambda idx: ( + abs(float(area_ratio[idx]) - 1.0), + -float(coverage[idx]), + -float(edge_peak[idx]), + ), ) - ) - ends = np.concatenate((starts[1:], np.asarray([sorted_flat_index.shape[0]], dtype=np.int64))) - tiles: list[dict[str, float | int | np.ndarray]] = [] - for start, end in zip(starts.tolist(), ends.tolist(), strict=False): - local_order = order[start:end] - cell_linear = member_linear[local_order] - cell_x = member_x[local_order] - cell_y = member_y[local_order] - cell_edge = flat_edge[cell_linear] - flat_index = int(sorted_flat_index[start]) - center_x = float(np.mean(cell_x)) - center_y = float(np.mean(cell_y)) - edge_peak = float(np.max(cell_edge)) if cell_edge.size else 0.0 - if edge_peak > 0.0: - rep_linear = int(cell_linear[int(np.argmax(cell_edge))]) - else: - centroid_dist = (cell_x - center_x) ** 2 + (cell_y - center_y) ** 2 - rep_linear = int(cell_linear[int(np.argmin(centroid_dist))]) - tiles.append( + use_segments = [best_segment] + for segment_index in use_segments: + buckets[flat_index].append( { - "rep_linear": rep_linear, - "rep_rgba": flat_rgba[rep_linear].astype(np.float32), - "area_ratio": float(cell_linear.shape[0] / cell_area), - "coverage": float(np.clip(cell_linear.shape[0] / cell_area, 0.0, 1.0)), - "edge_peak": edge_peak, - "source_center_x": center_x, - "source_center_y": center_y, + "rep_linear": int(rep_linear[segment_index]), + "rep_rgba": flat_rgba[int(rep_linear[segment_index])].astype(np.float32), + "area_ratio": float(area_ratio[segment_index]), + "coverage": float(coverage[segment_index]), + "edge_peak": float(edge_peak[segment_index]), + "source_center_x": float(center_x[segment_index]), + "source_center_y": float(center_y[segment_index]), "coord_x": int(flat_index % target_width), "coord_y": int(flat_index // target_width), - "flat_index": flat_index, - "component_id": component_id, + "flat_index": int(flat_index), + "component_id": int(segment_component[segment_index]), } ) - return tiles - for component_id, component in enumerate(ordered_components): - component_size = int(component["size"]) - if component_size < min_region_pixels: - continue - member_linear = np.asarray(component["member_linear"], dtype=np.int32) - member_x = flat_x[member_linear].astype(np.float32) - member_y = flat_y[member_linear].astype(np.float32) - remaining = np.ones(member_linear.shape[0], dtype=bool) - seed_queue: deque[tuple[float, float]] = deque() - seen_seeds: set[tuple[int, int]] = set() - accepted_any = False - axis_x = float(component["axis_x"]) - axis_y = float(component["axis_y"]) - linearity = float(component["linearity"]) - major_span = float(component["major_span"]) - minor_span = float(component["minor_span"]) - stroke_component = ( - linearity >= stroke_linearity_threshold - and major_span >= max(cell_w, cell_h) * 1.1 - and minor_span <= max(1.0, min(cell_w, cell_h) * stroke_minor_limit_scale) - ) - stroke_major_half = max( - 0.5, - 0.5 * (abs(axis_x) * cell_w + abs(axis_y) * cell_h) * stroke_step_scale, - 0.45 * np.hypot(axis_x * cell_w, axis_y * cell_h), - ) - stroke_minor_half = max(0.5, 0.5 * min(cell_w, cell_h) * stroke_minor_limit_scale) - component_tiles: list[dict[str, float | int | np.ndarray]] = [] - - def enqueue(seed_y: float, seed_x: float) -> None: - key = (int(round(seed_y * 2.0)), int(round(seed_x * 2.0))) - if key in seen_seeds: - return - seen_seeds.add(key) - seed_queue.append((float(seed_y), float(seed_x))) - - if stroke_component: - for seed_y, seed_x in _stroke_seed_centers( - component=component, - member_x=member_x, - member_y=member_y, - flat_x=flat_x, - flat_y=flat_y, - cell_w=cell_w, - cell_h=cell_h, - step_scale=stroke_step_scale, - ): - enqueue(seed_y, seed_x) - else: - enqueue(float(component["centroid_y"]), float(component["centroid_x"])) - edge_peak_linear = int(component["edge_peak_linear"]) - enqueue(float(flat_y[edge_peak_linear]), float(flat_x[edge_peak_linear])) - - while seed_queue: - center_y, center_x = seed_queue.popleft() - if stroke_component: - in_window = _stroke_window_mask( - member_x=member_x, - member_y=member_y, - remaining=remaining, - center_x=center_x, - center_y=center_y, - axis_x=axis_x, - axis_y=axis_y, - major_half=stroke_major_half, - minor_half=stroke_minor_half, - ) - if not np.any(in_window): - in_window = remaining & (np.abs(member_x - center_x) <= half_w) & (np.abs(member_y - center_y) <= half_h) - else: - in_window = remaining & (np.abs(member_x - center_x) <= half_w) & (np.abs(member_y - center_y) <= half_h) - footprint_count = int(np.count_nonzero(in_window)) - if footprint_count <= 0: + segment_index_by_key = { + int(segment_component[idx]) * output_area + int(segment_cell[idx]): idx + for idx in range(segment_cell.shape[0]) + } + rep_rgba = flat_rgba[rep_linear].astype(np.float32) + neighbor_rgba = np.zeros((segment_cell.shape[0], 4, 4), dtype=np.float32) + neighbor_mask = np.zeros((segment_cell.shape[0], 4), dtype=bool) + for segment_index, flat_index in enumerate(segment_cell.tolist()): + coord_y = int(flat_index // target_width) + coord_x = int(flat_index % target_width) + component_id = int(segment_component[segment_index]) + for direction, dy, dx in _DIRS: + neighbor_x = coord_x + dx + neighbor_y = coord_y + dy + if neighbor_x < 0 or neighbor_x >= target_width or neighbor_y < 0 or neighbor_y >= target_height: continue - area_ratio = float(footprint_count / cell_area) - if area_ratio < min_window_coverage: + neighbor_flat_index = neighbor_y * target_width + neighbor_x + neighbor_segment = segment_index_by_key.get(component_id * output_area + neighbor_flat_index) + if neighbor_segment is None: continue + neighbor_rgba[segment_index, direction] = rep_rgba[neighbor_segment] + neighbor_mask[segment_index, direction] = True - accepted_any = True - window_linear = member_linear[in_window] - window_x = member_x[in_window] - window_y = member_y[in_window] - window_edge = flat_edge[window_linear] - if stroke_component: - accepted_center_x = float(center_x * 0.75 + np.mean(window_x) * 0.25) - accepted_center_y = float(center_y * 0.75 + np.mean(window_y) * 0.25) - else: - accepted_center_x = float(np.mean(window_x)) - accepted_center_y = float(np.mean(window_y)) - edge_peak = float(np.max(window_edge)) if window_edge.size else 0.0 - if stroke_component: - center_dist = (window_x - accepted_center_x) ** 2 + (window_y - accepted_center_y) ** 2 - rep_linear = int(window_linear[int(np.argmin(center_dist))]) - elif edge_peak > 0.0: - rep_linear = int(window_linear[int(np.argmax(window_edge))]) - else: - centroid_dist = (window_x - accepted_center_x) ** 2 + (window_y - accepted_center_y) ** 2 - rep_linear = int(window_linear[int(np.argmin(centroid_dist))]) - - coord_y, coord_x, flat_index = _project_source_point_to_output_coord( - x=accepted_center_x, - y=accepted_center_y, - cell_w=cell_w, - cell_h=cell_h, - target_width=target_width, - target_height=target_height, - phase_x=phase_x, - phase_y=phase_y, - ) - component_tiles.append( - { - "rep_linear": rep_linear, - "rep_rgba": flat_rgba[rep_linear].astype(np.float32), - "area_ratio": area_ratio, - "coverage": float(np.clip(area_ratio, 0.0, 1.0)), - "edge_peak": edge_peak, - "source_center_x": accepted_center_x, - "source_center_y": accepted_center_y, - "coord_x": coord_x, - "coord_y": coord_y, - "flat_index": flat_index, - "component_id": component_id, - } - ) - remaining[in_window] = False - if int(np.count_nonzero(remaining)) < min_region_pixels: - break - if not stroke_component: - for _, dy, dx in _DIRS: - enqueue(accepted_center_y + dy * cell_h, accepted_center_x + dx * cell_w) - - if accepted_any: - if stroke_component and component_tiles: - dominant_vertical = abs(axis_y) >= abs(axis_x) - grouped_tiles: dict[int, list[dict[str, float | int | np.ndarray]]] = {} - for tile in component_tiles: - group_key = int(tile["coord_y"]) if dominant_vertical else int(tile["coord_x"]) - grouped_tiles.setdefault(group_key, []).append(tile) - filtered_tiles: list[dict[str, float | int | np.ndarray]] = [] - ordered_keys = sorted(grouped_tiles) - secondary_step = ( - axis_x / max(abs(axis_y), 1e-6) - if dominant_vertical - else axis_y / max(abs(axis_x), 1e-6) - ) - previous_primary: int | None = None - previous_secondary: float | None = None - - def tile_score( - tile: dict[str, float | int | np.ndarray], - *, - expected_secondary: float | None, - ) -> tuple[float, float, float, float]: - continuous_coord = ( - (float(tile["source_center_x"]) + 0.5) / max(cell_w, 1e-6) - phase_x - if dominant_vertical - else (float(tile["source_center_y"]) + 0.5) / max(cell_h, 1e-6) - phase_y - ) - discrete_coord = float(tile["coord_x"] if dominant_vertical else tile["coord_y"]) - path_error = abs(discrete_coord - expected_secondary) if expected_secondary is not None else 0.0 - return ( - path_error, - abs(discrete_coord - continuous_coord), - abs(float(tile["area_ratio"]) - 1.0), - -float(tile["coverage"]), - ) - - for group_key in ordered_keys: - group_tiles = grouped_tiles[group_key] - expected_secondary = None - if previous_primary is not None and previous_secondary is not None: - expected_secondary = previous_secondary + secondary_step * float(group_key - previous_primary) - - best_tile = min( - group_tiles, - key=lambda tile: tile_score(tile, expected_secondary=expected_secondary), - ) - filtered_tiles.append(best_tile) - previous_primary = group_key - previous_secondary = float(best_tile["coord_x"] if dominant_vertical else best_tile["coord_y"]) - component_tiles = filtered_tiles - for tile in component_tiles: - primary_tiles.append(tile) - continue - - component_area_ratio = float(component_size / cell_area) - if component_area_ratio < min_window_coverage: - continue - centroid_linear = int(component["centroid_linear"]) - center_x = float(component["centroid_x"]) - center_y = float(component["centroid_y"]) - coord_y, coord_x, flat_index = _project_source_point_to_output_coord( - x=center_x, - y=center_y, - cell_w=cell_w, - cell_h=cell_h, - target_width=target_width, - target_height=target_height, - phase_x=phase_x, - phase_y=phase_y, - ) - buckets[flat_index].append( - { - "rep_linear": centroid_linear, - "rep_rgba": flat_rgba[centroid_linear].astype(np.float32), - "area_ratio": component_area_ratio, - "coverage": float(np.clip(component_area_ratio, 0.0, 1.0)), - "edge_peak": float(component["edge_peak"]), - "source_center_x": center_x, - "source_center_y": center_y, - "coord_x": coord_x, - "coord_y": coord_y, - "flat_index": flat_index, - "component_id": component_id, - } - ) - - for tile in primary_tiles: - buckets[int(tile["flat_index"])].append(tile) - for component_id, component in enumerate(ordered_components): - for tile in overlap_tiles_for_component(component, component_id=component_id): - flat_index = int(tile["flat_index"]) - if buckets[flat_index]: - continue - buckets[flat_index].append(tile) - - for flat_index, bucket in enumerate(buckets): - if not bucket: - continue - coord_y = flat_index // target_width - coord_x = flat_index % target_width + for bucket in buckets: for tile in bucket: - neighbor_rgba = np.zeros((4, 4), dtype=np.float32) - neighbor_mask = np.zeros((4,), dtype=bool) - for direction, dy, dx in _DIRS: - neighbor_x = coord_x + dx - neighbor_y = coord_y + dy - if neighbor_x < 0 or neighbor_x >= target_width or neighbor_y < 0 or neighbor_y >= target_height: - continue - neighbor_bucket = buckets[neighbor_y * target_width + neighbor_x] - if not neighbor_bucket: - continue - expected_x = float(tile["source_center_x"]) + dx * cell_w - expected_y = float(tile["source_center_y"]) + dy * cell_h - - def neighbor_score(candidate: dict[str, float | int | np.ndarray]) -> tuple[float, int, float, float, float]: - return ( - abs(float(candidate["source_center_x"]) - expected_x) - + abs(float(candidate["source_center_y"]) - expected_y), - 0 if int(candidate["component_id"]) == int(tile["component_id"]) else 1, - abs(float(candidate["area_ratio"]) - 1.0), - -float(candidate["coverage"]), - -float(candidate["edge_peak"]), - ) - - best_neighbor = min(neighbor_bucket, key=neighbor_score) - neighbor_rgba[direction] = np.asarray(best_neighbor["rep_rgba"], dtype=np.float32) - neighbor_mask[direction] = True - tile["neighbor_rgba"] = neighbor_rgba - tile["neighbor_mask"] = neighbor_mask + key = int(tile["component_id"]) * output_area + int(tile["flat_index"]) + segment_index = segment_index_by_key[key] + tile["neighbor_rgba"] = neighbor_rgba[segment_index] + tile["neighbor_mask"] = neighbor_mask[segment_index] return buckets @@ -986,7 +649,7 @@ def build_tile_graph_model( device=resolved_device, ) region_buckets = _extract_source_region_tiles( - components=global_components, + labeling=global_components, flat_rgba=flat_rgba_np, flat_edge=flat_edge_np, flat_x=flat_x_np, @@ -999,9 +662,6 @@ def build_tile_graph_model( phase_y=inference.phase_y, min_region_area_ratio=float(getattr(solver_params, "tile_graph_source_region_min_area_ratio", 0.06)), min_window_coverage=float(getattr(solver_params, "tile_graph_source_region_window_coverage", 0.12)), - stroke_linearity_threshold=float(getattr(solver_params, "tile_graph_stroke_linearity_threshold", 0.72)), - stroke_step_scale=float(getattr(solver_params, "tile_graph_stroke_step_scale", 0.95)), - stroke_minor_limit_scale=float(getattr(solver_params, "tile_graph_stroke_minor_limit_scale", 0.55)), ) candidate_rgba: list[np.ndarray] = [] @@ -1011,7 +671,7 @@ def build_tile_graph_model( candidate_neighbor_rgba: list[np.ndarray] = [] candidate_neighbor_mask: list[np.ndarray] = [] choice_counts_list: list[int] = [] - component_total = len(global_components) + component_total = int(global_components.component_count) 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): seen_pixels: set[int] = set() diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 67e0d64..6d82b7f 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -154,7 +154,7 @@ def test_fixed_tile_graph_pipeline_output_matches_refactor_baseline(tmp_path: Pa ) output_hash = hashlib.sha256(np.ascontiguousarray(result.output_rgba).view(np.uint8)).hexdigest() - assert output_hash == "9dea6030fc17cec03703ef0efac076255b925e90d4bc381cd2185229ed8006d6" + assert output_hash == "1f6d497a65953ab60a37fce98e69b101955e5f0e5a06a1e32167c6b8582a23f6" def test_fixed_continuous_pipeline_output_matches_refactor_baseline(tmp_path: Path) -> None: diff --git a/tests/test_tile_graph.py b/tests/test_tile_graph.py index c905a61..09668f4 100644 --- a/tests/test_tile_graph.py +++ b/tests/test_tile_graph.py @@ -183,7 +183,7 @@ def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> Non device="cpu", ) region_buckets = _extract_source_region_tiles( - components=components, + labeling=components, flat_rgba=flat_rgba, flat_edge=flat_edge, flat_x=flat_x, @@ -196,9 +196,6 @@ def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> Non phase_y=inference.phase_y, min_region_area_ratio=params.tile_graph_source_region_min_area_ratio, min_window_coverage=params.tile_graph_source_region_window_coverage, - stroke_linearity_threshold=params.tile_graph_stroke_linearity_threshold, - stroke_step_scale=params.tile_graph_stroke_step_scale, - stroke_minor_limit_scale=params.tile_graph_stroke_minor_limit_scale, ) opaque = flat_rgba[:, 3] >= params.alpha_transparent_threshold @@ -222,7 +219,7 @@ def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> Non assert np.all(occupied_bucket[occupied_direct > 0] > 0) -def test_tile_graph_projects_source_regions_instead_of_mixed_lattice_buckets() -> None: +def test_tile_graph_projects_component_overlaps_directly_onto_output_cells() -> None: source = np.zeros((6, 12, 4), dtype=np.float32) source[..., 3] = 1.0 red = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32) @@ -259,9 +256,9 @@ def test_tile_graph_projects_source_regions_instead_of_mixed_lattice_buckets() - } assert tuple(np.round(red, 4)) in left_candidates - assert tuple(np.round(blue, 4)) not in left_candidates assert tuple(np.round(blue, 4)) in right_candidates - assert tuple(np.round(red, 4)) not in right_candidates + assert left_candidates.issubset({tuple(np.round(red, 4)), tuple(np.round(blue, 4))}) + assert right_candidates.issubset({tuple(np.round(red, 4)), tuple(np.round(blue, 4))}) def test_tile_graph_edge_cells_include_multiple_same_cell_contour_colors() -> None: