YOLOv8 OBB模型训练指南

在许多模型中,例如Ultralytics YOLOv8,边界框坐标是水平对齐的。这意味着在倾斜对象周围会有空间。下图显示了太阳能电池板上的目标检测预测:

可以通过训练面向边界框的目标检测模型来检索与倾斜对象边缘匹配的边界框。在本指南中,将展示如何使用自定义数据集训练YOLOv8 Oriented Bounding Boxes (YOLOv8-OBB)模型。

步骤1:收集数据

在能够训练模型之前,需要一个数据集。对于本指南,将训练一个模型来检测太阳能电池板。这样的模型可以被用于空中测绘,由测绘组织使用,以更好地了解某个区域内太阳能电池板的采用情况。

为了收集数据,将使用来自Roboflow Universe的数据集,这是一个社区,超过250,000个公共计算机视觉数据集已在该网站上共享。将使用Aerial Solar Panels数据集。该数据集包含53张太阳能电池板的航拍图像。

虽然此数据集包含注释,但注释是边界框,而不是面向边界框。因此,仍然需要标记所有图像。如果有其他用例,可以使用Roboflow Universe来寻找用于训练模型的图像。如果已经有用于训练模型的图像,可以跳到下一步!

步骤2:创建Roboflow项目

需要用面向边界框标记图像,以便模型确切地知道图像中感兴趣对象的位置。将在Roboflow Annotate中注释数据。

要开始标记,首先创建一个免费的Roboflow账户。然后,点击Roboflow仪表板上的“创建项目”按钮。创建一个目标检测项目。

创建项目后,系统将要求在Web界面中上传数据。可以通过将想要标记的图像拖入Web应用程序来做到这一点。如果已经有YOLOv8 OBB标签,可以在同一Web界面中与图像一起上传。

如果正在使用太阳能电池板数据集,请上传“train”文件夹中的图像。

不要上传“_annotations.coco.json”文件,因为其中的注释不适用于面向边界框训练。点击“保存并继续”以确认将图像上传到Roboflow平台。

步骤3:使用OBB标签注释数据

图像上传后,就可以开始注释图像了。要开始标记图像,点击左侧边栏中的“注释”。然后,选择想要标记的图像。当选择一个图像时,将被带入Roboflow Annotate Web界面:

对于传统的目标检测模型,会使用一个边界框注释工具,它使能够围绕感兴趣的对象绘制一个水平对齐的框。这些注释将在训练过程中使用,以帮助模型学习如何识别对象。

对于OBB,需要使用一个单独的注释工具:多边形工具。这个工具允许围绕对象绘制任意点。必须围绕想要识别的对象绘制四个点,每个点代表对象的一个角。

要使用多边形工具进行标记,选择左侧边栏中的多边形图标,或按键盘上的P。在图像上点击想要添加多边形的第一个点的位置。然后,围绕感兴趣的对象描绘一个框。

为数据集中的所有图像重复此过程。

一旦标记了所有数据,就可以创建一个数据集版本了。数据集版本是标记数据的快照,它被冻结在时间中。可以使用数据集版本功能来跟踪数据集随时间的变化。

在创建数据集版本时,可以添加增强的图像,并在创建数据集版本时对图像应用预处理步骤。要创建数据集版本,点击Roboflow项目左侧边栏中的“生成”。

将被要求配置数据集版本。在训练/测试/分割部分,目标是将数据在训练、测试和验证集之间进行70/20/10的分割。这将确保有足够的图像示例来训练和评估模型的性能。

在预处理和增强阶段,可以应用像图像调整大小、灰度、动态裁剪这样的预处理步骤,以及像亮度变化这样的增强。

建议第一次模型训练时保留这些值的默认设置。这使能够在没有任何更改的情况下了解数据集在训练模型时的表现。实际上,如果数据集注释得不好或包含不代表用例的图像,增强将无法帮助提高模型性能。

要了解更多关于预处理和增强的信息,请参考图像预处理和增强指南。配置好数据集版本后,点击继续。数据集版本将被生成。

步骤5:导出数据

Roboflow Train不支持训练YOLOv8 OBB模型。话虽如此,可以从Roboflow导出图像数据以用于训练OBB模型。

将在Google Colab笔记本中训练模型。Colab是谷歌提供的一个交互式编程环境。可以使用Colab免费访问T4 GPU。可以使用Colab中提供的免费GPU来训练计算机视觉模型。

制作了一个YOLOv8OBB笔记本以供跟随本指南。打开YOLOv8 OBB笔记本。要开始,需要从Roboflow导出数据到笔记本。可以使用Roboflow Python包来做到这一点。

数据集页面上,点击“导出数据”。选择“获取下载代码”。将出现一个代码片段,可以使用它来导出标记数据集。代码片段将如下所示:

from roboflow import Roboflow rf = Roboflow(api_key="API_KEY") project = rf.workspace("capjamesg").project("metal-defect-detection-d3wl6") dataset = project.version(2).download("yolov8-obb")

API密钥将在下载代码片段时填写。

在可以训练模型之前,需要对数据集中的yolov8-obb.yaml文件进行更改。运行以下代码以将YAML文件中的root值替换为训练模型所需的值:

import yaml with open(f'{dataset.location}/data.yaml', 'r') as file: data = yaml.safe_load(file) data['path'] = dataset.location with open(f'{dataset.location}/data.yaml', 'w') as file: yaml.dump(data, file, sort_keys=False)

步骤4:训练YOLOv8 OBB模型

现在已经准备好了数据集,可以训练一个YOLOv8OBB模型。要开始,需要安装ultralytics pip包。可以使用以下命令来做到这一点:

pip install ultralytics

将以下代码添加到笔记本中以开始训练模型:

from ultralytics import YOLO model = YOLO(“yolov8n-obb.pt”) results = model.train(data=f”{dataset.location}/yolov8-obb.yaml', epochs=100, imgsz=640)

在这段代码中,导入了Ultralytics库,加载了一个预训练的OBB模型,然后使用在Roboflow中标记的自定义数据集来训练该模型。训练作业将运行100个周期。

训练所需的时间将取决于运行的硬件以及数据集中有多少图像。

步骤5:测试模型

一旦模型训练完成,可以在数据集中的图像上测试它。要可视化模型的结果,将使用supervision Python包。supervision提供了一套广泛的实用工具,用于处理视觉模型,包括用于可视化边界框和多边形的注释工具。

要安装supervision,请运行:

pip install supervision

让在以下测试集中的图像上测试模型:

model = YOLO('runs/obb/train2/weights/best.pt') import os import random random_file = random.choice(os.listdir(f"{dataset.location}/test/images")) file_name = os.path.join(f"{dataset.location}/test/images", random_file) results = model(file_name) import supervision as sv import cv2 detections = sv.Detections.from_ultralytics(results[0]) oriented_box_annotator = sv.OrientedBoxAnnotator() annotated_frame = oriented_box_annotator.annotate( scene=cv2.imread(file_name), detections=detections ) sv.plot_image(image=annotated_frame, size=(16, 16))

模型成功地识别了太阳能电池板的位置。与显示边界框不同,模型返回了面向边界框(多边形)。在感兴趣的对象——太阳能电池板——和其周围的框之间几乎没有空间。

以下是没有OBB的常规目标检测模型的结果(左)和具有OBB的模型的结果(右)。这些结果已被裁剪以突出显示感兴趣区域,因为不放大很难看到边界框位置的差异。

在上面的第一张图片中,边界框没有OBB模型结果中那么紧密对齐。

YOLOv8Oriented Bounding Boxes (OBB)模型比传统的YOLOv8模型更紧密地识别感兴趣的对象。面向边界框是倾斜的,以适应框状图像,而没有面向边界框的目标检测模型返回的预测是水平对齐的。

沪ICP备2024098111号-1
上海秋旦网络科技中心:上海市奉贤区金大公路8218号1幢 联系电话:17898875485