FrameReader: use hwaccel auto (#36973)

* FrameReader: use hwaccel auto

* rm main block
This commit is contained in:
YassineYousfi
2026-01-01 19:06:31 -08:00
committed by GitHub
parent bb40d161e8
commit f62177a827
+13 -33
View File
@@ -2,7 +2,6 @@ import os
import subprocess
import json
import logging
from functools import cache
from collections.abc import Iterator
from collections import OrderedDict
@@ -17,19 +16,6 @@ HEVC_SLICE_B = 0
HEVC_SLICE_P = 1
HEVC_SLICE_I = 2
@cache
def get_hw_accel() -> list[str]:
"""Detect and return the best available ffmpeg hardware acceleration."""
priority = ("videotoolbox", "cuda", "vaapi", "d3d11va")
result = subprocess.run(["ffmpeg", "-hwaccels"], capture_output=True, text=True, timeout=5)
for accel in priority:
if accel in result.stdout.lower():
logger.info(f"HW accelerated video decode found, using ffmpeg's {accel}")
return ["-hwaccel", accel]
logger.warning("no HW accelerated video found with `ffmpeg -hwaccels`. falling back to ffmpeg CPU decode")
return []
class LRUCache:
def __init__(self, capacity: int):
self._cache: OrderedDict = OrderedDict()
@@ -47,7 +33,6 @@ class LRUCache:
def __contains__(self, key):
return key in self._cache
def assert_hvec(fn: str) -> None:
with FileReader(fn) as f:
header = f.read(4)
@@ -57,11 +42,11 @@ def assert_hvec(fn: str) -> None:
if 'hevc' not in fn:
raise NotImplementedError(fn)
def decompress_video_data(rawdat, w, h, pix_fmt="rgb24", vid_fmt='hevc') -> np.ndarray:
def decompress_video_data(rawdat, w, h, pix_fmt="rgb24", vid_fmt='hevc', loglevel="info") -> np.ndarray:
threads = os.getenv("FFMPEG_THREADS", "0")
args = ["ffmpeg", "-v", "quiet",
args = ["ffmpeg", "-v", loglevel,
"-threads", threads,
*get_hw_accel(),
"-hwaccel", "auto",
"-c:v", "hevc",
"-vsync", "0",
"-f", vid_fmt,
@@ -114,15 +99,15 @@ def get_video_index(fn):
'probe': probe
}
class FfmpegDecoder:
def __init__(self, fn: str, index_data: dict|None = None,
pix_fmt: str = "rgb24"):
pix_fmt: str = "rgb24", loglevel="quiet"):
self.fn = fn
self.index, self.prefix, self.w, self.h = get_index_data(fn, index_data)
self.frame_count = len(self.index) - 1 # sentinel row at the end
self.iframes = np.where(self.index[:, 0] == HEVC_SLICE_I)[0]
self.pix_fmt = pix_fmt
self.loglevel = loglevel
def _gop_bounds(self, frame_idx: int):
f_b = frame_idx
@@ -134,7 +119,7 @@ class FfmpegDecoder:
return f_b, f_e, self.index[f_b, 1], self.index[f_e, 1]
def _decode_gop(self, raw: bytes) -> Iterator[np.ndarray]:
yield from decompress_video_data(raw, self.w, self.h, self.pix_fmt)
yield from decompress_video_data(raw, self.w, self.h, self.pix_fmt, loglevel=self.loglevel)
def get_gop_start(self, frame_idx: int):
return self.iframes[np.searchsorted(self.iframes, frame_idx, side="right") - 1]
@@ -149,7 +134,7 @@ class FfmpegDecoder:
f.seek(off_b)
raw = self.prefix + f.read(off_e - off_b)
# number of frames to discard inside this GOP before the wanted one
for i, frm in enumerate(decompress_video_data(raw, self.w, self.h, self.pix_fmt)):
for i, frm in enumerate(decompress_video_data(raw, self.w, self.h, self.pix_fmt, loglevel=self.loglevel)):
fidx = f_b + i
if fidx >= end_fidx:
return
@@ -157,17 +142,16 @@ class FfmpegDecoder:
yield fidx, frm
fidx += 1
def FrameIterator(fn: str, index_data: dict|None=None,
pix_fmt: str = "rgb24",
start_fidx:int=0, end_fidx=None, frame_skip:int=1) -> Iterator[np.ndarray]:
dec = FfmpegDecoder(fn, pix_fmt=pix_fmt, index_data=index_data)
def FrameIterator(fn: str, index_data: dict|None=None, pix_fmt: str = "rgb24",
start_fidx:int=0, end_fidx=None, frame_skip:int=1, loglevel="quiet") -> Iterator[np.ndarray]:
dec = FfmpegDecoder(fn, pix_fmt=pix_fmt, index_data=index_data, loglevel=loglevel)
for _, frame in dec.get_iterator(start_fidx=start_fidx, end_fidx=end_fidx, frame_skip=frame_skip):
yield frame
class FrameReader:
def __init__(self, fn: str, index_data: dict|None = None,
cache_size: int = 30, pix_fmt: str = "rgb24"):
self.decoder = FfmpegDecoder(fn, index_data, pix_fmt)
def __init__(self, fn: str, index_data: dict|None = None, cache_size: int = 30,
pix_fmt: str = "rgb24", loglevel="quiet"):
self.decoder = FfmpegDecoder(fn, index_data, pix_fmt, loglevel=loglevel)
self.iframes = self.decoder.iframes
self._cache: LRUCache = LRUCache(cache_size)
self.w, self.h, self.frame_count, = self.decoder.w, self.decoder.h, self.decoder.frame_count
@@ -187,7 +171,3 @@ class FrameReader:
self.fidx, frame = next(self.it)
self._cache[self.fidx] = frame
return self._cache[fidx]
if __name__ == "__main__":
get_hw_accel()