Benutzer-Werkzeuge

Webseiten-Werkzeuge


yolov5

Inhaltsverzeichnis

YOLOv5 ist eine Implementierung des YOLO Paper (You only look once) zum Thema Object Detection von Ultralytics. Die aktuelle Version ist YOLOv8.

Image Detection

# pip install torch torchvision pandas ultralytics matplotlib
 
import torch
import matplotlib
import matplotlib.pyplot as plt
import sys
 
model = torch.hub.load('ultralytics/yolov5', 'yolov5x')
img = sys.argv[1]
results = model(img)
 
threshold = 0.6
 
filtered_results = []
for det in results.pred[0]:  
    if det[4] >= threshold:   
        filtered_results.append(det)
 
for det in filtered_results:
    x1, y1, x2, y2, prob, cls = det[:6]
    class_names = model.model.names
    class_name = class_names[int(cls)]
    print(f"{class_name} with {prob:.4f} probability at ({x1:.2f}, {y1:.2f}) - ({x2:.2f}, {y2:.2f})")
 
 
if filtered_results:
    results.pred[0] = torch.stack(filtered_results)
    results.show()  # results.save(save_dir='results') 

Video Detection

# pip install torch torchvision pandas ultralytics matplotlib
 
import torch
import matplotlib.pyplot as plt
import cv2
import sys
 
# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5x')
 
# Open video capture
video_path = sys.argv[1]
cap = cv2.VideoCapture(video_path)
 
target_frame = 4000  # Replace with the desired frame index to forward to
cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
 
threshold = 0.001
 
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
 
    # Perform object detection on the current frame
    results = model(frame)
    filtered_results = []
 
    for det in results.pred[0]:
        if det[4] >= threshold:
            filtered_results.append(det)
 
    if filtered_results:
        results.pred[0] = torch.stack(filtered_results)
        frame_with_results = results.render()[0]
 
        # Display the frame with detected objects
        cv2.imshow('Object Detection', frame_with_results)
 
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
 
cap.release()
cv2.destroyAllWindows()
yolov5.txt · Zuletzt geändert: 2023/08/13 16:12 (Externe Bearbeitung)