diff --git a/README.md b/README.md index a760ce9..e6ade72 100644 --- a/README.md +++ b/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: diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index 7544058..df61124 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -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 diff --git a/src/repixelizer/analysis.py b/src/repixelizer/analysis.py index de9def7..61e61ad 100644 --- a/src/repixelizer/analysis.py +++ b/src/repixelizer/analysis.py @@ -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 diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index 800e584..c9014d1 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -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 diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index 471386e..d1e2f06 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -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, diff --git a/src/repixelizer/source_reference.py b/src/repixelizer/source_reference.py index 64c98af..fcc067b 100644 --- a/src/repixelizer/source_reference.py +++ b/src/repixelizer/source_reference.py @@ -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) diff --git a/src/repixelizer/tile_graph.py b/src/repixelizer/tile_graph.py index b2877e7..3534f6b 100644 --- a/src/repixelizer/tile_graph.py +++ b/src/repixelizer/tile_graph.py @@ -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, diff --git a/tests/test_analysis.py b/tests/test_analysis.py new file mode 100644 index 0000000..b206397 --- /dev/null +++ b/tests/test_analysis.py @@ -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 diff --git a/tests/test_source_reference.py b/tests/test_source_reference.py index 8068b9b..0e93e0f 100644 --- a/tests/test_source_reference.py +++ b/tests/test_source_reference.py @@ -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) diff --git a/tests/test_tile_graph.py b/tests/test_tile_graph.py index c756ad2..72ac334 100644 --- a/tests/test_tile_graph.py +++ b/tests/test_tile_graph.py @@ -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"