Efficient Image Augmentation in Hugging Face Datasets

[fabiozappo]

Describe the bug

I'm using the Hugging Face datasets library to load images in batch and would like to apply a torchvision transform to solve the inconsistent image sizes in the dataset and apply some on the fly image augmentation. I can just think about using the collate_fn, but seems quite inefficient.

I'm new to the Hugging Face datasets library, I didn't find nothing in the documentation or the issues here on github.

Is there an existing way to add image transformations directly to the dataset loading pipeline?

Steps to reproduce the bug

from datasets import load_dataset
from torch.utils.data import DataLoader

def collate_fn(batch):
    images = [item['image'] for item in batch]
    texts = [item['text'] for item in batch]
    return {
        'images': images,
        'texts': texts
    }

dataset = load_dataset("Yuki20/pokemon_caption", split="train")
dataloader = DataLoader(dataset, batch_size=4, collate_fn=collate_fn)

# Output shows varying image sizes:
# [(1280, 1280), (431, 431), (789, 789), (769, 769)]

Expected behavior

I'm looking for a way to resize images on-the-fly when loading the dataset, similar to PyTorch's Dataset.getitem functionality. This would be more efficient than handling resizing in the collate_fn.

Environment info

• datasets version: 3.1.0
• Platform: Linux-6.5.0-41-generic-x86_64-with-glibc2.35
• Python version: 3.11.10
• huggingface_hub version: 0.26.2
• PyArrow version: 18.0.0
• Pandas version: 2.2.3
• fsspec version: 2024.9.0