video: fix VideoFrame(Ref)::plane_data() returning truncated buffer

Plane index and component index are not interchangeable for some video
formats, e.g., AV12. Each plane could contain more than one component.
Therefore, the height of each plane's buffer should be the aggregated height
of all its components.

Fixes #536

Part-of: <https://gitlab.freedesktop.org/gstreamer/gstreamer-rs/-/merge_requests/1607>
This commit is contained in:
Cheung Yik Pang 2024-12-09 08:22:36 +00:00
parent 4bb7208823
commit 1410e2f3a3
2 changed files with 151 additions and 120 deletions

View file

@ -370,6 +370,19 @@ impl VideoFormatInfo {
unsafe { &*(&self.0.tile_info[plane as usize] as *const _ as *const VideoTileInfo) }
}
#[cfg(feature = "v1_18")]
#[cfg_attr(docsrs, doc(cfg(feature = "v1_18")))]
#[doc(alias = "gst_video_format_info_component")]
pub fn component(&self, plane: u32) -> [i32; ffi::GST_VIDEO_MAX_COMPONENTS as usize] {
assert!(plane < self.n_planes());
let mut comp = [-1i32; ffi::GST_VIDEO_MAX_COMPONENTS as usize];
unsafe {
ffi::gst_video_format_info_component(self.to_glib_none().0, plane, comp.as_mut_ptr());
}
comp
}
}
unsafe impl Sync for VideoFormatInfo {}

View file

@ -19,6 +19,35 @@ unsafe impl<T> IsVideoFrame for VideoFrame<T> {
}
}
fn plane_buffer_info<T: IsVideoFrame>(
frame: &T,
plane: u32,
) -> Result<(usize, usize), glib::BoolError> {
skip_assert_initialized!();
if plane >= frame.n_planes() {
return Err(glib::bool_error!(
"Plane index higher than number of planes"
));
}
let format_info = frame.format_info();
// Just get the palette
if format_info.has_palette() && plane == 1 {
return Ok((1, 256 * 4));
}
let w = frame.plane_stride()[plane as usize] as u32;
let h = frame.plane_height(plane);
if w == 0 || h == 0 {
return Ok((0, 0));
}
Ok((plane as usize, (w * h) as usize))
}
pub struct VideoFrame<T> {
frame: ffi::GstVideoFrame,
buffer: gst::Buffer,
@ -146,6 +175,28 @@ pub trait VideoFrameExt: IsVideoFrame {
self.info().offset()
}
#[inline]
fn plane_height(&self, plane: u32) -> u32 {
cfg_if::cfg_if! {
if #[cfg(feature = "v1_18")] {
let comp = self.format_info().component(plane)[0];
if comp == -1 {
0
} else {
self.comp_height(comp as u32)
}
} else {
// FIXME: This assumes that the horizontal subsampling of all
// components in the plane is the same, which is probably safe
// Legacy implementation that does not support video formats
// where plane index and component index are not the same.
// See #536
self.format_info().scale_height(plane as u8, self.height())
}
}
}
#[inline]
fn comp_depth(&self, component: u32) -> u32 {
self.info().comp_depth(component as u8)
@ -246,40 +297,22 @@ impl<T> VideoFrame<T> {
}
pub fn plane_data(&self, plane: u32) -> Result<&[u8], glib::BoolError> {
if plane >= self.n_planes() {
return Err(glib::bool_error!(
"Plane index higher than number of planes"
));
}
let format_info = self.format_info();
// Just get the palette
if format_info.has_palette() && plane == 1 {
unsafe {
return Ok(slice::from_raw_parts(
self.frame.data[1] as *const u8,
256 * 4,
));
}
}
let w = self.plane_stride()[plane as usize] as u32;
// FIXME: This assumes that the horizontal subsampling of all
// components in the plane is the same, which is probably safe
let h = format_info.scale_height(plane as u8, self.height());
if w == 0 || h == 0 {
match plane_buffer_info(self, plane) {
Ok((plane, size)) => {
if size == 0 {
return Ok(&[]);
}
unsafe {
Ok(slice::from_raw_parts(
self.frame.data[plane as usize] as *const u8,
(w * h) as usize,
self.frame.data[plane] as *const u8,
size,
))
}
}
Err(err) => Err(err),
}
}
pub fn planes_data(&self) -> [&[u8]; 4] {
let mut planes = [[].as_slice(); 4];
@ -476,40 +509,22 @@ impl VideoFrame<Writable> {
}
pub fn plane_data_mut(&mut self, plane: u32) -> Result<&mut [u8], glib::BoolError> {
if plane >= self.n_planes() {
return Err(glib::bool_error!(
"Plane index higher than number of planes"
));
}
let format_info = self.format_info();
// Just get the palette
if format_info.has_palette() && plane == 1 {
unsafe {
return Ok(slice::from_raw_parts_mut(
self.frame.data[1] as *mut u8,
256 * 4,
));
}
}
let w = self.plane_stride()[plane as usize] as u32;
// FIXME: This assumes that the horizontal subsampling of all
// components in the plane is the same, which is probably safe
let h = format_info.scale_height(plane as u8, self.height());
if w == 0 || h == 0 {
match plane_buffer_info(self, plane) {
Ok((plane, size)) => {
if size == 0 {
return Ok(&mut []);
}
unsafe {
Ok(slice::from_raw_parts_mut(
self.frame.data[plane as usize] as *mut u8,
(w * h) as usize,
self.frame.data[plane] as *mut u8,
size,
))
}
}
Err(err) => Err(err),
}
}
pub fn planes_data_mut(&mut self) -> [&mut [u8]; 4] {
unsafe {
@ -615,40 +630,22 @@ impl<T> VideoFrameRef<T> {
}
pub fn plane_data(&self, plane: u32) -> Result<&[u8], glib::BoolError> {
if plane >= self.n_planes() {
return Err(glib::bool_error!(
"Plane index higher than number of planes"
));
}
let format_info = self.format_info();
// Just get the palette
if format_info.has_palette() && plane == 1 {
unsafe {
return Ok(slice::from_raw_parts(
self.frame.data[1] as *const u8,
256 * 4,
));
}
}
let w = self.plane_stride()[plane as usize] as u32;
// FIXME: This assumes that the horizontal subsampling of all
// components in the plane is the same, which is probably safe
let h = format_info.scale_height(plane as u8, self.height());
if w == 0 || h == 0 {
match plane_buffer_info(self, plane) {
Ok((plane, size)) => {
if size == 0 {
return Ok(&[]);
}
unsafe {
Ok(slice::from_raw_parts(
self.frame.data[plane as usize] as *const u8,
(w * h) as usize,
self.frame.data[plane] as *const u8,
size,
))
}
}
Err(err) => Err(err),
}
}
pub fn planes_data(&self) -> [&[u8]; 4] {
let mut planes = [[].as_slice(); 4];
@ -844,40 +841,22 @@ impl<'a> VideoFrameRef<&'a mut gst::BufferRef> {
}
pub fn plane_data_mut(&mut self, plane: u32) -> Result<&mut [u8], glib::BoolError> {
if plane >= self.n_planes() {
return Err(glib::bool_error!(
"Plane index higher than number of planes"
));
}
let format_info = self.format_info();
// Just get the palette
if format_info.has_palette() && plane == 1 {
unsafe {
return Ok(slice::from_raw_parts_mut(
self.frame.data[1] as *mut u8,
256 * 4,
));
}
}
let w = self.plane_stride()[plane as usize] as u32;
// FIXME: This assumes that the horizontal subsampling of all
// components in the plane is the same, which is probably safe
let h = format_info.scale_height(plane as u8, self.height());
if w == 0 || h == 0 {
match plane_buffer_info(self, plane) {
Ok((plane, size)) => {
if size == 0 {
return Ok(&mut []);
}
unsafe {
Ok(slice::from_raw_parts_mut(
self.frame.data[plane as usize] as *mut u8,
(w * h) as usize,
self.frame.data[plane] as *mut u8,
size,
))
}
}
Err(err) => Err(err),
}
}
pub fn planes_data_mut(&mut self) -> [&mut [u8]; 4] {
unsafe {
@ -1057,4 +1036,43 @@ mod tests {
assert!(frame.info() == &info);
}
}
#[cfg(feature = "v1_20")]
#[test]
fn test_plane_data() {
gst::init().unwrap();
let info = crate::VideoInfo::builder(crate::VideoFormat::Av12, 320, 240)
.build()
.unwrap();
let buffer = gst::Buffer::with_size(info.size()).unwrap();
let mut frame = VideoFrame::from_buffer_writable(buffer, &info).unwrap();
// Alpha plane
{
let mut frame = frame.as_mut_video_frame_ref();
let data = frame.plane_data_mut(2).unwrap();
assert_eq!(data.len(), 320 * 240);
data[0] = 42;
}
// UV plane
{
let mut frame = frame.as_mut_video_frame_ref();
let data = frame.plane_data_mut(1).unwrap();
assert_eq!(data.len(), 320 * 120);
data[0] = 42;
}
let frame = frame.into_buffer();
let frame = VideoFrame::from_buffer_readable(frame, &info).unwrap();
let alpha_data = frame.plane_data(2).unwrap();
assert_eq!(alpha_data.len(), 320 * 240);
assert_eq!(alpha_data[0], 42);
let uv_data = frame.plane_data(1).unwrap();
assert_eq!(uv_data.len(), 320 * 120);
assert_eq!(uv_data[0], 42);
}
}