From b811a39c8fb5ddb4d91b0ba45996ddcedc5453b1 Mon Sep 17 00:00:00 2001 From: Patrick Esser Date: Sun, 31 Jul 2022 23:24:45 +0000 Subject: [PATCH] make filtering more robust --- ldm/data/laion.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/ldm/data/laion.py b/ldm/data/laion.py index 7a41d4a..a32d699 100644 --- a/ldm/data/laion.py +++ b/ldm/data/laion.py @@ -185,10 +185,19 @@ class WebDataModuleFromConfig(pl.LightningDataModule): return loader def filter_size(self, x): - if self.min_size is None: - return True try: - return x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size and x['json']['pwatermark'] <= self.max_pwatermark + valid = True + if self.min_size is not None and self.min_size > 1: + try: + valid = valid and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size + except Exception: + valid = False + if self.max_pwatermark is not None and self.max_pwatermark < 1.0: + try: + valid = valid and x['json']['pwatermark'] <= self.max_pwatermark + except Exception: + valid = False + return valid except Exception: return False