| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| class RandomSubsetBatchDataset(Dataset): | |
| def __init__(self, pt_files, subset_size, seed=42): | |
| self.all_files = [str(Path(path)) for path in pt_files] | |
| if not self.all_files: | |
| raise ValueError("RandomSubsetBatchDataset requires at least one batch file") | |
| self.subset_size = max(1, min(int(subset_size), len(self.all_files))) | |
| self.seed = int(seed) | |
| self.epoch = 0 | |
| self.active_files = self.all_files[: self.subset_size] | |
| self.set_epoch(0) | |
| def set_epoch(self, epoch): | |
| self.epoch = int(epoch) | |
| if self.subset_size >= len(self.all_files): | |
| self.active_files = list(self.all_files) | |
| return | |
| rng = np.random.default_rng(self.seed + self.epoch) | |
| indices = rng.choice(len(self.all_files), size=self.subset_size, replace=False) | |
| self.active_files = [self.all_files[int(idx)] for idx in indices] | |
| def __len__(self): | |
| return len(self.active_files) | |
| def __getitem__(self, index): | |
| path = self.active_files[index] | |
| sample = torch.load(path, weights_only=False) | |
| sample["source_path"] = path | |
| sample["tile_name"] = Path(path).stem.split("_batch_")[0] | |
| return sample | |