Native video support for AV1

Cargo.lock

@@ -538,6 +538,12 @@ dependencies = [
  "zbus",
 ]
 
+[[package]]
+name = "assert_matches"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b34d609dfbaf33d6889b2b7106d3ca345eacad44200913df5ba02bfd31d2ba9"
+
 [[package]]
 name = "async-broadcast"
 version = "0.5.1"
@@ -757,6 +763,26 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1181e1e0d1fce796a03db1ae795d67167da795f9cf4a39c37589e85ef57f26d3"
 
+[[package]]
+name = "atomig"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0eaf2a17b11f7923d0abe0e07377c7f0f1255a8267f05c8a9cba8dc039acc821"
+dependencies = [
+ "atomig-macro",
+]
+
+[[package]]
+name = "atomig-macro"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a10c8c98ca4c65e4bdd6f1506beb768671f8dce3f5df4dd7d14632b6ecc6f43"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "atspi"
 version = "0.19.0"
@@ -811,6 +837,20 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
+[[package]]
+name = "av-data"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "124ae24335161b3d2226594640a67903da0866e2591312591fc8ddad64c1b38c"
+dependencies = [
+ "byte-slice-cast",
+ "bytes",
+ "num-derive",
+ "num-rational",
+ "num-traits",
+ "thiserror",
+]
+
 [[package]]
 name = "az"
 version = "1.2.1"
@@ -945,6 +985,12 @@ version = "3.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1"
 
+[[package]]
+name = "byte-slice-cast"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3ac9f8b63eca6fd385229b3675f6cc0dc5c8a5c8a54a59d4f52ffd670d87b0c"
+
 [[package]]
 name = "bytecount"
 version = "0.6.7"
@@ -1785,6 +1831,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "econtext"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "18aade80d5e09429040243ce1143ddc08a92d7a22820ac512610410a4dd5214f"
+
 [[package]]
 name = "eframe"
 version = "0.29.1"
@@ -3521,6 +3573,15 @@ dependencies = [
  "unicode-xid",
 ]
 
+[[package]]
+name = "nasm-rs"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12fcfa1bd49e0342ec1d07ed2be83b59963e7acbeb9310e1bb2c07b69dadd959"
+dependencies = [
+ "jobserver",
+]
+
 [[package]]
 name = "natord"
 version = "1.0.9"
@@ -4103,9 +4164,9 @@ checksum = "bb813b8af86854136c6922af0598d719255ecb2179515e6e7730d468f05c9cae"
 
 [[package]]
 name = "parking_lot"
-version = "0.12.1"
+version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f"
+checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27"
 dependencies = [
  "lock_api",
  "parking_lot_core",
@@ -4129,9 +4190,9 @@ dependencies = [
 
 [[package]]
 name = "paste"
-version = "1.0.12"
+version = "1.0.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f746c4065a8fa3fe23974dd82f15431cc8d40779821001404d10d2e79ca7d79"
+checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
 
 [[package]]
 name = "pathdiff"
@@ -4666,6 +4727,37 @@ dependencies = [
  "rand",
 ]
 
+[[package]]
+name = "rav1d"
+version = "1.0.0"
+source = "git+https://github.com/rerun-io/rav1d?branch=emilk/dav1d-interface#d3acfb077eac7759f94279ec9a81edd7d0c409ed"
+dependencies = [
+ "assert_matches",
+ "atomig",
+ "av-data",
+ "bitflags 2.6.0",
+ "cc",
+ "cfg-if",
+ "libc",
+ "nasm-rs",
+ "parking_lot",
+ "paste",
+ "raw-cpuid",
+ "static_assertions",
+ "strum",
+ "to_method",
+ "zerocopy",
+]
+
+[[package]]
+name = "raw-cpuid"
+version = "11.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ab240315c661615f2ee9f0f2cd32d5a7343a84d5ebcccb99d46e6637565e7b0"
+dependencies = [
+ "bitflags 2.6.0",
+]
+
 [[package]]
 name = "raw-window-handle"
 version = "0.5.2"
@@ -4930,6 +5022,7 @@ name = "re_crash_handler"
 version = "0.19.0-alpha.1+dev"
 dependencies = [
  "backtrace",
+ "econtext",
  "itertools 0.13.0",
  "libc",
  "parking_lot",
@@ -5856,8 +5949,17 @@ dependencies = [
 name = "re_video"
 version = "0.19.0-alpha.1+dev"
 dependencies = [
+ "crossbeam",
+ "econtext",
+ "indicatif",
  "itertools 0.13.0",
+ "ordered-float",
+ "parking_lot",
+ "rav1d",
+ "re_log",
  "re_mp4",
+ "re_tracing",
+ "re_video",
  "thiserror",
 ]
 
@@ -7309,6 +7411,12 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
+[[package]]
+name = "to_method"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7c4ceeeca15c8384bbc3e011dbd8fccb7f068a440b752b7d9b32ceb0ca0e2e8"
+
 [[package]]
 name = "tobj"
 version = "4.0.0"
@@ -8641,18 +8749,19 @@ dependencies = [
 
 [[package]]
 name = "zerocopy"
-version = "0.7.31"
+version = "0.7.35"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1c4061bedbb353041c12f413700357bec76df2c7e2ca8e4df8bac24c6bf68e3d"
+checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
 dependencies = [
+ "byteorder",
  "zerocopy-derive",
 ]
 
 [[package]]
 name = "zerocopy-derive"
-version = "0.7.31"
+version = "0.7.35"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b3c129550b3e6de3fd0ba67ba5c81818f9805e58b8d7fee80a3a59d2c9fc601a"
+checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
 dependencies = [
  "proc-macro2",
  "quote",

Cargo.toml

@@ -164,6 +164,7 @@ criterion = "0.5"
 crossbeam = "0.8"
 directories = "5"
 document-features = "0.2.8"
+econtext = "0.2"                                                   # Prints error contexts on crashes
 ehttp = "0.5.0"
 enumset = "1.0.12"
 env_logger = { version = "0.10", default-features = false }
@@ -308,6 +309,12 @@ zip = { version = "0.6", default-features = false }
 # Our dev profile has some optimizations turned on, as well as debug assertions.
 [profile.dev]
 opt-level = 1 # Make debug builds run faster
+[profile.dev.package.re_video]
+opt-level = 2 # Speed up CPU-side chroma-upsampling (TODO(#7298): move to GPU)
+[profile.dev.package.rav1d]
+opt-level = 2 # Speed up AV1 video decoding. TODO(emilk): this doesn't work. rav1d is dog slow in dev builds.
+[profile.dev.package.wgpu]
+opt-level = 2 # Speed up copying video textures
 
 # panic = "abort" leads to better optimizations and smaller binaries (and is the default in Wasm anyways),
 # but it also means backtraces don't work with the `backtrace` library (https://github.com/rust-lang/backtrace-rs/issues/397).

crates/store/re_types/src/image.rs

@@ -262,6 +262,7 @@ fn test_find_non_empty_dim_indices() {
 
 // ----------------------------------------------------------------------------
 
+// TODO(jan): there is a duplicate of this function in `crates/store/re_video/src/decode/av1.rs`
 /// Returns sRGB from YUV color.
 ///
 /// This conversion mirrors the function of the same name in `crates/viewer/re_renderer/shader/decodings.wgsl`

crates/store/re_video/Cargo.toml

@@ -24,8 +24,37 @@ no-default-features = true
 features = ["all"]
 
 [features]
+default = []
+
+## Opt-in to native AV1 decoding.
+## You need to install [nasm](https://nasm.us/) to compile with this feature.
+av1 = ["dep:dav1d"]
 
 [dependencies]
+re_log.workspace = true
+re_tracing.workspace = true
+
+crossbeam.workspace = true
+econtext.workspace = true
 itertools.workspace = true
+ordered-float.workspace = true
+parking_lot.workspace = true
 re_mp4.workspace = true
 thiserror.workspace = true
+
+# Native dependencies:
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+
+# If this package fails to build, install `nasm` locally, or build through `pixi`.
+# See https://github.com/rerun-io/rav1d/pull/1
+dav1d = { git = "https://github.com/rerun-io/rav1d", branch = "emilk/dav1d-interface", package = "rav1d", optional = true } # TODO(#7588): publish this fork
+
+# dav1d = { version = "0.10.3",  optional = true } # Requires more things to build, but is fast in debug builds. Useful for development.
+
+[dev-dependencies]
+indicatif.workspace = true
+re_video = { workspace = true, features = ["av1"] } # For the `frames` example
+
+
+[[example]]
+name = "frames"

crates/store/re_video/examples/frames.rs

@@ -0,0 +1,102 @@
+//! Decodes an mp4 with AV1 in it to a folder of images.
+
+#![allow(clippy::unwrap_used)]
+
+use std::{
+    fs::{File, OpenOptions},
+    io::Write as _,
+    path::{Path, PathBuf},
+    sync::Arc,
+    time::{Duration, Instant},
+};
+
+use indicatif::ProgressBar;
+use parking_lot::Mutex;
+
+use re_video::demux::mp4::load_mp4;
+
+fn main() {
+    // frames <video.mp4>
+    let args: Vec<_> = std::env::args().collect();
+    let Some(video_path) = args.get(1) else {
+        println!("Usage: frames <video.mp4>");
+        return;
+    };
+    let output_dir = PathBuf::new().join(Path::new(video_path).with_extension(""));
+
+    println!("Decoding {video_path}");
+
+    let video = std::fs::read(video_path).expect("failed to read video");
+    let video = load_mp4(&video).expect("failed to load video");
+
+    println!(
+        "{} {}x{}",
+        video.segments.len(),
+        video.config.coded_width,
+        video.config.coded_height
+    );
+
+    let progress = ProgressBar::new(video.samples.len() as u64).with_message("Decoding video");
+    progress.enable_steady_tick(Duration::from_millis(100));
+
+    let frames = Arc::new(Mutex::new(Vec::new()));
+    let mut decoder = re_video::decode::av1::Decoder::new("debug_name".to_owned(), {
+        let frames = frames.clone();
+        let progress = progress.clone();
+        move |frame| {
+            progress.inc(1);
+            frames.lock().push(frame);
+        }
+    });
+
+    let start = Instant::now();
+    for sample in &video.samples {
+        decoder.decode(video.get(sample).unwrap());
+    }
+
+    decoder.flush();
+    drop(decoder);
+    let end = Instant::now();
+    progress.finish();
+
+    let frames = frames.lock();
+
+    println!(
+        "Decoded {} frames in {:.2}ms",
+        frames.len(),
+        end.duration_since(start).as_secs_f64() * 1000.0
+    );
+
+    println!("Writing frames to {}", output_dir.display());
+    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() {
+        if let Ok(frame) = frame {
+            let mut file = OpenOptions::new()
+                .write(true)
+                .create(true)
+                .truncate(true)
+                .open(output_dir.join(format!("{i:0width$}.ppm")))
+                .expect("failed to open file");
+            write_binary_ppm(&mut file, frame.width, frame.height, &frame.data);
+        }
+    }
+}
+
+fn num_digits(n: usize) -> usize {
+    (n as f64).log10().floor() as usize + 1
+}
+
+fn write_binary_ppm(file: &mut File, width: u32, height: u32, rgba: &[u8]) {
+    let header = format!("P6\n{width} {height}\n255\n");
+
+    let mut data = Vec::with_capacity(header.len() + width as usize * height as usize * 3);
+    data.extend_from_slice(header.as_bytes());
+
+    for rgba in rgba.chunks(4) {
+        data.extend_from_slice(&[rgba[0], rgba[1], rgba[2]]);
+    }
+
+    file.write_all(&data).expect("failed to write frame data");
+}