From d112a58366d9d9a8830fa344efe6c50bf587cd00 Mon Sep 17 00:00:00 2001 From: Greg Hogan Date: Thu, 28 May 2020 17:33:22 -0700 Subject: [PATCH] FrameReader env vars FFMPEG_THREADS and FFMPEG_CUDA --- tools/lib/framereader.py | 44 +++++++++++++++++++--------------------- 1 file changed, 21 insertions(+), 23 deletions(-) diff --git a/tools/lib/framereader.py b/tools/lib/framereader.py index 9efe7961e..5a1403591 100644 --- a/tools/lib/framereader.py +++ b/tools/lib/framereader.py @@ -210,24 +210,24 @@ def rgb24toyuv420(rgb): return yuv420.clip(0,255).astype('uint8') -def decompress_video_data(rawdat, vid_fmt, w, h, pix_fmt, multithreaded=None): +def decompress_video_data(rawdat, vid_fmt, w, h, pix_fmt): # using a tempfile is much faster than proc.communicate for some reason - if multithreaded is None: - multithreaded = not os.getenv('SINGLETHREADED') - with tempfile.TemporaryFile() as tmpf: tmpf.write(rawdat) tmpf.seek(0) + threads = os.getenv("FFMPEG_THREADS", "0") + cuda = os.getenv("FFMPEG_CUDA", "0") == "1" proc = subprocess.Popen( ["ffmpeg", - "-threads", "0" if multithreaded else "1", + "-threads", threads, + "-c:v", "hevc" if not cuda else "hevc_cuvid", "-vsync", "0", "-f", vid_fmt, "-flags2", "showall", "-i", "pipe:0", - "-threads", "0" if multithreaded else "1", + "-threads", threads, "-f", "rawvideo", "-pix_fmt", pix_fmt, "pipe:1"], @@ -266,14 +266,14 @@ class BaseFrameReader: raise NotImplementedError -def FrameReader(fn, cache_prefix=None, readahead=False, readbehind=False, multithreaded=None, index_data=None): +def FrameReader(fn, cache_prefix=None, readahead=False, readbehind=False, index_data=None): frame_type = fingerprint_video(fn) if frame_type == FrameType.raw: return RawFrameReader(fn) elif frame_type in (FrameType.h265_stream,): if not index_data: index_data = get_video_index(fn, frame_type, cache_prefix) - return StreamFrameReader(fn, frame_type, index_data, readahead=readahead, readbehind=readbehind, multithreaded=multithreaded) + return StreamFrameReader(fn, frame_type, index_data, readahead=readahead, readbehind=readbehind) else: raise NotImplementedError(frame_type) @@ -325,15 +325,12 @@ class RawFrameReader(BaseFrameReader): class VideoStreamDecompressor: - def __init__(self, vid_fmt, w, h, pix_fmt, multithreaded=None): + def __init__(self, vid_fmt, w, h, pix_fmt): self.vid_fmt = vid_fmt self.w = w self.h = h self.pix_fmt = pix_fmt - if multithreaded is None: - multithreaded = not os.getenv('SINGLETHREADED') - if pix_fmt == "yuv420p": self.out_size = w*h*3//2 # yuv420p elif pix_fmt in ("rgb24", "yuv444p"): @@ -343,9 +340,12 @@ class VideoStreamDecompressor: self.out_q = queue.Queue() + threads = os.getenv("FFMPEG_THREADS", "0") + cuda = os.getenv("FFMPEG_CUDA", "0") == "1" self.proc = subprocess.Popen( ["ffmpeg", - "-threads", "0" if multithreaded else "1", + "-threads", threads, + "-c:v", "hevc" if not cuda else "hevc_cuvid", # "-avioflags", "direct", "-analyzeduration", "0", "-probesize", "32", @@ -354,7 +354,7 @@ class VideoStreamDecompressor: "-vsync", "0", "-f", vid_fmt, "-i", "pipe:0", - "-threads", "0" if multithreaded else "1", + "-threads", threads, "-f", "rawvideo", "-pix_fmt", pix_fmt, "pipe:1"], @@ -479,10 +479,9 @@ class StreamGOPReader(GOPReader): class GOPFrameReader(BaseFrameReader): #FrameReader with caching and readahead for formats that are group-of-picture based - def __init__(self, readahead=False, readbehind=False, multithreaded=None): + def __init__(self, readahead=False, readbehind=False): self.open_ = True - self.multithreaded = multithreaded self.readahead = readahead self.readbehind = readbehind self.frame_cache = LRU(64) @@ -542,8 +541,7 @@ class GOPFrameReader(BaseFrameReader): frame_b, num_frames, skip_frames, rawdat = self.get_gop(num) - ret = decompress_video_data(rawdat, self.vid_fmt, self.w, self.h, pix_fmt, - multithreaded=self.multithreaded) + ret = decompress_video_data(rawdat, self.vid_fmt, self.w, self.h, pix_fmt) ret = ret[skip_frames:] assert ret.shape[0] == num_frames @@ -573,17 +571,17 @@ class GOPFrameReader(BaseFrameReader): class StreamFrameReader(StreamGOPReader, GOPFrameReader): - def __init__(self, fn, frame_type, index_data, readahead=False, readbehind=False, multithreaded=None): + def __init__(self, fn, frame_type, index_data, readahead=False, readbehind=False): StreamGOPReader.__init__(self, fn, frame_type, index_data) - GOPFrameReader.__init__(self, readahead, readbehind, multithreaded) + GOPFrameReader.__init__(self, readahead, readbehind) -def GOPFrameIterator(gop_reader, pix_fmt, multithreaded=None): +def GOPFrameIterator(gop_reader, pix_fmt): # this is really ugly. ill think about how to refactor it when i can think good IN_FLIGHT_GOPS = 6 # should be enough that the stream decompressor starts returning data - with VideoStreamDecompressor(gop_reader.vid_fmt, gop_reader.w, gop_reader.h, pix_fmt, multithreaded) as dec: + with VideoStreamDecompressor(gop_reader.vid_fmt, gop_reader.w, gop_reader.h, pix_fmt) as dec: read_work = [] def readthing(): @@ -621,7 +619,7 @@ def GOPFrameIterator(gop_reader, pix_fmt, multithreaded=None): def FrameIterator(fn, pix_fmt, **kwargs): fr = FrameReader(fn, **kwargs) if isinstance(fr, GOPReader): - for v in GOPFrameIterator(fr, pix_fmt, kwargs.get("multithreaded", None)): + for v in GOPFrameIterator(fr, pix_fmt): yield v else: for i in range(fr.frame_count):