|
| 1 | +// Copyright 2025 Andrew D. Straw. |
| 2 | + |
| 3 | +use camino::Utf8PathBuf; |
| 4 | +use clap::{Parser, ValueEnum}; |
| 5 | +use eyre::{self, Context, Result}; |
| 6 | +use indicatif::{ProgressBar, ProgressStyle}; |
| 7 | +use serde::{Deserialize, Serialize}; |
| 8 | + |
| 9 | +use rusttype::Font; |
| 10 | + |
| 11 | +use font_drawing::stamp_frame; |
| 12 | + |
| 13 | +// trait DisplayTimestamp { |
| 14 | +// fn to_display(&self) -> String; |
| 15 | +// } |
| 16 | + |
| 17 | +// impl DisplayTimestamp for frame_source::Timestamp { |
| 18 | +// fn to_display(&self) -> String { |
| 19 | +// match self { |
| 20 | +// frame_source::Timestamp::Duration(dur) => { |
| 21 | +// format!("{:9.1}ms", dur.as_secs_f64() * 1000.0) |
| 22 | +// } |
| 23 | +// frame_source::Timestamp::Fraction(frac) => { |
| 24 | +// format!("{:2.1}%", frac * 100.0) |
| 25 | +// } |
| 26 | +// } |
| 27 | +// } |
| 28 | +// } |
| 29 | + |
| 30 | +// impl DisplayTimestamp for std::time::Duration { |
| 31 | +// fn to_display(&self) -> String { |
| 32 | +// frame_source::Timestamp::Duration(*self).to_display() |
| 33 | +// } |
| 34 | +// } |
| 35 | + |
| 36 | +// TODO: define SrtMsg only once in this codebase. |
| 37 | +#[derive(Serialize, Deserialize)] |
| 38 | +struct SrtMsg { |
| 39 | + timestamp: chrono::DateTime<chrono::Local>, |
| 40 | +} |
| 41 | + |
| 42 | +#[derive(Parser, Debug)] |
| 43 | +#[command(author, version, about)] |
| 44 | +struct Cli { |
| 45 | + /// Input MP4 video. |
| 46 | + #[arg(long)] |
| 47 | + input: camino::Utf8PathBuf, |
| 48 | + |
| 49 | + /// Output MP4 file. |
| 50 | + #[arg(long)] |
| 51 | + output: Utf8PathBuf, |
| 52 | + |
| 53 | + /// Source of timestamp. |
| 54 | + #[arg(long, value_enum, default_value_t)] |
| 55 | + timestamp_source: TimestampSource, |
| 56 | + |
| 57 | + /// Disable showing progress |
| 58 | + #[arg(short, long, default_value_t)] |
| 59 | + no_progress: bool, |
| 60 | +} |
| 61 | + |
| 62 | +#[derive(Default, Debug, Clone, Copy, ValueEnum, PartialEq)] |
| 63 | +enum TimestampSource { |
| 64 | + #[default] |
| 65 | + BestGuess, |
| 66 | + FrameInfoRecvTime, |
| 67 | + Mp4Pts, |
| 68 | + MispMicrosectime, |
| 69 | + SrtFile, |
| 70 | +} |
| 71 | + |
| 72 | +impl From<TimestampSource> for frame_source::TimestampSource { |
| 73 | + fn from(orig: TimestampSource) -> Self { |
| 74 | + match orig { |
| 75 | + TimestampSource::BestGuess => frame_source::TimestampSource::BestGuess, |
| 76 | + TimestampSource::FrameInfoRecvTime => frame_source::TimestampSource::FrameInfoRecvTime, |
| 77 | + TimestampSource::SrtFile => frame_source::TimestampSource::SrtFile, |
| 78 | + TimestampSource::Mp4Pts => frame_source::TimestampSource::Mp4Pts, |
| 79 | + TimestampSource::MispMicrosectime => frame_source::TimestampSource::MispMicrosectime, |
| 80 | + } |
| 81 | + } |
| 82 | +} |
| 83 | + |
| 84 | +fn main() -> Result<()> { |
| 85 | + env_tracing_logger::init(); |
| 86 | + let cli = Cli::parse(); |
| 87 | + |
| 88 | + // Load the font |
| 89 | + // let font_data = include_bytes!("../Roboto-Regular.ttf"); |
| 90 | + let font_data = ttf_firacode::REGULAR; |
| 91 | + // This only succeeds if collection consists of one font |
| 92 | + let font = Font::try_from_bytes(font_data as &[u8]).expect("Error constructing Font"); |
| 93 | + |
| 94 | + let input_path = cli.input; |
| 95 | + |
| 96 | + let mut srt_file_path = None; |
| 97 | + let is_file = std::fs::metadata(&input_path) |
| 98 | + .with_context(|| format!("While opening input {input_path}"))? |
| 99 | + .is_file(); |
| 100 | + if !is_file { |
| 101 | + eyre::bail!("Input path {input_path} is not a file."); |
| 102 | + } |
| 103 | + let file_ext = input_path.extension().map(|x| x.to_lowercase()); |
| 104 | + |
| 105 | + if file_ext != Some("mp4".into()) { |
| 106 | + eyre::bail!("Input file must be an MP4 file."); |
| 107 | + } |
| 108 | + |
| 109 | + let mut srt_path = input_path.clone(); |
| 110 | + srt_path.set_extension("srt"); |
| 111 | + if srt_path.exists() && std::fs::metadata(&srt_path)?.is_file() { |
| 112 | + srt_file_path = Some(srt_path); |
| 113 | + } else if cli.timestamp_source == TimestampSource::SrtFile { |
| 114 | + eyre::bail!("Source specified as SRT file, but {srt_path} is not a file."); |
| 115 | + } |
| 116 | + |
| 117 | + if !cli.no_progress { |
| 118 | + eprintln!("Performing initial open of \"{input_path}\"."); |
| 119 | + } |
| 120 | + |
| 121 | + let mut src = frame_source::FrameSourceBuilder::new(&input_path) |
| 122 | + .timestamp_source(cli.timestamp_source.into()) |
| 123 | + .srt_file_path(srt_file_path.map(Into::into)) |
| 124 | + .show_progress(!cli.no_progress) |
| 125 | + .build_source()?; |
| 126 | + |
| 127 | + let t0: chrono::DateTime<chrono::Utc> = src.frame0_time().unwrap().into(); |
| 128 | + if !cli.no_progress { |
| 129 | + eprintln!("Done with initial open."); |
| 130 | + } |
| 131 | + |
| 132 | + let mut ffmpeg_wtr = |
| 133 | + ffmpeg_writer::FfmpegWriter::new(cli.output.as_str(), Default::default(), None)?; |
| 134 | + |
| 135 | + let mut pb: Option<ProgressBar> = if !cli.no_progress { |
| 136 | + let (lower_bound, _upper_bound) = src.iter().size_hint(); |
| 137 | + |
| 138 | + // Custom progress bar with space at right end to prevent obscuring last |
| 139 | + // digit with cursor. |
| 140 | + let style = |
| 141 | + ProgressStyle::with_template("Burning timestamps {wide_bar} {pos}/{len} ETA: {eta} ")?; |
| 142 | + Some(ProgressBar::new(lower_bound.try_into().unwrap()).with_style(style)) |
| 143 | + } else { |
| 144 | + None |
| 145 | + }; |
| 146 | + |
| 147 | + for (idx, frame) in src.iter().enumerate() { |
| 148 | + let frame = frame?; |
| 149 | + |
| 150 | + if let Some(pb) = pb.as_mut() { |
| 151 | + pb.inc(1); |
| 152 | + } |
| 153 | + |
| 154 | + let im = if let Some(im) = frame.decoded() { |
| 155 | + im |
| 156 | + } else { |
| 157 | + eyre::bail!("Frame {idx} has no decoded image data.",); |
| 158 | + }; |
| 159 | + |
| 160 | + let text = match frame.timestamp() { |
| 161 | + frame_source::Timestamp::Duration(dur) => { |
| 162 | + format!( |
| 163 | + "{}", |
| 164 | + (t0 + dur).to_rfc3339_opts(chrono::format::SecondsFormat::Millis, true) |
| 165 | + ) |
| 166 | + } |
| 167 | + frame_source::Timestamp::Fraction(frac) => { |
| 168 | + format!("{:2.1}%", frac * 100.0) |
| 169 | + } |
| 170 | + }; |
| 171 | + |
| 172 | + let mut frame_rgb8 = im.into_pixel_format()?.owned(); |
| 173 | + stamp_frame(&mut frame_rgb8, &font, &text)?; |
| 174 | + |
| 175 | + let dy_im = strand_dynamic_frame::DynamicFrame::from_static_ref(&frame_rgb8); |
| 176 | + |
| 177 | + ffmpeg_wtr.write_dynamic_frame(&dy_im)?; |
| 178 | + } |
| 179 | + |
| 180 | + ffmpeg_wtr.close()?; |
| 181 | + if let Some(pb) = pb { |
| 182 | + pb.finish_and_clear(); |
| 183 | + } |
| 184 | + |
| 185 | + Ok(()) |
| 186 | +} |
0 commit comments