This commit is contained in:
rromb 2022-06-10 11:43:49 +02:00
parent 18a9202196
commit cbbef13b7d

View file

@ -154,6 +154,7 @@ class WebDataModuleFromConfig(pl.LightningDataModule):
dset = (dset dset = (dset
.select(self.filter_keys) .select(self.filter_keys)
.decode('pil', handler=wds.warn_and_continue) .decode('pil', handler=wds.warn_and_continue)
.select(self.filter_sizes)
.map_dict(**transform_dict, handler=wds.warn_and_continue) .map_dict(**transform_dict, handler=wds.warn_and_continue)
.batched(self.batch_size, partial=False, .batched(self.batch_size, partial=False,
collation_fn=dict_collation_fn) collation_fn=dict_collation_fn)
@ -174,7 +175,7 @@ class WebDataModuleFromConfig(pl.LightningDataModule):
def filter_keys(self, x): def filter_keys(self, x):
try: try:
return ("jpg" in x) and ("txt" in x) and self.filter_size(x) return ("jpg" in x) and ("txt" in x)
except Exception: except Exception:
return False return False