| import torchvision |
|
|
| """ |
| BabelImageNet from https://arxiv.org/pdf/2306.08658.pdf |
| Adapted from https://github.com/gregor-ge/Babel-ImageNet, thanks to the authors |
| """ |
| class BabelImageNet(torchvision.datasets.ImageNet): |
| def __init__(self, root: str, idxs, split: str = "val", download=None, **kwargs) -> None: |
| super().__init__(root, split, **kwargs) |
| examples_per_class = len(self.targets) // 1000 |
| select_idxs = [idx*examples_per_class + i for idx in idxs for i in range(examples_per_class)] |
| self.targets = [i for i in range(len(idxs)) for _ in range(examples_per_class)] |
| self.imgs = [self.imgs[i] for i in select_idxs] |
| self.samples = [self.samples[i] for i in select_idxs] |
| self.idxs = idxs |
|
|
| def __getitem__(self, i): |
| img, target = super().__getitem__(i) |
| target = self.idxs.index(target) |
| return img, target |
