diff --git a/frontend/src/recognition/ensemble.ts b/frontend/src/recognition/ensemble.ts new file mode 100644 index 0000000..ea16b41 --- /dev/null +++ b/frontend/src/recognition/ensemble.ts @@ -0,0 +1,121 @@ +import type { Stroke, RecognitionCandidate, EnsembleResult } from './types' +import * as webApiEngine from './webApiEngine' +import * as onnxEngine from './onnxEngine' +import * as tesseractEngine from './tesseractEngine' +import { recognizeStroke } from './strokeMatcher' + +export interface EnsembleSettings { + enableOnnx: boolean + enableTesseract: boolean + enableStrokeMatcher: boolean + autoFillThreshold: number + ambiguityThreshold: number +} + +// when the web API is available, it's so good we give it dominant weight +type EngineKey = 'webapi' | 'stroke' | 'onnx' | 'tesseract' + +const WEIGHTS_WITH_WEB_API = { webapi: 0.7, stroke: 0.15, onnx: 0.1, tesseract: 0.05 } satisfies Record +const WEIGHTS_FALLBACK = { onnx: 0.5, tesseract: 0.3, stroke: 0.2 } satisfies Partial> + +function weightFor(weights: Partial>, engine: EngineKey): number { + return weights[engine] ?? 0 +} + +export async function initEngines(settings: EnsembleSettings): Promise { + const inits: Promise[] = [] + inits.push(webApiEngine.init()) + if (settings.enableOnnx) inits.push(onnxEngine.init()) + if (settings.enableTesseract) inits.push(tesseractEngine.init()) + await Promise.all(inits) +} + +export async function recognizeEnsemble( + strokes: Stroke[], + settings: EnsembleSettings, +): Promise { + const useWebApi = webApiEngine.isReady() + const weights = useWebApi ? WEIGHTS_WITH_WEB_API : WEIGHTS_FALLBACK + + const results: { engine: EngineKey; candidates: RecognitionCandidate[] }[] = [] + const tasks: Promise[] = [] + + // web API - primary when available + if (useWebApi) { + tasks.push( + webApiEngine.recognize(strokes).then(c => { + if (c.length > 0) results.push({ engine: 'webapi', candidates: c }) + }) + ) + } + + if (settings.enableOnnx && onnxEngine.isReady()) { + tasks.push( + onnxEngine.recognize(strokes).then(c => { + if (c.length > 0) results.push({ engine: 'onnx', candidates: c }) + }) + ) + } + + if (settings.enableTesseract && tesseractEngine.isReady()) { + tasks.push( + tesseractEngine.recognize(strokes).then(c => { + if (c.length > 0) results.push({ engine: 'tesseract', candidates: c }) + }) + ) + } + + if (settings.enableStrokeMatcher) { + const c = recognizeStroke(strokes) + if (c.length > 0) results.push({ engine: 'stroke', candidates: c }) + } + + await Promise.all(tasks) + + if (results.length === 0) { + return { candidates: [], decision: 'redraw' } + } + + // calculate active weights + let totalWeight = 0 + for (const r of results) { + const w = weightFor(weights, r.engine) + totalWeight += w + } + + if (totalWeight === 0) { + return { candidates: [], decision: 'redraw' } + } + + // combine scores per letter (only alphabetic characters allowed) + const scores = new Map() + const isLetter = /^[a-zA-Z]$/ + for (const r of results) { + const w = weightFor(weights, r.engine) / totalWeight + for (const c of r.candidates) { + if (!isLetter.test(c.letter)) continue + scores.set(c.letter, (scores.get(c.letter) || 0) + c.confidence * w) + } + } + + const candidates = Array.from(scores.entries()) + .map(([letter, confidence]) => ({ letter, confidence })) + .sort((a, b) => b.confidence - a.confidence) + .slice(0, 5) + + if (candidates.length === 0) { + return { candidates: [], decision: 'redraw' } + } + + const topConfidence = candidates[0].confidence + let decision: EnsembleResult['decision'] + if (topConfidence >= settings.autoFillThreshold) { + decision = 'auto-fill' + } else if (topConfidence >= settings.ambiguityThreshold) { + decision = 'disambiguate' + } else { + decision = 'redraw' + } + + return { candidates, decision } +}