| 167 | |
| 168 | class ObjaverseDataModuleFromConfig(pl.LightningDataModule): |
| 169 | def __init__(self, root_dir, batch_size, total_view, train=None, validation=None, |
| 170 | test=None, num_workers=4, **kwargs): |
| 171 | super().__init__(self) |
| 172 | self.root_dir = root_dir |
| 173 | self.batch_size = batch_size |
| 174 | self.num_workers = num_workers |
| 175 | self.total_view = total_view |
| 176 | |
| 177 | if train is not None: |
| 178 | dataset_config = train |
| 179 | if validation is not None: |
| 180 | dataset_config = validation |
| 181 | |
| 182 | if 'image_transforms' in dataset_config: |
| 183 | image_transforms = [torchvision.transforms.Resize(dataset_config.image_transforms.size)] |
| 184 | else: |
| 185 | image_transforms = [] |
| 186 | image_transforms.extend([transforms.ToTensor(), |
| 187 | transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))]) |
| 188 | self.image_transforms = torchvision.transforms.Compose(image_transforms) |
| 189 | |
| 190 | |
| 191 | def train_dataloader(self): |