157 lines
4.7 KiB
Python
157 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from repixelizer.analysis import analyze_continuous_source
|
|
from repixelizer.continuous import (
|
|
_build_candidate_positions,
|
|
_build_source_reliability,
|
|
_line_pattern_loss,
|
|
_make_regular_uv,
|
|
optimize_lattice_pixels,
|
|
)
|
|
from repixelizer.metrics import source_lattice_consistency_breakdown
|
|
from repixelizer.params import SolverHyperParams
|
|
from repixelizer.source_reference import build_source_lattice_reference
|
|
from repixelizer.synthetic import fake_pixelize
|
|
from repixelizer.types import InferenceResult
|
|
|
|
|
|
def test_line_pattern_loss_penalizes_wobbling_line() -> None:
|
|
reference = torch.zeros((1, 5, 5, 4), dtype=torch.float32)
|
|
reference[:, :, 2, :] = 1.0
|
|
|
|
good = reference.clone()
|
|
bad = reference.clone()
|
|
bad[:, 2, 2, :] = 0.0
|
|
bad[:, 2, 3, :] = 1.0
|
|
|
|
good_loss = float(_line_pattern_loss(torch, good, reference).item())
|
|
bad_loss = float(_line_pattern_loss(torch, bad, reference).item())
|
|
|
|
assert good_loss <= 1e-6
|
|
assert bad_loss > good_loss + 0.01
|
|
|
|
|
|
def test_source_reliability_stays_high_for_high_contrast_edge_cells() -> None:
|
|
source = np.zeros((8, 8, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
source[:, 1, :3] = 1.0
|
|
|
|
analysis = analyze_continuous_source(source, seed=3)
|
|
reference = build_source_lattice_reference(
|
|
source,
|
|
target_width=4,
|
|
target_height=4,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
edge_hint=analysis.edge_map,
|
|
edge_grad_x_hint=np.gradient(analysis.edge_map, axis=1).astype(np.float32),
|
|
edge_grad_y_hint=np.gradient(analysis.edge_map, axis=0).astype(np.float32),
|
|
)
|
|
|
|
reliability = _build_source_reliability(reference, SolverHyperParams())
|
|
|
|
assert reference.cell_dispersion[1, 0] > reference.dispersion
|
|
assert reference.edge_strength[1, 0] > 0.0
|
|
assert reliability[1, 0] >= 0.7
|
|
|
|
|
|
def test_candidate_positions_include_sharp_and_edge_guided_samples() -> None:
|
|
source = np.zeros((6, 6, 4), dtype=np.float32)
|
|
source[..., 3] = 1.0
|
|
source[1:5, 2, :3] = 1.0
|
|
|
|
edge_hint = np.zeros((6, 6), dtype=np.float32)
|
|
edge_hint[1:5, 2] = 1.0
|
|
edge_grad_x = np.zeros((6, 6), dtype=np.float32)
|
|
edge_grad_x[1:5, 2] = 1.0
|
|
edge_grad_y = np.zeros((6, 6), dtype=np.float32)
|
|
|
|
reference = build_source_lattice_reference(
|
|
source,
|
|
target_width=2,
|
|
target_height=2,
|
|
phase_x=0.0,
|
|
phase_y=0.0,
|
|
edge_hint=edge_hint,
|
|
edge_grad_x_hint=edge_grad_x,
|
|
edge_grad_y_hint=edge_grad_y,
|
|
)
|
|
uv = torch.from_numpy(_make_regular_uv(6, 6, 2, 2, 0.0, 0.0)).to(dtype=torch.float32)
|
|
|
|
candidate_x, candidate_y, _, _ = _build_candidate_positions(
|
|
torch,
|
|
uv,
|
|
reference,
|
|
SolverHyperParams(),
|
|
width=6,
|
|
height=6,
|
|
cell_x=3.0,
|
|
cell_y=3.0,
|
|
base_fraction_values=np.asarray([0.0], dtype=np.float32),
|
|
)
|
|
|
|
cell_candidates = {
|
|
(int(y), int(x))
|
|
for y, x in zip(candidate_y[0, 0].cpu().tolist(), candidate_x[0, 0].cpu().tolist())
|
|
}
|
|
|
|
assert (int(reference.edge_peak_y[0, 0]), int(reference.edge_peak_x[0, 0])) in cell_candidates
|
|
assert len(cell_candidates) > 2
|
|
|
|
|
|
def test_optimize_lattice_pixels_does_not_regress_from_snap_on_thin_feature_case() -> 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_lattice_pixels(
|
|
fake,
|
|
inference=inference,
|
|
analysis=analyze_continuous_source(fake, seed=7),
|
|
steps=4,
|
|
seed=7,
|
|
device="cpu",
|
|
)
|
|
|
|
snap_score = source_lattice_consistency_breakdown(
|
|
fake,
|
|
artifacts.initial_rgba,
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
)["score"]
|
|
final_score = source_lattice_consistency_breakdown(
|
|
fake,
|
|
artifacts.target_rgba,
|
|
target_width=16,
|
|
target_height=16,
|
|
phase_x=0.18,
|
|
phase_y=-0.12,
|
|
)["score"]
|
|
|
|
assert final_score <= snap_score + 1e-6
|