| 44 | |
| 45 | |
| 46 | class WebDataModuleFromConfig(pl.LightningDataModule): |
| 47 | def __init__( |
| 48 | self, |
| 49 | tar_base, |
| 50 | batch_size, |
| 51 | image_size, |
| 52 | train=None, |
| 53 | validation=None, |
| 54 | test=None, |
| 55 | num_workers=4, |
| 56 | multinode=True, |
| 57 | is_video=False, |
| 58 | min_size=None, |
| 59 | max_pwatermark=1.0, |
| 60 | video_frames=0, |
| 61 | channel_last=False, |
| 62 | val_batch_size=None, |
| 63 | val_num_workers=None, |
| 64 | **kwargs, |
| 65 | ): |
| 66 | super().__init__() |
| 67 | print(f"Setting tar base to {tar_base}") |
| 68 | self.tar_base = tar_base |
| 69 | self.batch_size = batch_size |
| 70 | self.image_size = image_size |
| 71 | self.num_workers = num_workers |
| 72 | self.train = train |
| 73 | self.validation = validation |
| 74 | self.test = test |
| 75 | self.video_frames = video_frames |
| 76 | self.multinode = multinode |
| 77 | self.min_size = min_size # filter out very small images |
| 78 | self.max_pwatermark = max_pwatermark # filter out watermarked images |
| 79 | self.channel_last = channel_last |
| 80 | self.val_batch_size = ( |
| 81 | val_batch_size if val_batch_size is not None else batch_size |
| 82 | ) |
| 83 | self.val_num_workers = ( |
| 84 | val_num_workers if val_num_workers is not None else num_workers |
| 85 | ) |
| 86 | self.is_video = is_video |
| 87 | |
| 88 | def make_loader(self, dataset_config, train=True): |
| 89 | # change range from [0,1] to [-1,1] and put channel last or first |
| 90 | image_transforms = [] |
| 91 | if self.channel_last: |
| 92 | lambda_fn = lambda x: rearrange(x * 2.0 - 1.0, "c h w -> h w c") |
| 93 | else: |
| 94 | lambda_fn = lambda x: x * 2.0 - 1.0 |
| 95 | |
| 96 | image_transforms.extend( |
| 97 | [ |
| 98 | torchvision.transforms.ToTensor(), |
| 99 | torchvision.transforms.Resize(self.image_size, antialias=True), |
| 100 | torchvision.transforms.ConvertImageDtype(torch.float32), |
| 101 | torchvision.transforms.Lambda(lambda_fn), |
| 102 | ] |
| 103 | ) |
no outgoing calls