如何使用 MinIO 和 YOLO 的自定义数据集训练对象检测模型

计算机视觉仍然是人工智能极其引人注目的应用。无论是识别战场上的元素还是预测农作物产量，计算机视觉可以说是最具商业价值（也是社会重要）的人工智能领域之一。

然而，采用最佳计算机视觉功能的速度限制因素通常是与构建数据集和设计将在新图像上执行计算机视觉任务的简单端到端系统相关的复杂性。

在这篇博文中，我们将逐步了解如何使用 CVAT 和 MinIO Bucket Notifications 等一流工具解决这些问题。在本文结束时，您将能够在自定义数据集上训练对象检测模型，并在新图像出现时使用它进行预测。

前提

假设我们希望能够识别卫星图像中存在的飞机类型。我们还假设我们从头开始：没有预先构建的数据集，没有预先训练的模型。以下是我们想要在卫星图像中检测和识别的两架示例飞机：

第 1 步：获取和管理培训样本

对于这个项目，很大程度上由于没有按需成像卫星，我访问了谷歌地球上的机场，并对其中一些飞机可见的区域拍摄了多张屏幕截图。组装这组图像花了相当长的时间，因此我将它们全部存储在我的 MinIO 服务器上名为“对象检测”的存储桶中。在生产环境中，将收集的样本存储在 MinIO 上的好处变得更加有先见之明。主动-主动复制、最高级别的加密和超快速的 GET/PUT（仅举几例）意味着您辛勤收集的样本将具有高可用性、安全性和可靠性。

第 2 步：创建数据集

为了为您的用例训练对象检测模型，需要一个带标签（或“带注释”）的数据集。 OpenCV 的 CVAT 是一个很好的工具。一个很酷的功能是，CVAT 提供了一个实用程序，可以将您的 MinIO 存储桶连接为“云存储”，以便将存储桶的图像直接提供给数据集注释工具。为此，请确保 CVAT 服务器可以访问 MinIO Server 的主机，特别是如果您在本地或笔记本电脑上本地运行 MinIO Server。另外，请注意，有两种使用 CVAT 的方法：(1) 使用他们在app.cvat.ai上提供的 Web 应用程序或 (2) 在本地运行它。无论哪种情况，打开 CVAT 后，单击菜单栏中的“云存储”。从那里，您可以填写表格来附加您的（S3 兼容）MinIO 存储桶：

现在让我们在“任务”下创建新的标签任务：

系统应该提示您填写表格：

创建任务时，正确定义类标签非常重要（我定义了两个标题为“SU30”和“TU95”的矩形标签，对应于我想要检测的两个平面）：

现在剩下的步骤是将我们之前添加的 MinIO 存储桶附加为数据源。在“选择文件”下，单击“云存储”并填写您之前为该源提供的名称。我在上面使用了名称“minio-cv-bucket”。

上传过程将需要几分钟。完成后，您应该能够在“作业”下看到可用的注释作业。

现在，通过单击作业，您可以开始为每个图像添加注释。警告：这可能是一个非常耗时的过程。一般来说，在具有大量注释需求的生产环境中，最好将此任务转移给专门的内部团队或第三方数据标记公司。

完成注释后，以 YOLO 格式导出数据集。

第三步：整理训练数据

您导出的数据集将采用 zip 文件的形式。解压后，YOLO 格式的注释文本文件将位于随附的文件夹中。请随意看看它们。在 YOLO 格式中，每个图像的注释都位于文本文件中，其中每行包含边界框和类的两个角。类编号与您在创建任务时定义标签的顺序相对应。因此，在此示例中，0 对应于 Su-30，1 对应于 Tu-95。

此时，创建一个新的工作目录（或输入您已经创建的目录）。在此目录中，创建一个名为“dataset”的子目录。在“数据集”中，创建目录，使您的工作目录如下所示：

 my_cv_project (WORKING DIRECTORY) |---- dataset |----images |----train |----val |----test |----annotations |----train |----val |----test

您现在必须填充图像及其相应注释（文本文件）的 train、val 和 test 子目录。如何检索和分割样本取决于您。一个好的做法是将训练样本总量分为 80% 的训练、10% 的验证和 10% 的测试。确保在分区之前随机打乱图像。

就我个人而言，我在命令行中使用 MinIO Client 的mc cp来快速检索“对象检测”存储桶中的所有图像。或者，如果您已将所有示例图像存储在本地计算机上的一个位置，则可以直接使用它。将所有样本集中到一处后，我使用 Python 脚本对图像和注释进行洗牌、分割和移动到 train、val 和 test 目录。 这是为了方便起见提供的脚本。如果您对如何使用它有任何疑问，请随时在存储库中提出问题！

最后，确保对于放置在 images/train、images/val 或 images/test 中的每个图像，匹配的注释 .txt 文件也位于注释/目录中的相应子目录中。例如：

 my_cv_project (WORKING DIRECTORY) |---- dataset |----images |----train - 5.png - 3.png - 2.png |----val - 4.png |----test - 1.png |----annotations |----train - 5.txt - 3.txt - 2.txt |----val - 4.txt |----test - 1.txt

现在，我们的数据已经到位。是时候看看我们的对象检测模型并开始训练了。

第四步：物体检测模型

当前对象识别的黄金标准（在性能和易用性方面）是 YOLO（只看一次）模型类。在撰写本文时，YOLOv8 是最新版本，由 Ultralytics 开源维护。 YOLOv8 提供了一个简单的 API，我们可以利用它在新创建的注释上训练模型（并最终运行推理）。

让我们下载YOLOv8：

 $ pip install ultralytics

我们现在可以使用 YOLOv8 CLI 工具或 Python SDK 来训练、验证和预测。有关更多信息，请参阅 YOLOv8 文档。

第五步：培训

在您的工作目录中，定义一个 YAML 文件，该文件指定数据集的位置以及有关类的详细信息。请注意，这些路径与我之前在工作目录中创建的路径相同。我将我的文件命名为“ objdetect.yaml ”。另请注意，两个飞机类别标签的定义顺序必须与 CVAT 中的顺序相同。

 train: ./dataset/images/train/ val: ./dataset/images/val/ test: ./dataset/images/test/ # number of classes nc: 2 # class names names: ["SU-30","TU-95"]

使用以下命令（使用 YOLO CLI 工具）开始在我们的数据集上训练 YOLOv8 模型。请参阅 YOLO 文档，了解有关可以为训练配置的所有不同选项的更多信息。在这里，我开始训练 100 个 epoch，并将图像大小设置为 640 像素（我们所有的训练图像将在训练期间相应缩放）：

 $ yolo task=detect \ mode=train \ model=yolov8s.pt \ data=objdetect.yaml \ epochs=100 \ imgsz=640

培训需要一段时间，特别是如果您在笔记本电脑上工作（像我一样），所以现在是休息的好时机（或提前阅读😀）！

在训练循环结束时，经过训练的模型以及其他有趣的图形和图表将存储在名为“runs”的自动生成的目录中。终端输出（如下所示）将指示最近一次运行结果的具体位置。每次训练模型时，都会在“runs/detect/”中生成类似的目录。

 Results saved to runs/detect/train

步骤 5B：模型验证和测试

您可以使用以下命令运行验证：

 $ yolo task=detect \ mode=val \ model=path/to/best.pt \ data=objdetect.yaml

结果将自动存储在工作目录中的文件夹中，路径格式为“runs/detect/val”。

要对测试集进行推理，可以使用以下命令：

 $ yolo task=detect \ mode=predict \ model=path/to/best.pt \ conf=0.5 \ source=dataset/images/test

结果将存储在“运行/检测/预测”中。以下是测试集上的一些预测结果：

第 6 步：使用 MinIO Bucket 通知进行新图像推理

现在我们有了一个训练有素的模型，可以识别卫星图像中存在的某些飞机类型，我们如何以简单的方式将其用于新图像？

MinIO Bucket Notifications是一个完美的工具。我们可以构建一个系统，可以借助 Webhook 对放入存储桶中的新图像自动执行对象检测推理。

概括地说，我们有 3 个步骤。首先，我们需要定义一个端点，该端点可以用作网络钩子，以使用我们训练的模型对新图像执行对象检测。其次，我们需要为 MinIO 服务器部署配置一些环境变量，指示它在发生某些事件时访问我们的 webhook 端点。第三，我们需要配置要处理的存储桶事件类型（即 PUT）。让我们一步一步地看一下它。

以下是一个基于 Flask 的简单服务器 ( detector_server.py ) 的代码，该服务器对添加到 MinIO 存储桶的新图像运行推理：

 """ This is a simple Flask inference server implementation that serves as a webhook for the event of a new image being added to a MinIO bucket. Object detection using YOLO will be performed on that image and the resulting predictions will be returned. """ from flask import Flask, request, abort, make_response from ultralytics import YOLO import tempfile from minio import Minio # Make sure the following are populated with your MinIO details # (Best practice is to use environment variables!) MINIO_ENDPOINT = '' MINIO_ACCESS_KEY = '' MINIO_SECRET_KEY = '' model = YOLO('/PATH/TO/best.pt') # load a custom model (path to trained weights) client = Minio( MINIO_ENDPOINT, access_key=MINIO_ACCESS_KEY, secret_key=MINIO_SECRET_KEY, ) app = Flask(__name__) @app.route('/', methods=['POST']) async def inference_bucket_webhook(): """ This endpoint will be called when a new object is placed in your inference bucket """ if request.method == 'POST': # Get the request event from the 'POST' call event = request.json bucket = event['Records'][0]['s3']['bucket']['name'] obj_name = event['Records'][0]['s3']['object']['key'] with tempfile.TemporaryDirectory() as temp_dir: temp_file_name = temp_dir+'/'+obj_name client.fget_object(bucket, obj_name, temp_file_name) # See https://docs.ultralytics.com/modes/predict/ for more information about YOLO inference options results = model.predict(source=temp_file_name, conf=0.5, stream=False) # A list of bounding boxes (if any) is returned. # Each bounding box is in the format [x1, y1, x2, y2, probability, class]. result = {"results": results[0].boxes.data.tolist()} print(result) resp = make_response(result, 200) return resp else: abort(400) if __name__ == '__main__': app.run()

让我们启动推理服务器：

 $ python detection_server.py * Serving Flask app 'detection_server' * Debug mode: off * Running on http://127.0.0.1:5000 Press CTRL+C to quit

记下 Flask 应用程序运行的主机名和端口。

接下来，让我们开始在 MinIO 端配置 Webhooks。首先，设置以下环境变量。将 <YOURFUNCTIONNAME> 替换为您选择的函数名称。为了简单起见，我选择了“推理”。另外，请确保端点环境变量设置为推理服务器的正确主机和端口。在本例中，http://localhost:5000 是我们的 Flask 应用程序运行的位置。

 $ export MINIO_NOTIFY_WEBHOOK_ENABLE_<YOURFUNCTIONNAME>=on $ export MINIO_NOTIFY_WEBHOOK_ENDPOINT_<YOURFUNCTIONNAME>=http://localhost:5000

现在，使用命令mc admin service restart ALIAS重新启动 MinIO 服务器，或者如果您是第一次启动服务器，也可以只使用minio server命令。有关重新/启动 MinIO 服务器的更多信息，请查看 MinIO 文档。注意： ALIAS 应替换为您的 MinIO 服务器部署的别名。有关如何设置别名或查看现有别名的更多信息，请查看文档。

最后，让我们添加我们想要通知的存储桶和事件。在我们的例子中，我们希望收到有关存储桶中“ put”事件（创建新对象）的通知。为此，我制作了一个全新的空桶，名为“检测推理”，因此我将其替换为“桶”。

 $ mc event add ALIAS/BUCKET arn:minio:sqs::<YOURFUNCTIONNAME>:webhook --event put

您可以通过验证运行以下命令时是否输出“ s3:ObjectCreated:* ”来检查您是否为存储桶通知配置了正确的事件类型：

 $ mc event ls local/detect-inference arn:minio:sqs::<YOURFUNCTIONNAME>:webhook

有关将存储桶事件发布到 Webhook 的更详细说明，请查看文档。我们现在准备在全新图像上尝试对象检测！

尝试我们的推理系统

这是我想要进行推理的新图像（标题为“1.png”）：

我将新图像放入“检测推理”存储桶中：

几乎立刻，我就可以在 Flask 服务器上看到以下结果：

 $ python detection_server.py * Serving Flask app 'detection_server' * Debug mode: off * Running on http://127.0.0.1:5000 Press CTRL+C to quit image 1/1 /var/folders/xf/q7x0z8sn5nvckccp1g0m1vpm0000gn/T/tmpo6jx3w8u/1.png: 448x736 2 SU-30s, 101.0ms Speed: 4.1ms preprocess, 101.0ms inference, 5.8ms postprocess per image at shape (1, 3, 448, 736) {'results': [[1927.78369140625, 627.7123413085938, 1995.090576171875, 715.3443603515625, 0.8142037987709045, 0.0], [1735.740234375, 477.2108154296875, 1809.181640625, 555.767578125, 0.7766116261482239, 0.0]]} 127.0.0.1 - - [14/Sep/2023 15:39:21] "POST / HTTP/1.1" 200 -

请注意，结果列表中每个检测到的边界框均采用 YOLO 格式 [x1, y1, x2, y2, 概率, 类别]。以下是叠加在原始图像上的边界框和预测类别：

结论性思考

我们做到了！我们介绍了 MinIO 和 CVAT 如何结合在一起以保证我们收集的图像样本的安全和可用，以及如何创建自定义对象检测数据集。然后，我们为自定义任务训练了自己的自定义 YOLO 模型。最后，我们仅用 50 多行代码，就使用 MinIO Bucket Notifications 构建了一个推理服务器，该服务器可以通过我们自定义训练的对象检测模型运行新图像。

此外，对于计算机视觉的大多数关键任务应用，最好在边缘进行推理。否则，所述应用程序很容易受到与将新数据上传到公共云并等待云中的推理服务器返回答案相关的延迟的影响，更不用说网络连接错误的风险了。因此，以 MinIO 为中心的计算机视觉管道作为数据层更有意义。想象一下，一架无人机飞越机场，能够通过完全机载的硬件和软件捕获、存储和使用我们训练有素的模型来处理新图像。通过本地部署 MinIO 服务器，我们在文章末尾构建的基于 Bucket 通知的推理系统非常适合此场景以及无数其他类似场景。

如果您有任何疑问，请加入我们的Slack 频道或通过[email protected]给我们留言。我们是来帮你的。