import torch
import cv2
import RPi.GPIO as GPIO
import time

# ========== GPIO 설정 ==========
MOTOR_PIN = 18
GPIO.setmode(GPIO.BCM)
GPIO.setup(MOTOR_PIN, GPIO.OUT)

def stop_car():
    GPIO.output(MOTOR_PIN, GPIO.LOW)

def move_car():
    GPIO.output(MOTOR_PIN, GPIO.HIGH)

# ========== 모델 로드 (사전학습된 YOLOv5n) ==========
model = torch.hub.load('ultralytics/yolov5', 'yolov5n', pretrained=True)
model.conf = 0.5  # 신뢰도 임계값

# ========== 카메라 설정 ==========
cap = cv2.VideoCapture(0)

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

        results = model(frame)
        labels = results.pandas().xyxy[0]['name'].tolist()

        # 조건: 사람 or 차량이 있으면 정지
        if any(obj in labels for obj in ['person', 'car', 'truck', 'bus']):
            print(f"🛑 객체 감지됨: {labels} → 정지")
            stop_car()
        else:
            print("✅ 도로에 장애물 없음 → 직진")
            move_car()

        # 화면 출력
        annotated = results.render()[0]
        cv2.imshow("Object Detection", annotated)

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

finally:
    cap.release()
    cv2.destroyAllWindows()
    GPIO.cleanup()