mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix to image_folder_dataset to accomodate images with mismatched dimensions

Browse Source
This commit is contained in:
James Betker 2021-01-10 12:57:21 -07:00
parent 48f0d8964b
commit 4119cd6240
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
codes/data/image_folder_dataset.py
View File

@ -52,9 +52,7 @@ class ImageFolderDataset:
imgs = torch.load(cache_path)
else:
print("Building image folder cache, this can take some time for large datasets..")
imgs = []
for ext in supported_types:
imgs.extend(glob.glob(os.path.join(path, "*." + ext)))
imgs = util.get_image_paths('img', path)[0]
torch.save(imgs, cache_path)
for w in range(weight):
self.image_paths.extend(imgs)
@ -67,6 +65,7 @@ class ImageFolderDataset:
def resize_hq(self, imgs_hq):
# Enforce size constraints
h, w, _ = imgs_hq[0].shape
if self.target_hq_size is not None and self.target_hq_size != h:
hqs_adjusted = []
for hq in imgs_hq:
@ -114,6 +113,11 @@ class ImageFolderDataset:
if not self.disable_flip and random.random() < .5:
hq = hq[:, ::-1, :]
# We must convert the image into a square.
h, w, _ = hq.shape
dim = min(h, w)
hq = hq[(h - dim) // 2:dim + (h - dim) // 2, (w - dim) // 2:dim + (w - dim) // 2, :]
if self.labeler:
assert hq.shape[0] == hq.shape[1] # This just has not been accomodated yet.
dim = hq.shape[0]