由于 COVID-19 大流行和员工向在家工作的大规模过渡,音频/视频消息传递平台和流媒体服务的流行度迅速增加,因为它们有助于有效地组织远程工作。 员工必须学习新技术并使其适应远程工作。 在家或任何公共场所(如联合办公中心或咖啡馆)工作可能会造成破坏,这会影响您的效率并造成更大的压力。
通常没有高质量的照明(用于视频流),没有真正的绿屏背景,人们在背景中行走并进入画面,这会分散您的注意力,并且不会让您专注于工作。
最近开发了虚拟背景技术(也称为背景替换功能或背景模糊甚至智能色键,它不需要真正的单色背景)来解决这些问题。 该技术允许您实时替换或模糊视频源的背景,从而隐藏您的环境。
最早提供此功能的公司是 Zoom、Skype 和 Google。
虚拟背景替换的流行持续增长,并迅速成为任何视频通话软件的要求。 对于大多数公司而言,在内部构建此类功能是一项耗时且昂贵的任务,还需要开发团队具备高水平的专业知识。 当时市场上没有开箱即用的解决方案。
在本文中,我们将描述我们如何为实时网络摄像头视频源中的背景移除/替换创建跨平台解决方案。
删除背景最流行的方法之一是使用所谓的“绿屏”技术:用户在其后面放置一个纯色画布,软件会删除该颜色的所有像素,以及可选的一些彩色像素 它周围的空间,以防止由画布闪电缺陷和相机传感器噪声引起的颜色溢出效应。 此解决方案提供了良好的质量,但需要一些努力和正确的设置。 对于经常使用蒸汽的人来说,这不会是一个大问题,但是,对于临时使用或视频会议,最好在没有任何准备的情况下进行背景删除。
要移除背景,我们需要将每个像素分类为“人物”或“背景”。 这是语义分割任务的精确表述! 这就是卷积神经网络发挥作用的地方。 自从他们在 ImageNet 大规模视觉识别挑战赛上首次亮相以来,他们已经征服了许多与计算机视觉相关的任务。 今天,深度卷积模型在你能想象到的每个语义分割基准上都排名很高,这使得它们成为背景去除任务的明显选择。
在这里,我们面临两个主要障碍。 首先是准备一个可以解决我们目标的数据集。 第二个是开发一种架构,该架构将在低级和移动硬件上运行,例如功率受限的笔记本电脑 CPU 和运行 Android 和 iOS 的移动设备。
为了更好地理解数据训练的要求,我们检查了用于类似任务的现有数据集。
用于对象检测和语义分割模型基准测试的标准数据集。 它包含很多人的图像,但不幸的是,它们与通常在网络摄像头流和视频会议中看到的图像不同。 带注释的蒙版的质量也有很多不足之处,尤其是在蒙版的边缘。
该数据集是专门为图像中的人物分割而创建的。 图像是从照片库中收集的。 虽然这个数据集中的视觉标注质量相当高,但原始图像与网络摄像头通常可以看到的相差太远:背景模糊,焦点在人身上,场景光线充足,图像分辨率高 它几乎没有噪音。 在这个数据集上训练的模型不适合我们的需求。
数据集旨在解决另一个与我们非常接近的问题——人像抠图。 然而,他们解决的问题在两个主要方面与我们的不同。 首先,我们没有将每个像素分为两类(“人物”和“背景”),而是预测 [0; 1]范围。 这允许在复杂条件下提高背景去除的质量,例如头发分割或背景和人物形象之间的模糊性。 其次,数据集的作者只提供了肖像图像,因此姿势缺乏多样性,整个数据集中几乎看不到手臂。 基于此数据集训练的模型无法正确分割手臂。
在检查数据集和几次尝试训练具有可接受质量的模型的失败尝试后,我们决定实现我们目标的最简单方法是在来自目标域的图像上训练模型。 我们希望直接从网络摄像头收集数据集,以便用于模型训练的图像的视觉特征与模型在推理阶段使用的图像特征相似。 我们要求我们公司的员工在网络摄像头前录制简短的演示视频,同时采取各种姿势(手、头部和上身运动)。
我们数据集的数量和多样性不足以显着改善结果。 由于无法通过直接从最终用户收集数据来扩展数据集这一事实让我们感到沮丧,我们决定转向地球上最大的视频集合——YouTube。
我们从视频博客和流式录制中检索到大约 300 个视频,这些视频是在与我们的目标域相似的条件下拍摄的。 我们主要关注使用低质量网络摄像头在恶劣条件下拍摄的视频,并意识到专业的流媒体设置并不常见。 数据收集和标注随着模型的改进而迭代进行,每一轮新的标注都会在不修改模型的情况下改进结果。 目前,我们的数据集包含大约 12,000 张带有高质量蒙版的图像。
应用增强并没有显着影响指标,但提高了模型预测在自然环境中的稳定性。
尽管数据集是高质量背景去除的关键组成部分,但模型本身仍然发挥着重要作用。 如果您需要从解决方案中榨取所有性能,您不能只使用您最喜欢的文章中的任何常见深度架构。
我们所有的模型都基于“编码器-解码器”方案。 网络的第一部分(编码器)以阶段方式从输入图像中提取特征,在每个处理阶段将特征图的空间分辨率降低 2 倍。 生成的特征图分辨率可能比原始图像低 32 倍,因此很难从如此小的图像中预测出精细的细节。 这就是为什么我们需要网络的第二部分——解码器。 它负责使用编码器部分获得的特征恢复图像的原始分辨率。
该模型的第一个版本使用 DenseNet 的小型变体作为编码器,并使用与卷积混合的双线性插值序列作为解码器。 虽然模型性能很好,但它只能预测真人面具的粗略轮廓。 复杂的背景和可见的手臂使模型非常混乱,以至于无法使用。
接下来,我们尝试了一个普通的旧 UNet 模型。 它非常准确,但即使对于实时 GPU 推理来说也太重了。 我们尝试了不同的选项来天真地缩小模型:降低输入图像分辨率、降低特征图深度、减少解码阶段。 不幸的是,预测质量下降的速度快于推理速度的提高。 最后,我们用深度可分离选项替换了常规卷积。 这使得推理速度可以接受,而不会显着降低分割质量。 该模型展示了足以进行演示和部署的质量。
同时,我们评估了基于“扩张”卷积的轻量级模型。 但是,我们无法足够快地获得具有良好分割质量的快速模型,因此我们放弃了这条路。
作为模型后续迭代的灵感,我们使用了单目深度估计论文中提出的类 UNet 架构。 作者提出了具有轻编码器和重解码器的非对称架构。 本文的另一个重要想法是与输入相比以较低的分辨率进行预测。 我们采用了这些想法,并使用随处可见的轻量级瓶颈和简化的解码器构建了自己的模型。 我们放弃了多尺度损失评估,因为我们没有观察到它对预测质量的任何积极影响。 相反,我们添加了一个分支,它使用最后一个特征图仅估计掩码的边界,这有助于边界处预测的稳定性。
培训完成后,需要准备模型以部署到生产环境。 在 SDK 开发之初,我们计划将该模型部署为能够在 CPU 和 GPU 上运行的原生 Windows 应用程序的一部分。 不幸的是,由于在 CPU 上运行时推理速度慢,无法使用 PyTorch 部署模型。 此外,Windows 仅支持 Nvidia GPU,框架库二进制文件的大小约为 3 GB。 此外,CuDNN 必须由用户手动安装,因为在应用程序中将其作为依赖项分发是非法的。
接受了这个可悲的事实后,我们转向了专门为神经网络推理设计的其他解决方案。 幸运的是,有很多成熟的框架可以轻松完成这项任务。
首先,我们看到了 WinML——微软的神经网络推理框架,从 1803 版开始与 Windows 10 集成。它允许你在任何供应商的 CPU 和 GPU 上运行模型。 框架只能与以 ONNX 格式序列化的模型一起使用。 幸运的是,PyTorch 原生支持 ONNX 序列化。
使用 WinML,我们能够快速获得原型并继续检查其他工具以进行模型推理。
OpenVINO 是一个英特尔框架,旨在优化和推理 x86 CPU 和英特尔 iGPU 上的神经网络。 它实现了计算机视觉网络所需的大量流行算子,因此在模型转换过程中出现问题的可能性几乎为零。 它伴随着优秀的文档、模型测试和基准测试工具、大量现成的模型。 您不能直接从 PyTorch 导出模型,首先需要将其转换为 ONNX。 当然,也支持集成到 C++ 项目中。
在桌面 SDK 的初始测试版之后,我们收到了添加在移动设备和 Web 浏览器上运行模型的能力的请求。
我们从 Google 创建的 MediaPipe 库开始。 它是为在多媒体应用程序中实现流式视频和音频处理的管道而创建的。 您设置数据源(相机、麦克风)、指定数据操作顺序的图形,框架负责视频和音频帧处理。 由于现代应用程序可能会使用一些神经网络来处理花哨的东西,MediaPipe 包括用于神经网络推理的 TensorFlow Lite。 它支持桌面和移动平台,但有一些警告,我们将在下面讨论。
由于库本身不支持从 PyTorch 或 ONNX 导入,我们需要决定如何将模型转换为 TensorFlow Lite 允许的格式。 在研究了其他开发人员的经验后,我们选择了最简单但最可预测的路径——从头开始在 TensorFlow 中重新实现模型并从 PyTorch 导入权重。 这个过程很简单,而且大多是机械的。 最烦人的是,TF 使用的是 NHWC 张量格式,而 PyTorch 默认使用的是 NCHW 格式。 这种不兼容性可以通过在导入期间简单地转置卷积层的权重来解决。 我们确保模型正确(通过比较 PyTorch 和 TensorFlow 模型的输出)并将其从 TensorFlow 图转换为 TensorFlow Lite 格式,并尝试通过 MediaPipe 启动模型。
对于 Android,我们必须实现一些操作,例如使用 OpenGL 着色器将先前预测的蒙版与来自相机的当前图像连接起来,但整体模型移植大约需要 2 天。 在中端设备上,我们能够获得流畅的 25-30 FPS,这让我们很高兴。 此外,到那时我们有了进一步优化模型的新想法。
不幸的是,iOS 出了问题。 主要问题是 Apple 决定不在移动设备上实现 OpenGL 计算着色器,因此我们必须自己实现 Metal 和 OpenGL 之间的纹理转换,反之亦然。 我们无法足够快地实现这一点,因此我们在 iOS 上从 MediaPipe 切换到 CoreML。
CoreML 是用于神经网络推理的原生 Apple 框架。 它支持所有最新的 Apple 硬件,包括 GPU 和用于神经网络加速的专用 SoC 模块 - NPU。 由于 Apple 在其 coremltools 支持库中实现了 PyTorch 到 CoreML 的直接转换,因此移植模型的过程非常顺利。 该模型使用 iPhone X 上的移动 GPU 以 30 FPS 的速度运行。
今天,许多应用程序从一开始就被开发为面向 Web 的,因此我们不想错过将我们的模型交付给 Web 环境的机会。 为了完成这项任务,我们尝试了三个选项:ONNX.js、TensorFlow.js 和 TensorFlow Lite。 不幸的是,这种经历非常痛苦,目前只有一种解决方案令人满意。
在浏览器中启动模型的最简单方法(至少理论上是这样)是 ONNX.js,它允许您加载序列化到 ONNX 的模型并在 GPU 或 CPU 上执行它们。 遗憾的是,该库似乎没有得到积极维护,并且仅支持从旧版本的 ONNX IR 导入一组有限的运算符。 从较新版本的 PyTorch 导出模型时,这可能会导致问题。 在我们的例子中,问题在于转置卷积和上采样层。 同样,我们没有足够快地克服这个问题,所以我们去探索其他选择。
TensorFlow.js 是一个在 JavaScript 环境中支持神经网络训练和推理的框架,但也有一种方法可以将模型从 TensorFlow 直接导出到 TensorFlow.js。 我们采用了为移动版本准备的模型,并能够通过 WebGL 在 GPU 上运行它。 启用 XNNPACK 的 CPU 推理存在问题——转置卷积给出了不正确的结果,因此我们将其替换为双线性上采样。
我们查看的最后一个选项是使用 emscripten 编译为 WebAssembly 的 TensorFlow Lite。 WebAssembly 是一个标准,它定义了可以在沙箱中执行哪些与平台无关的字节码。 它主要是为了向浏览器环境提供高性能代码而创建的。 emscripten 工具链允许您将 C/C++ 代码编译为 WebAssembly。 幸运的是,TensorFlow Lite 是用 C++ 编写的,这让我们可以直接在浏览器中运行它。 将 TensorFlow Lite 编译为 WASM 并不是很曲折,并且表现出足以满足我们需求的性能。 此外,这种方法开辟了使用其他库的可能性,例如,我们乐于使用的用于图像预处理和后处理的 OpenCV。
集成的主要障碍是我们检查过的所有框架都缺乏内置的模型加密功能。 这就是为什么我们选择从源代码编译 TensorFlow Lite 并自己实现模型加密的原因。
目前,我们为各大平台准备了四个SDK版本,快速、稳定、高质量的背景分割。 已经有与商业项目的集成,我们一定会继续开发这个产品。