diff --git a/frontend/src/recognition/onnxEngine.ts b/frontend/src/recognition/onnxEngine.ts new file mode 100644 index 0000000..274f1b0 --- /dev/null +++ b/frontend/src/recognition/onnxEngine.ts @@ -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 { + 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 { + 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 [] + } +}