|
| 1 | + |
| 2 | +class AviMediaReader: |
| 3 | + """Read video frames |
| 4 | +
|
| 5 | + :param *urls: URLs of the media files to read. |
| 6 | + :type *urls: tuple(str) |
| 7 | + :param streams: list of file + stream specifiers or filtergraph label to output, alias of `map` option, |
| 8 | + defaults to None, which outputs at most one video and one audio, selected by FFmpeg |
| 9 | + :type streams: seq(str), optional |
| 10 | + :param progress: progress callback function, defaults to None |
| 11 | + :type progress: callable object, optional |
| 12 | + :param show_log: True to show FFmpeg log messages on the console, |
| 13 | + defaults to None (no show/capture) |
| 14 | + Ignored if stream format must be retrieved automatically. |
| 15 | + :type show_log: bool, optional |
| 16 | + :param use_ya8: True if piped video streams uses `ya8` pix_fmt instead of `gray16le`, default to None |
| 17 | + :type use_ya8: bool, optional |
| 18 | + :param \\**options: FFmpeg options, append '_in[input_url_id]' for input option names for specific |
| 19 | + input url or '_in' to be applied to all inputs. The url-specific option gets the |
| 20 | + preference (see :doc:`options` for custom options) |
| 21 | + :type \\**options: dict, optional |
| 22 | +
|
| 23 | + :return: frame rate and video frame data (dims: time x rows x cols x pix_comps) |
| 24 | + :rtype: (`fractions.Fraction`, `numpy.ndarray`) |
| 25 | +
|
| 26 | + Note: Only pass in multiple urls to implement complex filtergraph. It's significantly faster to run |
| 27 | + `ffmpegio.video.read()` for each url. |
| 28 | +
|
| 29 | +
|
| 30 | + Unlike :py:mod:`video` and :py:mod:`image`, video pixel formats are not autodetected. If output |
| 31 | + 'pix_fmt' option is not explicitly set, 'rgb24' is used. |
| 32 | +
|
| 33 | + For audio streams, if 'sample_fmt' output option is not specified, 's16le'. |
| 34 | +
|
| 35 | +
|
| 36 | + streams = ['0:v:0','1:a:3'] # pick 1st file's 1st video stream and 2nd file's 4th audio stream |
| 37 | +
|
| 38 | + """ |
| 39 | + |
| 40 | + def __init__(self, *urls, streams=None, progress=None, show_log=None, blocksize=None, **options): |
| 41 | + |
| 42 | + self.dtype = None # :numpy.dtype: output data type |
| 43 | + self.shape = ( |
| 44 | + None # :tuple of ints: dimension of each video frame or audio sample |
| 45 | + ) |
| 46 | + self.itemsize = None #:int: number of bytes of each video frame or audio sample |
| 47 | + self.blocksize = None #:positive int: number of video frames or audio samples to read when used as an iterator |
| 48 | + |
| 49 | + # get url/file stream |
| 50 | + url, stdin, input = configure.check_url(url, False) |
| 51 | + |
| 52 | + input_options = utils.pop_extra_options(options, "_in") |
| 53 | + |
| 54 | + ffmpeg_args = configure.empty() |
| 55 | + configure.add_url(ffmpeg_args, "input", url, input_options) |
| 56 | + configure.add_url(ffmpeg_args, "output", "-", options) |
| 57 | + |
| 58 | + # abstract method to finalize the options => sets self.dtype and self.shape if known |
| 59 | + self._finalize(ffmpeg_args) |
| 60 | + |
| 61 | + # create logger without assigning the source stream |
| 62 | + self._logger = _LogerThread(None, show_log) |
| 63 | + |
| 64 | + # start FFmpeg |
| 65 | + self._proc = ffmpegprocess.Popen( |
| 66 | + ffmpeg_args, |
| 67 | + stdin=stdin, |
| 68 | + progress=progress, |
| 69 | + capture_log=True, |
| 70 | + close_stdin=True, |
| 71 | + close_stdout=False, |
| 72 | + close_stderr=False, |
| 73 | + ) |
| 74 | + |
| 75 | + # set the log source and start the logger |
| 76 | + self._logger.stderr = self._proc.stderr |
| 77 | + self._logger.start() |
| 78 | + |
| 79 | + # if byte data is given, feed it |
| 80 | + if input is not None: |
| 81 | + self._proc.stdin.write(input) |
| 82 | + |
| 83 | + # wait until output stream log is captured if output format is unknown |
| 84 | + try: |
| 85 | + if self.dtype is None or self.shape is None: |
| 86 | + info = self._logger.output_stream() |
| 87 | + self._finalize_array(info) |
| 88 | + else: |
| 89 | + self._logger.index("Output") |
| 90 | + except: |
| 91 | + if self._proc.poll() is None: |
| 92 | + raise self._logger.Exception |
| 93 | + else: |
| 94 | + raise ValueError("failed retrieve output data format") |
| 95 | + |
| 96 | + self.itemsize = utils.get_itemsize(self.shape, self.dtype) |
| 97 | + |
| 98 | + self.blocksize = blocksize or max(1024 ** 2 // self.itemsize, 1) |
| 99 | + |
| 100 | + def close(self): |
| 101 | + """Flush and close this stream. This method has no effect if the stream is already |
| 102 | + closed. Once the stream is closed, any read operation on the stream will raise |
| 103 | + a ValueError. |
| 104 | +
|
| 105 | + As a convenience, it is allowed to call this method more than once; only the first call, |
| 106 | + however, will have an effect. |
| 107 | +
|
| 108 | + """ |
| 109 | + self._proc.stdout.close() |
| 110 | + self._proc.stderr.close() |
| 111 | + try: |
| 112 | + self._proc.terminate() |
| 113 | + except: |
| 114 | + pass |
| 115 | + self._logger.join() |
| 116 | + |
| 117 | + @property |
| 118 | + def closed(self): |
| 119 | + """:bool: True if the stream is closed.""" |
| 120 | + return self._proc.poll() is not None |
| 121 | + |
| 122 | + @property |
| 123 | + def lasterror(self): |
| 124 | + """:FFmpegError: Last error FFmpeg posted""" |
| 125 | + if self._proc.poll(): |
| 126 | + return self._logger.Exception() |
| 127 | + else: |
| 128 | + return None |
| 129 | + |
| 130 | + def __enter__(self): |
| 131 | + return self |
| 132 | + |
| 133 | + def __exit__(self, exc_type, exc_value, traceback): |
| 134 | + self.close() |
| 135 | + |
| 136 | + def __iter__(self): |
| 137 | + return self |
| 138 | + |
| 139 | + def __next__(self): |
| 140 | + try: |
| 141 | + return self.read(self.blocksize) |
| 142 | + except: |
| 143 | + raise StopIteration |
| 144 | + |
| 145 | + def readlog(self, n=None): |
| 146 | + if n is not None: |
| 147 | + self._logger.index(n) |
| 148 | + with self._logger._newline_mutex: |
| 149 | + return "\n".join(self._logger.logs or self._logger.logs[:n]) |
| 150 | + |
| 151 | + def read(self, n=-1): |
| 152 | + """Read and return numpy.ndarray with up to n frames/samples. If |
| 153 | + the argument is omitted, None, or negative, data is read and |
| 154 | + returned until EOF is reached. An empty bytes object is returned |
| 155 | + if the stream is already at EOF. |
| 156 | +
|
| 157 | + If the argument is positive, and the underlying raw stream is not |
| 158 | + interactive, multiple raw reads may be issued to satisfy the byte |
| 159 | + count (unless EOF is reached first). But for interactive raw streams, |
| 160 | + at most one raw read will be issued, and a short result does not |
| 161 | + imply that EOF is imminent. |
| 162 | +
|
| 163 | + A BlockingIOError is raised if the underlying raw stream is in non |
| 164 | + blocking-mode, and has no data available at the moment.""" |
| 165 | + |
| 166 | + b = self._proc.stdout.read(n * self.itemsize if n > 0 else n) |
| 167 | + if not len(b): |
| 168 | + self._proc.stdout.close() |
| 169 | + return _as_array(b, self.shape, self.dtype) |
| 170 | + |
| 171 | + def readinto(self, array): |
| 172 | + """Read bytes into a pre-allocated, writable bytes-like object array and |
| 173 | + return the number of bytes read. For example, b might be a bytearray. |
| 174 | +
|
| 175 | + Like read(), multiple reads may be issued to the underlying raw stream, |
| 176 | + unless the latter is interactive. |
| 177 | +
|
| 178 | + A BlockingIOError is raised if the underlying raw stream is in non |
| 179 | + blocking-mode, and has no data available at the moment.""" |
| 180 | + |
| 181 | + return self._proc.stdout.readinto(memoryview(array).cast("b")) // self.itemsize |
| 182 | + |
| 183 | + |
| 184 | +class SimpleVideoReader(SimpleReaderBase): |
| 185 | + def _finalize(self, ffmpeg_args): |
| 186 | + # finalize FFmpeg arguments and output array |
| 187 | + |
| 188 | + inurl, inopts = ffmpeg_args.get("inputs", [])[0] |
| 189 | + outopts = ffmpeg_args.get("outputs", [])[0][1] |
| 190 | + has_fg = configure.has_filtergraph(ffmpeg_args, "video") |
| 191 | + |
| 192 | + pix_fmt = outopts.get("pix_fmt", None) |
| 193 | + if pix_fmt is None or ( |
| 194 | + not has_fg |
| 195 | + and inurl not in ("-", "pipe:", "pipe:0") |
| 196 | + and not inopts.get("pix_fmt", None) |
| 197 | + ): |
| 198 | + # must assign output rgb/grayscale pixel format |
| 199 | + info = probe.video_streams_basic(inurl, 0)[0] |
| 200 | + pix_fmt_in = info["pix_fmt"] |
| 201 | + s_in = (info["width"], info["height"]) |
| 202 | + r_in = info["frame_rate"] |
| 203 | + else: |
| 204 | + pix_fmt_in = s_in = r_in = None |
| 205 | + |
| 206 | + ( |
| 207 | + self.dtype, |
| 208 | + self.shape, |
| 209 | + self.frame_rate, |
| 210 | + ) = configure.finalize_video_read_opts(ffmpeg_args, pix_fmt_in, s_in, r_in) |
| 211 | + |
| 212 | + # construct basic video filter if options specified |
| 213 | + configure.build_basic_vf( |
| 214 | + ffmpeg_args, utils.alpha_change(pix_fmt_in, pix_fmt, -1) |
| 215 | + ) |
| 216 | + |
| 217 | + def _finalize_array(self, info): |
| 218 | + # finalize array setup from FFmpeg log |
| 219 | + |
| 220 | + self.frame_rate = info["r"] |
| 221 | + self.dtype, ncomp, _ = utils.get_video_format(info["pix_fmt"]) |
| 222 | + self.shape = (*info["s"][::-1], ncomp) |
0 commit comments