Remove legacy continuous optimizer
This commit is contained in:
Meta Vens
parent
d9fa411daf
commit
0808271d05
19 changed files with 360 additions and 3434 deletions
|
|
@ -6,35 +6,17 @@ from pathlib import Path
|
|||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from repixelizer.analysis import analyze_source
|
||||
from repixelizer.continuous import (
|
||||
_build_candidate_positions,
|
||||
_build_source_detail_reference,
|
||||
_discrete_refine_output,
|
||||
_build_source_reliability,
|
||||
_edge_gradient_maps,
|
||||
_make_patch_offsets,
|
||||
_make_regular_uv,
|
||||
_relax_candidate_selection,
|
||||
_reference_match_energy,
|
||||
_representative_colors,
|
||||
_require_torch,
|
||||
_resolve_device,
|
||||
_sample_cell_patches,
|
||||
_select_colors,
|
||||
_snap_output_to_source_pixels,
|
||||
premultiply,
|
||||
)
|
||||
from repixelizer.analysis import analyze_phase_field_source
|
||||
from repixelizer.inference import infer_lattice
|
||||
from repixelizer.io import load_rgba
|
||||
from repixelizer.metrics import source_lattice_consistency_breakdown
|
||||
from repixelizer.params import SolverHyperParams
|
||||
from repixelizer.source_reference import build_source_lattice_reference
|
||||
from repixelizer.phase_field import optimize_phase_field
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Render a source/snap/relaxed/final crop sheet for a fixed output-grid region."
|
||||
description="Render a source/initial/final crop sheet for a fixed output-grid region."
|
||||
)
|
||||
parser.add_argument("--input", required=True, help="Input image path.")
|
||||
parser.add_argument("--out", required=True, help="Output PNG path.")
|
||||
|
|
@ -48,11 +30,11 @@ def _parse_args() -> argparse.Namespace:
|
|||
)
|
||||
parser.add_argument(
|
||||
"--panels",
|
||||
default="source,snap,relaxed,final",
|
||||
help="Comma-separated panel list. Supported: source,snap,relaxed,final",
|
||||
default="source,initial,final",
|
||||
help="Comma-separated panel list. Supported: source,initial,final. Legacy aliases: snap=initial, relaxed=final.",
|
||||
)
|
||||
parser.add_argument("--scale", type=int, default=16, help="Nearest-neighbor scale factor for each output cell.")
|
||||
parser.add_argument("--steps", type=int, default=64, help="Discrete refine iteration budget.")
|
||||
parser.add_argument("--steps", type=int, default=48, help="Phase-field iteration budget.")
|
||||
parser.add_argument("--device", default="auto", help="Torch device, usually auto/cpu/cuda.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
|
@ -93,215 +75,42 @@ def _build_focus_states(
|
|||
*,
|
||||
steps: int,
|
||||
device: str,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, tuple[float, float], dict[str, dict[str, float]]]:
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, tuple[float, float], dict[str, dict[str, float]]]:
|
||||
source = load_rgba(input_path)
|
||||
solver_params = SolverHyperParams()
|
||||
inference = infer_lattice(source, None, device=device)
|
||||
analysis = analyze_source(source, seed=7)
|
||||
inference = infer_lattice(source, device=device)
|
||||
analysis = analyze_phase_field_source(source, seed=7, device=device)
|
||||
artifacts = optimize_phase_field(
|
||||
source,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
steps=steps,
|
||||
seed=7,
|
||||
device=device,
|
||||
solver_params=SolverHyperParams(),
|
||||
)
|
||||
|
||||
torch, F = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
|
||||
premul = premultiply(source)
|
||||
height, width = premul.shape[:2]
|
||||
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(premul.transpose(2, 0, 1)[None, ...]).to(device=resolved_device, dtype=torch.float32)
|
||||
uv0 = _make_regular_uv(
|
||||
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_t = torch.from_numpy(uv0[None, ...]).to(device=resolved_device, dtype=torch.float32)
|
||||
offsets_t = torch.from_numpy(
|
||||
_make_patch_offsets(height=height, width=width, target_height=inference.target_height, target_width=inference.target_width)
|
||||
).to(device=resolved_device, dtype=torch.float32)
|
||||
source_edge = analysis.edge_map.astype(np.float32)
|
||||
edge_grad_x, edge_grad_y = _edge_gradient_maps(source_edge)
|
||||
source_lattice_reference = build_source_lattice_reference(
|
||||
source,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
alpha_threshold=solver_params.alpha_transparent_threshold,
|
||||
edge_hint=source_edge,
|
||||
edge_grad_x_hint=edge_grad_x,
|
||||
edge_grad_y_hint=edge_grad_y,
|
||||
)
|
||||
source_detail_reference = _build_source_detail_reference(source, source_lattice_reference, solver_params)
|
||||
initial_patches = _sample_cell_patches(F, source_t, uv0_t, offsets_t)
|
||||
representative_t, _ = _representative_colors(initial_patches, solver_params)
|
||||
representative_t = representative_t.detach()
|
||||
source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to(
|
||||
device=resolved_device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
source_reliability_t = torch.from_numpy(_build_source_reliability(source_lattice_reference, solver_params)[None, ...]).to(
|
||||
device=resolved_device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
source_delta_x_t = (
|
||||
torch.from_numpy((premultiply(source_detail_reference)[:, 1:, :] - premultiply(source_detail_reference)[:, :-1, :])[None, ...]).to(device=resolved_device, dtype=torch.float32)
|
||||
if inference.target_width > 1
|
||||
else None
|
||||
)
|
||||
source_delta_y_t = (
|
||||
torch.from_numpy((premultiply(source_detail_reference)[1:, :, :] - premultiply(source_detail_reference)[:-1, :, :])[None, ...]).to(device=resolved_device, dtype=torch.float32)
|
||||
if inference.target_height > 1
|
||||
else None
|
||||
)
|
||||
source_delta_diag_t = (
|
||||
torch.from_numpy((premultiply(source_detail_reference)[1:, 1:, :] - premultiply(source_detail_reference)[:-1, :-1, :])[None, ...]).to(device=resolved_device, dtype=torch.float32)
|
||||
if inference.target_height > 1 and inference.target_width > 1
|
||||
else None
|
||||
)
|
||||
source_delta_anti_t = (
|
||||
torch.from_numpy((premultiply(source_detail_reference)[1:, :-1, :] - premultiply(source_detail_reference)[:-1, 1:, :])[None, ...]).to(device=resolved_device, dtype=torch.float32)
|
||||
if inference.target_height > 1 and inference.target_width > 1
|
||||
else None
|
||||
)
|
||||
|
||||
snap_rgba = _snap_output_to_source_pixels(
|
||||
torch,
|
||||
source_t,
|
||||
uv0_t,
|
||||
representative_t,
|
||||
source_reference_t,
|
||||
source_reliability_t,
|
||||
source_lattice_reference,
|
||||
solver_params,
|
||||
cell_x=cell_x,
|
||||
cell_y=cell_y,
|
||||
)
|
||||
snap_t = torch.from_numpy(premultiply(snap_rgba).transpose(2, 0, 1)[None, ...]).to(
|
||||
device=resolved_device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
source_hw = source_t[0].permute(1, 2, 0)
|
||||
uv = uv0_t[0]
|
||||
representative = representative_t[0]
|
||||
anchor = snap_t[0].permute(1, 2, 0)
|
||||
output_height = representative.shape[0]
|
||||
output_width = representative.shape[1]
|
||||
|
||||
candidate_levels = max(3, int(solver_params.refine_candidate_levels))
|
||||
if candidate_levels % 2 == 0:
|
||||
candidate_levels += 1
|
||||
candidate_extent = max(0.05, float(solver_params.refine_candidate_extent))
|
||||
fraction_values = np.linspace(-candidate_extent, candidate_extent, candidate_levels, dtype=np.float32)
|
||||
candidate_x, candidate_y, candidate_offset_x, candidate_offset_y = _build_candidate_positions(
|
||||
torch,
|
||||
uv,
|
||||
source_lattice_reference,
|
||||
solver_params,
|
||||
width=width,
|
||||
height=height,
|
||||
cell_x=cell_x,
|
||||
cell_y=cell_y,
|
||||
base_fraction_values=fraction_values,
|
||||
)
|
||||
candidate_colors = source_hw[candidate_y, candidate_x]
|
||||
anchor_energy = (candidate_colors - anchor[..., None, :]).abs().mean(dim=-1)
|
||||
rep_energy = (candidate_colors - representative[..., None, :]).abs().mean(dim=-1)
|
||||
source_reference_energy = (candidate_colors - source_reference_t[0][..., None, :]).abs().mean(dim=-1)
|
||||
alpha_energy = (candidate_colors[..., 3] - anchor[..., None, 3]).abs()
|
||||
distance_energy = (candidate_offset_x / max(cell_x, 1e-4)).square() + (candidate_offset_y / max(cell_y, 1e-4)).square()
|
||||
match_energy = _reference_match_energy(
|
||||
rep_energy,
|
||||
source_reference_energy,
|
||||
source_reliability_t[0],
|
||||
representative_weight=solver_params.refine_representative_match_weight,
|
||||
source_weight=solver_params.refine_source_match_weight,
|
||||
)
|
||||
relax_base_energy = (
|
||||
anchor_energy * solver_params.refine_anchor_weight * solver_params.relax_anchor_scale
|
||||
+ match_energy * solver_params.refine_representative_weight
|
||||
+ alpha_energy * solver_params.refine_alpha_weight
|
||||
+ distance_energy * solver_params.refine_distance_weight
|
||||
)
|
||||
|
||||
anchor_delta_x = anchor[:, 1:, :] - anchor[:, :-1, :] if output_width > 1 else None
|
||||
anchor_delta_y = anchor[1:, :, :] - anchor[:-1, :, :] if output_height > 1 else None
|
||||
anchor_delta_diag = anchor[1:, 1:, :] - anchor[:-1, :-1, :] if output_height > 1 and output_width > 1 else None
|
||||
anchor_delta_anti = anchor[1:, :-1, :] - anchor[:-1, 1:, :] if output_height > 1 and output_width > 1 else None
|
||||
source_weight = solver_params.refine_source_delta_weight
|
||||
anchor_weight = 1.0 - source_weight
|
||||
|
||||
def blend(anchor_delta: torch.Tensor | None, source_delta_t: torch.Tensor | None) -> torch.Tensor | None:
|
||||
if anchor_delta is None:
|
||||
return None
|
||||
if source_delta_t is None:
|
||||
return anchor_delta
|
||||
return anchor_delta * anchor_weight + source_delta_t[0] * source_weight
|
||||
|
||||
selected, _, _, _ = _relax_candidate_selection(
|
||||
torch,
|
||||
candidate_colors,
|
||||
relax_base_energy,
|
||||
anchor,
|
||||
source_reference_t[0],
|
||||
blend(anchor_delta_x, source_delta_x_t),
|
||||
blend(anchor_delta_y, source_delta_y_t),
|
||||
blend(anchor_delta_diag, source_delta_diag_t),
|
||||
blend(anchor_delta_anti, source_delta_anti_t),
|
||||
solver_params,
|
||||
iterations=solver_params.relax_iterations,
|
||||
)
|
||||
relaxed_rgba = _select_colors(candidate_colors, selected).detach().cpu().numpy()
|
||||
|
||||
final_rgba, _ = _discrete_refine_output(
|
||||
torch,
|
||||
F,
|
||||
source_t,
|
||||
uv0_t,
|
||||
representative_t,
|
||||
source_reference_t,
|
||||
source_reliability_t,
|
||||
source_lattice_reference,
|
||||
snap_t,
|
||||
source_delta_x_t,
|
||||
source_delta_y_t,
|
||||
source_delta_diag_t,
|
||||
source_delta_anti_t,
|
||||
solver_params,
|
||||
cell_x=cell_x,
|
||||
cell_y=cell_y,
|
||||
iterations=steps,
|
||||
)
|
||||
fidelity = {}
|
||||
fidelity["snap"] = source_lattice_consistency_breakdown(
|
||||
source,
|
||||
snap_rgba,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
)
|
||||
fidelity["relaxed"] = source_lattice_consistency_breakdown(
|
||||
source,
|
||||
relaxed_rgba,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
)
|
||||
fidelity["final"] = source_lattice_consistency_breakdown(
|
||||
source,
|
||||
final_rgba,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
)
|
||||
if fidelity["final"]["score"] > fidelity["snap"]["score"] + 1e-6:
|
||||
final_rgba = snap_rgba
|
||||
fidelity["final"] = fidelity["snap"]
|
||||
return source, snap_rgba, relaxed_rgba, final_rgba, (cell_x, cell_y), fidelity
|
||||
fidelity = {
|
||||
"initial": source_lattice_consistency_breakdown(
|
||||
source,
|
||||
artifacts.initial_rgba,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
),
|
||||
"final": source_lattice_consistency_breakdown(
|
||||
source,
|
||||
artifacts.target_rgba,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
),
|
||||
}
|
||||
return source, artifacts.initial_rgba, artifacts.target_rgba, (cell_x, cell_y), fidelity
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
|
@ -310,13 +119,21 @@ def main() -> None:
|
|||
if x1 <= x0 or y1 <= y0:
|
||||
raise SystemExit("Invalid --cell-bbox: expected X1 > X0 and Y1 > Y0.")
|
||||
|
||||
source, snap_rgba, relaxed_rgba, final_rgba, (cell_x, cell_y), fidelity = _build_focus_states(
|
||||
source, initial_rgba, final_rgba, (cell_x, cell_y), fidelity = _build_focus_states(
|
||||
args.input,
|
||||
steps=args.steps,
|
||||
device=args.device,
|
||||
)
|
||||
|
||||
panel_names = [name.strip().lower() for name in args.panels.split(",") if name.strip()]
|
||||
normalized_panel_names = []
|
||||
for name in (token.strip().lower() for token in args.panels.split(",") if token.strip()):
|
||||
if name == "snap":
|
||||
normalized_panel_names.append("initial")
|
||||
elif name == "relaxed":
|
||||
normalized_panel_names.append("final")
|
||||
else:
|
||||
normalized_panel_names.append(name)
|
||||
|
||||
panels: list[Image.Image] = []
|
||||
source_bbox = (
|
||||
int(np.floor(x0 * cell_x)),
|
||||
|
|
@ -327,11 +144,10 @@ def main() -> None:
|
|||
target_size = (x1 - x0, y1 - y0)
|
||||
panel_map = {
|
||||
"source": _source_panel(source, source_bbox=source_bbox, target_size=target_size, scale=args.scale),
|
||||
"snap": _grid_panel(snap_rgba, "Snap", x0=x0, y0=y0, x1=x1, y1=y1, scale=args.scale),
|
||||
"relaxed": _grid_panel(relaxed_rgba, "Relaxed", x0=x0, y0=y0, x1=x1, y1=y1, scale=args.scale),
|
||||
"initial": _grid_panel(initial_rgba, "Initial", x0=x0, y0=y0, x1=x1, y1=y1, scale=args.scale),
|
||||
"final": _grid_panel(final_rgba, "Final", x0=x0, y0=y0, x1=x1, y1=y1, scale=args.scale),
|
||||
}
|
||||
for name in panel_names:
|
||||
for name in normalized_panel_names:
|
||||
if name not in panel_map:
|
||||
raise SystemExit(f"Unsupported panel name: {name}")
|
||||
panels.append(panel_map[name])
|
||||
|
|
@ -350,7 +166,7 @@ def main() -> None:
|
|||
{
|
||||
"cell_bbox": [x0, y0, x1, y1],
|
||||
"source_bbox": list(source_bbox),
|
||||
"panels": panel_names,
|
||||
"panels": normalized_panel_names,
|
||||
"scale": args.scale,
|
||||
"source_fidelity": fidelity,
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue