import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import torch

# Моделді жүктеу
model = None

def load_model(model_path="best.pt"):
    """Моделді жүктеу функциясы"""
    global model
    try:
        model = YOLO(model_path)
        return "Модель сәтті жүктелді!"
    except Exception as e:
        return f"Қате: {str(e)}"

def predict_image(image, conf_threshold=0.25, iou_threshold=0.45):
    """Суретті талдау және сегментация жасау"""
    if model is None:
        return None, "Модель жүктелмеген!"
    
    try:
        # Болжам жасау
        results = model.predict(
            source=image,
            conf=conf_threshold,
            iou=iou_threshold,
            save=False
        )
        
        # Нәтижені визуализациялау
        annotated_img = results[0].plot()
        annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
        
        # Анықталған объектілер туралы ақпарат
        detections = []
        if results[0].masks is not None:
            for i, (box, mask) in enumerate(zip(results[0].boxes, results[0].masks)):
                class_id = int(box.cls[0])
                confidence = float(box.conf[0])
                class_name = model.names[class_id]
                detections.append(f"{i+1}. {class_name}: {confidence:.2%}")
        
        info_text = "\n".join(detections) if detections else "Объектілер табылмады"
        
        return annotated_img, info_text
    
    except Exception as e:
        return None, f"Қате: {str(e)}"

# Gradio интерфейсі
def create_interface():
    with gr.Blocks(title="YOLOv8 Қазақша Сегментация", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎯 YOLOv8 Сегментация - Қазақша
        
        Бұл модель объектілерді анықтап, сегментациялайды. 
        Суретті жүктеп, параметрлерді орнатыңыз.
        """)
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="Суретті жүктеңіз",
                    type="numpy"
                )
                
                conf_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.25,
                    step=0.05,
                    label="Сенімділік табалдырығы (Confidence)"
                )
                
                iou_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.45,
                    step=0.05,
                    label="IoU табалдырығы"
                )
                
                predict_btn = gr.Button("🔍 Талдау", variant="primary")
            
            with gr.Column():
                output_image = gr.Image(
                    label="Нәтиже"
                )
                output_text = gr.Textbox(
                    label="Анықталған объектілер",
                    lines=10
                )
        
        gr.Markdown("""
        ### 📝 Нұсқаулық:
        1. Суретті жүктеңіз
        2. Қажет болса параметрлерді өзгертіңіз
        3. "Талдау" батырмасын басыңыз
        
        ### ⚙️ Параметрлер:
        - **Сенімділік табалдырығы**: Объектіді қабылдау үшін минималды сенімділік
        - **IoU табалдырығы**: Non-maximum suppression үшін қолданылады
        """)
        
        predict_btn.click(
            fn=predict_image,
            inputs=[input_image, conf_slider, iou_slider],
            outputs=[output_image, output_text]
        )
        
        # Мысал суреттер
        gr.Markdown("### 📸 Мысалдар:")
        gr.Examples(
            examples=[
                ["example1.jpg", 0.25, 0.45],
                ["example2.jpg", 0.25, 0.45],
            ],
            inputs=[input_image, conf_slider, iou_slider],
            outputs=[output_image, output_text],
            fn=predict_image,
            cache_examples=False,
        )
    
    return demo

# Моделді жүктеу
load_model()

# Интерфейсті іске қосу
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()