Refactor tile graph for full CUDA pipeline
This commit is contained in:
Meta Vens
parent
1b33050883
commit
124d31fa8a
10 changed files with 563 additions and 381 deletions
10
README.md
10
README.md
|
|
@ -38,14 +38,15 @@ Recent adjacency-focused status:
|
|||
|
||||
Current tile-graph status:
|
||||
|
||||
- it now builds literal source-pixel candidates from connected source clusters plus lattice-aligned fallback coverage
|
||||
- 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 local tile-graph solver now honors `--device` and runs through Torch, so its scoring/update loop can use CUDA instead of being hardwired to NumPy on the CPU
|
||||
- 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
|
||||
- accepted component candidates now consume a cell-sized source footprint, and each output cell only sees candidates attached to its own inferred lattice coord
|
||||
- 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
|
||||
- it already beats naive resize on the repo's synthetic thin-feature regression
|
||||
- this fixes the core design mismatch that had allowed repeated distant labels to create big same-color patches and opaque black background blocks
|
||||
- the remaining performance bottleneck is still candidate extraction/building on large real fixtures, so CUDA improves the solver stage but does not yet make badge-scale tile-graph runs fast end-to-end
|
||||
- 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
|
||||
|
||||
## Quickstart
|
||||
|
||||
|
|
@ -60,6 +61,7 @@ Run the optimizer:
|
|||
repixelize input.png --out output.png
|
||||
repixelize input.png --out output.png --diagnostics-dir diagnostics --device auto
|
||||
repixelize input.png --out output.png --reconstruction-mode tile-graph --diagnostics-dir diagnostics --device cpu
|
||||
repixelize input.png --out output.png --reconstruction-mode tile-graph --diagnostics-dir diagnostics --device cuda
|
||||
```
|
||||
|
||||
Run the optimizer plus baselines:
|
||||
|
|
|
|||
|
|
@ -8,11 +8,12 @@ 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
|
||||
- source-image structural analysis, now with an optional Torch/CUDA path
|
||||
- PyTorch-based continuous UV-field optimization
|
||||
- shared source-lattice reference built from actual inferred cell assignments
|
||||
- shared source-lattice reference built from actual inferred cell assignments, now with an optional Torch/CUDA path
|
||||
- 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
|
||||
|
|
@ -21,9 +22,10 @@ Implemented:
|
|||
|
||||
Verified:
|
||||
- local editable install in a dedicated venv
|
||||
- full test suite currently passing on the last completed pass
|
||||
- full test suite currently passing on the last completed pass (`54 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
|
||||
|
||||
## Status after adjacency-first pass
|
||||
|
||||
|
|
@ -57,28 +59,30 @@ The most provisional areas are:
|
|||
|
||||
### Active pass
|
||||
|
||||
This is the next solver-focused pass after the adjacency-first rewrite.
|
||||
This pass refactors the experimental tile-graph mode into a GPU-friendly codepath.
|
||||
|
||||
In scope now:
|
||||
- make source reliability edge-aware instead of dispersion-only so high-contrast thin features are not downweighted just because they are locally messy
|
||||
- expand snap/refine candidate generation beyond fixed offset grids by injecting sharp per-cell exemplars, edge peaks, and gradient-guided offsets
|
||||
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
|
||||
|
||||
Current implementation note:
|
||||
- raw image edges are the right signal for source-detail targeting; cluster boundaries were too broad and caused thin-feature washout when reused here
|
||||
- enriched edge candidates are currently gated to high-dispersion edge cells, and `snap` stays on the conservative local grid while `refine` gets the richer candidate pool
|
||||
- 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 extracts literal source-pixel candidates from connected source clusters, fills gaps with lattice-aligned fallback pixels, learns candidate adjacency, and places tiles with a soft propagation loop
|
||||
- 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
|
||||
- accepted component candidates now consume their cell-sized source footprint, and placement now uses per-cell local candidate sets instead of one reusable global candidate matrix
|
||||
- that fixes the core design mismatch behind the patchy repeated-color blocks and opaque-black background takeover: a candidate is now tied to its inferred output coord instead of being reusable across the whole raster
|
||||
- the local tile-graph solver now honors `device` and runs through Torch, so CUDA can accelerate the scoring/update loop
|
||||
- the remaining downside is performance in candidate extraction/building: large real-fixture runs are still noticeably slow because the tile-graph candidate pool scales with output area and is still assembled on the CPU
|
||||
- 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
|
||||
|
||||
Next after that:
|
||||
- rerank low-confidence top lattice candidates with short real solver probes instead of relying only on the cheapest preview
|
||||
- rerun tuning after the edge-aware/source-guided behavior lands
|
||||
- improve per-cell tile-graph proposal diversity so thin contours get more than “sharp vs edge-peak” in hard cells
|
||||
- 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
|
||||
|
||||
## High-value regression cases
|
||||
|
|
@ -117,8 +121,8 @@ What future work should improve here:
|
|||
- a retuning sweep after the edge-aware/source-guided changes land
|
||||
- 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 candidate coverage and stronger adjacency evidence are still needed even after switching candidates back to literal source pixels
|
||||
- for the experimental tile-graph path specifically, the next useful step is performance-aware extraction and propagation: keep the new per-cell candidate sets and CUDA-backed solver, but add pruning, caching, and cheaper component stepping so the badge-scale run is fast enough to iterate on
|
||||
- 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, 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
|
||||
|
|
|
|||
|
|
@ -6,6 +6,27 @@ from .metrics import luminance
|
|||
from .types import SourceAnalysis
|
||||
|
||||
|
||||
def _require_torch():
|
||||
try:
|
||||
import torch
|
||||
except ImportError as exc: # pragma: no cover - exercised only when torch missing
|
||||
raise RuntimeError(
|
||||
"PyTorch is required for GPU-accelerated source analysis. Install project dependencies first."
|
||||
) from exc
|
||||
return torch
|
||||
|
||||
|
||||
def _resolve_device(torch, requested: str) -> str:
|
||||
if requested == "auto":
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if requested == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
"CUDA was requested for source analysis, but this PyTorch build does not have a usable CUDA device. "
|
||||
"Install a CUDA-enabled PyTorch build or use --device cpu."
|
||||
)
|
||||
return requested
|
||||
|
||||
|
||||
def _compute_edge_map(rgba: np.ndarray) -> np.ndarray:
|
||||
lum = luminance(rgba)
|
||||
alpha = rgba[..., 3]
|
||||
|
|
@ -20,6 +41,20 @@ def _compute_edge_map(rgba: np.ndarray) -> np.ndarray:
|
|||
return edge.astype(np.float32)
|
||||
|
||||
|
||||
def _compute_edge_map_torch(torch, rgba_t):
|
||||
lum = rgba_t[..., 0] * 0.2126 + rgba_t[..., 1] * 0.7152 + rgba_t[..., 2] * 0.0722
|
||||
alpha = rgba_t[..., 3]
|
||||
dx = torch.zeros_like(lum, dtype=torch.float32)
|
||||
dy = torch.zeros_like(lum, dtype=torch.float32)
|
||||
dx[:, 1:] = (lum[:, 1:] - lum[:, :-1]).abs() + (alpha[:, 1:] - alpha[:, :-1]).abs()
|
||||
dy[1:, :] = (lum[1:, :] - lum[:-1, :]).abs() + (alpha[1:, :] - alpha[:-1, :]).abs()
|
||||
edge = torch.sqrt(dx * dx + dy * dy)
|
||||
max_edge = torch.max(edge)
|
||||
if float(max_edge.item()) > 0.0:
|
||||
edge = edge / max_edge
|
||||
return edge.to(dtype=torch.float32)
|
||||
|
||||
|
||||
def _kmeans(data: np.ndarray, k: int, seed: int, iterations: int = 12) -> tuple[np.ndarray, np.ndarray]:
|
||||
rng = np.random.default_rng(seed)
|
||||
if data.shape[0] <= k:
|
||||
|
|
@ -44,6 +79,31 @@ def _kmeans(data: np.ndarray, k: int, seed: int, iterations: int = 12) -> tuple[
|
|||
return centers.astype(np.float32), labels.astype(np.int32)
|
||||
|
||||
|
||||
def _kmeans_torch(torch, data_t, k: int, seed: int, iterations: int = 12):
|
||||
if data_t.shape[0] <= k:
|
||||
centers_t = data_t.clone()
|
||||
labels_t = torch.arange(data_t.shape[0], device=data_t.device, dtype=torch.long)
|
||||
return centers_t, labels_t
|
||||
generator = torch.Generator(device=data_t.device)
|
||||
generator.manual_seed(seed)
|
||||
picks = torch.randperm(data_t.shape[0], device=data_t.device, generator=generator)[:k]
|
||||
centers_t = data_t[picks].clone()
|
||||
for _ in range(iterations):
|
||||
distances = torch.linalg.vector_norm(data_t[:, None, :] - centers_t[None, :, :], dim=-1)
|
||||
labels_t = torch.argmin(distances, dim=1)
|
||||
new_centers = []
|
||||
for idx in range(k):
|
||||
members = data_t[labels_t == idx]
|
||||
if members.numel() == 0:
|
||||
new_centers.append(centers_t[idx])
|
||||
else:
|
||||
new_centers.append(members.mean(dim=0))
|
||||
centers_t = torch.stack(new_centers, dim=0).to(dtype=torch.float32)
|
||||
distances = torch.linalg.vector_norm(data_t[:, None, :] - centers_t[None, :, :], dim=-1)
|
||||
labels_t = torch.argmin(distances, dim=1)
|
||||
return centers_t, labels_t
|
||||
|
||||
|
||||
def _colorize_clusters(cluster_map: np.ndarray, centers: np.ndarray) -> np.ndarray:
|
||||
if centers.size == 0:
|
||||
return np.zeros((*cluster_map.shape, 4), dtype=np.float32)
|
||||
|
|
@ -53,7 +113,33 @@ def _colorize_clusters(cluster_map: np.ndarray, centers: np.ndarray) -> np.ndarr
|
|||
return np.concatenate([preview, alpha], axis=-1)
|
||||
|
||||
|
||||
def analyze_source(rgba: np.ndarray, seed: int, cluster_count: int = 6) -> SourceAnalysis:
|
||||
def analyze_source(rgba: np.ndarray, seed: int, cluster_count: int = 6, device: str | None = None) -> SourceAnalysis:
|
||||
if device is not None:
|
||||
torch = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
rgba_t = torch.from_numpy(rgba).to(device=resolved_device, dtype=torch.float32)
|
||||
edge_map_t = _compute_edge_map_torch(torch, rgba_t)
|
||||
alpha_t = rgba_t[..., 3]
|
||||
opaque_t = alpha_t > 0.05
|
||||
cluster_map_t = torch.full(alpha_t.shape, -1, device=resolved_device, dtype=torch.long)
|
||||
centers_t = torch.zeros((0, 4), device=resolved_device, dtype=torch.float32)
|
||||
if bool(torch.any(opaque_t).item()):
|
||||
samples_t = rgba_t[opaque_t]
|
||||
k = min(cluster_count, max(2, int(np.sqrt(int(samples_t.shape[0]) // 64 + 1))))
|
||||
centers_t, labels_t = _kmeans_torch(torch, samples_t, k=k, seed=seed)
|
||||
cluster_map_t[opaque_t] = labels_t
|
||||
preview_palette_t = centers_t[:, :3].clamp(0.0, 1.0) if centers_t.numel() else torch.zeros((1, 3), device=resolved_device, dtype=torch.float32)
|
||||
preview_indices_t = torch.clamp(cluster_map_t, 0, max(0, preview_palette_t.shape[0] - 1))
|
||||
preview_t = preview_palette_t[preview_indices_t]
|
||||
preview_t = torch.cat([preview_t, alpha_t[..., None]], dim=-1)
|
||||
return SourceAnalysis(
|
||||
edge_map=edge_map_t.detach().cpu().numpy().astype(np.float32),
|
||||
cluster_map=cluster_map_t.detach().cpu().numpy().astype(np.int32),
|
||||
cluster_centers=centers_t.detach().cpu().numpy().astype(np.float32),
|
||||
alpha_map=alpha_t.detach().cpu().numpy().astype(np.float32),
|
||||
cluster_preview=preview_t.detach().cpu().numpy().astype(np.float32),
|
||||
)
|
||||
|
||||
edge_map = _compute_edge_map(rgba)
|
||||
alpha = rgba[..., 3]
|
||||
opaque = alpha > 0.05
|
||||
|
|
|
|||
|
|
@ -78,6 +78,10 @@ class SolverHyperParams:
|
|||
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_area_weight: float = 0.03
|
||||
tile_graph_alpha_weight: float = 0.25
|
||||
tile_graph_coverage_weight: float = 0.05
|
||||
|
|
|
|||
|
|
@ -53,7 +53,7 @@ def run_pipeline(
|
|||
if strip_background:
|
||||
source = strip_edge_background(source)
|
||||
inference = infer_lattice(source, target_size=target_size, device=device)
|
||||
analysis = analyze_source(source, seed=seed)
|
||||
analysis = analyze_source(source, seed=seed, device=device if reconstruction_mode == "tile-graph" else None)
|
||||
inference = _select_phase_candidate(
|
||||
source,
|
||||
inference,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,27 @@ from .io import premultiply, unpremultiply
|
|||
from .types import SourceLatticeReference
|
||||
|
||||
|
||||
def _require_torch():
|
||||
try:
|
||||
import torch
|
||||
except ImportError as exc: # pragma: no cover - exercised only when torch missing
|
||||
raise RuntimeError(
|
||||
"PyTorch is required for GPU-accelerated source lattice references. Install project dependencies first."
|
||||
) from exc
|
||||
return torch
|
||||
|
||||
|
||||
def _resolve_device(torch, requested: str) -> str:
|
||||
if requested == "auto":
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if requested == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
"CUDA was requested for source lattice references, but this PyTorch build does not have a usable CUDA device. "
|
||||
"Install a CUDA-enabled PyTorch build or use --device cpu."
|
||||
)
|
||||
return requested
|
||||
|
||||
|
||||
def _default_edge_hint(source_rgba: np.ndarray) -> np.ndarray:
|
||||
rgb = source_rgba[..., :3]
|
||||
alpha = source_rgba[..., 3]
|
||||
|
|
@ -21,6 +42,21 @@ def _default_edge_hint(source_rgba: np.ndarray) -> np.ndarray:
|
|||
return edge.astype(np.float32)
|
||||
|
||||
|
||||
def _default_edge_hint_torch(torch, source_t):
|
||||
rgb = source_t[..., :3]
|
||||
alpha = source_t[..., 3]
|
||||
lum = rgb[..., 0] * 0.2126 + rgb[..., 1] * 0.7152 + rgb[..., 2] * 0.0722
|
||||
dx = torch.zeros_like(lum, dtype=torch.float32)
|
||||
dy = torch.zeros_like(lum, dtype=torch.float32)
|
||||
dx[:, 1:] = (lum[:, 1:] - lum[:, :-1]).abs() + (alpha[:, 1:] - alpha[:, :-1]).abs()
|
||||
dy[1:, :] = (lum[1:, :] - lum[:-1, :]).abs() + (alpha[1:, :] - alpha[:-1, :]).abs()
|
||||
edge = torch.sqrt(dx * dx + dy * dy)
|
||||
max_edge = torch.max(edge)
|
||||
if float(max_edge.item()) > 0.0:
|
||||
edge = edge / max_edge
|
||||
return edge.to(dtype=torch.float32)
|
||||
|
||||
|
||||
def lattice_indices(
|
||||
*,
|
||||
height: int,
|
||||
|
|
@ -39,6 +75,26 @@ def lattice_indices(
|
|||
return y_idx[:, None] * target_width + x_idx[None, :]
|
||||
|
||||
|
||||
def _lattice_indices_torch(
|
||||
torch,
|
||||
*,
|
||||
height: int,
|
||||
width: int,
|
||||
target_width: int,
|
||||
target_height: int,
|
||||
phase_x: float,
|
||||
phase_y: float,
|
||||
device: str,
|
||||
):
|
||||
cell_x = width / max(1, target_width)
|
||||
cell_y = height / max(1, target_height)
|
||||
xs = (torch.arange(width, device=device, dtype=torch.float32) + 0.5) / cell_x - phase_x
|
||||
ys = (torch.arange(height, device=device, dtype=torch.float32) + 0.5) / cell_y - phase_y
|
||||
x_idx = torch.floor(xs).to(dtype=torch.long).clamp(0, max(0, target_width - 1))
|
||||
y_idx = torch.floor(ys).to(dtype=torch.long).clamp(0, max(0, target_height - 1))
|
||||
return y_idx[:, None] * target_width + x_idx[None, :]
|
||||
|
||||
|
||||
def _reference_deltas(reference: np.ndarray) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]:
|
||||
height = reference.shape[0]
|
||||
width = reference.shape[1]
|
||||
|
|
@ -49,6 +105,35 @@ def _reference_deltas(reference: np.ndarray) -> tuple[np.ndarray | None, np.ndar
|
|||
return delta_x, delta_y, delta_diag, delta_anti
|
||||
|
||||
|
||||
def _unpremultiply_torch(torch, premul_t):
|
||||
alpha = premul_t[..., 3:4].clamp_min(1e-6)
|
||||
rgb = torch.where(alpha > 1e-6, premul_t[..., :3] / alpha, torch.zeros_like(premul_t[..., :3]))
|
||||
return torch.cat([rgb, premul_t[..., 3:4]], dim=-1).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _scatter_amax_1d(torch, indices_t, values_t, *, length: int):
|
||||
out_t = torch.full((length,), -float("inf"), device=values_t.device, dtype=torch.float32)
|
||||
out_t.scatter_reduce_(0, indices_t, values_t.to(dtype=torch.float32), reduce="amax", include_self=True)
|
||||
return torch.where(torch.isfinite(out_t), out_t, torch.zeros_like(out_t))
|
||||
|
||||
|
||||
def _argbest_linear_indices(torch, indices_t, scores_t, *, length: int, maximize: bool):
|
||||
if maximize:
|
||||
best_scores_t = torch.full((length,), -float("inf"), device=scores_t.device, dtype=torch.float32)
|
||||
best_scores_t.scatter_reduce_(0, indices_t, scores_t.to(dtype=torch.float32), reduce="amax", include_self=True)
|
||||
match_t = scores_t >= (best_scores_t[indices_t] - 1e-8)
|
||||
else:
|
||||
best_scores_t = torch.full((length,), float("inf"), device=scores_t.device, dtype=torch.float32)
|
||||
best_scores_t.scatter_reduce_(0, indices_t, scores_t.to(dtype=torch.float32), reduce="amin", include_self=True)
|
||||
match_t = scores_t <= (best_scores_t[indices_t] + 1e-8)
|
||||
linear_t = torch.arange(indices_t.shape[0], device=indices_t.device, dtype=torch.long)
|
||||
sentinel = torch.full_like(linear_t, indices_t.shape[0])
|
||||
masked_linear_t = torch.where(match_t, linear_t, sentinel)
|
||||
best_linear_t = torch.full((length,), indices_t.shape[0], device=indices_t.device, dtype=torch.long)
|
||||
best_linear_t.scatter_reduce_(0, indices_t, masked_linear_t, reduce="amin", include_self=True)
|
||||
return best_linear_t
|
||||
|
||||
|
||||
def build_source_lattice_reference(
|
||||
source_rgba: np.ndarray,
|
||||
*,
|
||||
|
|
@ -60,7 +145,146 @@ def build_source_lattice_reference(
|
|||
edge_hint: np.ndarray | None = None,
|
||||
edge_grad_x_hint: np.ndarray | None = None,
|
||||
edge_grad_y_hint: np.ndarray | None = None,
|
||||
device: str | None = None,
|
||||
) -> SourceLatticeReference:
|
||||
if device is not None:
|
||||
torch = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
source_t = torch.from_numpy(source_rgba).to(device=resolved_device, dtype=torch.float32)
|
||||
height = source_rgba.shape[0]
|
||||
width = source_rgba.shape[1]
|
||||
edge_hint_t = (
|
||||
torch.from_numpy(edge_hint.astype(np.float32)).to(device=resolved_device, dtype=torch.float32)
|
||||
if edge_hint is not None
|
||||
else _default_edge_hint_torch(torch, source_t)
|
||||
)
|
||||
indices_t = _lattice_indices_torch(
|
||||
torch,
|
||||
height=height,
|
||||
width=width,
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
phase_x=phase_x,
|
||||
phase_y=phase_y,
|
||||
device=resolved_device,
|
||||
)
|
||||
cell_count = max(1, target_width * target_height)
|
||||
premul_t = source_t.clone()
|
||||
premul_t[..., :3] = premul_t[..., :3] * premul_t[..., 3:4]
|
||||
flat_idx_t = indices_t.reshape(-1).to(dtype=torch.long)
|
||||
flat_premul_t = premul_t.reshape(-1, premul_t.shape[-1])
|
||||
flat_alpha_t = source_t.reshape(-1, source_t.shape[-1])[:, 3]
|
||||
flat_edge_t = edge_hint_t.reshape(-1)
|
||||
|
||||
counts_t = torch.bincount(flat_idx_t, minlength=cell_count).to(dtype=torch.float32)
|
||||
safe_counts_t = counts_t.clamp_min(1.0)
|
||||
channel_sums_t = torch.stack(
|
||||
[torch.bincount(flat_idx_t, weights=flat_premul_t[:, channel], minlength=cell_count) for channel in range(flat_premul_t.shape[-1])],
|
||||
dim=-1,
|
||||
).to(dtype=torch.float32)
|
||||
mean_premul_flat_t = channel_sums_t / safe_counts_t[:, None]
|
||||
mean_premul_t = mean_premul_flat_t.reshape(target_height, target_width, premul_t.shape[-1])
|
||||
mean_rgba_t = _unpremultiply_torch(torch, mean_premul_t)
|
||||
|
||||
per_pixel_mean_t = mean_premul_flat_t[flat_idx_t]
|
||||
pixel_diff_t = (flat_premul_t - per_pixel_mean_t).abs().mean(dim=-1)
|
||||
cell_diff_sums_t = torch.bincount(flat_idx_t, weights=pixel_diff_t, minlength=cell_count).to(dtype=torch.float32)
|
||||
cell_dispersion_t = (cell_diff_sums_t / safe_counts_t).reshape(target_height, target_width)
|
||||
|
||||
lattice_alpha_t = mean_rgba_t.reshape(-1, mean_rgba_t.shape[-1])[flat_idx_t, 3]
|
||||
support_mask_t = torch.maximum(flat_alpha_t, lattice_alpha_t) >= alpha_threshold
|
||||
if bool(torch.any(support_mask_t).item()):
|
||||
dispersion = float(pixel_diff_t[support_mask_t].mean().item())
|
||||
else:
|
||||
dispersion = float(pixel_diff_t.mean().item()) if pixel_diff_t.numel() else 0.0
|
||||
|
||||
cell_x = width / max(1, target_width)
|
||||
cell_y = height / max(1, target_height)
|
||||
default_x_t = torch.clamp(
|
||||
torch.round((torch.arange(target_width, device=resolved_device, dtype=torch.float32) + 0.5 + phase_x) * cell_x - 0.5),
|
||||
0,
|
||||
max(0, width - 1),
|
||||
).to(dtype=torch.long)
|
||||
default_y_t = torch.clamp(
|
||||
torch.round((torch.arange(target_height, device=resolved_device, dtype=torch.float32) + 0.5 + phase_y) * cell_y - 0.5),
|
||||
0,
|
||||
max(0, height - 1),
|
||||
).to(dtype=torch.long)
|
||||
sharp_x_flat_t = default_x_t[None, :].expand(target_height, target_width).reshape(-1).clone()
|
||||
sharp_y_flat_t = default_y_t[:, None].expand(target_height, target_width).reshape(-1).clone()
|
||||
|
||||
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_x_t = pixel_x_t.reshape(-1)
|
||||
flat_y_t = pixel_y_t.reshape(-1)
|
||||
exemplar_cost_t = pixel_diff_t + torch.clamp(per_pixel_mean_t[:, 3] - flat_premul_t[:, 3], min=0.0) * 1e-3
|
||||
|
||||
edge_score_t = flat_edge_t + flat_alpha_t.to(dtype=torch.float32) * 0.05
|
||||
edge_peak_x_flat_t = sharp_x_flat_t.clone()
|
||||
edge_peak_y_flat_t = sharp_y_flat_t.clone()
|
||||
edge_strength_flat_t = _scatter_amax_1d(torch, flat_idx_t, flat_edge_t, length=cell_count)
|
||||
edge_grad_x_flat_t = torch.zeros(cell_count, device=resolved_device, dtype=torch.float32)
|
||||
edge_grad_y_flat_t = torch.zeros(cell_count, device=resolved_device, dtype=torch.float32)
|
||||
best_edge_linear_t = _argbest_linear_indices(torch, flat_idx_t, edge_score_t, length=cell_count, maximize=True)
|
||||
valid_edge_mask_t = best_edge_linear_t < flat_idx_t.shape[0]
|
||||
valid_edge_cells_t = torch.nonzero(valid_edge_mask_t, as_tuple=False).reshape(-1)
|
||||
if valid_edge_cells_t.numel() > 0:
|
||||
valid_edge_linear_t = best_edge_linear_t[valid_edge_cells_t]
|
||||
edge_peak_x_flat_t[valid_edge_cells_t] = flat_x_t[valid_edge_linear_t]
|
||||
edge_peak_y_flat_t[valid_edge_cells_t] = flat_y_t[valid_edge_linear_t]
|
||||
if edge_grad_x_hint is not None or edge_grad_y_hint is not None:
|
||||
edge_peak_index_flat_t = edge_peak_y_flat_t * width + edge_peak_x_flat_t
|
||||
if edge_grad_x_hint is not None:
|
||||
edge_grad_x_hint_t = torch.from_numpy(edge_grad_x_hint.astype(np.float32)).to(device=resolved_device, dtype=torch.float32).reshape(-1)
|
||||
edge_grad_x_flat_t = edge_grad_x_hint_t[edge_peak_index_flat_t]
|
||||
if edge_grad_y_hint is not None:
|
||||
edge_grad_y_hint_t = torch.from_numpy(edge_grad_y_hint.astype(np.float32)).to(device=resolved_device, dtype=torch.float32).reshape(-1)
|
||||
edge_grad_y_flat_t = edge_grad_y_hint_t[edge_peak_index_flat_t]
|
||||
|
||||
sharp_premul_flat_t = mean_premul_flat_t.clone()
|
||||
best_exemplar_linear_t = _argbest_linear_indices(torch, flat_idx_t, exemplar_cost_t, length=cell_count, maximize=False)
|
||||
valid_exemplar_mask_t = best_exemplar_linear_t < flat_idx_t.shape[0]
|
||||
valid_exemplar_cells_t = torch.nonzero(valid_exemplar_mask_t, as_tuple=False).reshape(-1)
|
||||
if valid_exemplar_cells_t.numel() > 0:
|
||||
valid_exemplar_linear_t = best_exemplar_linear_t[valid_exemplar_cells_t]
|
||||
sharp_premul_flat_t[valid_exemplar_cells_t] = flat_premul_t[valid_exemplar_linear_t]
|
||||
sharp_x_flat_t[valid_exemplar_cells_t] = flat_x_t[valid_exemplar_linear_t]
|
||||
sharp_y_flat_t[valid_exemplar_cells_t] = flat_y_t[valid_exemplar_linear_t]
|
||||
sharp_premul_t = sharp_premul_flat_t.reshape(target_height, target_width, premul_t.shape[-1])
|
||||
sharp_rgba_t = _unpremultiply_torch(torch, sharp_premul_t)
|
||||
|
||||
cell_alpha_max_t = torch.zeros(cell_count, device=resolved_device, dtype=torch.float32)
|
||||
cell_alpha_max_t.scatter_reduce_(0, flat_idx_t, flat_alpha_t, reduce="amax", include_self=True)
|
||||
expected_pixels_per_cell = float(source_rgba.shape[0] * source_rgba.shape[1]) / float(cell_count)
|
||||
cell_support_t = (counts_t / max(expected_pixels_per_cell, 1.0)).reshape(target_height, target_width)
|
||||
|
||||
sharp_premul_np = sharp_premul_t.detach().cpu().numpy().astype(np.float32)
|
||||
delta_x, delta_y, delta_diag, delta_anti = _reference_deltas(sharp_premul_np)
|
||||
return SourceLatticeReference(
|
||||
mean_rgba=mean_rgba_t.detach().cpu().numpy().astype(np.float32),
|
||||
sharp_rgba=sharp_rgba_t.detach().cpu().numpy().astype(np.float32),
|
||||
dispersion=dispersion,
|
||||
lattice_indices=indices_t.detach().cpu().numpy().astype(np.int32),
|
||||
cell_dispersion=cell_dispersion_t.detach().cpu().numpy().astype(np.float32),
|
||||
cell_counts=counts_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32),
|
||||
cell_support=cell_support_t.detach().cpu().numpy().astype(np.float32),
|
||||
cell_alpha_max=cell_alpha_max_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32),
|
||||
sharp_x=sharp_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32),
|
||||
sharp_y=sharp_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32),
|
||||
edge_peak_x=edge_peak_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32),
|
||||
edge_peak_y=edge_peak_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.int32),
|
||||
edge_strength=edge_strength_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32),
|
||||
edge_grad_x=edge_grad_x_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32),
|
||||
edge_grad_y=edge_grad_y_flat_t.reshape(target_height, target_width).detach().cpu().numpy().astype(np.float32),
|
||||
delta_x=delta_x.astype(np.float32) if delta_x is not None else None,
|
||||
delta_y=delta_y.astype(np.float32) if delta_y is not None else None,
|
||||
delta_diag=delta_diag.astype(np.float32) if delta_diag is not None else None,
|
||||
delta_anti=delta_anti.astype(np.float32) if delta_anti is not None else None,
|
||||
)
|
||||
|
||||
height = source_rgba.shape[0]
|
||||
width = source_rgba.shape[1]
|
||||
edge_hint = edge_hint.astype(np.float32) if edge_hint is not None else _default_edge_hint(source_rgba)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -63,19 +62,7 @@ class TileGraphModel:
|
|||
component_count: int
|
||||
edge_density: float
|
||||
average_choices: float
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _Component:
|
||||
label: int
|
||||
coords: np.ndarray
|
||||
bbox: tuple[int, int, int, int]
|
||||
mask: np.ndarray
|
||||
available_mask: np.ndarray
|
||||
available_integral: np.ndarray
|
||||
area: int
|
||||
centroid_y: float
|
||||
centroid_x: float
|
||||
model_device: str
|
||||
|
||||
|
||||
def _make_regular_uv(height: int, width: int, target_height: int, target_width: int, phase_x: float, phase_y: float) -> np.ndarray:
|
||||
|
|
@ -92,210 +79,28 @@ def _make_regular_uv(height: int, width: int, target_height: int, target_width:
|
|||
return uv.astype(np.float32)
|
||||
|
||||
|
||||
def _component_integral(mask: np.ndarray) -> np.ndarray:
|
||||
integral = np.zeros((mask.shape[0] + 1, mask.shape[1] + 1), dtype=np.int32)
|
||||
integral[1:, 1:] = np.cumsum(np.cumsum(mask.astype(np.int32), axis=0), axis=1)
|
||||
return integral
|
||||
|
||||
|
||||
def _integral_sum(integral: np.ndarray, y0: int, x0: int, y1: int, x1: int) -> int:
|
||||
return int(integral[y1, x1] - integral[y0, x1] - integral[y1, x0] + integral[y0, x0])
|
||||
|
||||
|
||||
def _window_bounds(center_y: float, center_x: float, cell_h: float, cell_w: float, height: int, width: int) -> tuple[int, int, int, int]:
|
||||
y0 = int(np.floor(center_y - cell_h * 0.5))
|
||||
x0 = int(np.floor(center_x - cell_w * 0.5))
|
||||
y1 = int(np.ceil(center_y + cell_h * 0.5))
|
||||
x1 = int(np.ceil(center_x + cell_w * 0.5))
|
||||
y0 = max(0, min(height - 1, y0))
|
||||
x0 = max(0, min(width - 1, x0))
|
||||
y1 = max(y0 + 1, min(height, y1))
|
||||
x1 = max(x0 + 1, min(width, x1))
|
||||
return y0, x0, y1, x1
|
||||
|
||||
|
||||
def _source_to_cell_coord(
|
||||
center_y: float,
|
||||
center_x: float,
|
||||
*,
|
||||
cell_h: float,
|
||||
cell_w: float,
|
||||
target_height: int,
|
||||
target_width: int,
|
||||
phase_y: float,
|
||||
phase_x: float,
|
||||
) -> tuple[int, int]:
|
||||
x_idx = int(np.floor((center_x + 0.5) / max(cell_w, 1e-4) - phase_x))
|
||||
y_idx = int(np.floor((center_y + 0.5) / max(cell_h, 1e-4) - phase_y))
|
||||
x_idx = max(0, min(target_width - 1, x_idx))
|
||||
y_idx = max(0, min(target_height - 1, y_idx))
|
||||
return y_idx, x_idx
|
||||
|
||||
|
||||
def _sample_source_pixel(source_rgba: np.ndarray, center_y: float, center_x: float) -> np.ndarray:
|
||||
y = int(np.clip(np.rint(center_y), 0, max(0, source_rgba.shape[0] - 1)))
|
||||
x = int(np.clip(np.rint(center_x), 0, max(0, source_rgba.shape[1] - 1)))
|
||||
return source_rgba[y, x].astype(np.float32)
|
||||
|
||||
|
||||
def _window_component_coverage(component: _Component, center_y: float, center_x: float, cell_h: float, cell_w: float) -> float:
|
||||
y0, x0, y1, x1 = _window_bounds(
|
||||
center_y,
|
||||
center_x,
|
||||
cell_h,
|
||||
cell_w,
|
||||
component.available_mask.shape[0] + component.bbox[0],
|
||||
component.available_mask.shape[1] + component.bbox[1],
|
||||
)
|
||||
bbox_y0, bbox_x0, _, _ = component.bbox
|
||||
local_y0 = max(0, y0 - bbox_y0)
|
||||
local_x0 = max(0, x0 - bbox_x0)
|
||||
local_y1 = min(component.available_mask.shape[0], y1 - bbox_y0)
|
||||
local_x1 = min(component.available_mask.shape[1], x1 - bbox_x0)
|
||||
if local_y1 <= local_y0 or local_x1 <= local_x0:
|
||||
return 0.0
|
||||
covered = _integral_sum(component.available_integral, local_y0, local_x0, local_y1, local_x1)
|
||||
total = max(1, (local_y1 - local_y0) * (local_x1 - local_x0))
|
||||
return float(covered / total)
|
||||
|
||||
|
||||
def _nearest_component_pixel(component: _Component, center_y: float, center_x: float) -> tuple[float, float]:
|
||||
bbox_y0, bbox_x0, _, _ = component.bbox
|
||||
local = component.coords - np.asarray([bbox_y0, bbox_x0], dtype=np.int32)[None, :]
|
||||
available = component.available_mask[local[:, 0], local[:, 1]]
|
||||
coords = component.coords[available]
|
||||
if coords.shape[0] == 0:
|
||||
return float(center_y), float(center_x)
|
||||
deltas = coords.astype(np.float32) - np.asarray([center_y, center_x], dtype=np.float32)[None, :]
|
||||
best = int(np.argmin(np.sum(deltas * deltas, axis=1)))
|
||||
return float(coords[best, 0]), float(coords[best, 1])
|
||||
|
||||
|
||||
def _consume_component_window(component: _Component, center_y: float, center_x: float, cell_h: float, cell_w: float) -> None:
|
||||
y0, x0, y1, x1 = _window_bounds(
|
||||
center_y,
|
||||
center_x,
|
||||
cell_h,
|
||||
cell_w,
|
||||
component.available_mask.shape[0] + component.bbox[0],
|
||||
component.available_mask.shape[1] + component.bbox[1],
|
||||
)
|
||||
bbox_y0, bbox_x0, _, _ = component.bbox
|
||||
local_y0 = max(0, y0 - bbox_y0)
|
||||
local_x0 = max(0, x0 - bbox_x0)
|
||||
local_y1 = min(component.available_mask.shape[0], y1 - bbox_y0)
|
||||
local_x1 = min(component.available_mask.shape[1], x1 - bbox_x0)
|
||||
if local_y1 <= local_y0 or local_x1 <= local_x0:
|
||||
return
|
||||
component.available_mask[local_y0:local_y1, local_x0:local_x1] = False
|
||||
component.available_integral = _component_integral(component.available_mask)
|
||||
|
||||
|
||||
def _best_stepped_center(
|
||||
component: _Component,
|
||||
center_y: float,
|
||||
center_x: float,
|
||||
*,
|
||||
step_y: float,
|
||||
step_x: float,
|
||||
cell_h: float,
|
||||
cell_w: float,
|
||||
) -> tuple[float, float, float] | None:
|
||||
offsets_y = np.asarray([0.0, -step_y * 0.35, step_y * 0.35], dtype=np.float32)
|
||||
offsets_x = np.asarray([0.0, -step_x * 0.35, step_x * 0.35], dtype=np.float32)
|
||||
best: tuple[float, float, float] | None = None
|
||||
for dy in offsets_y:
|
||||
for dx in offsets_x:
|
||||
probe_y, probe_x = _nearest_component_pixel(component, center_y + float(dy), center_x + float(dx))
|
||||
coverage = _window_component_coverage(component, probe_y, probe_x, cell_h, cell_w)
|
||||
if best is None or coverage > best[2]:
|
||||
best = (probe_y, probe_x, coverage)
|
||||
return best
|
||||
|
||||
|
||||
def _directional_pixel(
|
||||
source_rgba: np.ndarray,
|
||||
component: _Component,
|
||||
center_y: float,
|
||||
center_x: float,
|
||||
*,
|
||||
dy: int,
|
||||
dx: int,
|
||||
step_y: float,
|
||||
step_x: float,
|
||||
cell_h: float,
|
||||
cell_w: float,
|
||||
coverage_threshold: float,
|
||||
) -> np.ndarray:
|
||||
best = _best_stepped_center(
|
||||
component,
|
||||
center_y + dy * step_y,
|
||||
center_x + dx * step_x,
|
||||
step_y=step_y,
|
||||
step_x=step_x,
|
||||
cell_h=cell_h,
|
||||
cell_w=cell_w,
|
||||
)
|
||||
if best is not None:
|
||||
probe_y, probe_x, coverage = best
|
||||
if coverage >= coverage_threshold:
|
||||
return _sample_source_pixel(source_rgba, probe_y, probe_x)
|
||||
return _sample_source_pixel(source_rgba, center_y + dy * step_y, center_x + dx * step_x)
|
||||
|
||||
|
||||
def _extract_components(cluster_map: np.ndarray, alpha_map: np.ndarray, alpha_threshold: float) -> list[_Component]:
|
||||
valid = alpha_map >= alpha_threshold
|
||||
visited = np.zeros(valid.shape, dtype=bool)
|
||||
components: list[_Component] = []
|
||||
height, width = valid.shape
|
||||
for y in range(height):
|
||||
for x in range(width):
|
||||
if not valid[y, x] or visited[y, x]:
|
||||
continue
|
||||
label = int(cluster_map[y, x])
|
||||
queue = deque([(y, x)])
|
||||
visited[y, x] = True
|
||||
coords: list[tuple[int, int]] = []
|
||||
min_y = max_y = y
|
||||
min_x = max_x = x
|
||||
while queue:
|
||||
cy, cx = queue.popleft()
|
||||
coords.append((cy, cx))
|
||||
if cy < min_y:
|
||||
min_y = cy
|
||||
if cy > max_y:
|
||||
max_y = cy
|
||||
if cx < min_x:
|
||||
min_x = cx
|
||||
if cx > max_x:
|
||||
max_x = cx
|
||||
for dy, dx in ((0, 1), (1, 0), (0, -1), (-1, 0)):
|
||||
ny = cy + dy
|
||||
nx = cx + dx
|
||||
if ny < 0 or ny >= height or nx < 0 or nx >= width:
|
||||
continue
|
||||
if visited[ny, nx] or not valid[ny, nx] or int(cluster_map[ny, nx]) != label:
|
||||
continue
|
||||
visited[ny, nx] = True
|
||||
queue.append((ny, nx))
|
||||
coords_array = np.asarray(coords, dtype=np.int32)
|
||||
mask = np.zeros((max_y - min_y + 1, max_x - min_x + 1), dtype=bool)
|
||||
local = coords_array - np.asarray([min_y, min_x], dtype=np.int32)[None, :]
|
||||
mask[local[:, 0], local[:, 1]] = True
|
||||
components.append(
|
||||
_Component(
|
||||
label=label,
|
||||
coords=coords_array,
|
||||
bbox=(min_y, min_x, max_y + 1, max_x + 1),
|
||||
mask=mask,
|
||||
available_mask=mask.copy(),
|
||||
available_integral=_component_integral(mask),
|
||||
area=int(coords_array.shape[0]),
|
||||
centroid_y=float(np.mean(coords_array[:, 0])),
|
||||
centroid_x=float(np.mean(coords_array[:, 1])),
|
||||
)
|
||||
)
|
||||
return components
|
||||
def _build_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 build_tile_graph_model(
|
||||
|
|
@ -304,7 +109,10 @@ def build_tile_graph_model(
|
|||
inference: InferenceResult,
|
||||
analysis: SourceAnalysis,
|
||||
solver_params: SolverHyperParams | None = None,
|
||||
device: str = "cpu",
|
||||
) -> TileGraphModel:
|
||||
torch = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
solver_params = solver_params or SolverHyperParams()
|
||||
source_reference = build_source_lattice_reference(
|
||||
source_rgba,
|
||||
|
|
@ -314,153 +122,153 @@ def build_tile_graph_model(
|
|||
phase_y=inference.phase_y,
|
||||
alpha_threshold=solver_params.alpha_transparent_threshold,
|
||||
edge_hint=analysis.edge_map,
|
||||
device=resolved_device,
|
||||
)
|
||||
height, width = source_rgba.shape[:2]
|
||||
cell_h = height / max(1, inference.target_height)
|
||||
cell_w = width / max(1, inference.target_width)
|
||||
target_area = cell_h * cell_w
|
||||
output_area = max(1, inference.target_height * inference.target_width)
|
||||
extracted_candidate_limit = max(
|
||||
int(solver_params.tile_graph_max_candidates),
|
||||
output_area * max(1, int(solver_params.tile_graph_max_candidates_per_coord)),
|
||||
)
|
||||
components = _extract_components(analysis.cluster_map, analysis.alpha_map, solver_params.alpha_transparent_threshold)
|
||||
components.sort(key=lambda component: (abs(component.area / max(target_area, 1e-4) - 1.0), -component.area))
|
||||
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)))
|
||||
|
||||
candidate_rgba: list[np.ndarray] = []
|
||||
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)
|
||||
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)
|
||||
edge_peak_x_flat_t = torch.from_numpy(source_reference.edge_peak_x.reshape(-1)).to(device=resolved_device, dtype=torch.long)
|
||||
edge_peak_y_flat_t = torch.from_numpy(source_reference.edge_peak_y.reshape(-1)).to(device=resolved_device, dtype=torch.long)
|
||||
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]
|
||||
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)
|
||||
color_error_t = (flat_rgba_t - sharp_rgba_flat_t[lattice_indices_t]).abs().mean(dim=-1)
|
||||
edge_cell_boost_t = (
|
||||
(1.0 - edge_focus_dist_t.clamp(0.0, 1.0))
|
||||
* (edge_strength_flat_t[lattice_indices_t] >= solver_params.source_edge_detail_threshold).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)
|
||||
- color_error_t * 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,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
sharp_linear_np = sharp_linear_t.detach().cpu().numpy().astype(np.int64)
|
||||
edge_linear_np = edge_linear_t.detach().cpu().numpy().astype(np.int64)
|
||||
ranked_pool_np = ranked_pool_t.detach().cpu().numpy().astype(np.int64)
|
||||
pixel_coverage_np = pixel_coverage_t.detach().cpu().numpy().astype(np.float32)
|
||||
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] = []
|
||||
candidate_deltas: list[np.ndarray] = []
|
||||
|
||||
extracted_coord_counts: dict[tuple[int, int], int] = {}
|
||||
step_y = max(1.0, cell_h)
|
||||
step_x = max(1.0, cell_w)
|
||||
for component in components:
|
||||
if not np.any(component.available_mask):
|
||||
continue
|
||||
area_ratio = float(component.area / max(target_area, 1e-4))
|
||||
if area_ratio < solver_params.tile_graph_component_min_area_ratio:
|
||||
continue
|
||||
per_component_limit = max(1, int(np.ceil(max(1.0, area_ratio))))
|
||||
seed_y, seed_x = _nearest_component_pixel(component, component.centroid_y, component.centroid_x)
|
||||
queue = deque([(seed_y, seed_x)])
|
||||
local_seen: set[tuple[int, int]] = set()
|
||||
added = 0
|
||||
while queue and added < per_component_limit and len(candidate_rgba) < extracted_candidate_limit:
|
||||
center_y, center_x = queue.popleft()
|
||||
coord = _source_to_cell_coord(
|
||||
center_y,
|
||||
center_x,
|
||||
cell_h=cell_h,
|
||||
cell_w=cell_w,
|
||||
target_height=inference.target_height,
|
||||
target_width=inference.target_width,
|
||||
phase_y=inference.phase_y,
|
||||
phase_x=inference.phase_x,
|
||||
)
|
||||
if coord in local_seen:
|
||||
continue
|
||||
if extracted_coord_counts.get(coord, 0) >= solver_params.tile_graph_max_candidates_per_coord:
|
||||
continue
|
||||
local_seen.add(coord)
|
||||
coverage = _window_component_coverage(component, center_y, center_x, cell_h, cell_w)
|
||||
if coverage < solver_params.tile_graph_window_coverage_threshold and added > 0:
|
||||
continue
|
||||
center_rgba = _sample_source_pixel(source_rgba, center_y, center_x)
|
||||
delta_features = np.zeros((4, 4), dtype=np.float32)
|
||||
for direction, dy, dx in _DIRS:
|
||||
neighbor_rgba = _directional_pixel(
|
||||
source_rgba,
|
||||
component,
|
||||
center_y,
|
||||
center_x,
|
||||
dy=dy,
|
||||
dx=dx,
|
||||
step_y=step_y,
|
||||
step_x=step_x,
|
||||
cell_h=cell_h,
|
||||
cell_w=cell_w,
|
||||
coverage_threshold=solver_params.tile_graph_window_coverage_threshold,
|
||||
)
|
||||
delta_features[direction] = neighbor_rgba - center_rgba
|
||||
candidate_rgba.append(center_rgba)
|
||||
candidate_coords.append(coord)
|
||||
candidate_area_ratio.append(area_ratio)
|
||||
candidate_coverage.append(coverage)
|
||||
candidate_deltas.append(delta_features)
|
||||
extracted_coord_counts[coord] = extracted_coord_counts.get(coord, 0) + 1
|
||||
_consume_component_window(component, center_y, center_x, cell_h, cell_w)
|
||||
added += 1
|
||||
if area_ratio < solver_params.tile_graph_large_component_ratio:
|
||||
continue
|
||||
for _, dy, dx in _DIRS:
|
||||
best = _best_stepped_center(
|
||||
component,
|
||||
center_y + dy * step_y,
|
||||
center_x + dx * step_x,
|
||||
step_y=step_y,
|
||||
step_x=step_x,
|
||||
cell_h=cell_h,
|
||||
cell_w=cell_w,
|
||||
)
|
||||
if best is None:
|
||||
continue
|
||||
probe_y, probe_x, probe_coverage = best
|
||||
if probe_coverage < solver_params.tile_graph_window_coverage_threshold:
|
||||
continue
|
||||
probe_coord = _source_to_cell_coord(
|
||||
probe_y,
|
||||
probe_x,
|
||||
cell_h=cell_h,
|
||||
cell_w=cell_w,
|
||||
target_height=inference.target_height,
|
||||
target_width=inference.target_width,
|
||||
phase_y=inference.phase_y,
|
||||
phase_x=inference.phase_x,
|
||||
)
|
||||
if (
|
||||
probe_coord not in local_seen
|
||||
and extracted_coord_counts.get(probe_coord, 0) < solver_params.tile_graph_max_candidates_per_coord
|
||||
):
|
||||
queue.append((probe_y, probe_x))
|
||||
|
||||
def add_reference_candidate(coord_y: int, coord_x: int) -> None:
|
||||
coord = (coord_y, coord_x)
|
||||
delta_features = np.zeros((4, 4), dtype=np.float32)
|
||||
center_rgba = source_reference.sharp_rgba[coord_y, coord_x].astype(np.float32)
|
||||
for direction, dy, dx in _DIRS:
|
||||
neighbor_y = min(max(coord_y + dy, 0), inference.target_height - 1)
|
||||
neighbor_x = min(max(coord_x + dx, 0), inference.target_width - 1)
|
||||
delta_features[direction] = source_reference.sharp_rgba[neighbor_y, neighbor_x] - center_rgba
|
||||
candidate_rgba.append(center_rgba)
|
||||
candidate_coords.append(coord)
|
||||
candidate_area_ratio.append(1.0)
|
||||
candidate_coverage.append(1.0)
|
||||
candidate_deltas.append(delta_features)
|
||||
|
||||
for coord_y in range(inference.target_height):
|
||||
for coord_x in range(inference.target_width):
|
||||
add_reference_candidate(coord_y, coord_x)
|
||||
|
||||
candidate_rgba_np = np.asarray(candidate_rgba, dtype=np.float32)
|
||||
candidate_coords_np = np.asarray(candidate_coords, dtype=np.int32)
|
||||
candidate_area_ratio_np = np.asarray(candidate_area_ratio, dtype=np.float32)
|
||||
candidate_coverage_np = np.asarray(candidate_coverage, dtype=np.float32)
|
||||
candidate_deltas_np = np.asarray(candidate_deltas, dtype=np.float32)
|
||||
|
||||
coord_to_indices: dict[tuple[int, int], list[int]] = {}
|
||||
for index, coord in enumerate(candidate_coords_np.tolist()):
|
||||
coord_to_indices.setdefault((int(coord[0]), int(coord[1])), []).append(index)
|
||||
cell_candidate_offsets = np.zeros(output_area + 1, dtype=np.int32)
|
||||
cell_candidate_indices: list[int] = []
|
||||
choice_counts_list: list[int] = []
|
||||
for flat_index in range(output_area):
|
||||
coord_y = flat_index // inference.target_width
|
||||
coord_x = flat_index % inference.target_width
|
||||
indices = coord_to_indices.get((coord_y, coord_x), [])
|
||||
cell_candidate_indices.extend(indices)
|
||||
cell_candidate_offsets[flat_index + 1] = len(cell_candidate_indices)
|
||||
cell_candidate_indices_np = np.asarray(cell_candidate_indices, dtype=np.int32)
|
||||
seen_pixels: set[int] = set()
|
||||
cell_start = len(candidate_linear)
|
||||
proposal_pixels = [int(sharp_linear_np[flat_index])]
|
||||
if edge_strength_np[flat_index] >= solver_params.source_edge_detail_threshold:
|
||||
proposal_pixels.append(int(edge_linear_np[flat_index]))
|
||||
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 edge_strength_np[flat_index] < solver_params.source_edge_detail_threshold
|
||||
):
|
||||
continue
|
||||
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) >= max_candidates_per_coord:
|
||||
break
|
||||
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])))
|
||||
choice_counts_list.append(len(candidate_linear) - cell_start)
|
||||
|
||||
candidate_linear_t = torch.as_tensor(candidate_linear, device=resolved_device, dtype=torch.long)
|
||||
center_rgba_t = flat_rgba_t[candidate_linear_t]
|
||||
center_y_t = flat_y_t[candidate_linear_t].to(dtype=torch.float32)
|
||||
center_x_t = flat_x_t[candidate_linear_t].to(dtype=torch.float32)
|
||||
candidate_deltas_t = torch.zeros((candidate_linear_t.shape[0], 4, 4), device=resolved_device, dtype=torch.float32)
|
||||
for direction, dy, dx in _DIRS:
|
||||
sample_y_t = torch.round(center_y_t + dy * cell_h).clamp(0, max(0, height - 1)).to(dtype=torch.long)
|
||||
sample_x_t = torch.round(center_x_t + dx * cell_w).clamp(0, max(0, width - 1)).to(dtype=torch.long)
|
||||
candidate_deltas_t[:, direction] = source_t[sample_y_t, sample_x_t] - center_rgba_t
|
||||
|
||||
candidate_rgba_np = center_rgba_t.detach().cpu().numpy().astype(np.float32)
|
||||
candidate_coords_np = np.asarray(candidate_coords, dtype=np.int32)
|
||||
candidate_area_ratio_np = np.asarray(candidate_area_ratio, dtype=np.float32)
|
||||
candidate_coverage_np = np.asarray(candidate_coverage, dtype=np.float32)
|
||||
candidate_deltas_np = candidate_deltas_t.detach().cpu().numpy().astype(np.float32)
|
||||
|
||||
cell_candidate_offsets = np.zeros(output_area + 1, dtype=np.int32)
|
||||
cell_candidate_offsets[1:] = np.cumsum(np.asarray(choice_counts_list, dtype=np.int32), dtype=np.int32)
|
||||
cell_candidate_indices_np = np.arange(candidate_rgba_np.shape[0], dtype=np.int32)
|
||||
choice_counts = np.diff(cell_candidate_offsets)
|
||||
average_choices = float(np.mean(choice_counts)) if choice_counts.size else 0.0
|
||||
|
||||
|
|
@ -495,9 +303,10 @@ 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),
|
||||
component_count=len(components),
|
||||
component_count=int(np.unique(analysis.cluster_map[analysis.alpha_map >= solver_params.alpha_transparent_threshold]).size),
|
||||
edge_density=edge_density,
|
||||
average_choices=average_choices,
|
||||
model_device=resolved_device,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -769,6 +578,7 @@ def optimize_tile_graph(
|
|||
inference=inference,
|
||||
analysis=analysis,
|
||||
solver_params=solver_params,
|
||||
device=resolved_device,
|
||||
)
|
||||
choice_indices_np, choice_mask_np = _build_choice_grid(model)
|
||||
choice_indices_t = torch.from_numpy(choice_indices_np).to(device=resolved_device, dtype=torch.long)
|
||||
|
|
@ -858,7 +668,9 @@ def optimize_tile_graph(
|
|||
target_rgba = _assignment_rgba(model, selected).astype(np.float32)
|
||||
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_component_count": model.component_count,
|
||||
"tile_graph_candidate_count": int(model.candidate_rgba.shape[0]),
|
||||
"tile_graph_edge_density": model.edge_density,
|
||||
|
|
|
|||
20
tests/test_analysis.py
Normal file
20
tests/test_analysis.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from repixelizer.analysis import analyze_source
|
||||
|
||||
|
||||
def test_analyze_source_cpu_device_matches_default_path() -> None:
|
||||
source = np.zeros((8, 8, 4), dtype=np.float32)
|
||||
source[..., 3] = 1.0
|
||||
source[2:6, 3:5, :3] = np.asarray([1.0, 0.8, 0.2], dtype=np.float32)
|
||||
|
||||
default = analyze_source(source, seed=7)
|
||||
accelerated = analyze_source(source, seed=7, device="cpu")
|
||||
|
||||
assert np.allclose(accelerated.edge_map, default.edge_map, atol=1e-5)
|
||||
assert accelerated.cluster_centers.shape == default.cluster_centers.shape
|
||||
assert np.allclose(accelerated.alpha_map, default.alpha_map, atol=1e-5)
|
||||
assert accelerated.cluster_map.shape == default.cluster_map.shape
|
||||
assert accelerated.cluster_preview.shape == default.cluster_preview.shape
|
||||
|
|
@ -78,3 +78,31 @@ def test_source_lattice_reference_computes_default_edge_metadata() -> None:
|
|||
|
||||
assert np.all(reference.edge_strength[:, 0] > 0.0)
|
||||
assert np.all(reference.edge_peak_x[:, 0] == 2)
|
||||
|
||||
|
||||
def test_source_lattice_reference_cpu_device_matches_default_path() -> None:
|
||||
source = np.zeros((8, 8, 4), dtype=np.float32)
|
||||
source[..., 3] = 1.0
|
||||
source[2:6, 4, :3] = np.asarray([1.0, 1.0, 1.0], dtype=np.float32)
|
||||
|
||||
default = build_source_lattice_reference(
|
||||
source,
|
||||
target_width=4,
|
||||
target_height=4,
|
||||
phase_x=0.0,
|
||||
phase_y=0.0,
|
||||
)
|
||||
accelerated = build_source_lattice_reference(
|
||||
source,
|
||||
target_width=4,
|
||||
target_height=4,
|
||||
phase_x=0.0,
|
||||
phase_y=0.0,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
assert np.allclose(accelerated.mean_rgba, default.mean_rgba, atol=1e-5)
|
||||
assert np.allclose(accelerated.sharp_rgba, default.sharp_rgba, atol=1e-5)
|
||||
assert np.allclose(accelerated.edge_strength, default.edge_strength, atol=1e-5)
|
||||
assert np.array_equal(accelerated.sharp_x, default.sharp_x)
|
||||
assert np.array_equal(accelerated.sharp_y, default.sharp_y)
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ from repixelizer.tile_graph import build_tile_graph_model, optimize_tile_graph
|
|||
from repixelizer.types import InferenceResult
|
||||
|
||||
|
||||
def test_tile_graph_model_extracts_candidates_and_adjacency_from_component_walk() -> None:
|
||||
def test_tile_graph_model_extracts_candidates_and_adjacency_from_lattice_proposals() -> None:
|
||||
lowres = np.zeros((4, 4, 4), dtype=np.float32)
|
||||
lowres[..., 3] = 1.0
|
||||
lowres[:, 1] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32)
|
||||
|
|
@ -30,11 +30,7 @@ def test_tile_graph_model_extracts_candidates_and_adjacency_from_component_walk(
|
|||
top_candidates=[],
|
||||
)
|
||||
|
||||
model = build_tile_graph_model(
|
||||
source,
|
||||
inference=inference,
|
||||
analysis=analyze_source(source, seed=7),
|
||||
)
|
||||
model = build_tile_graph_model(source, inference=inference, analysis=analyze_source(source, seed=7), device="cpu")
|
||||
|
||||
assert model.candidate_rgba.shape[0] > 4
|
||||
choice_counts = np.diff(model.cell_candidate_offsets)
|
||||
|
|
@ -42,6 +38,7 @@ def test_tile_graph_model_extracts_candidates_and_adjacency_from_component_walk(
|
|||
assert np.all(choice_counts >= 1)
|
||||
assert model.average_choices >= 1.0
|
||||
assert model.edge_density > 0.0
|
||||
assert model.model_device == "cpu"
|
||||
|
||||
|
||||
def test_tile_graph_candidates_use_literal_source_pixel_colors() -> None:
|
||||
|
|
@ -214,12 +211,13 @@ def test_tile_graph_cuda_matches_cpu_on_small_case() -> None:
|
|||
confidence=1.0,
|
||||
top_candidates=[],
|
||||
)
|
||||
analysis = analyze_source(fake, seed=7)
|
||||
cpu_analysis = analyze_source(fake, seed=7, device="cpu")
|
||||
cuda_analysis = analyze_source(fake, seed=7, device="cuda")
|
||||
|
||||
cpu_artifacts, cpu_diag = optimize_tile_graph(
|
||||
fake,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
analysis=cpu_analysis,
|
||||
steps=0,
|
||||
seed=7,
|
||||
device="cpu",
|
||||
|
|
@ -227,13 +225,15 @@ def test_tile_graph_cuda_matches_cpu_on_small_case() -> None:
|
|||
cuda_artifacts, cuda_diag = optimize_tile_graph(
|
||||
fake,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
analysis=cuda_analysis,
|
||||
steps=0,
|
||||
seed=7,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
assert cpu_diag["tile_graph_model_device"] == "cpu"
|
||||
assert cpu_diag["tile_graph_solver_device"] == "cpu"
|
||||
assert cuda_diag["tile_graph_model_device"] == "cuda"
|
||||
assert cuda_diag["tile_graph_solver_device"] == "cuda"
|
||||
assert np.allclose(cpu_artifacts.target_rgba, cuda_artifacts.target_rgba, atol=1e-5)
|
||||
|
||||
|
|
@ -252,6 +252,7 @@ def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(tmp_path: Pa
|
|||
diagnostics_dir=diagnostics_dir,
|
||||
steps=24,
|
||||
reconstruction_mode="tile-graph",
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
import json
|
||||
|
|
@ -260,6 +261,7 @@ def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(tmp_path: Pa
|
|||
assert output_path.exists()
|
||||
assert run_json["settings"]["reconstruction_mode"] == "tile-graph"
|
||||
assert run_json["reconstruction"]["mode"] == "tile-graph"
|
||||
assert run_json["reconstruction"]["tile_graph_model_device"] == "cpu"
|
||||
assert run_json["reconstruction"]["tile_graph_candidate_count"] > 0
|
||||
assert set(run_json["source_fidelity"].keys()) == {"snap_initial", "solver_target", "final_output"}
|
||||
assert result.diagnostics["reconstruction"]["mode"] == "tile-graph"
|
||||
|
|
|
|||
Loading…
Reference in a new issue