import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import os

# --- MODEL AYARLARI ---
MODEL_PATH_1 = "best.pt"            
MODEL_PATH_2 = "my_DeepYOLO11x.pt"  

models = {}

def load_yolo(path, name):
    if os.path.exists(path):
        try:
            models[name] = YOLO(path)
        except Exception as e:
            print(f"Hata: {e}")

load_yolo(MODEL_PATH_1, "ResYOLO11 (Hızlı)")
load_yolo(MODEL_PATH_2, "DeepYOLO11x (Uzman)")

MODEL_COLORS = {
    "ResYOLO11 (Hızlı)": (0, 255, 0),       
    "DeepYOLO11x (Uzman)": (0, 255, 255)             
}

# --- TAHMİN FONKSİYONU ---
def predict_fracture(input_image, conf_threshold, selected_models):
    empty_result = (None, "⚠️ Görüntü yok.")
    # Artık sadece 2 modelimiz var
    results = [empty_result] * 2

    if input_image is None or not selected_models:
        return results[0][0], results[0][1], results[1][0], results[1][1]

    if isinstance(input_image, dict):
        image_pil = input_image.get("composite", input_image.get("background"))
    else:
        image_pil = input_image

    if image_pil is None:
        return results[0][0], results[0][1], results[1][0], results[1][1]

    image_np = np.array(image_pil)
    
    if image_np.ndim == 3 and image_np.shape[-1] == 4:
        bg_mask = image_np[:, :, 3] == 0
        image_np[bg_mask] = [0, 0, 0, 255]
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
    elif image_np.ndim == 2:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)

    final_outputs = []
    model_names = ["ResYOLO11 (Hızlı)", "DeepYOLO11x (Uzman)"]

    for m_name in model_names:
        if m_name not in selected_models or m_name not in models or models[m_name] is None:
            # Siyah bir placeholder görüntü oluştur
            blank_image = np.zeros_like(image_np)
            cv2.putText(blank_image, "ATLANDI", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (100, 100, 100), 2)
            final_outputs.append((blank_image, f"⏳ {m_name} seçilmedi veya yüklenemedi."))
            continue

        model = models[m_name]
        
        # YOLO Tahmini
        res = model.predict(source=image_np, conf=conf_threshold, verbose=False, imgsz=1024)
        annotated_frame = cv2.cvtColor(res[0].plot(), cv2.COLOR_BGR2RGB)
        detections = len(res[0].boxes)

        info_text = f"✅ Tespit Edilen Bölge Sayısı: {detections}\n(Eşik: {conf_threshold:.2f})"
        final_outputs.append((annotated_frame, info_text))

    return final_outputs[0][0], final_outputs[0][1], final_outputs[1][0], final_outputs[1][1]

def create_demo():
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
            <h1 style="text-align: center;">🦴 Kırık Tespit ve Triage Analiz Paneli</h1>
            <p style="text-align: center; font-size: 16px;">Sadece yüksek performanslı YOLO modelleri devrede.</p>
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                paste_box = gr.Image(label="Panodan Yapıştır (Ctrl+V)", type="pil", interactive=True)
                image_editor = gr.ImageEditor(label="Röntgen Düzenleme Alanı", type="pil")
                paste_box.upload(fn=lambda x: x, inputs=paste_box, outputs=image_editor)

                model_selector = gr.CheckboxGroup(
                    label="Çalıştırılacak Modelleri Seçin",
                    choices=["ResYOLO11 (Hızlı)", "DeepYOLO11x (Uzman)"],
                    value=["ResYOLO11 (Hızlı)", "DeepYOLO11x (Uzman)"], 
                    interactive=True
                )

                conf_slider = gr.Slider(label="Güven Eşiği", minimum=0.0, maximum=1.0, value=0.15, step=0.05)
                analyze_btn = gr.Button("🔍 Seçili Modelleri Analiz Et", variant="primary", size="lg")

            with gr.Column(scale=2):
                with gr.Row():
                    with gr.Column():
                        output_image_1 = gr.Image(label="ResYOLO11", type="numpy")
                        info_box_1 = gr.Textbox(label="Sonuç", interactive=False)
                    with gr.Column():
                        output_image_2 = gr.Image(label="DeepYOLO11x", type="numpy")
                        info_box_2 = gr.Textbox(label="Sonuç", interactive=False)

        analyze_btn.click(
            fn=predict_fracture,
            inputs=[image_editor, conf_slider, model_selector],
            outputs=[output_image_1, info_box_1, output_image_2, info_box_2]
        )
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch()