Goodnotes 墨迹擦除功能,支持 Windows、Web 和 Android 平台,由 ONNX Runtime 提供支持

作者:

Pedro Gómez, Emma Ning

2024 年 11 月 18 日

Scribble to Erase feature on Goodnotes for Windows, Web, and Android

在过去的三年里,Goodnotes 工程团队一直致力于一个项目,旨在将这款成功的 iPad 笔记应用推广到其他平台,如 Windows、Web 和 Android。这篇文章介绍了这款 2022 年度 iPad 应用如何在三个平台上同时实现最受 AI 需求的功能之一:墨迹擦除,而这一切都得益于 ONNX Runtime。

📝 什么是墨迹擦除?

我们都是人,所以都会犯错。墨迹擦除是一个简单的功能,它允许用户通过在先前创建的内容上涂写来删除内容,而无需使用橡皮擦工具。

Stylus scribbling to erase the word Abracadabra, then rabbit ears, then a hat

用户之前写的任何笔记,无论写了什么,都可以通过简单的涂写手势删除。这个功能对于用户来说可能看起来很简单,但从工程角度来看却相当复杂。

Lorem Ipsum, with the ipsum scribbled out

这是 Goodnotes 工程团队借助 ONNX Runtime 为 Windows、Web 和 Android 发布的第一个使用人工智能的功能,ONNX Runtime 为边缘 AI 提供了跨平台的高性能模型推理。该团队为这个项目使用了内部训练的模型,并在三个不同的平台上进行了设备端评估。

🔍 如何检测涂写?

对于 Goodnotes 来说,涂写手势只不过是添加到文档中的另一个笔画,它遵循一种特殊的模式。一个笔画必须具备 2 个特征才能被认为是涂写

  • 作为笔画一部分的点数应该足够多。
  • 使用 ONNX Runtime 评估的 AI 模型应该将笔画识别为涂写。

对于工程团队来说,这意味着对于 Goodnotes 团队添加到文档中的每个笔记,他们都必须评估其大小和 AI 模型,以确定添加到任何特定文档中的这个新笔记是否是涂写。

一旦预处理阶段检查到笔记大小超过给定的阈值,就到了我们遵循经典 AI 模型评估流程的时候了,流程如下:

  • 从点中提取笔记特征。
  • 使用 ONNX Runtime 评估 AI 模型。

对于特征提取,Goondotes 团队将笔记区域包含的点进行归一化,并将用户手写笔生成的点列表转换为浮点数组。这个过程只不过是所有 AI 模型都遵循的经典特征提取,目的是将用户数据转换为 AI 模型能够理解的内容。

用于这个项目的 AI 模型是一个基于 LSTM 的监督模型,由团队精心制作并部署到每个平台,以便所有用户都可以在设备端评估它,即使没有互联网连接。

一旦这些点被表示为 AI 模型可以处理的内容,我们就可以使用 ONNX Runtime 并读取模型输出,将其作为一个分数,从而确定最近添加的笔记是否是涂写。如果笔画被认为是涂写,则其下方的所有笔记都将被自动删除。

Diagram of ONNX Runtime workflow for handwritten note recognition

🤝 为什么选择 ONNX Runtime?

当 Goodnotes 团队需要评估这个功能的实现时,他们需要做出的一个决定是如何评估 AI 模型。这是该团队的第一个使用 AI 的项目,而这款产品的 iOS 版本使用的是 CoreML,但 CoreML 与当前项目的技术栈不兼容,因为这种 Apple 技术在 iOS/MacOS SDK 之外是不可用的。因此,他们决定尝试一些不同的东西。

ONNX Runtime Logo

Goodnotes 在 Windows、Web 和 Android 平台的技术栈是基于 Web 技术的。在底层,该应用程序使用了一个 渐进式 Web 应用。当应用程序从 Microsoft Store 或任何其他商店(如 Google Play)安装时,该应用程序会使用原生包装器,但最终,这个项目是一个作为全屏原生应用运行的 Web 应用程序。这意味着团队必须使用的评估 AI 模型的技术必须与 Web 技术栈兼容,并且还必须足够高性能以满足他们的需求,并在可能的情况下启用硬件运行时。因此,在查看不同的替代方案时,团队发现了 ONNX 作为一种可移植格式,以及带有 Web 解决方案的 ONNX Runtime,并决定尝试一下。在进行了一些实验和使用 ONNX Runtime 实现该功能之前的 原型 创建之后,该团队认为这是要使用的正确技术!

ONNX Runtime web architecture diagram

团队决定使用 ONNX Runtime 而不是其他技术的原因有四个:

  • 开发的原型 证明 ONNX Runtime 的集成对我们来说非常简单,并且为我们提供了所需的所有功能。
  • ONNX 是一种可移植的格式,我们可以用它来导出我们当前的 CoreML 模型,使其可以在许多不同的操作系统中进行评估。
  • ONNX Runtime 中的执行提供程序 提供特定于硬件的加速,这使我们能够在评估模型时获得尽可能最佳的性能。特别是对于 Web 解决方案,它具有针对 CPU 执行的 WSAM 执行提供程序,以及用于进一步加速的 WebNN 和 WebGPU 执行提供程序,通过利用 GPU/NPU,这对我们来说是非常有趣的示例。
  • 与 AI 模型的 LSTM 设计兼容。

💻 我们的 ONNX Runtime 代码是什么样的?

Goodnotes 团队与 iOS/Mac 团队共享了 Goodnotes 应用程序的大部分业务逻辑代码。这意味着他们编译了原始的 Swift 代码库,处理笔画,并通过 WebAssembly 从 Swift 进行模型输出的后处理。但是,在执行堆栈中有一个点,Goodnotes 团队必须评估模型,在这个点上,团队将执行从 Swift 委托到使用 ONNX Runtime 的 Web 环境。

评估模型的 Typescript ONNX Runtime 代码类似于以下代码片段:

export class OnnxScribbleToEraseAIModel extends OnnxAIModel<Array<Array<number>>, EvaluationResult>

implements ScribbleToEraseAIModel
{
    getModelResource(): OnDemandResource {
        return OnDemandResource.ScribbleToErase;
    }


    async evaluateModel(input: Array<Array<number>>): Promise<EvaluationResult> {
        const startTime = performance.now();
        const { tensor, initializeTensorTime } = this.initializeTensor(input);
        const { evaluationScore, evaluateModelTime } = await this.runModel(tensor);
        const result = {
            score: evaluationScore ?? 0.0,
            timeToInitializeTensor: initializeTensorTime,
            timeToEvaluateTheModel: evaluateModelTime,
            totalExecutionTime: performance.now() - startTime,
        };
        return result;
    }
…..

正如您所看到的,该实现是您对任何 AI 功能所期望的经典代码。输入数据以特征数组的形式获得,我们稍后使用张量将其馈送到模型中。一旦完成评估,我们就会检查作为模型输出获得的分数,如果分数高于特定阈值,则将输入视为涂写。

正如您在代码中看到的,除了初始化张量和评估模型之外,我们还在跟踪执行时间,以便验证我们的实现,并更好地了解实际用户使用此功能时生产环境所需的资源。

private initializeTensor(input: number[][]) {
       const prepareTensorStartTime = performance.now();
       const modelInput = new Float32Array(input.flat());
       const tensor = new Tensor(modelInputTensorType, modelInput, modelInputDimensions);
       const initializeTensorTime = performance.now() - prepareTensorStartTime;
       return { tensor, initializeTensorTime };
   }


private async runModel(tensor: Tensor) {
    const evaluateModelStartTime = performance.now();
    const inferenceSession = this.session;
    const outputMap = await inferenceSession.run({ x: tensor });
    const outputTensor = outputMap[modelOutputName];
    const evaluationScore = outputTensor?.data[0] as number | undefined;
    const evaluateModelTime = performance.now() - evaluateModelStartTime;
    return { evaluationScore, evaluateModelTime };
}

最重要的是,在本例中,Goodnotes 团队决定从 Web Worker 加载 AI 模型并使用 ONNX Runtime 进行评估,并在 Web Worker 中运行推理会话,因为我们应用程序中的这个路径位于关键的 UX 流程中,我们希望最大限度地减少对用户性能的影响。

ort.env.logLevel = 'fatal';
ort.env.wasm.wasmPaths = '/onnx/providers/wasm/';
this.session = await InferenceSession.create(modelURL);

为此项目配置的执行提供程序是根据模型架构的 CPU 提供程序。这是一个轻量级模型,我们可以通过 WASM 在底层支持的默认 CPU 执行提供程序获得相当快的执行时间。我们计划使用 WebGPU 和 WebNN 执行提供程序来应对新的 AI 场景中更高级的模型。

🚀 部署和集成

由于团队使用的技术栈,Web 技术的使用使得 ONNX Runtime 的集成以及我们托管 AI 模型的方式值得一提。对于这个项目,Goodnotes 使用 Vite 作为前端工具,因此他们不得不稍微修改 Vite 配置,不仅分发我们的 AI 模型,还分发 CPU 执行提供程序所需的资源。这对团队来说没什么大不了的,因为 ONNX Runtime 文档已经涵盖了 bundler 的使用,但这对他们来说非常有趣,因为由于该应用程序是一个可以离线使用的 PWA,这种更改增加了捆绑包的大小,不仅包括模型二进制文件,还包括 ONNX Runtime 所需的所有资源。

📈 生产环境运行几个月后的结果

Goodnotes 在几个月前发布了这个功能。从第一天起,所有用户都开始透明地使用这个墨迹擦除模型。他们中的一些人主动开始书写涂写手势来删除内容,而另一些人则将这个功能视为一种自然的手势而发现了它。

自发布之日起,Goodnotes 团队使用 ONNX Runtime 评估 AI 模型已近 20 亿次!使用 CPU 执行提供程序并在 worker 中运行模型,团队获得了 P95 低于 16 毫秒、P99 低于 27 毫秒的评估时间!世界各地来自不同操作系统和平台的用户已经使用墨迹擦除功能修改了他们的笔记,团队对借助这款出色的 ONNX Runtime 解决方案所取得的技术成就感到非常自豪!