import cv2
import os
from ultralytics import YOLO
from datetime import datetime

# Load a YOLOv8 pretrained model (use 'yolov8n.pt' or 'yolov8s.pt' for speed)
model = YOLO('yolov8n.pt')

# Create directory to save car images
save_dir = 'captured_cars'
os.makedirs(save_dir, exist_ok=True)

# Open webcam (0 = default camera)
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("Failed to open webcam.")
    exit()

print("Press Q to quit.")

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # Run YOLO inference
    results = model(frame)

    for result in results:
        for box in result.boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])

            # Class 2 = car in COCO dataset
            if cls_id == 2 and conf > 0.5:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                car_crop = frame[y1:y2, x1:x2]

                # Save image
                timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
                filename = os.path.join(save_dir, f'car_{timestamp}.jpg')
                cv2.imwrite(filename, car_crop)
                print(f"Saved: {filename}")

                # Draw bounding box
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(frame, f"Car {conf:.2f}", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

    # Show the frame
    cv2.imshow("Car Detection", frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()