实战 | 本地GPU训练YOLOv8带方向的目标检测

前 言

目标检测是指在图像或视频帧内识别和定位物体的任务。定向目标检测具体涉及检测具有定义方向或旋转的物体，例如检测具有不同角度的车辆或检测自然场景中具有不同方向的文本。

此任务通常需要专门的算法或模型，能够准确检测和识别图像中的物体，无论其方向或旋转如何。定向物体检测可应用于各个领域，包括自动驾驶、文档分析和工业自动化。

在本文中，我们将学习如何使用 YOLOv8 进行有向物体检测。

官方文档可以在下面链接找到：

https://docs.ultralytics.com/tasks/obb/

从计算上来说，这比传统的物体检测稍微昂贵一些，因为传统的物体检测中所有框都是完美的矩形，没有任何方向。这是因为 OBB 除了需要检测物体外，还需要检测方向。

首先，使用 pip 安装 ultralytics：

pip install ultralytics

导入包和模型：

import ultralytics
from ultralytics import YOLO
# Load a model
model = YOLO('yolov8x-obb.pt')  # load an official model

这将导入 Ultralytics 包并下载预先训练的 yolov8 OBB 模型。在本例中，我使用了 yolov8x-obb.pt。其他模型可以参考链接：

https://docs.ultralytics.com/tasks/obb/#visual-samples

我们使用预训练模型对样本图像进行样本检测。请注意，这些模型是在DOTA数据集上训练的。

https://universe.roboflow.com/felipe-coradesque-6gmum/dota-aerial-images

# Predict with the model
results = model('sample.jpg',imgsz=640, save=True)  # predict on an image

默认情况下，预测结果会保存在 runs/obb/predict 文件夹中。请注意，每次运行的预测结果都会带有后缀，例如 predict1 或 predict2。

我们可以看到，预测结果有一个根据物体方向而定的有角度的矩形，而不是像传统物体检测方法中的直线矩形。

让我们在 DOTA 数据集上进一步训练模型，这样我们就可以了解训练数据和标注（ground truth）是什么样的。

为了使用 YOLOV8 训练模型，我们需要一个 .yml 文件。

我们可以从 Ultralytics Github Repo 中的ultralytics/cfg/datasets下载此 DOTA 数据集的文件。当我们运行以下代码时，这将自动完成。

运行之前，如果您有 GPU 并且想要使用 GPU 来运行训练，请按照以下额外步骤操作。首先，卸载 torch 和作为 ultralytics 的一部分下载的 torchvision。它们对应于 CPU 版本。仅当您计划运行 GPU 进行训练时，才运行以下代码。

pip uninstall torch torchvision

这将卸载 torch 和 torchvision。现在是时候安装支持 CUDA 的 torch 和 torchvision 了。

在我的计算机上，我使用的是 CUDA 12.1 版本，因此我使用此命令安装 Torch。

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

我将给出以下包含 CPU 和 GPU 的代码。首先，对于 CPU，这非常简单，我们可以直接运行训练而不通过设备，或者我们可以明确将其称为设备“cpu”

# CPU
results = model.train(data='DOTAv1.yaml', epochs=5, imgsz=640, device=cpu)

# GPU- Use the device=0 for Single GPU Training. 
if __name__=='__main__':
    results = model.train(data='DOTAv1.yaml', epochs=5, imgsz=640, device=0)

一旦我们运行此文件，训练和测试图像及其注释将自动下载。这是因为在“DOTAv1.yaml”文件中我们有一行用于下载文件。

Download script/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/DOTAv1.zip

OBB 标注文件

OBB 标注使用的格式如下，

class_index, x1, y1, x2, y2, x3, y3, x4, y4

图像显示了每个角的所有点的 x1、y1 等。在此示例中，我们使用基于点的 OBB。

请注意，YOLO OBB 格式通过四个角点指定边界框，其坐标在 0 到 1 之间标准化。因此，所有点的 x1、x2 等都将在 0 到 1 之间标准化。

所有训练结果都将存储在runs\obb\train7文件夹中。这里是 train 7，表示我的第 7 次训练。开始训练时将自动创建这些文件夹。

权重文件将自动存储在runs\obb\train7\weights文件夹中。

预测

我们运行下面的代码来获得预测：

from ultralytics import YOLO

# Load a model
model = YOLO('runs/obb/train7/weights/best.pt')  # load a custom model

# Predict with the model
results = model('test_prediction',imgsz=640, save=True)  # predict on an image

这里的“test_prediction”是包含我们要测试的预测图像的文件夹。

一旦我们运行预测，预测图像就会自动存储在runs\obb\predict3文件夹中。再次，这个 predict3 表示这是第三次运行预测函数。每个人的情况可能都不一样。

以下是预测结果：

由于我们又训练了 5个epoch，与预训练模型的预测相比，置信度得分略有增加。