diff --git a/testing/test_bucket_dataloader.py b/testing/test_bucket_dataloader.py index 3d8419b..cc6800f 100644 --- a/testing/test_bucket_dataloader.py +++ b/testing/test_bucket_dataloader.py @@ -32,7 +32,7 @@ bucket_tolerance = 64 batch_size = 4 dataset_config = DatasetConfig( - folder_path=dataset_folder, + dataset_path=dataset_folder, resolution=resolution, caption_ext='txt', default_caption='default', @@ -48,22 +48,22 @@ for batch in dataloader: batch: 'DataLoaderBatchDTO' img_batch = batch.tensor - chunks = torch.chunk(img_batch, batch_size, dim=0) - # put them so they are size by side - big_img = torch.cat(chunks, dim=3) - big_img = big_img.squeeze(0) - - min_val = big_img.min() - max_val = big_img.max() - - big_img = (big_img / 2 + 0.5).clamp(0, 1) - - # convert to image - img = transforms.ToPILImage()(big_img) - - show_img(img) - - time.sleep(1.0) + # chunks = torch.chunk(img_batch, batch_size, dim=0) + # # put them so they are size by side + # big_img = torch.cat(chunks, dim=3) + # big_img = big_img.squeeze(0) + # + # min_val = big_img.min() + # max_val = big_img.max() + # + # big_img = (big_img / 2 + 0.5).clamp(0, 1) + # + # # convert to image + # img = transforms.ToPILImage()(big_img) + # + # show_img(img) + # + # time.sleep(1.0) cv2.destroyAllWindows() diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 7fd5295..9aaa3a7 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -251,8 +251,8 @@ class ImageProcessingDTOMixin: transform: Union[None, transforms.Compose] ): # todo make sure this matches - img = Image.open(self.path).convert('RGB') try: + img = Image.open(self.path).convert('RGB') img = exif_transpose(img) except Exception as e: print(f"Error: {e}")