diff --git a/library/train_util.py b/library/train_util.py index 4c3728e14..ed546af32 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -1404,13 +1404,13 @@ 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}") - image_size = (0, 0) + # imagesize doesn't work for some images, so use PIL + 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}, {e}") + image_size = (0, 0) return image_size def load_image_with_face_info(self, subset: BaseSubset, image_path: str, alpha_mask=False):