import gradio as gr

# processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
# model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

from transformers import pipeline

import base64
import os
with open("Iso_Logotipo_Ceibal.png", "rb") as image_file:
    encoded_image = base64.b64encode(image_file.read()).decode()


classifier = pipeline(model="google/vit-base-patch16-224")
# classifier("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png")
def clasificador(image):
    results = classifier(image)
    result = {}
    for item in results:
        result[translate_text(item['label'])] = item['score']
    return result


es_en_translator = pipeline("translation",model = "Helsinki-NLP/opus-mt-es-en")
def translate_text(text):
    print(text)
    text = es_en_translator(text)[0].get("translation_text")
    print(text)
    return text

with gr.Blocks(title = "Uso de AI para la clasificación de imágenes.") as demo:
    gr.Markdown("""
    <center>
    <h1>
    Uso de AI para la clasificación de imágenes.
    </h1>
    <img src='data:image/jpg;base64,{}' width=200px>
    <h3>
    Con este espacio podrás clasificar imágenes y objetos a partir de una imagen.
    </h3>
    
    </center>
    """.format(encoded_image))
    with gr.Row():
        with gr.Column():
            inputt = gr.Image(type="pil", label="Ingresá la imagen a clasificar.")
            button = gr.Button(value="Clasificar")
            examples = gr.Examples(examples=[os.path.join(os.path.dirname(__file__), "palacio.jpeg")],inputs=[inputt])
        with gr.Column():
            output = gr.Label()
    button.click(clasificador,inputt,output)
    
demo.launch()