import gradio as gr
from transformers import pipeline
from transformers import BlipProcessor, BlipForConditionalGeneration

translator_ru = pipeline("translation_en_to_ru", model="Helsinki-NLP/opus-mt-en-ru")
translator_en = pipeline("translation_ru_to_en", model="Helsinki-NLP/opus-mt-ru-en")
image_classifier = pipeline("image-classification", model="microsoft/resnet-50")
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
captioner = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")


def translate_to_russian(text):
    try:
        translation = translator_ru(text)
        return translation[0]['translation_text']
    except Exception as e:
        print(f"Ошибка перевода на ru: {e}")
        return text
    
def translate_to_english(text):
    try:
        translation = translator_en(text)
        return translation[0]['translation_text']
    except Exception as e:
        print(f"Ошибка перевода на en: {e}")
        return text

def classify_image(image):
    results = image_classifier(image)
    output = {}
    for result in results:
        output[translate_to_russian((result['label'].split(','))[0])] = result['score']
    return output

def capt_image(image, text):
    if len(text) != 0:
        inputs = processor(image, translate_to_english(text), return_tensors="pt")
    else:
        inputs = processor(image, return_tensors="pt")
    out = captioner.generate(**inputs)
    return translate_to_russian(processor.decode(out[0], skip_special_tokens=True))
    
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("Перевод текста"):
        with gr.Row("en-ru"):
            text_input_1 = gr.Textbox(label="EN text", placeholder="Напишите сюда английский текст")
            text_output_1 = gr.Textbox(label="RU text", placeholder="Переведённый текст")
        btn_translate_en_ru = gr.Button("Перевести")
        btn_translate_en_ru.click(fn=translate_to_russian, inputs=text_input_1, outputs=text_output_1)
        with gr.Row("ru-en"):
            text_input_2 = gr.Textbox(label="RU text", placeholder="Напишите сюда русский текст")
            text_output_2 = gr.Textbox(label="EN text", placeholder="Переведённый текст")
        btn_translate_ru_en = gr.Button("Перевести")
        btn_translate_ru_en.click(fn=translate_to_english, inputs=text_input_2, outputs=text_output_2)
    with gr.Tab("Написание заголовков"):
        with gr.Row():
            with gr.Column():
                img_capt = gr.Image(type="pil", label="Загрузите изображение")
                pre_capt = gr.Textbox(label="Начало заголовка")
            with gr.Column():
                caption = gr.Label(label="Заголовок для картинки")
        btn_captioning = gr.Button("Написать заголовок")
        btn_captioning.click(fn=capt_image, inputs=[img_capt, pre_capt], outputs=caption)
if __name__ == "__main__":
    demo.launch()