来自WPS AI工程团队的熊龙飞、杜成、陈荣华、郑辉和胡旭华的客座文章

金山机器学习用例

WPS Office来自金山是一个服务全球1.5亿+用户的生产力工具。我们努力为用户提供最好的工具和功能，以便他们能够高效地完成文档处理。使用 TensorFlow，我们可以帮助用户在许多用例中更智能地处理文档，包括：

• 文档去扭曲：自动将扭曲的文档图像展平。我们将在下一部分详细说明这一点。
  
• 文档边缘检测：WPS 设计的基于 CNN 的矩形检测网络可以仅在设备上快速识别目标文档的边缘。
  
• 自动文档类别分类：我们使用 TensorFlow Lite 实现一个机器学习模型，能够根据文档图像自动对文档进行分类，并根据类别提供相应的功能。
  
• 在线 OCR：我们在服务器上部署 TensorFlow 模型，对用户文档进行旋转校正、文本行检测等，以节省编辑时间。
  
• 离线 OCR：我们使用 TensorFlow Lite 在移动设备上运行 OCR 模型，快速准确地定位文本，方便注重数据安全的用户。
  
• 布局分析：将图像转换为文档时，我们使用 TensorFlow 和 scikit-learn 来实现布局分析；这些框架大大减少了我们的工程工作量。
  

图 1. WPS 的文档布局分析

请注意，所有这些功能都在幕后运行，因此我们鼓励读者在 Google Play 和 Apple App Store 下载我们的应用程序，并试用一下。

今天，我们将重点关注一个用例 - 文档图像去扭曲。

什么是文档图像去扭曲？

阅读书籍时，许多用户在想要提取摘录以供个人笔记时会遇到问题。由于书籍尺寸较大，特定页面无法完全展平，这使得拍摄高质量照片以供稍后阅读成为一项挑战。

图 2. 扭曲的书页 （请注意，由于页面的曲率，文本不方便阅读）

可以肯定的是，你不可能随身携带扫描仪。那么，在这种情况下我们能做些什么呢？市场上的一些工具可以对来自以下扭曲视角的图像进行透视校正

图 3. 透视校正示例

这在某些情况下对用户有帮助，但它只解决了透视校正问题。如果图像存在曲率，则校正效果仍然不令人满意。

我们非常重视用户的这一痛点，并希望通过提供工具来帮助用户轻松获得完全平坦的电子文档，来帮助用户。我们花费了几个月时间将设备上 ML 与服务器端 ML 相结合，并构建了这样一个功能来执行几何校正和扭曲文档图像。现在，我们可以接收一个带弯曲页面的原始图像，并向用户显示一个完美展平的图像，如下所示（我们甚至会自动裁剪图像以供用户使用并去除背景）。

图 4. 去扭曲效果

处理后的图像现在看起来明显更清晰，更易于阅读和使用。能够记笔记或与朋友/同事分享此类图像，使我们的用户更具生产力和效率。

此外，曲率校正对于 OCR 等下游任务非常有用，因为扭曲的字符经常混淆 OCR 算法并导致识别失败。通过对图像进行去扭曲，我们能够显着提高 OCR 功能的准确性。

图 5. 去扭曲如何提高 OCR

现在，您已经对文档图像去扭曲的作用及其有用性有了更好的了解，让我们讨论一下我们如何在 TensorFlow 的帮助下实现它。

用于文档去扭曲的设备上 ML + 服务器端 ML

目前，文档去扭曲的主流方法，以及我们的方法，都遵循相同的模式

• 首先，检测文档的边缘或四个角。
  
• 其次，通过执行透视更改来进行几何校正。
  

设备上边缘检测

边缘检测 传统上非常昂贵，但我们设法使用 TensorFlow Lite 实时运行它，完全在移动设备上。此外，我们甚至可以显示平滑的检测预览。从技术上讲，我们也可以在云端执行此步骤，并且实现起来会容易得多。但是，在设备上运行可以提供更好的用户体验，因为用户可以实时看到检测结果。因此，我们突破了界限并在设备上实现了它。您可以查看以下实时检测演示

图 6. 边缘检测演示

从高层次来看，我们的边缘检测步骤实现如下

1. 我们使用 tf.keras 来构建一个深度神经网络来检测文档边缘。在我们的图像扫描算法中，我们的 CNN 文档检测网络的大小约为 5M。我们使用的骨干网络是 SqueezeNet（TensorFlow 实现是开源的，请查看 此处）。
   
2. 我们使用 fit_generator 来自 tf.keras 模块和 tf.data 模块来预处理训练数据。我们还根据业务需求自定义成本函数。
   
3. 训练后，我们使用 TensorFlow Lite（版本 1.x）附带的 TFLiteConverter 将模型转换为 FlatBuffer 格式并将其集成到我们的应用程序中。
   

当用户打开相机进行预览时，我们的深度学习模型会在后台运行，逐帧检测文档区域并渲染检测结果到屏幕上。为了实时实现这一点，我们不得不偏离传统的边缘检测算法而转向深度学习。随着我们对深度学习模型的迭代，模型的性能不断提高，但模型的尺寸也越来越大，这需要更多计算能力。为了适应中低端手机，我们通过模拟 SqueezeNet 的结构，设计了一个只有 7 层的卷积神经网络，而不是我们最初使用的更深层的网络。虽然新模型比较浅，但准确率只有轻微的损失，速度和内存使用量显着提高。这种权衡使我们能够在低端手机上以每秒几帧的速度运行检测模型，而在高端设备上以每秒 30 帧的速度运行。最后，我们在神经网络外部添加一个过滤器来补偿由不准确检测引起的抖动，这使得检测结果在视觉上更加稳定。

图 7. 我们从中汲取灵感并改进的原始 SqueezeNet 架构 （图片来源：https://www.researchgate.net/figure/SqueezeNet-CNN-Architecture_fig1_306003694）

完成检测模型和编码后，我们发现内部 SDK 包的大小有点太大。除了优化我们自己的代码和依赖项外，我们还需要简化推理框架。我们需要一个快速、轻量级的移动推理框架。我们进行了大量评估和基准测试，最终选择了 TensorFlow Lite。我们选择 TensorFlow Lite 是因为它具有以下优势

• 体积小：我们训练的 ML 模型可以轻松缩小尺寸并交付到移动设备。内存占用也很小，对用户非常友好。
  
• 跨平台支持：TensorFlow Lite 本身支持 iOS 和 Android，我们已经在数十台移动设备上验证了它的稳定性。
  
• 模型转换简便：使用 TensorFlow 训练模型并将其转换为 TensorFlow Lite 模型对我们来说非常顺利。
  
• 自然采用：TensorFlow Lite 的架构设计和工作流程与 TensorFlow 推理非常相似，并且它体积小巧，易于使用。因此，我们自然而然地采用了它。
  
• 充满活力的社区：TensorFlow 拥有非常活跃的社区和丰富的社区资源，因此我们可以从机器学习同行那里学习并找到对我们用例有价值的参考。
  

总的来说，采用 TensorFlow Lite 使我们能够更加专注于设备上的算法开发和模型优化。

云端几何校正

第二步中的几何校正需要巨大的计算能力，目前无法仅在设备上进行。因此，我们将其部署在我们的高性能集群中。我们压缩并加密原始图像，将其传输到我们的服务器，对图像进行去扭曲，然后将它们发送回用户的移动设备以显示。将来，我们也将探索如何将这一部分从云端迁移到设备上。

以下是我们最终去扭曲功能的演示

图 8. WPS 文档去扭曲功能演示

结论

在这篇文章中，我们讨论了我们如何采用混合方法（设备上和云端 ML）来开发 WPS 应用程序（Android 和 iOS）中的文档图像去扭曲功能，这是根据 ML 模型复杂性、硬件能力和业务需求而定的。我们相信，随着硬件变得越来越强大，越来越多的 ML 工作负载将完全在移动设备上运行。另一方面，5G 技术的到来将显着提高带宽并降低网络延迟，有望使云端 ML 更加便捷和用户友好。我们期待着它们，并将根据最新的技术进步不断改进 WPS 产品。