mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-07-15 14:22:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

FrameReader: add hwaccel arg and clear frames_cache (#36974)

Browse Source
This commit is contained in:
YassineYousfi
2026-01-02 11:29:45 -08:00
committed by GitHub
parent f62177a827
commit adbf68f771
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/lib/framereader.py
+10 -10
View File
@@ -42,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', loglevel="info") -> np.ndarray:
def decompress_video_data(rawdat, w, h, pix_fmt="rgb24", vid_fmt='hevc', hwaccel="auto", loglevel="info") -> np.ndarray:
threads = os.getenv("FFMPEG_THREADS", "0")
args = ["ffmpeg", "-v", loglevel,
"-threads", threads,
"-hwaccel", "auto",
"-hwaccel", hwaccel,
"-c:v", "hevc",
"-vsync", "0",
"-f", vid_fmt,
@@ -101,13 +101,13 @@ def get_video_index(fn):
class FfmpegDecoder:
def __init__(self, fn: str, index_data: dict|None = None,
pix_fmt: str = "rgb24", loglevel="quiet"):
pix_fmt: str = "rgb24", hwaccel="auto", 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
self.loglevel, self.hwaccel = loglevel, hwaccel
def _gop_bounds(self, frame_idx: int):
f_b = frame_idx
@@ -119,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, loglevel=self.loglevel)
yield from decompress_video_data(raw, self.w, self.h, pix_fmt=self.pix_fmt, hwaccel=self.hwaccel, loglevel=self.loglevel)
def get_gop_start(self, frame_idx: int):
return self.iframes[np.searchsorted(self.iframes, frame_idx, side="right") - 1]
@@ -134,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, loglevel=self.loglevel)):
for i, frm in enumerate(decompress_video_data(raw, self.w, self.h, self.pix_fmt, hwaccel=self.hwaccel, loglevel=self.loglevel)):
fidx = f_b + i
if fidx >= end_fidx:
return
@@ -143,15 +143,15 @@ class FfmpegDecoder:
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, loglevel="quiet") -> Iterator[np.ndarray]:
dec = FfmpegDecoder(fn, pix_fmt=pix_fmt, index_data=index_data, loglevel=loglevel)
start_fidx:int=0, end_fidx=None, frame_skip:int=1, hwaccel="auto", loglevel="quiet") -> Iterator[np.ndarray]:
dec = FfmpegDecoder(fn, pix_fmt=pix_fmt, index_data=index_data, hwaccel=hwaccel, 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", loglevel="quiet"):
self.decoder = FfmpegDecoder(fn, index_data, pix_fmt, loglevel=loglevel)
pix_fmt: str = "rgb24", hwaccel="auto", loglevel="quiet"):
self.decoder = FfmpegDecoder(fn, index_data=index_data, pix_fmt=pix_fmt, hwaccel=hwaccel, 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