用5行代码构建自定义训练的对象检测模型

如今，机器学习和计算机视觉已成为一种热潮。 我们都已经看到了有关自动驾驶汽车和面部识别的新闻，并且可能想象到建立我们自己的计算机视觉模型将会多么酷。 但是，进入该领域并不总是那么容易，尤其是在没有扎实的数学背景的情况下。 如果您想做的只是尝试一些小事情，像PyTorch和TensorFlow这样的库可能会很乏味。

在本教程中，我为所有人提供了一种仅需几行代码即可构建功能齐全的对象检测模型的简单方法。 更具体地说，我们将使用Detecto ，这是在PyTorch之上构建的Python软件包，它使该过程变得容易并且对所有级别的程序员开放。

快速简便的示例

为了演示使用Detecto有多么简单，让我们加载一个经过预训练的模型，并对下图进行推断：

首先，使用pip下载Detecto软件包：

pip3 install detecto

然后，将上面的图像另存为“ fruit.jpg”，并在与图像相同的文件夹中创建一个Python文件。 在Python文件中，编写以下5行代码：

from detecto import core, utils, visualize

image = utils.read_image( 'fruit.jpg' )
model = core.Model()

labels, boxes, scores = model.predict_top(image)
visualize.show_labeled_image(image, boxes, labels)

运行此文件后（如果您的计算机上没有启用CUDA的GPU，可能会花费几秒钟；稍后再进行介绍），您应该会看到类似以下内容的内容：

太棒了！ 我们仅用5行代码即可完成所有这些工作。 这是我们在每个过程中所做的：

1. 导入的Detecto模块
2. 读入图像
3. 初始化预训练模型
4. 在我们的图像上产生了最高的预测
5. 绘制我们的预测

Detecto使用PyTorch的模型动物园中的Faster R-CNN ResNet-50 FPN ，它能够检测大约80种不同的物体，例如动物，车辆，厨房用具等。但是，如果您想检测自定义物体（如可乐vs百事可乐罐头，斑马还是长颈鹿？

您会很高兴知道，在自定义数据集上训练Detecto模型同样简单； 同样，您只需要5行代码，以及现有的数据集或花一些时间标记图像。

建立自定义数据集

在本教程中，我们将从头开始构建自己的数据集。 我建议您执行相同的操作，但是如果要跳过此步骤，可以在此处下载示例数据集（从Stanford的Dog Dataset修改）。

对于我们的数据集，我们将训练模型以检测RoboSub比赛中的水下外星人，蝙蝠和女巫，如下所示：

理想情况下，每个类至少需要100张图像。 好消息是，每个图像中可以有多个对象，因此，如果每个图像都包含要检测的每个类别的对象，那么从理论上讲，您可以避免使用100张图像。 另外，如果您有视频素材，Detico可以轻松地将这些素材分割成图像，然后将其用于数据集：

from detecto.utils import split_video

split_video( 'video.mp4' , 'frames/' , step_size= 4 )

上面的代码在“ video.mp4”中每隔4帧拍摄一次，并将其另存为JPEG文件在“ frames”文件夹中。

生成训练数据集后，应该具有一个类似于以下内容的文件夹：

images/
|   image0.jpg
|   image1.jpg
|   image2.jpg
|   ...

如果需要，您还可以使用另一个文件夹，其中包含一组验证图像。

现在是耗时的部分：标记。 Detecto支持PASCAL VOC格式，您可以使用XML文件来包含图像中每个对象的标签和位置数据。 要创建这些XML文件，可以使用开源LabelImg工具，如下所示：

pip3 install labelImg    # Download LabelImg using pip
labelImg                 # Launch the application

现在，您应该会看到一个弹出窗口。 在左侧，单击“打开目录”按钮，然后选择要标记的图像文件夹。 如果一切正常，您应该会看到类似以下内容：

要绘制边界框，请单击左侧菜单栏中的图标（或使用键盘快捷键“ w”）。 然后，您可以在对象周围拖动一个框并编写/选择标签：

标记完图像后，请使用CTRL + S或CMD + S保存XML文件（为简便起见，您可以使用它们自动填充的默认文件位置和名称）。 要标记下一张图像，请单击“下一张图像”（或使用键盘快捷键“ d”）。

完成整个数据集后，您的文件夹应如下所示：

images/
|   image0.jpg
|   image0.xml
|   image1.jpg
|   image1.xml
|   ...

我们几乎准备开始训练我们的对象检测模型！

访问GPU

首先，检查您的计算机是否具有启用CUDA的GPU 。 由于深度学习需要大量处理能力，因此在典型的CPU上进行训练可能会非常缓慢。 值得庆幸的是，大多数现代深度学习框架（例如PyTorch和Tensorflow）都可以在GPU上运行，从而使处理速度更快。 确保已下载PyTorch（如果已安装Detecto，则应该已经下载了），然后运行以下两行代码：

import torch

print(torch.cuda.is_available())

如果打印True，那就太好了！ 您可以跳到下一部分。 如果打印为False，请不要烦恼。 请按照以下步骤创建Google Colaboratory笔记本，这是一个在线编码环境，带有免费的可用GPU。 对于本教程，您将只在Google云端硬盘文件夹中而不是在计算机上工作。

1.登录到Google云端硬盘

2.创建一个名为“ Detecto Tutorial”的文件夹并导航到该文件夹

3.将您的训练图像（和/或验证图像）上传到此文件夹

4.右键单击，转到“更多”，然后单击“ Google Colaboratory”：

现在，您应该看到如下界面：

5.如果需要，请给笔记本起个名字，然后转到“编辑”->“笔记本设置”->“硬件加速器”，然后选择“ GPU”

6.键入以下代码以“装入”驱动器，将目录更改为当前文件夹，然后安装Detecto：

import os
from google.colab import drive

drive.mount( '/content/drive' )

os.chdir( '/content/drive/My Drive/Detecto Tutorial' )

!pip install detecto

为了确保一切正常，您可以创建一个新的代码单元并键入!ls以检查您是否在正确的目录中。

训练自定义模型

最后，我们现在可以在自定义数据集上训练模型！ 如所承诺的，这是容易的部分。 它只需要4行代码：

from detecto import core, utils, visualize

dataset = core.Dataset( 'images/' )
model = core.Model([ 'alien' , 'bat' , 'witch' ])

model.fit(dataset)

让我们再次分解每一行代码的处理方式：

1. 导入的Detecto模块
2. 从“ images”文件夹（包含我们的JPEG和XML文件）创建了一个数据集
3. 初始化一个模型来检测我们的自定义对象（外星人，蝙蝠和女巫）
4. 在数据集上训练我们的模型

根据数据集的大小，这可能需要10分钟到1个小时以上的时间才能运行，因此请确保您的程序在完成上述语句后不会立即退出（即，您使用的是保留状态的Jupyter / Colab笔记本）活动时）。

使用训练好的模型

现在您已经有了训练有素的模型，让我们在一些图像上对其进行测试。 要从文件路径读取图像，可以使用detecto.utils模块中的read_image函数（也可以使用上面创建的数据 detecto.utils的图像）：

# Specify the path to your image
image = utils.read_image( 'images/image0.jpg' )
predictions = model.predict(image)

# predictions format: (labels, boxes, scores)
labels, boxes, scores = predictions

# ['alien', 'bat', 'bat']
print(labels) 

#           xmin       ymin       xmax       ymax
# tensor([[ 569.2125,  203.6702, 1003.4383,  658.1044],
#         [ 276.2478,  144.0074,  579.6044,  508.7444],
#         [ 277.2929,  162.6719,  627.9399,  511.9841]])
print(boxes)

# tensor([0.9952, 0.9837, 0.5153])
print(scores)

如您所见，模型的预测方法返回一个由3个元素组成的元组：标签，框和分数。 在上面的示例中，模型在坐标[569、204、1003、658]（ boxes[0] ）处预测了外星人（ labels[0] ），其置信度为0.995（ scores[0] ）。

根据这些预测，我们可以使用detecto.visualize模块绘制结果。 例如：

visualize.show_labeled_image(image, boxes, labels)

将上面的代码与收到的图像和预测一起运行将产生如下所示的内容：

如果您有视频，则可以在其上运行对象检测：

visualize.detect_video(model, 'input.mp4' , 'output.avi' )

这将获取一个名为“ input.mp4”的视频文件，并根据给定模型的预测结果生成一个“ output.avi”文件。 如果使用VLC或其他视频播放器打开此文件，应该会看到一些有希望的结果！

最后，您可以从文件中保存和加载模型，从而可以保存进度并稍后返回：

model.save( 'model_weights.pth' )

# ... Later ...

model = core.Model.load( 'model_weights.pth' , [ 'alien' , 'bat' , 'witch' ])

高级用法

您会很高兴知道Detecto不仅限于5行代码。 例如，假设该模型的运行不如您希望的那样。 我们可以尝试通过使用Torchvision转换来扩展我们的数据集并定义一个自定义DataLoader来提高其性能：

from torchvision import transforms

augmentations = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomHorizontalFlip( 0.5 ),
    transforms.ColorJitter(saturation= 0.5 ),
    transforms.ToTensor(),
    utils.normalize_transform(),
])

dataset = core.Dataset( 'images/' , transform=augmentations)

loader = core.DataLoader(dataset, batch_size= 2 , shuffle= True )

此代码对数据集中的图像应用了随机的水平翻转和饱和效果，从而增加了数据的多样性。 然后，我们使用batch_size=2定义一个DataLoader对象； 我们将其传递给model.fit而不是Dataset来告诉我们的模型对2张图像进行批处理，而不是默认的1张。

如果您之前创建了单独的验证数据集，那么现在是在培训期间加载它的时候了。 通过提供验证数据集， fit方法将返回每个时期的损失列表，如果verbose=True ，那么它还将在训练过程本身中将其打印出来。 以下代码块演示了这一点，并自定义了其他几个训练参数：

import matplotlib.pyplot as plt

val_dataset = core.Dataset( 'validation_images/' )

losses = model.fit(loader, val_dataset, epochs= 10 , learning_rate= 0.001 , 
                   lr_step_size= 5 , verbose= True )
                   
plt.plot(losses)
plt.show()

损失的结果图应或多或少地减少：

为了获得更大的灵活性和对模型的控制，您可以完全绕过Detecto。 model.get_internal_model方法返回所使用的基础火炬模型，您可以根据需要随意调整。

结论

在本教程中，我们展示了计算机视觉和对象检测不需要挑战。 您只需要一点时间和耐心就可以得出带有标签的数据集。

如果您有兴趣进一步探索，请在GitHub上查看Detecto或访问文档以获取更多教程和用例！

先前发布在https://medium.com/@alankbi/build-a-custom-trained-object-detection-model-with-5-lines-of-code-713ba7f6c0fb

From: https://hackernoon.com/build-a-custom-trained-object-detection-model-with-5-lines-of-code-y08n33vi