在工业环境中,确保工人佩戴适当的安全设备是至关重要的。本文将介绍如何使用OpenCV和ImageAI这两个强大的库来训练一个能够检测工人是否佩戴安全帽的AI模型。这不仅是一个学术练习,而是一个可以在现实世界中应用的端到端解决方案。
在开始之前,需要安装OpenCV和ImageAI。这两个库将是训练AI模型的基础。
为了训练模型,需要收集和准备训练数据。这包括将数据集分为训练集和验证集。
在开始训练自己的模型之前,可以使用OpenCV和ImageAI提供的预训练模型来检测对象。这可以帮助理解模型是如何工作的,并为自己的训练提供参考。
在训练模型之前,需要对图像进行预处理。这包括调整图像大小、归一化和增强图像,以提高模型的性能。
现在可以开始创建自定义的对象检测模型了。训练自定义检测模型的一般步骤如下:
让直接开始训练模型。创建一个新的代码块并输入以下内容:
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。训练任何模型的过程如下:
模型将开始训练,并为每个周期(或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
因为每次训练运行都会有所不同。此方法接受以下参数:
当对20次迭代的模型进行此验证时,得到了大约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个不同的模型,验证它们,并保存结果,所以如果运行它,请稍等一会儿。当这段代码块最终完成后,最后几行将给出结果: