Improve GUI

app.py

@@ -23,7 +23,9 @@ dataset = None
 def load_merit_dataset():
     global dataset
     if dataset is None:
-        dataset = load_dataset("de-Rodrigo/merit", name="en-digital-seq", split="test")
+        dataset = load_dataset(
+            "de-Rodrigo/merit", name="en-digital-seq", split="test", num_proc=8
+        )
     return dataset
 
 
@@ -98,34 +100,39 @@ def process_image_donut(model, processor, image):
 
 
 @spaces.GPU
-def process_image(model_name, dataset_image_index, custom_image=None):
-    if custom_image is not None:
-        image = custom_image
-    else:
+def process_image(model_name, image=None, dataset_image_index=None):
+    if dataset_image_index is not None:
         image = get_image_from_dataset(dataset_image_index)
 
     if model_name == "de-Rodrigo/donut-merit":
         model, processor = get_donut()
         result = process_image_donut(model, processor, image)
     else:
+        # Here you should implement processing for other models
         result = f"Processing for model {model_name} not implemented"
 
     return image, result
 
 
+def update_image(dataset_image_index):
+    return get_image_from_dataset(dataset_image_index)
+
+
 if __name__ == "__main__":
+    # Load the dataset
     load_merit_dataset()
+
     models = get_collection_models("saliency")
     models.append("de-Rodrigo/donut-merit")
 
     with gr.Blocks() as demo:
         gr.Markdown("# Document Understanding with Donut")
         gr.Markdown(
-            "Select an image from the dataset or optionally upload your own image."
+            "Select a model and an image from the dataset, or upload your own image."
         )
 
         with gr.Row():
-            with gr.Column(scale=2):
+            with gr.Column():
                 model_dropdown = gr.Dropdown(choices=models, label="Select Model")
                 dataset_slider = gr.Slider(
                     minimum=0,
@@ -133,11 +140,9 @@ if __name__ == "__main__":
                     step=1,
                     label="Dataset Image Index",
                 )
+                upload_image = gr.Image(type="pil", label="Or Upload Your Own Image")
 
-            with gr.Column(scale=1):
-                custom_image = gr.Image(
-                    type="pil", label="Optional: Upload Custom Image"
-                )
+            preview_image = gr.Image(label="Selected/Uploaded Image")
 
         process_button = gr.Button("Process Image")
 
@@ -145,19 +150,21 @@ if __name__ == "__main__":
             output_image = gr.Image(label="Processed Image")
             output_text = gr.Textbox(label="Result")
 
+        # Update preview image when slider changes
+        dataset_slider.change(
+            fn=update_image, inputs=[dataset_slider], outputs=[preview_image]
+        )
+
+        # Update preview image when an image is uploaded
+        upload_image.change(
+            fn=lambda x: x, inputs=[upload_image], outputs=[preview_image]
+        )
+
+        # Process image when button is clicked
         process_button.click(
             fn=process_image,
-            inputs=[model_dropdown, dataset_slider, custom_image],
+            inputs=[model_dropdown, upload_image, dataset_slider],
             outputs=[output_image, output_text],
         )
 
-        gr.Examples(
-            examples=[
-                ["de-Rodrigo/donut-merit", 0],
-                ["de-Rodrigo/donut-merit", 10],
-                ["de-Rodrigo/donut-merit", 20],
-            ],
-            inputs=[model_dropdown, dataset_slider],
-        )
-
     demo.launch()