Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 from datasets import load_dataset
 from PIL import Image, ImageDraw
+import os
 
-dataset = load_dataset("agentsea/wave-ui-25k", split="train")
-
+hf_token = os.getenv("HF_TOKEN")
+dataset = load_dataset("agentsea/wave-ui-25k", split="train", token=hf_token)
 current_index = 0
 
 def draw_bounding_box(image, bbox):
@@ -44,6 +45,13 @@ def previous_image():
     current_index = (current_index - 1) % len(dataset)
     return show_image(current_index)
 
+def change_index(new_index):
+    new_index = int(new_index)
+    if 0 <= new_index < len(dataset):
+        return show_image(new_index)
+    else:
+        return show_image(current_index), f"Invalid index. Please enter a number between 0 and {len(dataset) - 1}."
+
 shortcut_js = """
 <script>
 function shortcuts(e) {
@@ -67,7 +75,7 @@ document.addEventListener('keydown', shortcuts, false);
 """
 
 with gr.Blocks(head=shortcut_js) as app:
-    gr.Markdown("# Wave UI 25k Dataset Explorer")
+    gr.Markdown("# WaveUI 25k Dataset Explorer")
     
     with gr.Row():
         image_output = gr.Image(type="pil")
@@ -77,9 +85,14 @@ with gr.Blocks(head=shortcut_js) as app:
         prev_button = gr.Button("Previous", elem_id="prev_btn")
         next_button = gr.Button("Next", elem_id="next_btn")
     
+    with gr.Row():
+        index_input = gr.Number(label="Go to index (0-based)", value=0)
+        go_button = gr.Button("Go", elem_id="go_btn")
+    
     prev_button.click(previous_image, outputs=[image_output, label_output])
     next_button.click(next_image, outputs=[image_output, label_output])
+    go_button.click(change_index, inputs=[index_input], outputs=[image_output, label_output])
     
     app.load(show_image, inputs=[gr.Number(value=0, visible=False)], outputs=[image_output, label_output])
 
-app.launch()
+app.launch()