diff --git a/library/train_util.py b/library/train_util.py index 8b5cf214e..7f396d36e 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -1405,11 +1405,11 @@ def get_image_size(self, 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}") + 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