render_video/src/render/ffmpeg.rs

use super::{cmd, filter::Filter};
use crate::{
	render::filter::channel,
	time::{format_time, Time},
	Resolution
};
use anyhow::bail;
use camino::Utf8PathBuf as PathBuf;
use rational::Rational;
use std::{borrow::Cow, fmt::Write as _, process::Command};

pub(crate) struct FfmpegInput {
	pub(crate) concat: bool,
	pub(crate) loop_input: bool,
	pub(crate) fps: Option<Rational>,
	pub(crate) start: Option<Time>,
	pub(crate) duration: Option<Time>,
	pub(crate) path: PathBuf
}

impl FfmpegInput {
	pub(crate) fn new(path: PathBuf) -> Self {
		Self {
			concat: false,
			loop_input: false,
			fps: None,
			start: None,
			duration: None,
			path
		}
	}

	fn append_to_cmd(self, cmd: &mut Command) {
		if self.concat {
			cmd.arg("-f").arg("concat").arg("-safe").arg("0");
		}
		if self.loop_input {
			cmd.arg("-loop").arg("1");
		}
		if let Some(fps) = self.fps {
			cmd.arg("-r").arg(fps.to_string());
		}
		if let Some(start) = self.start {
			cmd.arg("-seek_streams_individually").arg("false");
			cmd.arg("-ss").arg(format_time(start));
		}
		if let Some(duration) = self.duration {
			cmd.arg("-t").arg(format_time(duration));
		}
		cmd.arg("-i").arg(self.path);
	}
}

pub(crate) struct FfmpegOutput {
	pub(crate) fps: Option<Rational>,
	pub(crate) duration: Option<Time>,
	pub(crate) time_base: Option<Rational>,
	pub(crate) fps_mode_vfr: bool,
	pub(crate) faststart: bool,
	pub(crate) path: PathBuf
}

impl FfmpegOutput {
	pub(crate) fn new(path: PathBuf) -> Self {
		Self {
			fps: None,
			duration: None,
			time_base: None,
			fps_mode_vfr: false,
			faststart: false,
			path
		}
	}

	pub(crate) fn enable_faststart(mut self) -> Self {
		self.faststart = true;
		self
	}

	fn append_to_cmd(self, cmd: &mut Command) {
		if let Some(fps) = self.fps {
			cmd.arg("-r").arg(fps.to_string());
		}
		if let Some(duration) = self.duration {
			cmd.arg("-t").arg(format_time(duration));
		}
		if let Some(time_base) = self.time_base {
			cmd.arg("-enc_time_base").arg(time_base.to_string());
		}
		if self.fps_mode_vfr {
			cmd.arg("-fps_mode").arg("vfr");
		}
		if self.faststart {
			cmd.arg("-movflags").arg("+faststart");
		}
		cmd.arg(self.path);
	}
}

enum FfmpegFilter {
	None,
	Filters {
		filters: Vec<Filter>,
		output: Cow<'static, str>
	},
	Loudnorm,
	Rescale(Resolution)
}

pub(crate) struct Ffmpeg {
	inputs: Vec<FfmpegInput>,
	filter: FfmpegFilter,
	video_bitrate: Option<u64>,
	output: FfmpegOutput,

	filter_idx: usize
}

impl Ffmpeg {
	pub fn new(output: FfmpegOutput) -> Self {
		Self {
			inputs: Vec::new(),
			filter: FfmpegFilter::None,
			video_bitrate: None,
			output,

			filter_idx: 0
		}
	}

	pub fn add_input(&mut self, input: FfmpegInput) -> String {
		self.inputs.push(input);
		(self.inputs.len() - 1).to_string()
	}

	pub fn add_filter(&mut self, filter: Filter) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => {
				self.filter = FfmpegFilter::Filters {
					filters: vec![filter],
					output: "0".into()
				}
			},
			FfmpegFilter::Filters { filters, .. } => filters.push(filter),
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn set_filter_output<T: Into<Cow<'static, str>>>(
		&mut self,
		filter_output: T
	) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => {
				self.filter = FfmpegFilter::Filters {
					filters: vec![],
					output: filter_output.into()
				}
			},
			FfmpegFilter::Filters { output, .. } => *output = filter_output.into(),
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn enable_loudnorm(&mut self) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => self.filter = FfmpegFilter::Loudnorm,
			FfmpegFilter::Loudnorm => {},
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn rescale_video(&mut self, res: Resolution) -> &mut Self {
		match &mut self.filter {
			FfmpegFilter::None => self.filter = FfmpegFilter::Rescale(res),
			FfmpegFilter::Loudnorm => {},
			_ => panic!("An incompatible type of filter has been set before")
		}
		self
	}

	pub fn set_video_bitrate(&mut self, bitrate: u64) -> &mut Self {
		self.video_bitrate = Some(bitrate);
		self
	}

	pub fn run(mut self) -> anyhow::Result<()> {
		let mut cmd = cmd();
		cmd.arg("ffmpeg").arg("-hide_banner").arg("-y");

		// determine whether the video need to be re-encoded
		// vdec is only true if the video should be decoded on hardware
		let (vdec, venc, aenc) = match &self.filter {
			FfmpegFilter::None => (false, false, false),
			FfmpegFilter::Filters { .. } => (false, true, true),
			FfmpegFilter::Loudnorm => (false, false, true),
			FfmpegFilter::Rescale(_) => (true, true, false)
		};

		// initialise a vaapi device if one exists
		let vaapi_device: PathBuf = "/dev/dri/renderD128".into();
		let vaapi = vaapi_device.exists();
		if vaapi && venc {
			if vdec {
				cmd.arg("-hwaccel").arg("vaapi");
				cmd.arg("-hwaccel_device").arg(vaapi_device);
				cmd.arg("-hwaccel_output_format").arg("vaapi");
			} else {
				cmd.arg("-vaapi_device").arg(&vaapi_device);
			}
		}

		// append all the inputs
		for i in self.inputs {
			i.append_to_cmd(&mut cmd);
		}

		// always try to synchronise audio
		cmd.arg("-async").arg("1");

		// apply filters
		match self.filter {
			FfmpegFilter::None => {},
			FfmpegFilter::Filters { filters, output } => {
				let mut complex = String::new();
				for filter in filters {
					filter.append_to_complex_filter(&mut complex, &mut self.filter_idx);
				}
				if vaapi {
					write!(complex, "{}format=nv12,hwupload[v]", channel('v', &output));
				} else {
					write!(complex, "{}null[v]", channel('v', &output));
				}
				cmd.arg("-filter_complex").arg(complex);
				cmd.arg("-map").arg("[v]");
				cmd.arg("-map").arg(channel('a', &output));
			},
			FfmpegFilter::Loudnorm => {
				cmd.arg("-af").arg(concat!(
					"pan=mono|c0=FR,",
					"loudnorm=dual_mono=true:print_format=summary,",
					"pan=stereo|c0=c0|c1=c0,",
					"aformat=sample_rates=48000"
				));
			},
			FfmpegFilter::Rescale(res) => {
				cmd.arg("-vf").arg(if vaapi {
					format!("scale_vaapi=w={}:h={}", res.width(), res.height())
				} else {
					format!("scale=w={}:h={}", res.width(), res.height())
				});
			}
		}

		// append encoding options
		const QUALITY: &str = "22";
		if venc {
			if vaapi {
				cmd.arg("-c:v").arg("h264_vaapi");
				if self.video_bitrate.is_none() {
					cmd.arg("-rc_mode").arg("CQP");
					cmd.arg("-global_quality").arg(QUALITY);
				}
			} else {
				cmd.arg("-c:v").arg("libx264");
				if self.video_bitrate.is_none() {
					cmd.arg("-crf").arg(QUALITY);
				}
			}
			if self.video_bitrate.is_some() {
				cmd.arg("-b:v").arg(self.video_bitrate.unwrap().to_string());
			}
		} else {
			cmd.arg("-c:v").arg("copy");
		}
		if aenc {
			cmd.arg("-c:a").arg("aac");
			cmd.arg("-b:a").arg("128000");
		} else {
			cmd.arg("-c:a").arg("copy");
		}

		self.output.append_to_cmd(&mut cmd);

		let status = cmd.status()?;
		if status.success() {
			Ok(())
		} else {
			bail!("ffmpeg failed with exit code {:?}", status.code())
		}
	}
}