update

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import os
 import threading
 import random
-from datasets import load_dataset, Dataset, Features, Value, concatenate_datasets
+from datasets import load_dataset, Dataset, DatasetDict, Features, Value, concatenate_datasets
 from huggingface_hub import login
 
 # Authenticate with Hugging Face
@@ -15,33 +15,46 @@ dataset_name = "GeorgeIbrahim/EGYCOCO"  # Replace with your dataset name
 
 # Load or create the dataset
 try:
-    dataset = load_dataset(dataset_name, split="train")
+    dataset = load_dataset(dataset_name)
     print("Loaded existing dataset:", dataset)
 except Exception as e:
-    # Create an empty dataset if it doesn't exist
+    # Create empty train and val datasets if they don't exist
     features = Features({
         'image_id': Value(dtype='string'),
         'caption': Value(dtype='string'),
     })
-    dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
-    dataset.push_to_hub(dataset_name)  # Push the empty dataset to Hugging Face
+    train_dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
+    val_dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
+    dataset = DatasetDict({"train": train_dataset, "validation": val_dataset})
+    dataset.push_to_hub(dataset_name)
 
 image_folder = "images"
 image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
 lock = threading.Lock()
 
+
+# Helper function to determine dataset split based on image filename
+def determine_split(image_id):
+    if "train" in image_id:
+        return "train"
+    elif "val" in image_id:
+        return "validation"
+    else:
+        return None
+
+
 # Function to get a random image that hasn’t been annotated or skipped
 def get_next_image(session_data):
     with lock:
-        annotated_images = set(dataset["image_id"])  # Set of annotated images
+        annotated_images = set(dataset["train"]["image_id"]) | set(dataset["validation"]["image_id"])
         available_images = [img for img in image_files if img not in annotated_images]
-        # Check if the user already has an image
+
         if session_data["current_image"] is None and available_images:
-            # Assign a new random image to the user
             session_data["current_image"] = random.choice(available_images)
     return os.path.join(image_folder, session_data["current_image"]) if session_data["current_image"] else None
 
-# Function to save the annotation to Hugging Face dataset and fetch the next image
+
+# Function to save the annotation to the correct split in the Hugging Face dataset and fetch the next image
 def save_annotation(caption, session_data):
     if session_data["current_image"] is None:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
@@ -49,20 +62,24 @@ def save_annotation(caption, session_data):
     with lock:
         image_id = session_data["current_image"]
 
+        # Determine the correct split for the image
+        split = determine_split(image_id)
+        if split is None:
+            return gr.update(value=None), gr.update(value="Error: Could not determine split.")
+
         # Save caption or "skipped" based on user input
         if caption.strip().lower() == "skip":
             caption = "skipped"
 
-        # Add the new annotation as a new row to the dataset
+        # Add the new annotation to the corresponding split
         new_data = Dataset.from_dict({"image_id": [image_id], "caption": [caption]})
-        global dataset
-        dataset = concatenate_datasets([dataset, new_data])
+        dataset[split] = concatenate_datasets([dataset[split], new_data])
 
         # Save updated dataset to Hugging Face
         dataset.push_to_hub(dataset_name)
-        print("Pushed updated dataset")
+        print(f"Pushed updated {split} dataset")
 
-        # Clear user's current image so they get a new one next time
+        # Clear user's current image to get a new one next time
         session_data["current_image"] = None
 
     # Fetch the next image
@@ -72,10 +89,12 @@ def save_annotation(caption, session_data):
     else:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
 
+
 # Function to skip the current image
 def skip_image(session_data):
     return save_annotation("skip", session_data)
 
+
 # Function to initialize the interface
 def initialize_interface(session_data):
     next_image = get_next_image(session_data)
@@ -84,10 +103,12 @@ def initialize_interface(session_data):
     else:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
 
+
 # Build the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Captioning Tool")
-    gr.Markdown("Please provide a caption for each image displayed. Click 'Submit' after writing your caption, or type 'skip' if you don’t want to annotate this image.")
+    gr.Markdown(
+        "Please provide a caption for each image displayed. Click 'Submit' after writing your caption, or type 'skip' if you don’t want to annotate this image.")
 
     session_data = gr.State({"current_image": None})  # Session-specific state
 
@@ -104,4 +125,4 @@ with gr.Blocks() as demo:
     # Load initial image
     demo.load(fn=initialize_interface, inputs=session_data, outputs=[image, caption])
 
-demo.launch(share=True)
+demo.launch(share=True)