​训练自定义数据集 YOLOv8 OBB - 定向边界框

点击下方卡片，关注“小白玩转Python”公众号

YOLOv8 定向目标检测

什么是定向边界框目标检测？

“定向边界框目标检测”指的是在图像或视频中识别和定位对象的同时，还估计它们的方向。这种方法在处理有一定旋转或方向变化的对象时特别有用，例如汽车、飞机或文本。

定向目标检测器的输出是一组旋转的边界框，这些框精确地包围图像中的对象，每个框都有类别标签和置信度分数。当您需要识别场景中感兴趣的对象，但不需要知道对象的确切位置或确切形状时，目标检测是一个不错的选择。

来源：https://mediaan.com/mediaan-blog/oriented-object-detection

如何加载 YOLOv8 OBB 模型？

# install ultralytics


pip3 install ultralytics

from ultralytics import YOLO


# Load a model
model = YOLO('yolov8n-obb.pt')  # load an official model


# Predict with the model
results = model('https://miro.medium.com/v2/resize:fit:1400/1*BybWklt3t0M8zdSXtkipRg.png', save=True)  # predict on an image

obb 结果

如何训练自定义数据集？

数据集格式

YOLO OBB 格式通过它们的四个角点指定边界框，坐标值在 0 和 1 之间归一化。它遵循这种格式：

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

内部，YOLO 以 xywhr 格式处理损失和输出，这表示边界框的中心点（xy）、宽度、高度和旋转。

来源：https://docs.ultralytics.com/datasets/obb/#yolo-obb-format

使用 Roboflow YOLOv8 OBB 数据集格式。在本教程中，我使用了一个公共数据集，并从 Roboflow Universe 生成了 obb 格式。

来源：https://universe.roboflow.com/ 

数据集链接：https://universe.roboflow.com/prata-vqsjq/dosr/dataset/3

导出数据集到 YOLOv8 定向边界框格式：

标签

0: cargo
  1: container
  2: cruise
  3: fishing boat
  4: military_ship
  5: tanker

如何训练？

# Train the model
epochs = 20


imgsz = 640
batch = 16
save_period = 5
device = 0
cache = False
pretrained = True
YAML_FILE = '/content/datasets/data.yaml'
project = 'obb'
name = 'custom_obb'


model.train(
    data=YAML_FILE,
    device=device,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch,
    cache=cache,
    save_period=save_period,
    project=project,
    name=name,
    patience=0,
    augment=False
)

训练过程将运行 20 个epoch，每批 16 张图像，图像大小为 640 像素。脚本每 5 个epoch保存一次训练好的模型，并使用 GPU 进行训练（device=0）。训练过程中不使用缓存（cache=False）。训练过程不使用预训练权重（pretrained=False）且不使用数据增强（augment=False）。项目名为‘obb’，训练好的模型将保存为‘custom_obb’。

结果

·  END  ·

🌟 想要变身计算机视觉小能手？快来「小白玩转Python」公众号！

回复“Python视觉实战项目”，解锁31个超有趣的视觉项目大礼包！🎁

本文仅供学习交流使用，如有侵权请联系作者删除