#!/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()