diff --git a/prepare_data.py b/prepare_data.py index 6f7f829..2ccb9fb 100755 --- a/prepare_data.py +++ b/prepare_data.py @@ -33,9 +33,12 @@ def resize_multiple( def resize_worker(img_file, sizes, resample): i, file = img_file - img = Image.open(file) - img = img.convert("RGB") - out = resize_multiple(img, sizes=sizes, resample=resample) + try: + img = Image.open(file) + img = img.convert("RGB") + out = resize_multiple(img, sizes=sizes, resample=resample) + except: + out = [] return i, out @@ -51,13 +54,14 @@ def prepare( with multiprocessing.Pool(n_worker) as pool: for i, imgs in tqdm(pool.imap_unordered(resize_fn, files)): - for size, img in zip(sizes, imgs): - key = f"{size}-{str(i).zfill(5)}".encode("utf-8") + if len(imgs) != 0: + for size, img in zip(sizes, imgs): + key = f"{size}-{str(i).zfill(5)}".encode("utf-8") - with env.begin(write=True) as txn: - txn.put(key, img) + with env.begin(write=True) as txn: + txn.put(key, img) - total += 1 + total += 1 with env.begin(write=True) as txn: txn.put("length".encode("utf-8"), str(total).encode("utf-8")) diff --git a/stylegan2/dataset.py b/stylegan2/dataset.py index 7713ea2..1624715 100755 --- a/stylegan2/dataset.py +++ b/stylegan2/dataset.py @@ -32,7 +32,9 @@ def __getitem__(self, index): with self.env.begin(write=False) as txn: key = f'{self.resolution}-{str(index).zfill(5)}'.encode('utf-8') img_bytes = txn.get(key) - + + + img_bytes.seek(0) buffer = BytesIO(img_bytes) img = Image.open(buffer) img = self.transform(img)