【计算机视觉】YOLOv9代码使用示例（含代码）

浅浅的花香味﹌ 2024-05-05 22:24 113阅读 0赞

#### 文章目录 ####

*  一、How to Train YOLOv9 on a Custom Dataset
 *   *  1.1 开始之前
     *  1.2 配置
     *  1.3 下载模型权重
     *  1.4 下载示例数据
     *  1.5 使用预训练的 COCO 模型进行检测
     *   *  1.5.1 gelan-c
         *  1.5.2 yolov9-e
     *  1.6 验证并下载数据集
     *  1.7 训练模型
     *  1.8 检查训练结果
     *  1.9 验证自定义模型
     *  1.10 推理
     *  1.11 使用推理部署 YOLOv9 模型

## 一、How to Train YOLOv9 on a Custom Dataset ##

### 1.1 开始之前 ###

让我们确保我们可以访问 GPU。 我们可以使用 nvidia-smi 命令来做到这一点。

!nvidia-smi

![在这里插入图片描述][ad1ad0fcd2b74a6d9a94e5d5b1b85860.png]  
注意：为了让我们更容易管理数据集、图像和模型，我们创建了一个 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)

![在这里插入图片描述][40ae2a83ee7041a796c8fa74ed6d1e1e.png]

#### 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)

![在这里插入图片描述][ae551f5f863640629c8357ec2b05c86f.png]

### 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)

![在这里插入图片描述][6f1e2e8e2afb47f8a34c672eef01ab50.png]

from IPython.display import Image
    
    Image(filename=f"{
          HOME}/yolov9/runs/train/exp/confusion_matrix.png", width=1000)

![在这里插入图片描述][f67d5b4cb33d480396d745b39b21edac.png]

from IPython.display import Image
    
    Image(filename=f"{
          HOME}/yolov9/runs/train/exp/val_batch0_pred.jpg", width=1000)

![在这里插入图片描述][e37a2ed9434d425fa17f4ca8a79a7057.png]

### 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))

![在这里插入图片描述][01019ce982a040fa893fa2708e430c84.png]

### 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)

![在这里插入图片描述][6f650590c3f0480d8cfdc0acfa281815.png]

[ad1ad0fcd2b74a6d9a94e5d5b1b85860.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/42b420f0a4f74445957d3d025a7ab0c3.png
[40ae2a83ee7041a796c8fa74ed6d1e1e.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/3a07386a43ff4b57a9758063feade43d.png
[ae551f5f863640629c8357ec2b05c86f.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/38fa25112ccd436eb496a3f68bc410e7.png
[6f1e2e8e2afb47f8a34c672eef01ab50.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/c7e412644ca94412ad2c2165227cb38c.png
[f67d5b4cb33d480396d745b39b21edac.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/1cfcef990b1e4a5890e6cab9d0ab0496.png
[e37a2ed9434d425fa17f4ca8a79a7057.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/f218b2b9afd044799a7ed619ee9fdb16.png
[01019ce982a040fa893fa2708e430c84.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/def069c62cb8438291838d9eac9793c7.png
[6f650590c3f0480d8cfdc0acfa281815.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/05/b21a3851c7034e7ea7296652b12091b7.png