如何把yolov5检测视频后生成图像保存

要将 YOLOv5 用于视频中的对象检测，并将检测到的目标保存为图像，请按照以下步骤操作：

1. 安装 YOLOv5 并下载预训练权重。可以从 YOLOv5 的 GitHub 存储库中下载预训练权重。

2. 使用 OpenCV 打开视频文件，可以按照以下示例代码执行此操作：

import cv2

# Open the video file
cap = cv2.VideoCapture('video.mp4')

3. 在视频帧上运行 YOLOv5 检测，可以按照以下示例代码执行此操作：

import torch
import numpy as np
from yolov5.models.experimental import attempt_load
from yolov5.utils.torch_utils import select_device
from yolov5.utils.general import non_max_suppression, scale_coords, plot_one_box

# Load YOLOv5 model
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))
model.eval()

# Set device
device = select_device('cpu')

# Define detection threshold
conf_thresh = 0.5

# Iterate over video frames
while cap.isOpened():
    # Read a frame from the video file
    ret, frame = cap.read()
    if not ret:
        break
    
    # Convert the frame from BGR to RGB
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    
    # Resize the frame to the size expected by the YOLOv5 model
    img = cv2.resize(frame, (640, 640))
    img = img.transpose(2, 0, 1)  # HWC to CHW
    img = np.expand_dims(img, axis=0)  # Add batch dimension
    
    # Convert the image to a PyTorch tensor
    img = torch.from_numpy(img).to(device)
    
    # Run YOLOv5 detection
    with torch.no_grad():
        detections = model(img)[0]
        detections = non_max_suppression(detections, conf_thresh, 0.4)
    
    # Iterate over the detected objects and draw bounding boxes on the frame
    for det in detections:
        if det is not None and len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
            for *xyxy, conf, cls in reversed(det):
                label = f'{model.names[int(cls)]} {conf:.2f}'
                plot_one_box(xyxy, frame, label=label, color=(0, 255, 0))
    
    # Display the frame with bounding boxes drawn
    cv2.imshow('frame', cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    
    # Press 'q' to quit
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

4. 将检测到的目标保存为图像，可以按照以下示例代码执行此操作：

# Iterate over the detected objects and save each one as an image
for i, det in enumerate(detections):
    if det is not None and len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
        for *xyxy, conf, cls in reversed(det):
            label = f'{model.names[int(cls)]} {conf:.2f}'
            # Save the object as an image
            crop_img = frame[int(xyxy[1]):int(xyxy[3]), int(xyxy[0]):int(xyxy[2])]
            cv2.imwrite(f'object_{i}.jpg', crop_img)

在这个示例代码中，我们遍历检测到的对象，使用 `cv2.imwrite()` 将每个对象保存为单独的图像文件。文件名中包含对象的编号。