diff --git a/anime_facesaver.py b/anime_facesaver.py index c71322e..f42b596 100644 --- a/anime_facesaver.py +++ b/anime_facesaver.py @@ -252,6 +252,61 @@ def face_qualifies(box, frame_shape, min_width, min_height, return True +class AvFrameReader: + """Forward-only frame reader with a cv2.VideoCapture-compatible + surface (read/release). PyAV ships its own ffmpeg build with libdav1d: + the ffmpeg bundled inside opencv-python has no software AV1 decoder, + and Alder Lake-N has no AV1 hw decode, so BD AV1 rips (e.g. SavI0r) + decode 0 frames through cv2. scenedetect likewise prefers its pyav + backend once av is installed, fixing scene detection on the same files. + """ + + def __init__(self, path): + import av + self._c = av.open(path) + self._vs = self._c.streams.video[0] + self._vs.thread_type = "AUTO" + self._iter = self._c.decode(self._vs) + + @property + def frame_count(self): + vs = self._vs + if vs.frames: + return int(vs.frames) + # MKV headers often leave stream.frames unset; estimate from the + # container duration (it is this file's total decode work anyway). + dur = self._c.duration # microseconds, or None + if dur and vs.average_rate: + return int(float(dur) / 1e6 * float(vs.average_rate)) + return 0 + + def read(self): + try: + frame = next(self._iter) + except StopIteration: + return False, None + return True, frame.to_ndarray(format="bgr24") + + def release(self): + try: + self._c.close() + except Exception: + pass + + +def open_reader(video_path): + """PyAV when available (AV1-capable), else plain cv2.""" + import cv2 + try: + return AvFrameReader(video_path) + except ImportError: + pass + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + return None + return cap + + def process_video(video_path, output_dir, args, model, device, draw, output_start=0): from scenedetect import open_video, SceneManager, ContentDetector @@ -260,7 +315,13 @@ def process_video(video_path, output_dir, args, model, device, draw, try: import cv2 - video = open_video(video_path) + try: + import av # noqa: F401 + # scenedetect defaults to backend="opencv", which cannot decode + # AV1 in this env; pyav ships its own ffmpeg with libdav1d. + video = open_video(video_path, backend="pyav") + except ImportError: + video = open_video(video_path) scene_manager = SceneManager() scene_manager.add_detector( ContentDetector(threshold=args.scene_threshold)) @@ -269,12 +330,14 @@ def process_video(video_path, output_dir, args, model, device, draw, f"{video_path}: {e}") return output_start - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): + cap = open_reader(video_path) + if cap is None: print(f"Error opening video file {video_path}") return output_start - frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + frame_count = getattr(cap, "frame_count", 0) + if not frame_count and hasattr(cap, "get"): + frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) try: scene_manager.detect_scenes(video=video)