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)