寒假学习3——YoLoV5视频帧处理

YOLOv5 是一种目标检测算法，它是由 Alexey Bochkovskiy 开发的 YOLO（You Only Look Once）系列算法的最新版本之一。YOLOv5 采用了单一网络结构，通过单次前向传播即可同时预测多个目标的位置和类别。YOLOv5 在处理视频帧时的基本原理：

视频帧输入：
- YOLOv5 可以接受视频流或视频文件作为输入。视频文件会被逐帧读取，每一帧作为模型的输入。
图像预处理：
- 输入的视频帧会被预处理，通常包括调整大小、归一化等操作，以符合模型的输入要求。这些操作可以确保模型能够正确地处理图像数据。
模型推断：
- 经过预处理的视频帧被输入到 YOLOv5 模型中进行推断。模型通过单次前向传播，将图像作为输入，同时输出目标的位置信息（bounding box）和类别信息。
后处理：
- 在模型输出之后，通常会应用一些后处理步骤，以提高检测结果的准确性和稳定性。这些后处理步骤可能包括去除重叠框、应用置信度阈值、非极大值抑制（NMS）等。
可视化：
- 最后，检测到的目标信息会被用来在视频帧上绘制边界框和类别标签，以便用户可视化检测结果。

总体来说，YOLOv5 在处理视频帧时，采用了类似于图像处理的流程，即输入预处理、模型推断、后处理和结果可视化。通过这些步骤，YOLOv5 能够高效地从视频中检测出目标的位置和类别信息。

示例代码：

import torch
from torchvision import transforms
from yolov5.models.experimental import attempt_load
from yolov5.utils.datasets import letterbox
from yolov5.utils.general import non_max_suppression, scale_coords
from yolov5.utils.plots import plot_one_box
import cv2

# 加载模型
model = attempt_load("yolov5s.pt", map_location=torch.device('cuda')) # 可根据实际模型文件路径进行修改
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device).eval()

# 图像预处理
img_size = 640
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize(img_size),
])

# 读取视频
video_path = "input_video.mp4" # 输入视频路径
cap = cv2.VideoCapture(video_path)

while cap.isOpened():
ret, frame = cap.read()
if not ret:
break

# 图像预处理
img = transforms.ToPILImage()(frame)
img = letterbox(img, new_shape=img_size)[0]
img = transform(img).unsqueeze(0).to(device)

# 模型推断
with torch.no_grad():
pred = model(img)[0]
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)[0]

if pred is not None and len(pred):
# 后处理
pred[:, :4] = scale_coords(img.shape[2:], pred[:, :4], frame.shape).round()

# 可视化
for det in pred:
plot_one_box(det[:4].cpu().numpy(), frame, color=(0, 255, 0), label=f'{int(det[5])}')

# 显示结果
cv2.imshow('YOLOv5', frame)
if cv2.waitKey(1) == ord('q'):
break

cap.release()
cv2.destroyAllWindows()