From 84ae3bae6aea3b29d801f23c883efb3f4574553b Mon Sep 17 00:00:00 2001 From: Andreas Reich Date: Thu, 10 Oct 2024 18:17:22 +0200 Subject: [PATCH] handle strides in YUV data --- crates/store/re_video/src/decode/av1.rs | 82 ++++++++++++++++++++----- crates/store/re_video/src/decode/mod.rs | 4 ++ 2 files changed, 72 insertions(+), 14 deletions(-) diff --git a/crates/store/re_video/src/decode/av1.rs b/crates/store/re_video/src/decode/av1.rs index a934b0cd13a4e..fd5c5bc52c3f7 100644 --- a/crates/store/re_video/src/decode/av1.rs +++ b/crates/store/re_video/src/decode/av1.rs @@ -2,9 +2,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; -use dav1d::{PixelLayout, PlanarImageComponent}; - use crate::Time; +use dav1d::{PixelLayout, PlanarImageComponent}; use super::{ Chunk, ColorPrimaries, Error, Frame, OutputCallback, PixelFormat, Result, SyncDecoder, @@ -117,18 +116,73 @@ fn output_picture(picture: &dav1d::Picture, on_output: &(dyn Fn(Result) + // - matrix_coefficients // - transfer_characteristics - let data = match picture.pixel_layout() { - PixelLayout::I400 => picture.plane(PlanarImageComponent::Y).to_vec(), - PixelLayout::I420 | PixelLayout::I422 | PixelLayout::I444 => { - let mut data = Vec::with_capacity( - picture.stride(PlanarImageComponent::Y) as usize - + picture.stride(PlanarImageComponent::U) as usize - + picture.stride(PlanarImageComponent::V) as usize, - ); - data.extend_from_slice(&picture.plane(PlanarImageComponent::Y)); - data.extend_from_slice(&picture.plane(PlanarImageComponent::U)); - data.extend_from_slice(&picture.plane(PlanarImageComponent::V)); - data // TODO: how badly does this break with hdr? + let data = { + re_tracing::profile_scope!("copy_picture_data"); + + match picture.pixel_layout() { + PixelLayout::I400 => picture.plane(PlanarImageComponent::Y).to_vec(), + PixelLayout::I420 | PixelLayout::I422 | PixelLayout::I444 => { + // TODO(#7594): If `picture.bit_depth()` isn't 8 we have a problem: + // We can't handle high bit depths yet and the YUV converter at the other side + // bases its opinion on what an acceptable number of incoming bytes is on this. + // So we just clamp to that expectation, ignoring `picture.stride(PlanarImageComponent::Y)` & friends. + // Note that `bit_depth` is either 8 or 16, which is semi-independent `bits_per_component` (which is None/8/10/12). + if picture.bit_depth() != 8 { + re_log::warn_once!( + "Video uses unsupported bit depth larger than 8 bit per color component which is currently unsupported." + ); + } + + let height_y = picture.height() as usize; + let height_uv = match picture.pixel_layout() { + PixelLayout::I400 => 0, + PixelLayout::I420 => height_y / 2, + PixelLayout::I422 | PixelLayout::I444 => height_y, + }; + + let packed_stride_y = picture.width() as usize; + let actual_stride_y = picture.stride(PlanarImageComponent::Y) as usize; + + let packed_stride_uv = match picture.pixel_layout() { + PixelLayout::I400 => 0, + PixelLayout::I420 | PixelLayout::I422 => packed_stride_y / 2, + PixelLayout::I444 => packed_stride_y, + }; + let actual_stride_uv = picture.stride(PlanarImageComponent::U) as usize; // U / V stride is always the same. + + let num_packed_bytes_y = packed_stride_y * height_y; + let num_packed_bytes_uv = packed_stride_uv * height_uv; + + let mut data = Vec::with_capacity(num_packed_bytes_y + num_packed_bytes_uv * 2); + + // We could make our image ingestion pipeline even more sophisticated and pass that stride information through. + // But given that this is a matter of replacing a single large memcpy with a few hundred _still_ quite large ones, + // this should not make a lot of difference (citation needed!). + { + let plane = picture.plane(PlanarImageComponent::Y); + if actual_stride_y != packed_stride_y { + for y in 0..height_y { + let offset = y * actual_stride_y; + data.extend_from_slice(&plane[offset..(offset + packed_stride_y)]); + } + } else { + data.extend_from_slice(&plane[0..num_packed_bytes_y]); + } + } + for comp in [PlanarImageComponent::U, PlanarImageComponent::V] { + let plane = picture.plane(comp); + if actual_stride_uv != packed_stride_uv { + for y in 0..height_uv { + let offset = y * actual_stride_uv; + data.extend_from_slice(&plane[offset..(offset + packed_stride_uv)]); + } + } else { + data.extend_from_slice(&plane[0..num_packed_bytes_uv]); + } + } + + data + } } }; diff --git a/crates/store/re_video/src/decode/mod.rs b/crates/store/re_video/src/decode/mod.rs index 21a455dfb36e5..77f641b41db91 100644 --- a/crates/store/re_video/src/decode/mod.rs +++ b/crates/store/re_video/src/decode/mod.rs @@ -60,6 +60,7 @@ pub struct Frame { } /// Pixel format/layout used by [`Frame::data`]. +#[derive(Debug)] pub enum PixelFormat { Rgb8Unorm, Rgba8Unorm, @@ -77,6 +78,7 @@ pub enum PixelFormat { /// /// For details see `re_renderer`'s `YuvPixelLayout` type. #[allow(non_camel_case_types)] +#[derive(Debug)] pub enum YuvPixelLayout { Y_U_V444, Y_U_V422, @@ -87,6 +89,7 @@ pub enum YuvPixelLayout { /// Yuv value range used by [`PixelFormat::Yuv`]. /// /// For details see `re_renderer`'s `YuvRange` type. +#[derive(Debug)] pub enum YuvRange { Limited, Full, @@ -95,6 +98,7 @@ pub enum YuvRange { /// Color primaries used by [`PixelFormat::Yuv`]. /// /// For details see `re_renderer`'s `ColorPrimaries` type. +#[derive(Debug)] pub enum ColorPrimaries { Bt601, Bt709,