import gradio as gr
from datasets import load_dataset

dataset = load_dataset("sayakpaul/nyu_depth_v2", split="train[:2%]", trust_remote_code=True)

def preview(index):
    sample = dataset[int(index)]
    return sample["image"], sample["depth"], sample["label"]

gr.Interface(
    fn=preview,
    inputs=gr.Slider(0, len(dataset)-1, step=1, label="Sample Index"),
    outputs=[
        gr.Image(type="pil", label="RGB Image"),
        gr.Image(type="numpy", label="Depth Map"),
        gr.Image(type="numpy", label="Segmentation Label")
    ],
    title="NYU Depth V2 Viewer"
).launch()