目标检测是一类检测目标在图像中的位置,以及在给定图像中对每个感兴趣的目标进行分类的任务。在计算机视觉领域,我们可将此技术应用于图片检索、监控摄像头和无人驾驶汽车中
就目标检测而言,深度卷积神经网络 (DNN) 家族中最负盛名的算法就是 YOLO (You Only Look Once)。
在本文中,我们将使用 TensorFlow 开发一种端到端解决方案, 用以在 Python 中训练自定义目标检测模型,然后将其投入生产,并通过 TensorFlow.js 在浏览器内运行实时推理。
-
项目代码库
https://github.com/hugozanini/TFJS-object-detection
本文的内容将分为以下四个步骤讲述:
目标检测流水线
准备数据
要想训练出优秀的模型,第一步便是获取优质数据。开发此项目时,我并未找到合适(并且足够小)的目标检测数据集,因此我决定自行创建。
我环顾四周,看到了卧室中的袋鼠标志,这是我为纪念在澳大利亚的生活而带回来的纪念品。于是,我决定构建一个袋鼠检测器。
为构建数据集,我使用图片搜索工具搜索并下载了 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 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
克隆、安装,以及测试 Tenso