From d66ce14660ec45459a99cede767aff9f05a8af23 Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Thu, 10 Oct 2024 08:29:54 +0200 Subject: [PATCH] Small video-related refactor (#7660) ### What Some small improvements in service of coming ffmpeg stuff * Extracted from https://github.com/rerun-io/rerun/pull/7658 * Part of https://github.com/rerun-io/rerun/issues/7606 ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested the web demo (if applicable): * Using examples from latest `main` build: [rerun.io/viewer](https://rerun.io/viewer/pr/7660?manifest_url=https://app.rerun.io/version/main/examples_manifest.json) * Using full set of examples from `nightly` build: [rerun.io/viewer](https://rerun.io/viewer/pr/7660?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json) * [x] The PR title and labels are set such as to maximize their usefulness for the next release's CHANGELOG * [x] If applicable, add a new check to the [release checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)! * [x] If have noted any breaking changes to the log API in `CHANGELOG.md` and the migration guide - [PR Build Summary](https://build.rerun.io/pr/7660) - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html) To run all checks from `main`, comment on the PR with `@rerun-bot full-check`. --- crates/store/re_video/examples/frames.rs | 36 +++++++++++++------ crates/store/re_video/src/decode/mod.rs | 6 ++++ crates/store/re_video/src/demux/mod.rs | 4 +++ crates/store/re_video/src/demux/mp4.rs | 1 + .../src/video/decoder/native_decoder.rs | 20 ++++++++--- 5 files changed, 53 insertions(+), 14 deletions(-) diff --git a/crates/store/re_video/examples/frames.rs b/crates/store/re_video/examples/frames.rs index a27ae350681c..877c20bdff1d 100644 --- a/crates/store/re_video/examples/frames.rs +++ b/crates/store/re_video/examples/frames.rs @@ -13,6 +13,8 @@ use std::{ use indicatif::ProgressBar; use parking_lot::Mutex; +use re_video::{decode::SyncDecoder, VideoData}; + fn main() { // frames let args: Vec<_> = std::env::args().collect(); @@ -27,10 +29,6 @@ fn main() { let video = std::fs::read(video_path).expect("failed to read video"); let video = re_video::VideoData::load_mp4(&video).expect("failed to load video"); - let sync_decoder = Box::new( - re_video::decode::av1::SyncDav1dDecoder::new().expect("Failed to start AV1 decoder"), - ); - println!( "{} {}x{}", video.gops.len(), @@ -38,6 +36,26 @@ fn main() { video.config.coded_height ); + let mut decoder = create_decoder(&video); + + write_video_frames(&video, decoder.as_mut(), &output_dir); +} + +fn create_decoder(video: &VideoData) -> Box { + if video.config.is_av1() { + Box::new( + re_video::decode::av1::SyncDav1dDecoder::new().expect("Failed to start AV1 decoder"), + ) + } else { + panic!("Unsupported codec: {}", video.human_readable_codec_string()); + } +} + +fn write_video_frames( + video: &re_video::VideoData, + decoder: &mut dyn re_video::decode::SyncDecoder, + output_dir: &PathBuf, +) { let progress = ProgressBar::new(video.samples.len() as u64).with_message("Decoding video"); progress.enable_steady_tick(Duration::from_millis(100)); @@ -50,16 +68,14 @@ fn main() { frames.lock().push(frame); } }; - let mut decoder = - re_video::decode::AsyncDecoder::new("debug_name".to_owned(), sync_decoder, on_output); let start = Instant::now(); for sample in &video.samples { - decoder.decode(video.get(sample).unwrap()); + let should_stop = std::sync::atomic::AtomicBool::new(false); + let chunk = video.get(sample).unwrap(); + decoder.submit_chunk(&should_stop, chunk, &on_output); } - decoder.flush(); - drop(decoder); let end = Instant::now(); progress.finish(); @@ -72,7 +88,7 @@ fn main() { ); println!("Writing frames to {}", output_dir.display()); - std::fs::create_dir_all(&output_dir).expect("failed to create output directory"); + std::fs::create_dir_all(output_dir).expect("failed to create output directory"); let width = num_digits(frames.len()); for (i, frame) in frames.iter().enumerate() { diff --git a/crates/store/re_video/src/decode/mod.rs b/crates/store/re_video/src/decode/mod.rs index 5b15fc30a2d8..7fd77960f6de 100644 --- a/crates/store/re_video/src/decode/mod.rs +++ b/crates/store/re_video/src/decode/mod.rs @@ -38,7 +38,12 @@ pub trait SyncDecoder { } /// One chunk of encoded video data; usually one frame. +/// +/// One loaded [`crate::Sample`]. pub struct Chunk { + /// The start of a new [`crate::demux::GroupOfPictures`]? + pub is_sync: bool, + pub data: Vec, pub timestamp: Time, pub duration: Time, @@ -55,5 +60,6 @@ pub struct Frame { } pub enum PixelFormat { + Rgb8Unorm, Rgba8Unorm, } diff --git a/crates/store/re_video/src/demux/mod.rs b/crates/store/re_video/src/demux/mod.rs index f2b495f087e2..f32e17d2ce17 100644 --- a/crates/store/re_video/src/demux/mod.rs +++ b/crates/store/re_video/src/demux/mod.rs @@ -144,6 +144,7 @@ impl VideoData { data: data.to_vec(), timestamp: sample.decode_timestamp, duration: sample.duration, + is_sync: sample.is_sync, }) } } @@ -174,6 +175,9 @@ impl GroupOfPictures { /// A single sample in a video. #[derive(Debug, Clone)] pub struct Sample { + /// Is t his the start of a new [`GroupOfPictures`]? + pub is_sync: bool, + /// Time at which this sample appears in the decoded bitstream, in time units. /// /// Samples should be decoded in this order. diff --git a/crates/store/re_video/src/demux/mp4.rs b/crates/store/re_video/src/demux/mp4.rs index 24659f078cb9..cd4d36ccbd00 100644 --- a/crates/store/re_video/src/demux/mp4.rs +++ b/crates/store/re_video/src/demux/mp4.rs @@ -58,6 +58,7 @@ impl VideoData { let byte_length = sample.size as u32; samples.push(Sample { + is_sync: sample.is_sync, decode_timestamp, composition_timestamp, duration, diff --git a/crates/viewer/re_renderer/src/video/decoder/native_decoder.rs b/crates/viewer/re_renderer/src/video/decoder/native_decoder.rs index 2e5176692524..a02cb4fbf807 100644 --- a/crates/viewer/re_renderer/src/video/decoder/native_decoder.rs +++ b/crates/viewer/re_renderer/src/video/decoder/native_decoder.rs @@ -129,6 +129,22 @@ fn copy_video_frame_to_texture( frame: &Frame, texture: &wgpu::Texture, ) -> Result<(), DecodingError> { + let format = match frame.format { + re_video::PixelFormat::Rgb8Unorm => { + return copy_video_frame_to_texture( + queue, + &Frame { + data: crate::pad_rgb_to_rgba(&frame.data, 255_u8), + format: re_video::PixelFormat::Rgba8Unorm, + ..*frame + }, + texture, + ); + } + + re_video::PixelFormat::Rgba8Unorm => wgpu::TextureFormat::Rgba8Unorm, + }; + re_tracing::profile_function!(); let size = wgpu::Extent3d { @@ -137,10 +153,6 @@ fn copy_video_frame_to_texture( depth_or_array_layers: 1, }; - let format = match frame.format { - re_video::PixelFormat::Rgba8Unorm => wgpu::TextureFormat::Rgba8Unorm, - }; - let width_blocks = frame.width / format.block_dimensions().0; #[allow(clippy::unwrap_used)] // block_copy_size can only fail for weird compressed formats