Prune tile-graph sidecar paths
This commit is contained in:
parent
3cff705d5e
commit
d33a68c310
8 changed files with 83 additions and 291 deletions
|
|
@ -36,7 +36,7 @@ def build_parser() -> argparse.ArgumentParser:
|
|||
target.add_argument(
|
||||
"--reconstruction-mode",
|
||||
default="continuous",
|
||||
choices=("continuous", "tile-graph", "hybrid"),
|
||||
choices=("continuous", "tile-graph"),
|
||||
help="Reconstruction engine to run after lattice estimation",
|
||||
)
|
||||
target.add_argument(
|
||||
|
|
|
|||
|
|
@ -92,8 +92,6 @@ class SolverHyperParams:
|
|||
tile_graph_coverage_weight: float = 0.05
|
||||
tile_graph_delta_weight: float = 0.45
|
||||
tile_graph_iterations: int = 10
|
||||
hybrid_geometry_match_weight: float = 0.35
|
||||
hybrid_geometry_edge_boost: float = 0.55
|
||||
|
||||
def to_dict(self) -> dict[str, float | int]:
|
||||
return asdict(self)
|
||||
|
|
|
|||
|
|
@ -32,10 +32,6 @@ from .tile_graph import optimize_tile_graph
|
|||
from .types import InferenceResult, RunResult
|
||||
|
||||
|
||||
def _reuse_phase_probe_reconstruction(reconstruction_mode: str) -> bool:
|
||||
return reconstruction_mode in {"tile-graph", "hybrid"}
|
||||
|
||||
|
||||
def run_pipeline(
|
||||
input_path: str | Path,
|
||||
output_path: str | Path,
|
||||
|
|
@ -83,8 +79,8 @@ def run_pipeline(
|
|||
device=device,
|
||||
)
|
||||
inference_mode = "fixed"
|
||||
analysis = analyze_source(source, seed=seed, device=device if reconstruction_mode in {"tile-graph", "hybrid"} else None)
|
||||
inference, cached_reconstruction = _select_phase_candidate_with_reconstruction(
|
||||
analysis = analyze_source(source, seed=seed, device=device if reconstruction_mode == "tile-graph" else None)
|
||||
inference = _select_phase_candidate(
|
||||
source,
|
||||
inference,
|
||||
analysis=analysis,
|
||||
|
|
@ -94,21 +90,16 @@ def run_pipeline(
|
|||
reconstruction_mode=reconstruction_mode,
|
||||
enable_phase_rerank=enable_phase_rerank,
|
||||
)
|
||||
if cached_reconstruction is not None:
|
||||
solver, reconstruction_diagnostics = cached_reconstruction
|
||||
reconstruction_diagnostics = dict(reconstruction_diagnostics)
|
||||
reconstruction_diagnostics["reused_phase_probe_reconstruction"] = True
|
||||
else:
|
||||
solver, reconstruction_diagnostics = _run_reconstruction(
|
||||
source,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
device=device,
|
||||
solver_params=solver_params,
|
||||
reconstruction_mode=reconstruction_mode,
|
||||
)
|
||||
solver, reconstruction_diagnostics = _run_reconstruction(
|
||||
source,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
device=device,
|
||||
solver_params=solver_params,
|
||||
reconstruction_mode=reconstruction_mode,
|
||||
)
|
||||
cleanup = cleanup_pixels(solver.target_rgba, source_guidance=solver.guidance_strength)
|
||||
palette = load_palette(palette_path) if palette_path else None
|
||||
palette_result = quantize_rgba(cleanup.cleaned_rgba, mode=palette_mode, palette=palette)
|
||||
|
|
@ -181,7 +172,7 @@ def _select_phase_candidate(
|
|||
reconstruction_mode: str = "continuous",
|
||||
enable_phase_rerank: bool = True,
|
||||
) -> InferenceResult:
|
||||
selected, _ = _select_phase_candidate_with_reconstruction(
|
||||
return _select_phase_candidate_with_reconstruction(
|
||||
source,
|
||||
inference,
|
||||
analysis=analysis,
|
||||
|
|
@ -191,7 +182,6 @@ def _select_phase_candidate(
|
|||
reconstruction_mode=reconstruction_mode,
|
||||
enable_phase_rerank=enable_phase_rerank,
|
||||
)
|
||||
return selected
|
||||
|
||||
|
||||
def _select_phase_candidate_with_reconstruction(
|
||||
|
|
@ -204,12 +194,14 @@ def _select_phase_candidate_with_reconstruction(
|
|||
solver_params: SolverHyperParams | None = None,
|
||||
reconstruction_mode: str = "continuous",
|
||||
enable_phase_rerank: bool = True,
|
||||
):
|
||||
) -> InferenceResult:
|
||||
solver_params = solver_params or SolverHyperParams()
|
||||
if not enable_phase_rerank:
|
||||
return inference, None
|
||||
return inference
|
||||
if reconstruction_mode != "continuous":
|
||||
return inference
|
||||
if len(inference.top_candidates) <= 1 or inference.confidence >= solver_params.phase_rerank_confidence_threshold:
|
||||
return inference, None
|
||||
return inference
|
||||
|
||||
top_score = float(inference.top_candidates[0].score)
|
||||
candidate_records: list[dict[str, float | InferenceResult]] = []
|
||||
|
|
@ -222,7 +214,7 @@ def _select_phase_candidate_with_reconstruction(
|
|||
confidence=inference.confidence,
|
||||
top_candidates=inference.top_candidates,
|
||||
)
|
||||
candidate_artifacts, candidate_diagnostics = _run_reconstruction(
|
||||
candidate_artifacts, _candidate_diagnostics = _run_reconstruction(
|
||||
source,
|
||||
inference=candidate_inference,
|
||||
analysis=analysis,
|
||||
|
|
@ -250,12 +242,11 @@ def _select_phase_candidate_with_reconstruction(
|
|||
"edge_concentration": foreground_edge_concentration(candidate_artifacts.target_rgba),
|
||||
"inference_penalty": top_score - float(candidate.score),
|
||||
"size_delta_ratio": _size_delta_ratio(inference, candidate_inference),
|
||||
"cached_reconstruction": (candidate_artifacts, candidate_diagnostics),
|
||||
}
|
||||
)
|
||||
|
||||
if not candidate_records:
|
||||
return inference, None
|
||||
return inference
|
||||
|
||||
support_penalty = _normalize_penalty(record["support_score"] for record in candidate_records)
|
||||
edge_position_penalty = _normalize_penalty(record["edge_position_error"] for record in candidate_records)
|
||||
|
|
@ -273,8 +264,6 @@ def _select_phase_candidate_with_reconstruction(
|
|||
baseline_rank: float | None = None
|
||||
best_rank = float("inf")
|
||||
best_candidate = inference
|
||||
best_cached_reconstruction = None
|
||||
baseline_cached_reconstruction = None
|
||||
annotated_candidates = []
|
||||
for index, record in enumerate(candidate_records):
|
||||
candidate_inference = record["inference"]
|
||||
|
|
@ -310,11 +299,9 @@ def _select_phase_candidate_with_reconstruction(
|
|||
)
|
||||
if baseline_rank is None:
|
||||
baseline_rank = rank
|
||||
baseline_cached_reconstruction = record.get("cached_reconstruction")
|
||||
if rank < best_rank:
|
||||
best_rank = rank
|
||||
best_candidate = record["inference"]
|
||||
best_cached_reconstruction = record.get("cached_reconstruction")
|
||||
annotated_candidates.sort(key=lambda candidate: float(candidate.breakdown.get("phase_rerank_score", float("inf"))))
|
||||
for rank, candidate in enumerate(annotated_candidates, start=1):
|
||||
candidate.breakdown["phase_rerank_rank"] = float(rank)
|
||||
|
|
@ -337,11 +324,9 @@ def _select_phase_candidate_with_reconstruction(
|
|||
confidence=inference.confidence,
|
||||
top_candidates=annotated_candidates,
|
||||
)
|
||||
cached = baseline_cached_reconstruction if _reuse_phase_probe_reconstruction(reconstruction_mode) else None
|
||||
return selected, cached
|
||||
return inference, None
|
||||
cached = best_cached_reconstruction if _reuse_phase_probe_reconstruction(reconstruction_mode) else None
|
||||
return best_candidate, cached
|
||||
return selected
|
||||
return inference
|
||||
return best_candidate
|
||||
|
||||
|
||||
def _resolve_requested_target_dims(
|
||||
|
|
@ -420,43 +405,6 @@ def _run_reconstruction(
|
|||
device=device,
|
||||
solver_params=solver_params,
|
||||
)
|
||||
if reconstruction_mode == "hybrid":
|
||||
geometry_solver = optimize_uv_field(
|
||||
source,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
steps=0,
|
||||
seed=seed,
|
||||
device=device,
|
||||
solver_params=solver_params,
|
||||
)
|
||||
hybrid_solver, hybrid_diagnostics = optimize_tile_graph(
|
||||
source,
|
||||
inference=inference,
|
||||
analysis=analysis,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
device=device,
|
||||
solver_params=solver_params,
|
||||
geometry_reference_rgba=geometry_solver.target_rgba,
|
||||
geometry_guidance_strength=geometry_solver.guidance_strength,
|
||||
)
|
||||
geometry_score = source_lattice_consistency_breakdown(
|
||||
source,
|
||||
geometry_solver.target_rgba,
|
||||
target_width=inference.target_width,
|
||||
target_height=inference.target_height,
|
||||
phase_x=inference.phase_x,
|
||||
phase_y=inference.phase_y,
|
||||
)["score"]
|
||||
diagnostics = dict(hybrid_diagnostics)
|
||||
diagnostics["mode"] = "hybrid"
|
||||
diagnostics["hybrid_geometry_prepass_mode"] = "continuous"
|
||||
diagnostics["hybrid_geometry_prepass_steps"] = 0.0
|
||||
diagnostics["hybrid_geometry_prepass_source_fidelity"] = float(geometry_score)
|
||||
diagnostics["hybrid_geometry_match_weight"] = float(solver_params.hybrid_geometry_match_weight)
|
||||
diagnostics["hybrid_geometry_edge_boost"] = float(solver_params.hybrid_geometry_edge_boost)
|
||||
return hybrid_solver, diagnostics
|
||||
solver = optimize_uv_field(
|
||||
source,
|
||||
inference=inference,
|
||||
|
|
|
|||
|
|
@ -68,8 +68,6 @@ class TileGraphModel:
|
|||
edge_density: float
|
||||
average_choices: float
|
||||
model_device: str
|
||||
geometry_reference_rgba: np.ndarray | None = None
|
||||
geometry_strength: np.ndarray | None = None
|
||||
cache_hit: bool = False
|
||||
|
||||
|
||||
|
|
@ -877,26 +875,10 @@ def build_tile_graph_model(
|
|||
analysis: SourceAnalysis,
|
||||
solver_params: SolverHyperParams | None = None,
|
||||
device: str = "cpu",
|
||||
geometry_reference_rgba: np.ndarray | None = None,
|
||||
geometry_guidance_strength: np.ndarray | None = None,
|
||||
) -> TileGraphModel:
|
||||
torch = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
solver_params = solver_params or SolverHyperParams()
|
||||
geometry_reference_np = None
|
||||
if geometry_reference_rgba is not None:
|
||||
geometry_reference_np = np.asarray(geometry_reference_rgba, dtype=np.float32)
|
||||
if geometry_reference_np.shape[:2] != (inference.target_height, inference.target_width):
|
||||
raise ValueError(
|
||||
"geometry_reference_rgba must match the inferred output size when provided to tile-graph."
|
||||
)
|
||||
geometry_strength_np = None
|
||||
if geometry_guidance_strength is not None:
|
||||
geometry_strength_np = np.asarray(geometry_guidance_strength, dtype=np.float32)
|
||||
if geometry_strength_np.shape != (inference.target_height, inference.target_width):
|
||||
raise ValueError(
|
||||
"geometry_guidance_strength must match the inferred output size when provided to tile-graph."
|
||||
)
|
||||
cache_key = _tile_graph_model_cache_key(
|
||||
source_rgba,
|
||||
inference=inference,
|
||||
|
|
@ -905,12 +887,7 @@ def build_tile_graph_model(
|
|||
)
|
||||
cached_model = _get_cached_tile_graph_model(cache_key)
|
||||
if cached_model is not None:
|
||||
return replace(
|
||||
cached_model,
|
||||
geometry_reference_rgba=geometry_reference_np,
|
||||
geometry_strength=geometry_strength_np,
|
||||
cache_hit=True,
|
||||
)
|
||||
return replace(cached_model, cache_hit=True)
|
||||
source_reference = build_source_lattice_reference(
|
||||
source_rgba,
|
||||
target_width=inference.target_width,
|
||||
|
|
@ -1106,17 +1083,10 @@ def build_tile_graph_model(
|
|||
edge_density=edge_density,
|
||||
average_choices=average_choices,
|
||||
model_device=resolved_device,
|
||||
geometry_reference_rgba=None,
|
||||
geometry_strength=None,
|
||||
cache_hit=False,
|
||||
)
|
||||
_store_cached_tile_graph_model(cache_key, base_model)
|
||||
return replace(
|
||||
base_model,
|
||||
geometry_reference_rgba=geometry_reference_np,
|
||||
geometry_strength=geometry_strength_np,
|
||||
cache_hit=False,
|
||||
)
|
||||
return replace(base_model, cache_hit=False)
|
||||
|
||||
|
||||
def _cell_candidate_span(model: TileGraphModel, y: int, x: int) -> tuple[int, int]:
|
||||
|
|
@ -1175,18 +1145,6 @@ def _tile_graph_unary_cost(model: TileGraphModel, solver_params: SolverHyperPara
|
|||
+ alpha_error * solver_params.tile_graph_alpha_weight
|
||||
+ coverage_error * solver_params.tile_graph_coverage_weight
|
||||
)
|
||||
if model.geometry_reference_rgba is not None:
|
||||
geometry_reference = model.geometry_reference_rgba[y, x]
|
||||
geometry_error = np.mean(np.abs(geometry_reference[None, :] - candidates), axis=-1)
|
||||
geometry_strength = 1.0
|
||||
if model.geometry_strength is not None:
|
||||
geometry_strength += (
|
||||
float(np.clip(model.geometry_strength[y, x], 0.0, 1.0))
|
||||
* solver_params.hybrid_geometry_edge_boost
|
||||
)
|
||||
total_cost = total_cost + geometry_error * (
|
||||
solver_params.hybrid_geometry_match_weight * geometry_strength
|
||||
)
|
||||
unary_cost[start:end] = total_cost.astype(np.float32)
|
||||
return unary_cost
|
||||
|
||||
|
|
@ -1208,17 +1166,6 @@ def _tile_graph_unary_cost_torch(
|
|||
reference_mean_t = torch.from_numpy(model.reference_mean_rgba).to(device=device, dtype=torch.float32)
|
||||
reference_edge_t = torch.from_numpy(model.reference_edge_rgba).to(device=device, dtype=torch.float32)
|
||||
edge_strength_t = torch.from_numpy(model.edge_strength).to(device=device, dtype=torch.float32)
|
||||
geometry_reference_t = (
|
||||
torch.from_numpy(model.geometry_reference_rgba).to(device=device, dtype=torch.float32)
|
||||
if model.geometry_reference_rgba is not None
|
||||
else None
|
||||
)
|
||||
geometry_strength_t = (
|
||||
torch.from_numpy(model.geometry_strength).to(device=device, dtype=torch.float32)
|
||||
if model.geometry_strength is not None
|
||||
else None
|
||||
)
|
||||
|
||||
choice_rgba_t = candidate_rgba_t[choice_indices_t]
|
||||
choice_deltas_t = candidate_deltas_t[choice_indices_t]
|
||||
sharp_error_t = (reference_sharp_t[..., None, :] - choice_rgba_t).abs().mean(dim=-1)
|
||||
|
|
@ -1240,14 +1187,6 @@ def _tile_graph_unary_cost_torch(
|
|||
+ alpha_error * solver_params.tile_graph_alpha_weight
|
||||
+ coverage_error * solver_params.tile_graph_coverage_weight
|
||||
)
|
||||
if geometry_reference_t is not None:
|
||||
geometry_error_t = (geometry_reference_t[..., None, :] - choice_rgba_t).abs().mean(dim=-1)
|
||||
geometry_scale_t = 1.0
|
||||
if geometry_strength_t is not None:
|
||||
geometry_scale_t = 1.0 + geometry_strength_t[..., None].clamp(0.0, 1.0) * solver_params.hybrid_geometry_edge_boost
|
||||
unary_cost_t = unary_cost_t + geometry_error_t * (
|
||||
solver_params.hybrid_geometry_match_weight * geometry_scale_t
|
||||
)
|
||||
unary_cost_t = unary_cost_t.masked_fill(~choice_mask_t, float("inf"))
|
||||
return unary_cost_t, candidate_rgba_t, candidate_deltas_t, choice_rgba_t, choice_deltas_t
|
||||
|
||||
|
|
@ -1420,8 +1359,6 @@ def optimize_tile_graph(
|
|||
seed: int,
|
||||
device: str,
|
||||
solver_params: SolverHyperParams | None = None,
|
||||
geometry_reference_rgba: np.ndarray | None = None,
|
||||
geometry_guidance_strength: np.ndarray | None = None,
|
||||
) -> tuple[SolverArtifacts, dict[str, Any]]:
|
||||
del steps
|
||||
torch = _require_torch()
|
||||
|
|
@ -1436,8 +1373,6 @@ def optimize_tile_graph(
|
|||
analysis=analysis,
|
||||
solver_params=solver_params,
|
||||
device=resolved_device,
|
||||
geometry_reference_rgba=geometry_reference_rgba,
|
||||
geometry_guidance_strength=geometry_guidance_strength,
|
||||
)
|
||||
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)
|
||||
|
|
@ -1561,7 +1496,6 @@ def optimize_tile_graph(
|
|||
"tile_graph_initial_source_fidelity": float(initial_source_fidelity),
|
||||
"tile_graph_final_source_fidelity": float(final_source_fidelity),
|
||||
"tile_graph_kept_initial_assignment": kept_initial_assignment,
|
||||
"tile_graph_geometry_prior": model.geometry_reference_rgba is not None,
|
||||
}
|
||||
return (
|
||||
SolverArtifacts(
|
||||
|
|
|
|||
Loading…
Reference in a new issue