60 lines
No EOL
2.3 KiB
Python
60 lines
No EOL
2.3 KiB
Python
from ultralytics import YOLO
|
|
import cv2
|
|
import os
|
|
|
|
# Load the model
|
|
# model = YOLO('../models/webcam-local-01/4/bangchak_poc/yolo11n.pt')
|
|
model = YOLO('yolo11m.pt')
|
|
|
|
def test_image(image_path):
|
|
"""Test a single image with YOLO model"""
|
|
if not os.path.exists(image_path):
|
|
print(f"Image not found: {image_path}")
|
|
return
|
|
|
|
# Run inference - filter for car class only (class 2 in COCO)
|
|
results = model(image_path, classes=[2, 5, 7]) # 2, 5, 7 = car, bus, truck in COCO dataset
|
|
|
|
# Display results
|
|
for r in results:
|
|
im_array = r.plot() # plot a BGR numpy array of predictions
|
|
|
|
# Resize image for display (max width/height 800px)
|
|
height, width = im_array.shape[:2]
|
|
max_dimension = 800
|
|
if width > max_dimension or height > max_dimension:
|
|
if width > height:
|
|
new_width = max_dimension
|
|
new_height = int(height * (max_dimension / width))
|
|
else:
|
|
new_height = max_dimension
|
|
new_width = int(width * (max_dimension / height))
|
|
im_array = cv2.resize(im_array, (new_width, new_height))
|
|
|
|
# Show image with predictions
|
|
cv2.imshow('YOLO Test - Car Detection Only', im_array)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
|
|
# Print detection info
|
|
print(f"\nDetections for {image_path}:")
|
|
if r.boxes is not None and len(r.boxes) > 0:
|
|
for i, box in enumerate(r.boxes):
|
|
cls = int(box.cls[0])
|
|
conf = float(box.conf[0])
|
|
original_class = model.names[cls] # Original class name (car/bus/truck)
|
|
# Get bounding box coordinates
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
|
# Rename all vehicle types to "car"
|
|
print(f"Detection {i+1}: car (was: {original_class}) - Confidence: {conf:.3f} - BBox: ({x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f})")
|
|
print(f"Total cars detected: {len(r.boxes)}")
|
|
else:
|
|
print("No cars detected in the image")
|
|
|
|
if __name__ == "__main__":
|
|
# Test with an image file
|
|
image_path = input("Enter image path (or press Enter for default test): ")
|
|
if not image_path:
|
|
image_path = "sample.png" # Default test image
|
|
|
|
test_image(image_path) |