update

app.py

@@ -4,6 +4,7 @@ import threading
 import random
 from datasets import load_dataset, Dataset, Features, Value, concatenate_datasets
 from huggingface_hub import login
+
 # Authenticate with Hugging Face
 token = os.getenv("HUGGINGFACE_TOKEN")
 if token:
@@ -12,40 +13,27 @@ else:
     print("HUGGINGFACE_TOKEN environment variable not set.")
 dataset_name = "GeorgeIbrahim/EGYCOCO"  # Replace with your dataset name
 
-# Load or create the dataset with train and val splits
+# Load or create the dataset
 try:
-    dataset = load_dataset(dataset_name)
+    dataset = load_dataset(dataset_name, split="train")
     print("Loaded existing dataset:", dataset)
 except Exception as e:
-    # Create empty datasets for train and val splits if they don't exist
+    # Create an empty dataset if it doesn't exist
     features = Features({
         'image_id': Value(dtype='string'),
         'caption': Value(dtype='string'),
     })
-    train_dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
-    val_dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
-    dataset = {"train": train_dataset, "val": val_dataset}
-    # Push empty splits to Hugging Face
-    dataset["train"].push_to_hub(f"{dataset_name}", split="train")
-    dataset["val"].push_to_hub(f"{dataset_name}", split="val")
-
-image_folder = "test"
+    dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
+    dataset.push_to_hub(dataset_name)  # Push the empty dataset to Hugging Face
+
+image_folder = "images"
 image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
 lock = threading.Lock()
 
-# Function to get the appropriate split from the image ID
-def get_split_from_image_id(image_id):
-    if "train" in image_id:
-        return "train"
-    elif "val" in image_id:
-        return "val"
-    else:
-        raise ValueError("Image ID does not contain a valid split identifier (train/val).")
-
 # 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["train"]["image_id"]) | set(dataset["val"]["image_id"])  # Combine annotated image IDs from both splits
+        annotated_images = set(dataset["image_id"])  # Set of annotated images
         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:
@@ -53,27 +41,26 @@ def get_next_image(session_data):
             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 the correct split and fetch the next image
+# Function to save the annotation to 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!")
 
     with lock:
         image_id = session_data["current_image"]
-        split = get_split_from_image_id(image_id)
 
         # 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 appropriate split
+        # Add the new annotation as a new row to the dataset
         new_data = Dataset.from_dict({"image_id": [image_id], "caption": [caption]})
         global dataset
-        dataset[split] = concatenate_datasets([dataset[split], new_data])
+        dataset = concatenate_datasets([dataset, new_data])
 
-        # Save updated split to Hugging Face
-        dataset[split].push_to_hub(dataset_name, split=split)
-        print(f"Pushed updated dataset for split: {split}")
+        # Save updated dataset to Hugging Face
+        dataset.push_to_hub(dataset_name)
+        print("Pushed updated dataset")
 
         # Clear user's current image so they get a new one next time
         session_data["current_image"] = None