| from pathlib import Path | |
| import random | |
| import torch | |
| from torch.utils.data import Dataset | |
| class DeepChoiceDataset(Dataset): | |
| def __init__(self, pt_files, shuffle=False): | |
| self.pt_files = [str(Path(path)) for path in pt_files] | |
| if shuffle: | |
| random.shuffle(self.pt_files) | |
| def __len__(self): | |
| return len(self.pt_files) | |
| def __getitem__(self, index): | |
| path = self.pt_files[index] | |
| sample = torch.load(path, weights_only=False) | |
| sample["source_path"] = path | |
| sample["tile_name"] = Path(path).stem.split("_batch_")[0] | |
| return sample | |