keep pil data in loading?

[rxqy]

Hi, I'm quite new to litdata and am looking into the imagenet demo here, just wondering if it is possible to return the original pil image here instead of pt tensor? Not sure what's the correct way to past code from studio here, but I'm following the demo here.

# stream/lightning_data.py
    def __getitem__(self, index):
        # Note: If torchvision is installed, we return a tensor image instead of a pil image as it is much faster. 
        img, class_index = super().__getitem__(index) # <- Whatever you returned from the DatasetOptimizer prepare_item method.
        return self.transform(to_rgb(img)), int(class_index)    

[deependujha]

Hi @rxqy, sure you can.

Just be sure that when you optimized your dataset (optimize), you returned a pil image rather than a tensor.

For details on optimize fn, refer readme

streamingDataset will yield in the respective format.

[rxqy]

Hi @deependujha, in the imagenet demo we are returning the pil img right? But I'm still getting pt tensor in my streaming dataset. It seems that litdata is automatically doing the convertion for me.

I made some small modification to the optimize_fn so I can double check the result, should make no difference here.

# running with python lightning_data.py --input_dir ILSVRC2012/val --output_dir ./data/
def optimize_fn(data):
    filepath, class_index = data 
    img = Image.open(filepath)
    return img

inputs = get_inputs(args.input)
sample = optimize_fn(inputs[0])
print(sample) 
# -> gives <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x375 at 0x7F1C40363BE0>

While in my streaming dataset:

data = super().__getitem__(index)
print(data.shape)
# -> gives a pt tensor, torch.Size([3, 375, 500]))

[deependujha]

when you call optimize function, data is serialized before writing to the chunks. src/litdata/streaming/writer.py file checks which serializer can serialize the data (serializers.py file)

In your case, JPEGSerializer might be the best fit, and if you look at its deserialize code

def deserialize(self, data: bytes) -> Union["JpegImageFile", torch.Tensor]:
        if _TORCH_VISION_AVAILABLE:
            from torchvision.io import decode_jpeg
            from torchvision.transforms.functional import pil_to_tensor

            array = torch.frombuffer(data, dtype=torch.uint8)
            # Note: Some datasets like Imagenet contains some PNG images with JPEG extension, so we fallback to PIL
            with suppress(RuntimeError):
                return decode_jpeg(array)

        img = PILSerializer.deserialize(data)
        if _TORCH_VISION_AVAILABLE:
            img = pil_to_tensor(img)
        return img

If torchvision is available, it tries to use it as it is much faster.

btw, you can convert back from torchvision to pil image using: to_pil_image

[rxqy]

@deependujha I see. Many thanks for the help and detailed explain!
My problem here is whether I can keep the data in pil format (without the to_pil_image and pil_to_tensor convertion, which should be slow I guess?). I understand that torchvisions' v2 transforms is much faster, but our old data transform code must take pil_image as input and we have to stick to that.

[rxqy]

Oh I found the PILSerializer here. A quick fix would be return the img in PIL.Image.Image format (not PIL.JpegImagePlugin.JpegImageFile). Many thanks!

def optimize_fn(data):
    filepath, class_index = data 
    img = Image.open(filepath).convert("RGB")
    return img