import gradio as gr
from datasets import load_dataset
import numpy as np

# Load your dataset only once
ds = load_dataset("Devenarya/Microsoft100", split="train")

# Generator as before
def get_image_and_mask(index):
    # This assumes dataset alternates: image, mask, image, mask, ...
    images = [item['image'] for item in ds]
    img = images[2 * index].convert('RGB')
    mask = images[2 * index + 1].convert('L')
    return np.array(img), np.array(mask)

def demo_fn(index):
    # Output two images: the photo and its segmentation
    return get_image_and_mask(index)

demo = gr.Interface(
    fn=demo_fn,
    inputs=gr.Slider(0, (len(ds)//2)-1, step=1, label="Image Index"),
    outputs=[gr.Image(label="Original Image"), gr.Image(label="Segmentation Mask")],
    title="Microsoft100 Image & Segmentation Explorer"
)

if __name__ == "__main__":
    demo.launch()