diff --git a/mokuro/manga_page_ocr.py b/mokuro/manga_page_ocr.py index 2157bb4..b5ccacd 100644 --- a/mokuro/manga_page_ocr.py +++ b/mokuro/manga_page_ocr.py @@ -1,5 +1,7 @@ +import os import cv2 import numpy as np +import pillow_avif # noqa: F401 from PIL import Image from loguru import logger from scipy.signal.windows import gaussian @@ -45,7 +47,17 @@ def __init__( self.mocr = MangaOcr(pretrained_model_name_or_path, force_cpu) def __call__(self, img_path): - img = imread(img_path) + # Check if the file is in AVIF format + # Create a temporary png path + img_path = str(img_path) + if img_path.lower().endswith('.avif'): + temp_path = img_path.replace('.avif', '_temp.png') + avif_img = Image.open(img_path) + avif_img.save(temp_path, 'PNG') + img = imread(temp_path) + os.remove(temp_path) + else: + img = imread(img_path) if img is None: raise InvalidImage() H, W, *_ = img.shape diff --git a/mokuro/volume.py b/mokuro/volume.py index 489752b..2dbaad3 100644 --- a/mokuro/volume.py +++ b/mokuro/volume.py @@ -106,7 +106,7 @@ def get_img_paths(self): img_paths = natsorted( p.relative_to(self.path_in) for p in self.path_in.glob("**/*") - if p.is_file() and p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp") + if p.is_file() and p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp", ".avif") ) img_paths = {p.with_suffix(""): p for p in img_paths} return img_paths diff --git a/pyproject.toml b/pyproject.toml index 7583161..4be27f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ "numpy", "opencv-python>=4.1.2", "Pillow>=7.1.2", + "pillow-avif-plugin >= 1.4.6", "pyclipper", "requests", "scipy",