#!/usr/bin/env python3 """ Video output writer for faceswap.py converter """ from __future__ import annotations import os import typing as T from math import ceil from subprocess import CalledProcessError, check_output, STDOUT import imageio import imageio_ffmpeg as im_ffm import numpy as np from ._base import Output, logger if T.TYPE_CHECKING: from collections.abc import Generator class Writer(Output): """ Video output writer using imageio-ffmpeg. Parameters ---------- output_folder: str The folder to save the output video to total_count: int The total number of frames to be converted frame_ranges: list or ``None`` List of tuples for starting and end values of each frame range to be converted or ``None`` if all frames are to be converted source_video: str The full path to the source video for obtaining fps and audio kwargs: dict Any additional standard :class:`plugins.convert.writer._base.Output` key word arguments. """ def __init__(self, output_folder: str, total_count: int, frame_ranges: list[tuple[int, int]] | None, source_video: str, **kwargs) -> None: super().__init__(output_folder, **kwargs) logger.debug("total_count: %s, frame_ranges: %s, source_video: '%s'", total_count, frame_ranges, source_video) self._source_video: str = source_video self._output_filename: str = self._get_output_filename() self._frame_ranges: list[tuple[int, int]] | None = frame_ranges self.frame_order: list[int] = self._set_frame_order(total_count) self._output_dimensions: str | None = None # Fix dims on 1st received frame # Need to know dimensions of first frame, so set writer then self._writer: Generator[None, np.ndarray, None] | None = None @property def _valid_tunes(self) -> dict: """ dict: Valid tune selections for libx264 and libx265 codecs. """ return {"libx264": ["film", "animation", "grain", "stillimage", "fastdecode", "zerolatency"], "libx265": ["grain", "fastdecode", "zerolatency"]} @property def _video_fps(self) -> float: """ float: The fps of the source video. """ reader = imageio.get_reader(self._source_video, "ffmpeg") # type:ignore[arg-type] retval = reader.get_meta_data()["fps"] reader.close() logger.debug(retval) return retval @property def _output_params(self) -> list[str]: """ list: The FFMPEG Output parameters """ codec = self.config["codec"] tune = self.config["tune"] # Force all frames to the same size output_args = ["-vf", f"scale={self._output_dimensions}"] output_args.extend(["-crf", str(self.config["crf"])]) output_args.extend(["-preset", self.config["preset"]]) if tune is not None and tune in self._valid_tunes[codec]: output_args.extend(["-tune", tune]) if codec == "libx264" and self.config["profile"] != "auto": output_args.extend(["-profile:v", self.config["profile"]]) if codec == "libx264" and self.config["level"] != "auto": output_args.extend(["-level", self.config["level"]]) logger.debug(output_args) return output_args @property def _audio_codec(self) -> str | None: """ str or ``None``: The audio codec to use. This will either be ``"copy"`` (the default) or ``None`` if skip muxing has been selected in configuration options, or if frame ranges have been passed in the command line arguments. """ retval: str | None = "copy" if self.config["skip_mux"]: logger.info("Skipping audio muxing due to configuration settings.") retval = None elif self._frame_ranges is not None: logger.warning("Muxing audio is not supported for limited frame ranges." "The output video will be created but you will need to mux audio " "manually.") retval = None elif not self._test_for_audio_stream(): logger.warning("No audio stream could be found in the source video '%s'. Muxing audio " "will be disabled.", self._source_video) retval = None logger.debug("Audio codec: %s", retval) return retval def _test_for_audio_stream(self) -> bool: """ Check whether the source video file contains an audio stream. If we attempt to mux audio from a source video that does not contain an audio stream ffmpeg will crash faceswap in a fairly ugly manner. Returns ------- bool ``True`` if an audio stream is found in the source video file, otherwise ``False`` Raises ------ ValueError If a subprocess error is raised scanning the input video file """ exe = im_ffm.get_ffmpeg_exe() cmd = [exe, "-hide_banner", "-i", self._source_video, "-f", "ffmetadata", "-"] try: out = check_output(cmd, stderr=STDOUT) except CalledProcessError as err: err_out = err.output.decode(errors="ignore") msg = f"Error checking audio stream. Status: {err.returncode}\n{err_out}" raise ValueError(msg) from err retval = False for line in out.splitlines(): if not line.strip().startswith(b"Stream #"): continue logger.debug("scanning Stream line: %s", line.decode(errors="ignore").strip()) if b"Audio" in line: retval = True break logger.debug("Audio found: %s", retval) return retval def _get_output_filename(self) -> str: """ Return full path to video output file. The filename is the same as the input video with `"_converted"` appended to the end. The file extension is as selected in the plugin settings. If a file already exists with the given filename, then `"_1"` is appended to the end of the filename. This number iterates until a valid filename that does not exist is found. Returns ------- str The full path to the output video filename """ filename = os.path.basename(self._source_video) filename = os.path.splitext(filename)[0] ext = self.config["container"] idx = 0 while True: out_file = f"{filename}_converted{'' if idx == 0 else f'_{idx}'}.{ext}" retval = os.path.join(self.output_folder, out_file) if not os.path.exists(retval): break idx += 1 logger.info("Outputting to: '%s'", retval) return retval def _set_frame_order(self, total_count: int) -> list[int]: """ Obtain the full list of frames to be converted in order. Parameters ---------- total_count: int The total number of frames to be converted Returns ------- list Full list of all frame indices to be converted """ if self._frame_ranges is None: retval = list(range(1, total_count + 1)) else: retval = [] for rng in self._frame_ranges: retval.extend(list(range(rng[0], rng[1] + 1))) logger.debug("frame_order: %s", retval) return retval def _get_writer(self, frame_dims: tuple[int, int]) -> Generator[None, np.ndarray, None]: """ Add the requested encoding options and return the writer. Parameters ---------- frame_dims: tuple The (rows, colums) shape of the input image Returns ------- generator The imageio ffmpeg writer """ audio_codec = self._audio_codec audio_path = None if audio_codec is None else self._source_video logger.debug("writer config: %s, audio_path: '%s'", self.config, audio_path) retval = im_ffm.write_frames(self._output_filename, size=(frame_dims[1], frame_dims[0]), fps=self._video_fps, quality=None, codec=self.config["codec"], macro_block_size=8, ffmpeg_log_level="error", ffmpeg_timeout=10, output_params=self._output_params, audio_path=audio_path, audio_codec=audio_codec) logger.debug("FFMPEG Writer created: %s", retval) retval.send(None) return retval def write(self, filename: str, image: np.ndarray) -> None: """ Frames come from the pool in arbitrary order, so frames are cached for writing out in the correct order. Parameters ---------- filename: str The incoming frame filename. image: :class:`numpy.ndarray` The converted image to be written """ logger.trace("Received frame: (filename: '%s', shape: %s", # type:ignore[attr-defined] filename, image.shape) if not self._output_dimensions: input_dims = T.cast(tuple[int, int], image.shape[:2]) self._set_dimensions(input_dims) self._writer = self._get_writer(input_dims) self.cache_frame(filename, image) self._save_from_cache() def _set_dimensions(self, frame_dims: tuple[int, int]) -> None: """ Set the attribute :attr:`_output_dimensions` based on the first frame received. This protects against different sized images coming in and ensures all images are written to ffmpeg at the same size. Dimensions are mapped to a macro block size 8. Parameters ---------- frame_dims: tuple The (rows, colums) shape of the input image """ logger.debug("input dimensions: %s", frame_dims) self._output_dimensions = (f"{int(ceil(frame_dims[1] / 8) * 8)}:" f"{int(ceil(frame_dims[0] / 8) * 8)}") logger.debug("Set dimensions: %s", self._output_dimensions) def _save_from_cache(self) -> None: """ Writes any consecutive frames to the video container that are ready to be output from the cache. """ assert self._writer is not None while self.frame_order: if self.frame_order[0] not in self.cache: logger.trace("Next frame not ready. Continuing") # type:ignore[attr-defined] break save_no = self.frame_order.pop(0) save_image = self.cache.pop(save_no) logger.trace("Rendering from cache. Frame no: %s", # type:ignore[attr-defined] save_no) self._writer.send(np.ascontiguousarray(save_image[:, :, ::-1])) logger.trace("Current cache size: %s", len(self.cache)) # type:ignore[attr-defined] def close(self) -> None: """ Close the ffmpeg writer and mux the audio """ if self._writer is not None: self._writer.close()