513 lines
18 KiB
Python
513 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from repixelizer.analysis import analyze_source
|
|
from repixelizer.baselines import naive_resize_baseline
|
|
from repixelizer.io import save_rgba
|
|
from repixelizer.metrics import source_lattice_consistency_breakdown
|
|
from repixelizer.params import SolverHyperParams
|
|
from repixelizer.pipeline import run_pipeline
|
|
from repixelizer.synthetic import fake_pixelize, make_emblem
|
|
from repixelizer.tile_graph import (
|
|
_extract_source_region_tiles,
|
|
_segment_atomic_source_regions,
|
|
build_tile_graph_model,
|
|
clear_tile_graph_model_cache,
|
|
optimize_tile_graph,
|
|
)
|
|
from repixelizer.types import InferenceResult
|
|
|
|
|
|
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)
|
|
source = np.repeat(np.repeat(lowres, 8, axis=0), 8, axis=1)
|
|
inference = InferenceResult(
|
|
target_width=4,
|
|
target_height=4,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
|
|
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)
|
|
assert choice_counts.shape[0] == inference.target_width * inference.target_height
|
|
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:
|
|
source = np.zeros((8, 8, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
palette = (
|
|
np.asarray([1.0, 0.1, 0.1, 1.0], dtype=np.float32),
|
|
np.asarray([0.1, 1.0, 0.1, 1.0], dtype=np.float32),
|
|
np.asarray([0.1, 0.1, 1.0, 1.0], dtype=np.float32),
|
|
np.asarray([1.0, 1.0, 0.1, 1.0], dtype=np.float32),
|
|
)
|
|
for y in range(source.shape[0]):
|
|
for x in range(source.shape[1]):
|
|
source[y, x] = palette[(2 * (y % 2)) + (x % 2)]
|
|
|
|
inference = InferenceResult(
|
|
target_width=4,
|
|
target_height=4,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
model = build_tile_graph_model(
|
|
source,
|
|
inference=inference,
|
|
analysis=analyze_source(source, seed=13),
|
|
solver_params=SolverHyperParams(tile_graph_max_candidates=32, tile_graph_max_candidates_per_coord=2),
|
|
)
|
|
|
|
source_colors = {tuple(np.round(color, 4)) for color in source.reshape(-1, 4)}
|
|
candidate_colors = {tuple(np.round(color, 4)) for color in model.candidate_rgba}
|
|
|
|
assert candidate_colors.issubset(source_colors | {tuple(np.zeros(4, dtype=np.float32))})
|
|
|
|
|
|
def test_tile_graph_candidates_are_scoped_to_their_output_coord() -> None:
|
|
source = np.zeros((8, 8, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
source[:, 2:4] = np.asarray([1.0, 0.9, 0.2, 1.0], dtype=np.float32)
|
|
source[2:6, 5:7] = np.asarray([0.2, 0.8, 1.0, 1.0], dtype=np.float32)
|
|
inference = InferenceResult(
|
|
target_width=4,
|
|
target_height=4,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
|
|
model = build_tile_graph_model(
|
|
source,
|
|
inference=inference,
|
|
analysis=analyze_source(source, seed=11),
|
|
)
|
|
|
|
for y in range(inference.target_height):
|
|
for x in range(inference.target_width):
|
|
flat = y * inference.target_width + x
|
|
start = int(model.cell_candidate_offsets[flat])
|
|
end = int(model.cell_candidate_offsets[flat + 1])
|
|
coords = model.candidate_coords[model.cell_candidate_indices[start:end]]
|
|
assert np.all(coords[:, 0] == y)
|
|
assert np.all(coords[:, 1] == x)
|
|
|
|
|
|
def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> None:
|
|
source = np.zeros((64, 64, 4), dtype=np.float32)
|
|
white = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32)
|
|
source[6:10, 6:58] = white
|
|
source[14:18, 6:50] = white
|
|
source[22:26, 14:58] = white
|
|
source[30:34, 6:50] = white
|
|
source[38:42, 14:58] = white
|
|
source[46:50, 6:58] = white
|
|
source[6:50, 6:10] = white
|
|
source[14:58, 54:58] = white
|
|
|
|
inference = InferenceResult(
|
|
target_width=8,
|
|
target_height=8,
|
|
phase_x=0.0,
|
|
phase_y=-0.2,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
params = SolverHyperParams()
|
|
analysis = analyze_source(source, seed=7, device="cpu")
|
|
cell_h = source.shape[0] / inference.target_height
|
|
cell_w = source.shape[1] / inference.target_width
|
|
source_region_stride = 2
|
|
sampled_rgba = source[::source_region_stride, ::source_region_stride].astype(np.float32)
|
|
sampled_edge = analysis.edge_map[::source_region_stride, ::source_region_stride].astype(np.float32)
|
|
sampled_y_coords = np.arange(0, source.shape[0], source_region_stride, dtype=np.int32)
|
|
sampled_x_coords = np.arange(0, source.shape[1], source_region_stride, dtype=np.int32)
|
|
sampled_y_grid, sampled_x_grid = np.meshgrid(sampled_y_coords, sampled_x_coords, indexing="ij")
|
|
sampled_flat_rgba = sampled_rgba.reshape(-1, sampled_rgba.shape[-1]).astype(np.float32)
|
|
sampled_flat_edge = sampled_edge.reshape(-1).astype(np.float32)
|
|
sampled_flat_y = sampled_y_grid.reshape(-1).astype(np.int32)
|
|
sampled_flat_x = sampled_x_grid.reshape(-1).astype(np.int32)
|
|
|
|
components = _segment_atomic_source_regions(
|
|
source_rgba=sampled_rgba,
|
|
edge_map=sampled_edge,
|
|
alpha_floor=params.alpha_transparent_threshold,
|
|
color_threshold=params.tile_graph_component_color_threshold,
|
|
alpha_threshold=params.tile_graph_component_alpha_threshold,
|
|
device="cpu",
|
|
)
|
|
region_buckets = _extract_source_region_tiles(
|
|
components=components,
|
|
flat_rgba=sampled_flat_rgba,
|
|
flat_edge=sampled_flat_edge,
|
|
flat_x=sampled_flat_x,
|
|
flat_y=sampled_flat_y,
|
|
cell_w=cell_w,
|
|
cell_h=cell_h,
|
|
sample_area=float(source_region_stride * source_region_stride),
|
|
target_width=inference.target_width,
|
|
target_height=inference.target_height,
|
|
phase_x=inference.phase_x,
|
|
phase_y=inference.phase_y,
|
|
min_region_area_ratio=params.tile_graph_source_region_min_area_ratio,
|
|
min_window_coverage=params.tile_graph_source_region_window_coverage,
|
|
stroke_linearity_threshold=params.tile_graph_stroke_linearity_threshold,
|
|
stroke_step_scale=params.tile_graph_stroke_step_scale,
|
|
stroke_minor_limit_scale=params.tile_graph_stroke_minor_limit_scale,
|
|
)
|
|
|
|
opaque = sampled_flat_rgba[:, 3] >= params.alpha_transparent_threshold
|
|
coord_x = np.clip(
|
|
np.floor((sampled_flat_x.astype(np.float32) + 0.5) / cell_w - inference.phase_x).astype(np.int32),
|
|
0,
|
|
inference.target_width - 1,
|
|
)
|
|
coord_y = np.clip(
|
|
np.floor((sampled_flat_y.astype(np.float32) + 0.5) / cell_h - inference.phase_y).astype(np.int32),
|
|
0,
|
|
inference.target_height - 1,
|
|
)
|
|
direct_indices = coord_y * inference.target_width + coord_x
|
|
occupied_direct = np.bincount(
|
|
direct_indices[opaque],
|
|
minlength=inference.target_width * inference.target_height,
|
|
)
|
|
occupied_bucket = np.asarray([len(bucket) for bucket in region_buckets], dtype=np.int32)
|
|
|
|
assert np.all(occupied_bucket[occupied_direct > 0] > 0)
|
|
|
|
|
|
def test_tile_graph_projects_source_regions_instead_of_mixed_lattice_buckets() -> None:
|
|
source = np.zeros((6, 12, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
red = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
|
blue = np.asarray([0.0, 0.0, 1.0, 1.0], dtype=np.float32)
|
|
source[:, :6] = red
|
|
source[:, 6:] = blue
|
|
inference = InferenceResult(
|
|
target_width=2,
|
|
target_height=1,
|
|
phase_x=0.5,
|
|
phase_y=0.0,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
|
|
model = build_tile_graph_model(
|
|
source,
|
|
inference=inference,
|
|
analysis=analyze_source(source, seed=19),
|
|
device="cpu",
|
|
)
|
|
|
|
left_start = int(model.cell_candidate_offsets[0])
|
|
left_end = int(model.cell_candidate_offsets[1])
|
|
right_start = int(model.cell_candidate_offsets[1])
|
|
right_end = int(model.cell_candidate_offsets[2])
|
|
left_candidates = {
|
|
tuple(np.round(color, 4))
|
|
for color in model.candidate_rgba[model.cell_candidate_indices[left_start:left_end]]
|
|
}
|
|
right_candidates = {
|
|
tuple(np.round(color, 4))
|
|
for color in model.candidate_rgba[model.cell_candidate_indices[right_start:right_end]]
|
|
}
|
|
|
|
assert tuple(np.round(red, 4)) in left_candidates
|
|
assert tuple(np.round(blue, 4)) not in left_candidates
|
|
assert tuple(np.round(blue, 4)) in right_candidates
|
|
assert tuple(np.round(red, 4)) not in right_candidates
|
|
|
|
|
|
def test_tile_graph_edge_cells_include_multiple_same_cell_contour_colors() -> None:
|
|
source = np.zeros((8, 8, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
source[:4, 0] = np.asarray([1.0, 0.9, 0.2, 1.0], dtype=np.float32)
|
|
source[:4, 1] = np.asarray([0.85, 0.6, 0.1, 1.0], dtype=np.float32)
|
|
source[:4, 2] = np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
|
source[:4, 3] = np.asarray([0.9, 0.2, 0.1, 1.0], dtype=np.float32)
|
|
source[:4, 4:] = np.asarray([0.2, 0.8, 0.2, 1.0], dtype=np.float32)
|
|
source[4:, :] = np.asarray([0.2, 0.2, 0.8, 1.0], dtype=np.float32)
|
|
|
|
inference = InferenceResult(
|
|
target_width=2,
|
|
target_height=2,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
model = build_tile_graph_model(
|
|
source,
|
|
inference=inference,
|
|
analysis=analyze_source(source, seed=17),
|
|
device="cpu",
|
|
)
|
|
|
|
start = int(model.cell_candidate_offsets[0])
|
|
end = int(model.cell_candidate_offsets[1])
|
|
candidates = {
|
|
tuple(np.round(color, 4))
|
|
for color in model.candidate_rgba[model.cell_candidate_indices[start:end]]
|
|
}
|
|
non_black_candidates = {color for color in candidates if color[:3] != (np.float32(0.0), np.float32(0.0), np.float32(0.0))}
|
|
|
|
assert len(candidates) >= 3
|
|
assert tuple(np.round(np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32), 4)) in candidates
|
|
assert len(non_black_candidates) >= 2
|
|
|
|
|
|
def test_tile_graph_reconstructs_synthetic_thin_feature_better_than_naive() -> None:
|
|
source = np.zeros((16, 16, 4), dtype=np.float32)
|
|
source[4:12, 8] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
|
|
source[4, 7:10] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
|
|
source[11, 7] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
|
|
source[10, 6] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
|
|
source[9, 5] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
|
|
|
|
fake = fake_pixelize(
|
|
source,
|
|
upscale=9,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
blur_radius=0.45,
|
|
warp_strength=0.2,
|
|
warp_detail=5,
|
|
seed=11,
|
|
)
|
|
inference = InferenceResult(
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
artifacts, _ = optimize_tile_graph(
|
|
fake,
|
|
inference=inference,
|
|
analysis=analyze_source(fake, seed=7),
|
|
steps=0,
|
|
seed=7,
|
|
device="cpu",
|
|
)
|
|
naive = naive_resize_baseline(fake, width=16, height=16)
|
|
|
|
tile_score = source_lattice_consistency_breakdown(
|
|
fake,
|
|
artifacts.target_rgba,
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
)["score"]
|
|
initial_tile_score = source_lattice_consistency_breakdown(
|
|
fake,
|
|
artifacts.initial_rgba,
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
)["score"]
|
|
naive_score = source_lattice_consistency_breakdown(
|
|
fake,
|
|
naive,
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
)["score"]
|
|
|
|
assert tile_score <= initial_tile_score + 1e-6
|
|
assert tile_score <= naive_score
|
|
|
|
|
|
def test_tile_graph_keeps_shallow_connected_stroke_from_fanning_out() -> None:
|
|
source = np.zeros((48, 48, 4), dtype=np.float32)
|
|
for y in range(48):
|
|
x = int(round(8 + y * 0.65))
|
|
source[max(0, y - 1) : min(48, y + 2), max(0, x - 1) : min(48, x + 2)] = np.asarray(
|
|
[1.0, 1.0, 1.0, 1.0],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
artifacts, _ = optimize_tile_graph(
|
|
source,
|
|
inference=InferenceResult(
|
|
target_width=6,
|
|
target_height=6,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
),
|
|
analysis=analyze_source(source, seed=7),
|
|
steps=0,
|
|
seed=7,
|
|
device="cpu",
|
|
)
|
|
|
|
occupied_per_row = [
|
|
int(np.count_nonzero(artifacts.target_rgba[y, :, 3] > 0.5))
|
|
for y in range(artifacts.target_rgba.shape[0])
|
|
]
|
|
|
|
assert sum(count > 1 for count in occupied_per_row) <= 1
|
|
assert sum(occupied_per_row) <= artifacts.target_rgba.shape[0] + 1
|
|
|
|
|
|
def test_tile_graph_preserves_transparent_background_on_sparse_detail() -> None:
|
|
source = np.zeros((12, 12, 4), dtype=np.float32)
|
|
source[5:7, 5:7] = np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
|
|
|
fake = fake_pixelize(
|
|
source,
|
|
upscale=8,
|
|
phase_x=0.12,
|
|
phase_y=-0.08,
|
|
blur_radius=0.35,
|
|
warp_strength=0.1,
|
|
warp_detail=4,
|
|
seed=5,
|
|
)
|
|
artifacts, _ = optimize_tile_graph(
|
|
fake,
|
|
inference=InferenceResult(
|
|
target_width=12,
|
|
target_height=12,
|
|
phase_x=0.12,
|
|
phase_y=-0.08,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
),
|
|
analysis=analyze_source(fake, seed=7),
|
|
steps=0,
|
|
seed=7,
|
|
device="cpu",
|
|
)
|
|
|
|
alpha_zero_ratio = float((artifacts.target_rgba[..., 3] <= 0.05).mean())
|
|
|
|
assert alpha_zero_ratio >= 0.75
|
|
|
|
|
|
def test_tile_graph_cuda_matches_cpu_on_small_case() -> None:
|
|
torch = pytest.importorskip("torch")
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA is not available in this environment")
|
|
|
|
source = make_emblem(16, 16)
|
|
fake = fake_pixelize(source, upscale=8, phase_x=0.1, phase_y=-0.05, blur_radius=0.4, warp_strength=0.15, warp_detail=4)
|
|
inference = InferenceResult(
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.1,
|
|
phase_y=-0.05,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
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=cpu_analysis,
|
|
steps=0,
|
|
seed=7,
|
|
device="cpu",
|
|
)
|
|
cuda_artifacts, cuda_diag = optimize_tile_graph(
|
|
fake,
|
|
inference=inference,
|
|
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)
|
|
|
|
|
|
def test_tile_graph_model_cache_reuses_fixed_lattice_build() -> None:
|
|
clear_tile_graph_model_cache()
|
|
source = make_emblem(16, 16)
|
|
fake = fake_pixelize(source, upscale=8, phase_x=0.1, phase_y=-0.05, blur_radius=0.4, warp_strength=0.15, warp_detail=4)
|
|
inference = InferenceResult(
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.1,
|
|
phase_y=-0.05,
|
|
confidence=1.0,
|
|
top_candidates=[],
|
|
)
|
|
analysis = analyze_source(fake, seed=7, device="cpu")
|
|
|
|
first = build_tile_graph_model(
|
|
fake,
|
|
inference=inference,
|
|
analysis=analysis,
|
|
device="cpu",
|
|
)
|
|
second = build_tile_graph_model(
|
|
fake,
|
|
inference=inference,
|
|
analysis=analysis,
|
|
device="cpu",
|
|
)
|
|
|
|
assert first.cache_hit is False
|
|
assert second.cache_hit is True
|
|
assert np.array_equal(first.candidate_rgba, second.candidate_rgba)
|
|
|
|
|
|
def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(tmp_path: Path) -> None:
|
|
source = make_emblem(24, 24)
|
|
fake = fake_pixelize(source, upscale=10, phase_x=0.15, phase_y=0.2, blur_radius=0.6, warp_strength=0.2, warp_detail=5)
|
|
input_path = tmp_path / "input.png"
|
|
output_path = tmp_path / "nested" / "output.png"
|
|
diagnostics_dir = tmp_path / "diagnostics"
|
|
|
|
save_rgba(input_path, fake)
|
|
result = run_pipeline(
|
|
input_path,
|
|
output_path,
|
|
diagnostics_dir=diagnostics_dir,
|
|
steps=24,
|
|
reconstruction_mode="tile-graph",
|
|
device="cpu",
|
|
)
|
|
|
|
import json
|
|
|
|
run_json = json.loads((diagnostics_dir / "run.json").read_text(encoding="utf-8"))
|
|
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 "tile_graph_model_cache_hit" in run_json["reconstruction"]
|
|
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"
|