import gradio as gr
from ultralytics import YOLO
from PIL import Image

# Pretrained YOLOv8n моделін жүктеу
model = YOLO('yolov8n.pt')

# Қоқыс жәшігін тексеру және дабыл шығару функциясы
def check_trash(img):
    results = model(img)
    detected_bins = []
    detected_trash = []

    for box in results[0].boxes:
        label = results[0].names[int(box.cls.item())]
        conf = box.conf.item()
        if conf < 0.5:
            continue
        if 'bin' in label or 'container' in label:  # Жәшік
            detected_bins.append(label)
        else:  # Қоқыс
            detected_trash.append(label)

    output = []
    # Толы / бос жәшік
    if detected_bins:
        if detected_trash:
            output.append("Жәшік: Толы ✅")
        else:
            output.append("Жәшік: Бос ⚪")
    else:
        output.append("Жәшік табылмады ❌")

    # Сыртқа шыққан қоқыс
    if detected_trash:
        output.append(f"Сыртқа шыққан қоқыс табылды 🚨: {', '.join(detected_trash)}")

    return "\n".join(output)

# Gradio интерфейсі
app = gr.Interface(
    fn=check_trash,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(),
    title="SmartCity Waste Detection",
    description="Қоқыс жәшігін, сыртқа төгілген қоқысты және толық/бос жәшікті анықтайды"
)

app.launch()