{"version":3,"file":"index.mjs","names":[],"sources":["../src/index.ts"],"sourcesContent":["import type { GraphModel, Tensor } from '@tensorflow/tfjs'\r\nimport type { HandwritingRecognizer, RecognizerInitOptions } from '@zh-keyboard/core'\r\nimport { loadGraphModel } from '@tensorflow/tfjs-converter'\r\nimport * as tf from '@tensorflow/tfjs-core'\r\nimport '@tensorflow/tfjs-backend-cpu'\r\n\r\nexport interface RecognizerOptions {\r\n  /**\r\n   * 模型路径\r\n   */\r\n  modelPath: string\r\n  /**\r\n   * 字典路径\r\n   */\r\n  dictPath: string\r\n  /**\r\n   * 后端类型\r\n   */\r\n  backend?: 'webgl' | 'cpu'\r\n}\r\n\r\nexport class ZhkRecognizer implements HandwritingRecognizer {\r\n  private model?: GraphModel\r\n  private dict: string[] = []\r\n  private canvas: HTMLCanvasElement\r\n  private ctx: CanvasRenderingContext2D\r\n  private modelPath: string\r\n  private dictPath: string\r\n  private backend: 'webgl' | 'cpu'\r\n\r\n  constructor(options: RecognizerOptions) {\r\n    this.modelPath = options.modelPath\r\n    this.dictPath = options.dictPath\r\n    this.backend = options.backend || 'cpu'\r\n    this.canvas = document.createElement('canvas')\r\n    this.canvas.width = this.canvas.height = 64\r\n    this.ctx = this.canvas.getContext('2d', { willReadFrequently: true })!\r\n  }\r\n\r\n  async initialize(options?: RecognizerInitOptions) {\r\n    const text = await fetch(this.dictPath).then(r => r.text())\r\n    this.dict = text.split('\\n')\r\n    this.model = await loadGraphModel(this.modelPath, {\r\n      streamWeights: true,\r\n      onProgress: options?.onProgress,\r\n    })\r\n    // 如果后端为webgl，则需要进行预热\r\n    if (this.backend === 'webgl') {\r\n      await tf.setBackend('webgl')\r\n      await tf.ready()\r\n      await this.recognize([10, 10, 0, 20, 20, 1])\r\n    } else {\r\n      await tf.setBackend('cpu')\r\n    }\r\n    return true\r\n  }\r\n\r\n  async recognize(strokeData: number[]): Promise<string[]> {\r\n    if (!this.model) {\r\n      throw new Error('Model not initialized')\r\n    }\r\n    const { canvas, ctx, model, dict } = this\r\n    ctx.fillStyle = 'white'\r\n    ctx.fillRect(0, 0, canvas.width, canvas.height)\r\n\r\n    const n = strokeData.length / 3\r\n    const strokes = Array.from({ length: n }, (_, i) => ({\r\n      x: strokeData[3 * i],\r\n      y: strokeData[3 * i + 1],\r\n      isEnd: strokeData[3 * i + 2] === 1,\r\n    }))\r\n\r\n    let minX = Infinity\r\n    let minY = Infinity\r\n    let maxX = -Infinity\r\n    let maxY = -Infinity\r\n    for (const { x, y } of strokes) {\r\n      if (x < minX) {\r\n        minX = x\r\n      }\r\n      if (x > maxX) {\r\n        maxX = x\r\n      }\r\n      if (y < minY) {\r\n        minY = y\r\n      }\r\n      if (y > maxY) {\r\n        maxY = y\r\n      }\r\n    }\r\n\r\n    const w = maxX - minX || 1\r\n    const h = maxY - minY || 1\r\n    const cx = (minX + maxX) / 2\r\n    const cy = (minY + maxY) / 2\r\n    const scale = Math.min(canvas.width * 0.9 / w, canvas.height * 0.9 / h)\r\n\r\n    ctx.strokeStyle = 'black'\r\n    ctx.lineWidth = 2\r\n    ctx.lineCap = 'round'\r\n    ctx.lineJoin = 'round'\r\n\r\n    let last: { x: number, y: number, isEnd: boolean } | null = null\r\n    for (const s of strokes) {\r\n      const x = canvas.width / 2 + (s.x - cx) * scale\r\n      const y = canvas.height / 2 + (s.y - cy) * scale\r\n      if (last && !last.isEnd) {\r\n        ctx.beginPath()\r\n        ctx.moveTo(canvas.width / 2 + (last.x - cx) * scale, canvas.height / 2 + (last.y - cy) * scale)\r\n        ctx.lineTo(x, y)\r\n        ctx.stroke()\r\n      } else {\r\n        ctx.beginPath()\r\n        ctx.moveTo(x, y)\r\n      }\r\n      last = s\r\n    }\r\n\r\n    return tf.tidy(() => {\r\n      const image = tf.browser.fromPixels(canvas, 3)\r\n      const floatImage = tf.cast(image, 'float32')\r\n      const normalizedImage = tf.div(floatImage, 255)\r\n      const batchedImage = tf.expandDims(normalizedImage, 0)\r\n\r\n      const probs = (model!.predict(batchedImage) as Tensor).dataSync()\r\n      const idxs = Array.from(probs.keys()).sort((a, b) => probs[b] - probs[a]).slice(0, 10)\r\n\r\n      return idxs.map(i => (i < dict.length ? dict[i] : '')).filter(Boolean)\r\n    })\r\n  }\r\n\r\n  async close() {\r\n    this.model?.dispose()\r\n    this.model = undefined\r\n  }\r\n}\r\n"],"mappings":";;;;AAqBA,IAAa,gBAAb,MAA4D;CAC1D;CACA,OAAyB,EAAE;CAC3B;CACA;CACA;CACA;CACA;CAEA,YAAY,SAA4B;AACtC,OAAK,YAAY,QAAQ;AACzB,OAAK,WAAW,QAAQ;AACxB,OAAK,UAAU,QAAQ,WAAW;AAClC,OAAK,SAAS,SAAS,cAAc,SAAS;AAC9C,OAAK,OAAO,QAAQ,KAAK,OAAO,SAAS;AACzC,OAAK,MAAM,KAAK,OAAO,WAAW,MAAM,EAAE,oBAAoB,MAAM,CAAC;;CAGvE,MAAM,WAAW,SAAiC;AAEhD,OAAK,QADQ,MAAM,MAAM,KAAK,SAAS,CAAC,MAAK,MAAK,EAAE,MAAM,CAAC,EAC1C,MAAM,KAAK;AAC5B,OAAK,QAAQ,MAAM,eAAe,KAAK,WAAW;GAChD,eAAe;GACf,YAAY,SAAS;GACtB,CAAC;AAEF,MAAI,KAAK,YAAY,SAAS;AAC5B,SAAM,GAAG,WAAW,QAAQ;AAC5B,SAAM,GAAG,OAAO;AAChB,SAAM,KAAK,UAAU;IAAC;IAAI;IAAI;IAAG;IAAI;IAAI;IAAE,CAAC;QAE5C,OAAM,GAAG,WAAW,MAAM;AAE5B,SAAO;;CAGT,MAAM,UAAU,YAAyC;AACvD,MAAI,CAAC,KAAK,MACR,OAAM,IAAI,MAAM,wBAAwB;EAE1C,MAAM,EAAE,QAAQ,KAAK,OAAO,SAAS;AACrC,MAAI,YAAY;AAChB,MAAI,SAAS,GAAG,GAAG,OAAO,OAAO,OAAO,OAAO;EAE/C,MAAM,IAAI,WAAW,SAAS;EAC9B,MAAM,UAAU,MAAM,KAAK,EAAE,QAAQ,GAAG,GAAG,GAAG,OAAO;GACnD,GAAG,WAAW,IAAI;GAClB,GAAG,WAAW,IAAI,IAAI;GACtB,OAAO,WAAW,IAAI,IAAI,OAAO;GAClC,EAAE;EAEH,IAAI,OAAO;EACX,IAAI,OAAO;EACX,IAAI,OAAO;EACX,IAAI,OAAO;AACX,OAAK,MAAM,EAAE,GAAG,OAAO,SAAS;AAC9B,OAAI,IAAI,KACN,QAAO;AAET,OAAI,IAAI,KACN,QAAO;AAET,OAAI,IAAI,KACN,QAAO;AAET,OAAI,IAAI,KACN,QAAO;;EAIX,MAAM,IAAI,OAAO,QAAQ;EACzB,MAAM,IAAI,OAAO,QAAQ;EACzB,MAAM,MAAM,OAAO,QAAQ;EAC3B,MAAM,MAAM,OAAO,QAAQ;EAC3B,MAAM,QAAQ,KAAK,IAAI,OAAO,QAAQ,KAAM,GAAG,OAAO,SAAS,KAAM,EAAE;AAEvE,MAAI,cAAc;AAClB,MAAI,YAAY;AAChB,MAAI,UAAU;AACd,MAAI,WAAW;EAEf,IAAI,OAAwD;AAC5D,OAAK,MAAM,KAAK,SAAS;GACvB,MAAM,IAAI,OAAO,QAAQ,KAAK,EAAE,IAAI,MAAM;GAC1C,MAAM,IAAI,OAAO,SAAS,KAAK,EAAE,IAAI,MAAM;AAC3C,OAAI,QAAQ,CAAC,KAAK,OAAO;AACvB,QAAI,WAAW;AACf,QAAI,OAAO,OAAO,QAAQ,KAAK,KAAK,IAAI,MAAM,OAAO,OAAO,SAAS,KAAK,KAAK,IAAI,MAAM,MAAM;AAC/F,QAAI,OAAO,GAAG,EAAE;AAChB,QAAI,QAAQ;UACP;AACL,QAAI,WAAW;AACf,QAAI,OAAO,GAAG,EAAE;;AAElB,UAAO;;AAGT,SAAO,GAAG,WAAW;GACnB,MAAM,QAAQ,GAAG,QAAQ,WAAW,QAAQ,EAAE;GAC9C,MAAM,aAAa,GAAG,KAAK,OAAO,UAAU;GAC5C,MAAM,kBAAkB,GAAG,IAAI,YAAY,IAAI;GAC/C,MAAM,eAAe,GAAG,WAAW,iBAAiB,EAAE;GAEtD,MAAM,QAAS,MAAO,QAAQ,aAAa,CAAY,UAAU;AAGjE,UAFa,MAAM,KAAK,MAAM,MAAM,CAAC,CAAC,MAAM,GAAG,MAAM,MAAM,KAAK,MAAM,GAAG,CAAC,MAAM,GAAG,GAAG,CAE1E,KAAI,MAAM,IAAI,KAAK,SAAS,KAAK,KAAK,GAAI,CAAC,OAAO,QAAQ;IACtE;;CAGJ,MAAM,QAAQ;AACZ,OAAK,OAAO,SAAS;AACrB,OAAK,QAAQ,KAAA"}