from datasets import load_dataset
import gradio as gr
import os
import random

wmtis = load_dataset("nlphuji/wmtis-identify")['test']
print(f"Loaded WMTIS identify, first example:")
print(wmtis[0])
dataset_size = len(wmtis) - 1

NORMAL_IMAGE = 'normal_image'
STRANGE_IMAGE = 'strange_image'
def func(index):
    example = wmtis[index]
    return example['normal_image'], example['normal_hash'], example['strange_image'], example['strange_hash']

demo = gr.Blocks()

with demo:
    gr.Markdown("# Slide to iterate WMTIS: Normal vs. Strange Images")

    with gr.Column():
        slider = gr.Slider(minimum=0, maximum=dataset_size)
        with gr.Row():
            index = random.choice(range(0, dataset_size))
            with gr.Column():
                i1 = gr.Image(value=wmtis[index]["normal_image"], label='Normal Image')
                t1 = gr.Textbox(value=wmtis[index]["normal_hash"], label='Image ID')
            with gr.Column():
                i2 = gr.Image(value=wmtis[index]["strange_image"], label='Strange Image')
                t2 = gr.Textbox(value=wmtis[index]["strange_hash"], label='Image ID')

    slider.change(func, inputs=[slider], outputs=[i1, t1, i2, t2])

demo.launch()