import json
import gradio as gr

with open("img_id_to_info.json", "r") as f:
    img_id_to_info = json.load(f)
    
def output_id_info(img_id):
    if img_id is int:
        img_id = str(img_id)
    if img_id not in img_id_to_info.keys():
        print("wrong")
        return
    this_info = img_id_to_info[img_id]
    img_url = this_info['coco_url']
    is_train = this_info['is_train']
    caps = this_info['caps']
    return img_url, is_train, [[cap] for cap in caps]
    
demo = gr.Interface(
    title="COCO Information Retriever",
    description="Given COCO Image ID, this app will give you the corresponding image, captions and place in the dataset.",
    fn=output_id_info,
    inputs=gr.Text(label="img_id"),
    outputs=[
    gr.Image(label="Img"),
    gr.Checkbox(label="Is it in training set?"),
    gr.DataFrame(
        label="Caps",
        headers=["Captions"],
        datatype=['str'],
        row_count=5,
        col_count=1)
    ],
    examples=['391895',
    '522418',
    '184613',
    '318219',
    '554625']
)

demo.launch()