# !pip install ultralytics
import gradio as gr
import cv2
from ultralytics import YOLO
import torch
import os
import numpy as np 


loaded_model = YOLO('best.pt')

# Проверка устройства
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Используется устройство: {device}')
loaded_model.to(device)


# # --- Глобальные переменные для оптимизации ---
# # Пропуск кадров для увеличения производительности
# FRAME_SKIP = 2 # Обрабатывать каждый 2-й кадр
# frame_counter = 0
# last_annotated_frame = None # Храним последний обработанный кадр


# --- Функция для обработки веб-камеры ---
def predict_webcam(img, conf_threshold, iou_threshold):
    # global frame_counter, last_annotated_frame

    if img is None:
        return last_annotated_frame # Возвращаем последний кадр, если новый не пришел

    # Исправляем зеркальное отражение
    img = cv2.flip(img, 1)

    # frame_counter += 1
    # if frame_counter % FRAME_SKIP != 0 and last_annotated_frame is not None:
    #      # Если пропускаем кадр, возвращаем предыдущий результат, чтобы не было "замирания"
    #     return last_annotated_frame

    # Пропускаем кадр через модель
    results = loaded_model.track( # ИЗМЕНЕНО: .predict() на .track() для стабильности
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=640,
        verbose=False,
        persist=True, # ВАЖНО для трекера: сохранять треки между кадрами
        tracker='botsort.yaml' # Включаем трекер
    )

    # results - это список, содержащий один объект Results
    # plot() возвращает аннотированное изображение в формате numpy array (BGR)
    annotated_frame = results[0].plot()

    # Преобразование цвета из BGR (OpenCV) в RGB (Gradio)
    # annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)

    # Сохраняем последний успешно обработанный кадр
    # last_annotated_frame = annotated_frame_rgb
    last_annotated_frame = annotated_frame
    
    # return annotated_frame_rgb
    return annotated_frame

# --- Определение интерфейса Gradio ---
iface = gr.Interface(
    fn=predict_webcam,
    inputs=[
        gr.Image(type="numpy", label="Webcam Feed", streaming=True),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
    ],
    outputs=gr.Image(type="numpy", label="Result"),
    live=True,
    title="Impurity Detection (Webcam)",
    description="Real-time impurity detection using your webcam. Adjust confidence and IoU thresholds below.",
)

# # --- Запуск интерфейса Gradio ---
# if __name__ == "__main__":
#     # РЕШЕНИЕ №4 (дополнительно): .queue() для более отзывчивого интерфейса
#     iface.queue().launch(share=True, debug=True)


if __name__ == "__main__":
    print("\nЗапуск Gradio интерфейса. Ожидайте появления публичной ссылки...")
    iface.launch(share=True) # Устанавливаем share=True для запуска в Colab