Fixed the image selection logic.

app.py

@@ -130,13 +130,6 @@ def save_image(
     print("Image saved successfully!")
 
 
-# Function to handle image selection and enable the save button
-def select_image(selected_image_index, images):
-    # Get the selected image from its index
-    selected_image = images[selected_image_index]
-    return gr.update(interactive=True), selected_image
-
-
 with gr.Blocks() as demo:
     selected_image = gr.State(None)
     current_episode = gr.State(0)
@@ -189,11 +182,10 @@ with gr.Blocks() as demo:
     )
 
     # When an image is clicked
-    images.select(
-        select_image,
-        inputs=[gr.Number(), images],
-        outputs=[save_button, selected_image],
-    )
+    def get_select_index(evt: gr.SelectData, images):
+        return images[evt.index]
+
+    images.select(get_select_index, images, selected_image)
 
     save_button.click(
         save_image,