65.9K
CodeProject 正在变化。 阅读更多。
Home

使用 OpenCV 和 ImageAI 训练自定义模型

格伦·普林斯
starIconstarIconstarIconstarIconstarIcon

5.00/5 (7投票s)

2020年6月10日

CPOL

5分钟阅读

viewsIcon

46954

downloadIcon

371

在本文中，我们将开始创建自定义对象检测模型的过程。

在本系列中，我们将学习如何使用Python、OpenCV（一个开源计算机视觉库）和ImageAI（一个用于视觉的深度学习库）来训练AI，以检测工人是否佩戴安全帽。在此过程中，我们将创建一个可用于现实生活中的端到端解决方案——这不仅仅是一项学术练习！

这是一个重要的用例，因为许多公司必须确保工人拥有合适的安全设备。但我们学到的知识不仅可以用于检测安全帽。在本系列结束时，您将能够使用AI检测图像或视频流中的几乎任何类型的对象。

您当前在6篇文章中的第5篇

  1. 为对象检测安装 OpenCV 和 ImageAI
  2. 为 OpenCV 和 ImageAI 对象检测查找训练数据
  3. 使用预训练模型通过OpenCV和ImageAI检测对象
  4. 为使用 OpenCV 和 ImageAI 的对象检测准备图像
  5. 使用OpenCV和ImageAI训练自定义模型
  6. 使用OpenCV和ImageAI检测自定义模型对象

在上一篇文章中，我们清理了数据并将其分为训练集和验证集。

现在我们可以开始创建自定义对象检测模型的过程了。训练自定义检测模型的一般步骤是

  1. 训练模型
  2. 验证模型；如果验证效果不佳，则进行调整并重新训练
  3. 直观地测试模型和结果
  4. 部署模型

训练我们的模型

让我们直接开始训练模型。创建一个新的代码块并输入以下内容

from imageai.Detection.Custom import DetectionModelTrainer

trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="hardhat")

trainer.setTrainConfig(object_names_array=["person hardhat"], batch_size=4, num_experiments=20, 
                       train_from_pretrained_model="yolo.h5")

trainer.trainModel()



此代码块利用了ImageAI的检测类中的一个新方法，即DetectionModelTrainer。训练任何模型的过程是

  • 定义一个新的DetectionModelTrainer()方法。
  • 将模型类型设置为YOLOv3。目前，YOLOv3是唯一可以使用ImageAI训练的模型类型。
  • 设置包含您的数据的目录。请注意，此文件夹必须包含一个名为“train”的文件夹和一个名为“validation”的文件夹。这些文件夹中必须分别包含一个名为“images”的文件夹和另一个名为“annotations”的文件夹。
  • 按如下方式设置训练器的配置
    • 指定图像中使用的标注的名称。在我们的例子中，我们只使用“person hardhat”。
    • 定义批次大小，在本例中为四个。这决定了模型将在每个批次中训练多少张图像。批次大小越大，模型训练效果越好，但请记住，您需要更强大的GPU。
    • 使用num_experiments指定要执行的模型迭代次数。迭代次数越多，最终结果越好，但所需时间也越长。
    • 如果需要，可以指定一个预训练模型进行迁移学习，以更快地获得更好的结果。
  • 使用trainModel()开始模型训练过程。

模型将开始训练，并为每个周期（或epoch）输出一个状态。对于每个周期，都会报告一个损失值，该值决定模型是否比前一个周期更好。如果是，则将保存该模型，因此请确保您有足够的磁盘空间！

验证模型

训练模型可能需要很长时间。此模型进行了20次迭代，训练时间略超过四个小时。一些模型训练建议超过200小时。在我们将计算机放在那里几天来训练模型之前，让我们看看在我们的模型训练了20次迭代后创建了什么。

在您的“hardhat”目录中，您会看到创建了一些额外的目录：“cache”、“json”、“logs”和“models”。这里最重要的两个目录是“json”和“models”。“json”目录包含使用模型所需的JSON配置文件。“model”目录包含许多相当大的模型文件，带有增量编号。这些文件中的每一个都是比上一个更好的模型训练迭代的结果。

因此，我们有许多模型，从理论上讲，它们随着周期号的增加而越来越好。让我们通过验证它们来测试它们。开始一个新的代码块并输入以下内容

trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-020--loss-0008.462.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)

您需要做的唯一更改就是模型路径，使用字符串：hardhat\models\detection_model-ex-020--loss-0008.462.h5，因为每次训练运行都会不同。此方法接受以下参数

  • model_path – 指定您希望对其运行验证的模型
  • json_path – 指定模型训练的配置文件
  • iou_threshold – 表示预测的边界框与实际的边界框的交并比
  • object_threshold – 用于删除检测的置信度级别
  • nms_threshold – 检测到多个边界框时的置信度级别。

当我们对20次迭代的模型进行此验证时，我们获得的平均精度为0.84464，即大约84%，这还不错。但与其他模型相比如何？让我们将代码块扩展为以下内容

model05 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-005--loss-0014.238.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)
model10 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-010--loss-0011.053.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)
model15 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-015--loss-0009.620.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)
model20 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-020--loss-0008.462.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)

print('---------------------------------------------------------')
print('Iteration 05:', model05[0]['average_precision']['person hardhat'])
print('Iteration 10:', model10[0]['average_precision']['person hardhat'])
print('Iteration 15:', model15[0]['average_precision']['person hardhat'])
print('Iteration 20:', model20[0]['average_precision']['person hardhat'])
print('---------------------------------------------------------')



此代码块运行需要一些时间，因为它需要加载4个不同的模型、验证它们并保存结果，所以如果您运行此代码，请离开并稍后再回来。当此代码块最终完成时，最后几行将给出结果

  • 5次迭代 – 71%
  • 10次迭代 – 78%
  • 15次迭代 – 83%
  • 20次迭代 – 84%

因此，从这些结果来看，我们迭代的次数越多，模型的效果就越好。然而，在某些时候，回报会递减，所以您在训练模型时也需要考虑到这一点。

接下来

在本文中，我们学习了如何使用准备好的数据集训练自定义图像检测模型。

接下来，我们将研究如何使用此模型检测人们是否佩戴安全帽。

© . All rights reserved.