【HuggingFace 如何上传数据集 (3) 】Dataset-稳定上传与下载大量（上万）图片等格式的数据

最新推荐文章于 2025-01-21 15:14:36 发布

多恩Stone

最新推荐文章于 2025-01-21 15:14:36 发布

阅读量759

点赞数 5

分类专栏：编程学习文章标签： AIGC 深度学习 python hugging face 数据集

本文链接：https://blog.csdn.net/weixin_44212848/article/details/143036101

版权

编程学习专栏收录该内容

84 篇文章 4 订阅

订阅专栏

系列文章目录

文章目录

前言

【HuggingFace 如何上传数据集】快速上传图片、文本等各种格式的数据可能遇到网络不稳定的情况。
【HuggingFace 如何上传数据集 (2) 】国内网络-稳定上传图片、文本等各种格式的数据可能遇到传输限制文件数量不能超过 10000 条的情况。
本文则通过 datasets 库来进行大量的图像数据上传。

一、上传

from datasets import Dataset, Features, Array3D, Value
import os
import numpy as np
from PIL import Image

# Define image path
image_dir = "/path/to/local_img_dir"

# Load and convert images to numpy arrays with resizing
def load_image(image_path):
    image = Image.open(image_path).resize((512, 512))  # Resize to 512x512
    return np.array(image)

# Generate image file paths list
image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(".png")]

# Create dataset, returning a dictionary with keys 'image' and 'filename'
def generate_dataset():
    for path in image_paths:
        yield {"image": load_image(path), "filename": os.path.basename(path)}

dataset = Dataset.from_generator(
    generate_dataset,
    features=Features({
        "image": Array3D(dtype="uint8", shape=(512, 512, 3)),  # Updated shape
        "filename": Value("string")
    })
)

# Set batch size to avoid memory overflow
dataset = dataset.with_format("numpy", writer_batch_size=50)  # Adjust as necessary

# Push dataset to Hugging Face Hub
dataset.push_to_hub("your/data_repo")

注意每个 batch 不能超过 2GB，所以可能需要调节下 writer_batch_size。
Array3D(dtype=“uint8”, shape=(512, 512, 3) shape 的第二个维度不能为 None（动态轴）。

二、下载

保证服务器能够 curl 到 huggingface 是最关键的。
通过 ping google.com 或 curl https://huggingface.co 来测下网络有没有问题。

from datasets import load_dataset
from PIL import Image
import numpy as np
import os

HF_TOKEN = 'hf_yourtoken' # 如果是 private 需要用特定的 token 来鉴权

# Load the dataset from Hugging Face Hub
dataset = load_dataset("your/data_repo", token=HF_TOKEN)

# Directory where the images will be saved
save_dir = "/local/dir"

# Ensure the directory exists
os.makedirs(save_dir, exist_ok=True)

# Function to convert numpy array to PIL image
def convert_to_image(image_array):
    return Image.fromarray(np.uint8(image_array))

# 用于测试，先下载前 10 张看看，不过 load_dataset 会直接将所有数据（.parquet）都下载下来
# Save the first 10 images 
for i, example in enumerate(dataset['train'].select(range(10))):  # Select the first 10 examples
    image_array = example['image']  # Get the numpy array for the image
    image = convert_to_image(image_array)  # Convert to PIL image
    filename = example['filename']  # Get the filename (without extension)
    
    # Construct the full file path and save the image as a PNG
    image_path = os.path.join(save_dir, f"{filename}.png")
    image.save(image_path)
    
    print(f"Saved {filename}.png to {save_dir}")