# ==============================================================================
# 🚨 МӘЖБҮРЛЕП ОРНАТУ (Қажет болса): 'transformers' табылмаса, оны орнатады.
# ==============================================================================
import os
import subprocess

try:
    from transformers import pipeline
except ImportError:
    print("WARNING: transformers табылмады. Қазір орнатылуда...")
    subprocess.check_call([os.sys.executable, "-m", "pip", "install", "transformers", "accelerate", "torch", "gradio", "Pillow"])
    from transformers import pipeline

# ----------------------------------------------------------------------
# 1. КОНФИГУРАЦИЯЛАР (Қолданушы Атын Қою)
# ----------------------------------------------------------------------

import gradio as gr 

# ⭐ ҚОЛДАНУШЫ АТЫ ⭐
HF_USERNAME = "Nonabzbssbbsbs"

TRASH_MODEL_PATH = f"{HF_USERNAME}/Trash-Net-Fine-Tuned-Demo"
CIFAR_MODEL_PATH = f"{HF_USERNAME}/CIFAR10-Kazakh-Fast-Demo"

# 2. Қазақша Белгілер
CIFAR_KAZAKH_LABELS = {
    0: "Ұшақ", 1: "Автомобиль", 2: "Құс", 3: "Мысық", 4: "Бұғы", 
    5: "Ит", 6: "Бақа", 7: "Жылқы", 8: "Кеме", 9: "Жүк көлігі"
}

# ----------------------------------------------------------------------
# 2. PIPELINES ЖҮКТЕУ (Екі Модельді Де Жүктеу)
# ----------------------------------------------------------------------

print("Модельдер жүктелуде...")

# 1-ші Модель (Trash/Food)
try:
    trash_classifier = pipeline(
        "image-classification", 
        model=TRASH_MODEL_PATH
    )
except Exception as e:
    print(f"Trash-Net моделін жүктеу қатесі: {e}")
    trash_classifier = None 

# 2-ші Модель (CIFAR-10)
try:
    cifar_classifier = pipeline(
        "image-classification", 
        model=CIFAR_MODEL_PATH
    )
except Exception as e:
    print(f"CIFAR-10 моделін жүктеу қатесі: {e}")
    cifar_classifier = None 

print("Модельдерді жүктеу әрекеті аяқталды.")

# ----------------------------------------------------------------------
# 3. ЖІКТЕУ ФУНКЦИЯЛАРЫ
# ----------------------------------------------------------------------

# CIFAR-10 жіктеу
def classify_cifar(image):
    if image is None:
        return "Сурет жүктеңіз."
    if cifar_classifier is None:
        return "CIFAR-10 моделін жүктеу мүмкін болмады."
        
    results = cifar_classifier(image)
    output = {}
    
    for item in results:
        try:
            label_id = int(item['label'])
            kazakh_label = CIFAR_KAZAKH_LABELS.get(label_id, item['label']) 
            output[kazakh_label] = item['score']
        except ValueError:
            output[item['label']] = item['score']

    return output

# Trash/Food жіктеу
def classify_trash(image):
    if image is None:
        return "Сурет жүктеңіз."
    if trash_classifier is None:
        return "Trash/Food моделін жүктеу мүмкін болмады."
        
    results = trash_classifier(image)
    output = {item['label']: item['score'] for item in results}
    return output

# ----------------------------------------------------------------------
# 4. GRADIO ИНТЕРФЕЙСІ
# ----------------------------------------------------------------------

cifar_iface = gr.Interface(
    fn=classify_cifar,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Қазақша CIFAR-10 Жіктеуі",
    description="CIFAR-10 моделі (10 класс). Нәтижелер қазақшаға аударылған.",
    allow_flagging="never"
)

trash_iface = gr.Interface(
    fn=classify_trash,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Trash/Food Жіктеуі",
    description="Trash-Net негізіндегі модель Food101-де оқытылған (101 класс).",
    allow_flagging="never"
)

demo = gr.TabbedInterface(
    [cifar_iface, trash_iface], 
    ["1. CIFAR10 (Қазақша)", "2. Trash/Food (101 класс)"],
    title="Екі Компьютерлік Көру Моделінің Демонстрациясы"
)

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