import gradio as gr
from transformers import pipeline

translator = pipeline("translation_en_to_ru", model="Helsinki-NLP/opus-mt-en-ru")
image_classifier = pipeline("image-classification", model="microsoft/resnet-50")

# 3. Функция для перевода текста
def translate_to_russian(text):
    try:
        translation = translator(text)
        return translation[0]['translation_text']
    except Exception as e:
        print(f"Ошибка перевода: {e}")
        return text

def classify_image(image):
    results = image_classifier(image)
    output = {}
    for result in results:
        output[translate_to_russian(result['label'])] = result['score']
    return output
    
with gr.Blocks() as demo:
    gr.Markdown("Домашнее задание для курса ML2")
    with gr.Tab("Классификация изображений"):
        with gr.Row():
            img_input = gr.Image(type="pil", label="Загрузите изображение")
            label_output = gr.Label(label="Результаты классификации")
        btn_classify = gr.Button("Классифицировать")
        btn_classify.click(fn=classify_image, inputs=img_input, outputs=label_output)
    with gr.Tab("Перевод текста en-ru"):
        with gr.Row():
            text_input = gr.Textbox(label="EN text", placeholder="Напишите сюда английский текст")
            text_output = gr.Textbox(label="RU text", placeholder="Переведённый текст")
        btn_translate = gr.Button("Перевести")
        btn_translate.click(fn=translate_to_russian, inputs=text_input, outputs=text_output)

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