定向边界框(Oriented Bounding Boxes)是一种旋转的边界框,用于更精确地适应倾斜角度的物体。例如,在药片检测数据集中,使用YOLOv5-OBB可以更紧密、更准确地检测到倾斜的药片,避免将多个药片或其他物体错误地包含在一个边界框内。
本教程将指导如何通过导出已经标注的数据集来使用YOLOv5-OBB。将使用Roboflow平台来管理和标注数据集,并将其转换为适用于不同模型的26种以上格式。Roboflow免费提供高达10,000张图片的服务,基于云端,团队使用起来非常方便。
可以使用Roboflow来管理数据集、标注数据,并将其转换为YOLOv5-OBB所需的格式。Roboflow提供了免费的服务,并且支持团队协作。也可以使用这里使用的多边形标注的美国手语数据集,或者查看Roboflow Universe以获取更多的计算机视觉数据集。
为了跟随本教程,可以在新的标签页中打开这个。
在确认版本要求后,将运行以下代码来安装YOLOv5-OBB所需的依赖。首先,需要安装特定版本的PyTorch和相关库:
pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
nvcc -V
python
然后,需要导入torch库并确认CUDA运行时api版本至少为10.0:
import torch
torch.version.cuda
接下来,将克隆YOLOv5-OBB的代码库,并安装所需的依赖:
!git clone https://github.com/hukaixuan19970627/yolov5_obb.git
%cd /content/yolov5_obb/
!pip install -r requirements.txt
在utils/nms_rotated文件夹中运行setup.py develop:
%cd utils/nms_rotated
!python setup.py develop #或 "pip install -v -e ."
接下来,将创建一个名为datasets的目录,并进入该目录以运行代码片段。将使用Roboflow来导出数据集。以下是使用实例分割的美国手语数据集的示例,可以将其替换为自己的数据集信息,方法是导出YOLOv5定向边界框格式的数据集。
%cd /content/datasets/
!pip install roboflow
from roboflow import Roboflow
rf = Roboflow(api_key="YOUR_API_KEY")
project = rf.workspace("YOUR_WORKSPACE_NAME").project("YOUR_PROJECT_NAME")
dataset = project.version(4).download("yolov5-obb")
运行上述代码后,当运行!python train.py时,预期的数据集名称将是roboflow。因此,请确保将数据集文件夹重命名为roboflow以避免触发任何错误。换句话说,确保数据集结构如下:
├── yolov5
└── datasets
└── roboflow
├── test
├── train
└── valid
├── images
|────1.jpg
|────...
└────10000.jpg
├── labelTxt
|────1.txt
|────...
└────10000.txt
└────README.roboflow.txt
└────data.yaml
在/content/datasets/roboflow/目录下运行训练命令:
%cd /content/yolov5_obb/
!python train.py --data /content/datasets/roboflow/data.yaml --epochs 10 --batch-size 1 --img 1024 --device 0 --exist-ok
训练完成后,使用以下代码运行detect.py:
!python detect.py --weights 'runs/train/exp/weights/best.pt' --source '../datasets/roboflow/valid/images' --conf-thres 0.01 --img 416 --device 0 --agnostic
detect.py完成后,选择一些随机图像以查看exp文件夹中给定图像上的定向边界框检测结果。
import os, random
random_file = random.choice(os.listdir("/content/yolov5_obb/runs/detect/exp"))
from IPython.display import Image
Image(random_file)