more_fps/
ffmpeg.rs

1use crate::command;
2use crate::Error;
3use crate::TimeRange;
4use crate::TimeRanges;
5use log::debug;
6use regex::Regex;
7use rust_decimal::Decimal;
8use std::fs;
9use std::io;
10use std::io::Write;
11use std::num::NonZeroUsize;
12use std::path::Path;
13use std::path::PathBuf;
14
15pub fn ffmpeg<T: AsRef<str>>(requirements: T) -> Result<String, Error> {
16    command::run(
17        "ffmpeg",
18        requirements.as_ref().try_into()?,
19        Error::FfmpegCommand,
20    )
21}
22
23#[derive(Debug)]
24pub struct FfmpegStepper {
25    input_file: PathBuf,
26    frames_dir: PathBuf,
27    videos_dir: PathBuf,
28    scene_file: PathBuf,
29    /// text file
30    concat_file: PathBuf,
31    /// the video file generated from concatting the videos_dir
32    video_file: PathBuf,
33    crf: NonZeroUsize,
34    fps: NonZeroUsize,
35    input_extension: String,
36}
37
38impl FfmpegStepper {
39    pub fn try_new(
40        temp_dir: &Path,
41        input_file: PathBuf,
42        crf: NonZeroUsize,
43        fps: NonZeroUsize,
44    ) -> Result<Self, Error> {
45        let frames_dir = temp_dir.join("frames");
46        dir_exists_or_create(&frames_dir)?;
47
48        let videos_dir = temp_dir.join("videos");
49        dir_exists_or_create(&videos_dir)?;
50
51        let input_extension = get_extension(&input_file)?.to_owned();
52
53        let scene_file = temp_dir.join("scene_timestamps.txt");
54        let concat_file = temp_dir.join("concat.txt");
55        let video_file = temp_dir.join(format!("video.{input_extension}"));
56
57        Ok(Self {
58            input_file,
59            frames_dir,
60            scene_file,
61            concat_file,
62            video_file,
63            videos_dir,
64            crf,
65            fps,
66            input_extension,
67        })
68    }
69
70    pub fn frames_dir(&self) -> &Path {
71        &self.frames_dir
72    }
73
74    /// Returns the number of video pieces we've already extracted
75    pub fn existing_video_count(&self) -> Result<usize, Error> {
76        let count = fs::read_dir(&self.videos_dir)
77            .map_err(|_| Error::ReadDir(self.videos_dir.to_path_buf()))?
78            .count();
79        Ok(count)
80    }
81
82    /// Returns a frame extractor which can be used to extract frames
83    // todo if scene_file doesn't exist yet, we should clear the video_dir
84    pub fn flattened_time_ranges(
85        &self,
86        max_step_size: NonZeroUsize,
87        scene_gt: &str,
88    ) -> Result<Vec<TimeRange>, Error> {
89        let mut timestamps: Vec<Decimal> = vec![Decimal::ZERO];
90        timestamps.append(&mut find_scene_timestamps(
91            &self.input_file,
92            scene_gt,
93            &self.scene_file,
94        )?);
95        let time_ranges = timestamps
96            .as_slice()
97            .windows(2)
98            .map(|pair| (pair[0], pair[1]))
99            .map(|(start, end)| {
100                TimeRanges::try_new(start, max_step_size, end)
101                    .ok_or(Error::UnableToCreateTimeRanges(start, max_step_size, end))
102            })
103            .collect::<Result<Vec<TimeRanges>, _>>()?;
104
105        let flattened_time_ranges = time_ranges
106            .into_iter()
107            .flatten()
108            // skipping because we already have existing videos
109            .collect();
110
111        Ok(flattened_time_ranges)
112    }
113
114    pub fn clear_frames_dir(&self) -> Result<(), Error> {
115        if self.frames_dir.exists() {
116            fs::remove_dir_all(&self.frames_dir)?;
117            fs::create_dir_all(&self.frames_dir)?;
118        }
119        Ok(())
120    }
121
122    pub fn extract_frames(&self, time_range: &TimeRange) -> Result<&Path, Error> {
123        self.clear_frames_dir()?;
124
125        extract_frames(
126            &time_range.start,
127            &self.input_file,
128            &time_range.duration(),
129            &self.frames_dir,
130        )?;
131        Ok(&self.frames_dir)
132    }
133
134    /// Takes the extracted frames when calling `extract_frames` and creates a video in the
135    /// `video_dir`
136    pub fn frames_to_video(&self, video_number: usize, input_dir: PathBuf) -> Result<(), Error> {
137        let video_path = self
138            .videos_dir
139            .join(format!("{video_number}.{}", self.input_extension));
140
141        // maybe this will work for windows?
142        #[cfg(target_os="windows")]
143        let args = format!("-y framerate {} -{} -pattern_type sequence -i %08d.png -crf {} -c:v libx264 -pix_fmt yuv420p {}",
144            self.fps,
145            self.crf,
146            video_path.display()
147        );
148        #[cfg(not(target_os="windows"))]
149        let args = format!("-y -framerate {} -pattern_type glob -i '*.png' -crf {} -c:v libx264 -pix_fmt yuv420p {}",
150            self.fps,
151            self.crf,
152            video_path.display()
153        );
154        let requirements = command::Requirements {
155            args: &args,
156            current_dir: input_dir,
157        };
158        command::run("ffmpeg", requirements, Error::FfmpegCommand)?;
159        Ok(())
160    }
161    /// When you're done extracting frames, call this function and we'll aggregate the
162    /// videos + audio + subtitles into the output file provided
163    pub fn aggregate(&self, output_file: &Path) -> Result<(), Error> {
164        concat_videos(&self.concat_file, &self.videos_dir, &self.video_file)?;
165
166        // Need -max_interleave_delta:
167        // https://trac.ffmpeg.org/ticket/6037
168        let args = format!("-ignore_unknown -y -i {} -vn -i {} -map 0 -c:v copy -map 1 -c:a copy -c:s copy -map_chapters 1 -max_interleave_delta 0 {}",
169            &self.video_file.display(),
170            &self.input_file.display(),
171            output_file.display()
172        );
173        ffmpeg(args)?;
174        Ok(())
175    }
176}
177
178/// Give a path to create the concat file for ffmpeg to reference
179/// This concat file will have all of the files in the videos_dir provided
180/// Then use ffmpeg to concat the videos into the final output_file path
181pub fn concat_videos(
182    concat_file_path: &Path,
183    videos_dir: &Path,
184    output_file: &Path,
185) -> Result<(), Error> {
186    if concat_file_path.exists() {
187        fs::remove_file(concat_file_path)?;
188    }
189
190    // todo look into not collecting so many times...
191    let lines = fs::read_dir(videos_dir)?
192        .collect::<Result<Vec<_>, _>>()?
193        .iter()
194        .map(fs::DirEntry::path)
195        .map(fs::canonicalize)
196        .collect::<Result<Vec<_>, _>>()?
197        .iter()
198        .map(|path| format!("file {}", path.display()))
199        .rev()
200        .collect::<Vec<_>>()
201        .join("\n");
202
203    let mut concat_file = fs::File::create(concat_file_path)?;
204    write!(concat_file, "{lines}")?;
205
206    let args = format!(
207        "-y -f concat -safe 0 -i {} -c copy {}",
208        concat_file_path.display(),
209        output_file.display()
210    );
211    ffmpeg(args)?;
212    Ok(())
213}
214
215/// Extracts audio from `input_file` to the audio file you pass in
216pub fn extract_audio(input_file: &Path, audio_file: &Path) -> Result<(), Error> {
217    let args = format!(
218        "-y -i {} -map 0:a -c copy {}",
219        input_file.display(),
220        audio_file.display()
221    );
222    ffmpeg(args)?;
223    Ok(())
224}
225
226fn ffprobe<T: AsRef<str>>(requirements: T) -> Result<String, Error> {
227    command::run(
228        "ffprobe",
229        requirements.as_ref().try_into()?,
230        Error::FfprobeCommand,
231    )
232}
233
234fn parse_timestamps(lines: &str) -> Result<Vec<Decimal>, Error> {
235    Regex::new(r"best_effort_timestamp_time=(\d+.\d+)|")
236        .unwrap()
237        .captures_iter(lines)
238        .filter_map(|caps| caps.get(1))
239        .map(|m| m.as_str())
240        .map(Decimal::from_str_exact)
241        .collect::<Result<Vec<_>, _>>()
242        .map_err(Error::from)
243}
244
245// https://superuser.com/questions/819573/split-up-a-video-using-ffmpeg-through-scene-detection
246pub fn find_scene_timestamps(
247    input_file: &Path,
248    scene_gt: &str,
249    scene_file: &Path,
250) -> Result<Vec<Decimal>, Error> {
251    if scene_file.exists() {
252        debug!("{scene_file:?} exists, so using data in that file");
253        let decimals = fs::read_to_string(scene_file)?
254            .split('\n')
255            .map(|s| Decimal::from_str_exact(s).map_err(Error::from))
256            .collect::<Result<Vec<_>, Error>>()?;
257
258        return Ok(decimals);
259    }
260    debug!("creating file: {scene_file:?}");
261
262    let args = format!(
263        r#"-show_frames -of compact=p=0 -f lavfi "movie={},select=gt(scene\,{scene_gt})""#,
264        input_file.display()
265    );
266
267    let stdout = ffprobe(args)?;
268    let decimals = parse_timestamps(&stdout)?;
269    let mut f = fs::File::create(scene_file)?;
270    f.write_all(
271        decimals
272            .iter()
273            .map(|e| e.to_string())
274            .collect::<Vec<String>>()
275            .join("\n")
276            .as_str()
277            .as_bytes(),
278    )?;
279    Ok(decimals)
280}
281
282pub fn extract_frames(
283    start: &Decimal,
284    input_file: &Path,
285    duration: &Decimal,
286    output_dir: &Path,
287) -> Result<(), Error> {
288    let args = format!(
289        "-ss {start} -i {} -t {duration} -frame_pts true {}",
290        input_file.display(),
291        output_dir.join("frame_%08d.png").display()
292    );
293    ffmpeg(args)?;
294    Ok(())
295}
296
297fn get_extension(path: &Path) -> Result<&str, Error> {
298    let extension = path
299        .extension()
300        .ok_or(Error::MissingExtension(path.to_path_buf()))?;
301    let extension = extension
302        .to_str()
303        .ok_or(Error::InvalidUnicode(extension.to_os_string()))?;
304    Ok(extension)
305}
306
307fn dir_exists_or_create(path: &Path) -> Result<(), io::Error> {
308    if !path.exists() {
309        fs::create_dir_all(path)?;
310    }
311    Ok(())
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    const SCENE_TIMESTAMPS: &str = "
319media_type=video|stream_index=0|key_frame=1|pkt_pts=9760|pkt_pts_time=9.760000|pkt_dts=9760|pkt_dts_time=9.760000|best_effort_timestamp=9760|best_effort_timestamp_time=9.760000|pkt_duration=N/A|pkt_duration_time=N/A|pkt_pos=858320|pkt_size=6220800|width=1920|height=1080|pix_fmt=yuv420p10le|sample_aspect_ratio=1:1|pict_type=I|coded_picture_number=0|display_picture_number=0|interlaced_frame=0|top_field_first=0|repeat_pict=0|color_range=unknown|color_space=unknown|color_primaries=unknown|color_transfer=unknown|chroma_location=unspecified|tag:lavfi.scene_score=0.504959
320media_type=video|stream_index=0|key_frame=1|pkt_pts=13513|pkt_pts_time=13.513000|pkt_dts=13513|pkt_dts_time=13.513000|best_effort_timestamp=13513|best_effort_timestamp_time=13.513000|pkt_duration=N/A|pkt_duration_time=N/A|pkt_pos=1176734|pkt_size=6220800|width=1920|height=1080|pix_fmt=yuv420p10le|sample_aspect_ratio=1:1|pict_type=I|coded_picture_number=0|display_picture_number=0|interlaced_frame=0|top_field_first=0|repeat_pict=0|color_range=unknown|color_space=unknown|color_primaries=unknown|color_transfer=unknown|chroma_location=unspecified|tag:lavfi.scene_score=0.428674
321media_type=video|stream_index=0|key_frame=1|pkt_pts=18936|pkt_pts_time=18.936000|pkt_dts=18936|pkt_dts_time=18.936000|best_effort_timestamp=18936|best_effort_timestamp_time=18.936000|pkt_duration=N/A|pkt_duration_time=N/A|pkt_pos=1694070|pkt_size=6220800|width=1920|height=1080|pix_fmt=yuv420p10le|sample_aspect_ratio=1:1|pict_type=I|coded_picture_number=0|display_picture_number=0|interlaced_frame=0|top_field_first=0|repeat_pict=0|color_range=unknown|color_space=unknown|color_primaries=unknown|color_transfer=unknown|chroma_location=unspecified|tag:lavfi.scene_score=0.990057
322media_type=video|stream_index=0|key_frame=1|pkt_pts=22105|pkt_pts_time=22.105000|pkt_dts=22105|pkt_dts_time=22.105000|best_effort_timestamp=22105|best_effort_timestamp_time=22.105000|pkt_duration=N/A|pkt_duration_time=N/A|pkt_pos=2498438|pkt_size=6220800|width=1920|height=1080|pix_fmt=yuv420p10le|sample_aspect_ratio=1:1|pict_type=I|coded_picture_number=0|display_picture_number=0|interlaced_frame=0|top_field_first=0|repeat_pict=0|color_range=unknown|color_space=unknown|color_primaries=unknown|color_transfer=unknown|chroma_location=unspecified|tag:lavfi.scene_score=0.547889";
323
324    #[test]
325    fn scene_time_stamps() {
326        let actual = parse_timestamps(SCENE_TIMESTAMPS).unwrap();
327        let expected = vec![
328            Decimal::from_str_exact("9.76").unwrap(),
329            Decimal::from_str_exact("13.513").unwrap(),
330            Decimal::from_str_exact("18.936").unwrap(),
331            Decimal::from_str_exact("22.105").unwrap(),
332        ];
333        assert_eq!(actual, expected);
334    }
335}
more_fps/ffmpeg.rs

more_fps/
ffmpeg.rs
