如何在本地GPU上使用YOLOv8进行目标检测

2024年03月19日 由 alex 发表 155 0

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


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


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


2


官方文档见 https://docs.ultralytics.com/tasks/obb/


在计算上,这比传统的对象检测(所有方框都是完美的矩形,没有任何方向)要昂贵一些。这是因为 OBB 除了检测物体外,还需要检测方向。


首先,用 pip 安装 ultralytics


pip install ultralytics


在本文中,我使用的是 Ultralytics 8.1.29 版。


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 中还指定了其他模型。


3


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


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


4


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


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


让我们在 DOTA 数据集上进一步训练模型,这样我们就能了解训练数据和注释(地面实况)是怎样的。


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


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


在运行之前,如果你有 GPU 并希望使用 GPU 运行训练,请遵循以下额外步骤。首先,卸载作为 ultralytics 一部分下载的 torch 和 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)


5


运行该文件后,训练图像和测试图像及其注释将自动下载。这是因为我们在 "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


6


图像显示了每个角上所有点的 x1、y1 等值。在本例中,我们使用的是基于点的 OBB。


因此,所有点的 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 表示这是第三次运行预测函数。每个人的情况可能不同。


以下是预测图像、


7


由于我们多训练了 5 个历元,因此置信度得分比使用预训练模型预测时略有提高。

文章来源:https://medium.com/gitconnected/train-oriented-object-detection-using-yolo-v8-on-local-gpu-55461a3da2fd
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消