import gradio as gr
import cv2
from ultralytics import YOLO

# Загружаем модель YOLO
model_path = "sss.pt"  # Замените на путь к вашей модели
model = YOLO(model_path)

def process_frame(frame):
    # Выполняем обнаружение объектов
    results = model(frame, stream=True)  # Используем stream=True для повышения скорости
    
    # Рисуем боксы и метки на кадре
    for result in results:
        for box in result.boxes.data.tolist():
            x1, y1, x2, y2, confidence, class_id = box
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            label = model.names[int(class_id)]  # Убираем уверенность, оставляем только имя объекта

            # Рисуем прямоугольник вокруг объекта
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Добавляем текст
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return frame

def capture_video():
    # Открываем видеопоток с камеры
    cap = cv2.VideoCapture(0)
    
    if not cap.isOpened():
        print("Не удалось открыть камеру")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # Обрабатываем кадр
        processed_frame = process_frame(frame)

        # Gradio ожидает изображение в формате NumPy, поэтому передаем его так
        yield processed_frame

    cap.release()

# Создаем интерфейс Gradio
iface = gr.Interface(fn=capture_video, 
                     inputs=[], 
                     outputs=gr.Image(type="numpy", label="Object Detection"),  # Ожидаем NumPy массив
                     live=True)

# Запуск интерфейса
iface.launch()