| 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 | ) |
| 104 | |
| 105 | if "image_transforms" in dataset_config: |
| 106 | image_transforms.extend( |
| 107 | [instantiate_from_config(tt) for tt in dataset_config.image_transforms] |
| 108 | ) |
| 109 | image_transforms = torchvision.transforms.Compose(image_transforms) |
| 110 | |
| 111 | if "transforms" in dataset_config: |
| 112 | transforms_config = OmegaConf.to_container(dataset_config.transforms) |
| 113 | else: |
| 114 | transforms_config = dict() |
| 115 | |
| 116 | transform_dict = { |
| 117 | dkey: ( |
| 118 | load_partial_from_config(transforms_config[dkey]) |
| 119 | if transforms_config[dkey] != "identity" |
| 120 | else identity |
| 121 | ) |
| 122 | for dkey in transforms_config |
| 123 | } |
| 124 | # this is crucial to set correct image key to get the transofrms applied correctly |
| 125 | img_key = dataset_config.get("image_key", "image.png") |
| 126 | transform_dict.update({img_key: image_transforms}) |
| 127 | |
| 128 | if "postprocess" in dataset_config: |
| 129 | postprocess = instantiate_from_config(dataset_config["postprocess"]) |
| 130 | else: |
| 131 | postprocess = None |
| 132 | print("No postprocess") |
| 133 | |
| 134 | # some illustration how shuffle works in webdataset |
| 135 | # https://github.com/webdataset/webdataset/issues/71 |
| 136 | # TL;DR: set the shuffle as big as you can afford ->len(files) |
| 137 | shuffle = dataset_config.get("shuffle", 0) |
| 138 | shardshuffle = shuffle > 0 |
| 139 | |
| 140 | nodesplitter = ( |
| 141 | wds.shardlists.split_by_node |
| 142 | if self.multinode |
| 143 | else wds.shardlists.single_node_only |
| 144 | ) |
| 145 | |