updates

app.py

@@ -26,8 +26,9 @@ except Exception as e:
     })
     train_dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
     val_dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
-    train_dataset.push_to_hub(dataset_name, split="train")  # Push the empty train dataset to Hugging Face
-    val_dataset.push_to_hub(dataset_name, split="validation")  # Push the empty validation dataset to Hugging Face
+    dataset_dict = {"train": train_dataset, "validation": val_dataset}
+    for split_name, split_dataset in dataset_dict.items():
+        split_dataset.push_to_hub(dataset_name, split=split_name)
 
 image_folder = "images"
 image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
@@ -38,13 +39,12 @@ def get_next_image(session_data):
     with lock:
         annotated_images = set(train_dataset["image_id"]) | set(val_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:
             # 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 datasets and fetch the next image
+# Function to save the annotation to the correct split in Hugging Face datasets 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!")
@@ -66,11 +66,11 @@ def save_annotation(caption, session_data):
         else:
             return gr.update(visible=False), gr.update(value="Unknown dataset split for image!")
 
-        # Add the new annotation as a new row to the appropriate dataset
+        # Add the new annotation to the correct dataset without overwriting
         new_data = Dataset.from_dict({"image_id": [image_id], "caption": [caption]})
         target_dataset = concatenate_datasets([target_dataset, new_data])
 
-        # Save updated dataset to Hugging Face
+        # Save the updated split to Hugging Face without overwriting
         target_dataset.push_to_hub(dataset_name, split=split_name)
         print(f"Pushed updated {split_name} dataset")