Error in the Getting Started Example -- default_collate cannot handle PIL.Images

[JacobARose]

🐛 Bug

Hello! Excited about making use of litdata.optimize to solve some long standing data loading issues Ive encountered many times before.

While trying to get it working with my dataset, I tried running through the getting started example that creates binary files of a fake dataset consisting of random noise images, but pretty quickly ran into an error that should've been caught.

Running litdata.optimize() runs without an issue, and the streaming.py also technically runs without issue. But, if I try to do anything with the DataLoader defined at the end of streaming.py, I get a default_collate error because it is incapable of batching PIL.Images.

Given that pretty much all of the examples I could find in the documentation rely heavily on PIL.Image, I'm wondering

A. Can we update the getting started example so that it at least defines a DataLoader that wont immediately throw an error?
B. Can we get some more explanation/exploration into whether it's best to save images as PIL.Image within the optimized files, or if it's better to convert them to torch.Tensor/np.Arrays in advance so they're not repeatedly converted every epoch (My intuition is leaning towards the latter)

To Reproduce

If desired, I can provide more details on the environment I used to produce this error, but given that it's already based on the minimal working example I anticipate it shouldnt require much imagination to reproduce.

[github-actions]

Hi! thanks for your contribution!, great first issue!

[tchaton]

Hey @JacobARose.

Would you mind to share the exact script that you run ?

[JacobARose]

Hey @JacobARose.

Would you mind to share the exact script that you run ?

Thanks @tchaton for the quick reply! You can find the exact code I ran in the generate.py and streaming.py files located in the directory found at the following commit:
https://github.com/Lightning-AI/litData/tree/9fb801916a6acae9e944fa64c3549009e0c73566/examples/getting_started

[bhimrazy]

Thanks for pointing this out! @JacobARose The given sample is just a minimal demo — in a real-world case, you’d typically work with actual .jpg/.png images, which get saved and loaded as tensors directly.

If you’re saving PIL images directly, you can handle them by either subclassing the streaming dataset to apply transforms or by passing a custom collate_fn like:

def collate_fn(batch):
    return {
        "image": [sample["image"] for sample in batch],
        "class": [sample["class"] for sample in batch],
    }

train_dataloader = ld.StreamingDataLoader(train_dataset, collate_fn=collate_fn)

As for whether to save images as PIL or preconvert to tensors/arrays — it’s generally recommended to pass as JPEG for better storage optimization and compression. You can also apply resizing and quality settings at this stage for even more space savings and faster loading downstream.

Totally agree the example should be updated so it works smoothly out of the box.