Fix Dataset Loading and Streamline Code

- Add global dataset variable and load_merit_dataset() function
- Implement get_image_from_dataset() to ensure dataset is loaded before access
- Load dataset at the start of main block
- Update Gradio slider to use dataset length for maximum value
- Remove unused methods and imports
- Refactor process_image() to handle dataset image selection
- Adjust main block to initialize dataset before Gradio interface creation

This commit resolves the NameError related to undefined 'dataset' and
ensures proper dataset loading and access throughout the application.
It also removes unnecessary code, improving overall efficiency and readability.

app.py

@@ -1,76 +1,43 @@
-import io
-import requests
 import gradio as gr
 from huggingface_hub import list_models
-from datasets import load_dataset
 from typing import List
-from PIL import Image
 import torch
 from transformers import DonutProcessor, VisionEncoderDecoderModel
+from PIL import Image
 import json
 import re
 import logging
+from datasets import load_dataset
 
 # Logging configuration
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global variables for Donut model and processor
+# Global variables for Donut model, processor, and dataset
 donut_model = None
 donut_processor = None
+dataset = None
 
 
-def get_image_names(dataset):
-    return [str(i) for i in range(len(dataset))]
+def load_merit_dataset():
+    global dataset
+    if dataset is None:
+        dataset = load_dataset("de-Rodrigo/merit", name="en-digital-seq", split="train")
+    return dataset
 
 
 def get_image_from_dataset(index):
+    global dataset
+    if dataset is None:
+        dataset = load_merit_dataset()
     image_data = dataset[int(index)]["image"]
     return image_data
 
 
-def process_image(image=None, dataset_image_index=None):
-    if dataset_image_index:
-        image = get_image_from_dataset(dataset_image_index)
-
-    return image
-
-
-def create_interface(tag, image_indices):
-    """Create Gradio interface"""
-    iface = gr.Interface(
-        fn=process_image,
-        inputs=[
-            gr.Dropdown(choices=get_collection_models(tag), label="Select Model"),
-            gr.Image(type="pil", label="Upload Image"),
-            gr.Dropdown(
-                choices=image_indices, label="Select one from MERIT Dataset test-set"
-            ),
-        ],
-        outputs=gr.Image(label="Output Image"),
-        title="Saliency Visualization",
-        description="Upload your image or select one from the MERIT Dataset test-set.",
-    )
-    return iface
-
-
 def get_collection_models(tag: str) -> List[str]:
     """Get a list of models from a specific Hugging Face collection."""
     models = list_models(author="de-Rodrigo")
-
-    model_names = []
-    for model in models:
-        if tag in model.tags:
-            model_names.append(model.modelId)
-
-    return model_names
-
-
-def load_model(model_name: str):
-    """Load a model from Hugging Face Hub."""
-    model = AutoModel.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    return model, tokenizer
+    return [model.modelId for model in models if tag in model.tags]
 
 
 def get_donut():
@@ -145,6 +112,9 @@ def process_image(model_name, image=None, dataset_image_index=None):
 
 
 if __name__ == "__main__":
+    # Load the dataset
+    load_merit_dataset()
+
     models = get_collection_models("saliency")
     models.append("de-Rodrigo/donut-merit")
 
@@ -153,7 +123,9 @@ if __name__ == "__main__":
         inputs=[
             gr.Dropdown(choices=models, label="Select Model"),
             gr.Image(type="pil", label="Upload Image"),
-            gr.Slider(minimum=0, maximum=99, step=1, label="Dataset Image Index"),
+            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",
@@ -161,12 +133,3 @@ if __name__ == "__main__":
     )
 
     demo.launch()
-
-dataset_name = "de-Rodrigo/merit"
-dataset = load_dataset(dataset_name, name="en-digital-seq", split="train", num_proc=8)
-image_indices = get_image_names(dataset)
-
-models_tag = "saliency-merit"
-
-iface = create_interface(models_tag, image_indices)
-iface.launch()