文章目录
一、How to Train YOLOv9 on a Custom Dataset
1.1 开始之前
让我们确保我们可以访问 GPU。 我们可以使用 nvidia-smi 命令来做到这一点。
!nvidia-smi
注意:为了让我们更容易管理数据集、图像和模型,我们创建了一个 HOME 常量。
import os
HOME = os.getcwd()
print(HOME)
1.2 配置
注意:YOLOv9 非常新。 目前,我们建议使用主存储库的分支。 detector.py 脚本包含一个阻止推理的错误。 该错误已在 fork 中修复。
!git clone https://github.com/SkalskiP/yolov9.git
%cd yolov9
!pip install -r requirements.txt -q
注意:让我们安装 roboflow 包,我们将使用它从 Roboflow Universe 下载数据集。
!pip install -q roboflow
1.3 下载模型权重
注意:在 YOLOv9 论文中,还提到了 yolov9-s 和 yolov9-m 版本,但 YOLOv9 存储库中尚未提供这些模型的权重。
!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c.pt
!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt
!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c.pt
!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-e.pt
!ls -la {HOME}/weights
输出为:
total 402440
drwxr-xr-x 2 root root 4096 Feb 29 11:06 .
drwxr-xr-x 1 root root 4096 Feb 29 11:06 ..
-rw-r--r-- 1 root root 51508261 Feb 18 12:36 gelan-c.pt
-rw-r--r-- 1 root root 117203713 Feb 18 12:36 gelan-e.pt
-rw-r--r-- 1 root root 103153312 Feb 18 12:36 yolov9-c.pt
-rw-r--r-- 1 root root 140217688 Feb 18 12:36 yolov9-e.pt
1.4 下载示例数据
注意:如果您想使用自己的文件作为输入来运行推理,只需将图像上传到 Google Colab 并使用指向您的文件的路径更新 SOURCE_IMAGE_PATH 即可。
!wget -P {HOME}/data -q https://media.roboflow.com/notebooks/examples/dog.jpeg
SOURCE_IMAGE_PATH = f"{HOME}/dog.jpeg"
1.5 使用预训练的 COCO 模型进行检测
1.5.1 gelan-c
!python detect.py --weights {HOME}/weights/gelan-c.pt --conf 0.1 --source {HOME}/data/dog.jpeg --device 0
输出为:
detect: weights=['/content/weights/gelan-c.pt'], source=/content/data/dog.jpeg, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.1, iou_thres=0.45, max_det=1000, device=0, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5 🚀 1e33dbb Python-3.10.12 torch-2.1.0+cu121 CUDA:0 (Tesla T4, 15102MiB)
Fusing layers...
Model summary: 467 layers, 25472640 parameters, 0 gradients, 102.8 GFLOPs
image 1/1 /content/data/dog.jpeg: 640x384 1 person, 1 car, 1 dog, 1 handbag, 235.8ms
Speed: 0.6ms pre-process, 235.8ms inference, 704.3ms NMS per image at shape (1, 3, 640, 640)
Results saved to runs/detect/exp
注意:默认情况下,每个后续推理会话的结果保存在 {HOME}/yolov9/runs/detect/ 中名为 exp、exp2、exp3… 的目录中。您可以使用 --name 参数覆盖此行为 。
from IPython.display import Image
Image(filename=f"{HOME}/yolov9/runs/detect/exp/dog.jpeg", width=600)
1.5.2 yolov9-e
!python detect.py --weights {HOME}/weights/yolov9-e.pt --conf 0.1 --source {HOME}/data/dog.jpeg --device 0
输出为:
detect: weights=['/content/weights/yolov9-e.pt'], source=/content/data/dog.jpeg, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.1, iou_thres=0.45, max_det=1000, device=0, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5 🚀 1e33dbb Python-3.10.12 torch-2.1.0+cu121 CUDA:0 (Tesla T4, 15102MiB)
Fusing layers...
Model summary: 1119 layers, 69470144 parameters, 0 gradients, 244.0 GFLOPs
image 1/1 /content/data/dog.jpeg: 640x384 1 person, 1 car, 1 dog, 1 backpack, 1 handbag, 169.4ms
Speed: 0.5ms pre-process, 169.4ms inference, 489.2ms NMS per image at shape (1, 3, 640, 640)
Results saved to runs/detect/exp2
from IPython.display import Image
Image(filename=f"{HOME}/yolov9/runs/detect/exp2/dog.jpeg", width=600)
1.6 验证并下载数据集
注意:数据集必须保存在{HOME}/yolov9目录中,否则训练将不会成功。
%cd {HOME}/yolov9
注意:在本教程中,我将使用足球运动员检测数据集。 请随意将其替换为 YOLO 格式的数据集,或使用 Roboflow Universe 上提供的其他数据集。 此外,如果您计划在训练后将模型部署到 Roboflow,请确保您是数据集的所有者,并且没有模型与您要训练的数据集版本关联。
import roboflow
roboflow.login()
rf = roboflow.Roboflow()
project = rf.workspace("roboflow-jvuqo").project("football-players-detection-3zvbc")
version = project.version(8)
dataset = version.download("yolov9")
1.7 训练模型
%cd {HOME}/yolov9
!python train.py \
--batch 16 --epochs 25 --img 640 --device 0 --min-items 0 --close-mosaic 15 \
--data {dataset.location}/data.yaml \
--weights {HOME}/weights/gelan-c.pt \
--cfg models/detect/gelan-c.yaml \
--hyp hyp.scratch-high.yaml
1.8 检查训练结果
注意:默认情况下,每个后续训练课程的结果都保存在 {HOME}/yolov9/runs/train/ 中名为 exp、exp2、exp3… 的目录中。您可以使用 --name 参数覆盖此行为 。
!ls {HOME}/yolov9/runs/train/exp/
输出为:
confusion_matrix.png P_curve.png train_batch2.jpg
events.out.tfevents.1709205293.f80636ca2632.2700.0 PR_curve.png val_batch0_labels.jpg
F1_curve.png R_curve.png val_batch0_pred.jpg
hyp.yaml results.csv val_batch1_labels.jpg
labels_correlogram.jpg results.png val_batch1_pred.jpg
labels.jpg train_batch0.jpg weights
opt.yaml train_batch1.jpg
from IPython.display import Image
Image(filename=f"{HOME}/yolov9/runs/train/exp/results.png", width=1000)
from IPython.display import Image
Image(filename=f"{HOME}/yolov9/runs/train/exp/confusion_matrix.png", width=1000)
from IPython.display import Image
Image(filename=f"{HOME}/yolov9/runs/train/exp/val_batch0_pred.jpg", width=1000)
1.9 验证自定义模型
%cd {HOME}/yolov9
!python val.py \
--img 640 --batch 32 --conf 0.001 --iou 0.7 --device 0 \
--data {dataset.location}/data.yaml \
--weights {HOME}/yolov9/runs/train/exp/weights/best.pt
输出为:
/content/yolov9
val: data=/content/yolov9/football-players-detection-6/data.yaml, weights=['/content/yolov9/runs/train/exp/weights/best.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.7, max_det=300, task=val, device=0, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=runs/val, name=exp, exist_ok=False, half=False, dnn=False, min_items=0
YOLOv5 🚀 1e33dbb Python-3.10.12 torch-2.1.0+cu121 CUDA:0 (Tesla T4, 15102MiB)
Fusing layers...
gelan-c summary: 467 layers, 25414044 parameters, 0 gradients, 102.5 GFLOPs
val: Scanning /content/yolov9/football-players-detection-6/valid/labels.cache... 38 images, 0 backgrounds, 0 corrupt: 100% 38/38 [00:00<?, ?it/s]
Class Images Instances P R mAP50 mAP50-95: 100% 2/2 [00:04<00:00, 2.22s/it]
all 38 905 0.91 0.689 0.759 0.528
ball 38 35 1 0.156 0.322 0.0889
goalkeeper 38 27 0.967 0.889 0.937 0.746
player 38 754 0.948 0.942 0.981 0.758
referee 38 89 0.725 0.771 0.797 0.517
Speed: 0.1ms pre-process, 31.6ms inference, 23.8ms NMS per image at shape (32, 3, 640, 640)
Results saved to runs/val/exp
1.10 推理
!python detect.py \
--img 1280 --conf 0.1 --device 0 \
--weights {HOME}/yolov9/runs/train/exp/weights/best.pt \
--source {dataset.location}/test/images
注意:就像以前一样,推理结果已保存在 {HOME}/yolov9/runs/detect/ 内的相应目录中。 让我们检查其中的一些结果。
import glob
from IPython.display import Image, display
for image_path in glob.glob(f'{HOME}/yolov9/runs/detect/exp3/*.jpg')[:2]:
display(Image(filename=image_path, width=600))
1.11 使用推理部署 YOLOv9 模型
注意:要部署模型并显示推理结果,我们需要两个额外的包 - 推理和监督。 让我们安装并导入它们!
!pip install -q inference supervision
import cv2
import random
import getpass
import supervision as sv
from inference import get_model
%matplotlib inline
注意:在推理中使用模型之前,您首先需要将权重上传到 Roboflow Universe。 确保指定正确的 model_type - yolov9,并且项目版本与您用于训练的数据集版本匹配,由 [1] 表示。 就我而言,是 6。
version.deploy(model_type="yolov9", model_path=f"{HOME}/yolov9/runs/train/exp")
输出为:
View the status of your deployment at: https://app.roboflow.com/roboflow-jvuqo/football-players-detection-3zvbc/6
Share your model with the world at: https://universe.roboflow.com/roboflow-jvuqo/football-players-detection-3zvbc/model/6
注意:现在我们可以使用 [2] 指定的 model_id 在任何地方下载我们的模型。 就我而言,足球运动员检测-3zvbc/6。 要下载模型,您需要 ROBOFLOW_API_KEY。
ROBOFLOW_API_KEY = getpass.getpass()
model = get_model(model_id="football-players-detection-3zvbc/8", api_key=ROBOFLOW_API_KEY)
注意:让我们从测试子集中选择随机图像,并使用新微调的模型检测对象。
image_paths = sv.list_files_with_extensions(
directory=f"{dataset.location}/test/images",
extensions=['png', 'jpg', 'jpeg']
)
image_path = random.choice(image_paths)
image = cv2.imread(image_path)
result = model.infer(image, confidence=0.1)[0]
detections = sv.Detections.from_inference(result)
label_annotator = sv.LabelAnnotator(text_color=sv.Color.BLACK)
bounding_box_annotator = sv.BoundingBoxAnnotator()
annotated_image = image.copy()
annotated_image = bounding_box_annotator.annotate(scene=annotated_image, detections=detections)
annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
sv.plot_image(annotated_image)