av1an_core/
ffmpeg.rs

1use std::{
2    ffi::OsStr,
3    path::{Path, PathBuf},
4    process::{Command, Stdio},
5    str::FromStr,
6};
7
8use anyhow::bail;
9use av_format::rational::Rational64;
10use path_abs::{PathAbs, PathInfo};
11use serde::{Deserialize, Serialize};
12use tracing::warn;
13use vapoursynth::format::PresetFormat;
14
15use crate::{into_array, into_vec, ClipInfo, InputPixelFormat};
16
17#[inline]
18pub fn compose_ffmpeg_pipe<S: Into<String>>(
19    params: impl IntoIterator<Item = S>,
20    pix_format: FFPixelFormat,
21) -> Vec<String> {
22    let mut p: Vec<String> =
23        into_vec!["ffmpeg", "-y", "-hide_banner", "-loglevel", "error", "-i", "-",];
24
25    p.extend(params.into_iter().map(Into::into));
26
27    p.extend(into_array![
28        "-pix_fmt",
29        pix_format.to_pix_fmt_string(),
30        "-strict",
31        "-1",
32        "-f",
33        "yuv4mpegpipe",
34        "-"
35    ]);
36
37    p
38}
39
40#[derive(Debug, Clone, Deserialize)]
41struct FfProbeInfo {
42    pub streams: Vec<FfProbeStreamInfo>,
43}
44
45#[derive(Debug, Clone, Deserialize)]
46struct FfProbeStreamInfo {
47    pub width:          u32,
48    pub height:         u32,
49    pub pix_fmt:        String,
50    pub color_transfer: Option<String>,
51    pub avg_frame_rate: String,
52    pub nb_frames:      Option<String>,
53}
54
55#[inline]
56pub fn get_clip_info(source: &Path) -> anyhow::Result<ClipInfo> {
57    let output = Command::new("ffprobe")
58        .arg("-v")
59        .arg("quiet")
60        .arg("-select_streams")
61        .arg("v:0")
62        .arg("-print_format")
63        .arg("json")
64        .arg("-show_entries")
65        .arg("stream=width,height,pix_fmt,avg_frame_rate,nb_frames,color_transfer")
66        .arg(source)
67        .output()?
68        .stdout;
69    let ffprobe_info: FfProbeInfo = serde_json::from_slice(&output)?;
70    let stream_info = ffprobe_info
71        .streams
72        .first()
73        .ok_or_else(|| anyhow::anyhow!("no video streams found in source file"))?;
74
75    Ok(ClipInfo {
76        format_info:              InputPixelFormat::FFmpeg {
77            format: FFPixelFormat::from_str(&stream_info.pix_fmt)?,
78        },
79        frame_rate:               parse_frame_rate(&stream_info.avg_frame_rate)?,
80        resolution:               (stream_info.width, stream_info.height),
81        transfer_characteristics: match stream_info.color_transfer.as_deref() {
82            Some("smpte2084") => av1_grain::TransferFunction::SMPTE2084,
83            _ => av1_grain::TransferFunction::BT1886,
84        },
85        num_frames:               match stream_info.nb_frames.as_deref().map(str::parse) {
86            Some(Ok(nb_frames)) => nb_frames,
87            _ => get_num_frames(source)?,
88        },
89    })
90}
91
92/// Get frame count using FFmpeg
93#[inline]
94pub fn get_num_frames(source: &Path) -> anyhow::Result<usize> {
95    let output = Command::new("ffprobe")
96        .arg("-v")
97        .arg("error")
98        .arg("-select_streams")
99        .arg("v:0")
100        .arg("-count_packets")
101        .arg("-show_entries")
102        .arg("stream=nb_read_packets")
103        .arg("-print_format")
104        .arg("csv=p=0")
105        .arg(source)
106        .output()?
107        .stdout;
108    match String::from_utf8_lossy(&output).trim().parse::<usize>() {
109        Ok(x) if x > 0 => Ok(x),
110        _ => {
111            // If we got empty output or a 0 frame count, try using the slower
112            // but more reliable method
113            get_num_frames_slow(source)
114        },
115    }
116}
117
118/// Slower but more reliable frame count method
119fn get_num_frames_slow(source: &Path) -> anyhow::Result<usize> {
120    let output = Command::new("ffprobe")
121        .arg("-v")
122        .arg("error")
123        .arg("-count_frames")
124        .arg("-select_streams")
125        .arg("v:0")
126        .arg("-show_entries")
127        .arg("stream=nb_read_frames")
128        .arg("-print_format")
129        .arg("default=noprint_wrappers=1:nokey=1")
130        .arg(source)
131        .output()?
132        .stdout;
133    Ok(String::from_utf8_lossy(&output).trim().parse::<usize>()?)
134}
135
136fn parse_frame_rate(rate: &str) -> anyhow::Result<Rational64> {
137    let (numer, denom) = rate
138        .split_once('/')
139        .ok_or_else(|| anyhow::anyhow!("failed to parse frame rate from ffprobe output"))?;
140    Ok(Rational64::new(
141        numer.parse::<i64>()?,
142        denom.parse::<i64>()?,
143    ))
144}
145
146#[derive(Debug, Clone, Deserialize)]
147struct FfProbeKeyframesData {
148    pub frames: Vec<FfProbeKeyframeFrame>,
149}
150
151#[derive(Debug, Clone, Deserialize)]
152struct FfProbeKeyframeFrame {
153    // 0 or 1
154    pub key_frame: u8,
155}
156
157/// Returns vec of all keyframes
158#[tracing::instrument(level = "debug")]
159#[inline]
160pub fn get_keyframes(source: &Path) -> anyhow::Result<Vec<usize>> {
161    // This is slow because it has to iterate through the whole video,
162    // but it is the best suggestion that reliably worked
163    // since not all codecs code "coded_picture_number" into the frames.
164    let output = Command::new("ffprobe")
165        .arg("-v")
166        .arg("quiet")
167        .arg("-print_format")
168        .arg("json")
169        .arg("-show_frames")
170        .arg("-select_streams")
171        .arg("v:0")
172        .arg("-show_entries")
173        .arg("frame=key_frame")
174        .arg(source)
175        .output()?
176        .stdout;
177    let frames = serde_json::from_slice::<FfProbeKeyframesData>(&output)?.frames;
178    Ok(frames
179        .into_iter()
180        .enumerate()
181        .filter_map(|(i, frame)| (frame.key_frame > 0).then_some(i))
182        .collect())
183}
184
185/// Returns true if input file have audio in it
186#[inline]
187pub fn has_audio(file: &Path) -> anyhow::Result<bool> {
188    let output = Command::new("ffprobe")
189        .arg("-v")
190        .arg("error")
191        .arg("-select_streams")
192        .arg("a")
193        .arg("-show_entries")
194        .arg("stream=index")
195        .arg("-of")
196        .arg("csv=p=0")
197        .arg(file)
198        .output()?
199        .stdout;
200    let output = String::from_utf8_lossy(&output);
201    Ok(!output.trim().is_empty())
202}
203
204/// Encodes the audio using FFmpeg, blocking the current thread.
205///
206/// This function returns `Some(output)` if the audio exists and the audio
207/// successfully encoded, or `None` otherwise.
208#[inline]
209pub fn encode_audio<S: AsRef<OsStr>>(
210    input: impl AsRef<Path> + std::fmt::Debug,
211    temp: impl AsRef<Path> + std::fmt::Debug,
212    audio_params: &[S],
213) -> anyhow::Result<Option<PathBuf>> {
214    let input = input.as_ref();
215    let temp = temp.as_ref();
216
217    if has_audio(input)? {
218        let audio_file = Path::new(temp).join("audio.mkv");
219        let mut encode_audio = Command::new("ffmpeg");
220
221        encode_audio.stdout(Stdio::piped());
222        encode_audio.stderr(Stdio::piped());
223
224        encode_audio.args(["-y", "-hide_banner", "-loglevel", "error"]);
225        encode_audio.args(["-i", &input.to_string_lossy()]);
226        encode_audio.args(["-map_metadata", "0"]);
227        encode_audio.args(["-map", "0", "-c", "copy", "-vn", "-dn"]);
228
229        encode_audio.args(audio_params);
230        encode_audio.arg(&audio_file);
231
232        let output = encode_audio.output()?;
233
234        if !output.status.success() {
235            warn!("FFmpeg failed to encode audio!\n{output:#?}\nParams: {encode_audio:?}");
236            return Ok(None);
237        }
238
239        Ok(Some(audio_file))
240    } else {
241        Ok(None)
242    }
243}
244
245/// Escapes paths in ffmpeg filters if on windows
246#[inline]
247pub fn escape_path_in_filter(path: impl AsRef<Path>) -> anyhow::Result<String> {
248    Ok(if cfg!(windows) {
249        PathAbs::new(path.as_ref())?
250            .to_string_lossy()
251            // This is needed because of how FFmpeg handles absolute file paths on Windows.
252            // https://stackoverflow.com/questions/60440793/how-can-i-use-windows-absolute-paths-with-the-movie-filter-on-ffmpeg
253            .replace('\\', "/")
254            .replace(':', r"\\:")
255    } else {
256        PathAbs::new(path.as_ref())?.to_string_lossy().to_string()
257    }
258    .replace('[', r"\[")
259    .replace(']', r"\]")
260    .replace(',', "\\,"))
261}
262
263/// Pixel formats supported by ffmpeg
264#[derive(Eq, PartialEq, Copy, Clone, Debug, Serialize, Deserialize)]
265pub enum FFPixelFormat {
266    GBRP,
267    GBRP10LE,
268    GBRP12L,
269    GBRP12LE,
270    GRAY10LE,
271    GRAY12L,
272    GRAY12LE,
273    GRAY8,
274    NV12,
275    NV16,
276    NV20LE,
277    NV21,
278    YUV420P,
279    YUV420P10LE,
280    YUV420P12LE,
281    YUV422P,
282    YUV422P10LE,
283    YUV422P12LE,
284    YUV440P,
285    YUV440P10LE,
286    YUV440P12LE,
287    YUV444P,
288    YUV444P10LE,
289    YUV444P12LE,
290    YUVA420P,
291    YUVJ420P,
292    YUVJ422P,
293    YUVJ444P,
294}
295
296impl FFPixelFormat {
297    /// The string to be used with ffmpeg's `-pix_fmt` argument.
298    #[inline]
299    pub fn to_pix_fmt_string(&self) -> &'static str {
300        match self {
301            FFPixelFormat::GBRP => "gbrp",
302            FFPixelFormat::GBRP10LE => "gbrp10le",
303            FFPixelFormat::GBRP12L => "gbrp12l",
304            FFPixelFormat::GBRP12LE => "gbrp12le",
305            FFPixelFormat::GRAY10LE => "gray10le",
306            FFPixelFormat::GRAY12L => "gray12l",
307            FFPixelFormat::GRAY12LE => "gray12le",
308            FFPixelFormat::GRAY8 => "gray",
309            FFPixelFormat::NV12 => "nv12",
310            FFPixelFormat::NV16 => "nv16",
311            FFPixelFormat::NV20LE => "nv20le",
312            FFPixelFormat::NV21 => "nv21",
313            FFPixelFormat::YUV420P => "yuv420p",
314            FFPixelFormat::YUV420P10LE => "yuv420p10le",
315            FFPixelFormat::YUV420P12LE => "yuv420p12le",
316            FFPixelFormat::YUV422P => "yuv422p",
317            FFPixelFormat::YUV422P10LE => "yuv422p10le",
318            FFPixelFormat::YUV422P12LE => "yuv422p12le",
319            FFPixelFormat::YUV440P => "yuv440p",
320            FFPixelFormat::YUV440P10LE => "yuv440p10le",
321            FFPixelFormat::YUV440P12LE => "yuv440p12le",
322            FFPixelFormat::YUV444P => "yuv444p",
323            FFPixelFormat::YUV444P10LE => "yuv444p10le",
324            FFPixelFormat::YUV444P12LE => "yuv444p12le",
325            FFPixelFormat::YUVA420P => "yuva420p",
326            FFPixelFormat::YUVJ420P => "yuvj420p",
327            FFPixelFormat::YUVJ422P => "yuvj422p",
328            FFPixelFormat::YUVJ444P => "yuvj444p",
329        }
330    }
331
332    #[inline]
333    pub fn to_vapoursynth_format(&self) -> anyhow::Result<PresetFormat> {
334        Ok(match self {
335            // Vapoursynth doesn't have a Gray10/Gray12 so use the next best thing.
336            // No quality loss from 10/12->16 but might be slightly slower.
337            FFPixelFormat::GRAY10LE => PresetFormat::Gray16,
338            FFPixelFormat::GRAY12L => PresetFormat::Gray16,
339            FFPixelFormat::GRAY12LE => PresetFormat::Gray16,
340            FFPixelFormat::GRAY8 => PresetFormat::Gray8,
341            FFPixelFormat::YUV420P => PresetFormat::YUV420P8,
342            FFPixelFormat::YUV420P10LE => PresetFormat::YUV420P10,
343            FFPixelFormat::YUV420P12LE => PresetFormat::YUV420P12,
344            FFPixelFormat::YUV422P => PresetFormat::YUV422P8,
345            FFPixelFormat::YUV422P10LE => PresetFormat::YUV422P10,
346            FFPixelFormat::YUV422P12LE => PresetFormat::YUV422P12,
347            FFPixelFormat::YUV440P => PresetFormat::YUV440P8,
348            FFPixelFormat::YUV444P => PresetFormat::YUV444P8,
349            FFPixelFormat::YUV444P10LE => PresetFormat::YUV444P10,
350            FFPixelFormat::YUV444P12LE => PresetFormat::YUV444P12,
351            FFPixelFormat::YUVJ420P => PresetFormat::YUV420P8,
352            FFPixelFormat::YUVJ422P => PresetFormat::YUV422P8,
353            FFPixelFormat::YUVJ444P => PresetFormat::YUV444P8,
354            x => bail!(
355                "pixel format {} cannot be converted to Vapoursynth format",
356                x.to_pix_fmt_string()
357            ),
358        })
359    }
360}
361
362impl FromStr for FFPixelFormat {
363    type Err = anyhow::Error;
364
365    #[inline]
366    fn from_str(s: &str) -> Result<Self, Self::Err> {
367        Ok(match s {
368            "gbrp" => FFPixelFormat::GBRP,
369            "gbrp10le" => FFPixelFormat::GBRP10LE,
370            "gbrp12l" => FFPixelFormat::GBRP12L,
371            "gbrp12le" => FFPixelFormat::GBRP12LE,
372            "gray10le" => FFPixelFormat::GRAY10LE,
373            "gray12l" => FFPixelFormat::GRAY12L,
374            "gray12le" => FFPixelFormat::GRAY12LE,
375            "gray" => FFPixelFormat::GRAY8,
376            "nv12" => FFPixelFormat::NV12,
377            "nv16" => FFPixelFormat::NV16,
378            "nv20le" => FFPixelFormat::NV20LE,
379            "nv21" => FFPixelFormat::NV21,
380            "yuv420p" => FFPixelFormat::YUV420P,
381            "yuv420p10le" => FFPixelFormat::YUV420P10LE,
382            "yuv420p12le" => FFPixelFormat::YUV420P12LE,
383            "yuv422p" => FFPixelFormat::YUV422P,
384            "yuv422p10le" => FFPixelFormat::YUV422P10LE,
385            "yuv422p12le" => FFPixelFormat::YUV422P12LE,
386            "yuv440p" => FFPixelFormat::YUV440P,
387            "yuv440p10le" => FFPixelFormat::YUV440P10LE,
388            "yuv440p12le" => FFPixelFormat::YUV440P12LE,
389            "yuv444p" => FFPixelFormat::YUV444P,
390            "yuv444p10le" => FFPixelFormat::YUV444P10LE,
391            "yuv444p12le" => FFPixelFormat::YUV444P12LE,
392            "yuva420p" => FFPixelFormat::YUVA420P,
393            "yuvj420p" => FFPixelFormat::YUVJ420P,
394            "yuvj422p" => FFPixelFormat::YUVJ422P,
395            "yuvj444p" => FFPixelFormat::YUVJ444P,
396            s => bail!("Unsupported pixel format string: {s}"),
397        })
398    }
399}
av1an_core/ffmpeg.rs

av1an_core/
ffmpeg.rs
