Handle capital letter image extensions (#120)

PixArt-alpha · Jun 28, 2024 · 2b08ce5 · 2b08ce5
1 parent cefaa95
commit 2b08ce5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/diffusion/data/datasets/InternalData.py b/diffusion/data/datasets/InternalData.py
@@ -14,7 +14,8 @@
 
 # helper function for replacing image extensions with an another
 def replace_img_ext(path, dst_ext: str) -> str:
-    return path.replace('.png', dst_ext).replace('.jpg', dst_ext).replace('.webp', dst_ext)
+    return path.replace('.png', dst_ext).replace('.jpg', dst_ext).replace('.webp', dst_ext).replace('.jpeg', dst_ext).\
+        replace('.JPEG', dst_ext).replace('.JPG', dst_ext)
 
 
 @DATASETS.register_module()

diff --git a/tools/convert_images_to_json.py b/tools/convert_images_to_json.py
@@ -51,7 +51,7 @@ def main():
                 # check if an image exists for this caption
                 image_filename = filename[:-len(caption_extension)]
 
-                for image_extension in ['.jpg', '.png', '.jpeg', 'webp']:
+                for image_extension in ['.jpg', '.png', '.jpeg', 'webp', '.JPEG', '.JPG']:
                     image_path = Path(dirpath).joinpath(image_filename + image_extension)
                     if path.exists(image_path):
                         write_entry(json_entries, dirpath, image_path, Path(dirpath).joinpath(filename), image_filename + image_extension, intern_imgs_folder)