调用已有模型，检测视频中的物体并保存结果

1 用途

我们有时需要调用已有模型，帮助我们完成检测或者识别任务，并将结果保存

2 举例

已Yolo V9为例:

#导入所需要的库
cv2
from ultralytics import YOLO

# 加载模型，模型需要提前下载好
model = YOLO("yolov9c.pt")

#定义一个用于推理预测的函数
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])), (0, 0, 255), 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_SIMPLEX, 0.5, (0, 0, 255), text_thickness)
    return img, results

#读取一幅图像
image = cv2.imread("t01.jpg")
#设置图像的宽度
new_width = 640
new_height = 640

#使用cv2.resize()改变图像的分辨率
image = cv2.resize(image, (new_width, new_height))
#检测图像的内容
result_img, _ = predict_and_detect(model, image, classes=[], conf=0.5)
#显示并保存图像
cv2.imshow("Image", result_img)
cv2.imwrite("YourSavePath", result_img)
cv2.waitKey(0)

3、视频的的检测与保存

#定义一个用于保存视频的对象
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))
    # initialize the FourCC and a video writer object
    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    writer = cv2.VideoWriter(output_filename, fourcc, fps,
                             (frame_width, frame_height))
    return writer


#读取视频
video_path = r"t03.mp4"
cap = cv2.VideoCapture(video_path)

# 定义输出的视频
output_filename = "output.mp4"
writer = create_video_writer(cap, output_filename)


while True:
#读取视频
    success, img = cap.read()
    if not success:
        break
    #推理检测
    result_img, _ = predict_and_detect(model, img, classes=[0], conf=0.5)
    
    #写入视频
    writer.write(result_img)
    cv2.imshow("Image", result_img)
    
    #每1ms 检测一次安检
    key=cv2.waitKey(1)
    if key == 27:  # 如果按下Esc键
        break

writer.release()

4. 运行结果