import os
import gradio as gr
from gradio_client import Client

MY_HF_TOKEN_KEY = os.environ["MY_HF_TOKEN_KEY"]

client = Client("Ghana-NLP/Khaya-image-caption-backend-infer-api",hf_token=MY_HF_TOKEN_KEY,upload_files=True)

def generate(filepath,language):
    output = client.predict(filepath,language,api_name="/predict")
    return output

with gr.Blocks() as demo:
    title = gr.Markdown(
    """
    # African Language Image Captioning
    Based on the <a href="https://translation.ghananlp.org/">Khaya AI Translation API</a>, <a href="https://lesan.ai/">Lesan AI</a>, Google Translate API and the BLIP model. Lesan is used for <b>Amharic</b> and <b>Tigrinya</b>, <a href="https://translation.ghananlp.org/">Khaya AI</a> is used for <b>Twi</b>, <b>Dagbani</b>, <b>Ewe</b>, <b>Ga</b>, <b>Gurene</b>, <b>Fante</b>, <b>Kikuyu</b>, <b>Kimeru</b>, <b>Luo</b>, <b>Yoruba</b> and Google is used for <b>Hausa</b>, <b>Swahili</b> and <b>Shona</b>.
    """)
    with gr.Row():
        with gr.Column(scale=1):
            language_selector = gr.Dropdown(["Twi","Amharic","Dagbani","Ewe","Ga","Gurene","Fante","Hausa", "Kikuyu", "Kimeru", "Luo","Shona","Swahili","Tigrinya","Yoruba"],value="Twi",label="Choose Language! (default is Twi)", info="Language:")
            inputs = [gr.Image(type="filepath"),language_selector]
            examples = [["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png"],
                ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/rotated_image.jpeg"]]
        with gr.Column(scale=1):
            outputs = [gr.Textbox(label="English Caption"), gr.Textbox(label="African Language Caption"),
                       gr.Image(label="OUT", type="filepath")]
    btn = gr.Button("Generate African Language Caption")
    btn.click(generate, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()