dataset.py

import torch

class RetailDataset(torch.utils.data.Dataset):
    def __init__(self, data, labels=None, transform=None):
        self.data = data
        self.labels = labels
        self.num_classes = len(set(labels))
        self.transform = transform

    def __getitem__(self, idx):
        item = {key: val[idx].detach().clone() for key, val in self.data.items()}
        item['labels'] = self.labels[idx]
        return item

    def __len__(self):
        return len(self.labels)

    def __repr__(self):
        return 'RetailDataset'

    def __str__(self):
        return str({
            'data': self.data['pixel_values'].shape,
            'labels': self.labels.shape,
            'num_classes': self.num_classes,
            'num_samples': len(self.labels)
        })