new tabs

app.py

@@ -1,25 +1,44 @@
 import gradio as gr
 from transformers import pipeline
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
-#sentencepiece no installed??
-
-translator = pipeline("translation_en_to_ru", model="Helsinki-NLP/opus-mt-en-ru")
+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(text)
+        translation = translator_ru(text)
         return translation[0]['translation_text']
     except Exception as e:
-        print(f"Ошибка перевода: {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'])] = result['score']
+        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")
@@ -29,12 +48,25 @@ with gr.Blocks() as demo:
             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.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():
-            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)
-
+            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()