From 22913cbc94d63c2caf159e0f2a4efcc70007e34d Mon Sep 17 00:00:00 2001 From: Scott Lamb Date: Tue, 16 Jul 2024 09:29:42 -0700 Subject: [PATCH] rkmpp encoder example --- Cargo.lock | 152 +++++++++++++++++++++++++++++++-- rkmpp/Cargo.toml | 4 + rkmpp/examples/encode.rs | 176 +++++++++++++++++++++++++++++++++++++++ rkmpp/src/encoder.rs | 69 ++++++++++++++- 4 files changed, 391 insertions(+), 10 deletions(-) create mode 100644 rkmpp/examples/encode.rs diff --git a/Cargo.lock b/Cargo.lock index 0bee893..71f5ef7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,12 +17,55 @@ dependencies = [ "winapi", ] +[[package]] +name = "anstream" +version = "0.6.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418c75fa768af9c03be99d17643f93f79bbba589895012a80e3452a19ddda15b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + [[package]] name = "anstyle" version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7079075b41f533b8c61d2a4d073c4676e1f8b249ff94a393b0595db304e0dd87" +[[package]] +name = "anstyle-parse" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c03a11a9034d92058ceb6ee011ce58af4a9bf61491aa7e1e59ecd24bd40d22d4" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad186efb764318d35165f1758e7dcef3b10628e26d41a44bc5550652e6804391" +dependencies = [ + "windows-sys 0.52.0", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61a38449feb7068f52bb06c12759005cf459ee52bb4adc1d5a7c4322d716fb19" +dependencies = [ + "anstyle", + "windows-sys 0.52.0", +] + [[package]] name = "anyhow" version = "1.0.81" @@ -228,7 +271,7 @@ dependencies = [ "ansi_term", "atty", "bitflags 1.3.2", - "strsim", + "strsim 0.8.0", "textwrap 0.11.0", "unicode-width", "vec_map", @@ -236,28 +279,43 @@ dependencies = [ [[package]] name = "clap" -version = "4.4.11" +version = "4.5.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfaff671f6b22ca62406885ece523383b9b64022e341e53e009a62ebc47a45f2" +checksum = "64acc1846d54c1fe936a78dc189c34e28d3f5afc348403f28ecf53660b9b8462" dependencies = [ "clap_builder", + "clap_derive", ] [[package]] name = "clap_builder" -version = "4.4.11" +version = "4.5.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a216b506622bb1d316cd51328dce24e07bdff4a6128a47c7e7fad11878d5adbb" +checksum = "6fb8393d67ba2e7bfaf28a23458e4e2b543cc73a99595511eb207fdb8aede942" dependencies = [ + "anstream", "anstyle", "clap_lex", + "strsim 0.11.1", +] + +[[package]] +name = "clap_derive" +version = "4.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bac35c6dafb060fd4d275d9a4ffae97917c13a6327903a8be2153cd964f7085" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.58", ] [[package]] name = "clap_lex" -version = "0.6.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1" +checksum = "4b82cf0babdbd58558212896d1a4272303a57bdb245c2bf1147185fb45640e70" [[package]] name = "cmake" @@ -268,6 +326,25 @@ dependencies = [ "cc", ] +[[package]] +name = "colorchoice" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b6a852b24ab71dffc585bcb46eaf7959d175cb865a7152e35b348d1b2960422" + +[[package]] +name = "console" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e1f83fc076bd6dd27517eacdf25fef6c4dfe5f1d7448bafaaf3a26f13b5e4eb" +dependencies = [ + "encode_unicode", + "lazy_static", + "libc", + "unicode-width", + "windows-sys 0.52.0", +] + [[package]] name = "cookie" version = "0.16.2" @@ -377,7 +454,7 @@ dependencies = [ "anes", "cast", "ciborium", - "clap 4.4.11", + "clap 4.5.9", "criterion-plot", "is-terminal", "itertools", @@ -453,6 +530,12 @@ version = "1.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91" +[[package]] +name = "encode_unicode" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" + [[package]] name = "encoding_rs" version = "0.8.32" @@ -672,6 +755,12 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.1.19" @@ -809,6 +898,19 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "indicatif" +version = "0.17.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "763a5a8f45087d6bcea4222e7b72c291a054edf80e4ef6efd2a4979878c7bea3" +dependencies = [ + "console", + "instant", + "number_prefix", + "portable-atomic", + "unicode-width", +] + [[package]] name = "instant" version = "0.1.12" @@ -846,6 +948,12 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8478577c03552c21db0e2724ffb8986a5ce7af88107e6be5d2ee6e158c12800" + [[package]] name = "itertools" version = "0.10.5" @@ -1085,6 +1193,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "once_cell" version = "1.17.1" @@ -1199,6 +1313,12 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "portable-atomic" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0" + [[package]] name = "pretty-hex" version = "0.4.1" @@ -1395,6 +1515,8 @@ name = "rkmpp" version = "0.1.0" dependencies = [ "av-traits", + "clap 4.5.9", + "indicatif", "rkmpp-sys", "snafu", ] @@ -1578,7 +1700,7 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "080c44971436b1af15d6f61ddd8b543995cf63ab8e677d46b00cc06f4ef267a0" dependencies = [ - "heck", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.58", @@ -1622,6 +1744,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "1.0.109" @@ -1922,6 +2050,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "v4l2" version = "0.1.0" diff --git a/rkmpp/Cargo.toml b/rkmpp/Cargo.toml index 38d134d..5677d51 100644 --- a/rkmpp/Cargo.toml +++ b/rkmpp/Cargo.toml @@ -7,3 +7,7 @@ edition = "2021" av-traits = { path = "../av-traits" } snafu = { version = "0.8.0", default-features = false } rkmpp-sys = { path = "rkmpp-sys" } + +[dev-dependencies] +clap = { version = "4.5.9", features = ["derive"] } +indicatif = "0.17.8" diff --git a/rkmpp/examples/encode.rs b/rkmpp/examples/encode.rs new file mode 100644 index 0000000..9d21c40 --- /dev/null +++ b/rkmpp/examples/encode.rs @@ -0,0 +1,176 @@ +//! Encodes raw video to H.264/H.265. + +use av_traits::VideoEncoder as _; +use clap::Parser; +use std::io::{Read as _, Write as _}; +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +#[derive(Parser)] +struct Args { + #[arg(long)] + width: u16, + + #[arg(long)] + height: u16, + + #[arg(long)] + fps: f64, + + #[arg(long)] + bitrate: Option, + + #[arg(long)] + keyframe_interval: Option, + + #[arg(long)] + codec: rkmpp::RkMppEncoderCodec, + + #[arg(long)] + pixel_format: rkmpp::RkMppEncoderInputFormat, + + #[arg(long)] + input: PathBuf, + + #[arg(long)] + output: PathBuf, +} + +struct Frame { + data: Vec, + plane_offsets: [usize; 4], +} + +impl Frame { + fn new(width: u16, height: u16, format: rkmpp::RkMppEncoderInputFormat) -> Self { + match format { + rkmpp::RkMppEncoderInputFormat::Yuv420Planar => { + let luma_dims = usize::from(width) * usize::from(height); + // TODO: not right for odd widths. + let chroma_dims = luma_dims >> 2; + let plane_offsets = [0, luma_dims, luma_dims + chroma_dims, luma_dims + 2 * chroma_dims]; + Self { + data: vec![0; luma_dims + 2 * chroma_dims], + plane_offsets, + } + } + _ => unimplemented!(), + } + } +} + +impl av_traits::RawVideoFrame for Frame { + fn samples(&self, plane: usize) -> &[u8] { + &self.data[self.plane_offsets[plane]..self.plane_offsets[plane + 1]] + } +} + +fn main() { + let args = Args::parse(); + + let mut infile = std::fs::File::open(&args.input).unwrap(); + let mut outfile = std::fs::File::create(&args.output).unwrap(); + + let config = rkmpp::RkMppEncoderConfig { + width: args.width, + height: args.height, + fps: args.fps, + bitrate: args.bitrate, + keyframe_interval: args.keyframe_interval, + codec: args.codec, + input_format: args.pixel_format, + }; + let mut encoder = rkmpp::RkMppEncoder::new(config).unwrap(); + + let len = infile.metadata().unwrap().len(); + + // yuck. + let frame_len = u64::try_from(Frame::new(args.width, args.height, args.pixel_format).data.len()).unwrap(); + let frames_in = len / frame_len; + if len % frame_len != 0 { + panic!("Input length {len} is not an even multiple of frame length {frame_len}; check width/height/format"); + } + + let mut frames_out = 0; + let pb = + indicatif::ProgressBar::new(frames_in).with_style(indicatif::ProgressStyle::with_template("encoding at {msg} {elapsed} {wide_bar} {eta}").unwrap()); + + let (raw_frame_tx, raw_frame_rx) = std::sync::mpsc::sync_channel(3); + let (encoded_frame_tx, encoded_frame_rx) = std::sync::mpsc::sync_channel::>(3); + std::thread::scope(|s| { + let pb = &pb; + let before = Instant::now(); + std::thread::Builder::new() + .name("reader".to_owned()) + .spawn_scoped(s, move || { + let mut last = before; + let mut reading = Duration::ZERO; + let mut sending = Duration::ZERO; + for _ in 0..frames_in { + let mut frame = Frame::new(args.width, args.height, args.pixel_format); + infile.read_exact(&mut frame.data[..]).unwrap(); + let pre_send = Instant::now(); + reading += pre_send.duration_since(last); + raw_frame_tx.send(frame).unwrap(); + last = Instant::now(); + sending += last.duration_since(pre_send); + } + drop(raw_frame_tx); + pb.println(format!("read thread: reading={reading:?} sending={sending:?}")); + }) + .unwrap(); + + std::thread::Builder::new() + .name("writer".to_owned()) + .spawn_scoped(s, move || { + let mut last = Instant::now(); + let mut receiving = Duration::ZERO; + let mut writing = Duration::ZERO; + while let Ok(frame) = encoded_frame_rx.recv() { + let pre_write = Instant::now(); + receiving += pre_write.duration_since(last); + outfile.write_all(&frame.encoded_frame.unwrap().data).unwrap(); + last = Instant::now(); + frames_out += 1; + pb.set_message(format!("{:.02}x", (frames_out as f64) / args.fps / (last - before).as_secs_f64())); + pb.inc(1); + } + let pre_drop = Instant::now(); + receiving += pre_drop.duration_since(last); + drop(outfile); + writing += Instant::now().duration_since(pre_drop); + pb.println(format!("write thread: receiving={receiving:?} writing={writing:?}")); + }) + .unwrap(); + + let mut encoding = Duration::ZERO; + let mut frame_type = av_traits::EncodedFrameType::Key; // first frame must be key. + while let Ok(frame) = raw_frame_rx.recv() { + let pre = Instant::now(); + let output = encoder.encode(frame, frame_type).unwrap(); + encoding += Instant::now().duration_since(pre); + if let Some(o) = output { + encoded_frame_tx.send(o).unwrap(); + pb.inc(1); + frames_out += 1; + } + frame_type = av_traits::EncodedFrameType::Auto; + } + loop { + let pre = Instant::now(); + let output = encoder.flush().unwrap(); + encoding += Instant::now().duration_since(pre); + let Some(output) = output else { + break; + }; + encoded_frame_tx.send(output).unwrap(); + pb.inc(1); + frames_out += 1; + } + drop(encoded_frame_tx); + let total = Instant::now().duration_since(before); + pb.println(format!("encode thread: encoding={encoding:?} total={total:?}")); + }); + assert_eq!(frames_in, frames_out); + pb.finish(); +} diff --git a/rkmpp/src/encoder.rs b/rkmpp/src/encoder.rs index 4913dd8..8992f02 100644 --- a/rkmpp/src/encoder.rs +++ b/rkmpp/src/encoder.rs @@ -26,12 +26,44 @@ pub struct RkMppEncoder { frames_emitted: u64, } +#[derive(Debug)] +pub struct BadParameterError { + name: &'static str, + value: String, +} + +impl std::fmt::Display for BadParameterError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "bad rkmpp encoder {}={:?}", self.name, &self.value) + } +} + +impl std::error::Error for BadParameterError {} + #[derive(Clone, Copy, Debug)] +#[non_exhaustive] pub enum RkMppEncoderInputFormat { Yuv420Planar, Bgra, } +impl std::str::FromStr for RkMppEncoderInputFormat { + type Err = BadParameterError; + + fn from_str(s: &str) -> core::result::Result { + Ok(match s { + "yuv420p" => Self::Yuv420Planar, + "bgra" => Self::Bgra, + _ => { + return Err(BadParameterError { + name: "input-format", + value: s.to_owned(), + }) + } + }) + } +} + impl RkMppEncoderInputFormat { fn to_mpp(self) -> sys::MppFrameFormat { match self { @@ -48,6 +80,39 @@ impl RkMppEncoderInputFormat { } } +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub enum RkMppEncoderCodec { + H264, + H265, +} + +impl std::str::FromStr for RkMppEncoderCodec { + type Err = BadParameterError; + + fn from_str(s: &str) -> core::result::Result { + Ok(match s { + "h264" => Self::H264, + "h265" => Self::H265, + _ => { + return Err(BadParameterError { + name: "codec", + value: s.to_owned(), + }) + } + }) + } +} + +impl RkMppEncoderCodec { + fn to_mpp(self) -> sys::MppFrameFormat { + match self { + Self::H264 => sys::MppCodingType_MPP_VIDEO_CodingAVC, + Self::H265 => sys::MppCodingType_MPP_VIDEO_CodingHEVC, + } + } +} + #[derive(Clone, Debug)] pub struct RkMppEncoderConfig { pub height: u16, @@ -56,6 +121,7 @@ pub struct RkMppEncoderConfig { pub bitrate: Option, pub keyframe_interval: Option, pub input_format: RkMppEncoderInputFormat, + pub codec: RkMppEncoderCodec, } impl RkMppEncoderConfig { @@ -78,7 +144,7 @@ impl RkMppEncoder { pub fn new(config: RkMppEncoderConfig) -> Result { let lib = mpp::Lib::new().ok_or(RkMppEncoderError::LibraryUnavailable)?; - let mut context = lib.new_context(sys::MppCodingType_MPP_VIDEO_CodingAVC)?; + let mut context = lib.new_context(config.codec.to_mpp())?; let mut mpp_config = lib.new_config()?; context.get_config(&mut mpp_config)?; @@ -250,6 +316,7 @@ mod test { bitrate: Some(50000), keyframe_interval: Some(120), input_format: RkMppEncoderInputFormat::Yuv420Planar, + codec: RkMppEncoderCodec::H264, }) .unwrap();