diff options
Diffstat (limited to 'rust/crates/media/src/lib.rs')
| -rw-r--r-- | rust/crates/media/src/lib.rs | 610 |
1 files changed, 610 insertions, 0 deletions
diff --git a/rust/crates/media/src/lib.rs b/rust/crates/media/src/lib.rs new file mode 100644 index 00000000..239f5c75 --- /dev/null +++ b/rust/crates/media/src/lib.rs @@ -0,0 +1,610 @@ +use std::collections::HashMap; +use std::path::Path; + +use ffmpeg_next as ffmpeg; +use ffmpeg::codec::context::Context as CodecContext; +use ffmpeg::format::context::Input; +use ffmpeg::format::{Pixel, Sample, sample}; +use ffmpeg::software::{resampling, scaling}; +use ffmpeg::{ChannelLayout, Rational, media}; +pub use time::MediaTime; + +const MICROS_PER_SECOND: i64 = 1_000_000; +const AUDIO_PREROLL_SECONDS: f64 = 0.5; +const FORWARD_DECODE_LIMIT_SECONDS: f64 = 10.0; +const EOF_RESEEK_MARGIN_SECONDS: f64 = 2.0; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("ffmpeg error: {0}")] + Ffmpeg(#[from] ffmpeg::Error), + #[error("no video stream found")] + NoVideoStream, + #[error("no audio stream found")] + NoAudioStream, + #[error("stream index {0} does not exist")] + StreamNotFound(usize), + #[error("stream index {0} is not a video stream")] + NotAVideoStream(usize), + #[error("stream index {0} is not an audio stream")] + NotAnAudioStream(usize), + #[error("no frames could be decoded from stream {0}")] + NoFramesDecoded(usize), + #[error("invalid time range: start must be before end")] + InvalidTimeRange, + #[error("bucket count must be at least 1")] + InvalidBucketCount, + #[error("max dimension must be at least 1")] + InvalidMaxDimension, +} + +pub type Result<T, E = Error> = std::result::Result<T, E>; + +#[derive(Debug, Clone)] +pub struct VideoStreamInfo { + pub index: usize, + pub width: u32, + pub height: u32, + pub frame_rate: f64, + pub frame_count: Option<u64>, +} + +#[derive(Debug, Clone)] +pub struct AudioStreamInfo { + pub index: usize, + pub channels: u32, + pub sample_rate: u32, +} + +#[derive(Debug, Clone)] +pub struct MediaInfo { + pub duration: MediaTime, + pub video_streams: Vec<VideoStreamInfo>, + pub audio_streams: Vec<AudioStreamInfo>, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct VideoFrame { + pub width: u32, + pub height: u32, + pub data: Vec<u8>, +} + +#[derive(Debug, Clone)] +pub struct AudioSamples { + pub sample_rate: u32, + pub channels: u32, + pub samples: Vec<f32>, +} + +fn seconds_to_ts(seconds: f64, time_base: Rational) -> i64 { + (seconds * f64::from(time_base.denominator()) / f64::from(time_base.numerator())).round() + as i64 +} + +fn ts_to_seconds(ts: i64, time_base: Rational) -> f64 { + ts as f64 * f64::from(time_base.numerator()) / f64::from(time_base.denominator()) +} + +fn seconds_to_micros(seconds: f64) -> i64 { + (seconds * MICROS_PER_SECOND as f64) as i64 +} + +fn effective_layout(layout: ChannelLayout, channels: u16) -> ChannelLayout { + if layout.is_empty() { + ChannelLayout::default(i32::from(channels)) + } else { + layout + } +} + +struct VideoPipeline { + stream_index: usize, + decoder: ffmpeg::codec::decoder::Video, + time_base: Rational, + last: Option<(i64, ffmpeg::frame::Video)>, + drained: bool, + scaler: Option<(Pixel, u32, u32, scaling::Context)>, +} + +impl VideoPipeline { + fn open(input: &Input, stream_index: usize) -> Result<Self> { + let stream = input + .streams() + .find(|stream| stream.index() == stream_index) + .ok_or(Error::StreamNotFound(stream_index))?; + if stream.parameters().medium() != media::Type::Video { + return Err(Error::NotAVideoStream(stream_index)); + } + let decoder = CodecContext::from_parameters(stream.parameters())? + .decoder() + .video()?; + Ok(Self { + stream_index, + decoder, + time_base: stream.time_base(), + last: None, + drained: false, + scaler: None, + }) + } + + fn frame_at( + &mut self, + input: &mut Input, + target_ts: i64, + seek_seconds: f64, + ) -> Result<ffmpeg::frame::Video> { + if let Some((last_pts, last_frame)) = &self.last { + if target_ts >= *last_pts { + if self.drained { + return Ok(last_frame.clone()); + } + let forward_limit = + seconds_to_ts(FORWARD_DECODE_LIMIT_SECONDS, self.time_base); + if target_ts - last_pts <= forward_limit { + let candidate = self.last.take(); + return self.decode_forward(input, target_ts, candidate); + } + } + } + self.decoder.flush(); + self.drained = false; + self.last = None; + let micros = seconds_to_micros(seek_seconds.max(0.0)); + input.seek(micros, ..micros)?; + match self.decode_forward(input, target_ts, None) { + Err(Error::NoFramesDecoded(_)) => { + self.decoder.flush(); + input.seek(0, ..0)?; + self.decode_forward(input, i64::MAX, None) + } + result => result, + } + } + + fn decode_forward( + &mut self, + input: &mut Input, + target_ts: i64, + mut candidate: Option<(i64, ffmpeg::frame::Video)>, + ) -> Result<ffmpeg::frame::Video> { + let mut frame = ffmpeg::frame::Video::empty(); + for (stream, packet) in input.packets() { + if stream.index() != self.stream_index { + continue; + } + self.decoder.send_packet(&packet)?; + while self.decoder.receive_frame(&mut frame).is_ok() { + let pts = frame.pts().unwrap_or(i64::MIN); + if pts <= target_ts { + candidate = Some((pts, frame.clone())); + } else { + self.last = Some((pts, frame.clone())); + return match candidate { + Some((_, frame)) => Ok(frame), + None => Ok(frame.clone()), + }; + } + frame = ffmpeg::frame::Video::empty(); + } + } + self.decoder.send_eof()?; + while self.decoder.receive_frame(&mut frame).is_ok() { + let pts = frame.pts().unwrap_or(i64::MIN); + if pts <= target_ts { + candidate = Some((pts, frame.clone())); + } else { + self.last = Some((pts, frame.clone())); + return match candidate { + Some((_, frame)) => Ok(frame), + None => Ok(frame.clone()), + }; + } + frame = ffmpeg::frame::Video::empty(); + } + self.drained = true; + match candidate { + Some((pts, frame)) => { + self.last = Some((pts, frame.clone())); + Ok(frame) + } + None => Err(Error::NoFramesDecoded(self.stream_index)), + } + } + + fn rgba_frame(&mut self, frame: &ffmpeg::frame::Video) -> Result<VideoFrame> { + convert_to_rgba( + &mut self.scaler, + frame, + frame.width(), + frame.height(), + ) + } +} + +fn convert_to_rgba( + cache: &mut Option<(Pixel, u32, u32, scaling::Context)>, + frame: &ffmpeg::frame::Video, + width: u32, + height: u32, +) -> Result<VideoFrame> { + let key = (frame.format(), frame.width(), frame.height()); + let matches = cache + .as_ref() + .is_some_and(|(format, w, h, _)| (*format, *w, *h) == key); + if !matches { + let scaler = scaling::Context::get( + frame.format(), + frame.width(), + frame.height(), + Pixel::RGBA, + width, + height, + scaling::Flags::BILINEAR, + )?; + *cache = Some((key.0, key.1, key.2, scaler)); + } + let scaler = &mut cache.as_mut().expect("scaler just inserted").3; + let mut converted = ffmpeg::frame::Video::new(Pixel::RGBA, width, height); + scaler.run(frame, &mut converted)?; + let stride = converted.stride(0); + let row_bytes = width as usize * 4; + let plane = converted.data(0); + let mut data = Vec::with_capacity(row_bytes * height as usize); + if stride == row_bytes { + data.extend_from_slice(&plane[..row_bytes * height as usize]); + } else { + for row in 0..height as usize { + data.extend_from_slice(&plane[row * stride..row * stride + row_bytes]); + } + } + Ok(VideoFrame { width, height, data }) +} + +pub struct Decoder { + input: Input, + video: HashMap<usize, VideoPipeline>, +} + +impl Decoder { + pub fn open(path: impl AsRef<Path>) -> Result<Self> { + let input = ffmpeg::format::input(path.as_ref())?; + Ok(Self { + input, + video: HashMap::new(), + }) + } + + pub fn info(&self) -> MediaInfo { + let mut video_streams = Vec::new(); + let mut audio_streams = Vec::new(); + let mut stream_duration_seconds = 0.0_f64; + for stream in self.input.streams() { + let duration = stream.duration(); + if duration > 0 { + stream_duration_seconds = + stream_duration_seconds.max(ts_to_seconds(duration, stream.time_base())); + } + match stream.parameters().medium() { + media::Type::Video => { + let Ok(context) = CodecContext::from_parameters(stream.parameters()) else { + continue; + }; + let Ok(decoder) = context.decoder().video() else { + continue; + }; + let rate = stream.avg_frame_rate(); + let frame_rate = if rate.denominator() == 0 { + 0.0 + } else { + f64::from(rate.numerator()) / f64::from(rate.denominator()) + }; + let frames = stream.frames(); + video_streams.push(VideoStreamInfo { + index: stream.index(), + width: decoder.width(), + height: decoder.height(), + frame_rate, + frame_count: (frames > 0).then_some(frames as u64), + }); + } + media::Type::Audio => { + let Ok(context) = CodecContext::from_parameters(stream.parameters()) else { + continue; + }; + let Ok(decoder) = context.decoder().audio() else { + continue; + }; + audio_streams.push(AudioStreamInfo { + index: stream.index(), + channels: u32::from(decoder.channels()), + sample_rate: decoder.rate(), + }); + } + _ => {} + } + } + let duration_micros = self.input.duration(); + let duration_seconds = if duration_micros > 0 { + duration_micros as f64 / MICROS_PER_SECOND as f64 + } else { + stream_duration_seconds + }; + MediaInfo { + duration: MediaTime::from_seconds_f64(duration_seconds).unwrap_or(MediaTime::ZERO), + video_streams, + audio_streams, + } + } + + pub fn decode_video_frame(&mut self, stream_index: usize, time: MediaTime) -> Result<VideoFrame> { + let raw = self.decode_video_raw(stream_index, time)?; + let pipeline = self + .video + .get_mut(&stream_index) + .expect("pipeline created by decode_video_raw"); + pipeline.rgba_frame(&raw) + } + + fn decode_video_raw( + &mut self, + stream_index: usize, + time: MediaTime, + ) -> Result<ffmpeg::frame::Video> { + if !self.video.contains_key(&stream_index) { + let pipeline = VideoPipeline::open(&self.input, stream_index)?; + self.video.insert(stream_index, pipeline); + } + let target_seconds = time.to_seconds_f64().max(0.0); + let duration_seconds = self.input.duration().max(0) as f64 / MICROS_PER_SECOND as f64; + let pipeline = self + .video + .get_mut(&stream_index) + .expect("pipeline just inserted"); + let mut target_ts = seconds_to_ts(target_seconds, pipeline.time_base); + let mut seek_seconds = target_seconds; + if duration_seconds > 0.0 && target_seconds > duration_seconds { + target_ts = i64::MAX; + seek_seconds = (duration_seconds - EOF_RESEEK_MARGIN_SECONDS).max(0.0); + } + pipeline.frame_at(&mut self.input, target_ts, seek_seconds) + } + + pub fn decode_audio_range( + &mut self, + stream_index: usize, + start: MediaTime, + end: MediaTime, + ) -> Result<AudioSamples> { + self.decode_audio_range_inner(stream_index, start, end, true) + } + + fn decode_audio_range_inner( + &mut self, + stream_index: usize, + start: MediaTime, + end: MediaTime, + pad_to_range: bool, + ) -> Result<AudioSamples> { + if end <= start { + return Err(Error::InvalidTimeRange); + } + let stream = self + .input + .streams() + .find(|stream| stream.index() == stream_index) + .ok_or(Error::StreamNotFound(stream_index))?; + if stream.parameters().medium() != media::Type::Audio { + return Err(Error::NotAnAudioStream(stream_index)); + } + let time_base = stream.time_base(); + let mut decoder = CodecContext::from_parameters(stream.parameters())? + .decoder() + .audio()?; + let rate = decoder.rate(); + let channels = u32::from(decoder.channels()); + let layout = effective_layout(decoder.channel_layout(), decoder.channels()); + let mut resampler = resampling::Context::get( + decoder.format(), + layout, + rate, + Sample::F32(sample::Type::Packed), + layout, + rate, + )?; + let start_seconds = start.to_seconds_f64().max(0.0); + let end_seconds = end.to_seconds_f64(); + let seek_seconds = (start_seconds - AUDIO_PREROLL_SECONDS).max(0.0); + let seek_micros = seconds_to_micros(seek_seconds); + self.input.seek(seek_micros, ..seek_micros)?; + + let rate_f64 = f64::from(rate); + let channel_count = channels as usize; + let mut anchor_seconds: Option<f64> = None; + let mut produced: i64 = 0; + let mut out_samples: Vec<f32> = Vec::new(); + let mut done = false; + + let mut frame = ffmpeg::frame::Audio::empty(); + for (packet_stream, packet) in self.input.packets() { + if packet_stream.index() != stream_index { + continue; + } + decoder.send_packet(&packet)?; + while decoder.receive_frame(&mut frame).is_ok() { + if frame.channel_layout().is_empty() { + frame.set_channel_layout(layout); + } + if anchor_seconds.is_none() { + anchor_seconds = Some(match frame.pts() { + Some(pts) => ts_to_seconds(pts, time_base), + None => seek_seconds, + }); + } + let mut converted = ffmpeg::frame::Audio::empty(); + resampler.run(&frame, &mut converted)?; + append_resampled( + &converted, + anchor_seconds.expect("anchor set above"), + start_seconds, + end_seconds, + rate_f64, + channel_count, + &mut produced, + &mut out_samples, + ); + frame = ffmpeg::frame::Audio::empty(); + let anchor = anchor_seconds.expect("anchor set above"); + let end_index = ((end_seconds - anchor) * rate_f64).round() as i64; + if produced >= end_index { + done = true; + break; + } + } + if done { + break; + } + } + if !done { + decoder.send_eof()?; + while decoder.receive_frame(&mut frame).is_ok() { + if frame.channel_layout().is_empty() { + frame.set_channel_layout(layout); + } + if anchor_seconds.is_none() { + anchor_seconds = Some(match frame.pts() { + Some(pts) => ts_to_seconds(pts, time_base), + None => seek_seconds, + }); + } + let mut converted = ffmpeg::frame::Audio::empty(); + resampler.run(&frame, &mut converted)?; + append_resampled( + &converted, + anchor_seconds.expect("anchor set above"), + start_seconds, + end_seconds, + rate_f64, + channel_count, + &mut produced, + &mut out_samples, + ); + frame = ffmpeg::frame::Audio::empty(); + } + let mut converted = ffmpeg::frame::Audio::empty(); + if resampler.flush(&mut converted).is_ok() && converted.samples() > 0 { + append_resampled( + &converted, + anchor_seconds.unwrap_or(seek_seconds), + start_seconds, + end_seconds, + rate_f64, + channel_count, + &mut produced, + &mut out_samples, + ); + } + } + if pad_to_range && let Some(anchor) = anchor_seconds { + let start_index = (((start_seconds - anchor) * rate_f64).round() as i64).max(0); + let end_index = (((end_seconds - anchor) * rate_f64).round() as i64).max(start_index); + let expected = ((end_index - start_index) as usize) * channel_count; + if out_samples.len() < expected { + out_samples.resize(expected, 0.0); + } + } + Ok(AudioSamples { + sample_rate: rate, + channels, + samples: out_samples, + }) + } +} + +fn append_resampled( + converted: &ffmpeg::frame::Audio, + anchor_seconds: f64, + start_seconds: f64, + end_seconds: f64, + rate: f64, + channels: usize, + produced: &mut i64, + out_samples: &mut Vec<f32>, +) { + let count = converted.samples() as i64; + if count == 0 { + return; + } + let start_index = (((start_seconds - anchor_seconds) * rate).round() as i64).max(0); + let end_index = (((end_seconds - anchor_seconds) * rate).round() as i64).max(start_index); + let low = (*produced).max(start_index); + let high = (*produced + count).min(end_index); + if low < high { + if out_samples.is_empty() && low > start_index { + out_samples.resize(((low - start_index) as usize) * channels, 0.0); + } + let data = converted.data(0); + let first = (low - *produced) as usize * channels; + let last = (high - *produced) as usize * channels; + for bytes in data[first * 4..last * 4].chunks_exact(4) { + out_samples.push(f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])); + } + } + *produced += count; +} + +pub fn thumbnail(path: impl AsRef<Path>, max_dimension: u32) -> Result<VideoFrame> { + if max_dimension == 0 { + return Err(Error::InvalidMaxDimension); + } + let mut decoder = Decoder::open(path.as_ref())?; + let video_stream = decoder + .info() + .video_streams + .first() + .cloned() + .ok_or(Error::NoVideoStream)?; + let raw = decoder.decode_video_raw(video_stream.index, MediaTime::ZERO)?; + let source_max = raw.width().max(raw.height()); + let scale = (f64::from(max_dimension) / f64::from(source_max)).min(1.0); + let width = ((f64::from(raw.width()) * scale).round() as u32).max(1); + let height = ((f64::from(raw.height()) * scale).round() as u32).max(1); + let mut cache = None; + convert_to_rgba(&mut cache, &raw, width, height) +} + +pub fn waveform_summary(path: impl AsRef<Path>, bucket_count: usize) -> Result<Vec<(f32, f32)>> { + if bucket_count == 0 { + return Err(Error::InvalidBucketCount); + } + let mut decoder = Decoder::open(path.as_ref())?; + let info = decoder.info(); + let audio_stream = info.audio_streams.first().ok_or(Error::NoAudioStream)?; + let samples = decoder.decode_audio_range_inner( + audio_stream.index, + MediaTime::ZERO, + info.duration, + false, + )?; + let channels = samples.channels as usize; + let frame_count = samples.samples.len() / channels; + let bucket_size = frame_count.div_ceil(bucket_count).max(1); + Ok((0..bucket_count) + .map(|bucket| { + let start_frame = bucket * bucket_size; + let end_frame = (start_frame + bucket_size).min(frame_count); + if start_frame >= end_frame { + return (0.0, 0.0); + } + let slice = &samples.samples[start_frame * channels..end_frame * channels]; + let mut low = f32::INFINITY; + let mut high = f32::NEG_INFINITY; + for &sample in slice { + low = low.min(sample); + high = high.max(sample); + } + (low, high) + }) + .collect()) +} |