Update Space Visualization

app.py

@@ -9,11 +9,17 @@ import json
 import re
 import logging
 from datasets import load_dataset
+import os
 
 # Logging configuration
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
+# Paths to the static image and GIF
+README_IMAGE_PATH = os.path.join("figs", "saliencies-merit-dataset.png")
+GIF_PATH = os.path.join("figs", "demo_samples.gif")
+
 # Global variables for Donut model, processor, and dataset
 donut_model = None
 donut_processor = None
@@ -126,45 +132,68 @@ if __name__ == "__main__":
     models.append("de-Rodrigo/donut-merit")
 
     with gr.Blocks() as demo:
-        gr.Markdown("# Document Understanding with Donut")
-        gr.Markdown(
-            "Select a model and an image from the dataset, or upload your own image."
-        )
+        gr.Markdown("# Saliency Maps with the MERIT Dataset 🎒📃🏆")
 
-        with gr.Row():
-            with gr.Column():
-                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",
-                )
-                upload_image = gr.Image(type="pil", label="Or Upload Your Own Image")
+        gr.Image(value=README_IMAGE_PATH, label="Example Document")
 
-            preview_image = gr.Image(label="Selected/Uploaded Image")
+        with gr.Tab("Introduction"):
+            gr.Markdown(
+                """
+            ## Welcome to Saliency Maps with the [MERIT Dataset](https://huggingface.co/datasets/de-Rodrigo/merit)
 
-        process_button = gr.Button("Process Image")
+            This space demonstrates the capabilities of different Vision Language models 
+            for document understanding tasks.
 
-        with gr.Row():
-            output_image = gr.Image(label="Processed Image")
-            output_text = gr.Textbox(label="Result")
+            ### Key Features:
+            - Process images from the [MERIT Dataset](https://huggingface.co/datasets/de-Rodrigo/merit) or upload your own image.
+            - Use a fine-tuned version of the models availabe to extract grades from documents.
+            - Visualize saliency maps to understand where the model is looking (WIP 🛠️).
+            """
+            )
 
-        # Update preview image when slider changes
-        dataset_slider.change(
-            fn=update_image, inputs=[dataset_slider], outputs=[preview_image]
-        )
+            gr.Image(value=GIF_PATH, label="Document Understanding Process")
 
-        # Update preview image when an image is uploaded
-        upload_image.change(
-            fn=lambda x: x, inputs=[upload_image], outputs=[preview_image]
-        )
+        with gr.Tab("Try It Yourself"):
+            gr.Markdown(
+                "Select a model and an image from the dataset, or upload your own image."
+            )
 
-        # Process image when button is clicked
-        process_button.click(
-            fn=process_image,
-            inputs=[model_dropdown, upload_image, dataset_slider],
-            outputs=[output_image, output_text],
-        )
+            with gr.Row():
+                with gr.Column():
+                    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",
+                    )
+                    upload_image = gr.Image(
+                        type="pil", label="Or Upload Your Own Image"
+                    )
+
+                preview_image = gr.Image(label="Selected/Uploaded Image")
+
+            process_button = gr.Button("Process Image")
+
+            with gr.Row():
+                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, upload_image, dataset_slider],
+                outputs=[output_image, output_text],
+            )
 
     demo.launch()