add upload to dataset

utils/upload-to-dataset.py

@@ -0,0 +1,76 @@
+from datasets import Dataset, Features, Value, Image
+from huggingface_hub import HfApi
+import os
+from collections import defaultdict
+import pandas as pd
+
+# Define the path to your images
+IMAGE_DIR = "../../background-removal-arena-v0/train/data/resized"
+
+# Define the dataset features with dedicated columns for each model
+features = Features({
+    "original_image": Image(),  # Original image feature
+    "clipdrop_image": Image(),  # Clipdrop segmented image
+    "bria_image": Image(),      # Bria segmented image
+    "photoroom_image": Image(), # Photoroom segmented image
+    "removebg_image": Image(),  # RemoveBG segmented image
+    "original_filename": Value("string")  # Original filename
+})
+
+# Load image paths and metadata
+data = defaultdict(lambda: {
+    "clipdrop_image": None,
+    "bria_image": None,
+    "photoroom_image": None,
+    "removebg_image": None
+})
+
+# Walk into the web-original-images folder
+web_original_images_dir = os.path.join(IMAGE_DIR, "web-original-images")
+for root, _, files in os.walk(web_original_images_dir):
+    for f in files:
+        if f.endswith(('.png', '.jpg', '.jpeg')):
+            original_image_path = os.path.join(root, f)
+            data[f]["original_image"] = original_image_path
+            data[f]["original_filename"] = f
+
+            # Check for corresponding images in other directories
+            for source in ["clipdrop", "bria", "photoroom", "removebg"]:
+                # Check for processed images ending in .png or .jpg
+                for ext in ['.png', '.jpg']:
+                    processed_image_filename = os.path.splitext(f)[0] + ext
+                    source_image_path = os.path.join(IMAGE_DIR, source, processed_image_filename)
+    
+                    if os.path.exists(source_image_path):
+                        data[f][f"{source}_image"] = source_image_path
+                        break  # Stop checking other extensions if a file is found
+
+# Convert the data to a dictionary of lists
+dataset_dict = {
+    "original_image": [],
+    "clipdrop_image": [],
+    "bria_image": [],
+    "photoroom_image": [],
+    "removebg_image": [],
+    "original_filename": []
+}
+
+for filename, entry in data.items():
+    if "original_image" in entry:
+        dataset_dict["original_image"].append(entry["original_image"])
+        dataset_dict["clipdrop_image"].append(entry["clipdrop_image"])
+        dataset_dict["bria_image"].append(entry["bria_image"])
+        dataset_dict["photoroom_image"].append(entry["photoroom_image"])
+        dataset_dict["removebg_image"].append(entry["removebg_image"])
+        dataset_dict["original_filename"].append(filename)
+
+# Save the data dictionary to a CSV file for inspection
+df = pd.DataFrame.from_dict(dataset_dict)
+df.to_csv("image_data.csv", index=False)
+
+# Create a Dataset
+dataset = Dataset.from_dict(dataset_dict, features=features)
+
+# Push the dataset to Hugging Face Hub
+api = HfApi()
+dataset.push_to_hub("bgsys/background-removal-arena-test", token=api.token)