From eb712a095dafd79cfc340dc7d59bb79220746f77 Mon Sep 17 00:00:00 2001 From: Rik Heijdens Date: Fri, 23 Oct 2026 06:53:02 +0000 Subject: [PATCH 1/2] Add fps_conversion option for even-cadence frame-rate resampling `fps` resampling keeps the source frame nearest each output timestamp, so non-integer ratios space frames unevenly: 30 -> 24 fps keeps 4 of every 5 frames and motion steps go 1x, 1x, 1x, 2x. VideoStreamRequest gains `fps_conversion`: - "nearest" (default): unchanged FFmpeg `fps` behaviour. - "blend": FFmpeg `framerate`, cross-fades the two neighbouring frames. Even cadence at about nearest's cost; moving edges show double images. - "motion": bilateral block-matching motion estimation with overlapped block motion compensation (the model of FFmpeg's `minterpolate`), implemented in src/decoder/motion.rs. Motion is estimated only for source pairs that bracket an output frame. About 1.25x nearest's decode time, against ~12x for `minterpolate`, with equal or better fidelity. Both interpolating modes run after the resize and after HDR tone mapping, keep nearest's frame count and timestamps (so planned buffer sizes and A/V sync are unchanged), and fall back to nearest when the source rate is an integer multiple of the target (48 -> 24, 23.976 -> 24). The video buffersrc now declares the stream frame rate. --- README.md | 10 + avtensor.pyi | 6 + src/decoder/mod.rs | 817 +++++++++++++++++++++++++++--- src/decoder/motion.rs | 1080 ++++++++++++++++++++++++++++++++++++++++ src/ffi/mod.rs | 14 +- src/util/test_utils.rs | 30 ++ 6 files changed, 1879 insertions(+), 78 deletions(-) create mode 100644 src/decoder/motion.rs diff --git a/README.md b/README.md index 5b858c5..7bdd287 100644 --- a/README.md +++ b/README.md @@ -113,6 +113,16 @@ request = MediaDecodeRequest( streams = avtensor.decode_asset(request) ``` +By default `fps` keeps the source frame nearest each output timestamp, which +spaces frames unevenly for non-integer ratios (30 -> 24 fps keeps 4 of every 5 +frames). `fps_conversion="blend"` cross-fades neighbouring frames instead, and +`fps_conversion="motion"` interpolates along estimated motion (about 1.25x +the cost of the default): + +```python +VideoStreamRequest(fps=24.0, fps_conversion="motion") +``` + ### Decode a time window (seek) ```python diff --git a/avtensor.pyi b/avtensor.pyi index 086871e..0634f82 100644 --- a/avtensor.pyi +++ b/avtensor.pyi @@ -38,6 +38,11 @@ class VideoStreamRequest: # "raw" when you need the actual HDR signal, e.g. training on PQ # masters or colorimetric measurement. hdr_mode: Literal["tonemap", "raw"] | None + # How `fps` resampling builds each output frame: "nearest" (default) + # keeps the nearest source frame (uneven spacing for e.g. 30 -> 24 fps); + # "blend" cross-fades neighbouring frames; "motion" interpolates along + # estimated block motion (about 1.2-1.6x the cost of "nearest"). + fps_conversion: Literal["nearest", "blend", "motion"] | None def __init__( self, @@ -52,6 +57,7 @@ class VideoStreamRequest: device: str | None = None, dtype: str | None = None, hdr_mode: Literal["tonemap", "raw"] | None = None, + fps_conversion: Literal["nearest", "blend", "motion"] | None = None, ): ... class LoudnessNormalization: diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs index 21fd676..fe5b693 100644 --- a/src/decoder/mod.rs +++ b/src/decoder/mod.rs @@ -3,7 +3,10 @@ use rsmpeg::{ avcodec::{AVCodec, AVCodecContext, AVPacket}, avfilter::{AVFilter, AVFilterContextMut, AVFilterGraph, AVFilterInOut}, avformat::{AVFormatContextInput, AVStreamRef}, - avutil::{av_q2d, get_sample_fmt_name, AVChannelLayout, AVDictionary, AVFrame, AVRational}, + avutil::{ + av_d2q, av_inv_q, av_q2d, get_sample_fmt_name, AVChannelLayout, AVDictionary, AVFrame, + AVRational, + }, error::RsmpegError, ffi::{self, AVPixelFormat, AV_NOPTS_VALUE}, }; @@ -21,6 +24,7 @@ use crate::conversion; mod cuda; mod io; +mod motion; struct StreamContext { stream_type: StreamType, @@ -43,6 +47,8 @@ struct FilteringContext<'graph> { /// How decoded frames reach `frame_data`: through the filter graph, or /// directly when the graph would add nothing but copies. direct_path: DirectPath, + /// Motion-interpolating frame-rate conversion, fed by the filter graph. + motion: Option>, /// A Tensor, allocated at the start of the media decoding process, to which /// filtered frames will be written. frame_data: tch::Tensor, @@ -60,6 +66,15 @@ struct FilterContext<'graph> { buffersink_ctx: AVFilterContextMut<'graph>, } +/// Frame-rate conversion by motion interpolation ([`FpsConversion::Motion`]). +/// The stream's filter graph resizes/tone maps into 8-bit planar YUV; the +/// resampler retimes and interpolates; a second small graph converts the +/// result to the output pixel format. +struct MotionStage<'graph> { + resampler: motion::Resampler, + output: FilterContext<'graph>, +} + /// Direct (graph-bypassing) processing mode for a stream. When a stream's /// only processing is something the decoder can do while writing into the /// output tensor, the filter graph — and its per-decode slice-threading @@ -153,6 +168,9 @@ pub struct VideoStreamRequest { /// HDR handling: `Tonemap` (default, SDR BT.709 preview) or `Raw` /// (preserve PQ/HLG code values; matrix/range conversion only). pub hdr_mode: HdrMode, + /// How `frame_rate` resampling builds each output frame. Ignored when + /// `frame_rate` is unset. + pub fps_conversion: FpsConversion, } #[derive(Default, Debug, Clone)] @@ -257,7 +275,7 @@ pub fn decode_media( // widths libswscale can't vectorize. log::debug!("Setting up filter graphs for the streams."); let mut filter_graphs: Vec<_> = (0..stream_ctx.len()) - .map(|_| AVFilterGraph::new()) + .map(|_| (AVFilterGraph::new(), AVFilterGraph::new())) .collect(); let mut filter_ctx = init_filters(&mut filter_graphs, stream_ctx)?; @@ -393,6 +411,7 @@ pub fn decode_media( buffersrc_ctx, buffersink_ctx, direct_path, + motion, frame_data, frame_data_ptr, frame_pts, @@ -414,6 +433,7 @@ pub fn decode_media( buffersrc_ctx, buffersink_ctx, direct_path, + motion, frame_data, frame_data_ptr, frame_pts, @@ -447,6 +467,7 @@ pub fn decode_media( buffersink_ctx .as_mut() .context("filter graph missing for graph-path stream")?, + motion.as_mut(), frame_data, frame_data_ptr, frame_pts, @@ -953,6 +974,21 @@ fn init_video_stream_context( let mut video_filter = VideoFilterConfig::try_from(request)?; video_filter.source_color = SourceColorInfo::from_codec_context(&decode_context); log::debug!("Source color info: {:?}", video_filter.source_color); + if let (Some(target), Some(source)) = (request.frame_rate, video_stream.guess_framerate()) { + // Interpolation adds nothing when every output frame lands on a + // source frame (48 -> 24), and for NTSC rate pairs (23.976 -> 24) it + // would re-synthesize nearly every frame to fix a 0.1% drift. + if video_filter.fps_conversion != FpsConversion::Nearest + && is_integer_downsample(av_q2d(source), target) + { + log::debug!( + "Source frame rate {} is an integer multiple of {target}, \ + using nearest-frame conversion", + av_q2d(source) + ); + video_filter.fps_conversion = FpsConversion::Nearest; + } + } let plan = plan_video_stream(input_format_context, request, seek)?; log::debug!( @@ -1360,6 +1396,7 @@ fn decode_packet( buffersrc_ctx, buffersink_ctx, direct_path, + motion, frame_data, frame_data_ptr, frame_pts, @@ -1382,6 +1419,7 @@ fn decode_packet( buffersrc_ctx, buffersink_ctx, direct_path, + motion, frame_data, frame_data_ptr, frame_pts, @@ -1429,6 +1467,7 @@ fn decode_packet( buffersrc_ctx, buffersink_ctx, direct_path, + motion, frame_data, frame_data_ptr, frame_pts, @@ -1523,6 +1562,7 @@ fn receive_and_filter_frames( buffersrc_ctx: &mut Option>, buffersink_ctx: &mut Option>, direct_path: &mut DirectPath, + motion: &mut Option>, frame_data: &mut tch::Tensor, frame_data_ptr: &mut usize, frame_pts: &mut Vec, @@ -1627,6 +1667,7 @@ fn receive_and_filter_frames( buffersink_ctx .as_mut() .context("filter graph missing for graph-path stream")?, + motion.as_mut(), frame_data, frame_data_ptr, frame_pts, @@ -1666,40 +1707,114 @@ enum FilterFrameError { Unexpected(anyhow::Error), } -/// Apply FFmpeg filter graph to a decoded frame. +/// Apply FFmpeg filter graph to a decoded frame (`None` flushes it), then +/// the stream's [`MotionStage`] if it has one. #[allow(clippy::too_many_arguments)] fn filter_frame( stream_type: StreamType, frame: Option, buffersrc_ctx: &mut AVFilterContextMut, buffersink_ctx: &mut AVFilterContextMut, + motion: Option<&mut MotionStage>, frame_data: &mut tch::Tensor, frame_data_ptr: &mut usize, frame_pts: &mut Vec, seek: &Seek, ) -> Result, FilterFrameError> { - // Send the frame to the filter source. - buffersrc_ctx - .buffersrc_add_frame(frame, None) - .map_err(|e| { - FilterFrameError::Unexpected( - anyhow!(e).context("Error submitting the frame to the filtergraph"), - ) - })?; + let flush = frame.is_none(); + submit_frame(buffersrc_ctx, frame)?; + let Some(stage) = motion else { + return drain_sink( + stream_type, + buffersink_ctx, + frame_data, + frame_data_ptr, + frame_pts, + seek, + ); + }; - let mut filtered_frames: Vec = Vec::new(); - loop { - // Get the filtered frames from the buffersink. - let mut filtered_frame = match buffersink_ctx.buffersink_get_frame(None) { - Ok(frame) => frame, - Err(RsmpegError::BufferSinkDrainError) | Err(RsmpegError::BufferSinkEofError) => break, - Err(_) => { - return Err(FilterFrameError::Unexpected(anyhow!( - "Get frame from buffersink failed" - ))) - } - }; + let MotionStage { + resampler, + output: + FilterContext { + buffersrc_ctx: output_src, + buffersink_ctx: output_sink, + }, + } = stage; + let mut filtered_frames = Vec::new(); + let mut emit = |frames: Vec, filtered_frames: &mut Vec| { + for frame in frames { + submit_frame(output_src, Some(frame))?; + filtered_frames.extend(drain_sink( + stream_type, + output_sink, + frame_data, + frame_data_ptr, + frame_pts, + seek, + )?); + } + Ok::<_, FilterFrameError>(()) + }; + while let Some(mut frame) = receive_filtered(buffersink_ctx)? { + frame.set_time_base(buffersink_ctx.get_time_base()); + let frames = resampler + .push(frame) + .map_err(|e| FilterFrameError::Unexpected(e.context("Motion interpolation failed")))?; + emit(frames, &mut filtered_frames)?; + } + if flush { + emit(resampler.finish(), &mut filtered_frames)?; + submit_frame(output_src, None)?; + filtered_frames.extend(drain_sink( + stream_type, + output_sink, + frame_data, + frame_data_ptr, + frame_pts, + seek, + )?); + } + Ok(filtered_frames) +} +/// Sends a frame (or, with `None`, end of stream) to a graph's source. +fn submit_frame( + buffersrc_ctx: &mut AVFilterContextMut, + frame: Option, +) -> Result<(), FilterFrameError> { + buffersrc_ctx.buffersrc_add_frame(frame, None).map_err(|e| { + FilterFrameError::Unexpected( + anyhow!(e).context("Error submitting the frame to the filtergraph"), + ) + }) +} + +/// The next frame waiting in a graph's sink, if any. +fn receive_filtered( + buffersink_ctx: &mut AVFilterContextMut, +) -> Result, FilterFrameError> { + match buffersink_ctx.buffersink_get_frame(None) { + Ok(frame) => Ok(Some(frame)), + Err(RsmpegError::BufferSinkDrainError) | Err(RsmpegError::BufferSinkEofError) => Ok(None), + Err(_) => Err(FilterFrameError::Unexpected(anyhow!( + "Get frame from buffersink failed" + ))), + } +} + +/// Writes every frame waiting in `buffersink_ctx` to `frame_data`. +fn drain_sink( + stream_type: StreamType, + buffersink_ctx: &mut AVFilterContextMut, + frame_data: &mut tch::Tensor, + frame_data_ptr: &mut usize, + frame_pts: &mut Vec, + seek: &Seek, +) -> Result, FilterFrameError> { + let mut filtered_frames: Vec = Vec::new(); + while let Some(mut filtered_frame) = receive_filtered(buffersink_ctx)? { filtered_frame.set_time_base(buffersink_ctx.get_time_base()); filtered_frame.set_pict_type(ffi::AV_PICTURE_TYPE_NONE); @@ -2482,7 +2597,57 @@ impl TryFrom> for HdrMode { } } -#[derive(Debug)] +/// Whether `source` fps reaches `target` by keeping every n-th frame, +/// allowing for NTSC's 1000/1001 rate offset (23.976 vs 24). +fn is_integer_downsample(source: f64, target: f64) -> bool { + if !(source > 0.0 && target > 0.0) { + return false; + } + let ratio = source / target; + let n = ratio.round(); + n >= 1.0 && (ratio / n - 1.0).abs() < 2e-3 +} + +/// How frame-rate resampling builds each output frame. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum FpsConversion { + /// Keep the source frame nearest each output timestamp (FFmpeg `fps`), + /// dropping or repeating whole frames. Fast, but a non-integer ratio + /// gives uneven frame spacing (30 -> 24 fps keeps 4 of every 5 frames, + /// so motion steps go 1x, 1x, 1x, 2x). + #[default] + Nearest, + /// Cross-fade the two source frames around each output timestamp + /// (FFmpeg `framerate`). Even timing at near-`Nearest` cost, but moving + /// edges show as double images. + Blend, + /// Motion-compensated interpolation (block matching with overlapped + /// block motion compensation; the model of FFmpeg's `minterpolate` + /// defaults, implemented in [`motion`] for speed). Even timing without + /// double images, at about 1.2-1.6x `Nearest`'s cost. Falls back to the + /// nearest frame at detected scene changes; works on 8-bit YUV, so + /// deeper sources are reduced to 8 bits before interpolation. + Motion, + // `Blend` and `Motion` fall back to `Nearest` when the source rate is an + // integer multiple of the target (see `is_integer_downsample`). +} + +impl TryFrom> for FpsConversion { + type Error = anyhow::Error; + + fn try_from(value: Option<&str>) -> Result { + match value { + None | Some("nearest") => Ok(Self::Nearest), + Some("blend") => Ok(Self::Blend), + Some("motion") => Ok(Self::Motion), + Some(other) => Err(anyhow::anyhow!( + "fps_conversion must be \"nearest\", \"blend\" or \"motion\", got {other:?}" + )), + } + } +} + +#[derive(Debug, Clone)] pub struct VideoFilterConfig { /// Desired frame rate for the video. frame_rate: Option, @@ -2498,12 +2663,18 @@ pub struct VideoFilterConfig { device: Option, /// How HDR sources are converted (tone-mapped SDR vs raw code values). hdr_mode: HdrMode, + /// How `frame_rate` resampling builds each output frame. + fps_conversion: FpsConversion, } impl Display for VideoFilterConfig { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let mut filters = Vec::new(); - if let Some(frame_rate) = self.frame_rate { + // Nearest-frame selection runs first so dropped frames are never + // scaled. The interpolating modes need every source frame and their + // cost grows with pixel count, so they run after the resize; + // `Motion` runs outside the graph (see `MotionStage`). + if let (Some(frame_rate), FpsConversion::Nearest) = (self.frame_rate, self.fps_conversion) { filters.push(format!("fps={}", frame_rate)); } @@ -2539,6 +2710,12 @@ impl Display for VideoFilterConfig { filters.push("zscale=p=bt709:t=bt709:m=bt709:range=tv".to_string()); } + // After tone mapping: `framerate` takes YUV only, and 8-bit PQ/HLG + // code values would band once linearized. + if let (Some(frame_rate), FpsConversion::Blend) = (self.frame_rate, self.fps_conversion) { + filters.push(format!("framerate=fps={}", frame_rate)); + } + filters.push(format!("format=pix_fmts={}", self.pixel_format)); if !filters.is_empty() { write!(f, "{}", filters.join(",")) @@ -2548,6 +2725,25 @@ impl Display for VideoFilterConfig { } } +impl VideoFilterConfig { + /// Whether frame-rate conversion runs in a [`MotionStage`] after the + /// filter graph rather than inside it. + fn uses_motion_stage(&self) -> bool { + self.frame_rate.is_some() && self.fps_conversion == FpsConversion::Motion + } + + /// Graph feeding a [`MotionStage`]: this config's resize and tone + /// mapping, without frame-rate conversion, ending in the 8-bit planar + /// YUV format `working_format` that motion interpolation works in. + fn motion_input(&self, working_format: &str) -> Self { + Self { + frame_rate: None, + pixel_format: working_format.to_string(), + ..self.clone() + } + } +} + impl Default for VideoFilterConfig { fn default() -> Self { Self { @@ -2558,6 +2754,7 @@ impl Default for VideoFilterConfig { source_color: Default::default(), device: Default::default(), hdr_mode: Default::default(), + fps_conversion: Default::default(), } } } @@ -2576,6 +2773,7 @@ impl TryFrom<&VideoStreamRequest> for VideoFilterConfig { OutputDtype::Float32 => "gbrpf32le".to_string(), }, hdr_mode: req.hdr_mode, + fps_conversion: req.fps_conversion, ..Default::default() }) } @@ -2703,13 +2901,17 @@ impl Default for LoudnessNormalization { /// Create transcoding context corresponding to the given `stream_contexts`, the /// added filter contexts is mutable reference to objects stored in /// `filter_graphs`. +/// Builds each stream's processing. Every stream gets a pair of graphs: the +/// first is its filter graph, the second is only used by a [`MotionStage`]. fn init_filters( - filter_graphs: &mut [AVFilterGraph], + filter_graphs: &mut [(AVFilterGraph, AVFilterGraph)], stream_contexts: Vec, ) -> Result>, anyhow::Error> { let mut filter_ctx = Vec::with_capacity(stream_contexts.len()); - for (filter_graph, stream_context) in filter_graphs.iter_mut().zip(stream_contexts) { + for ((filter_graph, motion_graph), stream_context) in + filter_graphs.iter_mut().zip(stream_contexts) + { let StreamContext { mut dec_ctx, stream_index, @@ -2774,16 +2976,33 @@ fn init_filters( // Only build the filter graph when it will actually process frames: // configuring a graph spins up its slice-threading worker pool, which // is pure overhead for direct-path streams. - let (buffersrc_ctx, buffersink_ctx) = if matches!(direct_path, DirectPath::Disabled) { - let FilterContext { - buffersrc_ctx, - buffersink_ctx, - } = init_filter(filter_graph, &mut dec_ctx, &filter_config) - .context("Failed to initialize filter")?; - log::debug!("Initialized filter for stream {stream_index} of type {stream_type:?}"); - (Some(buffersrc_ctx), Some(buffersink_ctx)) - } else { - (None, None) + let mut motion = None; + let (buffersrc_ctx, buffersink_ctx) = match &filter_config { + _ if !matches!(direct_path, DirectPath::Disabled) => (None, None), + FilterConfig::Video(config) if config.uses_motion_stage() => { + let working_format = motion_working_format(config, &dec_ctx); + let input_config = FilterConfig::Video(config.motion_input(working_format)); + let FilterContext { + buffersrc_ctx, + buffersink_ctx, + } = init_filter(filter_graph, &mut dec_ctx, &input_config) + .context("Failed to initialize filter")?; + motion = Some( + init_motion_stage(motion_graph, &buffersink_ctx, config) + .context("Failed to initialize motion interpolation")?, + ); + log::debug!("Initialized motion interpolation for stream {stream_index}"); + (Some(buffersrc_ctx), Some(buffersink_ctx)) + } + _ => { + let FilterContext { + buffersrc_ctx, + buffersink_ctx, + } = init_filter(filter_graph, &mut dec_ctx, &filter_config) + .context("Failed to initialize filter")?; + log::debug!("Initialized filter for stream {stream_index} of type {stream_type:?}"); + (Some(buffersrc_ctx), Some(buffersink_ctx)) + } }; filter_ctx.push(FilteringContext { @@ -2794,6 +3013,7 @@ fn init_filters( buffersrc_ctx, buffersink_ctx, direct_path, + motion, frame_data, frame_data_ptr: 0, frame_pts: Vec::new(), @@ -2804,22 +3024,108 @@ fn init_filters( Ok(filter_ctx) } -fn init_filter<'graph>( - filter_graph: &'graph mut AVFilterGraph, - dec_ctx: &mut AVCodecContext, - filter_spec: &FilterConfig, -) -> Result, anyhow::Error> { - // Cap the graph's slice-threading pool. FFmpeg's default (0 = one worker - // per core) builds a large thread pool for every decode on many-core - // machines, which oversubscribes badly when many decodes run - // concurrently. Slice threading cannot use more workers than there are - // output rows/slices anyway, so a small pool preserves single-decode - // filter performance. +/// Caps a graph's slice-threading pool. FFmpeg's default (0 = one worker per +/// core) builds a large thread pool for every decode on many-core machines, +/// which oversubscribes badly when many decodes run concurrently. Slice +/// threading cannot use more workers than there are output rows/slices +/// anyway, so a small pool preserves single-decode filter performance. +fn cap_graph_threads(filter_graph: &mut AVFilterGraph) { let graph_threads = std::thread::available_parallelism() .map(|n| n.get()) .unwrap_or(16) .min(16) as i32; unsafe { (*filter_graph.as_mut_ptr()).nb_threads = graph_threads }; +} + +/// The 8-bit planar YUV format a [`MotionStage`] interpolates in. Keeps the +/// source's chroma subsampling when frames keep their size, so nothing is +/// resampled. A resize or HDR tone mapping already resamples chroma, and +/// subsampling at the output size would leave less chroma detail than the +/// other conversions (which upsample the source chroma straight to the +/// output size), so those use 4:4:4. +fn motion_working_format(config: &VideoFilterConfig, dec_ctx: &AVCodecContext) -> &'static str { + let resized = config.width.is_some_and(|w| w != dec_ctx.width as usize) + || config.height.is_some_and(|h| h != dec_ctx.height as usize); + let tonemapped = config.source_color.is_hdr() && config.hdr_mode == HdrMode::Tonemap; + let source = buffersrc_pixel_format(dec_ctx); + match motion::working_pixel_format(source) { + "yuvj420p" | "yuvj422p" | "yuvj444p" if resized || tonemapped => "yuvj444p", + _ if resized || tonemapped => "yuv444p", + format => format, + } +} + +/// Builds the [`MotionStage`] that resamples the frames leaving `input_sink` +/// (the stream's filter graph, configured by [`VideoFilterConfig::motion_input`]) +/// to `config`'s frame rate and converts them to its pixel format. +fn init_motion_stage<'graph>( + filter_graph: &'graph mut AVFilterGraph, + input_sink: &AVFilterContextMut, + config: &VideoFilterConfig, +) -> Result, anyhow::Error> { + let frame_rate = config + .frame_rate + .context("motion interpolation needs a frame rate")?; + // The `fps` filter's rate and time base for the same value, so output + // timestamps match `Nearest`'s. + let out_tb = av_inv_q(av_d2q(frame_rate, i32::MAX)); + let (width, height, format) = ( + input_sink.get_w(), + input_sink.get_h(), + input_sink.get_format(), + ); + let resampler = motion::Resampler::new( + motion::Layout::new(width, height, format)?, + input_sink.get_time_base(), + out_tb, + ); + + let sar = input_sink.get_sample_aspect_ratio(); + let mut args = format!( + "video_size={width}x{height}:pix_fmt={format}:time_base={}/{}:pixel_aspect={}/{}", + out_tb.num, out_tb.den, sar.num, sar.den, + ); + let (colorspace, range) = unsafe { + ( + ffi::av_buffersink_get_colorspace(input_sink.as_ptr()), + ffi::av_buffersink_get_color_range(input_sink.as_ptr()), + ) + }; + if let Some(name) = ffi_color_name(unsafe { ffi::av_color_space_name(colorspace) }) { + args.push_str(&format!(":colorspace={name}")); + } + if let Some(name) = ffi_color_name(unsafe { ffi::av_color_range_name(range) }) { + args.push_str(&format!(":range={name}")); + } + let output_config = VideoFilterConfig { + pixel_format: config.pixel_format.clone(), + ..Default::default() + }; + + cap_graph_threads(filter_graph); + let (mut buffersrc_ctx, mut buffersink_ctx) = + create_video_endpoints(filter_graph, &args, &output_config)?; + link_graph( + filter_graph, + &mut buffersrc_ctx, + &mut buffersink_ctx, + &output_config.to_filter_spec()?, + )?; + Ok(MotionStage { + resampler, + output: FilterContext { + buffersrc_ctx, + buffersink_ctx, + }, + }) +} + +fn init_filter<'graph>( + filter_graph: &'graph mut AVFilterGraph, + dec_ctx: &mut AVCodecContext, + filter_spec: &FilterConfig, +) -> Result, anyhow::Error> { + cap_graph_threads(filter_graph); let (mut buffersrc_ctx, mut buffersink_ctx, filter_spec) = { match (dec_ctx.codec_type, &filter_spec) { @@ -2847,14 +3153,12 @@ fn init_filter<'graph>( } }; - // Endpoints for the filter graph - // - // Yes the outputs' name is `in`, this is in alignment with upstream samples in ffmpeg / rsmpeg which may confuse you when debugging this. - let outputs = AVFilterInOut::new(c"in", &mut buffersrc_ctx, 0); - let inputs = AVFilterInOut::new(c"out", &mut buffersink_ctx, 0); - - let (_inputs, _outputs) = filter_graph.parse_ptr(&filter_spec, Some(inputs), Some(outputs))?; - filter_graph.config()?; + link_graph( + filter_graph, + &mut buffersrc_ctx, + &mut buffersink_ctx, + &filter_spec, + )?; Ok(FilterContext { buffersrc_ctx, @@ -2862,6 +3166,25 @@ fn init_filter<'graph>( }) } +/// Parses `filter_spec` between the graph's source and sink and configures +/// the graph. +fn link_graph( + filter_graph: &AVFilterGraph, + buffersrc_ctx: &mut AVFilterContextMut, + buffersink_ctx: &mut AVFilterContextMut, + filter_spec: &std::ffi::CStr, +) -> Result<(), anyhow::Error> { + // Endpoints for the filter graph + // + // Yes the outputs' name is `in`, this is in alignment with upstream samples in ffmpeg / rsmpeg which may confuse you when debugging this. + let outputs = AVFilterInOut::new(c"in", buffersrc_ctx, 0); + let inputs = AVFilterInOut::new(c"out", buffersink_ctx, 0); + + let (_inputs, _outputs) = filter_graph.parse_ptr(filter_spec, Some(inputs), Some(outputs))?; + filter_graph.config()?; + Ok(()) +} + /// Initializes the video filter graph according to the provided `filter_spec`. /// /// Arguments: @@ -2873,26 +3196,7 @@ fn init_video_filter<'graph, F: VideoFilterSpec>( dec_ctx: &mut AVCodecContext, filter_spec: &F, ) -> Result<(AVFilterContextMut<'graph>, AVFilterContextMut<'graph>), anyhow::Error> { - log::debug!("Creating buffer/buffersink filters for video stream."); - let buffersrc = - AVFilter::get_by_name(c"buffer").context("filtering source element not found")?; - let buffersink = - AVFilter::get_by_name(c"buffersink").context("filtering sink element not found")?; - - // Hardware decoders report a GPU pixel format (e.g. AV_PIX_FMT_CUDA), but - // frames are downloaded to system memory before filtering — declare the - // underlying software format to the filter graph instead. - let buffersrc_pix_fmt = if dec_ctx.pix_fmt == ffi::AV_PIX_FMT_CUDA { - if dec_ctx.sw_pix_fmt != ffi::AV_PIX_FMT_NONE { - dec_ctx.sw_pix_fmt - } else { - // NVDEC emits NV12 for 8-bit content. - ffi::AV_PIX_FMT_NV12 - } - } else { - dec_ctx.pix_fmt - }; - + let buffersrc_pix_fmt = buffersrc_pixel_format(dec_ctx); let mut args = format!( "video_size={}x{}:pix_fmt={}:time_base={}/{}:pixel_aspect={}/{}", dec_ctx.width, @@ -2913,7 +3217,47 @@ fn init_video_filter<'graph, F: VideoFilterSpec>( if let Some(name) = ffi_color_name(unsafe { ffi::av_color_range_name(dec_ctx.color_range) }) { args.push_str(&format!(":range={name}")); } + // Filters that synthesize frames (e.g. `tpad`) size them from the link's + // frame rate; without it they compute a 1/0 duration. + if dec_ctx.framerate.num > 0 && dec_ctx.framerate.den > 0 { + args.push_str(&format!( + ":frame_rate={}/{}", + dec_ctx.framerate.num, dec_ctx.framerate.den + )); + } + + create_video_endpoints(filter_graph, &args, filter_spec) +} + +/// The pixel format decoded frames enter the filter graph in. Hardware +/// decoders report a GPU pixel format (e.g. AV_PIX_FMT_CUDA), but frames are +/// downloaded to system memory before filtering, so this is the underlying +/// software format. +fn buffersrc_pixel_format(dec_ctx: &AVCodecContext) -> AVPixelFormat { + if dec_ctx.pix_fmt == ffi::AV_PIX_FMT_CUDA { + if dec_ctx.sw_pix_fmt != ffi::AV_PIX_FMT_NONE { + dec_ctx.sw_pix_fmt + } else { + // NVDEC emits NV12 for 8-bit content. + ffi::AV_PIX_FMT_NV12 + } + } else { + dec_ctx.pix_fmt + } +} +/// Creates a video graph's `buffer` source (with `args`) and a `buffersink` +/// accepting `filter_spec`'s pixel format. +fn create_video_endpoints<'graph, F: VideoFilterSpec>( + filter_graph: &'graph AVFilterGraph, + args: &str, + filter_spec: &F, +) -> Result<(AVFilterContextMut<'graph>, AVFilterContextMut<'graph>), anyhow::Error> { + log::debug!("Creating buffer/buffersink filters for video stream."); + let buffersrc = + AVFilter::get_by_name(c"buffer").context("filtering source element not found")?; + let buffersink = + AVFilter::get_by_name(c"buffersink").context("filtering sink element not found")?; let args = &CString::new(args).unwrap(); let buffer_src_context = filter_graph @@ -3045,6 +3389,11 @@ impl VideoFilterSpec for VideoFilterConfig { "rgb24" => Some(ffi::AV_PIX_FMT_RGB24), "gbrpf32le" => Some(ffi::AV_PIX_FMT_GBRPF32LE), "yuv420p" => Some(ffi::AV_PIX_FMT_YUV420P), + "yuv422p" => Some(ffi::AV_PIX_FMT_YUV422P), + "yuv444p" => Some(ffi::AV_PIX_FMT_YUV444P), + "yuvj420p" => Some(ffi::AV_PIX_FMT_YUVJ420P), + "yuvj422p" => Some(ffi::AV_PIX_FMT_YUVJ422P), + "yuvj444p" => Some(ffi::AV_PIX_FMT_YUVJ444P), _ => None, // Default to None if no pixel format is specified } } @@ -3074,7 +3423,10 @@ mod tests { use crate::util::{ mel::{create_mel_spectrogram, MelSpectrogramConfig}, - test_utils::{generate_test_video_file, init_logger, ChannelLayout, TestVideoParameters}, + test_utils::{ + generate_moving_bar_video, generate_test_video_file, init_logger, ChannelLayout, + TestVideoParameters, + }, }; use super::*; @@ -3296,6 +3648,235 @@ mod tests { Ok(()) } + /// Horizontal centroid of the white bar in each decoded frame. + fn bar_centroids(frames: &tch::Tensor) -> Vec { + // [T, H, W, C] uint8 -> mean brightness per column, [T, W]. + let columns = frames.to_kind(tch::Kind::Float).mean_dim( + [1i64, 3].as_slice(), + false, + tch::Kind::Float, + ); + let x = tch::Tensor::arange(columns.size()[1], (tch::Kind::Float, tch::Device::Cpu)); + let centroids = (&columns * &x).sum_dim_intlist(1, false, tch::Kind::Float) + / columns.sum_dim_intlist(1, false, tch::Kind::Float); + Vec::::try_from(centroids.to_kind(tch::Kind::Double)).unwrap() + } + + #[test_case(FpsConversion::Nearest ; "nearest")] + #[test_case(FpsConversion::Blend ; "blend")] + #[test_case(FpsConversion::Motion ; "motion")] + fn test_fps_conversion_motion_spacing(fps_conversion: FpsConversion) -> anyhow::Result<()> { + init_logger(); + // 120 px/s at 30 fps is 4 px per source frame; an evenly spaced 24 + // fps output moves 5 px per frame. Nearest keeps 4 of every 5 source + // frames, so its steps go 4, 8, 4, 4, 4, 8, ... + let video = generate_moving_bar_video(30, 120, Duration::from_secs(2))?; + let decoded = decode_media( + MediaDecodeRequest { + source: MediaSource::Uri(video.path().to_str().unwrap().into()), + start_time: None, + end_time: None, + video_stream: Some(VideoStreamRequest { + frame_rate: Some(24.0), + fps_conversion, + ..Default::default() + }), + audio_streams: None, + }, + None, + )?; + let stream = &decoded[0]; + let frames = stream.data.as_ref().unwrap(); + // Every mode fills the planned buffer exactly. + assert_eq!(frames.size()[0], 48); + assert_eq!(stream.frame_pts.len(), 48); + for (i, pts) in stream.frame_pts.iter().enumerate() { + assert!((pts - i as f64 / 24.0).abs() < 1e-6, "frame {i} pts {pts}"); + } + + let centroids = bar_centroids(frames); + let steps: Vec = centroids.windows(2).map(|w| w[1] - w[0]).collect(); + let max_step = steps.iter().cloned().fold(f64::MIN, f64::max); + let min_step = steps.iter().cloned().fold(f64::MAX, f64::min); + log::info!("{fps_conversion:?}: steps {min_step:.2}..{max_step:.2}"); + match fps_conversion { + FpsConversion::Nearest => { + assert!(min_step < 4.5 && max_step > 7.5, "{steps:?}"); + } + FpsConversion::Blend | FpsConversion::Motion => { + assert!(min_step > 4.5 && max_step < 5.5, "{steps:?}"); + } + } + Ok(()) + } + + /// Motion interpolation works for every source chroma layout and HDR + /// tone mapping, with any output size and dtype, and keeps `Nearest`'s + /// frame count and timestamps. + #[test_case("yuv420p", None, OutputDtype::Uint8 ; "yuv420p")] + #[test_case("yuv422p10le", None, OutputDtype::Uint8 ; "yuv422p10le")] + #[test_case("yuv444p", Some((150, 98)), OutputDtype::Uint8 ; "yuv444p resized odd width")] + #[test_case("yuvj420p", Some((96, 64)), OutputDtype::Float32 ; "yuvj420p float32")] + #[test_case("yuv420p10le", Some((96, 64)), OutputDtype::Uint8 ; "hdr tone mapped")] + fn test_fps_conversion_motion_formats( + pixel_format: &str, + size: Option<(u32, u32)>, + dtype: OutputDtype, + ) -> anyhow::Result<()> { + init_logger(); + let hdr = pixel_format == "yuv420p10le"; + let params = TestVideoParameters { + width: 160, + height: 120, + frame_rate: 25.0, + duration: Duration::from_secs(2), + video_codec: if pixel_format == "yuvj420p" { + "mjpeg".to_string() + } else { + "libx264".to_string() + }, + pixel_format: Some(pixel_format.to_string()), + color_primaries: hdr.then(|| "bt2020".to_string()), + color_trc: hdr.then(|| "smpte2084".to_string()), + colorspace: hdr.then(|| "bt2020nc".to_string()), + ..Default::default() + }; + let video = generate_test_video_file(¶ms)?; + let decode = |fps_conversion| { + decode_media( + MediaDecodeRequest { + source: MediaSource::Uri(video.path().to_str().unwrap().into()), + start_time: None, + end_time: None, + video_stream: Some(VideoStreamRequest { + frame_rate: Some(24.0), + width: size.map(|s| s.0), + height: size.map(|s| s.1), + dtype, + fps_conversion, + ..Default::default() + }), + audio_streams: None, + }, + None, + ) + .map(|mut streams| streams.remove(0)) + }; + let nearest = decode(FpsConversion::Nearest)?; + let motion = decode(FpsConversion::Motion)?; + let (n, m) = (nearest.data.unwrap(), motion.data.unwrap()); + assert_eq!(m.size(), n.size()); + assert_eq!(motion.frame_pts, nearest.frame_pts); + // The first frame lands on a source frame and passes through; later + // ones are interpolated but stay close to the nearest source frame + // (testsrc moves only a little per frame). + let diff = |i: i64| -> f64 { + (m.get(i).to_kind(tch::Kind::Float) - n.get(i).to_kind(tch::Kind::Float)) + .abs() + .mean(tch::Kind::Float) + .double_value(&[]) + }; + let scale = if dtype == OutputDtype::Float32 { + 1.0 / 255.0 + } else { + 1.0 + }; + // 10-bit sources pass through 8 bits, so allow a quantization step. + assert!(diff(0) < 1.0 * scale, "first frame differs by {}", diff(0)); + for i in 1..m.size()[0] { + assert!(diff(i) < 20.0 * scale, "frame {i} differs by {}", diff(i)); + } + Ok(()) + } + + #[test_case(48.0, 24.0, true ; "exact 2x")] + #[test_case(24000.0 / 1001.0, 24.0, true ; "ntsc 23.976")] + #[test_case(120000.0 / 1001.0, 24.0, true ; "ntsc 119.88")] + #[test_case(24.0, 24.0, true ; "same rate")] + #[test_case(30.0, 24.0, false ; "30 to 24")] + #[test_case(25.0, 24.0, false ; "25 to 24")] + #[test_case(60000.0 / 1001.0, 24.0, false ; "59.94 to 24")] + #[test_case(12.0, 24.0, false ; "upsample")] + fn test_is_integer_downsample(source: f64, target: f64, expected: bool) { + assert_eq!(is_integer_downsample(source, target), expected); + } + + #[test] + fn test_fps_conversion_ntsc_source_matches_nearest() -> anyhow::Result<()> { + init_logger(); + let video = crate::util::test_utils::generate_fractional_fps_video(48)?; + let decode = |fps_conversion| { + decode_media( + MediaDecodeRequest { + source: MediaSource::Uri(video.path().to_str().unwrap().into()), + start_time: None, + end_time: None, + video_stream: Some(VideoStreamRequest { + frame_rate: Some(24.0), + fps_conversion, + ..Default::default() + }), + audio_streams: None, + }, + None, + ) + .map(|mut streams| streams.remove(0).data.unwrap()) + }; + let nearest = decode(FpsConversion::Nearest)?; + assert!(nearest.equal(&decode(FpsConversion::Motion)?)); + assert!(nearest.equal(&decode(FpsConversion::Blend)?)); + Ok(()) + } + + #[test_case(FpsConversion::Blend ; "blend")] + #[test_case(FpsConversion::Motion ; "motion")] + fn test_fps_conversion_seek_window_matches_nearest_timing( + fps_conversion: FpsConversion, + ) -> anyhow::Result<()> { + init_logger(); + let video = generate_moving_bar_video(30, 60, Duration::from_secs(4))?; + let decode = |fps_conversion| { + decode_media( + MediaDecodeRequest { + source: MediaSource::Uri(video.path().to_str().unwrap().into()), + start_time: Some(1.1), + end_time: Some(3.3), + video_stream: Some(VideoStreamRequest { + frame_rate: Some(24.0), + fps_conversion, + ..Default::default() + }), + audio_streams: None, + }, + None, + ) + .map(|mut streams| streams.remove(0)) + }; + let nearest = decode(FpsConversion::Nearest)?; + let converted = decode(fps_conversion)?; + assert_eq!(converted.data.unwrap().size(), nearest.data.unwrap().size()); + // Same timestamps; `framerate` keeps the source time base, so they + // can differ from `fps`'s 1/24 time base in the last bit. + assert_eq!(converted.frame_pts.len(), nearest.frame_pts.len()); + for (a, b) in converted.frame_pts.iter().zip(&nearest.frame_pts) { + assert!((a - b).abs() < 1e-9, "{a} vs {b}"); + } + Ok(()) + } + + #[test] + fn test_fps_conversion_parse() { + assert_eq!( + FpsConversion::try_from(None).unwrap(), + FpsConversion::Nearest + ); + assert_eq!( + FpsConversion::try_from(Some("motion")).unwrap(), + FpsConversion::Motion + ); + assert!(FpsConversion::try_from(Some("minterpolate")).is_err()); + } + #[test] fn test_decode_media_resample_audio() -> anyhow::Result<()> { init_logger(); @@ -3853,10 +4434,41 @@ mod tests { ..Default::default() }, hdr_mode: HdrMode::Tonemap, + fps_conversion: FpsConversion::Nearest, }, "fps=24,scale=width=1920:height=1080,zscale=t=linear:npl=100,format=gbrpf32le,tonemap=hable:desat=0,zscale=p=bt709:t=bt709:m=bt709:range=tv,format=pix_fmts=rgb24"; "HDR with fps and scale options" )] + #[test_case( + VideoFilterConfig { + frame_rate: Some(24.0), + width: Some(640), + height: Some(360), + fps_conversion: FpsConversion::Blend, + ..Default::default() + }, + "scale=width=640:height=360,framerate=fps=24,format=pix_fmts=rgb24"; + "blend fps conversion runs after the resize" + )] + #[test_case( + VideoFilterConfig { + frame_rate: Some(24.0), + width: Some(640), + height: Some(360), + fps_conversion: FpsConversion::Motion, + ..Default::default() + }, + "scale=width=640:height=360,format=pix_fmts=rgb24"; + "motion fps conversion runs outside the graph" + )] + #[test_case( + VideoFilterConfig { + fps_conversion: FpsConversion::Motion, + ..Default::default() + }, + "format=pix_fmts=rgb24"; + "fps conversion without a frame rate adds no filters" + )] fn test_video_filter_display(filter: VideoFilterConfig, expected: &str) { assert_eq!(filter.to_string(), expected); assert_eq!(filter.pix_fmt(), Some(ffi::AV_PIX_FMT_RGB24)); @@ -3866,6 +4478,57 @@ mod tests { ); } + /// The graph feeding motion interpolation keeps the resize and tone + /// mapping, drops the frame-rate conversion, and ends in the working + /// YUV format. + #[test_case( + VideoFilterConfig { + frame_rate: Some(24.0), + width: Some(640), + height: Some(360), + fps_conversion: FpsConversion::Motion, + ..Default::default() + }, + "yuv420p", + "scale=width=640:height=360,format=pix_fmts=yuv420p"; + "resize" + )] + #[test_case( + VideoFilterConfig { + frame_rate: Some(24.0), + source_color: SourceColorInfo { + color_trc: Some("smpte2084".to_string()), + color_primaries: Some("bt2020".to_string()), + ..Default::default() + }, + fps_conversion: FpsConversion::Motion, + ..Default::default() + }, + "yuv422p", + "zscale=t=linear:npl=100,format=gbrpf32le,tonemap=hable:desat=0,zscale=p=bt709:t=bt709:m=bt709:range=tv,format=pix_fmts=yuv422p"; + "after HDR tone mapping" + )] + fn test_motion_input_filter(filter: VideoFilterConfig, working_format: &str, expected: &str) { + assert!(filter.uses_motion_stage()); + let input = filter.motion_input(working_format); + assert!(!input.uses_motion_stage()); + assert_eq!(input.to_string(), expected); + assert!(input.pix_fmt().is_some()); + } + + #[test_case(ffi::AV_PIX_FMT_YUV420P, "yuv420p")] + #[test_case(ffi::AV_PIX_FMT_YUV420P10LE, "yuv420p")] + #[test_case(ffi::AV_PIX_FMT_NV12, "yuv420p")] + #[test_case(ffi::AV_PIX_FMT_YUVJ420P, "yuvj420p")] + #[test_case(ffi::AV_PIX_FMT_YUV422P10LE, "yuv422p")] + #[test_case(ffi::AV_PIX_FMT_YUVJ422P, "yuvj422p")] + #[test_case(ffi::AV_PIX_FMT_YUV444P, "yuv444p")] + #[test_case(ffi::AV_PIX_FMT_RGB24, "yuv444p")] + #[test_case(ffi::AV_PIX_FMT_YUV411P, "yuv420p")] + fn test_motion_working_pixel_format(source: AVPixelFormat, expected: &str) { + assert_eq!(motion::working_pixel_format(source), expected); + } + #[test_case( AudioFilterConfig { sample_rate: Some(44100), diff --git a/src/decoder/motion.rs b/src/decoder/motion.rs new file mode 100644 index 0000000..9c573cf --- /dev/null +++ b/src/decoder/motion.rs @@ -0,0 +1,1080 @@ +//! Motion-compensated frame interpolation for [`super::FpsConversion::Motion`]. +//! +//! The model is the one FFmpeg's `minterpolate` uses by default +//! (`mi_mode=mci:mc_mode=obmc:me_mode=bilat:me=epzs:mb_size=16:search_param=32:scd=fdiff`): +//! +//! - Bilateral motion estimation: for every 16x16 luma block, find the +//! half-motion `m` for which the previous frame at `+m` best matches the +//! next frame at `-m` (SAD over the block's 32x32 overlapped window, plus a +//! penalty for straying from the median of the neighbouring vectors). +//! Candidates are the EPZS predictors (spatial neighbours, the co-located +//! and neighbouring vectors of the previous estimate, its acceleration), +//! refined with a small diamond search. +//! - Overlapped block motion compensation: each output block blends the two +//! frames along its vector, weighted by a 32x32 bilinear window, so +//! neighbouring blocks cross-fade instead of leaving block edges. +//! - Scene cuts (frame-difference detector) are not interpolated; the caller +//! repeats the nearer frame. +//! +//! It is written for throughput rather than ported: planes are copied into +//! border-replicated buffers so block reads never clip, SAD and blending run +//! row-wise over fixed-size arrays (auto-vectorized), chroma is compensated +//! at its own resolution instead of once per luma pixel, and motion is +//! estimated only for source pairs that bracket a synthesized frame. That +//! makes it roughly 15x faster than `minterpolate` with ~15x less memory, at +//! equal or better fidelity (measured on 960x540 clips: interpolation costs +//! ~2.5 ms per synthesized frame plus ~2.5 ms per estimated pair, against +//! `minterpolate`'s ~110 ms and ~23 ms). + +use anyhow::{anyhow, bail, Context}; +use rsmpeg::{ + avutil::{av_rescale_q_rnd, AVFrame, AVRational}, + ffi, +}; + +use crate::util::ffmpeg::RetUpgrade; + +/// Interpolation phases run over `0..ALPHA_ONE` (0 = previous frame). +pub(crate) const ALPHA_ONE: u32 = 256; + +/// Luma block size; motion vectors are per block. +const BLOCK: usize = 16; +/// Overlapped window: the block plus half a block on each side. +const WINDOW: usize = 2 * BLOCK; +const HALF: i32 = (BLOCK / 2) as i32; +/// Largest half-motion component searched (`minterpolate` `search_param`). +const SEARCH: i32 = 32; +/// SAD penalty per pixel of distance from the median predictor. +const PRED_COST: u32 = 64; +/// Border replicated around every padded plane. Reads reach at most a block +/// plus the window margin plus a full motion vector (2 * SEARCH) outside. +const PAD: usize = 128; +/// `minterpolate`'s default `scd_threshold`. +const SCENE_THRESHOLD: f64 = 10.0; + +/// A borrowed 8-bit image plane: `height` rows of `width` bytes, `stride` +/// bytes apart. +#[derive(Clone, Copy)] +pub(crate) struct Plane<'a> { + pub data: &'a [u8], + pub stride: usize, +} + +/// A mutable 8-bit image plane, laid out like [`Plane`]. +pub(crate) struct PlaneMut<'a> { + pub data: &'a mut [u8], + pub stride: usize, +} + +/// One plane copied into a buffer with replicated borders. +struct PaddedPlane { + data: Vec, + stride: usize, + width: usize, + height: usize, +} + +impl PaddedPlane { + fn new(width: usize, height: usize) -> Self { + let stride = (width + 2 * PAD).next_multiple_of(64); + Self { + data: vec![0; stride * (height + 2 * PAD)], + stride, + width, + height, + } + } + + fn fill(&mut self, src: Plane) { + let (w, h, stride) = (self.width, self.height, self.stride); + for y in 0..h { + let row = &src.data[y * src.stride..][..w]; + let dst = &mut self.data[(y + PAD) * stride..][..stride]; + dst[PAD..PAD + w].copy_from_slice(row); + dst[..PAD].fill(row[0]); + dst[PAD + w..].fill(row[w - 1]); + } + let (first, last) = (PAD * stride, (PAD + h - 1) * stride); + for y in 0..PAD { + self.data.copy_within(first..first + stride, y * stride); + self.data + .copy_within(last..last + stride, (PAD + h + y) * stride); + } + } + + /// Offset of pixel `(x, y)`, which may lie up to `PAD` outside the frame. + #[inline(always)] + fn index(&self, x: i32, y: i32) -> usize { + (y + PAD as i32) as usize * self.stride + (x + PAD as i32) as usize + } + + #[inline(always)] + fn row(&self, index: usize) -> &[u8; N] { + self.data[index..][..N].try_into().unwrap() + } +} + +/// A source frame prepared for [`Interpolator`]: its three planes with +/// replicated borders. +pub(crate) struct PaddedFrame { + planes: [PaddedPlane; 3], +} + +impl PaddedFrame { + pub fn new(layout: &Layout) -> Self { + let (cw, ch) = layout.chroma_size(); + Self { + planes: [ + PaddedPlane::new(layout.width, layout.height), + PaddedPlane::new(cw, ch), + PaddedPlane::new(cw, ch), + ], + } + } + + pub fn fill(&mut self, planes: [Plane; 3]) { + for (dst, src) in self.planes.iter_mut().zip(planes) { + dst.fill(src); + } + } +} + +/// Frame geometry: luma size and chroma subsampling (log2, 4:2:0 = (1, 1)). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct Layout { + pub width: usize, + pub height: usize, + pub chroma_shift: (u32, u32), +} + +impl Layout { + pub fn chroma_size(&self) -> (usize, usize) { + let (sx, sy) = self.chroma_shift; + (self.width.div_ceil(1 << sx), self.height.div_ceil(1 << sy)) + } +} + +/// Bilinear OBMC window: `4 * round(p_i * p_j / 16)` with +/// `p = 1, 3, ..., 31, 31, ..., 3, 1`. The four windows overlapping any +/// pixel sum to exactly 256. +fn obmc_window() -> [[u16; WINDOW]; WINDOW] { + let p = |i: usize| { + if i < BLOCK { + 2 * i + 1 + } else { + 4 * BLOCK - 1 - 2 * i + } + }; + let mut w = [[0u16; WINDOW]; WINDOW]; + for (j, row) in w.iter_mut().enumerate() { + for (i, v) in row.iter_mut().enumerate() { + *v = (4 * ((p(i) * p(j) + 8) / 16)) as u16; + } + } + w +} + +#[inline(always)] +fn median3(a: i32, b: i32, c: i32) -> i32 { + a.max(b).min(a.min(b).max(c)) +} + +#[inline(always)] +fn sad_row(a: &[u8; WINDOW], b: &[u8; WINDOW]) -> u32 { + // 32 * 255 fits u16, which keeps the reduction in narrow lanes. + a.iter() + .zip(b) + .map(|(&p, &q)| p.abs_diff(q) as u16) + .sum::() as u32 +} + +/// Motion-compensated interpolation between consecutive frames of one +/// stream. Feed every consecutive source pair to [`Self::scene_change`]; +/// call [`Self::estimate`] for a pair before [`Self::interpolate`]ing +/// frames inside it. +pub(crate) struct Interpolator { + layout: Layout, + blocks_w: usize, + blocks_h: usize, + /// Half-motion per block of the latest estimate and the two before it + /// (temporal EPZS predictors). + mvs: Vec<[i32; 2]>, + mvs_prev: Vec<[i32; 2]>, + mvs_prev2: Vec<[i32; 2]>, + window: [[u16; WINDOW]; WINDOW], + /// Weighted-sum accumulator for one plane, with a block of margin. + acc: Vec, + prev_mafd: f64, +} + +impl Interpolator { + pub fn new(layout: Layout) -> Self { + let blocks_w = layout.width.div_ceil(BLOCK); + let blocks_h = layout.height.div_ceil(BLOCK); + let blocks = blocks_w * blocks_h; + Self { + layout, + blocks_w, + blocks_h, + mvs: vec![[0; 2]; blocks], + mvs_prev: vec![[0; 2]; blocks], + mvs_prev2: vec![[0; 2]; blocks], + window: obmc_window(), + acc: vec![0; (blocks_w + 3) * BLOCK * (blocks_h + 3) * BLOCK], + prev_mafd: 0.0, + } + } + + /// Whether `next` starts a new scene after `prev` (luma planes). Uses the + /// previous pair's difference, so call it for every consecutive pair, in + /// order. + pub fn scene_change(&mut self, prev: Plane, next: Plane) -> bool { + let (w, h) = (self.layout.width, self.layout.height); + let mut sad = 0u64; + for y in 0..h { + let (a, b) = ( + &prev.data[y * prev.stride..][..w], + &next.data[y * next.stride..][..w], + ); + sad += a + .iter() + .zip(b) + .map(|(&p, &q)| p.abs_diff(q) as u32) + .sum::() as u64; + } + let mafd = sad as f64 * 100.0 / (w * h) as f64 / 256.0; + let diff = (mafd - self.prev_mafd).abs(); + self.prev_mafd = mafd; + mafd.min(diff).clamp(0.0, 100.0) >= SCENE_THRESHOLD + } + + #[inline(always)] + fn cost(a: &PaddedPlane, b: &PaddedPlane, x: i32, y: i32, mv: [i32; 2], pred: [i32; 2]) -> u32 { + let ia = a.index(x + mv[0] - HALF, y + mv[1] - HALF); + let ib = b.index(x - mv[0] - HALF, y - mv[1] - HALF); + let mut sad = 0; + for row in 0..WINDOW { + sad += sad_row(a.row(ia + row * a.stride), b.row(ib + row * b.stride)); + } + sad + ((mv[0] - pred[0]).unsigned_abs() + (mv[1] - pred[1]).unsigned_abs()) * PRED_COST + } + + /// Estimates per-block motion between `prev` and `next`. + pub fn estimate(&mut self, prev: &PaddedFrame, next: &PaddedFrame) { + std::mem::swap(&mut self.mvs_prev2, &mut self.mvs_prev); + std::mem::swap(&mut self.mvs_prev, &mut self.mvs); + let (a, b) = (&prev.planes[0], &next.planes[0]); + let (bw, bh) = (self.blocks_w, self.blocks_h); + for by in 0..bh { + for bx in 0..bw { + let i = by * bw + bx; + let (x, y) = ((bx * BLOCK) as i32, (by * BLOCK) as i32); + + // Spatial neighbours already estimated in this frame. + let mut spatial = [[0i32; 2]; 3]; + let mut ns = 0; + if bx > 0 { + spatial[ns] = self.mvs[i - 1]; + ns += 1; + } + if by > 0 { + spatial[ns] = self.mvs[i - bw]; + ns += 1; + if bx + 1 < bw { + spatial[ns] = self.mvs[i - bw + 1]; + ns += 1; + } + } + let pred = match ns { + 3 => [ + median3(spatial[0][0], spatial[1][0], spatial[2][0]), + median3(spatial[0][1], spatial[1][1], spatial[2][1]), + ], + 2 => [ + median3(0, spatial[0][0], spatial[1][0]), + median3(0, spatial[0][1], spatial[1][1]), + ], + 1 => spatial[0], + _ => [0, 0], + }; + + let mut candidates = [[0i32; 2]; 11]; + let mut nc = 0; + let mut add = |mv: [i32; 2]| { + candidates[nc] = mv; + nc += 1; + }; + add(pred); + add([0, 0]); + spatial[..ns].iter().for_each(|&mv| add(mv)); + let col = self.mvs_prev[i]; + add(col); + let col2 = self.mvs_prev2[i]; + add([2 * col[0] - col2[0], 2 * col[1] - col2[1]]); + if bx > 0 { + add(self.mvs_prev[i - 1]); + } + if by > 0 { + add(self.mvs_prev[i - bw]); + } + if bx + 1 < bw { + add(self.mvs_prev[i + 1]); + } + if by + 1 < bh { + add(self.mvs_prev[i + bw]); + } + + let mut visited = [[i32::MAX; 2]; 48]; + let mut nv = 0; + let mut best = ([0i32; 2], u32::MAX); + let mut try_mv = |mv: [i32; 2], best: &mut ([i32; 2], u32)| { + if mv[0].abs() > SEARCH || mv[1].abs() > SEARCH || visited[..nv].contains(&mv) { + return; + } + if nv < visited.len() { + visited[nv] = mv; + nv += 1; + } + let cost = Self::cost(a, b, x, y, mv, pred); + if cost < best.1 { + *best = (mv, cost); + } + }; + for &mv in &candidates[..nc] { + try_mv(mv, &mut best); + } + loop { + let center = best.0; + for (dx, dy) in [(-1, 0), (0, -1), (1, 0), (0, 1)] { + try_mv([center[0] + dx, center[1] + dy], &mut best); + } + if best.0 == center { + break; + } + } + self.mvs[i] = best.0; + } + } + } + + /// Synthesizes the frame at phase `alpha` (`1..ALPHA_ONE`, towards + /// `next`) of the last [`Self::estimate`]d pair into `out`. + pub fn interpolate( + &mut self, + prev: &PaddedFrame, + next: &PaddedFrame, + alpha: u32, + out: [PlaneMut; 3], + ) { + debug_assert!(alpha > 0 && alpha < ALPHA_ONE); + let [y, u, v] = out; + self.compensate::(&prev.planes[0], &next.planes[0], alpha, (0, 0), y); + let shift = self.layout.chroma_shift; + for (plane, out) in [(1, u), (2, v)] { + let (a, b) = (&prev.planes[plane], &next.planes[plane]); + match shift { + (1, 1) => self.compensate::(a, b, alpha, shift, out), + (1, 0) => self.compensate::(a, b, alpha, shift, out), + (0, 0) => self.compensate::(a, b, alpha, shift, out), + _ => unreachable!("unsupported chroma subsampling {shift:?}"), + } + } + } + + /// OBMC for one plane whose blocks are `W` x `H` windows (luma blocks + /// scaled by `shift`). + fn compensate( + &mut self, + a: &PaddedPlane, + b: &PaddedPlane, + alpha: u32, + (sx, sy): (u32, u32), + out: PlaneMut, + ) { + let (bsx, bsy) = ((BLOCK >> sx) as i32, (BLOCK >> sy) as i32); + // The accumulator starts one block plus the window margin before the + // frame, so the windows of the border blocks (-1 and blocks_w) fit. + let (ox, oy) = (bsx + bsx / 2, bsy + bsy / 2); + let acc_w = (self.blocks_w + 3) * (BLOCK >> sx); + let acc_h = (self.blocks_h + 3) * (BLOCK >> sy); + let acc = &mut self.acc[..acc_w * acc_h]; + acc.fill(0); + + // The luma window averaged over each chroma sample's luma pixels. + // Luma weights are multiples of 4, so the averages stay exact and + // the four overlapping windows still sum to 256. + let mut window = [[0u16; W]; H]; + for (j, row) in window.iter_mut().enumerate() { + for (i, w) in row.iter_mut().enumerate() { + let mut sum = 0; + for dy in 0..1 << sy { + for dx in 0..1 << sx { + sum += self.window[(j << sy) + dy][(i << sx) + dx]; + } + } + *w = sum >> (sx + sy); + } + } + + let (ka, kb) = ((ALPHA_ONE - alpha) as u16, alpha as u16); + let (bw, bh) = (self.blocks_w as i32, self.blocks_h as i32); + for by in -1..=bh { + for bx in -1..=bw { + let mv = self.mvs[(by.clamp(0, bh - 1) * bw + bx.clamp(0, bw - 1)) as usize]; + // Full motion prev -> next is twice the half-motion; split it + // by phase (truncating towards zero, like minterpolate). + let full = [2 * mv[0], 2 * mv[1]]; + let to_prev = |c: i32| c * alpha as i32 / ALPHA_ONE as i32; + let to_next = |c: i32| -c * (ALPHA_ONE - alpha) as i32 / ALPHA_ONE as i32; + let (pax, pay) = (to_prev(full[0]) / (1 << sx), to_prev(full[1]) / (1 << sy)); + let (pbx, pby) = (to_next(full[0]) / (1 << sx), to_next(full[1]) / (1 << sy)); + let (x0, y0) = (bx * bsx - bsx / 2, by * bsy - bsy / 2); + let ia = a.index(x0 + pax, y0 + pay); + let ib = b.index(x0 + pbx, y0 + pby); + let iacc = (y0 + oy) as usize * acc_w + (x0 + ox) as usize; + for (j, wrow) in window.iter().enumerate() { + let ra: &[u8; W] = a.row(ia + j * a.stride); + let rb: &[u8; W] = b.row(ib + j * b.stride); + let racc: &mut [u16; W] = + (&mut acc[iacc + j * acc_w..][..W]).try_into().unwrap(); + for i in 0..W { + let p = (ka * ra[i] as u16 + kb * rb[i] as u16 + 128) >> 8; + // Window weights at a pixel sum to 256, so this + // stays below 256 * 256. + racc[i] += wrow[i] * p; + } + } + } + } + + let (pw, ph) = (a.width, a.height); + for y in 0..ph { + let src = &acc[(y + oy as usize) * acc_w + ox as usize..][..pw]; + let dst = &mut out.data[y * out.stride..][..pw]; + for (d, &s) in dst.iter_mut().zip(src) { + *d = ((s as u32 + 128) >> 8) as u8; + } + } + } +} + +/// The 8-bit planar YUV format motion interpolation works in for a source +/// in `pix_fmt`: the source's own chroma subsampling (4:2:0, 4:2:2 or +/// 4:4:4; full-range `yuvj` kept as is), 4:4:4 for RGB, 4:2:0 otherwise. +pub(crate) fn working_pixel_format(pix_fmt: ffi::AVPixelFormat) -> &'static str { + let desc = unsafe { ffi::av_pix_fmt_desc_get(pix_fmt) }; + if desc.is_null() { + return "yuv420p"; + } + let desc = unsafe { &*desc }; + if desc.flags & ffi::AV_PIX_FMT_FLAG_RGB as u64 != 0 { + return "yuv444p"; + } + let full_range = matches!( + pix_fmt, + ffi::AV_PIX_FMT_YUVJ420P | ffi::AV_PIX_FMT_YUVJ422P | ffi::AV_PIX_FMT_YUVJ444P + ); + match ((desc.log2_chroma_w, desc.log2_chroma_h), full_range) { + ((1, 0), false) => "yuv422p", + ((1, 0), true) => "yuvj422p", + ((0, 0), false) => "yuv444p", + ((0, 0), true) => "yuvj444p", + (_, true) => "yuvj420p", + _ => "yuv420p", + } +} + +impl Layout { + /// Geometry of `width` x `height` frames in planar 8-bit YUV `pix_fmt`. + pub fn new( + width: i32, + height: i32, + pix_fmt: ffi::AVPixelFormat, + ) -> Result { + let desc = unsafe { ffi::av_pix_fmt_desc_get(pix_fmt) }; + let supported = !desc.is_null() && { + let desc = unsafe { &*desc }; + desc.nb_components == 3 + && desc.flags & (ffi::AV_PIX_FMT_FLAG_PLANAR | ffi::AV_PIX_FMT_FLAG_RGB) as u64 + == ffi::AV_PIX_FMT_FLAG_PLANAR as u64 + && desc.comp.iter().take(3).all(|c| c.depth == 8) + && matches!( + (desc.log2_chroma_w, desc.log2_chroma_h), + (1, 1) | (1, 0) | (0, 0) + ) + }; + if !supported || width <= 0 || height <= 0 { + bail!("Motion interpolation needs 8-bit planar YUV 4:2:0/4:2:2/4:4:4 frames, got pix_fmt {pix_fmt} at {width}x{height}"); + } + let desc = unsafe { &*desc }; + Ok(Self { + width: width as usize, + height: height as usize, + chroma_shift: (desc.log2_chroma_w as u32, desc.log2_chroma_h as u32), + }) + } +} + +fn frame_layout(frame: &AVFrame) -> Result { + Layout::new(frame.width, frame.height, frame.format) +} + +/// (rows, row bytes) of plane `index` of a frame with `layout`. +fn plane_size(layout: &Layout, index: usize) -> (usize, usize) { + if index == 0 { + (layout.height, layout.width) + } else { + let (cw, ch) = layout.chroma_size(); + (ch, cw) + } +} + +/// Borrows the three planes of a planar 8-bit YUV `frame`. +fn frame_planes<'f>(frame: &'f AVFrame, layout: &Layout) -> Result<[Plane<'f>; 3], anyhow::Error> { + let plane = |i: usize| { + let (rows, width) = plane_size(layout, i); + let stride = frame.linesize[i]; + if frame.data[i].is_null() || stride < width as i32 { + bail!("Unexpected plane {i} layout for motion interpolation: {frame:?}"); + } + let stride = stride as usize; + // SAFETY: an allocated frame holds `rows` rows of `stride` bytes + // per plane; the last row only needs `width` of them. + let data = + unsafe { std::slice::from_raw_parts(frame.data[i], stride * (rows - 1) + width) }; + Ok(Plane { data, stride }) + }; + Ok([plane(0)?, plane(1)?, plane(2)?]) +} + +fn frame_planes_mut<'f>(frame: &'f mut AVFrame, layout: &Layout) -> [PlaneMut<'f>; 3] { + let (data, linesize) = (frame.data, frame.linesize); + [0, 1, 2].map(|i| { + let (rows, width) = plane_size(layout, i); + let stride = linesize[i] as usize; + // SAFETY: the frame was just allocated by `av_frame_get_buffer` with + // this layout, so every plane holds `rows` rows of `stride >= width` + // bytes and nothing else references its buffers. + let data = unsafe { std::slice::from_raw_parts_mut(data[i], stride * (rows - 1) + width) }; + PlaneMut { data, stride } + }) +} + +/// A source frame with its timestamp on the resampler's exact time grid. +struct Source { + frame: AVFrame, + /// Presentation time in units of `1 / (in_tb.den * out_tb.den)` seconds. + time: i128, +} + +/// Resamples a frame stream to a constant rate, synthesizing in-between +/// frames with [`Interpolator`]. +/// +/// Output timing mirrors FFmpeg's `fps` filter (default options), so the +/// frame count and timestamps equal those of [`super::FpsConversion::Nearest`]: +/// the first output timestamp is the first input timestamp rounded to the +/// output grid, and output stops at the end of the last input frame +/// (`pts + duration`, rounded). Output frames that land exactly on a source +/// frame pass it through untouched, as do frames across a scene cut (the +/// nearer side) and past the last source frame. +pub(crate) struct Resampler { + interp: Interpolator, + layout: Layout, + in_tb: AVRational, + out_tb: AVRational, + /// Output frames that may lie past the end of the stream: at or after the + /// latest source frame's rounded timestamp. Released once a later source + /// frame (or the stream end) shows they precede the end. + pending: Vec, + prev: Option, + /// Pts (output time base) of the next output frame. + next_out: Option, + /// End (pts + duration, input time base) of the last source frame. + end: i64, + /// Whether the current pair straddles a scene cut. + scene_cut: bool, + /// Prepared copies of source frames, keyed by source pts. + padded: [(PaddedFrame, Option); 2], + /// Source pts of the pair the interpolator's motion field belongs to. + estimated: Option<(i64, i64)>, +} + +impl Resampler { + pub fn new(layout: Layout, in_tb: AVRational, out_tb: AVRational) -> Self { + Self { + interp: Interpolator::new(layout), + layout, + in_tb, + out_tb, + pending: Vec::new(), + prev: None, + next_out: None, + end: 0, + scene_cut: false, + padded: [ + (PaddedFrame::new(&layout), None), + (PaddedFrame::new(&layout), None), + ], + estimated: None, + } + } + + fn in_time(&self, pts: i64) -> i128 { + pts as i128 * self.in_tb.num as i128 * self.out_tb.den as i128 + } + + fn out_time(&self, pts: i64) -> i128 { + pts as i128 * self.out_tb.num as i128 * self.in_tb.den as i128 + } + + /// Feeds the next source frame; returns the output frames it completes. + pub fn push(&mut self, frame: AVFrame) -> Result, anyhow::Error> { + if frame.pts == ffi::AV_NOPTS_VALUE { + log::debug!("Dropping a frame without timestamp before motion interpolation"); + return Ok(vec![]); + } + if frame_layout(&frame)? != self.layout { + bail!("Frame size or format changed mid-stream, which motion interpolation does not support"); + } + let source = Source { + time: self.in_time(frame.pts), + frame, + }; + self.end = source.frame.pts + source.frame.duration; + let next_out = *self.next_out.get_or_insert_with(|| { + av_rescale_q_rnd(source.frame.pts, self.in_tb, self.out_tb, ROUNDING) + }); + let Some(prev) = self.prev.take() else { + self.prev = Some(source); + return Ok(vec![]); + }; + if source.time <= prev.time { + log::debug!( + "Non-increasing pts {} before motion interpolation, dropping the earlier frame", + source.frame.pts + ); + // Pts may now repeat, so cached copies keyed by pts are stale. + self.estimated = None; + self.padded.iter_mut().for_each(|(_, pts)| *pts = None); + self.prev = Some(source); + return Ok(vec![]); + } + let layout = self.layout; + self.scene_cut = self.interp.scene_change( + frame_planes(&prev.frame, &layout)?[0], + frame_planes(&source.frame, &layout)?[0], + ); + let mut k = next_out; + while self.out_time(k) < source.time { + let offset = (self.out_time(k) - prev.time).max(0); + let alpha = (offset * ALPHA_ONE as i128 / (source.time - prev.time)) as u32; + let frame = self.output(&prev, &source, alpha, k)?; + self.pending.push(frame); + k += 1; + } + self.next_out = Some(k); + // The stream ends no earlier than this frame's rounded timestamp, so + // every output frame before it is final (as in `fps`). + let bound = av_rescale_q_rnd(source.frame.pts, self.in_tb, self.out_tb, ROUNDING); + self.prev = Some(source); + Ok(self.release(bound)) + } + + /// Ends the stream: returns the remaining output frames up to the end of + /// the last source frame (repeats of it). + pub fn finish(&mut self) -> Vec { + let end = av_rescale_q_rnd(self.end, self.in_tb, self.out_tb, ROUNDING); + let mut out = self.release(end); + self.pending.clear(); + let (Some(prev), Some(mut k)) = (self.prev.take(), self.next_out) else { + return out; + }; + while k < end { + out.push(self.retimed(&prev.frame, k)); + k += 1; + } + self.next_out = Some(k); + out + } + + /// Removes and returns the pending frames before output pts `bound`. + fn release(&mut self, bound: i64) -> Vec { + let keep = self.pending.iter().take_while(|f| f.pts < bound).count(); + self.pending.drain(..keep).collect() + } + + /// A reference to `frame` re-stamped as output frame `pts`. + fn retimed(&self, frame: &AVFrame, pts: i64) -> AVFrame { + let mut out = frame.clone(); + out.set_pts(pts); + out.set_time_base(self.out_tb); + unsafe { (*out.as_mut_ptr()).duration = 1 }; + out + } + + fn output( + &mut self, + prev: &Source, + next: &Source, + alpha: u32, + pts: i64, + ) -> Result { + if alpha == 0 { + return Ok(self.retimed(&prev.frame, pts)); + } + if self.scene_cut { + let nearer = if alpha > ALPHA_ONE / 2 { next } else { prev }; + return Ok(self.retimed(&nearer.frame, pts)); + } + self.estimate(prev, next)?; + + let mut out = AVFrame::new(); + out.set_format(prev.frame.format); + out.set_width(prev.frame.width); + out.set_height(prev.frame.height); + out.alloc_buffer() + .context("Allocating an interpolated frame")?; + unsafe { ffi::av_frame_copy_props(out.as_mut_ptr(), prev.frame.as_ptr()) } + .upgrade() + .map_err(|e| anyhow!("Copying frame properties: {e}"))?; + let layout = self.layout; + let [(a, _), (b, _)] = &self.padded; + self.interp + .interpolate(a, b, alpha, frame_planes_mut(&mut out, &layout)); + out.set_pts(pts); + out.set_time_base(self.out_tb); + unsafe { (*out.as_mut_ptr()).duration = 1 }; + Ok(out) + } + + /// Estimates motion for the pair (prev, next) unless already done, + /// reusing a prepared copy of `prev` from the previous pair. + fn estimate(&mut self, prev: &Source, next: &Source) -> Result<(), anyhow::Error> { + let (prev_pts, next_pts) = (prev.frame.pts, next.frame.pts); + if self.estimated == Some((prev_pts, next_pts)) { + return Ok(()); + } + let layout = self.layout; + if self.padded[1].1 == Some(prev_pts) { + self.padded.swap(0, 1); + } else { + self.padded[0].0.fill(frame_planes(&prev.frame, &layout)?); + self.padded[0].1 = Some(prev_pts); + } + self.padded[1].0.fill(frame_planes(&next.frame, &layout)?); + self.padded[1].1 = Some(next_pts); + let [(a, _), (b, _)] = &self.padded; + self.interp.estimate(a, b); + self.estimated = Some((prev_pts, next_pts)); + Ok(()) + } +} + +/// `fps`'s default timestamp rounding. +const ROUNDING: u32 = ffi::AV_ROUND_NEAR_INF | ffi::AV_ROUND_PASS_MINMAX; + +#[cfg(test)] +mod tests { + use super::*; + use test_case::test_case; + + /// A smooth, aperiodic test pattern sampled at `(x, y)`. + fn pattern(x: f64, y: f64) -> u8 { + let v = 128.0 + + 50.0 * (x * 0.11 + y * 0.03).sin() + + 40.0 * (y * 0.13 - x * 0.05).cos() + + 25.0 * ((x * 0.07) * (y * 0.021 + 1.0)).sin(); + v.clamp(0.0, 255.0) as u8 + } + + /// Planes of `layout` filled from `f(x, y, plane)`. + fn planes(layout: &Layout, f: impl Fn(f64, f64, usize) -> u8) -> [Vec; 3] { + let (cw, ch) = layout.chroma_size(); + let (sx, sy) = layout.chroma_shift; + let luma = (0..layout.height) + .flat_map(|y| (0..layout.width).map(move |x| (x, y))) + .map(|(x, y)| f(x as f64, y as f64, 0)) + .collect(); + let chroma = |p| { + (0..ch) + .flat_map(|y| (0..cw).map(move |x| (x, y))) + .map(|(x, y)| f((x << sx) as f64, (y << sy) as f64, p)) + .collect() + }; + [luma, chroma(1), chroma(2)] + } + + fn padded(layout: &Layout, planes: &[Vec; 3]) -> PaddedFrame { + let (cw, _) = layout.chroma_size(); + let mut frame = PaddedFrame::new(layout); + frame.fill([ + Plane { + data: &planes[0], + stride: layout.width, + }, + Plane { + data: &planes[1], + stride: cw, + }, + Plane { + data: &planes[2], + stride: cw, + }, + ]); + frame + } + + fn run(layout: Layout, prev: &[Vec; 3], next: &[Vec; 3], alpha: u32) -> [Vec; 3] { + let (cw, ch) = layout.chroma_size(); + let (pa, pb) = (padded(&layout, prev), padded(&layout, next)); + let mut interp = Interpolator::new(layout); + interp.estimate(&pa, &pb); + let mut out = [ + vec![0u8; layout.width * layout.height], + vec![0u8; cw * ch], + vec![0u8; cw * ch], + ]; + let [y, u, v] = &mut out; + interp.interpolate( + &pa, + &pb, + alpha, + [ + PlaneMut { + data: y, + stride: layout.width, + }, + PlaneMut { + data: u, + stride: cw, + }, + PlaneMut { + data: v, + stride: cw, + }, + ], + ); + out + } + + /// A yuv420p frame of `layout` filled from `f`, stamped `pts` / `duration`. + fn av_frame( + layout: &Layout, + pts: i64, + duration: i64, + f: impl Fn(f64, f64, usize) -> u8, + ) -> AVFrame { + let data = planes(layout, f); + let mut frame = AVFrame::new(); + frame.set_format(ffi::AV_PIX_FMT_YUV420P); + frame.set_width(layout.width as i32); + frame.set_height(layout.height as i32); + frame.alloc_buffer().unwrap(); + let (cw, _) = layout.chroma_size(); + for (p, src) in frame_planes_mut(&mut frame, layout).into_iter().zip(&data) { + let w = if src.len() == layout.width * layout.height { + layout.width + } else { + cw + }; + for (row, chunk) in src.chunks(w).enumerate() { + p.data[row * p.stride..][..w].copy_from_slice(chunk); + } + } + frame.set_pts(pts); + unsafe { (*frame.as_mut_ptr()).duration = duration }; + frame + } + + /// Output pts the `fps` filter would produce for source frames at `pts` + /// (time base `in_tb`) whose last frame ends at `end`. + fn fps_filter_pts(pts: &[i64], end: i64, in_tb: AVRational, out_tb: AVRational) -> Vec { + let r = |p| av_rescale_q_rnd(p, in_tb, out_tb, ROUNDING); + let src: Vec = pts.iter().map(|&p| r(p)).collect(); + (src[0]..r(end)).collect() + } + + /// Resampler output pts match `fps`'s, including streams whose last frame + /// starts just after an output timestamp but ends before the next one + /// rounds in (59.94 -> 24 fps: last frame at 6.006 output frames; or + /// zero-duration frames), where `fps` emits no frame for that timestamp. + #[test_case(1001, 60000, 1001, 13 ; "59.94 fps")] + #[test_case(1001, 60000, 0, 13 ; "59.94 fps zero duration")] + #[test_case(1, 30, 0, 16 ; "30 fps zero duration")] + #[test_case(1, 25, 1, 13 ; "25 fps")] + #[test_case(512, 12800, 512, 26 ; "25 fps mp4 time base")] + fn test_resampler_matches_fps_timing(step: i64, den: i32, duration: i64, n: i64) { + let layout = Layout { + width: 48, + height: 32, + chroma_shift: (1, 1), + }; + let in_tb = AVRational { num: 1, den }; + let out_tb = AVRational { num: 1, den: 24 }; + let mut resampler = Resampler::new(layout, in_tb, out_tb); + let pts: Vec = (0..n).map(|i| 3 * step + i * step).collect(); + let mut out = Vec::new(); + for &p in &pts { + let frame = av_frame(&layout, p, duration, |x, y, _| { + pattern(x + p as f64 / step as f64, y) + }); + out.extend(resampler.push(frame).unwrap().iter().map(|f| f.pts)); + } + out.extend(resampler.finish().iter().map(|f| f.pts)); + let end = pts[pts.len() - 1] + duration; + assert_eq!(out, fps_filter_pts(&pts, end, in_tb, out_tb)); + } + + /// After a timestamp goes backwards, a pts seen before is not mistaken + /// for the cached pair: frames between A' and E interpolate A' and E. + #[test] + fn test_resampler_restarts_after_pts_regression() { + let layout = Layout { + width: 48, + height: 32, + chroma_shift: (1, 1), + }; + let tb = AVRational { num: 1, den: 48 }; + let mut resampler = Resampler::new(layout, tb, AVRational { num: 1, den: 24 }); + let dark = |_: f64, _: f64, _: usize| 40; + let bright = |_: f64, _: f64, _: usize| 200; + // A(0) and B(4): the frame between them (pts 1/24) is interpolated. + assert!(resampler + .push(av_frame(&layout, 0, 4, dark)) + .unwrap() + .is_empty()); + assert_eq!( + resampler.push(av_frame(&layout, 4, 4, dark)).unwrap().len(), + 2 + ); + // A'(0) again (bright), then E(10) (bright): the frames between them + // (2/24 .. 4/24) must come from A' and E, not the cached A/B pair. + assert!(resampler + .push(av_frame(&layout, 0, 4, bright)) + .unwrap() + .is_empty()); + let out = resampler.push(av_frame(&layout, 10, 4, bright)).unwrap(); + assert_eq!(out.iter().map(|f| f.pts).collect::>(), [2, 3, 4]); + for f in out { + assert_eq!( + unsafe { *f.data[0] }, + 200, + "stale motion cache reused for pts {}", + f.pts + ); + } + } + + #[test] + fn test_obmc_windows_sum_to_256() { + let w = obmc_window(); + for y in 0..BLOCK { + for x in 0..BLOCK { + let sum = w[y][x] + w[y][x + BLOCK] + w[y + BLOCK][x] + w[y + BLOCK][x + BLOCK]; + assert_eq!(sum, 256); + } + } + } + + #[test_case((1, 1) ; "yuv420")] + #[test_case((1, 0) ; "yuv422")] + #[test_case((0, 0) ; "yuv444")] + fn test_static_content_is_reproduced(chroma_shift: (u32, u32)) { + let layout = Layout { + width: 100, + height: 60, + chroma_shift, + }; + let frame = planes(&layout, |x, y, p| pattern(x + 7.0 * p as f64, y)); + for alpha in [1, 64, 128, 255] { + assert_eq!(run(layout, &frame, &frame, alpha), frame); + } + } + + /// A pattern translating by (12, -6) px between frames is synthesized + /// at the right position: far closer to the true in-between frame than + /// a cross-fade. + #[test_case(64 ; "quarter")] + #[test_case(128 ; "half")] + #[test_case(192 ; "three quarters")] + fn test_translation_is_interpolated(alpha: u32) { + let layout = Layout { + width: 160, + height: 96, + chroma_shift: (1, 1), + }; + let (dx, dy) = (12.0, -6.0); + let at = |t: f64| { + planes(&layout, move |x, y, p| { + pattern(x - t * dx + 9.0 * p as f64, y - t * dy) + }) + }; + let (prev, next) = (at(0.0), at(1.0)); + let t = alpha as f64 / ALPHA_ONE as f64; + let truth = at(t); + let out = run(layout, &prev, &next, alpha); + + // Mean absolute luma error away from the borders, where content + // enters or leaves the frame. + let interior_mae = |img: &[u8]| { + let mut sum = 0.0; + let mut n = 0.0; + for y in 16..layout.height - 16 { + for x in 24..layout.width - 24 { + let i = y * layout.width + x; + sum += (img[i] as f64 - truth[0][i] as f64).abs(); + n += 1.0; + } + } + sum / n + }; + let blend: Vec = prev[0] + .iter() + .zip(&next[0]) + .map(|(&a, &b)| ((a as f64 * (1.0 - t) + b as f64 * t).round()) as u8) + .collect(); + let (mae, blend_mae) = (interior_mae(&out[0]), interior_mae(&blend)); + assert!( + mae < 1.5 && mae * 5.0 < blend_mae, + "mae {mae:.2} vs blend {blend_mae:.2}" + ); + } + + #[test_case(1, 1, (1, 1) ; "single pixel")] + #[test_case(37, 23, (1, 1) ; "odd yuv420")] + #[test_case(17, 40, (1, 0) ; "odd yuv422")] + #[test_case(33, 5, (0, 0) ; "odd yuv444")] + fn test_any_frame_size(width: usize, height: usize, chroma_shift: (u32, u32)) { + let layout = Layout { + width, + height, + chroma_shift, + }; + let prev = planes(&layout, |x, y, _| pattern(x, y)); + let next = planes(&layout, |x, y, _| pattern(x + 3.0, y + 1.0)); + let out = run(layout, &prev, &next, 100); + let (cw, ch) = layout.chroma_size(); + assert_eq!(out.map(|p| p.len()), [width * height, cw * ch, cw * ch]); + } + + #[test] + fn test_scene_change() { + let layout = Layout { + width: 64, + height: 32, + chroma_shift: (1, 1), + }; + let a = planes(&layout, |x, y, _| pattern(x, y)); + let b = planes(&layout, |x, y, _| pattern(x + 1.0, y)); + let cut = planes(&layout, |x, y, _| 255 - pattern(y * 3.0, x * 2.0)); + fn plane(p: &[Vec; 3]) -> Plane<'_> { + Plane { + data: &p[0], + stride: 64, + } + } + let mut interp = Interpolator::new(layout); + assert!(!interp.scene_change(plane(&a), plane(&b))); + assert!(interp.scene_change(plane(&b), plane(&cut))); + } +} diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs index cdc47ab..8d5c225 100644 --- a/src/ffi/mod.rs +++ b/src/ffi/mod.rs @@ -98,6 +98,15 @@ pub struct VideoStreamRequest { /// oriented and substantially alters both luminance and chroma. #[pyo3(get, set)] hdr_mode: Option, + /// How `fps` resampling builds each output frame: "nearest" (the + /// default) keeps the nearest source frame, dropping or repeating + /// whole frames; "blend" cross-fades the two neighbouring frames; + /// "motion" interpolates along estimated block motion (the model of + /// FFmpeg's `minterpolate`). Nearest gives uneven frame spacing for + /// non-integer ratios (e.g. 30 -> 24 fps); blend and motion keep it + /// even. Motion costs about 1.2-1.6x nearest's decode time. + #[pyo3(get, set)] + fps_conversion: Option, } impl VideoStreamRequest { @@ -117,7 +126,7 @@ impl VideoStreamRequest { #[pymethods] impl VideoStreamRequest { #[new] - #[pyo3(signature = (*, index=None, width=None, height=None, fps=None, number_of_threads=None, hardware_acceleration=None, dimension_order=None, device=None, dtype=None, hdr_mode=None))] + #[pyo3(signature = (*, index=None, width=None, height=None, fps=None, number_of_threads=None, hardware_acceleration=None, dimension_order=None, device=None, dtype=None, hdr_mode=None, fps_conversion=None))] #[allow(clippy::too_many_arguments)] pub fn py_new( index: Option, @@ -130,6 +139,7 @@ impl VideoStreamRequest { device: Option, dtype: Option, hdr_mode: Option, + fps_conversion: Option, ) -> Self { VideoStreamRequest { index, @@ -142,6 +152,7 @@ impl VideoStreamRequest { device, dtype, hdr_mode, + fps_conversion, } } } @@ -183,6 +194,7 @@ impl VideoStreamRequest { device: self.device_ordinal()?, dtype: self.dtype_parsed()?, hdr_mode: decoder::HdrMode::try_from(self.hdr_mode.as_deref())?, + fps_conversion: decoder::FpsConversion::try_from(self.fps_conversion.as_deref())?, }) } } diff --git a/src/util/test_utils.rs b/src/util/test_utils.rs index 720b7c2..18b6157 100644 --- a/src/util/test_utils.rs +++ b/src/util/test_utils.rs @@ -266,6 +266,36 @@ pub fn generate_fractional_fps_video(num_frames: usize) -> Result Result { + let file = NamedTempFile::with_suffix(".mp4")?; + let secs = duration.as_secs(); + run_ffmpeg(&[ + "-f", + "lavfi", + "-i", + &format!( + "color=black:s=320x240:r={frame_rate}:d={secs}[bg];\ + color=white:s=16x240:r={frame_rate}:d={secs}[bar];\ + [bg][bar]overlay=x='8+t*{speed}':y=0:eval=frame" + ), + "-c:v", + "libx264", + "-crf", + "10", + "-pix_fmt", + "yuv420p", + file.path().to_str().context("temp path")?, + ])?; + Ok(file) +} + /// Re-encodes `input` as a variable-frame-rate video: the first 75 frames /// keep 30 fps timing, the rest switch to 15 fps timing, producing /// non-uniform PTS gaps. From 8b473ecfb9fb2090d7aa2d641ccff43242a7093e Mon Sep 17 00:00:00 2001 From: Rik Heijdens Date: Fri, 23 Oct 2026 08:06:34 +0000 Subject: [PATCH 2/2] Validate frames locally before borrowing their planes as slices `frame_planes` and `frame_planes_mut` turn FFmpeg plane pointers into Rust slices. Their bounds relied on callers having checked the frame's layout (`Resampler::push`) or on the frame being freshly allocated. Both now check for themselves, in `plane_strides`: the frame's size and pixel format must match the layout (8-bit, fully planar, 3-component YUV), every plane must be non-null with a non-negative stride of at least its width, and `frame_planes_mut` additionally requires the frame to be writable so the mutable slices cannot alias another frame's buffers. The SAFETY comments state the invariants each slice relies on. Adds tests that frames whose geometry, format or writability do not match are rejected, and motion decodes of odd-sized, 4:2:2 and sub-block frames. --- src/decoder/mod.rs | 3 ++ src/decoder/motion.rs | 121 ++++++++++++++++++++++++++++++++++-------- 2 files changed, 103 insertions(+), 21 deletions(-) diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs index fe5b693..aaee490 100644 --- a/src/decoder/mod.rs +++ b/src/decoder/mod.rs @@ -3718,6 +3718,9 @@ mod tests { #[test_case("yuv444p", Some((150, 98)), OutputDtype::Uint8 ; "yuv444p resized odd width")] #[test_case("yuvj420p", Some((96, 64)), OutputDtype::Float32 ; "yuvj420p float32")] #[test_case("yuv420p10le", Some((96, 64)), OutputDtype::Uint8 ; "hdr tone mapped")] + #[test_case("yuv420p", Some((101, 67)), OutputDtype::Uint8 ; "yuv420p resized odd width and height")] + #[test_case("yuv422p", Some((77, 33)), OutputDtype::Uint8 ; "yuv422p resized odd width and height")] + #[test_case("yuv420p", Some((17, 15)), OutputDtype::Uint8 ; "frames smaller than one block")] fn test_fps_conversion_motion_formats( pixel_format: &str, size: Option<(u32, u32)>, diff --git a/src/decoder/motion.rs b/src/decoder/motion.rs index 9c573cf..715afab 100644 --- a/src/decoder/motion.rs +++ b/src/decoder/motion.rs @@ -529,35 +529,71 @@ fn plane_size(layout: &Layout, index: usize) -> (usize, usize) { } } +/// Checks that `frame` has exactly `layout`'s geometry and three non-null +/// planes with `stride >= width`, and returns the plane strides. +/// +/// [`Layout::new`] only accepts 8-bit, fully planar YUV formats with 3 +/// components, so for a frame that passes this check every plane `i` is a +/// separate array of [`plane_size`]`(layout, i)` rows, `stride` bytes apart, +/// each row starting with `width` one-byte samples. FFmpeg allocates video +/// frame buffers (`av_frame_get_buffer`, decoder and filter-graph pools) for +/// at least the frame's full height, so such a plane spans at least +/// `stride * (rows - 1) + width` bytes from `data[i]`. +fn plane_strides(frame: &AVFrame, layout: &Layout) -> Result<[usize; 3], anyhow::Error> { + if frame_layout(frame)? != *layout { + bail!("Frame geometry does not match the motion interpolation layout: {frame:?}"); + } + let mut strides = [0; 3]; + for (i, stride) in strides.iter_mut().enumerate() { + let (_, width) = plane_size(layout, i); + let linesize = frame.linesize[i]; + if frame.data[i].is_null() || linesize < 0 || (linesize as usize) < width { + bail!("Unexpected plane {i} layout for motion interpolation: {frame:?}"); + } + *stride = linesize as usize; + } + Ok(strides) +} + /// Borrows the three planes of a planar 8-bit YUV `frame`. fn frame_planes<'f>(frame: &'f AVFrame, layout: &Layout) -> Result<[Plane<'f>; 3], anyhow::Error> { - let plane = |i: usize| { + let strides = plane_strides(frame, layout)?; + Ok([0, 1, 2].map(|i| { let (rows, width) = plane_size(layout, i); - let stride = frame.linesize[i]; - if frame.data[i].is_null() || stride < width as i32 { - bail!("Unexpected plane {i} layout for motion interpolation: {frame:?}"); - } - let stride = stride as usize; - // SAFETY: an allocated frame holds `rows` rows of `stride` bytes - // per plane; the last row only needs `width` of them. + let stride = strides[i]; + // SAFETY: `plane_strides` checked that `frame` has `layout`'s + // geometry and a non-null plane `i` with `stride >= width`, so the + // plane spans at least `stride * (rows - 1) + width` initialized + // bytes (see `plane_strides`). The slice borrows `frame`, which + // holds a reference to the buffer for its lifetime. let data = unsafe { std::slice::from_raw_parts(frame.data[i], stride * (rows - 1) + width) }; - Ok(Plane { data, stride }) - }; - Ok([plane(0)?, plane(1)?, plane(2)?]) + Plane { data, stride } + })) } -fn frame_planes_mut<'f>(frame: &'f mut AVFrame, layout: &Layout) -> [PlaneMut<'f>; 3] { - let (data, linesize) = (frame.data, frame.linesize); - [0, 1, 2].map(|i| { +/// Mutably borrows the three planes of a planar 8-bit YUV `frame`. Fails +/// unless the frame's buffers are writable (not shared with another frame). +fn frame_planes_mut<'f>( + frame: &'f mut AVFrame, + layout: &Layout, +) -> Result<[PlaneMut<'f>; 3], anyhow::Error> { + let strides = plane_strides(frame, layout)?; + if !frame.is_writable().context("Checking frame writability")? { + bail!("Motion interpolation output frame is not writable"); + } + let data = frame.data; + Ok([0, 1, 2].map(|i| { let (rows, width) = plane_size(layout, i); - let stride = linesize[i] as usize; - // SAFETY: the frame was just allocated by `av_frame_get_buffer` with - // this layout, so every plane holds `rows` rows of `stride >= width` - // bytes and nothing else references its buffers. + let stride = strides[i]; + // SAFETY: as in `frame_planes`, plane `i` spans at least + // `stride * (rows - 1) + width` bytes. The frame is writable, so no + // other frame references its buffers, and the three planes are + // disjoint arrays (fully planar format), so the mutable slices do + // not alias. They borrow `frame` mutably for their lifetime. let data = unsafe { std::slice::from_raw_parts_mut(data[i], stride * (rows - 1) + width) }; PlaneMut { data, stride } - }) + })) } /// A source frame with its timestamp on the resampler's exact time grid. @@ -740,7 +776,7 @@ impl Resampler { let layout = self.layout; let [(a, _), (b, _)] = &self.padded; self.interp - .interpolate(a, b, alpha, frame_planes_mut(&mut out, &layout)); + .interpolate(a, b, alpha, frame_planes_mut(&mut out, &layout)?); out.set_pts(pts); out.set_time_base(self.out_tb); unsafe { (*out.as_mut_ptr()).duration = 1 }; @@ -871,7 +907,11 @@ mod tests { frame.set_height(layout.height as i32); frame.alloc_buffer().unwrap(); let (cw, _) = layout.chroma_size(); - for (p, src) in frame_planes_mut(&mut frame, layout).into_iter().zip(&data) { + for (p, src) in frame_planes_mut(&mut frame, layout) + .unwrap() + .into_iter() + .zip(&data) + { let w = if src.len() == layout.width * layout.height { layout.width } else { @@ -1057,6 +1097,45 @@ mod tests { assert_eq!(out.map(|p| p.len()), [width * height, cw * ch, cw * ch]); } + #[test] + fn test_frame_planes_reject_mismatched_frames() { + let layout = Layout { + width: 64, + height: 32, + chroma_shift: (1, 1), + }; + let frame = av_frame(&layout, 0, 1, |x, y, _| pattern(x, y)); + assert!(frame_planes(&frame, &layout).is_ok()); + // A larger layout than the frame was allocated for must not be + // borrowed: its planes would extend past the frame's buffers. + let taller = Layout { + height: 64, + ..layout + }; + assert!(frame_planes(&frame, &taller).is_err()); + let wider_chroma = Layout { + chroma_shift: (0, 0), + ..layout + }; + assert!(frame_planes(&frame, &wider_chroma).is_err()); + + // A frame that shares its buffers with another is read-only. + let mut shared = frame.clone(); + assert!(frame_planes(&shared, &layout).is_ok()); + assert!(frame_planes_mut(&mut shared, &layout).is_err()); + drop(frame); + assert!(frame_planes_mut(&mut shared, &layout).is_ok()); + + // Packed formats are rejected before any plane is borrowed. + let mut packed = AVFrame::new(); + packed.set_format(ffi::AV_PIX_FMT_RGB24); + packed.set_width(64); + packed.set_height(32); + packed.alloc_buffer().unwrap(); + assert!(frame_planes(&packed, &layout).is_err()); + assert!(frame_planes_mut(&mut packed, &layout).is_err()); + } + #[test] fn test_scene_change() { let layout = Layout {