diff --git a/library/train_util.py b/library/train_util.py index 8b5cf214e..a5d6fdd21 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -1404,12 +1404,12 @@ def get_image_size(self, image_path): # return imagesize.get(image_path) image_size = imagesize.get(image_path) if image_size[0] <= 0: - # imagesize doesn't work for some images, so use cv2 - img = cv2.imread(image_path) - if img is not None: - image_size = (img.shape[1], img.shape[0]) - else: - logger.warning(f"failed to get image size: {image_path}") + # imagesize doesn't work for some images, so use PIL as a fallback + try: + with Image.open(image_path) as img: + image_size = img.size + except Exception as e: + logger.warning(f"failed to get image size: {image_path}, error: {e}") image_size = (0, 0) return image_size