如何调用yolov5训练好的结果进行摄像头目标检测
时间: 2023-08-11 22:28:34 浏览: 85
首先,你需要安装yolov5和OpenCV的Python包。然后,你需要调用训练好的模型权重文件和模型配置文件。你可以使用以下代码来进行摄像头目标检测:
```
import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device
# Load YOLOv5s model
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))
# Set device to CPU
device = select_device('cpu')
# Initialize VideoCapture object
cap = cv2.VideoCapture(0)
# Set video resolution
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
# Define classes
classes = ['person', 'car', 'truck', 'bus']
while True:
# Read frame from video
ret, frame = cap.read()
if not ret:
break
# Convert frame from BGR to RGB
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Resize frame to model input size
img = cv2.resize(frame, (640, 640))
# Normalize pixel values
img = img / 255.0
# Convert image to PyTorch tensor
img = torch.from_numpy(img).permute(2, 0, 1).float().to(device)
# Run inference on model
pred = model(img, augment=False)[0]
# Apply non-maximum suppression
pred = non_max_suppression(pred, conf_thres=0.4, iou_thres=0.5)
# Process detections
for i, det in enumerate(pred):
if det is not None and len(det):
# Rescale box coordinates to original image size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
# Loop over detections
for *xyxy, conf, cls in reversed(det):
# Draw bounding box and label
label = f'{classes[int(cls)]} {conf:.2f}'
cv2.rectangle(frame, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (255, 0, 0), 2)
cv2.putText(frame, label, (int(xyxy[0]), int(xyxy[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
# Convert frame back to BGR
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
# Display frame
cv2.imshow('frame', frame)
# Press 'q' to quit
if cv2.waitKey(1) == ord('q'):
break
# Release video capture and destroy all windows
cap.release()
cv2.destroyAllWindows()
```
在这个代码中,我们首先加载了YOLOv5s模型,然后初始化了一个VideoCapture对象来捕捉摄像头的视频流。我们将视频帧转换为模型输入大小的图像,并将其归一化。然后,我们对模型进行推理,并应用非最大抑制来剔除低置信度的检测结果。最后,我们将边界框和标签绘制在原始帧上,并将其显示在屏幕上。