import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from PIL import Image
import torch

# ----------------------------
# Модельдер: ашық қолжетімді
# ----------------------------

# 1️⃣ Жеміс/өсімдік классификациясы
FRUIT_MODEL = "nateraw/vit-base-fruit-classification"
fruit_extractor = AutoFeatureExtractor.from_pretrained(FRUIT_MODEL)
fruit_model = AutoModelForImageClassification.from_pretrained(FRUIT_MODEL)
fruit_classifier = pipeline("image-classification", model=fruit_model, feature_extractor=fruit_extractor)

# 2️⃣ Қоқыс классификациясы
TRASH_MODEL = "keremberke/vit-base-patch16-224-trash"
trash_extractor = AutoFeatureExtractor.from_pretrained(TRASH_MODEL)
trash_model = AutoModelForImageClassification.from_pretrained(TRASH_MODEL)
trash_classifier = pipeline("image-classification", model=trash_model, feature_extractor=trash_extractor)

# 3️⃣ Аударма (2 модель)
# Қазақша - ағылшынша
TRANSLATE_MODEL_KZ_EN = "Helsinki-NLP/opus-mt-kaz-en"
tokenizer_kz_en = AutoTokenizer.from_pretrained(TRANSLATE_MODEL_KZ_EN)
model_kz_en = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATE_MODEL_KZ_EN)
translator_kz_en = pipeline("translation", model=model_kz_en, tokenizer=tokenizer_kz_en)

# Қазақша - Монголша
TRANSLATE_MODEL_KZ_MN = "Helsinki-NLP/opus-mt-kaz-mon"
tokenizer_kz_mn = AutoTokenizer.from_pretrained(TRANSLATE_MODEL_KZ_MN)
model_kz_mn = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATE_MODEL_KZ_MN)
translator_kz_mn = pipeline("translation", model=model_kz_mn, tokenizer=tokenizer_kz_mn)

# ----------------------------
# Функциялар
# ----------------------------

def classify_fruit(img):
    results = fruit_classifier(img)
    return {r['label']: float(r['score']) for r in results}

def classify_trash(img):
    results = trash_classifier(img)
    return {r['label']: float(r['score']) for r in results}

def translate_text(text):
    en = translator_kz_en(text)[0]['translation_text']
    mn = translator_kz_mn(text)[0]['translation_text']
    return en, mn

# ----------------------------
# Gradio интерфейс
# ----------------------------

with gr.Blocks() as demo:
    gr.Markdown("## 🌱 EcoVision: Жеміс/өсімдік, қоқыс классификациясы + аударма")

    with gr.Tab("Суретке түсіру"):
        with gr.Row():
            fruit_input = gr.Image(label="Жеміс/өсімдік суреті", type="pil")
            fruit_output = gr.Label(label="Жеміс/өсімдік нәтижесі")
        fruit_btn = gr.Button("Тексеру")
        fruit_btn.click(classify_fruit, inputs=fruit_input, outputs=fruit_output)

        with gr.Row():
            trash_input = gr.Image(label="Қоқыс суреті", type="pil")
            trash_output = gr.Label(label="Қоқыс нәтижесі")
        trash_btn = gr.Button("Тексеру")
        trash_btn.click(classify_trash, inputs=trash_input, outputs=trash_output)

    with gr.Tab("Аударма"):
        text_input = gr.Textbox(label="Мәтін (қазақша)")
        en_output = gr.Textbox(label="Ағылшынша")
        mn_output = gr.Textbox(label="Монголша")
        trans_btn = gr.Button("Аудару")
        trans_btn.click(translate_text, inputs=text_input, outputs=[en_output, mn_output])

demo.launch()