使用 TensorFlow.js 在浏览器中进行自定义目标检测
2021 年 1 月 22 日

Hugo Zanini 的客座文章,机器学习工程师

目标检测是在图像中检测目标所在位置并对图像中所有感兴趣的目标进行分类的任务。在计算机视觉中,该技术应用于图片检索、安全摄像头和自动驾驶等领域。

最著名的深度卷积神经网络 (DNN) 家族之一是 YOLO (You Only Look Once)。

在这篇文章中,我们将使用 TensorFlow 开发一个端到端解决方案,以在 Python 中训练自定义目标检测模型,然后将其投入生产,并通过 TensorFlow.js 在浏览器中运行实时推断。

这篇文章将分为以下四个步骤

目标检测流程

准备数据

训练一个优秀模型的第一步是拥有高质量的数据。在开发这个项目时,我没有找到合适的(且足够小)的目标检测数据集,所以我决定自己创建数据集。

我四处寻找,发现我卧室里有一个袋鼠标志——我买来纪念我在澳洲的时光。所以我决定构建一个袋鼠检测器。

为了构建我的数据集,我从图像搜索中下载了 350 张袋鼠图像,并使用 LabelImg 应用程序手工标注了所有图像。因为每张图像中可能有多个动物,所以最终标注了 520 只袋鼠。

标注示例

在这种情况下,我只选择了一个类别,但该软件也可以用于标注多个类别。它将为每张图像生成一个 XML 文件(Pascal VOC 格式),其中包含所有标注和边界框。

<annotation>
    <folder>images</folder>
    <filename>kangaroo-0.jpg</filename>
    <path>/home/hugo/Documents/projects/tfjs/dataset/images/kangaroo-0.jpg</path>
  <source>
    <database>Unknown</database>
  </source>
  <size>
    <width>3872</width>
    <height>2592</height>
    <depth>3</depth>
  </size>
  <segmented>0</segmented>
  <object>
    <name>kangaroo</name>
    <pose>Unspecified</pose>
    <truncated>0</truncated>
    <difficult>0</difficult>
    <bndbox>
      <xmin>60</xmin>
      <ymin>367</ymin>
      <xmax>2872</xmax>
      <ymax>2399</ymax>
    </bndbox>
  </object>
</annotation> 

XML 标注示例

为了方便转换为 TF.record 格式(如下),我将上面程序的 XML 转换为两个 CSV 文件,其中数据已经按训练集和测试集(80%-20%)拆分。这些文件有 9 列

  • 文件名:图像名称
  • 宽度:图像宽度
  • 高度:图像高度
  • 类别:图像类别(袋鼠)
  • xmin:边界框的最小 x 坐标值
  • ymin:边界框的最小 y 坐标值
  • xmax:边界框的 x 坐标的最大值
  • ymax:边界框的 y 坐标的最大值
  • 来源:图像来源

使用 LabelImg 可以轻松创建自己的数据集,但也可以随意使用 我的袋鼠数据集,我已经将它上传到 Kaggle 上

袋鼠数据集

训练模型

有了良好的数据集,就可以考虑模型了。TensorFlow 2 提供了一个 目标检测 API,它可以轻松构建、训练和部署目标检测模型。在这个项目中,我们将使用这个 API 并使用 Google Colaboratory 笔记本 训练模型。本节的其余部分解释了如何设置环境、模型选择和训练。如果你想直接跳到 Colab 笔记本,请点击这里

设置环境

创建 一个新的 Google Colab 笔记本,并选择 GPU 作为硬件加速器

 Runtime > Change runtime type > Hardware accelerator: GPU 

克隆、安装并测试 TensorFlow 目标检测 API

获取和处理数据

如前所述,该模型将使用 Kaggle 上的 袋鼠数据集 进行训练。如果你也想要使用它,则需要创建一个用户,进入 Kaggle 的帐户部分,并获取一个 API Token

获取 API Token

然后,就可以下载数据了

现在,需要创建一个 labelmap 文件来定义要使用的类别。袋鼠是唯一的类别,所以在 Google Colab文件 部分右键点击并创建一个名为 labelmap.pbtxt新文件,内容如下

 item {
    name: "kangaroo"
    id: 1
}

最后一步是将数据转换为一系列 二进制记录,以便它们可以被 TensorFlow 的目标检测 API 馈送。为此,可以使用 袋鼠数据集 中提供的 generate_tf_records.py 脚本将数据转换为 TFRecord 格式。

选择模型

现在,我们准备选择一个将成为袋鼠检测器的模型。TensorFlow 2 在 COCO 2017 数据集 上提供了 40 个预训练检测模型。该集合是 TensorFlow 2 检测模型动物园,可以从 这里 访问。

每个模型都有速度、平均精度均值 (mAP) 和输出。通常,较高的 mAP 意味着速度较低,但由于此项目基于单类目标检测问题,因此速度较快的模型 (SSD MobileNet v2 320x320) 应该足够了。

除了模型动物园,TensorFlow 还提供了一个 模型配置存储库。在那里,可以获取在训练之前必须修改的配置文件。让我们下载这些文件

配置训练

如前所述,下载的权重是在 COCO 2017 数据集 上预训练的,但这里的重点是训练模型以识别一个类别,因此这些权重将仅用于初始化网络——这种技术被称为 迁移学习,它通常用于加速学习过程。

现在,需要设置 mobilenet_v2.config 文件,并开始训练。我强烈建议阅读 MobileNetV2 论文 (Sandler, Mark 等人 - 2018) 以了解体系结构的要点。

选择最佳超参数是一个需要一些实验的任务。由于 Google Colab 中的资源有限,我将使用与论文中相同的批次大小,设置一些步骤以获得合理的低损失,并将所有其他值保留为默认值。如果你想尝试更复杂的方法来查找超参数,我推荐使用 Keras Tuner——一个易于使用的框架,它应用了贝叶斯优化、Hyperband 和随机搜索算法。

设置好参数后,开始训练

为了识别训练进展如何,我们使用损失值。损失是一个数字,表示模型在训练样本上的预测有多糟糕。如果模型的预测是完美的,则损失为零;否则,损失更大。训练模型的目标是找到一组权重和偏差,这些权重和偏差在所有样本上的平均损失较低 (深入机器学习:训练和损失 | 机器学习速成课程)。

从日志中,可以看到值呈下降趋势,所以我们说“模型正在收敛”。在下一节中,我们将绘制所有训练步骤的这些值,趋势将更加清晰。

该模型大约训练了 4 小时(使用 Colab GPU),但通过设置不同的参数,可以使训练过程更快或更慢。一切取决于你使用的类别数量和你的精度/召回目标。一个高度准确的网络,它可以识别多个类别,将需要更多步骤,并且需要更详细的参数调整。

验证模型

现在,让我们使用测试数据评估训练好的模型

评估在 89 张图像上完成,并基于 COCO 检测评估指标 提供三个指标:精度、召回率和损失。

召回率衡量模型在命中正类方面有多好。也就是说,在所有正样本中,算法正确识别了多少?

召回率

精度定义了你对正类预测的信任度有多高:在模型预测为正类的样本中,实际有多少是正类?

精度

设置一个实际例子:假设我们有一张包含 10 只袋鼠的图像,我们的模型返回 5 个检测结果,其中 3 个是真正的袋鼠 (TP = 3, FN =7),2 个是错误的检测结果 (FP = 2)。在这种情况下,我们的召回率为 30%(模型检测到图像中的 10 只袋鼠中的 3 只),精度为 60%(在 5 个检测结果中,3 个是正确的)。

精度和召回率按交并比 (IoU) 阈值划分。IoU 定义为预测边界框 (B) 与真实框 (B) 的交集面积除以它们的并集面积 (Zeng, N. - 2018)

交并比

为简单起见,可以认为 IoU 阈值用于确定检测结果是真阳性 (TP)、假阳性 (FP) 还是假阴性 (FN)。以下是一个示例

IoU 阈值示例

有了这些概念,我们可以分析评估中获得的一些指标。从TensorFlow 2 检测模型库中,SSD MobileNet v2 320x320 的 mAP 为 0.202。我们的模型针对不同的 IoU 呈现了以下平均精度 (AP)

 
AP@[IoU=0.50:0.95 | area=all | maxDets=100] = 0.222
AP@[IoU=0.50      | area=all | maxDets=100] = 0.405
AP@[IoU=0.75      | area=all | maxDets=100] = 0.221

这相当不错!我们可以将获得的 AP 与COCO 数据集文档中的 SSD MobileNet v2 320x320 mAP 进行比较。

我们不区分 AP 和 mAP(以及 AR 和 mAR),并假设差异从上下文中可以清楚地了解。

平均召回率 (AR) 按每张图像的最大检测数(1、10、100)进行划分。当每张图像只有一个袋鼠时,召回率约为 30%,而当最多有 100 只袋鼠时,召回率约为 51%。这些值不是很好,但对于我们试图解决的问题类型来说是合理的。

 
(AR)@[ IoU=0.50:0.95 | area=all | maxDets=  1] = 0.293
(AR)@[ IoU=0.50:0.95 | area=all | maxDets= 10] = 0.414
(AR)@[ IoU=0.50:0.95 | area=all | maxDets=100] = 0.514

损失分析非常简单,我们有 4 个值。

 
INFO:tensorflow: + Loss/localization_loss: 0.345804
INFO:tensorflow: + Loss/classification_loss: 1.496982
INFO:tensorflow: + Loss/regularization_loss: 0.130125
INFO:tensorflow: + Loss/total_loss: 1.972911

定位损失计算预测边界框与标记边界框之间的差异。分类损失指示边界框类别是否与预测类别匹配。正则化损失由网络的正则化函数生成,有助于将优化算法引导到正确的方向。最后一项是总损失,是前三项的总和。

Tensorflow 提供了一个工具,可以轻松地可视化所有这些指标。它被称为TensorBoard,可以通过以下命令初始化。

 
%load_ext tensorboard
%tensorboard --logdir '/content/training/'

这将被显示出来,您可以探索所有训练和评估指标。

Tensorboard - 损失

在 IMAGES 选项卡中,可以找到预测与真实值并排的一些比较。在验证过程中,这是一个非常有趣的探索资源。

Tensorboard - 测试图像

导出模型

现在训练已经验证,是时候导出模型了。我们将把训练检查点转换为 protobuf (pb) 文件。此文件将包含模型的图形定义和权重。

由于我们将使用 TensorFlow.js 部署模型,并且 Google Colab 的最大生存期限制为 12 小时,因此让我们下载训练好的权重并将其保存在本地。运行命令 files.download('/content/saved_model.zip') 时,colab 会自动提示文件。

如果要检查模型是否已正确保存,请加载并测试它。我创建了一些函数来简化此过程,因此您可以随意克隆 inferenceutils.py 文件 从我的 GitHub中测试一些图像。

一切正常,因此我们可以准备将模型投入生产。

部署模型

模型将以一种方式部署,即任何人都可以通过 Web 浏览器打开 PC 或移动摄像头并实时执行推理。为此,我们将把保存的模型转换为 Tensorflow.js 层格式,将模型加载到 JavaScript 应用程序中,并将所有内容提供到Glitch 上。

转换模型

此时,您应该在本地保存了类似于这种结构的内容。

 
├── inference-graph
│ ├── saved_model
│ │ ├── assets
│ │ ├── saved_model.pb
│ │ ├── variables
│ │ ├── variables.data-00000-of-00001
│ │ └── variables.index

在我们开始之前,让我们创建一个隔离的 Python 环境,以便在空工作区中工作,避免任何库冲突。安装 virtualenv,然后在 inference-graph 文件夹中打开一个终端,创建并激活一个新的虚拟环境。

 
virtualenv -p python3 venv
source venv/bin/activate

安装TensorFlow.js 转换器

  pip install tensorflowjs[wizard]

启动转换向导

  tensorflowjs_wizard

现在,该工具将引导您完成转换过程,并为您需要做出的每个选择提供解释。下图显示了转换模型时所做的所有选择。大多数是标准的,但可以根据您的需要更改诸如分片大小和压缩之类的选项。

为了让浏览器自动缓存权重,建议将它们分成大约 4MB 的分片文件。为了确保转换能正常工作,也不要跳过操作验证,因为并非所有 TensorFlow 操作都受支持,因此某些模型可能与 TensorFlow.js 不兼容 - 请参阅此列表以了解哪些操作当前受支持。

使用 Tensorflow.js 转换器进行模型转换(完整分辨率图像在此

如果一切正常,您将在 web_model 目录中获得转换为 Tensorflow.js 层格式的模型。该文件夹包含一个 model.json 文件和一组以二进制格式存储的分片权重文件。model.json 包含模型拓扑(也称为“架构”或“图形”:对层及其连接方式的描述)以及权重文件的清单 (Lin, Tsung-Yi, et al)。

  
└ web_model
  ├── group1-shard1of5.bin
  ├── group1-shard2of5.bin
  ├── group1-shard3of5.bin
  ├── group1-shard4of5.bin
  ├── group1-shard5of5.bin
  └── model.json

配置应用程序

模型已准备好加载到 JavaScript 中。我创建了一个应用程序来直接从浏览器执行推理。让我们 克隆存储库,以了解如何在实时情况下使用转换后的模型。这是项目结构。

 
├── models
│   └── kangaroo-detector
│       ├── group1-shard1of5.bin
│       ├── group1-shard2of5.bin
│       ├── group1-shard3of5.bin
│       ├── group1-shard4of5.bin
│       ├── group1-shard5of5.bin
│       └── model.json
├── package.json
├── package-lock.json
├── public
│   └── index.html
├── README.MD
└── src
    ├── index.js
    └── styles.css

为了简便起见,我在 models 文件夹中提供了一个已转换的 kangaroo-detector 模型。但是,让我们将上一节生成的 web_model 放入 models 文件夹中并进行测试。

首先要做的就是定义模型将如何在 load_model 函数中加载(文件 src>index.js 中的第 10-15 行)。有两种选择。

第一种选择是在本地创建一个 HTTP 服务器,它将使模型在 URL 中可用,从而允许请求并被视为 REST API。加载模型时,TensorFlow.js 将执行以下请求。

 
GET /model.json
GET /group1-shard1of5.bin
GET /group1-shard2of5.bin
GET /group1-shard3of5.bin
GET /group1-shardo4f5.bin
GET /group1-shardo5f5.bin

如果选择此选项,请按如下方式定义 load_model 函数。

  async function load_model() {
    // It's possible to load the model locally or from a repo
    // You can choose whatever IP and PORT you want in the "http://127.0.0.1:8080/model.json"     just set it before in your https server
    const model = await loadGraphModel("http://127.0.0.1:8080/model.json");
    //const model = await loadGraphModel("https://raw.githubusercontent.com/hugozanini/TFJS-object-detection/master/models/web_model/model.json");
    return model;
}

然后安装http-server

  npm install http-server -g

转到 models > web_model 并运行以下命令,使模型在 http://127.0.0.1:8080 上可用。当您想要将模型权重保存在安全的地方并控制谁可以向其请求推理时,这是一个不错的选择。-c1 参数用于禁用缓存,--cors 标志启用跨域资源共享,允许托管文件由给定域的客户端 JavaScript 使用。

  http-server -c1 --cors .

或者,您可以将模型文件上传到某个地方,就我而言,我选择了自己的 GitHub 存储库,并在 load_model 函数中引用了 model.json URL。

 

async function load_model() {
    // It's possible to load the model locally or from a repo
    //const model = await loadGraphModel("http://127.0.0.1:8080/model.json");
    const model = await loadGraphModel("https://raw.githubusercontent.com/hugozanini/TFJS-object-detection/master/models/web_model/model.json");
    return model;
}

这是一个不错的选择,因为它为应用程序提供了更大的灵活性,并使其更容易在某些平台上运行,例如Glitch

本地运行

要在本地运行应用程序,请安装所需的包。

 npm install
 And start:
 npm start

应用程序将在 http://localhost:3000 上运行,您应该看到类似于此的内容。

应用程序在本地运行

模型需要 1 到 2 秒才能加载,之后您可以向相机显示袋鼠图像,应用程序将在它们周围绘制边界框。

在 Glitch 上发布

Glitch 是一个用于创建 Web 应用程序的简单工具,我们可以在其中上传代码并使应用程序在网上对每个人可用。将模型文件上传到 GitHub 存储库并在 load_model 函数中引用它们,我们可以简单地登录 Glitch,单击 New project > Import from Github 并选择应用程序存储库。

等待几分钟安装软件包,您的应用程序将在公共 URL 中可用。单击 Show > In a new window,一个选项卡将打开。复制此 URL 并将其粘贴到任何 Web 浏览器(PC 或移动设备)中,您的对象检测将准备就绪。在下面的视频中查看一些示例。

在不同设备上运行模型

首先,我进行了一个测试,向一个袋鼠标志展示了一个袋鼠标志,以验证应用程序的稳健性。这表明模型专门关注袋鼠特征,并且没有专门针对许多图像中存在的无关特征(例如浅色或灌木丛)。

然后,我在我的手机上打开了应用程序,并展示了一些测试集中的图像。模型运行流畅,识别出大多数袋鼠。如果您想测试我的实时应用程序,它在此处可用(glitch 需要几分钟才能启动)。

除了准确性之外,这些实验中一个有趣的部分是推理时间 - 一切都通过 JavaScript 在浏览器中实时运行。能够在浏览器中运行并使用少量计算资源的良好对象检测模型在许多应用程序中必不可少,尤其是在工业中。将机器学习模型放在客户端意味着降低成本并使应用程序更安全,因为用户隐私得到保护,因为无需将信息发送到任何服务器来执行推理。

下一步

浏览器中的对象检测可以解决许多现实世界的问题,我希望本文能作为涉及计算机视觉、Python、TensorFlow 和 JavaScript 的新项目的参考基础。

作为下一步,我想进行更多详细的训练实验。由于资源有限,我无法尝试许多不同的参数,我相信模型还有很大的改进空间。

我更专注于模型的训练,但我希望看到一个更好的应用程序用户界面。如果有人有兴趣为项目做出贡献,请随时在项目存储库中创建拉取请求。创建一个更友好的应用程序将非常棒。

如果您有任何问题或建议,可以通过Linkedin联系我。感谢您的阅读!

下一篇文章
Custom object detection in the browser using TensorFlow.js

由机器学习工程师 Hugo Zanini 撰写的一篇客座文章 对象检测是指检测图像中物体的位置并对给定图像中的每个感兴趣物体进行分类的任务。在计算机视觉中,此技术用于图像检索、安全摄像头和自动驾驶等应用程序。最著名的深度卷积神经网络 (…