File size: 3,161 Bytes
836049a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
#!/usr/bin/env python3
"""
Simple example script for using the Highway Vehicle Detection model.
This script demonstrates basic usage of the trained YOLOv8 model.
"""
from ultralytics import YOLO
import cv2
import os
def detect_vehicles(image_path, model_path="models/yolov8m_stage2_improved_best.pt"):
"""
Detect vehicles in an image using the trained model.
Args:
image_path (str): Path to the input image
model_path (str): Path to the trained model
Returns:
results: YOLO detection results
"""
# Load the trained model
model = YOLO(model_path)
# Run inference
results = model(image_path)
# Print detection results
for result in results:
boxes = result.boxes
if boxes is not None:
print(f"\nDetected {len(boxes)} vehicles:")
for i, box in enumerate(boxes):
x1, y1, x2, y2 = box.xyxy[0]
conf = box.conf[0]
cls = int(box.cls[0])
class_name = model.names[cls]
print(f" {i+1}. {class_name} (confidence: {conf:.2f})")
else:
print("No vehicles detected.")
return results
def process_video(video_path, output_path=None, model_path="models/yolov8m_stage2_improved_best.pt"):
"""
Process a video file and save results.
Args:
video_path (str): Path to the input video
output_path (str): Path to save the output video (optional)
model_path (str): Path to the trained model
"""
# Load the trained model
model = YOLO(model_path)
# Process video
if output_path:
results = model(video_path, save=True, project="output", name="detection")
print(f"Results saved to: output/detection/")
else:
results = model(video_path)
return results
def main():
"""Main function with example usage."""
print("Highway Vehicle Detection - Example Usage")
print("=" * 50)
# Check if model exists
model_path = "models/yolov8m_stage2_improved_best.pt"
if not os.path.exists(model_path):
print(f"Error: Model file not found at {model_path}")
print("Please make sure you've downloaded the model files.")
return
# Example 1: Process an image
print("\n1. Image Detection Example:")
print(" To detect vehicles in an image:")
print(" results = detect_vehicles('path/to/image.jpg')")
# Example 2: Process a video
print("\n2. Video Processing Example:")
print(" To process a video:")
print(" results = process_video('path/to/video.mp4', 'output.mp4')")
# Example 3: Direct model usage
print("\n3. Direct Model Usage:")
print(" from ultralytics import YOLO")
print(" model = YOLO('models/yolov8m_stage2_improved_best.pt')")
print(" results = model('path/to/image.jpg')")
print(" results[0].show()")
print("\n" + "=" * 50)
print("For more advanced features, use main.py")
if __name__ == "__main__":
main()
|