59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
|
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()
|