File size: 4,934 Bytes
c134030 e587a3a c134030 e587a3a c134030 e587a3a c134030 e587a3a c134030 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
from datasets import Dataset, Features, Value, Image
from huggingface_hub import HfApi
import os
from collections import defaultdict
import pandas as pd
import argparse
from PIL import Image as PILImage
import sys
def upload_to_dataset(original_images_dir, processed_images_dir, dataset_name, dry_run=False):
# 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 original images folder
for root, _, files in os.walk(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 processed 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(processed_images_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": []
}
errors = []
for filename, entry in data.items():
if "original_image" in entry:
# Check if all images have the same size
try:
original_size = PILImage.open(entry["original_image"]).size
for source in ["clipdrop_image", "bria_image", "photoroom_image", "removebg_image"]:
if entry[source] is not None:
processed_size = PILImage.open(entry[source]).size
if processed_size != original_size:
errors.append(f"Size mismatch for {filename}: {source} image size {processed_size} does not match original size {original_size}.")
except Exception as e:
errors.append(f"Error processing {filename}: {e}")
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)
if errors:
for error in errors:
print(error)
sys.exit(1)
# 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)
if dry_run:
print("Dry run: Dataset prepared but not pushed to Hugging Face Hub.")
print(df.head()) # Display the first few rows for inspection
else:
# Push the dataset to Hugging Face Hub in a private way
api = HfApi()
dataset.push_to_hub(dataset_name, token=api.token, private=True)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Upload images to a Hugging Face dataset.")
parser.add_argument("original_images_dir", type=str, help="Directory containing the original images.")
parser.add_argument("processed_images_dir", type=str, help="Directory containing the processed images with subfolders for each model.")
parser.add_argument("dataset_name", type=str, help="Name of the dataset to upload to Hugging Face Hub.")
parser.add_argument("--dry-run", action="store_true", help="Perform a dry run without uploading to the hub.")
args = parser.parse_args()
upload_to_dataset(args.original_images_dir, args.processed_images_dir, args.dataset_name, dry_run=args.dry_run)
|