recognition ensemble

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<EngineKey, number>
+const WEIGHTS_FALLBACK = { onnx: 0.5, tesseract: 0.3, stroke: 0.2 } satisfies Partial<Record<EngineKey, number>>
+
+function weightFor(weights: Partial<Record<EngineKey, number>>, engine: EngineKey): number {
+  return weights[engine] ?? 0
+}
+
+export async function initEngines(settings: EnsembleSettings): Promise<void> {
+  const inits: Promise<void>[] = []
+  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<EnsembleResult> {
+  const useWebApi = webApiEngine.isReady()
+  const weights = useWebApi ? WEIGHTS_WITH_WEB_API : WEIGHTS_FALLBACK
+
+  const results: { engine: EngineKey; candidates: RecognitionCandidate[] }[] = []
+  const tasks: Promise<void>[] = []
+
+  // 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<string, number>()
+  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 }
+}