yolo5图片、视频推理demo
![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/f84b757854684ff78e4529d370727da2.png)
图片
import torch
model = torch.hub.load('./yolo5', 'custom', path='yolov5s.pt', source='local')
img = '1.jpg'
results = model(img)
detections = results.xyxy[0].cpu().numpy()
for detection in detections:
x1, y1, x2, y2, confidence, cls = detection
print(f"Class: {model.names[int(cls)]}, Confidence: {confidence:.2f}, Box: [{x1}, {y1}, {x2}, {y2}]")
results.show()
视频
import cv2
import torch
model = torch.hub.load('../yolo5', 'custom', path='yolov5s.pt', source='local')
video_path = '1.mp4'
cap = cv2.VideoCapture(video_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
output_path = 'output_video.mp4'
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = model(img_rgb)
detections = results.xyxy[0].cpu().numpy()
for detection in detections:
x1, y1, x2, y2, confidence, cls = detection
label = f"{model.names[int(cls)]} {confidence:.2f}"
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
print(f"Class: {model.names[int(cls)]}, Confidence: {confidence:.2f}, Box: [{x1}, {y1}, {x2}, {y2}]")
cv2.imshow('YOLOv5 Detection', frame)
out.write(frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
out.release()
cv2.destroyAllWindows()