forked from scott/YOLOv8_test_example
37 lines
953 B
Python
Executable File
37 lines
953 B
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import sys
|
|
|
|
from ultralytics import YOLO
|
|
|
|
image_path = sys.argv[1]
|
|
|
|
# Load the small pretrained YOLOv8 object-detection model.
|
|
model = YOLO("yolov8n.pt")
|
|
|
|
# Run object detection on a test image.
|
|
results = model(image_path)
|
|
|
|
has_bus = False
|
|
max_confidence = 0.0
|
|
|
|
for result in results:
|
|
class_names = result.names
|
|
detected_classes = result.boxes.cls.tolist()
|
|
confidences = result.boxes.conf.tolist()
|
|
|
|
# Check each detected object and its matching confidence score
|
|
for cls_id, conf in zip(detected_classes, confidences):
|
|
if class_names[int(cls_id)] == "bus":
|
|
has_bus = True
|
|
# Keep track of the highest confidence if there are multiple buses
|
|
if conf > max_confidence:
|
|
max_confidence = conf
|
|
|
|
# Output the Yes/No answer and the confidence percentage
|
|
if has_bus:
|
|
print("Yes")
|
|
print(f"Confidence: {max_confidence * 100:.1f}%")
|
|
else:
|
|
print("No")
|