diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d10de2d..e78661c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,8 +54,8 @@ jobs: "$prefix/bin/ffmpeg" -version - name: Static checks run: | - ruff check src/tensorcodec av/python tests scripts benchmarks/image_codecs.py - ruff format --check src/tensorcodec av/python tests scripts benchmarks/image_codecs.py + ruff check src/tensorcodec av/python tests scripts benchmarks/image_codecs.py benchmarks/transform_bench.py + ruff format --check src/tensorcodec av/python tests scripts benchmarks/image_codecs.py benchmarks/transform_bench.py cargo fmt --manifest-path av/Cargo.toml --check cargo clippy --manifest-path av/Cargo.toml --locked -- -D warnings - name: Build tensorcodec-av against prebuilt FFmpeg and install tensorcodec diff --git a/README.md b/README.md index 656be41..cb19b11 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,7 @@ with VideoDecoder("video.mkv", seek_mode="timestamp") as decoder: frames = decoder.get_frames_played_at([10.0, 10.1, 10.2]) ``` -Decoder transforms (since v0.1.5) resize and crop inside the decoder, as TorchCodec's do: +Decoder transforms (since v0.1.5) resize and crop inside the decoder: ```python from tensorcodec.transforms import CenterCrop, Resize @@ -69,6 +69,10 @@ with VideoDecoder("video.mp4", transforms=[Resize((256, 340)), CenterCrop((224, frames = decoder.get_frames_at([0, 10]) # (2, 3, 224, 224) ``` +Crops select the same pixels as TorchCodec's. Since v0.4.2, `Resize` scales in YUV while converting +to RGB, so it costs no more than native-size output; its pixels differ from TorchCodec's by about +0.9 levels on average (see [resizing](docs/compatibility.md#resizing)). + This TensorCodec extension selects by actual PTS and retries seeks that overshoot. It supports time queries, including ranges with explicit `fps`, but not frame indices, `len(decoder)`, or `get_all_frames()`. See [the contract](docs/compatibility.md#timestamp-mode). diff --git a/av/Cargo.lock b/av/Cargo.lock index 5ae0aee..d00c733 100644 --- a/av/Cargo.lock +++ b/av/Cargo.lock @@ -449,7 +449,7 @@ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" [[package]] name = "tensorcodec-av" -version = "0.4.1" +version = "0.4.2" dependencies = [ "ffmpeg-sys-next", "libc", diff --git a/av/Cargo.toml b/av/Cargo.toml index f6656af..024ff5b 100644 --- a/av/Cargo.toml +++ b/av/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tensorcodec-av" -version = "0.4.1" +version = "0.4.2" edition = "2021" rust-version = "1.88" # ffmpeg-sys-next 9 builds with let chains license = "Apache-2.0" diff --git a/av/src/ffmpeg.rs b/av/src/ffmpeg.rs index a53a356..381c46b 100644 --- a/av/src/ffmpeg.rs +++ b/av/src/ffmpeg.rs @@ -43,8 +43,8 @@ pub enum VideoRequest { Timestamps(Vec), } -/// A geometric step on RGB frames; a rotation (counterclockwise quarter turns) comes first, then -/// crops and resizes in the rotated frame's coordinates. +/// A geometric step; a rotation (counterclockwise quarter turns) comes first, then crops and resizes in the +/// rotated frame's coordinates. `Pipeline` decides which run on the decoded planes and which on RGB. #[derive(Clone, Copy, PartialEq)] pub enum Op { Rotate { @@ -62,21 +62,51 @@ pub enum Op { }, } -/// One resize step's swscale context (bilinear, RGB to RGB at the same depth) and its output frame. +/// One resize step: a bilinear swscale pass from a source region straight to RGB at the output size. swscale +/// writes planar RGB with SIMD (its packed full-chroma writer is C only), interleaved into `frame` afterwards. struct Resizer { scale: *mut av::SwsContext, - config: (i32, i32, i32, i32, i32), + config: [i32; 8], + planar: *mut av::AVFrame, frame: *mut av::AVFrame, } +impl Resizer { + fn new() -> Result { + let (planar, frame) = unsafe { (av::av_frame_alloc(), av::av_frame_alloc()) }; + // Owned before the check, so Drop frees whichever allocation succeeded. + let resizer = Self { + scale: ptr::null_mut(), + config: [0; 8], + planar, + frame, + }; + if planar.is_null() || frame.is_null() { + return Err(failure("cannot allocate resize frames")); + } + Ok(resizer) + } +} impl Drop for Resizer { fn drop(&mut self) { unsafe { av::sws_freeContext(self.scale); + av::av_frame_free(&mut self.planar); av::av_frame_free(&mut self.frame); } } } +/// Planes of a region to resize: decoded YUV/RGB planes or an RGB intermediate. `chroma_position` is swscale's +/// `src_h_chr_pos`/`src_v_chr_pos` (1/256 luma pixel; -513 for its default, centered siting). +struct Region { + data: [*const u8; 8], + linesize: [i32; 8], + format: av::AVPixelFormat, + width: i32, + height: i32, + chroma_position: (i32, i32), +} + enum Selection { Frame { pts: i64, key: i64, exact: bool }, Timestamp(f64), @@ -214,7 +244,7 @@ pub struct Decoder { frame: *mut av::AVFrame, rgb_frame: *mut av::AVFrame, scale: *mut av::SwsContext, - scale_config: Option<(i32, i32, i32, i32)>, + scale_config: Option<[i32; 4]>, resizers: Vec, rotated: Vec, io: *mut av::AVIOContext, @@ -892,6 +922,7 @@ impl Decoder { let mut pts = vec![0.; length]; let mut durations = vec![0.; length]; let mut active_key = None; + let pipeline = Pipeline::new(ops, (self.video_layout.0, self.video_layout.1)); let stride = count * if high_depth { 2 } else { 1 }; for (selection, positions) in requests { match selection { @@ -910,95 +941,53 @@ impl Decoder { return Err(failure("dynamic frame dimensions are unsupported")); } let input_format: av::AVPixelFormat = std::mem::transmute(frame.format); - let output_frame = if native { + let format = if high_depth { + av::AVPixelFormat::AV_PIX_FMT_RGB48LE + } else { + av::AVPixelFormat::AV_PIX_FMT_RGB24 + }; + let (data, linesize, size, rest, slot) = if native { if input_format != source_format { return Err(Error("pixel format changed within stream".into(), true)); } - frame + let size = (frame.width, frame.height); + ( + frame.data[0] as *const u8, + frame.linesize[0], + size, + &[][..], + 0, + ) + } else if let Some((region, (h, w))) = pipeline + .first + .and_then(|(rect, size)| Some((region_planes(frame, rect)?, size))) + { + let (data, linesize) = + self.resize(0, ®ion, (w, h), high_depth, Some(frame))?; + (data, linesize, (w, h), &pipeline.after_first[..], 1) } else { - let output_format = if high_depth { - av::AVPixelFormat::AV_PIX_FMT_RGB48LE - } else { - av::AVPixelFormat::AV_PIX_FMT_RGB24 - }; - let config = ( + // Full-size color conversion, then every op on RGB. + let config = [ frame.width, frame.height, input_format as i32, - output_format as i32, - ); + format as i32, + ]; if self.scale_config != Some(config) { av::sws_freeContext(self.scale); - self.scale = av::sws_getContext( - frame.width, - frame.height, - input_format, - frame.width, - frame.height, - output_format, + self.scale = scale_context( + (frame.width, frame.height, input_format), + (frame.width, frame.height, format), 0, - ptr::null_mut(), - ptr::null_mut(), - ptr::null(), + (-513, -513), ); self.scale_config = Some(config); } if self.scale.is_null() { return Err(failure("cannot initialize color conversion")); } - let mut inverse = ptr::null_mut(); - let mut table = ptr::null_mut(); - let ( - mut source_range, - mut destination_range, - mut brightness, - mut contrast, - mut saturation, - ) = (0, 0, 0, 0, 0); - check( - av::sws_getColorspaceDetails( - self.scale, - &mut inverse, - &mut source_range, - &mut table, - &mut destination_range, - &mut brightness, - &mut contrast, - &mut saturation, - ), - "read color conversion settings", - )?; - if frame.color_range != av::AVColorRange::AVCOL_RANGE_UNSPECIFIED { - source_range = - i32::from(frame.color_range == av::AVColorRange::AVCOL_RANGE_JPEG); - } - let coefficients = av::sws_getCoefficients(frame.colorspace as i32); - check( - av::sws_setColorspaceDetails( - self.scale, - coefficients, - source_range, - coefficients, - destination_range, - brightness, - contrast, - saturation, - ), - "configure color conversion", - )?; - if (*self.rgb_frame).width != frame.width - || (*self.rgb_frame).height != frame.height - || (*self.rgb_frame).format != output_format as i32 - { - av::av_frame_unref(self.rgb_frame); - (*self.rgb_frame).width = frame.width; - (*self.rgb_frame).height = frame.height; - (*self.rgb_frame).format = output_format as i32; - check( - av::av_frame_get_buffer(self.rgb_frame, 32), - "allocate RGB frame", - )?; - } + configure_colors(self.scale, frame)?; + ensure_frame(self.rgb_frame, frame.width, frame.height, format)?; let rows = av::sws_scale( self.scale, frame.data.as_ptr() as *const *const u8, @@ -1011,18 +1000,18 @@ impl Decoder { if rows != frame.height { return Err(failure("color conversion failed")); } - &*self.rgb_frame - }; - let (data, linesize) = if ops.is_empty() { - (output_frame.data[0] as *const u8, output_frame.linesize[0]) - } else { - let format = if high_depth { - av::AVPixelFormat::AV_PIX_FMT_RGB48LE - } else { - av::AVPixelFormat::AV_PIX_FMT_RGB24 - }; - self.transform(output_frame, format, ops)? + let rgb = &*self.rgb_frame; + let size = (frame.width, frame.height); + ( + rgb.data[0] as *const u8, + rgb.linesize[0], + size, + &pipeline.ops[..], + 0, + ) }; + let (data, linesize) = + self.transform((data, linesize), size, format, rest, slot)?; let row_bytes = width * channels * if high_depth { 2 } else { 1 }; if data.is_null() || (linesize.unsigned_abs() as usize) < row_bytes { return Err(failure("invalid decoded frame stride")); @@ -1080,22 +1069,18 @@ impl Decoder { }) } - /// Applies `ops` to an RGB frame: a crop moves the view, a resize scales it into its own frame. - /// Returns the result's first row and stride. + /// Applies `ops` to an RGB image (first row, stride) of `size` (width, height): a crop moves the view, a + /// resize scales it into resize slot `slot` onwards. Returns the result's first row and stride. unsafe fn transform( &mut self, - frame: &av::AVFrame, + (mut data, mut linesize): (*const u8, i32), + (mut width, mut height): (i32, i32), format: av::AVPixelFormat, ops: &[Op], + mut slot: usize, ) -> Result<(*const u8, i32)> { - let pixel_bytes: isize = if format == av::AVPixelFormat::AV_PIX_FMT_RGB24 { - 3 - } else { - 6 - }; - let (mut data, mut linesize) = (frame.data[0] as *const u8, frame.linesize[0]); - let (mut width, mut height) = (frame.width, frame.height); - let mut resize = 0; + let high_depth = format == av::AVPixelFormat::AV_PIX_FMT_RGB48LE; + let pixel_bytes: isize = if high_depth { 6 } else { 3 }; for op in ops { match *op { Op::Rotate { turns } => { @@ -1146,70 +1131,98 @@ impl Decoder { height: h, width: w, } => { - let config = (width, height, w, h, format as i32); - if self.resizers.len() <= resize { - let frame = av::av_frame_alloc(); - if frame.is_null() { - return Err(failure("cannot allocate resize frame")); - } - self.resizers.push(Resizer { - scale: ptr::null_mut(), - config: (0, 0, 0, 0, 0), - frame, - }); - } - let resizer = &mut self.resizers[resize]; - if resizer.scale.is_null() || resizer.config != config { - av::sws_freeContext(resizer.scale); - resizer.scale = av::sws_getContext( - width, - height, - format, - w, - h, - format, - av::SWS_BILINEAR, - ptr::null_mut(), - ptr::null_mut(), - ptr::null(), - ); - if resizer.scale.is_null() { - return Err(failure("cannot initialize resize")); - } - resizer.config = config; - av::av_frame_unref(resizer.frame); - (*resizer.frame).width = w; - (*resizer.frame).height = h; - (*resizer.frame).format = format as i32; - check( - av::av_frame_get_buffer(resizer.frame, 32), - "allocate resize frame", - )?; - } - let source = [data, ptr::null(), ptr::null(), ptr::null()]; - let strides = [linesize, 0, 0, 0]; - let rows = av::sws_scale( - resizer.scale, - source.as_ptr(), - strides.as_ptr(), - 0, + let mut planes = [ptr::null(); 8]; + let mut strides = [0; 8]; + (planes[0], strides[0]) = (data, linesize); + let region = Region { + data: planes, + linesize: strides, + format, + width, height, - (*resizer.frame).data.as_ptr(), - (*resizer.frame).linesize.as_ptr(), - ); - if rows != h { - return Err(failure("resize failed")); - } - data = (*resizer.frame).data[0]; - linesize = (*resizer.frame).linesize[0]; + chroma_position: (-513, -513), + }; + (data, linesize) = self.resize(slot, ®ion, (w, h), high_depth, None)?; (width, height) = (w, h); - resize += 1; + slot += 1; } } } Ok((data, linesize)) } + /// Resizes `region` to `size` (width, height) RGB in resize slot `slot`; `colors` is the decoded frame + /// whose color space and range apply when the region holds its planes. Returns the first row and stride. + unsafe fn resize( + &mut self, + slot: usize, + region: &Region, + (width, height): (i32, i32), + high_depth: bool, + colors: Option<&av::AVFrame>, + ) -> Result<(*const u8, i32)> { + while self.resizers.len() <= slot { + self.resizers.push(Resizer::new()?); + } + let resizer = &mut self.resizers[slot]; + let (planar, packed, bytes) = if high_depth { + ( + av::AVPixelFormat::AV_PIX_FMT_GBRP16LE, + av::AVPixelFormat::AV_PIX_FMT_RGB48LE, + 2, + ) + } else { + ( + av::AVPixelFormat::AV_PIX_FMT_GBRP, + av::AVPixelFormat::AV_PIX_FMT_RGB24, + 1, + ) + }; + let config = [ + region.width, + region.height, + region.format as i32, + width, + height, + planar as i32, + region.chroma_position.0, + region.chroma_position.1, + ]; + if resizer.scale.is_null() || resizer.config != config { + av::sws_freeContext(resizer.scale); + // Full-resolution chroma on both sides: interpolated from subsampled planes, not decimated from RGB. + resizer.scale = scale_context( + (region.width, region.height, region.format), + (width, height, planar), + av::SWS_BILINEAR | av::SWS_FULL_CHR_H_INT | av::SWS_FULL_CHR_H_INP, + region.chroma_position, + ); + if resizer.scale.is_null() { + return Err(failure("cannot initialize resize")); + } + resizer.config = config; + } + if let Some(frame) = colors { + configure_colors(resizer.scale, frame)?; + } + ensure_frame(resizer.planar, width, height, planar)?; + ensure_frame(resizer.frame, width, height, packed)?; + let rows = av::sws_scale( + resizer.scale, + region.data.as_ptr(), + region.linesize.as_ptr(), + 0, + region.height, + (*resizer.planar).data.as_ptr(), + (*resizer.planar).linesize.as_ptr(), + ); + if rows != height { + return Err(failure("resize failed")); + } + interleave(&*resizer.planar, &*resizer.frame, bytes); + Ok(((*resizer.frame).data[0], (*resizer.frame).linesize[0])) + } + pub fn audio(&mut self, rate: i32, channels: i32, stop: Option) -> Result