File size: 598 Bytes
989c6ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | """Indexed deterministic cubed-sphere samples for training and validation."""
from torch.utils.data import Dataset
from .fake_data import make_fake_batch
class FakeCubeSphereDataset(Dataset):
def __init__(self, samples=64, seed=7, channels=2, height=8, width=8):
self.samples, self.seed = samples, seed
self.shape = (channels, 6, height, width)
def __len__(self):
return self.samples
def __getitem__(self, index):
return make_fake_batch(batch=1, channels=self.shape[0], faces=6, height=self.shape[2], width=self.shape[3], seed=self.seed + index)[0]
|