使用 TensorFlow.js 在浏览器中自定义目标检测

目标检测是一类检测目标在图像中的位置,以及在给定图像中对每个感兴趣的目标进行分类的任务。在计算机视觉领域,我们可将此技术应用于图片检索、监控摄像头和无人驾驶汽车中

就目标检测而言,深度卷积神经网络 (DNN) 家族中最负盛名的算法就是 YOLO (You Only Look Once)。

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

  • 项目代码库
    https://github.com/hugozanini/TFJS-object-detection

本文的内容将分为以下四个步骤讲述:

使用 TensorFlow.js 在浏览器中自定义目标检测_第1张图片

目标检测流水线

准备数据

要想训练出优秀的模型,第一步便是获取优质数据。开发此项目时,我并未找到合适(并且足够小)的目标检测数据集,因此我决定自行创建。

我环顾四周,看到了卧室中的袋鼠标志,这是我为纪念在澳大利亚的生活而带回来的纪念品。于是,我决定构建一个袋鼠检测器。

为构建数据集,我使用图片搜索工具搜索并下载了 350 张袋鼠图像,并使用 LabelImg 应用手动标记了所有图像。由于每张图像中不止有一只袋鼠,所以我最后标记了 520 只袋鼠。

使用 TensorFlow.js 在浏览器中自定义目标检测_第2张图片

使用 TensorFlow.js 在浏览器中自定义目标检测_第3张图片

标记示例

  • LabelImg
    https://github.com/tzutalin/labelImg

在本例中,我只选择了一个类别,但我也可以使用软件为多种类别添加注释。接下来是为每张图像生成包含所有注释和边界框的 XML 文件(Pascal VOC 格式)。


    images
    kangaroo-0.jpg
    /home/hugo/Documents/projects/tfjs/dataset/images/kangaroo-0.jpg
  
    Unknown
  
  
    3872
    2592
    3
  
  0
  
    kangaroo
    Unspecified
    0
    0
    
      60
      367
      2872
      2399
    
  

XML 注释示例

为了方便转换为 TF.record 格式(如下方所示),我之后将上述计划的 XML 文件转换为了两个 CSV 文件,其中包含在训练和测试 (80%-20%) 中已拆分的数据。这些文件有 9 列内容:

  • 文件名:图像名

  • 宽度:图像宽度

  • 高度:图像高度

  • 类别:图像类别(袋鼠)

  • xmin:边界框 x 坐标的最小值

  • ymin:边界框 y 坐标的最小值

  • xmax:边界框 x 坐标的最大值

  • ymax:边界框 y 坐标的最大值

  • 来源:图像来源

借助 LabelImg 可以十分便捷地创建自己的数据集,但也欢迎直接使用我的袋鼠数据集,我已将其上传至 Kaggle。

使用 TensorFlow.js 在浏览器中自定义目标检测_第4张图片

使用 TensorFlow.js 在浏览器中自定义目标检测_第5张图片

  • 袋鼠数据集
    https://www.kaggle.com/hugozanini1/kangaroodataset

训练模型

有了优质的数据集,下面就该考虑模型的问题了。TensorFlow 2 所提供的目标检测 API 能让我们轻松构建、训练,以及部署目标检测模型。在本项目中,我们将采用此 API,并使用 Google Colaboratory Notebook 训练模型。本部分的剩余内容将解释如何设置环境、选择模型以及开展训练。如果您想直接跳转至对应案例的 Colab Notebook,请访问此处。

  • 目标检测 API 
    https://github.com/tensorflow/models/tree/master/research/object_detection

  • Google Colaboratory Notebook 
    https://colab.research.google.com/drive/1dyuGpeLh1K4rqOkxELLTTlSqhfbTnFac?usp=sharing

  • 此处
    https://colab.research.google.com/drive/1MdzgmdYJk947sXyls45V7auMPttHelBZ?usp=sharing

设置环境

创建一个新的 Google Colab Notebook,然后选择 GPU 作为硬件加速器:

Runtime > Change runtime type > Hardware accelerator: GPU
  • 创建
    https://colab.research.google.com/

克隆、安装,以及测试 TensorFlow 目标检测 API:

你可能感兴趣的:(TensorFlow,实操,TensorFlow,TensorFlow,js.)