Report lattice search and rerank progress
This commit is contained in:
Meta Vens
parent
99c4f54f7a
commit
a2cf44b3ba
6 changed files with 579 additions and 23 deletions
|
|
@ -41,10 +41,29 @@ type LogItem = {
|
|||
detail: string;
|
||||
};
|
||||
|
||||
type LatticeSearchInfo = {
|
||||
candidateCount: number;
|
||||
completedCandidates: number;
|
||||
phaseSampleCount: number;
|
||||
device: string | null;
|
||||
currentTargetWidth: number | null;
|
||||
currentTargetHeight: number | null;
|
||||
bestScore: number | null;
|
||||
};
|
||||
|
||||
type PhaseRerankInfo = {
|
||||
previewSteps: number;
|
||||
candidateCount: number;
|
||||
confidence: number;
|
||||
completedCandidates: number;
|
||||
currentCandidateIndex: number;
|
||||
currentTargetWidth: number | null;
|
||||
currentTargetHeight: number | null;
|
||||
currentPhaseX: number | null;
|
||||
currentPhaseY: number | null;
|
||||
currentStep: number;
|
||||
currentTotalSteps: number;
|
||||
currentLoss: number | null;
|
||||
};
|
||||
|
||||
type PhaseFieldPrepInfo = {
|
||||
|
|
@ -67,6 +86,7 @@ type AppState = {
|
|||
guidanceImage: ImageAsset | null;
|
||||
inference: InferencePayload | null;
|
||||
inferenceMode: string | null;
|
||||
latticeSearch: LatticeSearchInfo | null;
|
||||
phaseRerank: PhaseRerankInfo | null;
|
||||
phaseFieldPrep: PhaseFieldPrepInfo | null;
|
||||
frames: FramePayload[];
|
||||
|
|
@ -156,8 +176,13 @@ const eventTypes = [
|
|||
"stage_started",
|
||||
"source_loaded",
|
||||
"preprocess_completed",
|
||||
"lattice_search_started",
|
||||
"lattice_search_progress",
|
||||
"inference_candidates_ready",
|
||||
"phase_rerank_started",
|
||||
"phase_rerank_candidate_started",
|
||||
"phase_rerank_candidate_step",
|
||||
"phase_rerank_candidate_completed",
|
||||
"phase_selection_completed",
|
||||
"analysis_completed",
|
||||
"phase_field_prepared",
|
||||
|
|
@ -182,6 +207,7 @@ const state: AppState = {
|
|||
guidanceImage: null,
|
||||
inference: null,
|
||||
inferenceMode: null,
|
||||
latticeSearch: null,
|
||||
phaseRerank: null,
|
||||
phaseFieldPrep: null,
|
||||
frames: [],
|
||||
|
|
@ -285,6 +311,7 @@ function resetRunArtifacts(options: { preserveSourceImage?: boolean } = {}): voi
|
|||
state.guidanceImage = null;
|
||||
state.inference = null;
|
||||
state.inferenceMode = null;
|
||||
state.latticeSearch = null;
|
||||
state.phaseRerank = null;
|
||||
state.phaseFieldPrep = null;
|
||||
state.frames = [];
|
||||
|
|
@ -439,7 +466,21 @@ function renderStatusMetrics(): void {
|
|||
statusMetrics.innerHTML = "";
|
||||
const items: Array<[string, string]> = [];
|
||||
const frame = getSelectedFrame();
|
||||
if (state.stageKey === "inference" || state.stageKey === "analysis" || state.stageKey === "selection") {
|
||||
if (state.stageKey === "inference") {
|
||||
if (state.latticeSearch) {
|
||||
items.push(["Sizes", `${state.latticeSearch.completedCandidates} / ${state.latticeSearch.candidateCount}`]);
|
||||
items.push(["Phase samples", String(state.latticeSearch.phaseSampleCount)]);
|
||||
if (state.latticeSearch.currentTargetWidth && state.latticeSearch.currentTargetHeight) {
|
||||
items.push(["Current grid", `${state.latticeSearch.currentTargetWidth} x ${state.latticeSearch.currentTargetHeight}`]);
|
||||
}
|
||||
if (state.latticeSearch.device) {
|
||||
items.push(["Device", state.latticeSearch.device]);
|
||||
}
|
||||
if (state.latticeSearch.bestScore !== null) {
|
||||
items.push(["Best score", formatNumber(state.latticeSearch.bestScore, 3)]);
|
||||
}
|
||||
}
|
||||
} else if (state.stageKey === "analysis" || state.stageKey === "selection") {
|
||||
if (state.inference) {
|
||||
items.push(["Mode", state.inferenceMode === "fixed" ? "pinned" : state.inferenceMode ?? "search"]);
|
||||
items.push(["Grid", `${state.inference.target_width} x ${state.inference.target_height}`]);
|
||||
|
|
@ -448,8 +489,25 @@ function renderStatusMetrics(): void {
|
|||
}
|
||||
} else if (state.stageKey === "rerank") {
|
||||
if (state.phaseRerank) {
|
||||
items.push(["Preview steps", String(state.phaseRerank.previewSteps)]);
|
||||
items.push(["Candidates", String(state.phaseRerank.candidateCount)]);
|
||||
items.push(["Candidates", `${state.phaseRerank.completedCandidates} / ${state.phaseRerank.candidateCount}`]);
|
||||
if (state.phaseRerank.currentCandidateIndex > 0) {
|
||||
items.push(["Active candidate", `${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}`]);
|
||||
}
|
||||
items.push([
|
||||
"Preview step",
|
||||
state.phaseRerank.currentTotalSteps > 0
|
||||
? `${state.phaseRerank.currentStep} / ${state.phaseRerank.currentTotalSteps}`
|
||||
: String(state.phaseRerank.previewSteps),
|
||||
]);
|
||||
if (state.phaseRerank.currentTargetWidth && state.phaseRerank.currentTargetHeight) {
|
||||
items.push(["Grid", `${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight}`]);
|
||||
}
|
||||
if (state.phaseRerank.currentPhaseX !== null && state.phaseRerank.currentPhaseY !== null) {
|
||||
items.push(["Phase", `${formatNumber(state.phaseRerank.currentPhaseX, 2)}, ${formatNumber(state.phaseRerank.currentPhaseY, 2)}`]);
|
||||
}
|
||||
if (state.phaseRerank.currentLoss !== null) {
|
||||
items.push(["Loss", formatNumber(state.phaseRerank.currentLoss, 4)]);
|
||||
}
|
||||
items.push(["Confidence", formatNumber(state.phaseRerank.confidence, 3)]);
|
||||
}
|
||||
} else if (state.stageKey === "solver") {
|
||||
|
|
@ -892,6 +950,48 @@ async function handleEvent(eventName: string, payload: JsonRecord): Promise<void
|
|||
state.preprocessedImage = payload.sourceImage as ImageAsset;
|
||||
addLog("Preprocess", "Stripped edge-connected background noise.");
|
||||
break;
|
||||
case "lattice_search_started":
|
||||
state.latticeSearch = {
|
||||
candidateCount: Number(payload.candidateCount ?? 0),
|
||||
completedCandidates: 0,
|
||||
phaseSampleCount: Number(payload.phaseSampleCount ?? 0),
|
||||
device: typeof payload.device === "string" ? payload.device : null,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
bestScore: null,
|
||||
};
|
||||
setStage(
|
||||
"inference",
|
||||
"Lattice search",
|
||||
`Testing ${state.latticeSearch.candidateCount} size candidates with ${state.latticeSearch.phaseSampleCount} phase samples each.`,
|
||||
);
|
||||
addLog("Search", `Started lattice search across ${String(payload.candidateCount)} size candidates.`);
|
||||
break;
|
||||
case "lattice_search_progress":
|
||||
if (!state.latticeSearch) {
|
||||
state.latticeSearch = {
|
||||
candidateCount: Number(payload.totalCandidates ?? 0),
|
||||
completedCandidates: 0,
|
||||
phaseSampleCount: Number(payload.phaseSampleCount ?? 0),
|
||||
device: null,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
bestScore: null,
|
||||
};
|
||||
}
|
||||
state.latticeSearch.completedCandidates = Number(payload.completedCandidates ?? state.latticeSearch.completedCandidates);
|
||||
state.latticeSearch.candidateCount = Number(payload.totalCandidates ?? state.latticeSearch.candidateCount);
|
||||
state.latticeSearch.phaseSampleCount = Number(payload.phaseSampleCount ?? state.latticeSearch.phaseSampleCount);
|
||||
state.latticeSearch.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.latticeSearch.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.latticeSearch.bestScore =
|
||||
payload.bestScore === null || payload.bestScore === undefined ? state.latticeSearch.bestScore : Number(payload.bestScore);
|
||||
setStage(
|
||||
"inference",
|
||||
"Lattice search",
|
||||
`Size ${state.latticeSearch.completedCandidates} / ${state.latticeSearch.candidateCount}: scored ${state.latticeSearch.currentTargetWidth} x ${state.latticeSearch.currentTargetHeight} over ${state.latticeSearch.phaseSampleCount} phase samples.`,
|
||||
);
|
||||
break;
|
||||
case "inference_candidates_ready":
|
||||
state.inference = payload.inference as InferencePayload;
|
||||
state.inferenceMode = typeof payload.inferenceMode === "string" ? payload.inferenceMode : state.inferenceMode;
|
||||
|
|
@ -902,6 +1002,15 @@ async function handleEvent(eventName: string, payload: JsonRecord): Promise<void
|
|||
previewSteps: Number(payload.previewSteps ?? 0),
|
||||
candidateCount: Number(payload.candidateCount ?? 0),
|
||||
confidence: Number(payload.confidence ?? 0),
|
||||
completedCandidates: 0,
|
||||
currentCandidateIndex: 0,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
currentPhaseX: null,
|
||||
currentPhaseY: null,
|
||||
currentStep: 0,
|
||||
currentTotalSteps: 0,
|
||||
currentLoss: null,
|
||||
};
|
||||
setStage(
|
||||
"rerank",
|
||||
|
|
@ -913,6 +1022,62 @@ async function handleEvent(eventName: string, payload: JsonRecord): Promise<void
|
|||
`Running ${String(payload.previewSteps)} preview steps across ${String(payload.candidateCount)} low-confidence candidates.`,
|
||||
);
|
||||
break;
|
||||
case "phase_rerank_candidate_started":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.currentStep = 0;
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.previewSteps ?? state.phaseRerank.previewSteps);
|
||||
state.phaseRerank.currentLoss = null;
|
||||
setStage(
|
||||
"rerank",
|
||||
"Phase rerank",
|
||||
`Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}: previewing ${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight} at phase ${formatNumber(state.phaseRerank.currentPhaseX, 2)}, ${formatNumber(state.phaseRerank.currentPhaseY, 2)}.`,
|
||||
);
|
||||
}
|
||||
break;
|
||||
case "phase_rerank_candidate_step":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.currentStep = Number(payload.step ?? 0);
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.totalSteps ?? 0);
|
||||
state.phaseRerank.currentLoss =
|
||||
payload.loss === null || payload.loss === undefined ? null : Number(payload.loss);
|
||||
setStage(
|
||||
"rerank",
|
||||
"Phase rerank",
|
||||
state.phaseRerank.currentTotalSteps <= 0
|
||||
? `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}: preview solve skipped.`
|
||||
: `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}, preview step ${state.phaseRerank.currentStep} / ${state.phaseRerank.currentTotalSteps} on ${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight}.`,
|
||||
);
|
||||
}
|
||||
break;
|
||||
case "phase_rerank_candidate_completed":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.completedCandidates = Number(payload.completedCandidates ?? state.phaseRerank.completedCandidates);
|
||||
state.phaseRerank.currentStep = Number(payload.totalSteps ?? state.phaseRerank.currentStep);
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.totalSteps ?? state.phaseRerank.currentTotalSteps);
|
||||
state.phaseRerank.currentLoss =
|
||||
payload.finalLoss === null || payload.finalLoss === undefined ? state.phaseRerank.currentLoss : Number(payload.finalLoss);
|
||||
setStage(
|
||||
"rerank",
|
||||
"Phase rerank",
|
||||
`Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount} scored. ${state.phaseRerank.completedCandidates} of ${state.phaseRerank.candidateCount} previews finished.`,
|
||||
);
|
||||
}
|
||||
break;
|
||||
case "phase_selection_completed":
|
||||
state.inference = payload.inference as InferencePayload;
|
||||
state.inferenceMode = typeof payload.inferenceMode === "string" ? payload.inferenceMode : state.inferenceMode;
|
||||
|
|
|
|||
|
|
@ -208,6 +208,55 @@ class GuiJob:
|
|||
"candidateCount": int(payload["candidate_count"]),
|
||||
"confidence": float(payload["confidence"]),
|
||||
}
|
||||
if event == "lattice_search_started":
|
||||
return {
|
||||
"candidateCount": int(payload["candidate_count"]),
|
||||
"phaseSampleCount": int(payload["phase_sample_count"]),
|
||||
"device": str(payload["device"]),
|
||||
}
|
||||
if event == "lattice_search_progress":
|
||||
return {
|
||||
"completedCandidates": int(payload["completed_candidates"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseSampleCount": int(payload["phase_sample_count"]),
|
||||
"bestScore": None if payload.get("best_score") is None else float(payload["best_score"]),
|
||||
}
|
||||
if event == "phase_rerank_candidate_started":
|
||||
return {
|
||||
"candidateIndex": int(payload["candidate_index"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseX": float(payload["phase_x"]),
|
||||
"phaseY": float(payload["phase_y"]),
|
||||
"previewSteps": int(payload["preview_steps"]),
|
||||
}
|
||||
if event == "phase_rerank_candidate_step":
|
||||
return {
|
||||
"candidateIndex": int(payload["candidate_index"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseX": float(payload["phase_x"]),
|
||||
"phaseY": float(payload["phase_y"]),
|
||||
"step": int(payload["step"]),
|
||||
"totalSteps": int(payload["total_steps"]),
|
||||
"loss": None if payload.get("loss") is None else float(payload["loss"]),
|
||||
}
|
||||
if event == "phase_rerank_candidate_completed":
|
||||
return {
|
||||
"candidateIndex": int(payload["candidate_index"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"completedCandidates": int(payload["completed_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseX": float(payload["phase_x"]),
|
||||
"phaseY": float(payload["phase_y"]),
|
||||
"totalSteps": int(payload["total_steps"]),
|
||||
"finalLoss": None if payload.get("final_loss") is None else float(payload["final_loss"]),
|
||||
}
|
||||
if event == "stage_started":
|
||||
return {
|
||||
"stage": str(payload["stage"]),
|
||||
|
|
|
|||
|
|
@ -61,8 +61,13 @@ const eventTypes = [
|
|||
"stage_started",
|
||||
"source_loaded",
|
||||
"preprocess_completed",
|
||||
"lattice_search_started",
|
||||
"lattice_search_progress",
|
||||
"inference_candidates_ready",
|
||||
"phase_rerank_started",
|
||||
"phase_rerank_candidate_started",
|
||||
"phase_rerank_candidate_step",
|
||||
"phase_rerank_candidate_completed",
|
||||
"phase_selection_completed",
|
||||
"analysis_completed",
|
||||
"phase_field_prepared",
|
||||
|
|
@ -86,6 +91,7 @@ const state = {
|
|||
guidanceImage: null,
|
||||
inference: null,
|
||||
inferenceMode: null,
|
||||
latticeSearch: null,
|
||||
phaseRerank: null,
|
||||
phaseFieldPrep: null,
|
||||
frames: [],
|
||||
|
|
@ -178,6 +184,7 @@ function resetRunArtifacts(options = {}) {
|
|||
state.guidanceImage = null;
|
||||
state.inference = null;
|
||||
state.inferenceMode = null;
|
||||
state.latticeSearch = null;
|
||||
state.phaseRerank = null;
|
||||
state.phaseFieldPrep = null;
|
||||
state.frames = [];
|
||||
|
|
@ -322,7 +329,22 @@ function renderStatusMetrics() {
|
|||
statusMetrics.innerHTML = "";
|
||||
const items = [];
|
||||
const frame = getSelectedFrame();
|
||||
if (state.stageKey === "inference" || state.stageKey === "analysis" || state.stageKey === "selection") {
|
||||
if (state.stageKey === "inference") {
|
||||
if (state.latticeSearch) {
|
||||
items.push(["Sizes", `${state.latticeSearch.completedCandidates} / ${state.latticeSearch.candidateCount}`]);
|
||||
items.push(["Phase samples", String(state.latticeSearch.phaseSampleCount)]);
|
||||
if (state.latticeSearch.currentTargetWidth && state.latticeSearch.currentTargetHeight) {
|
||||
items.push(["Current grid", `${state.latticeSearch.currentTargetWidth} x ${state.latticeSearch.currentTargetHeight}`]);
|
||||
}
|
||||
if (state.latticeSearch.device) {
|
||||
items.push(["Device", state.latticeSearch.device]);
|
||||
}
|
||||
if (state.latticeSearch.bestScore !== null) {
|
||||
items.push(["Best score", formatNumber(state.latticeSearch.bestScore, 3)]);
|
||||
}
|
||||
}
|
||||
}
|
||||
else if (state.stageKey === "analysis" || state.stageKey === "selection") {
|
||||
if (state.inference) {
|
||||
items.push(["Mode", state.inferenceMode === "fixed" ? "pinned" : state.inferenceMode ?? "search"]);
|
||||
items.push(["Grid", `${state.inference.target_width} x ${state.inference.target_height}`]);
|
||||
|
|
@ -332,8 +354,25 @@ function renderStatusMetrics() {
|
|||
}
|
||||
else if (state.stageKey === "rerank") {
|
||||
if (state.phaseRerank) {
|
||||
items.push(["Preview steps", String(state.phaseRerank.previewSteps)]);
|
||||
items.push(["Candidates", String(state.phaseRerank.candidateCount)]);
|
||||
items.push(["Candidates", `${state.phaseRerank.completedCandidates} / ${state.phaseRerank.candidateCount}`]);
|
||||
if (state.phaseRerank.currentCandidateIndex > 0) {
|
||||
items.push(["Active candidate", `${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}`]);
|
||||
}
|
||||
items.push([
|
||||
"Preview step",
|
||||
state.phaseRerank.currentTotalSteps > 0
|
||||
? `${state.phaseRerank.currentStep} / ${state.phaseRerank.currentTotalSteps}`
|
||||
: String(state.phaseRerank.previewSteps),
|
||||
]);
|
||||
if (state.phaseRerank.currentTargetWidth && state.phaseRerank.currentTargetHeight) {
|
||||
items.push(["Grid", `${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight}`]);
|
||||
}
|
||||
if (state.phaseRerank.currentPhaseX !== null && state.phaseRerank.currentPhaseY !== null) {
|
||||
items.push(["Phase", `${formatNumber(state.phaseRerank.currentPhaseX, 2)}, ${formatNumber(state.phaseRerank.currentPhaseY, 2)}`]);
|
||||
}
|
||||
if (state.phaseRerank.currentLoss !== null) {
|
||||
items.push(["Loss", formatNumber(state.phaseRerank.currentLoss, 4)]);
|
||||
}
|
||||
items.push(["Confidence", formatNumber(state.phaseRerank.confidence, 3)]);
|
||||
}
|
||||
}
|
||||
|
|
@ -745,6 +784,40 @@ async function handleEvent(eventName, payload) {
|
|||
state.preprocessedImage = payload.sourceImage;
|
||||
addLog("Preprocess", "Stripped edge-connected background noise.");
|
||||
break;
|
||||
case "lattice_search_started":
|
||||
state.latticeSearch = {
|
||||
candidateCount: Number(payload.candidateCount ?? 0),
|
||||
completedCandidates: 0,
|
||||
phaseSampleCount: Number(payload.phaseSampleCount ?? 0),
|
||||
device: typeof payload.device === "string" ? payload.device : null,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
bestScore: null,
|
||||
};
|
||||
setStage("inference", "Lattice search", `Testing ${state.latticeSearch.candidateCount} size candidates with ${state.latticeSearch.phaseSampleCount} phase samples each.`);
|
||||
addLog("Search", `Started lattice search across ${String(payload.candidateCount)} size candidates.`);
|
||||
break;
|
||||
case "lattice_search_progress":
|
||||
if (!state.latticeSearch) {
|
||||
state.latticeSearch = {
|
||||
candidateCount: Number(payload.totalCandidates ?? 0),
|
||||
completedCandidates: 0,
|
||||
phaseSampleCount: Number(payload.phaseSampleCount ?? 0),
|
||||
device: null,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
bestScore: null,
|
||||
};
|
||||
}
|
||||
state.latticeSearch.completedCandidates = Number(payload.completedCandidates ?? state.latticeSearch.completedCandidates);
|
||||
state.latticeSearch.candidateCount = Number(payload.totalCandidates ?? state.latticeSearch.candidateCount);
|
||||
state.latticeSearch.phaseSampleCount = Number(payload.phaseSampleCount ?? state.latticeSearch.phaseSampleCount);
|
||||
state.latticeSearch.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.latticeSearch.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.latticeSearch.bestScore =
|
||||
payload.bestScore === null || payload.bestScore === undefined ? state.latticeSearch.bestScore : Number(payload.bestScore);
|
||||
setStage("inference", "Lattice search", `Size ${state.latticeSearch.completedCandidates} / ${state.latticeSearch.candidateCount}: scored ${state.latticeSearch.currentTargetWidth} x ${state.latticeSearch.currentTargetHeight} over ${state.latticeSearch.phaseSampleCount} phase samples.`);
|
||||
break;
|
||||
case "inference_candidates_ready":
|
||||
state.inference = payload.inference;
|
||||
state.inferenceMode = typeof payload.inferenceMode === "string" ? payload.inferenceMode : state.inferenceMode;
|
||||
|
|
@ -755,10 +828,63 @@ async function handleEvent(eventName, payload) {
|
|||
previewSteps: Number(payload.previewSteps ?? 0),
|
||||
candidateCount: Number(payload.candidateCount ?? 0),
|
||||
confidence: Number(payload.confidence ?? 0),
|
||||
completedCandidates: 0,
|
||||
currentCandidateIndex: 0,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
currentPhaseX: null,
|
||||
currentPhaseY: null,
|
||||
currentStep: 0,
|
||||
currentTotalSteps: 0,
|
||||
currentLoss: null,
|
||||
};
|
||||
setStage("rerank", "Phase rerank", `Previewing ${state.phaseRerank.candidateCount} low-confidence candidates over ${state.phaseRerank.previewSteps} short solver steps.`);
|
||||
addLog("Rerank", `Running ${String(payload.previewSteps)} preview steps across ${String(payload.candidateCount)} low-confidence candidates.`);
|
||||
break;
|
||||
case "phase_rerank_candidate_started":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.currentStep = 0;
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.previewSteps ?? state.phaseRerank.previewSteps);
|
||||
state.phaseRerank.currentLoss = null;
|
||||
setStage("rerank", "Phase rerank", `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}: previewing ${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight} at phase ${formatNumber(state.phaseRerank.currentPhaseX, 2)}, ${formatNumber(state.phaseRerank.currentPhaseY, 2)}.`);
|
||||
}
|
||||
break;
|
||||
case "phase_rerank_candidate_step":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.currentStep = Number(payload.step ?? 0);
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.totalSteps ?? 0);
|
||||
state.phaseRerank.currentLoss =
|
||||
payload.loss === null || payload.loss === undefined ? null : Number(payload.loss);
|
||||
setStage("rerank", "Phase rerank", state.phaseRerank.currentTotalSteps <= 0
|
||||
? `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}: preview solve skipped.`
|
||||
: `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}, preview step ${state.phaseRerank.currentStep} / ${state.phaseRerank.currentTotalSteps} on ${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight}.`);
|
||||
}
|
||||
break;
|
||||
case "phase_rerank_candidate_completed":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.completedCandidates = Number(payload.completedCandidates ?? state.phaseRerank.completedCandidates);
|
||||
state.phaseRerank.currentStep = Number(payload.totalSteps ?? state.phaseRerank.currentStep);
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.totalSteps ?? state.phaseRerank.currentTotalSteps);
|
||||
state.phaseRerank.currentLoss =
|
||||
payload.finalLoss === null || payload.finalLoss === undefined ? state.phaseRerank.currentLoss : Number(payload.finalLoss);
|
||||
setStage("rerank", "Phase rerank", `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount} scored. ${state.phaseRerank.completedCandidates} of ${state.phaseRerank.candidateCount} previews finished.`);
|
||||
}
|
||||
break;
|
||||
case "phase_selection_completed":
|
||||
state.inference = payload.inference;
|
||||
state.inferenceMode = typeof payload.inferenceMode === "string" ? payload.inferenceMode : state.inferenceMode;
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import numpy as np
|
|||
|
||||
from .io import premultiply
|
||||
from .metrics import source_lattice_evidence_breakdown
|
||||
from .observe import PipelineObserver, emit_observer
|
||||
from .types import InferenceCandidate, InferenceResult
|
||||
|
||||
|
||||
|
|
@ -471,7 +472,12 @@ def _score_phase_group(
|
|||
return candidates
|
||||
|
||||
|
||||
def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str = "auto") -> InferenceResult:
|
||||
def infer_lattice(
|
||||
rgba: np.ndarray,
|
||||
target_size: int | None = None,
|
||||
device: str = "auto",
|
||||
observer: PipelineObserver | None = None,
|
||||
) -> InferenceResult:
|
||||
torch, _ = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
height, width = rgba.shape[:2]
|
||||
|
|
@ -479,21 +485,40 @@ def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str
|
|||
hinted_sizes = _hint_target_sizes_from_spacing(width, height, spacing_x, spacing_y)
|
||||
prior_cell_x, prior_cell_y, prior_reliability = _estimate_lattice_prior_details(rgba)
|
||||
phase_values = np.linspace(-0.4, 0.4, num=5, dtype=np.float32)
|
||||
candidate_dims = _candidate_dims(width, height, target_size, hinted_sizes=hinted_sizes)
|
||||
phase_sample_count = int(phase_values.size * phase_values.size)
|
||||
|
||||
emit_observer(
|
||||
observer,
|
||||
"lattice_search_started",
|
||||
candidate_count=len(candidate_dims),
|
||||
phase_sample_count=phase_sample_count,
|
||||
device=resolved_device,
|
||||
)
|
||||
|
||||
candidates: list[InferenceCandidate] = []
|
||||
for target_width, target_height in _candidate_dims(width, height, target_size, hinted_sizes=hinted_sizes):
|
||||
candidates.extend(
|
||||
_score_phase_group(
|
||||
rgba,
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
prior_cell_x=prior_cell_x,
|
||||
prior_cell_y=prior_cell_y,
|
||||
prior_reliability=prior_reliability,
|
||||
phase_x_values=phase_values,
|
||||
phase_y_values=phase_values,
|
||||
device=resolved_device,
|
||||
)
|
||||
for candidate_index, (target_width, target_height) in enumerate(candidate_dims, start=1):
|
||||
scored_group = _score_phase_group(
|
||||
rgba,
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
prior_cell_x=prior_cell_x,
|
||||
prior_cell_y=prior_cell_y,
|
||||
prior_reliability=prior_reliability,
|
||||
phase_x_values=phase_values,
|
||||
phase_y_values=phase_values,
|
||||
device=resolved_device,
|
||||
)
|
||||
candidates.extend(scored_group)
|
||||
emit_observer(
|
||||
observer,
|
||||
"lattice_search_progress",
|
||||
completed_candidates=candidate_index,
|
||||
total_candidates=len(candidate_dims),
|
||||
target_width=int(target_width),
|
||||
target_height=int(target_height),
|
||||
phase_sample_count=phase_sample_count,
|
||||
best_score=None if not scored_group else float(max(candidate.score for candidate in scored_group)),
|
||||
)
|
||||
|
||||
candidates.sort(key=lambda item: item.score, reverse=True)
|
||||
|
|
|
|||
|
|
@ -125,7 +125,7 @@ def run_pipeline_rgba(
|
|||
label="Lattice search",
|
||||
detail="Searching for the output size and phase that best fit the input.",
|
||||
)
|
||||
inference = infer_lattice(source, target_size=target_size, device=device)
|
||||
inference = infer_lattice(source, target_size=target_size, device=device, observer=observer)
|
||||
inference_mode = "searched"
|
||||
else:
|
||||
emit_observer(
|
||||
|
|
@ -373,8 +373,9 @@ def _select_phase_candidate_with_reconstruction(
|
|||
candidate_count=min(8, len(inference.top_candidates)),
|
||||
confidence=float(inference.confidence),
|
||||
)
|
||||
total_candidates = min(8, len(inference.top_candidates))
|
||||
candidate_records: list[dict[str, float | InferenceResult]] = []
|
||||
for candidate in inference.top_candidates[:8]:
|
||||
for candidate_index, candidate in enumerate(inference.top_candidates[:total_candidates], start=1):
|
||||
candidate_inference = InferenceResult(
|
||||
target_width=candidate.target_width,
|
||||
target_height=candidate.target_height,
|
||||
|
|
@ -383,6 +384,39 @@ def _select_phase_candidate_with_reconstruction(
|
|||
confidence=inference.confidence,
|
||||
top_candidates=inference.top_candidates,
|
||||
)
|
||||
emit_observer(
|
||||
observer,
|
||||
"phase_rerank_candidate_started",
|
||||
candidate_index=candidate_index,
|
||||
total_candidates=total_candidates,
|
||||
target_width=int(candidate.target_width),
|
||||
target_height=int(candidate.target_height),
|
||||
phase_x=float(candidate.phase_x),
|
||||
phase_y=float(candidate.phase_y),
|
||||
preview_steps=preview_steps,
|
||||
)
|
||||
|
||||
def rerank_observer(event: str, payload: dict[str, Any]) -> None:
|
||||
if observer is None:
|
||||
return
|
||||
common = {
|
||||
"candidate_index": candidate_index,
|
||||
"total_candidates": total_candidates,
|
||||
"target_width": int(candidate.target_width),
|
||||
"target_height": int(candidate.target_height),
|
||||
"phase_x": float(candidate.phase_x),
|
||||
"phase_y": float(candidate.phase_y),
|
||||
}
|
||||
if event in {"phase_field_initial", "phase_field_step"}:
|
||||
emit_observer(
|
||||
observer,
|
||||
"phase_rerank_candidate_step",
|
||||
step=int(payload["step"]),
|
||||
total_steps=int(payload["total_steps"]),
|
||||
loss=None if payload.get("loss") is None else float(payload["loss"]),
|
||||
**common,
|
||||
)
|
||||
|
||||
candidate_artifacts, _candidate_diagnostics = _run_reconstruction(
|
||||
source,
|
||||
inference=candidate_inference,
|
||||
|
|
@ -391,7 +425,22 @@ def _select_phase_candidate_with_reconstruction(
|
|||
seed=seed,
|
||||
device=device,
|
||||
solver_params=solver_params,
|
||||
observer=None,
|
||||
observer=rerank_observer if observer is not None else None,
|
||||
)
|
||||
emit_observer(
|
||||
observer,
|
||||
"phase_rerank_candidate_completed",
|
||||
candidate_index=candidate_index,
|
||||
total_candidates=total_candidates,
|
||||
completed_candidates=candidate_index,
|
||||
target_width=int(candidate.target_width),
|
||||
target_height=int(candidate.target_height),
|
||||
phase_x=float(candidate.phase_x),
|
||||
phase_y=float(candidate.phase_y),
|
||||
total_steps=preview_steps,
|
||||
final_loss=(
|
||||
None if not candidate_artifacts.loss_history else float(candidate_artifacts.loss_history[-1])
|
||||
),
|
||||
)
|
||||
support = source_lattice_consistency_breakdown(
|
||||
source,
|
||||
|
|
|
|||
|
|
@ -4,11 +4,16 @@ import importlib.util
|
|||
import io
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from starlette.datastructures import UploadFile
|
||||
|
||||
import repixelizer.inference as inference_module
|
||||
import repixelizer.pipeline as pipeline_module
|
||||
from repixelizer.pipeline import run_pipeline_rgba
|
||||
from repixelizer.synthetic import fake_pixelize, make_emblem
|
||||
from repixelizer.params import SolverHyperParams
|
||||
from repixelizer.types import InferenceCandidate, InferenceResult, PhaseFieldSourceAnalysis, SolverArtifacts
|
||||
|
||||
|
||||
def test_run_pipeline_rgba_emits_observer_events_for_gui() -> None:
|
||||
|
|
@ -81,3 +86,140 @@ def test_create_job_upload_field_validates_without_forward_ref_errors() -> None:
|
|||
|
||||
assert errors == []
|
||||
assert value.filename == "tiny.png"
|
||||
|
||||
|
||||
def test_infer_lattice_emits_search_progress_events(monkeypatch) -> None:
|
||||
class _FakeCuda:
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
return False
|
||||
|
||||
class _FakeTorch:
|
||||
cuda = _FakeCuda()
|
||||
|
||||
monkeypatch.setattr(inference_module, "_require_torch", lambda: (_FakeTorch(), object()))
|
||||
monkeypatch.setattr(inference_module, "_estimate_lattice_spacing", lambda rgba: ((None, 0.0), (None, 0.0)))
|
||||
monkeypatch.setattr(inference_module, "_hint_target_sizes_from_spacing", lambda width, height, spacing_x, spacing_y: [])
|
||||
monkeypatch.setattr(inference_module, "_estimate_lattice_prior_details", lambda rgba: (4.0, 4.0, 0.5))
|
||||
monkeypatch.setattr(
|
||||
inference_module,
|
||||
"_candidate_dims",
|
||||
lambda width, height, target_size, *, hinted_sizes=None: [(10, 8), (12, 10)],
|
||||
)
|
||||
|
||||
def fake_score_phase_group(
|
||||
rgba,
|
||||
*,
|
||||
target_width,
|
||||
target_height,
|
||||
prior_cell_x,
|
||||
prior_cell_y,
|
||||
prior_reliability,
|
||||
phase_x_values,
|
||||
phase_y_values,
|
||||
device,
|
||||
):
|
||||
del rgba, prior_cell_x, prior_cell_y, prior_reliability, phase_x_values, phase_y_values, device
|
||||
return [
|
||||
InferenceCandidate(
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
phase_x=0.0,
|
||||
phase_y=0.0,
|
||||
score=0.9 if target_width == 12 else 0.6,
|
||||
breakdown={},
|
||||
)
|
||||
]
|
||||
|
||||
monkeypatch.setattr(inference_module, "_score_phase_group", fake_score_phase_group)
|
||||
monkeypatch.setattr(inference_module, "_top_candidates_by_size", lambda candidates, limit: candidates[:limit])
|
||||
monkeypatch.setattr(
|
||||
inference_module,
|
||||
"_rerank_size_candidates_with_source_evidence",
|
||||
lambda rgba, candidates: sorted(candidates, key=lambda candidate: candidate.score, reverse=True),
|
||||
)
|
||||
|
||||
events: list[tuple[str, dict[str, object]]] = []
|
||||
|
||||
def observer(event: str, payload: dict[str, object]) -> None:
|
||||
events.append((event, payload))
|
||||
|
||||
result = inference_module.infer_lattice(np.zeros((16, 20, 4), dtype=np.uint8), observer=observer)
|
||||
|
||||
assert result.target_width == 12
|
||||
assert [event for event, _payload in events] == [
|
||||
"lattice_search_started",
|
||||
"lattice_search_progress",
|
||||
"lattice_search_progress",
|
||||
]
|
||||
assert events[0][1]["candidate_count"] == 2
|
||||
assert events[1][1]["completed_candidates"] == 1
|
||||
assert events[2][1]["completed_candidates"] == 2
|
||||
assert events[2][1]["target_width"] == 12
|
||||
|
||||
|
||||
def test_phase_rerank_emits_candidate_progress_events(monkeypatch) -> None:
|
||||
inference = InferenceResult(
|
||||
target_width=16,
|
||||
target_height=16,
|
||||
phase_x=0.0,
|
||||
phase_y=0.0,
|
||||
confidence=0.0,
|
||||
top_candidates=[
|
||||
InferenceCandidate(target_width=16, target_height=16, phase_x=0.0, phase_y=0.0, score=0.75, breakdown={}),
|
||||
InferenceCandidate(target_width=18, target_height=18, phase_x=0.1, phase_y=-0.1, score=0.73, breakdown={}),
|
||||
],
|
||||
)
|
||||
analysis = PhaseFieldSourceAnalysis(edge_map=np.zeros((4, 4), dtype=np.float32))
|
||||
|
||||
def fake_run_reconstruction(source, *, observer=None, inference, **kwargs):
|
||||
del source, kwargs
|
||||
if observer is not None:
|
||||
observer("phase_field_initial", {"step": 0, "total_steps": 2})
|
||||
observer("phase_field_step", {"step": 1, "total_steps": 2, "loss": 0.25})
|
||||
observer("phase_field_step", {"step": 2, "total_steps": 2, "loss": 0.125})
|
||||
target_rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.uint8)
|
||||
artifacts = SolverArtifacts(
|
||||
target_rgba=target_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=target_rgba.copy(),
|
||||
loss_history=[0.25, 0.125],
|
||||
)
|
||||
return artifacts, {}
|
||||
|
||||
monkeypatch.setattr(pipeline_module, "_run_reconstruction", fake_run_reconstruction)
|
||||
monkeypatch.setattr(pipeline_module, "source_lattice_consistency_breakdown", lambda *args, **kwargs: {"score": 0.8})
|
||||
monkeypatch.setattr(
|
||||
pipeline_module,
|
||||
"nearest_resize",
|
||||
lambda rgba, width, height: np.zeros((height, width, 4), dtype=np.uint8),
|
||||
)
|
||||
monkeypatch.setattr(pipeline_module, "foreground_edge_position_error", lambda preview, source: 0.12)
|
||||
monkeypatch.setattr(pipeline_module, "foreground_stroke_wobble_error", lambda preview, source: 0.08)
|
||||
monkeypatch.setattr(pipeline_module, "foreground_edge_concentration", lambda rgba: 0.9)
|
||||
|
||||
events: list[tuple[str, dict[str, object]]] = []
|
||||
|
||||
def observer(event: str, payload: dict[str, object]) -> None:
|
||||
events.append((event, payload))
|
||||
|
||||
pipeline_module._select_phase_candidate_with_reconstruction(
|
||||
np.zeros((16, 16, 4), dtype=np.uint8),
|
||||
inference,
|
||||
analysis=analysis,
|
||||
steps=2,
|
||||
seed=1,
|
||||
device="cpu",
|
||||
solver_params=SolverHyperParams(phase_rerank_preview_steps=2, phase_rerank_confidence_threshold=1.0),
|
||||
observer=observer,
|
||||
)
|
||||
|
||||
event_names = [event for event, _payload in events]
|
||||
assert event_names[0] == "phase_rerank_started"
|
||||
assert event_names.count("phase_rerank_candidate_started") == 2
|
||||
assert event_names.count("phase_rerank_candidate_step") == 6
|
||||
assert event_names.count("phase_rerank_candidate_completed") == 2
|
||||
first_step_payload = next(payload for event, payload in events if event == "phase_rerank_candidate_step")
|
||||
assert first_step_payload["candidate_index"] == 1
|
||||
assert first_step_payload["total_steps"] == 2
|
||||
|
|
|
|||
Loading…
Reference in a new issue