File size: 1,016 Bytes
158f4dc 0f734ea 158f4dc 0f734ea 158f4dc 0f734ea 158f4dc 0f734ea 158f4dc 0f734ea 158f4dc 0f734ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import torch
class RetailDataset(torch.utils.data.Dataset):
def __init__(self, data, labels=None, transform=None, device=None):
self.data = data
self.labels = labels
self.num_classes = len(set(labels))
self.transform = transform
self.device = device if device is not None else torch.device("cpu")
def __getitem__(self, idx):
item = {
key: torch.tensor(val[idx].detach().clone(), device=self.device)
for key, val in self.data.items()
}
item["labels"] = torch.tensor(self.labels[idx], device=self.device)
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),
}
)
|