使用YOLOv9进行图像与视频检测

大家好,YOLOv9 与其前身v8一样,专注于识别和精确定位图像和视频中的对象。本文将介绍如何使用YOLOv9进行图像与视频检测,自动驾驶汽车、安全系统和高级图像搜索等应用在很大程度上依赖于此功能,YOLOv9 引入了比 YOLOv8 更令人印象深刻的创新点。

1.安装必要的库

python 复制代码
pip install opencv-python ultralytics

2.导入库

python 复制代码
import cv2
from ultralytics import YOLO

3.选择模型型号尺寸

python 复制代码
model = YOLO("yolov9c.pt")

这里我们选择yolov9c.pt,大家可以选择不同的模型尺寸进行检测,并比较不同的型号并权衡它们各自的优缺点。

4.编写函数预测和检测图像和视频中的对象

python 复制代码
def predict(chosen_model, img, classes=[], conf=0.5):
    if classes:
        results = chosen_model.predict(img, classes=classes, conf=conf)
    else:
        results = chosen_model.predict(img, conf=conf)

    return results

def predict_and_detect(chosen_model, img, classes=[], conf=0.5, rectangle_thickness=2, text_thickness=1):
    results = predict(chosen_model, img, classes, conf=conf)
    for result in results:
        for box in result.boxes:
            cv2.rectangle(img, (int(box.xyxy[0][0]), int(box.xyxy[0][1])),
                          (int(box.xyxy[0][2]), int(box.xyxy[0][3])), (255, 0, 0), rectangle_thickness)
            cv2.putText(img, f"{result.names[int(box.cls[0])]}",
                        (int(box.xyxy[0][0]), int(box.xyxy[0][1]) - 10),
                        cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), text_thickness)
    return img, results

predict() 这个函数采用三个参数:

  • chosen_model :用于预测的训练模型

  • img :要进行预测的图像

  • classes :(可选)要将预测筛选到的类名列表

  • conf :(可选)要考虑的预测的最小置信度阈值

函数首先检查是否提供classes参数。如果是,则使用classes参数调用该chosen_model.predict() 方法,该参数仅将预测筛选为这些类。否则,将调用该 chosen_model.predict() 方法时不带 classes 参数,该参数将返回所有预测。

conf 参数用于筛选出置信度分数低于指定阈值的预测。这对于消除误报很有用。

该函数返回预测结果列表,其中每个结果都包含以下信息:

  • name :预测类的名称

  • conf :预测的置信度分数

  • box :预测对象的边界框

predict_and_detect() 函数采用与 predict() 函数相同的参数,但除了预测结果外,它还返回带注释的图像。

该函数首先调用该 predict() 函数以获取预测结果。然后,它循环访问预测结果,并在每个预测对象周围绘制一个边界框。预测类的名称也写在边界框上方。

该函数返回一个包含带注释的图像和预测结果的元组。

以下是这两个函数之间差异的摘要:

  • predict() 函数仅返回预测结果,而该 predict_and_detect() 函数还返回带注释的图像。

  • predict_and_detect() 函数是 predict() 函数的包装器,这意味着它在内部调用函数 predict()

5.使用 YOLOv9 检测图像

python 复制代码
# read the image
image = cv2.imread("YourImagePath")
result_img, _ = predict_and_detect(model, image, classes=[], conf=0.5)

如果要检测特定类,只需在类列表classes中输入对象的 ID 号即可。

6.保存并绘制结果图像

python 复制代码
cv2.imshow("Image", result_img)
cv2.imwrite("YourSavePath", result_img)
cv2.waitKey(0)

7.使用 YOLOv9 检测视频

python 复制代码
video_path = r"YourVideoPath"
cap = cv2.VideoCapture(video_path)
while True:
    success, img = cap.read()
    if not success:
        break
    result_img, _ = predict_and_detect(model, img, classes=[], conf=0.5)
    cv2.imshow("Image", result_img)
    
    cv2.waitKey(1)

8.保存结果视频

python 复制代码
# 定义保存函数
def create_video_writer(video_cap, output_filename):
    # grab the width, height, and fps of the frames in the video stream.
    frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(video_cap.get(cv2.CAP_PROP_FPS))
    #初始化
    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    writer = cv2.VideoWriter(output_filename, fourcc, fps,
                             (frame_width, frame_height))
    return writer

只需使用上面的函数和代码即可:

python 复制代码
output_filename = "YourFilename"
writer = create_video_writer(cap, output_filename)

video_path = r"YourVideoPath"
cap = cv2.VideoCapture(video_path)
while True:
    success, img = cap.read()
    if not success:
        break
    result_img, _ = predict_and_detect(model, img, classes=[], conf=0.5)
    writer.write(result_img)
    cv2.imshow("Image", result_img)
    
    cv2.waitKey(1)
writer.release()
相关推荐
湫ccc17 分钟前
《Python基础》之pip换国内镜像源
开发语言·python·pip
fhvyxyci18 分钟前
【C++之STL】摸清 string 的模拟实现(下)
开发语言·c++·string
hakesashou19 分钟前
Python中常用的函数介绍
java·网络·python
qq_4597300320 分钟前
C 语言面向对象
c语言·开发语言
菜鸟的人工智能之路28 分钟前
极坐标气泡图:医学数据分析的可视化新视角
python·数据分析·健康医疗
菜鸟学Python29 分钟前
Python 数据分析核心库大全!
开发语言·python·数据挖掘·数据分析
小白不太白95031 分钟前
设计模式之 责任链模式
python·设计模式·责任链模式
喜欢猪猪36 分钟前
Django:从入门到精通
后端·python·django
一个小坑货36 分钟前
Cargo Rust 的包管理器
开发语言·后端·rust
bluebonnet2741 分钟前
【Rust练习】22.HashMap
开发语言·后端·rust