onnx recognition engine

2026-02-23 22:56:55 +02:00
parent 495ead0d13
commit f9f90f552b
1 changed files with 214 additions and 0 deletions
@@ -0,0 +1,214 @@
+import type { Stroke, RecognitionCandidate } from './types'
+import { strokesToImage28 } from './preprocess'
+
+// Pure JS CNN inference - no ONNX Runtime needed
+// Model: 4 conv layers (with ReLU, MaxPool, Dropout) + 2 linear layers
+
+interface WeightTensor {
+  dims: number[]
+  data: Float32Array
+}
+
+interface Weights {
+  [name: string]: WeightTensor
+}
+
+let weights: Weights | null = null
+let loading = false
+let initFailed = false
+
+const LETTERS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+
+export function isReady(): boolean {
+  return weights !== null
+}
+
+export async function init(): Promise<void> {
+  if (weights || loading || initFailed) return
+  loading = true
+
+  try {
+    const res = await fetch('/models/emnist-weights.bin')
+    if (!res.ok) throw new Error('weights not found')
+    const buf = await res.arrayBuffer()
+    weights = parseWeights(buf)
+    console.log('CNN engine ready (pure JS)')
+  } catch (e) {
+    console.warn('CNN engine init failed:', e)
+    weights = null
+    initFailed = true
+  }
+  loading = false
+}
+
+function parseWeights(buf: ArrayBuffer): Weights {
+  const view = new DataView(buf)
+  let offset = 0
+
+  function readU32() { const v = view.getUint32(offset, true); offset += 4; return v }
+  function readF32() { const v = view.getFloat32(offset, true); offset += 4; return v }
+
+  const numTensors = readU32()
+  const w: Weights = {}
+
+  for (let t = 0; t < numTensors; t++) {
+    const nameLen = readU32()
+    const nameBytes = new Uint8Array(buf, offset, nameLen)
+    const name = new TextDecoder().decode(nameBytes)
+    offset += nameLen
+
+    const ndims = readU32()
+    const dims: number[] = []
+    for (let d = 0; d < ndims; d++) dims.push(readU32())
+
+    const dataLen = readU32()
+    const data = new Float32Array(dataLen)
+    for (let i = 0; i < dataLen; i++) data[i] = readF32()
+
+    w[name] = { dims, data }
+  }
+  return w
+}
+
+// CNN operations
+
+interface TensorPlane {
+  data: Float32Array
+  h: number
+  w: number
+}
+
+function conv2d(input: Float32Array, inC: number, inH: number, inW: number,
+                weight: Float32Array, bias: Float32Array, outC: number, kH: number, kW: number, pad: number): TensorPlane {
+  const outH = inH + 2 * pad - kH + 1
+  const outW = inW + 2 * pad - kW + 1
+  const output = new Float32Array(outC * outH * outW)
+
+  for (let oc = 0; oc < outC; oc++) {
+    for (let oh = 0; oh < outH; oh++) {
+      for (let ow = 0; ow < outW; ow++) {
+        let sum = bias[oc]
+        for (let ic = 0; ic < inC; ic++) {
+          for (let kh = 0; kh < kH; kh++) {
+            for (let kw = 0; kw < kW; kw++) {
+              const ih = oh - pad + kh
+              const iw = ow - pad + kw
+              if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
+                const inputIdx = ic * inH * inW + ih * inW + iw
+                const weightIdx = oc * inC * kH * kW + ic * kH * kW + kh * kW + kw
+                sum += input[inputIdx] * weight[weightIdx]
+              }
+            }
+          }
+        }
+        output[oc * outH * outW + oh * outW + ow] = sum
+      }
+    }
+  }
+  return { data: output, h: outH, w: outW }
+}
+
+function relu(data: Float32Array): Float32Array {
+  const out = new Float32Array(data.length)
+  for (let i = 0; i < data.length; i++) out[i] = data[i] > 0 ? data[i] : 0
+  return out
+}
+
+function maxpool2d(input: Float32Array, c: number, h: number, w: number, k: number): TensorPlane {
+  const oh = Math.floor(h / k)
+  const ow = Math.floor(w / k)
+  const output = new Float32Array(c * oh * ow)
+  for (let ch = 0; ch < c; ch++) {
+    for (let i = 0; i < oh; i++) {
+      for (let j = 0; j < ow; j++) {
+        let max = -Infinity
+        for (let ki = 0; ki < k; ki++) {
+          for (let kj = 0; kj < k; kj++) {
+            const idx = ch * h * w + (i * k + ki) * w + (j * k + kj)
+            if (input[idx] > max) max = input[idx]
+          }
+        }
+        output[ch * oh * ow + i * ow + j] = max
+      }
+    }
+  }
+  return { data: output, h: oh, w: ow }
+}
+
+function linear(input: Float32Array, weight: Float32Array, bias: Float32Array, outSize: number, inSize: number): Float32Array {
+  const output = new Float32Array(outSize)
+  for (let o = 0; o < outSize; o++) {
+    let sum = bias[o]
+    for (let i = 0; i < inSize; i++) {
+      sum += input[i] * weight[o * inSize + i]
+    }
+    output[o] = sum
+  }
+  return output
+}
+
+function forward(input: Float32Array, w: Weights): Float32Array {
+  // features.0: Conv2d(1, 32, 3, padding=1) + ReLU
+  let { data, h, w: width } = conv2d(input, 1, 28, 28,
+    w['features.0.weight'].data, w['features.0.bias'].data, 32, 3, 3, 1)
+  data = relu(data)
+
+  // features.2: Conv2d(32, 32, 3, padding=1) + ReLU
+  ;({ data, h, w: width } = conv2d(data, 32, h, width,
+    w['features.2.weight'].data, w['features.2.bias'].data, 32, 3, 3, 1))
+  data = relu(data)
+
+  // features.4: MaxPool2d(2)
+  ;({ data, h, w: width } = maxpool2d(data, 32, h, width, 2))
+  // features.5: Dropout2d - skip at inference
+
+  // features.6: Conv2d(32, 64, 3, padding=1) + ReLU
+  ;({ data, h, w: width } = conv2d(data, 32, h, width,
+    w['features.6.weight'].data, w['features.6.bias'].data, 64, 3, 3, 1))
+  data = relu(data)
+
+  // features.8: Conv2d(64, 64, 3, padding=1) + ReLU
+  ;({ data, h, w: width } = conv2d(data, 64, h, width,
+    w['features.8.weight'].data, w['features.8.bias'].data, 64, 3, 3, 1))
+  data = relu(data)
+
+  // features.10: MaxPool2d(2)
+  ;({ data, h, w: width } = maxpool2d(data, 64, h, width, 2))
+  // features.11: Dropout2d - skip at inference
+
+  // classifier.1: Linear(3136, 256) + ReLU
+  data = linear(data, w['classifier.1.weight'].data, w['classifier.1.bias'].data, 256, 64 * h * width)
+  data = relu(data)
+  // classifier.3: Dropout - skip at inference
+
+  // classifier.4: Linear(256, 26)
+  data = linear(data, w['classifier.4.weight'].data, w['classifier.4.bias'].data, 26, 256)
+
+  return data
+}
+
+export async function recognize(strokes: Stroke[]): Promise<RecognitionCandidate[]> {
+  if (!weights) return []
+
+  try {
+    const imageData = strokesToImage28(strokes)
+    const logits = forward(imageData, weights)
+
+    // softmax
+    const maxVal = Math.max(...logits)
+    const exps = Array.from(logits).map(v => Math.exp(v - maxVal))
+    const sum = exps.reduce((a, b) => a + b, 0)
+    const probs = exps.map(e => e / sum)
+
+    const candidates: RecognitionCandidate[] = probs.map((confidence, i) => ({
+      letter: LETTERS[i],
+      confidence,
+    }))
+
+    candidates.sort((a, b) => b.confidence - a.confidence)
+    return candidates.slice(0, 5)
+  } catch (e) {
+    console.warn('CNN recognition failed:', e)
+    return []
+  }
+}