From 0747c1a42a84d8122e1dd1c0b844f47371afd13b Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 23 Apr 2026 23:24:28 +0100 Subject: [PATCH] Add experimental phase-field optimizer --- README.md | 3 + docs/implementation-plan.md | 28 +- docs/lean-optimizer-algorithm-map.md | 13 + src/repixelizer/cli.py | 2 +- src/repixelizer/params.py | 10 + src/repixelizer/phase_field.py | 439 +++++++++++++++++++++++++++ src/repixelizer/pipeline.py | 18 ++ tests/test_phase_field.py | 74 +++++ tests/test_pipeline.py | 85 ++++++ 9 files changed, 667 insertions(+), 5 deletions(-) create mode 100644 src/repixelizer/phase_field.py create mode 100644 tests/test_phase_field.py diff --git a/README.md b/README.md index df873c8..91e59e4 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,7 @@ Current tile-graph status: - the optimizer now emits stage displacement diagnostics too, so `run.json` can show what snap, relax handoff, relax mode, and final refine actually did to the per-cell source-pixel displacement field instead of just reporting a single fidelity score - those new diagnostics also settled the design argument about relax: on the pinned badge check under `artifacts/optimizer-relax-purpose-check/`, relax improves source-fidelity a bit (`0.07549` without relax to `0.07451` with it), but it does not act like a strong broad-swatch phase smoother; the displacement field stays roughly as jittery, so the current relax stage is better described as a soft consensus / handoff stabilizer over fixed local candidates - that diagnosis finally forced the obvious architectural verdict: the current optimizer is still a tray-based discrete chooser in a fake moustache, so the repo now also carries `docs/lean-optimizer-algorithm-map.md`, a replacement target for a real displacement-field optimizer with one field, one objective, and one final sample +- the first version of that replacement now exists behind `--reconstruction-mode phase-field`; it is a deliberately lean displacement-field solver in `src/repixelizer/phase_field.py`, and the first pinned badge smoke under `artifacts/phase-field-v1-badge-126/` lands at `0.1461`, which is not competitive yet but at least comes from a machine that is no longer lying about what it is - after the first pruning pass, tile-graph no longer participates in pipeline phase-rerank probes and no longer carries hybrid geometry priors through its unary cost; the path is now one lattice-conditioned candidate generator plus one local discrete solver - this fixes the core design mismatch that had allowed repeated distant labels to create big same-color patches and opaque black background blocks - on the current `24x24` emblem smoke case, an end-to-end `tile-graph` run dropped from about `2.57s` on CPU to `0.61s` on CUDA on this machine @@ -83,6 +84,7 @@ Run the optimizer: ```powershell 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 phase-field --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 repixelize input.png --out output.png --reconstruction-mode tile-graph --target-width 126 --target-height 126 --phase-x 0.0 --phase-y -0.2 --device cuda @@ -154,6 +156,7 @@ Core modules: - `inference.py`: target size and phase inference - `continuous.py`: source-aware lattice snapping and local discrete refinement +- `phase_field.py`: experimental lean displacement-field optimizer - `metrics.py`: fidelity, adjacency, and motif metrics - `benchmark.py`: corpus benchmark runner - `tuning.py`: black-box hyperparameter search harness diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index cfb4979..8e238cf 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -204,10 +204,30 @@ Why this is the right next cut: Next implementation steps: -1. create a brand-new optimizer module beside `continuous.py` -2. wire it up behind an experimental reconstruction mode or internal switch -3. start with only the five objective terms listed above -4. compare it against pinned badge and emblem runs before porting any old heuristics +1. compare the new `phase-field` mode against pinned badge and emblem runs before porting any old heuristics +2. improve the local evidence term so the field is pulled harder into solid fake-pixel interiors +3. tune the projection / spacing story so the field can drift farther without collapsing or freezing +4. only then consider importing a small number of old heuristics if they clearly help + +What has already landed: + +- a brand-new `src/repixelizer/phase_field.py` module +- an experimental `--reconstruction-mode phase-field` pipeline path +- a direct displacement-field optimizer using only: + - local coherence + - local edge avoidance + - edge-aware smoothness + - anti-collapse spacing + - displacement magnitude regularization +- pinned-badge smoke artifacts under `artifacts/phase-field-v1-badge-126/` + +First honest result: + +- on the cleaned badge at fixed `126x126`, phase `(0.0, -0.2)`, `steps=48`, CUDA +- the new `phase-field` mode lands at `0.1461` final source-fidelity +- that is much worse than the current continuous path (`0.0745` on the same pinned case), but dramatically better than the earlier broken tile-graph collapse and, more importantly, it is being produced by a machine that actually matches its own description + +That is a respectable first brick: not good enough, but finally the right shape. ## High-value regression case diff --git a/docs/lean-optimizer-algorithm-map.md b/docs/lean-optimizer-algorithm-map.md index 423aa8a..a5b71fc 100644 --- a/docs/lean-optimizer-algorithm-map.md +++ b/docs/lean-optimizer-algorithm-map.md @@ -14,6 +14,19 @@ It is a map of the optimizer we should replace it with: That is the whole dream. No watercolor portrait. No accountant's ledger. No snap religion and refine religion living in the same cathedral. Just one field, one objective, one final sample. +The first implementation of this restart now exists in `src/repixelizer/phase_field.py` behind `--reconstruction-mode phase-field`. + +That implementation is intentionally small and slightly stupid: + +- one displacement vector per output cell +- one local solidity term +- one edge-aware smoothness term +- one anti-collapse spacing term +- one displacement magnitude prior +- one final nearest source sample + +It is not yet beating the old continuous optimizer on the pinned badge stress case, but it is finally the right kind of machine to improve. + ## One-sentence machine The lean optimizer should be: diff --git a/src/repixelizer/cli.py b/src/repixelizer/cli.py index 70fe631..32ac118 100644 --- a/src/repixelizer/cli.py +++ b/src/repixelizer/cli.py @@ -36,7 +36,7 @@ def build_parser() -> argparse.ArgumentParser: target.add_argument( "--reconstruction-mode", default="continuous", - choices=("continuous", "tile-graph"), + choices=("continuous", "phase-field", "tile-graph"), help="Reconstruction engine to run after lattice estimation", ) target.add_argument( diff --git a/src/repixelizer/params.py b/src/repixelizer/params.py index ba86143..6e171df 100644 --- a/src/repixelizer/params.py +++ b/src/repixelizer/params.py @@ -5,6 +5,16 @@ from dataclasses import asdict, dataclass @dataclass(slots=True) class SolverHyperParams: + phase_field_patch_extent: float = 0.18 + phase_field_data_coherence_weight: float = 1.00 + phase_field_data_edge_weight: float = 0.35 + phase_field_smoothness_weight: float = 0.28 + phase_field_edge_gate_strength: float = 6.0 + phase_field_collapse_weight: float = 1.20 + phase_field_min_spacing_ratio: float = 0.18 + phase_field_magnitude_weight: float = 0.08 + phase_field_learning_rate: float = 0.10 + phase_field_max_displacement_ratio: float = 0.48 representative_softmax_scale: float = 18.0 source_edge_reliability_gain: float = 1.6 source_edge_reliability_floor: float = 0.25 diff --git a/src/repixelizer/phase_field.py b/src/repixelizer/phase_field.py new file mode 100644 index 0000000..058476f --- /dev/null +++ b/src/repixelizer/phase_field.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .io import premultiply +from .types import ContinuousSourceAnalysis, InferenceResult, SolverArtifacts +from .params import SolverHyperParams + + +@dataclass(slots=True) +class _PhaseFieldPrep: + source_t: object + edge_t: object + uv0_px_t: object + uv0_norm: np.ndarray + base_x_t: object + base_y_t: object + guidance: np.ndarray + cell_x: float + cell_y: float + patch_offsets_t: object + width: int + height: int + + +def _require_torch(): + try: + import torch + import torch.nn.functional as F + except ImportError as exc: # pragma: no cover - exercised only when torch missing + raise RuntimeError( + "PyTorch is required for the phase-field optimization stage. Install project dependencies first." + ) from exc + return torch, F + + +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, but this PyTorch build does not have a usable CUDA device. " + "Install a CUDA-enabled PyTorch build or use --device cpu." + ) + return requested + + +def _make_regular_uv_px( + *, + height: int, + width: int, + target_height: int, + target_width: int, + phase_x: float, + phase_y: float, +) -> np.ndarray: + cell_x = width / target_width + cell_y = height / target_height + xs = (np.arange(target_width, dtype=np.float32) + 0.5 + phase_x) * cell_x - 0.5 + ys = (np.arange(target_height, dtype=np.float32) + 0.5 + phase_y) * cell_y - 0.5 + xs = np.clip(xs, 0.0, max(0.0, width - 1)) + ys = np.clip(ys, 0.0, max(0.0, height - 1)) + grid_x, grid_y = np.meshgrid(xs, ys) + return np.stack([grid_x, grid_y], axis=-1).astype(np.float32) + + +def _pixel_to_normalized(coords_px, *, width: int, height: int): + coords = coords_px.clone() + coords[..., 0] = (coords[..., 0] / max(1.0, float(width - 1))) * 2.0 - 1.0 + coords[..., 1] = (coords[..., 1] / max(1.0, float(height - 1))) * 2.0 - 1.0 + return coords + + +def _sample_rgba(F, source_t, coords_px, *, width: int, height: int): + grid = _pixel_to_normalized(coords_px, width=width, height=height).clamp(-1.0, 1.0) + sampled = F.grid_sample(source_t, grid, align_corners=True, mode="bilinear", padding_mode="border") + return sampled.permute(0, 2, 3, 1) + + +def _sample_scalar(F, scalar_t, coords_px, *, width: int, height: int): + grid = _pixel_to_normalized(coords_px, width=width, height=height).clamp(-1.0, 1.0) + sampled = F.grid_sample(scalar_t, grid, align_corners=True, mode="bilinear", padding_mode="border") + return sampled[:, 0] + + +def _sample_patch_rgba(F, source_t, coords_px, offsets_t, *, width: int, height: int): + patch_px = coords_px[:, :, :, None, :] + offsets_t[None, None, None, :, :] + patch_px = patch_px.new_empty((*patch_px.shape[:-1], 2)) + patch_px[..., 0] = (coords_px[:, :, :, None, 0] + offsets_t[None, None, None, :, 0]).clamp( + 0.0, max(0.0, float(width - 1)) + ) + patch_px[..., 1] = (coords_px[:, :, :, None, 1] + offsets_t[None, None, None, :, 1]).clamp( + 0.0, max(0.0, float(height - 1)) + ) + patch_grid = _pixel_to_normalized(patch_px, width=width, height=height).clamp(-1.0, 1.0) + batch, out_h, out_w, samples, _ = patch_grid.shape + flattened = patch_grid.reshape(batch, out_h, out_w * samples, 2) + sampled = F.grid_sample(source_t, flattened, align_corners=True, mode="bilinear", padding_mode="border") + sampled = sampled.permute(0, 2, 3, 1) + return sampled.reshape(batch, out_h, out_w, samples, sampled.shape[-1]) + + +def _sample_patch_scalar(F, scalar_t, coords_px, offsets_t, *, width: int, height: int): + patch_px = coords_px[:, :, :, None, :] + offsets_t[None, None, None, :, :] + patch_px = patch_px.new_empty((*patch_px.shape[:-1], 2)) + patch_px[..., 0] = (coords_px[:, :, :, None, 0] + offsets_t[None, None, None, :, 0]).clamp( + 0.0, max(0.0, float(width - 1)) + ) + patch_px[..., 1] = (coords_px[:, :, :, None, 1] + offsets_t[None, None, None, :, 1]).clamp( + 0.0, max(0.0, float(height - 1)) + ) + patch_grid = _pixel_to_normalized(patch_px, width=width, height=height).clamp(-1.0, 1.0) + batch, out_h, out_w, samples, _ = patch_grid.shape + flattened = patch_grid.reshape(batch, out_h, out_w * samples, 2) + sampled = F.grid_sample(scalar_t, flattened, align_corners=True, mode="bilinear", padding_mode="border") + sampled = sampled[:, 0] + return sampled.reshape(batch, out_h, out_w, samples) + + +def _displacement_diagnostics(uv_field: np.ndarray, selected_x, selected_y, *, width: int, height: int) -> dict[str, np.ndarray | float]: + uv_x = (uv_field[..., 0] + 1.0) * 0.5 * max(1.0, float(width - 1)) + uv_y = (uv_field[..., 1] + 1.0) * 0.5 * max(1.0, float(height - 1)) + displacement_x = selected_x.astype(np.float32) - uv_x.astype(np.float32) + displacement_y = selected_y.astype(np.float32) - uv_y.astype(np.float32) + magnitude = np.sqrt(np.square(displacement_x) + np.square(displacement_y)).astype(np.float32) + + orthogonal_jitter_terms: list[np.ndarray] = [] + local_residual_terms: list[np.ndarray] = [] + for field in (displacement_x, displacement_y): + if field.shape[1] > 1: + orthogonal_jitter_terms.append(np.abs(field[:, 1:] - field[:, :-1])) + if field.shape[0] > 1: + orthogonal_jitter_terms.append(np.abs(field[1:, :] - field[:-1, :])) + blurred = field.copy() + if field.shape[1] > 2: + blurred[:, 1:-1] = (field[:, :-2] + field[:, 1:-1] + field[:, 2:]) / 3.0 + if field.shape[0] > 2: + blurred[1:-1, :] = (blurred[:-2, :] + blurred[1:-1, :] + blurred[2:, :]) / 3.0 + local_residual_terms.append(np.abs(field - blurred)) + + orthogonal_jitter = ( + float(np.mean(np.concatenate([term.reshape(-1) for term in orthogonal_jitter_terms]))) + if orthogonal_jitter_terms + else 0.0 + ) + local_residual = ( + float(np.mean(np.concatenate([term.reshape(-1) for term in local_residual_terms]))) + if local_residual_terms + else 0.0 + ) + rounded_offsets = np.stack([np.rint(displacement_x), np.rint(displacement_y)], axis=-1).reshape(-1, 2) + if rounded_offsets.size == 0: + dominant_offset_ratio = 1.0 + else: + _, counts = np.unique(rounded_offsets, axis=0, return_counts=True) + dominant_offset_ratio = float(np.max(counts) / max(1, rounded_offsets.shape[0])) + return { + "displacement_x": displacement_x, + "displacement_y": displacement_y, + "mean_magnitude_px": float(np.mean(magnitude)), + "orthogonal_jitter_px": orthogonal_jitter, + "local_residual_px": local_residual, + "dominant_offset_ratio": dominant_offset_ratio, + } + + +def _project_displacements_in_place(torch, disp_t, base_x_t, base_y_t, *, min_dx: float, min_dy: float, max_dx: float, max_dy: float, width: int, height: int) -> None: + with torch.no_grad(): + pos_x = (base_x_t + disp_t[..., 0]).clamp(0.0, max(0.0, float(width - 1))) + pos_y = (base_y_t + disp_t[..., 1]).clamp(0.0, max(0.0, float(height - 1))) + + for index in range(1, pos_x.shape[2]): + pos_x[:, :, index] = torch.maximum(pos_x[:, :, index], pos_x[:, :, index - 1] + min_dx) + for index in range(pos_x.shape[2] - 2, -1, -1): + pos_x[:, :, index] = torch.minimum(pos_x[:, :, index], pos_x[:, :, index + 1] - min_dx) + for index in range(1, pos_y.shape[1]): + pos_y[:, index, :] = torch.maximum(pos_y[:, index, :], pos_y[:, index - 1, :] + min_dy) + for index in range(pos_y.shape[1] - 2, -1, -1): + pos_y[:, index, :] = torch.minimum(pos_y[:, index, :], pos_y[:, index + 1, :] - min_dy) + + pos_x = pos_x.clamp(0.0, max(0.0, float(width - 1))) + pos_y = pos_y.clamp(0.0, max(0.0, float(height - 1))) + disp_t[..., 0].copy_((pos_x - base_x_t).clamp(-max_dx, max_dx)) + disp_t[..., 1].copy_((pos_y - base_y_t).clamp(-max_dy, max_dy)) + + +def _prepare_phase_field( + torch, + F, + rgba: np.ndarray, + inference: InferenceResult, + analysis: ContinuousSourceAnalysis, + solver_params: SolverHyperParams, + *, + device: str, +) -> _PhaseFieldPrep: + source = premultiply(rgba) + height, width = source.shape[:2] + cell_x = width / max(1, inference.target_width) + cell_y = height / max(1, inference.target_height) + source_t = torch.from_numpy(source.transpose(2, 0, 1)[None, ...]).to(device=device, dtype=torch.float32) + edge_t = torch.from_numpy(analysis.edge_map[None, None, ...]).to(device=device, dtype=torch.float32) + uv0_px = _make_regular_uv_px( + height=height, + width=width, + target_height=inference.target_height, + target_width=inference.target_width, + phase_x=inference.phase_x, + phase_y=inference.phase_y, + ) + uv0_px_t = torch.from_numpy(uv0_px[None, ...]).to(device=device, dtype=torch.float32) + base_x_t = uv0_px_t[..., 0] + base_y_t = uv0_px_t[..., 1] + uv0_norm = uv0_px.copy() + uv0_norm[..., 0] = (uv0_norm[..., 0] / max(1.0, float(width - 1))) * 2.0 - 1.0 + uv0_norm[..., 1] = (uv0_norm[..., 1] / max(1.0, float(height - 1))) * 2.0 - 1.0 + offsets = np.asarray( + [ + [0.0, 0.0], + [-solver_params.phase_field_patch_extent * cell_x, 0.0], + [solver_params.phase_field_patch_extent * cell_x, 0.0], + [0.0, -solver_params.phase_field_patch_extent * cell_y], + [0.0, solver_params.phase_field_patch_extent * cell_y], + [-solver_params.phase_field_patch_extent * cell_x, -solver_params.phase_field_patch_extent * cell_y], + [solver_params.phase_field_patch_extent * cell_x, -solver_params.phase_field_patch_extent * cell_y], + [-solver_params.phase_field_patch_extent * cell_x, solver_params.phase_field_patch_extent * cell_y], + [solver_params.phase_field_patch_extent * cell_x, solver_params.phase_field_patch_extent * cell_y], + ], + dtype=np.float32, + ) + patch_offsets_t = torch.from_numpy(offsets).to(device=device, dtype=torch.float32) + guide_small = F.interpolate( + edge_t, + size=(inference.target_height, inference.target_width), + mode="bilinear", + align_corners=True, + ) + return _PhaseFieldPrep( + source_t=source_t, + edge_t=edge_t, + uv0_px_t=uv0_px_t, + uv0_norm=uv0_norm, + base_x_t=base_x_t, + base_y_t=base_y_t, + guidance=guide_small[0, 0].detach().cpu().numpy().astype(np.float32), + cell_x=cell_x, + cell_y=cell_y, + patch_offsets_t=patch_offsets_t, + width=width, + height=height, + ) + + +def _phase_field_loss(torch, F, prep: _PhaseFieldPrep, disp_t, solver_params: SolverHyperParams): + pos_px = prep.uv0_px_t + disp_t + sampled_rgba = _sample_rgba(F, prep.source_t, pos_px, width=prep.width, height=prep.height) + patch_rgba = _sample_patch_rgba( + F, + prep.source_t, + pos_px, + prep.patch_offsets_t, + width=prep.width, + height=prep.height, + ) + patch_edge = _sample_patch_scalar( + F, + prep.edge_t, + pos_px, + prep.patch_offsets_t, + width=prep.width, + height=prep.height, + ) + + center_rgba = patch_rgba[..., 0, :] + neighbor_rgba = patch_rgba[..., 1:, :] + local_coherence = (neighbor_rgba - center_rgba[..., None, :]).abs().mean() + local_edge = patch_edge.mean() + + disp_norm = disp_t.clone() + disp_norm[..., 0] = disp_norm[..., 0] / max(prep.cell_x, 1e-4) + disp_norm[..., 1] = disp_norm[..., 1] / max(prep.cell_y, 1e-4) + pos_mid_x = (pos_px[:, :, 1:, :] + pos_px[:, :, :-1, :]) * 0.5 if pos_px.shape[2] > 1 else None + pos_mid_y = (pos_px[:, 1:, :, :] + pos_px[:, :-1, :, :]) * 0.5 if pos_px.shape[1] > 1 else None + + smoothness = sampled_rgba.new_tensor(0.0) + if pos_mid_x is not None: + edge_x = _sample_scalar(F, prep.edge_t, pos_mid_x, width=prep.width, height=prep.height) + weight_x = torch.exp(-solver_params.phase_field_edge_gate_strength * edge_x) + delta_x = disp_norm[:, :, 1:, :] - disp_norm[:, :, :-1, :] + smoothness = smoothness + (weight_x * torch.sqrt(delta_x.square().sum(dim=-1) + 1e-6)).mean() + if pos_mid_y is not None: + edge_y = _sample_scalar(F, prep.edge_t, pos_mid_y, width=prep.width, height=prep.height) + weight_y = torch.exp(-solver_params.phase_field_edge_gate_strength * edge_y) + delta_y = disp_norm[:, 1:, :, :] - disp_norm[:, :-1, :, :] + smoothness = smoothness + (weight_y * torch.sqrt(delta_y.square().sum(dim=-1) + 1e-6)).mean() + + collapse = sampled_rgba.new_tensor(0.0) + min_dx = solver_params.phase_field_min_spacing_ratio * prep.cell_x + min_dy = solver_params.phase_field_min_spacing_ratio * prep.cell_y + if pos_px.shape[2] > 1: + step_x = pos_px[:, :, 1:, 0] - pos_px[:, :, :-1, 0] + collapse = collapse + torch.relu(min_dx - step_x).square().mean() + if pos_px.shape[1] > 1: + step_y = pos_px[:, 1:, :, 1] - pos_px[:, :-1, :, 1] + collapse = collapse + torch.relu(min_dy - step_y).square().mean() + + magnitude = ( + (disp_t[..., 0] / max(prep.cell_x, 1e-4)).square() + + (disp_t[..., 1] / max(prep.cell_y, 1e-4)).square() + ).mean() + + loss = ( + local_coherence * solver_params.phase_field_data_coherence_weight + + local_edge * solver_params.phase_field_data_edge_weight + + smoothness * solver_params.phase_field_smoothness_weight + + collapse * solver_params.phase_field_collapse_weight + + magnitude * solver_params.phase_field_magnitude_weight + ) + terms = { + "local_coherence": float(local_coherence.detach().cpu().item()), + "local_edge": float(local_edge.detach().cpu().item()), + "smoothness": float(smoothness.detach().cpu().item()), + "collapse": float(collapse.detach().cpu().item()), + "magnitude": float(magnitude.detach().cpu().item()), + } + return loss, sampled_rgba, terms + + +def _nearest_source_rgba(source_rgba: np.ndarray, sample_x: np.ndarray, sample_y: np.ndarray) -> np.ndarray: + return source_rgba[sample_y, sample_x] + + +def optimize_phase_field( + rgba: np.ndarray, + inference: InferenceResult, + analysis: ContinuousSourceAnalysis, + steps: int, + seed: int, + device: str, + solver_params: SolverHyperParams | None = None, +) -> SolverArtifacts: + torch, F = _require_torch() + device = _resolve_device(torch, device) + solver_params = solver_params or SolverHyperParams() + torch.manual_seed(seed) + if device == "cuda": + torch.cuda.manual_seed_all(seed) + + prep = _prepare_phase_field( + torch, + F, + rgba, + inference, + analysis, + solver_params, + device=device, + ) + + disp_t = torch.zeros_like(prep.uv0_px_t, device=device, dtype=torch.float32, requires_grad=True) + optimizer = torch.optim.Adam([disp_t], lr=solver_params.phase_field_learning_rate) + + initial_x = np.rint(prep.uv0_px_t[0, ..., 0].detach().cpu().numpy()).astype(np.int32).clip(0, prep.width - 1) + initial_y = np.rint(prep.uv0_px_t[0, ..., 1].detach().cpu().numpy()).astype(np.int32).clip(0, prep.height - 1) + initial_rgba = _nearest_source_rgba(rgba, initial_x, initial_y) + + loss_history: list[float] = [] + final_terms: dict[str, float] = {} + min_dx = solver_params.phase_field_min_spacing_ratio * prep.cell_x + min_dy = solver_params.phase_field_min_spacing_ratio * prep.cell_y + max_dx = solver_params.phase_field_max_displacement_ratio * prep.cell_x + max_dy = solver_params.phase_field_max_displacement_ratio * prep.cell_y + + for _ in range(max(0, steps)): + optimizer.zero_grad(set_to_none=True) + loss, _sampled_rgba, terms = _phase_field_loss(torch, F, prep, disp_t, solver_params) + loss.backward() + optimizer.step() + _project_displacements_in_place( + torch, + disp_t, + prep.base_x_t, + prep.base_y_t, + min_dx=min_dx, + min_dy=min_dy, + max_dx=max_dx, + max_dy=max_dy, + width=prep.width, + height=prep.height, + ) + loss_history.append(float(loss.detach().cpu().item())) + final_terms = terms + + with torch.no_grad(): + if steps <= 0: + loss, _sampled_rgba, final_terms = _phase_field_loss(torch, F, prep, disp_t, solver_params) + loss_history.append(float(loss.detach().cpu().item())) + final_px = prep.uv0_px_t + disp_t + final_x = torch.round(final_px[..., 0]).clamp(0, prep.width - 1).to(dtype=torch.long) + final_y = torch.round(final_px[..., 1]).clamp(0, prep.height - 1).to(dtype=torch.long) + final_pos_norm = _pixel_to_normalized(final_px.clone(), width=prep.width, height=prep.height)[0].detach().cpu().numpy() + + final_x_np = final_x[0].detach().cpu().numpy().astype(np.int32) + final_y_np = final_y[0].detach().cpu().numpy().astype(np.int32) + target_rgba = _nearest_source_rgba(rgba, final_x_np, final_y_np) + final_disp = disp_t[0].detach().cpu().numpy().astype(np.float32) + stage_diagnostics = { + "displacements": { + "initial_output": _displacement_diagnostics( + prep.uv0_norm, + initial_x, + initial_y, + width=prep.width, + height=prep.height, + ), + "final_output": _displacement_diagnostics( + prep.uv0_norm, + final_x_np, + final_y_np, + width=prep.width, + height=prep.height, + ), + }, + "phase_field": { + "max_abs_dx_px": float(np.max(np.abs(final_disp[..., 0]))), + "max_abs_dy_px": float(np.max(np.abs(final_disp[..., 1]))), + "mean_displacement_px": float(np.mean(np.sqrt(np.square(final_disp[..., 0]) + np.square(final_disp[..., 1])))), + **final_terms, + }, + } + return SolverArtifacts( + target_rgba=target_rgba, + uv_field=final_pos_norm, + guidance_strength=prep.guidance, + initial_rgba=initial_rgba, + loss_history=loss_history, + stage_diagnostics=stage_diagnostics, + ) diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index 272637c..5f65fb9 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -27,6 +27,7 @@ from .metrics import ( source_lattice_consistency_breakdown, ) from .params import SolverHyperParams +from .phase_field import optimize_phase_field from .palette import load_palette, quantize_rgba, save_palette_report from .preprocess import strip_edge_background from .tile_graph import optimize_tile_graph @@ -440,6 +441,23 @@ def _run_reconstruction( device=device, solver_params=solver_params, ) + if reconstruction_mode == "phase-field": + solver = optimize_phase_field( + source, + inference=inference, + analysis=analysis, + steps=steps, + seed=seed, + device=device, + solver_params=solver_params, + ) + return solver, { + "mode": "phase-field", + **{ + f"phase_field_{key}": value + for key, value in solver.stage_diagnostics.get("phase_field", {}).items() + }, + } solver = optimize_lattice_pixels( source, inference=inference, diff --git a/tests/test_phase_field.py b/tests/test_phase_field.py new file mode 100644 index 0000000..d307c88 --- /dev/null +++ b/tests/test_phase_field.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import numpy as np + +from repixelizer.analysis import analyze_continuous_source +from repixelizer.phase_field import optimize_phase_field +from repixelizer.synthetic import fake_pixelize, make_emblem +from repixelizer.types import InferenceResult + + +def test_optimize_phase_field_emits_displacement_diagnostics_and_source_colors() -> None: + source = np.zeros((12, 12, 4), dtype=np.float32) + source[..., 3] = 1.0 + source[2:10, 2:10, :3] = np.asarray([0.8, 0.2, 0.1], dtype=np.float32) + source[4:8, 4:8, :3] = np.asarray([0.1, 0.1, 0.9], dtype=np.float32) + + inference = InferenceResult( + target_width=6, + target_height=6, + phase_x=0.0, + phase_y=0.0, + confidence=1.0, + top_candidates=[], + ) + artifacts = optimize_phase_field( + source, + inference=inference, + analysis=analyze_continuous_source(source, seed=7), + steps=4, + seed=7, + device="cpu", + ) + + displacement = artifacts.stage_diagnostics["displacements"] + assert {"initial_output", "final_output"} <= set(displacement.keys()) + assert displacement["final_output"]["displacement_x"].shape == (6, 6) + assert displacement["final_output"]["displacement_y"].shape == (6, 6) + assert artifacts.stage_diagnostics["phase_field"]["mean_displacement_px"] >= 0.0 + + source_colors = {tuple(px.tolist()) for px in np.clip(np.rint(source * 255.0), 0, 255).astype(np.uint8).reshape(-1, 4)} + output_colors = {tuple(px.tolist()) for px in np.clip(np.rint(artifacts.target_rgba * 255.0), 0, 255).astype(np.uint8).reshape(-1, 4)} + assert output_colors <= source_colors + + +def test_optimize_phase_field_moves_off_zero_when_phase_is_wrong() -> None: + source = make_emblem(18, 18) + fake = fake_pixelize( + source, + upscale=8, + phase_x=0.18, + phase_y=-0.14, + blur_radius=0.35, + warp_strength=0.12, + warp_detail=4, + seed=4, + ) + inference = InferenceResult( + target_width=18, + target_height=18, + phase_x=0.0, + phase_y=0.0, + confidence=1.0, + top_candidates=[], + ) + artifacts = optimize_phase_field( + fake, + inference=inference, + analysis=analyze_continuous_source(fake, seed=7), + steps=6, + seed=7, + device="cpu", + ) + + assert artifacts.stage_diagnostics["phase_field"]["mean_displacement_px"] > 0.01 diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 858a171..91f94a5 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -354,6 +354,35 @@ def test_select_phase_candidate_skips_phase_rerank_for_tile_graph(monkeypatch) - assert selected.target_height == candidate_a.target_height +def test_select_phase_candidate_skips_phase_rerank_for_phase_field(monkeypatch) -> None: + source = np.zeros((8, 8, 4), dtype=np.float32) + candidate_a = InferenceCandidate(target_width=8, target_height=8, phase_x=0.0, phase_y=0.0, score=0.9, breakdown={}) + candidate_b = InferenceCandidate(target_width=10, target_height=10, phase_x=0.2, phase_y=0.2, score=0.89, breakdown={}) + inference = InferenceResult( + target_width=8, + target_height=8, + phase_x=0.0, + phase_y=0.0, + confidence=0.0, + top_candidates=[candidate_a, candidate_b], + ) + + def fail(*args, **kwargs): + raise AssertionError("phase-field phase rerank probe should not run") + + monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fail) + selected = _select_phase_candidate( + source, + inference, + analysis=object(), + seed=7, + device="cpu", + reconstruction_mode="phase-field", + ) + assert selected.target_width == candidate_a.target_width + assert selected.target_height == candidate_a.target_height + + def test_pipeline_runs_tile_graph_once_without_phase_probe_reuse(tmp_path: Path, monkeypatch) -> None: source = np.zeros((6, 6, 4), dtype=np.float32) source[..., 3] = 1.0 @@ -431,6 +460,62 @@ def test_pipeline_runs_tile_graph_once_without_phase_probe_reuse(tmp_path: Path, assert result.output_rgba.shape[:2] == (3, 3) +def test_pipeline_runs_phase_field_mode(tmp_path: Path, monkeypatch) -> None: + source = np.zeros((6, 6, 4), dtype=np.float32) + source[..., 3] = 1.0 + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, source) + + inference = InferenceResult( + target_width=3, + target_height=3, + phase_x=0.0, + phase_y=0.0, + confidence=1.0, + top_candidates=[], + ) + + class DummyAnalysis: + edge_map = np.zeros((6, 6), dtype=np.float32) + + def fake_optimize_phase_field(source_rgba, inference, analysis, steps, seed, device, solver_params=None): + del source_rgba, analysis, steps, seed, device, solver_params + rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32) + rgba[..., 1] = 1.0 + rgba[..., 3] = 1.0 + return SolverArtifacts( + target_rgba=rgba, + uv_field=np.zeros((inference.target_height, inference.target_width, 2), dtype=np.float32), + guidance_strength=np.zeros((inference.target_height, inference.target_width), dtype=np.float32), + initial_rgba=rgba.copy(), + loss_history=[0.0], + stage_diagnostics={ + "displacements": {}, + "phase_field": {"mean_displacement_px": 0.0}, + }, + ) + + monkeypatch.setattr("repixelizer.pipeline.infer_lattice", lambda *args, **kwargs: inference) + monkeypatch.setattr("repixelizer.pipeline.analyze_continuous_source", lambda *args, **kwargs: DummyAnalysis()) + monkeypatch.setattr("repixelizer.pipeline.optimize_phase_field", fake_optimize_phase_field) + + result = run_pipeline( + input_path, + output_path, + reconstruction_mode="phase-field", + steps=8, + device="cpu", + ) + + assert output_path.exists() + assert result.output_rgba.shape[:2] == (3, 3) + assert result.diagnostics["reconstruction"]["mode"] == "phase-field" + + def test_pipeline_fixed_target_and_phase_skip_search_inference(tmp_path: Path, monkeypatch) -> None: source = np.zeros((6, 6, 4), dtype=np.float32) source[..., 3] = 1.0