diff --git a/.github/workflows/reusable_bench.yml b/.github/workflows/reusable_bench.yml index c7a2576f71f2..e5f55fccdf70 100644 --- a/.github/workflows/reusable_bench.yml +++ b/.github/workflows/reusable_bench.yml @@ -75,18 +75,30 @@ jobs: workload_identity_provider: ${{ secrets.GOOGLE_WORKLOAD_IDENTITY_PROVIDER }} service_account: ${{ secrets.GOOGLE_SERVICE_ACCOUNT }} + - uses: prefix-dev/setup-pixi@v0.8.1 + with: + pixi-version: v0.25.0 + # Only has the deps for round-trips. Not all examples. + environments: wheel-test-min + + - name: Download test assets + run: pixi run -e wheel-test-min python ./tests/assets/download_test_assets.py + - name: Add SHORT_SHA env property with commit short sha run: echo "SHORT_SHA=`echo ${{github.sha}} | cut -c1-7`" >> $GITHUB_ENV - name: Run benchmark # Use bash shell so we get pipefail behavior with tee + # Running under `pixi` so we get `nasm` run: | - cargo bench \ + pixi run -e wheel-test-min \ + cargo bench \ --all-features \ -p re_entity_db \ -p re_log_encoding \ -p re_query \ -p re_tuid \ + -p re_video \ -- --output-format=bencher | tee /tmp/${{ env.SHORT_SHA }} - name: "Set up Cloud SDK" diff --git a/Cargo.lock b/Cargo.lock index 7ec45bf32de6..f723db98fd95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5555,8 +5555,7 @@ dependencies = [ [[package]] name = "re_mp4" version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d1e30657b1ae7f0dd3428a59dc8140732b74a22cc07763606c9ec4054138731" +source = "git+https://github.com/rerun-io/re_mp4?rev=7d38361ee5b05f5a2b83a8029057c8a24d2e9023#7d38361ee5b05f5a2b83a8029057c8a24d2e9023" dependencies = [ "byteorder", "bytes", @@ -6182,6 +6181,7 @@ name = "re_video" version = "0.20.0-alpha.1+dev" dependencies = [ "cfg_aliases 0.2.1", + "criterion", "crossbeam", "econtext", "indicatif", diff --git a/Cargo.toml b/Cargo.toml index ea41a9a9f7e7..ed6f82b13c0c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -565,3 +565,6 @@ missing_errors_doc = "allow" re_arrow2 = { git = "https://github.com/rerun-io/re_arrow2", rev = "e4717d6debc6d4474ec10db8f629f823f57bad07" } # dav1d = { path = "/home/cmc/dev/rerun-io/rav1d", package = "re_rav1d", version = "0.1.1" } + +# Commit on `main` branch of `re_mp4` +re_mp4 = { git = "https://github.com/rerun-io/re_mp4", rev = "7d38361ee5b05f5a2b83a8029057c8a24d2e9023" } diff --git a/crates/store/re_video/Cargo.toml b/crates/store/re_video/Cargo.toml index a9ccd76883c5..c65854fc3d2f 100644 --- a/crates/store/re_video/Cargo.toml +++ b/crates/store/re_video/Cargo.toml @@ -62,7 +62,7 @@ dav1d = { workspace = true, optional = true, default-features = false, features [dev-dependencies] indicatif.workspace = true - +criterion.workspace = true # For build.rs: [build-dependencies] @@ -71,3 +71,8 @@ cfg_aliases.workspace = true [[example]] name = "frames" + + +[[bench]] +name = "video_load_bench" +harness = false diff --git a/crates/store/re_video/benches/video_load_bench.rs b/crates/store/re_video/benches/video_load_bench.rs new file mode 100644 index 000000000000..5786b717c212 --- /dev/null +++ b/crates/store/re_video/benches/video_load_bench.rs @@ -0,0 +1,24 @@ +#![allow(clippy::unwrap_used)] // acceptable in benchmarks + +use std::path::Path; + +use criterion::{criterion_group, criterion_main, Criterion}; + +fn video_load(c: &mut Criterion) { + let video_path = Path::new(env!("CARGO_MANIFEST_DIR")) + .ancestors() + .nth(3) + .unwrap() + .join("tests/assets/video/Big_Buck_Bunny_1080_10s_av1.mp4"); + let video = std::fs::read(video_path).unwrap(); + c.bench_function("video_load", |b| { + b.iter_batched( + || {}, + |()| re_video::VideoData::load_from_bytes(&video, "video/mp4"), + criterion::BatchSize::LargeInput, + ); + }); +} + +criterion_group!(benches, video_load); +criterion_main!(benches); diff --git a/crates/store/re_video/examples/frames.rs b/crates/store/re_video/examples/frames.rs index 79f1ac37025e..c03d6a637643 100644 --- a/crates/store/re_video/examples/frames.rs +++ b/crates/store/re_video/examples/frames.rs @@ -24,8 +24,8 @@ fn main() { println!("Decoding {video_path}"); - 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 video_blob = std::fs::read(video_path).expect("failed to read video"); + let video = re_video::VideoData::load_mp4(&video_blob).expect("failed to load video"); println!( "{} {}x{}", @@ -37,11 +37,12 @@ fn main() { let mut decoder = re_video::decode::new_decoder(video_path.to_string(), &video) .expect("Failed to create decoder"); - write_video_frames(&video, decoder.as_mut(), &output_dir); + write_video_frames(&video, &video_blob, decoder.as_mut(), &output_dir); } fn write_video_frames( video: &re_video::VideoData, + video_blob: &[u8], decoder: &mut dyn re_video::decode::SyncDecoder, output_dir: &PathBuf, ) { @@ -61,7 +62,7 @@ fn write_video_frames( let start = Instant::now(); for sample in &video.samples { let should_stop = std::sync::atomic::AtomicBool::new(false); - let chunk = video.get(sample).unwrap(); + let chunk = sample.get(video_blob).unwrap(); decoder.submit_chunk(&should_stop, chunk, &on_output); } diff --git a/crates/store/re_video/src/demux/mod.rs b/crates/store/re_video/src/demux/mod.rs index 79c1e83307a1..d93083a21cb1 100644 --- a/crates/store/re_video/src/demux/mod.rs +++ b/crates/store/re_video/src/demux/mod.rs @@ -60,9 +60,6 @@ pub struct VideoData { /// and should be presented in composition-timestamp order. pub samples: Vec, - /// This array stores all data used by samples. - pub data: Vec, - /// All the tracks in the mp4; not just the video track. /// /// Can be nice to show in a UI. @@ -245,25 +242,6 @@ impl VideoData { .sorted() }) } - - /// Returns `None` if the sample is invalid/out-of-range. - pub fn get(&self, sample: &Sample) -> Option { - let byte_offset = sample.byte_offset as usize; - let byte_length = sample.byte_length as usize; - - if self.data.len() < byte_offset + byte_length { - None - } else { - let data = &self.data[byte_offset..byte_offset + byte_length]; - - Some(Chunk { - data: data.to_vec(), - composition_timestamp: sample.composition_timestamp, - duration: sample.duration, - is_sync: sample.is_sync, - }) - } - } } /// A Group of Pictures (GOP) always starts with an I-frame, followed by delta-frames. @@ -311,13 +289,34 @@ pub struct Sample { /// Duration of the sample, in time units. pub duration: Time, - /// Offset into [`VideoData::data`] + /// Offset into the video data. pub byte_offset: u32, /// Length of sample starting at [`Sample::byte_offset`]. pub byte_length: u32, } +impl Sample { + /// Read the sample from the video data. + /// + /// Note that `data` _must_ be a reference to the original MP4 file + /// from which the [`VideoData`] was loaded. + /// + /// Returns `None` if the sample is out of bounds, which can only happen + /// if `data` is not the original video data. + pub fn get(&self, data: &[u8]) -> Option { + let data = data + .get(self.byte_offset as usize..(self.byte_offset + self.byte_length) as usize)? + .to_vec(); + Some(Chunk { + data, + composition_timestamp: self.composition_timestamp, + duration: self.duration, + is_sync: self.is_sync, + }) + } +} + /// Configuration of a video. #[derive(Debug, Clone)] pub struct Config { @@ -385,7 +384,6 @@ impl std::fmt::Debug for VideoData { "samples", &self.samples.iter().enumerate().collect::>(), ) - .field("data", &self.data.len()) .finish() } } diff --git a/crates/store/re_video/src/demux/mp4.rs b/crates/store/re_video/src/demux/mp4.rs index 58d0c4551b27..3bff80f0a63e 100644 --- a/crates/store/re_video/src/demux/mp4.rs +++ b/crates/store/re_video/src/demux/mp4.rs @@ -41,7 +41,6 @@ impl VideoData { let mut samples = Vec::::new(); let mut gops = Vec::::new(); let mut gop_sample_start_index = 0; - let data = track.data.clone(); for sample in &track.samples { if sample.is_sync && !samples.is_empty() { @@ -86,7 +85,6 @@ impl VideoData { duration, gops, samples, - data, mp4_tracks, }) } diff --git a/crates/viewer/re_data_ui/src/blob.rs b/crates/viewer/re_data_ui/src/blob.rs index 503ebb942b84..220c1aec7e13 100644 --- a/crates/viewer/re_data_ui/src/blob.rs +++ b/crates/viewer/re_data_ui/src/blob.rs @@ -129,6 +129,7 @@ pub fn blob_preview_and_save_ui( ui_layout, &video_result, video_timestamp, + blob, ); } @@ -175,6 +176,7 @@ fn show_video_blob_info( ui_layout: UiLayout, video_result: &Result, video_timestamp: Option, + blob: &re_types::datatypes::Blob, ) { #[allow(clippy::match_same_arms)] match video_result { @@ -262,7 +264,12 @@ fn show_video_blob_info( ui.id().with("video_player").value(), ); - match video.frame_at(render_ctx, decode_stream_id, timestamp_in_seconds) { + match video.frame_at( + render_ctx, + decode_stream_id, + timestamp_in_seconds, + blob.as_slice(), + ) { Ok(VideoFrameTexture { texture, time_range, diff --git a/crates/viewer/re_renderer/src/video/decoder/mod.rs b/crates/viewer/re_renderer/src/video/decoder/mod.rs index f82026315f47..a0fc1413d1ff 100644 --- a/crates/viewer/re_renderer/src/video/decoder/mod.rs +++ b/crates/viewer/re_renderer/src/video/decoder/mod.rs @@ -189,6 +189,7 @@ impl VideoDecoder { &mut self, render_ctx: &RenderContext, presentation_timestamp_s: f64, + video_data: &[u8], ) -> Result { if presentation_timestamp_s < 0.0 { return Err(DecodingError::NegativeTimestamp); @@ -197,7 +198,7 @@ impl VideoDecoder { let presentation_timestamp = presentation_timestamp.min(self.data.duration); // Don't seek past the end of the video. let error_on_last_frame_at = self.last_error.is_some(); - let result = self.frame_at_internal(render_ctx, presentation_timestamp); + let result = self.frame_at_internal(render_ctx, presentation_timestamp, video_data); match result { Ok(()) => { @@ -248,6 +249,7 @@ impl VideoDecoder { &mut self, render_ctx: &RenderContext, presentation_timestamp: Time, + video_data: &[u8], ) -> Result<(), DecodingError> { re_tracing::profile_function!(); @@ -322,12 +324,12 @@ impl VideoDecoder { if requested_gop_idx != self.current_gop_idx { if self.current_gop_idx.saturating_add(1) == requested_gop_idx { // forward seek to next GOP - queue up the one _after_ requested - self.enqueue_gop(requested_gop_idx + 1)?; + self.enqueue_gop(requested_gop_idx + 1, video_data)?; } else { // forward seek by N>1 OR backward seek across GOPs - reset self.reset()?; - self.enqueue_gop(requested_gop_idx)?; - self.enqueue_gop(requested_gop_idx + 1)?; + self.enqueue_gop(requested_gop_idx, video_data)?; + self.enqueue_gop(requested_gop_idx + 1, video_data)?; } } else if requested_sample_idx != self.current_sample_idx { // special case: handle seeking backwards within a single GOP @@ -335,8 +337,8 @@ impl VideoDecoder { // while maintaining a buffer of only 2 GOPs if requested_sample_idx < self.current_sample_idx { self.reset()?; - self.enqueue_gop(requested_gop_idx)?; - self.enqueue_gop(requested_gop_idx + 1)?; + self.enqueue_gop(requested_gop_idx, video_data)?; + self.enqueue_gop(requested_gop_idx + 1, video_data)?; } } @@ -384,7 +386,7 @@ impl VideoDecoder { /// Enqueue all samples in the given GOP. /// /// Does nothing if the index is out of bounds. - fn enqueue_gop(&mut self, gop_idx: usize) -> Result<(), DecodingError> { + fn enqueue_gop(&mut self, gop_idx: usize, video_data: &[u8]) -> Result<(), DecodingError> { let Some(gop) = self.data.gops.get(gop_idx) else { return Ok(()); }; @@ -392,7 +394,7 @@ impl VideoDecoder { let samples = &self.data.samples[gop.range()]; for (i, sample) in samples.iter().enumerate() { - let chunk = self.data.get(sample).ok_or(DecodingError::BadData)?; + let chunk = sample.get(video_data).ok_or(DecodingError::BadData)?; let is_keyframe = i == 0; self.chunk_decoder.decode(chunk, is_keyframe)?; } diff --git a/crates/viewer/re_renderer/src/video/mod.rs b/crates/viewer/re_renderer/src/video/mod.rs index 75eca41bc7b5..db1ebe4d8c3c 100644 --- a/crates/viewer/re_renderer/src/video/mod.rs +++ b/crates/viewer/re_renderer/src/video/mod.rs @@ -203,6 +203,7 @@ impl Video { render_context: &RenderContext, decoder_stream_id: VideoDecodingStreamId, presentation_timestamp_s: f64, + video_data: &[u8], ) -> FrameDecodingResult { re_tracing::profile_function!(); @@ -233,7 +234,7 @@ impl Video { decoder_entry.frame_index = render_context.active_frame_idx(); decoder_entry .decoder - .frame_at(render_context, presentation_timestamp_s) + .frame_at(render_context, presentation_timestamp_s, video_data) } /// Removes all decoders that have been unused in the last frame. diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/videos.rs b/crates/viewer/re_space_view_spatial/src/visualizers/videos.rs index 12ece086966c..e452f2fd586c 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/videos.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/videos.rs @@ -167,7 +167,7 @@ impl VideoFrameReferenceVisualizer { let video_reference: EntityPath = video_references .and_then(|v| v.first().map(|e| e.as_str().into())) .unwrap_or_else(|| self.fallback_for(ctx).as_str().into()); - let video = latest_at_query_video_from_datastore(ctx.viewer_ctx, &video_reference); + let query_result = latest_at_query_video_from_datastore(ctx.viewer_ctx, &video_reference); let world_from_entity = spatial_ctx .transform_info @@ -179,7 +179,7 @@ impl VideoFrameReferenceVisualizer { // Note that this area is also used for the bounding box which is important for the 2D view to determine default bounds. let mut video_resolution = glam::vec2(1280.0, 720.0); - match video.as_ref().map(|v| v.as_ref()) { + match query_result { None => { self.show_video_error( ctx, @@ -191,80 +191,87 @@ impl VideoFrameReferenceVisualizer { ); } - Some(Ok(video)) => { - video_resolution = glam::vec2(video.width() as _, video.height() as _); - - match video.frame_at(render_ctx, decode_stream_id, video_timestamp.as_seconds()) { - Ok(VideoFrameTexture { - texture, - time_range: _, // TODO(emilk): maybe add to `PickableTexturedRect` and `PickingHitType::TexturedRect` so we can show on hover? - is_pending, - show_spinner, - }) => { - // Make sure to use the video instead of texture size here, - // since the texture may be a placeholder which doesn't have the full size yet. - let top_left_corner_position = - world_from_entity.transform_point3(glam::Vec3::ZERO); - let extent_u = - world_from_entity.transform_vector3(glam::Vec3::X * video_resolution.x); - let extent_v = - world_from_entity.transform_vector3(glam::Vec3::Y * video_resolution.y); - - if is_pending { - // Keep polling for a fresh texture - ctx.viewer_ctx.egui_ctx.request_repaint(); - } - - if show_spinner { - // Show loading rectangle: - self.data.loading_spinners.push(LoadingSpinner { - center: top_left_corner_position + 0.5 * (extent_u + extent_v), - half_extent_u: 0.5 * extent_u, - half_extent_v: 0.5 * extent_v, + Some((video, video_data)) => match video.as_ref() { + Ok(video) => { + video_resolution = glam::vec2(video.width() as _, video.height() as _); + + match video.frame_at( + render_ctx, + decode_stream_id, + video_timestamp.as_seconds(), + video_data.as_slice(), + ) { + Ok(VideoFrameTexture { + texture, + time_range: _, // TODO(emilk): maybe add to `PickableTexturedRect` and `PickingHitType::TexturedRect` so we can show on hover? + is_pending, + show_spinner, + }) => { + // Make sure to use the video instead of texture size here, + // since the texture may be a placeholder which doesn't have the full size yet. + let top_left_corner_position = + world_from_entity.transform_point3(glam::Vec3::ZERO); + let extent_u = world_from_entity + .transform_vector3(glam::Vec3::X * video_resolution.x); + let extent_v = world_from_entity + .transform_vector3(glam::Vec3::Y * video_resolution.y); + + if is_pending { + // Keep polling for a fresh texture + ctx.viewer_ctx.egui_ctx.request_repaint(); + } + + if show_spinner { + // Show loading rectangle: + self.data.loading_spinners.push(LoadingSpinner { + center: top_left_corner_position + 0.5 * (extent_u + extent_v), + half_extent_u: 0.5 * extent_u, + half_extent_v: 0.5 * extent_v, + }); + } + + let textured_rect = TexturedRect { + top_left_corner_position, + extent_u, + extent_v, + colormapped_texture: ColormappedTexture::from_unorm_rgba(texture), + options: RectangleOptions { + texture_filter_magnification: TextureFilterMag::Nearest, + texture_filter_minification: TextureFilterMin::Linear, + outline_mask: spatial_ctx.highlight.overall, + ..Default::default() + }, + }; + self.data.pickable_rects.push(PickableTexturedRect { + ent_path: entity_path.clone(), + textured_rect, + source_data: PickableRectSourceData::Video, }); } - let textured_rect = TexturedRect { - top_left_corner_position, - extent_u, - extent_v, - colormapped_texture: ColormappedTexture::from_unorm_rgba(texture), - options: RectangleOptions { - texture_filter_magnification: TextureFilterMag::Nearest, - texture_filter_minification: TextureFilterMin::Linear, - outline_mask: spatial_ctx.highlight.overall, - ..Default::default() - }, - }; - self.data.pickable_rects.push(PickableTexturedRect { - ent_path: entity_path.clone(), - textured_rect, - source_data: PickableRectSourceData::Video, - }); - } - - Err(err) => { - self.show_video_error( - ctx, - spatial_ctx, - world_from_entity, - err.to_string(), - video_resolution, - entity_path, - ); + Err(err) => { + self.show_video_error( + ctx, + spatial_ctx, + world_from_entity, + err.to_string(), + video_resolution, + entity_path, + ); + } } } - } - Some(Err(err)) => { - self.show_video_error( - ctx, - spatial_ctx, - world_from_entity, - err.to_string(), - video_resolution, - entity_path, - ); - } + Err(err) => { + self.show_video_error( + ctx, + spatial_ctx, + world_from_entity, + err.to_string(), + video_resolution, + entity_path, + ); + } + }, } if spatial_ctx.space_view_class_identifier == SpatialSpaceView2D::identifier() { @@ -394,7 +401,7 @@ impl VideoFrameReferenceVisualizer { fn latest_at_query_video_from_datastore( ctx: &ViewerContext<'_>, entity_path: &EntityPath, -) -> Option>> { +) -> Option<(Arc>, Blob)> { let query = ctx.current_query(); let results = ctx.recording().query_caches().latest_at( @@ -408,7 +415,7 @@ fn latest_at_query_video_from_datastore( let blob = results.component_instance::(0)?; let media_type = results.component_instance::(0); - Some(ctx.cache.entry(|c: &mut VideoCache| { + let video = ctx.cache.entry(|c: &mut VideoCache| { let debug_name = entity_path.to_string(); c.entry( debug_name, @@ -417,7 +424,8 @@ fn latest_at_query_video_from_datastore( media_type.as_ref(), ctx.app_options.video_decoder_hw_acceleration, ) - })) + }); + Some((video, blob)) } impl TypedComponentFallbackProvider