import sys import time import pathlib import subprocess from functools import lru_cache from collections import defaultdict from ._utils import get_ffmpeg_exe, _popen_kwargs, logger from ._parsing import LogCatcher, parse_ffmpeg_header, cvsecs ISWIN = sys.platform.startswith("win") h264_encoder_preference = defaultdict(lambda: -1) # The libx264 was the default encoder for a longe time with imageio h264_encoder_preference["libx264"] = 100 # Encoder with the nvidia graphics card dedicated hardware h264_encoder_preference["h264_nvenc"] = 90 # Deprecated names for the same encoder h264_encoder_preference["nvenc_h264"] = 90 h264_encoder_preference["nvenc"] = 90 # vaapi provides hardware encoding with intel integrated graphics chipsets h264_encoder_preference["h264_vaapi"] = 80 # openh264 is cisco's open source encoder h264_encoder_preference["libopenh264"] = 70 h264_encoder_preference["libx264rgb"] = 50 def ffmpeg_test_encoder(encoder): # Use the null streams to validate if we can encode anything # https://trac.ffmpeg.org/wiki/Null cmd = [ get_ffmpeg_exe(), "-hide_banner", "-f", "lavfi", "-i", "nullsrc=s=256x256:d=8", "-vcodec", encoder, "-f", "null", "-", ] p = subprocess.run( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) return p.returncode == 0 def get_compiled_h264_encoders(): cmd = [get_ffmpeg_exe(), "-hide_banner", "-encoders"] p = subprocess.run( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) stdout = p.stdout.decode().replace("\r", "") # 2022/04/08: hmaarrfk # I couldn't find a good way to get the list of available encoders from # the ffmpeg command # The ffmpeg command return a table that looks like # Notice the leading space at the very beginning # On ubuntu with libffmpeg-nvenc-dev we get # $ ffmpeg -hide_banner -encoders | grep -i h.264 # # Encoders: # V..... = Video # A..... = Audio # S..... = Subtitle # .F.... = Frame-level multithreading # ..S... = Slice-level multithreading # ...X.. = Codec is experimental # ....B. = Supports draw_horiz_band # .....D = Supports direct rendering method 1 # ------ # V..... libx264 libx264 H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10 (codec h264) # V..... libx264rgb libx264 H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10 RGB (codec h264) # V....D h264_nvenc NVIDIA NVENC H.264 encoder (codec h264) # V..... h264_omx OpenMAX IL H.264 video encoder (codec h264) # V..... h264_qsv H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10 (Intel Quick Sync Video acceleration) (codec h264) # V..... h264_v4l2m2m V4L2 mem2mem H.264 encoder wrapper (codec h264) # V....D h264_vaapi H.264/AVC (VAAPI) (codec h264) # V..... nvenc NVIDIA NVENC H.264 encoder (codec h264) # V..... nvenc_h264 NVIDIA NVENC H.264 encoder (codec h264) # # However, just because ffmpeg was compiled with the options enabled # it doesn't mean that it will be successful header_footer = stdout.split("------") footer = header_footer[1].strip("\n") encoders = [] for line in footer.split("\n"): # Strip to remove any leading spaces line = line.strip() encoder = line.split(" ")[1] if encoder in h264_encoder_preference: # These encoders are known to support H.264 # We forcibly include them in case their description changes to # not include the string "H.264" encoders.append(encoder) elif (line[0] == "V") and ("H.264" in line): encoders.append(encoder) encoders.sort(reverse=True, key=lambda x: h264_encoder_preference[x]) if "h264_nvenc" in encoders: # Remove deprecated names for the same encoder for encoder in ["nvenc", "nvenc_h264"]: if encoder in encoders: encoders.remove(encoder) # Return an immutable tuple to avoid users corrupting the lru_cache return tuple(encoders) @lru_cache() def get_first_available_h264_encoder(): compiled_encoders = get_compiled_h264_encoders() for encoder in compiled_encoders: if ffmpeg_test_encoder(encoder): return encoder else: raise RuntimeError( "No valid H.264 encoder was found with the ffmpeg installation" ) def count_frames_and_secs(path): """ Get the number of frames and number of seconds for the given video file. Note that this operation can be quite slow for large files. Disclaimer: I've seen this produce different results from actually reading the frames with older versions of ffmpeg (2.x). Therefore I cannot say with 100% certainty that the returned values are always exact. """ # https://stackoverflow.com/questions/2017843/fetch-frame-count-with-ffmpeg if isinstance(path, pathlib.PurePath): path = str(path) if not isinstance(path, str): raise TypeError("Video path must be a string or pathlib.Path.") cmd = [ get_ffmpeg_exe(), "-i", path, "-map", "0:v:0", "-c", "copy", "-f", "null", "-", ] try: out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, **_popen_kwargs()) except subprocess.CalledProcessError as err: out = err.output.decode(errors="ignore") raise RuntimeError("FFMEG call failed with {}:\n{}".format(err.returncode, out)) # Note that other than with the subprocess calls below, ffmpeg wont hang here. # Worst case Python will stop/crash and ffmpeg will continue running until done. nframes = nsecs = None for line in reversed(out.splitlines()): if line.startswith(b"frame="): line = line.decode(errors="ignore") i = line.find("frame=") if i >= 0: s = line[i:].split("=", 1)[-1].lstrip().split(" ", 1)[0].strip() nframes = int(s) i = line.find("time=") if i >= 0: s = line[i:].split("=", 1)[-1].lstrip().split(" ", 1)[0].strip() nsecs = cvsecs(*s.split(":")) return nframes, nsecs raise RuntimeError("Could not get number of frames") # pragma: no cover def read_frames( path, pix_fmt="rgb24", bpp=None, input_params=None, output_params=None, bits_per_pixel=None, ): """ Create a generator to iterate over the frames in a video file. It first yields a small metadata dictionary that contains: * ffmpeg_version: the ffmpeg version in use (as a string). * codec: a hint about the codec used to encode the video, e.g. "h264". * source_size: the width and height of the encoded video frames. * size: the width and height of the frames that will be produced. * fps: the frames per second. Can be zero if it could not be detected. * duration: duration in seconds. Can be zero if it could not be detected. After that, it yields frames until the end of the video is reached. Each frame is a bytes object. This function makes no assumptions about the number of frames in the data. For one because this is hard to predict exactly, but also because it may depend on the provided output_params. If you want to know the number of frames in a video file, use count_frames_and_secs(). It is also possible to estimate the number of frames from the fps and duration, but note that even if both numbers are present, the resulting value is not always correct. Example: gen = read_frames(path) meta = gen.__next__() for frame in gen: print(len(frame)) Parameters: path (str): the filename of the file to read from. pix_fmt (str): the pixel format of the frames to be read. The default is "rgb24" (frames are uint8 RGB images). input_params (list): Additional ffmpeg input command line parameters. output_params (list): Additional ffmpeg output command line parameters. bits_per_pixel (int): The number of bits per pixel in the output frames. This depends on the given pix_fmt. Default is 24 (RGB) bpp (int): DEPRECATED, USE bits_per_pixel INSTEAD. The number of bytes per pixel in the output frames. This depends on the given pix_fmt. Some pixel formats like yuv420p have 12 bits per pixel and cannot be set in bytes as integer. For this reason the bpp argument is deprecated. """ # ----- Input args if isinstance(path, pathlib.PurePath): path = str(path) if not isinstance(path, str): raise TypeError("Video path must be a string or pathlib.Path.") # Note: Dont check whether it exists. The source could be e.g. a camera. pix_fmt = pix_fmt or "rgb24" bpp = bpp or 3 bits_per_pixel = bits_per_pixel or bpp * 8 input_params = input_params or [] output_params = output_params or [] assert isinstance(pix_fmt, str), "pix_fmt must be a string" assert isinstance(bits_per_pixel, int), "bpp and bits_per_pixel must be an int" assert isinstance(input_params, list), "input_params must be a list" assert isinstance(output_params, list), "output_params must be a list" # ----- Prepare pre_output_params = ["-pix_fmt", pix_fmt, "-vcodec", "rawvideo", "-f", "image2pipe"] cmd = [get_ffmpeg_exe()] cmd += input_params + ["-i", path] cmd += pre_output_params + output_params + ["-"] p = subprocess.Popen( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **_popen_kwargs(prevent_sigint=True) ) log_catcher = LogCatcher(p.stderr) # Init policy by which to terminate ffmpeg. May be set to "kill" later. stop_policy = "timeout" # not wait; ffmpeg should be able to quit quickly # Enter try block directly after opening the process. # We terminate ffmpeg in the finally clause. # Generators are automatically closed when they get deleted, # so the finally block is guaranteed to run. try: # ----- Load meta data # Wait for the log catcher to get the meta information etime = time.time() + 10.0 while log_catcher.is_alive() and not log_catcher.header and time.time() < etime: time.sleep(0.01) # Check whether we have the information if not log_catcher.header: err2 = log_catcher.get_text(0.2) fmt = "Could not load meta information\n=== stderr ===\n{}" raise IOError(fmt.format(err2)) elif "No such file or directory" in log_catcher.header: raise IOError("{} not found! Wrong path?".format(path)) meta = parse_ffmpeg_header(log_catcher.header) yield meta # ----- Read frames w, h = meta["size"] framesize_bits = w * h * bits_per_pixel framesize_bytes = framesize_bits / 8 assert ( framesize_bytes.is_integer() ), "incorrect bits_per_pixel, framesize in bytes must be an int" framesize_bytes = int(framesize_bytes) framenr = 0 while True: framenr += 1 try: bb = bytes() while len(bb) < framesize_bytes: extra_bytes = p.stdout.read(framesize_bytes - len(bb)) if not extra_bytes: if len(bb) == 0: return else: raise RuntimeError( "End of file reached before full frame could be read." ) bb += extra_bytes yield bb except Exception as err: err1 = str(err) err2 = log_catcher.get_text(0.4) fmt = "Could not read frame {}:\n{}\n=== stderr ===\n{}" raise RuntimeError(fmt.format(framenr, err1, err2)) except GeneratorExit: # Note that GeneratorExit does not inherit from Exception but BaseException pass except Exception: # Normal exceptions fall through raise except BaseException: # Detect KeyboardInterrupt / SystemExit: don't wait for ffmpeg to quit stop_policy = "kill" raise finally: # Stop the LogCatcher thread, which reads from stderr. log_catcher.stop_me() # Make sure that ffmpeg is terminated. if p.poll() is None: # Ask ffmpeg to quit try: # I read somewhere that modern ffmpeg on Linux prefers a # "ctrl-c", but tests so far suggests sending q is more robust. # > p.send_signal(signal.SIGINT) # Sending q via communicate works, but can hang (see #17) # > p.communicate(b"q") # So let's do similar to what communicate does, but without # reading stdout (which may block). It looks like only closing # stdout is enough (tried Windows+Linux), but let's play safe. # Found that writing to stdin can cause "Invalid argument" on # Windows # and "Broken Pipe" on Unix. # p.stdin.write(b"q") # commented out in v0.4.1 p.stdout.close() p.stdin.close() # p.stderr.close() -> not here, the log_catcher closes it except Exception as err: # pragma: no cover logger.warning("Error while attempting stop ffmpeg (r): " + str(err)) if stop_policy == "timeout": # Wait until timeout, produce a warning and kill if it still exists try: etime = time.time() + 1.5 while time.time() < etime and p.poll() is None: time.sleep(0.01) finally: if p.poll() is None: # pragma: no cover logger.warning("We had to kill ffmpeg to stop it.") p.kill() else: # stop_policy == "kill" # Just kill it p.kill() def write_frames( path, size, pix_fmt_in="rgb24", pix_fmt_out="yuv420p", fps=16, quality=5, bitrate=None, codec=None, macro_block_size=16, ffmpeg_log_level="warning", ffmpeg_timeout=None, input_params=None, output_params=None, audio_path=None, audio_codec=None, ): """ Create a generator to write frames (bytes objects) into a video file. The frames are written by using the generator's `send()` method. Frames can be anything that can be written to a file. Typically these are bytes objects, but c-contiguous Numpy arrays also work. Example: gen = write_frames(path, size) gen.send(None) # seed the generator for frame in frames: gen.send(frame) gen.close() # don't forget this Parameters: path (str): the filename to write to. size (tuple): the width and height of the frames. pix_fmt_in (str): the pixel format of incoming frames. E.g. "gray", "gray8a", "rgb24", or "rgba". Default "rgb24". pix_fmt_out (str): the pixel format to store frames. Default yuv420p". fps (float): The frames per second. Default 16. quality (float): A measure for quality between 0 and 10. Default 5. Ignored if bitrate is given. bitrate (str): The bitrate, e.g. "192k". The defaults are pretty good. codec (str): The codec. Default "libx264" for .mp4 (if available from the ffmpeg executable) or "msmpeg4" for .wmv. macro_block_size (int): You probably want to align the size of frames to this value to avoid image resizing. Default 16. Can be set to 1 to avoid block alignment, though this is not recommended. ffmpeg_log_level (str): The ffmpeg logging level. Default "warning". ffmpeg_timeout (float): Timeout in seconds to wait for ffmpeg process to finish. Value of 0 or None will wait forever (default). The time that ffmpeg needs depends on CPU speed, compression, and frame size. input_params (list): Additional ffmpeg input command line parameters. output_params (list): Additional ffmpeg output command line parameters. audio_path (str): A input file path for encoding with an audio stream. Default None, no audio. audio_codec (str): The audio codec to use if audio_path is provided. "copy" will try to use audio_path's audio codec without re-encoding. Default None, but some formats must have certain codecs specified. """ # ----- Input args if isinstance(path, pathlib.PurePath): path = str(path) if not isinstance(path, str): raise TypeError("Video path must be a string or pathlib.Path.") # The pix_fmt_out yuv420p is the best for the outpur to work in # QuickTime and most other players. These players only support # the YUV planar color space with 4:2:0 chroma subsampling for # H.264 video. Otherwise, depending on the source, ffmpeg may # output to a pixel format that may be incompatible with these # players. See https://trac.ffmpeg.org/wiki/Encode/H.264#Encodingfordumbplayers pix_fmt_in = pix_fmt_in or "rgb24" pix_fmt_out = pix_fmt_out or "yuv420p" fps = fps or 16 # bitrate, codec, macro_block_size can all be None or ... macro_block_size = macro_block_size or 16 ffmpeg_log_level = ffmpeg_log_level or "warning" input_params = input_params or [] output_params = output_params or [] ffmpeg_timeout = ffmpeg_timeout or 0 floatish = float, int if isinstance(size, (tuple, list)): assert len(size) == 2, "size must be a 2-tuple" assert isinstance(size[0], int) and isinstance( size[1], int ), "size must be ints" sizestr = "{:d}x{:d}".format(*size) # elif isinstance(size, str): # assert "x" in size, "size as string must have format NxM" # sizestr = size else: assert False, "size must be str or tuple" assert isinstance(pix_fmt_in, str), "pix_fmt_in must be str" assert isinstance(pix_fmt_out, str), "pix_fmt_out must be str" assert isinstance(fps, floatish), "fps must be float" if quality is not None: assert isinstance(quality, floatish), "quality must be float" assert 1 <= quality <= 10, "quality must be between 1 and 10 inclusive" assert isinstance(macro_block_size, int), "macro_block_size must be int" assert isinstance(ffmpeg_log_level, str), "ffmpeg_log_level must be str" assert isinstance(ffmpeg_timeout, floatish), "ffmpeg_timeout must be float" assert isinstance(input_params, list), "input_params must be a list" assert isinstance(output_params, list), "output_params must be a list" # ----- Prepare # Get parameters if not codec: if path.lower().endswith(".wmv"): # This is a safer default codec on windows to get videos that # will play in powerpoint and other apps. H264 is not always # available on windows. codec = "msmpeg4" else: codec = get_first_available_h264_encoder() audio_params = ["-an"] if audio_path is not None and not path.lower().endswith(".gif"): audio_params = ["-i", audio_path] if audio_codec is not None: output_params += ["-acodec", audio_codec] output_params += ["-map", "0:v:0", "-map", "1:a:0"] # Get command cmd = [ get_ffmpeg_exe(), "-y", "-f", "rawvideo", "-vcodec", "rawvideo", "-s", sizestr, ] cmd += ["-pix_fmt", pix_fmt_in, "-r", "{:.02f}".format(fps)] + input_params cmd += ["-i", "-"] + audio_params cmd += ["-vcodec", codec, "-pix_fmt", pix_fmt_out] # Add fixed bitrate or variable bitrate compression flags if bitrate is not None: cmd += ["-b:v", str(bitrate)] elif quality is not None: # If None, then we don't add anything quality = 1 - quality / 10.0 if codec == "libx264": # crf ranges 0 to 51, 51 being worst. quality = int(quality * 51) cmd += ["-crf", str(quality)] # for h264 else: # Many codecs accept q:v # q:v range can vary, 1-31, 31 being worst # But q:v does not always have the same range. # May need a way to find range for any codec. quality = int(quality * 30) + 1 cmd += ["-qscale:v", str(quality)] # for others # Note, for most codecs, the image dimensions must be divisible by # 16 the default for the macro_block_size is 16. Check if image is # divisible, if not have ffmpeg upsize to nearest size and warn # user they should correct input image if this is not desired. if macro_block_size > 1: if size[0] % macro_block_size > 0 or size[1] % macro_block_size > 0: out_w = size[0] out_h = size[1] if size[0] % macro_block_size > 0: out_w += macro_block_size - (size[0] % macro_block_size) if size[1] % macro_block_size > 0: out_h += macro_block_size - (size[1] % macro_block_size) cmd += ["-vf", "scale={}:{}".format(out_w, out_h)] logger.warning( "IMAGEIO FFMPEG_WRITER WARNING: input image is not" " divisible by macro_block_size={}, resizing from {} " "to {} to ensure video compatibility with most codecs " "and players. To prevent resizing, make your input " "image divisible by the macro_block_size or set the " "macro_block_size to 1 (risking incompatibility).".format( macro_block_size, size[:2], (out_w, out_h) ) ) # Rather than redirect stderr to a pipe, just set minimal # output from ffmpeg by default. That way if there are warnings # the user will see them. cmd += ["-v", ffmpeg_log_level] cmd += output_params cmd.append(path) cmd_str = " ".join(cmd) if any( [level in ffmpeg_log_level for level in ("info", "verbose", "debug", "trace")] ): logger.info("RUNNING FFMPEG COMMAND: " + cmd_str) # Launch process p = subprocess.Popen( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=None, **_popen_kwargs(prevent_sigint=True) ) # Note that directing stderr to a pipe on windows will cause ffmpeg # to hang if the buffer is not periodically cleared using # StreamCatcher or other means. # Setting bufsize to 0 or a small value does not seem to have much effect # (tried on Windows and Linux). I suspect that ffmpeg buffers # multiple frames (before encoding in a batch). # Init policy by which to terminate ffmpeg. May be set to "kill" later. stop_policy = "timeout" if not ffmpeg_timeout: stop_policy = "wait" # ----- Write frames # Enter try block directly after opening the process. # We terminate ffmpeg in the finally clause. # Generators are automatically closed when they get deleted, # so the finally block is guaranteed to run. try: # Just keep going until the generator.close() is called (raises GeneratorExit). # This could also happen when the generator is deleted somehow. nframes = 0 while True: # Get frame bb = yield # framesize = size[0] * size[1] * depth * bpp # assert isinstance(bb, bytes), "Frame must be send as bytes" # assert len(bb) == framesize, "Frame must have width*height*depth*bpp bytes" # Actually, we accept anything that can be written to file. # This e.g. allows writing numpy arrays without having to make a copy ... # Write try: p.stdin.write(bb) except Exception as err: # Show the command and stderr from pipe msg = ( "{0:}\n\nFFMPEG COMMAND:\n{1:}\n\nFFMPEG STDERR " "OUTPUT:\n".format(err, cmd_str) ) raise IOError(msg) nframes += 1 except GeneratorExit: # Note that GeneratorExit does not inherit from Exception but BaseException # Detect premature closing if nframes == 0: logger.warning("No frames have been written; the written video is invalid.") except Exception: # Normal exceptions fall through raise except BaseException: # Detect KeyboardInterrupt / SystemExit: don't wait for ffmpeg to quit stop_policy = "kill" raise finally: # Make sure that ffmpeg is terminated. if p.poll() is None: # Tell ffmpeg that we're done try: p.stdin.close() except Exception as err: # pragma: no cover logger.warning("Error while attempting stop ffmpeg (w): " + str(err)) if stop_policy == "timeout": # Wait until timeout, produce a warning and kill if it still exists try: etime = time.time() + ffmpeg_timeout while (time.time() < etime) and p.poll() is None: time.sleep(0.01) finally: if p.poll() is None: # pragma: no cover logger.warning( "We had to kill ffmpeg to stop it. " + "Consider increasing ffmpeg_timeout, " + "or setting it to zero (no timeout)." ) p.kill() elif stop_policy == "wait": # Wait forever, kill if it if we're interrupted try: while p.poll() is None: time.sleep(0.01) finally: # the above can raise e.g. by ctrl-c or systemexit if p.poll() is None: # pragma: no cover p.kill() else: # stop_policy == "kill": # Just kill it p.kill() # Just to be safe, wrap in try/except try: p.stdout.close() except Exception: pass