Improve GUI

app.py

@@ -23,7 +23,7 @@ dataset = None
 def load_merit_dataset():
     global dataset
     if dataset is None:
-        dataset = load_dataset("de-Rodrigo/merit", name="en-digital-seq", split="train")
+        dataset = load_dataset("de-Rodrigo/merit", name="en-digital-seq", split="test")
     return dataset
 
 
@@ -98,39 +98,66 @@ def process_image_donut(model, processor, image):
 
 
 @spaces.GPU
-def process_image(model_name, image=None, dataset_image_index=None):
-    if dataset_image_index is not None:
+def process_image(model_name, dataset_image_index, custom_image=None):
+    if custom_image is not None:
+        image = custom_image
+    else:
         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
 
 
 if __name__ == "__main__":
-    # Load the dataset
     load_merit_dataset()
-
     models = get_collection_models("saliency")
     models.append("de-Rodrigo/donut-merit")
 
-    demo = gr.Interface(
-        fn=process_image,
-        inputs=[
-            gr.Dropdown(choices=models, label="Select Model"),
-            gr.Image(type="pil", label="Upload Image"),
-            gr.Slider(
-                minimum=0, maximum=len(dataset) - 1, step=1, label="Dataset Image Index"
-            ),
-        ],
-        outputs=[gr.Image(label="Processed Image"), gr.Textbox(label="Result")],
-        title="Document Understanding with Donut",
-        description="Upload an image or select one from the dataset to process with the selected model.",
-    )
+    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."
+        )
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                model_dropdown = gr.Dropdown(choices=models, label="Select Model")
+                dataset_slider = gr.Slider(
+                    minimum=0,
+                    maximum=len(dataset) - 1,
+                    step=1,
+                    label="Dataset Image Index",
+                )
+
+            with gr.Column(scale=1):
+                custom_image = gr.Image(
+                    type="pil", label="Optional: Upload Custom Image"
+                )
+
+        process_button = gr.Button("Process Image")
+
+        with gr.Row():
+            output_image = gr.Image(label="Processed Image")
+            output_text = gr.Textbox(label="Result")
+
+        process_button.click(
+            fn=process_image,
+            inputs=[model_dropdown, dataset_slider, custom_image],
+            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()