import gradio as gr
from generic import save_image_vis



def generate_viz(URL, question):
    print(f"Input Image URL: {URL}")
    print(f"Target Class: {question}")
    oring, img, html_page, answer = save_image_vis(URL, question)
    return answer, oring, img, html_page

title = "Explain LXMERT 🤯"


iface = gr.Interface(fn=generate_viz, inputs=[
        gr.Text(label="Input Image URL"),
        gr.Text(label="Input the Question"),
    ],
    outputs=[gr.Text(label="Answer"),
                              gr.Image(label="Oringinal Image", width=450, height=300),
                              gr.Image(label="EXplain Image", width=450, height=300),
                              gr.HTML(label="Visualize Sentence Importance"), ],
    title=title,
    allow_flagging="never",
    cache_examples=True,
    examples=[
        ["./lxmert/experiments/paper/COCO_val2014_000000185590/COCO_val2014_000000185590.jpg", "is the animal eating?"],
        ["./lxmert/experiments/paper/COCO_val2014_000000127510/COCO_val2014_000000127510.jpg", "did he catch the ball?"],
        ["./lxmert/experiments/paper/COCO_val2014_000000324266/COCO_val2014_000000324266.jpg", "is the tub white ?"],
        ["./lxmert/experiments/paper/COCO_val2014_000000200717/COCO_val2014_000000200717.jpg", "did the man just catch the frisbee?"],
    ],
)

iface.launch(debug=True)