import deeplake
import torch
from PIL import Image
import numpy as np
from torch.utils.data import Dataset

class FlickrStreamer(Dataset):
    def __init__(self, limit=1000, transform=None):
        # Streams data from cloud. No massive download needed.
        print(f"Connecting to Flickr30k (First {limit} images)...")
        # Windows sometimes has permission issues with cache, so we explicitly set access_method
        self.ds = deeplake.load('hub://activeloop/flickr30k', access_method="stream")
        self.limit = limit
        self.transform = transform

    def __len__(self):
        return self.limit

    def __getitem__(self, idx):
        sample = self.ds[idx]
        image_data = sample.image.numpy()
        
        try:
            image = Image.fromarray(image_data)
        except:
            image = Image.fromarray(image_data.astype('uint8'))
            
        if self.transform:
            image = self.transform(image)
            
        return image