From c35fbc1dbffab9d57697cdea3c811992e62e3d5e Mon Sep 17 00:00:00 2001 From: edgar Date: Tue, 17 Sep 2024 01:20:09 +0200 Subject: [PATCH 1/7] video is recording --- crates/kornia-io/Cargo.toml | 3 +- crates/kornia-io/src/stream/mod.rs | 4 + crates/kornia-io/src/stream/video.rs | 162 +++++++++++++++++++++++++++ examples/video_write/Cargo.toml | 14 +++ examples/video_write/README.md | 21 ++++ examples/video_write/src/main.rs | 94 ++++++++++++++++ 6 files changed, 297 insertions(+), 1 deletion(-) create mode 100644 crates/kornia-io/src/stream/video.rs create mode 100644 examples/video_write/Cargo.toml create mode 100644 examples/video_write/README.md create mode 100644 examples/video_write/src/main.rs diff --git a/crates/kornia-io/Cargo.toml b/crates/kornia-io/Cargo.toml index c6648841..78a3f71f 100644 --- a/crates/kornia-io/Cargo.toml +++ b/crates/kornia-io/Cargo.toml @@ -26,6 +26,7 @@ thiserror = "1" futures = { version = "0.3.1", optional = true } gst = { version = "0.23.0", package = "gstreamer", optional = true } gst-app = { version = "0.23.0", package = "gstreamer-app", optional = true } +gst-video = { version = "0.23.0", package = "gstreamer-video", optional = true } memmap2 = "0.9.4" tokio = { version = "1", features = ["full"], optional = true } turbojpeg = { version = "1.0.0", optional = true } @@ -36,7 +37,7 @@ tempfile = "3.10" kornia = { workspace = true, features = ["jpegturbo"] } [features] -gstreamer = ["futures", "gst", "gst-app", "tokio"] +gstreamer = ["futures", "gst", "gst-app", "tokio", "gst-video"] jpegturbo = ["turbojpeg"] [[bench]] diff --git a/crates/kornia-io/src/stream/mod.rs b/crates/kornia-io/src/stream/mod.rs index 73572544..1cd11975 100644 --- a/crates/kornia-io/src/stream/mod.rs +++ b/crates/kornia-io/src/stream/mod.rs @@ -13,8 +13,12 @@ pub mod rtsp; /// A module for capturing video streams from v4l2 cameras. pub mod v4l2; +/// A module for capturing video streams from video files. +pub mod video; + pub use crate::stream::camera::{CameraCapture, CameraCaptureConfig}; pub use crate::stream::capture::StreamCapture; pub use crate::stream::error::StreamCaptureError; pub use crate::stream::rtsp::RTSPCameraConfig; pub use crate::stream::v4l2::V4L2CameraConfig; +pub use crate::stream::video::VideoWriter; diff --git a/crates/kornia-io/src/stream/video.rs b/crates/kornia-io/src/stream/video.rs new file mode 100644 index 00000000..de232de7 --- /dev/null +++ b/crates/kornia-io/src/stream/video.rs @@ -0,0 +1,162 @@ +use std::{path::Path, sync::Arc}; + +use futures::prelude::*; +use gst::{buffer, prelude::*}; + +use kornia_image::{Image, ImageSize}; +use tokio::sync::Mutex; + +use super::StreamCaptureError; + +/// A struct for writing video files. +pub struct VideoWriter { + pipeline: gst::Pipeline, + appsrc: gst_app::AppSrc, + fps: f32, + counter: u64, + handle: Option>, +} + +impl VideoWriter { + /// Create a new VideoWriter. + /// + /// # Arguments + /// + /// * `path` - The path to save the video file. + /// * `fps` - The frames per second of the video. + /// * `size` - The size of the video. + pub fn new(path: &Path, fps: f32, size: ImageSize) -> Result { + gst::init()?; + + let pipeline_str = format!( + "appsrc name=src ! \ + videoconvert ! video/x-raw,format=I420 ! \ + x264enc ! \ + video/x-h264,profile=main ! \ + h264parse ! \ + mp4mux ! \ + filesink location={}", + path.to_string_lossy() + ); + + println!("Pipeline: {}", pipeline_str); + + let pipeline = gst::parse::launch(&pipeline_str)? + .dynamic_cast::() + .map_err(StreamCaptureError::DowncastPipelineError)?; + + let appsrc = pipeline + .by_name("src") + .unwrap() + .dynamic_cast::() + .unwrap(); + + appsrc.set_format(gst::Format::Time); + + let caps = gst::Caps::builder("video/x-raw") + .field("format", "RGB") + .field("width", size.width as i32) + .field("height", size.height as i32) + .field("framerate", gst::Fraction::new(fps as i32, 1)) + .build(); + + appsrc.set_caps(Some(&caps)); + + appsrc.set_is_live(true); + appsrc.set_property("block", false); + + Ok(Self { + pipeline, + appsrc, + fps, + counter: 0, + handle: None, + }) + } + + /// Start the video writer + pub fn start(&mut self) -> Result<(), StreamCaptureError> { + self.pipeline.set_state(gst::State::Playing)?; + + let bus = self.pipeline.bus().ok_or(StreamCaptureError::BusError)?; + let mut messages = bus.stream(); + + let handle = tokio::spawn(async move { + while let Some(msg) = messages.next().await { + match msg.view() { + gst::MessageView::Eos(..) => { + println!("EOS"); + break; + } + gst::MessageView::Error(err) => { + eprintln!( + "Error from {:?}: {} ({:?})", + msg.src().map(|s| s.path_string()), + err.error(), + err.debug() + ); + } + _ => {} + } + } + }); + + self.handle = Some(handle); + + Ok(()) + } + + /// Stop the video writer + pub fn stop(&mut self) -> Result<(), StreamCaptureError> { + // Send end of stream to the appsrc + self.appsrc + .end_of_stream() + .expect("Failed to send end of stream"); + + // Take the handle and await it + // TODO: This is a blocking call, we need to make it non-blocking + if let Some(handle) = self.handle.take() { + tokio::task::block_in_place(|| { + tokio::runtime::Handle::current().block_on(async { + if let Err(e) = handle.await { + eprintln!("Error waiting for handle: {:?}", e); + } + }); + }); + } + + // Set the pipeline to null + self.pipeline.set_state(gst::State::Null)?; + + Ok(()) + } + + /// Write an image to the video file. + /// + /// # Arguments + /// + /// * `img` - The image to write to the video file. + pub fn write(&mut self, img: Image) -> Result<(), StreamCaptureError> { + let mut buffer = gst::Buffer::with_size(img.storage.len())?; + + let pts = gst::ClockTime::from_nseconds(self.counter * 1_000_000_000 / self.fps as u64); + let duration = gst::ClockTime::from_nseconds(1_000_000_000 / self.fps as u64); + + { + let buffer_ref = buffer.get_mut().expect("Failed to get buffer"); + buffer_ref.set_pts(Some(pts)); + buffer_ref.set_duration(Some(duration)); + + let mut map = buffer_ref.map_writable()?; + map.copy_from_slice(img.as_slice()); + } + + self.counter += 1; + + if let Err(err) = self.appsrc.push_buffer(buffer) { + return Err(StreamCaptureError::InvalidConfig(err.to_string())); + } + + Ok(()) + } +} diff --git a/examples/video_write/Cargo.toml b/examples/video_write/Cargo.toml new file mode 100644 index 00000000..3f8d5d00 --- /dev/null +++ b/examples/video_write/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "video_write" +version = "0.1.0" +authors = ["Edgar Riba "] +license = "Apache-2.0" +edition = "2021" +publish = false + +[dependencies] +clap = { version = "4.5.4", features = ["derive"] } +ctrlc = "3.4.4" +kornia = { workspace = true, features = ["gstreamer"] } +rerun = "0.18" +tokio = { version = "1" } diff --git a/examples/video_write/README.md b/examples/video_write/README.md new file mode 100644 index 00000000..d762593e --- /dev/null +++ b/examples/video_write/README.md @@ -0,0 +1,21 @@ +An example showing how to use the webcam with the `kornia::io` module with the ability to cancel the feed after a certain amount of time. This example will display the webcam feed in a [`rerun`](https://github.com/rerun-io/rerun) window. + +NOTE: This example requires the gstremer backend to be enabled. To enable the gstreamer backend, use the `gstreamer` feature flag when building the `kornia` crate and its dependencies. + +```bash +Usage: webcam [OPTIONS] + +Options: + -c, --camera-id [default: 0] + -f, --fps [default: 30] + -d, --duration + -h, --help Print help +``` + +Example: + +```bash +cargo run --bin webcam --release -- --camera-id 0 --duration 5 --fps 30 +``` + +![Screenshot from 2024-08-28 18-33-56](https://github.com/user-attachments/assets/783619e4-4867-48bc-b7d2-d32a133e4f5a) diff --git a/examples/video_write/src/main.rs b/examples/video_write/src/main.rs new file mode 100644 index 00000000..8ff993d0 --- /dev/null +++ b/examples/video_write/src/main.rs @@ -0,0 +1,94 @@ +use clap::Parser; +use std::{ + path::Path, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, +}; +use tokio::signal; + +use kornia::{ + image::{ops, Image, ImageSize}, + imgproc, + io::{ + fps_counter::FpsCounter, + stream::{StreamCaptureError, V4L2CameraConfig, VideoWriter}, + }, +}; + +#[derive(Parser)] +struct Args { + #[arg(short, long, default_value = "0")] + camera_id: u32, + + #[arg(short, long, default_value = "30")] + fps: u32, + + #[arg(short, long)] + duration: Option, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let args = Args::parse(); + + // start the recording stream + let rec = rerun::RecordingStreamBuilder::new("Kornia Webcapture App").spawn()?; + + // allocate the image buffers + let frame_size = ImageSize { + width: 640, + height: 480, + }; + + // create a webcam capture object with camera id 0 + // and force the image size to 640x480 + let webcam = V4L2CameraConfig::new() + .with_camera_id(args.camera_id) + .with_fps(args.fps) + .with_size(frame_size) + .build()?; + + // start the video writer + let mut video_writer = VideoWriter::new(Path::new("output.mp4"), args.fps as f32, frame_size)?; + video_writer.start()?; + + let video_writer = Arc::new(Mutex::new(video_writer)); + + // start grabbing frames from the camera + webcam + .run_with_termination( + |img| { + let rec = rec.clone(); + let video_writer = video_writer.clone(); + async move { + // write the image to the video writer + video_writer.lock().unwrap().write(img)?; + //println!("Wrote frame"); + + // log the image + //rec.log_static( + // "image", + // &rerun::Image::from_elements( + // img.as_slice(), + // img.size().into(), + // rerun::ColorModel::RGB, + // ), + //)?; + + Ok(()) + } + }, + async { + signal::ctrl_c().await.expect("Failed to listen for Ctrl+C"); + println!("👋 Finished recording. Closing app."); + }, + ) + .await?; + + // stop the video writer + video_writer.lock().unwrap().stop()?; + + Ok(()) +} From f6f96f685bf9453df3169c6865e50c11912c5826 Mon Sep 17 00:00:00 2001 From: edgar Date: Tue, 17 Sep 2024 12:09:24 +0200 Subject: [PATCH 2/7] improvements --- crates/kornia-io/Cargo.toml | 3 +- crates/kornia-io/src/stream/video.rs | 47 +++++++-- examples/video_write tasks/Cargo.toml | 14 +++ examples/video_write tasks/README.md | 21 ++++ examples/video_write tasks/src/main.rs | 127 +++++++++++++++++++++++++ examples/video_write/src/main.rs | 60 ++++++------ 6 files changed, 228 insertions(+), 44 deletions(-) create mode 100644 examples/video_write tasks/Cargo.toml create mode 100644 examples/video_write tasks/README.md create mode 100644 examples/video_write tasks/src/main.rs diff --git a/crates/kornia-io/Cargo.toml b/crates/kornia-io/Cargo.toml index 78a3f71f..c6648841 100644 --- a/crates/kornia-io/Cargo.toml +++ b/crates/kornia-io/Cargo.toml @@ -26,7 +26,6 @@ thiserror = "1" futures = { version = "0.3.1", optional = true } gst = { version = "0.23.0", package = "gstreamer", optional = true } gst-app = { version = "0.23.0", package = "gstreamer-app", optional = true } -gst-video = { version = "0.23.0", package = "gstreamer-video", optional = true } memmap2 = "0.9.4" tokio = { version = "1", features = ["full"], optional = true } turbojpeg = { version = "1.0.0", optional = true } @@ -37,7 +36,7 @@ tempfile = "3.10" kornia = { workspace = true, features = ["jpegturbo"] } [features] -gstreamer = ["futures", "gst", "gst-app", "tokio", "gst-video"] +gstreamer = ["futures", "gst", "gst-app", "tokio"] jpegturbo = ["turbojpeg"] [[bench]] diff --git a/crates/kornia-io/src/stream/video.rs b/crates/kornia-io/src/stream/video.rs index de232de7..127917f5 100644 --- a/crates/kornia-io/src/stream/video.rs +++ b/crates/kornia-io/src/stream/video.rs @@ -1,18 +1,23 @@ -use std::{path::Path, sync::Arc}; +use std::path::Path; use futures::prelude::*; -use gst::{buffer, prelude::*}; +use gst::prelude::*; use kornia_image::{Image, ImageSize}; -use tokio::sync::Mutex; use super::StreamCaptureError; +/// The codec to use for the video writer. +pub enum VideoWriterCodec { + /// H.264 codec. + H264, +} + /// A struct for writing video files. pub struct VideoWriter { pipeline: gst::Pipeline, appsrc: gst_app::AppSrc, - fps: f32, + fps: i32, counter: u64, handle: Option>, } @@ -23,11 +28,28 @@ impl VideoWriter { /// # Arguments /// /// * `path` - The path to save the video file. + /// * `codec` - The codec to use for the video writer. /// * `fps` - The frames per second of the video. /// * `size` - The size of the video. - pub fn new(path: &Path, fps: f32, size: ImageSize) -> Result { + pub fn new( + path: &Path, + codec: VideoWriterCodec, + fps: i32, + size: ImageSize, + ) -> Result { gst::init()?; + // TODO: Add support for other codecs + #[allow(unreachable_patterns)] + let _codec = match codec { + VideoWriterCodec::H264 => "x264enc", + _ => { + return Err(StreamCaptureError::InvalidConfig( + "Unsupported codec".to_string(), + )) + } + }; + let pipeline_str = format!( "appsrc name=src ! \ videoconvert ! video/x-raw,format=I420 ! \ @@ -39,8 +61,6 @@ impl VideoWriter { path.to_string_lossy() ); - println!("Pipeline: {}", pipeline_str); - let pipeline = gst::parse::launch(&pipeline_str)? .dynamic_cast::() .map_err(StreamCaptureError::DowncastPipelineError)?; @@ -49,7 +69,7 @@ impl VideoWriter { .by_name("src") .unwrap() .dynamic_cast::() - .unwrap(); + .map_err(StreamCaptureError::DowncastPipelineError)?; appsrc.set_format(gst::Format::Time); @@ -57,7 +77,7 @@ impl VideoWriter { .field("format", "RGB") .field("width", size.width as i32) .field("height", size.height as i32) - .field("framerate", gst::Fraction::new(fps as i32, 1)) + .field("framerate", gst::Fraction::new(fps, 1)) .build(); appsrc.set_caps(Some(&caps)); @@ -136,7 +156,8 @@ impl VideoWriter { /// # Arguments /// /// * `img` - The image to write to the video file. - pub fn write(&mut self, img: Image) -> Result<(), StreamCaptureError> { + // TODO: support write_async + pub fn write(&mut self, img: &Image) -> Result<(), StreamCaptureError> { let mut buffer = gst::Buffer::with_size(img.storage.len())?; let pts = gst::ClockTime::from_nseconds(self.counter * 1_000_000_000 / self.fps as u64); @@ -160,3 +181,9 @@ impl VideoWriter { Ok(()) } } + +impl Drop for VideoWriter { + fn drop(&mut self) { + self.stop().unwrap(); + } +} diff --git a/examples/video_write tasks/Cargo.toml b/examples/video_write tasks/Cargo.toml new file mode 100644 index 00000000..995b9a17 --- /dev/null +++ b/examples/video_write tasks/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "video_write_tasks" +version = "0.1.0" +authors = ["Edgar Riba "] +license = "Apache-2.0" +edition = "2021" +publish = false + +[dependencies] +clap = { version = "4.5.4", features = ["derive"] } +ctrlc = "3.4.4" +kornia = { workspace = true, features = ["gstreamer"] } +rerun = "0.18" +tokio = { version = "1" } diff --git a/examples/video_write tasks/README.md b/examples/video_write tasks/README.md new file mode 100644 index 00000000..d762593e --- /dev/null +++ b/examples/video_write tasks/README.md @@ -0,0 +1,21 @@ +An example showing how to use the webcam with the `kornia::io` module with the ability to cancel the feed after a certain amount of time. This example will display the webcam feed in a [`rerun`](https://github.com/rerun-io/rerun) window. + +NOTE: This example requires the gstremer backend to be enabled. To enable the gstreamer backend, use the `gstreamer` feature flag when building the `kornia` crate and its dependencies. + +```bash +Usage: webcam [OPTIONS] + +Options: + -c, --camera-id [default: 0] + -f, --fps [default: 30] + -d, --duration + -h, --help Print help +``` + +Example: + +```bash +cargo run --bin webcam --release -- --camera-id 0 --duration 5 --fps 30 +``` + +![Screenshot from 2024-08-28 18-33-56](https://github.com/user-attachments/assets/783619e4-4867-48bc-b7d2-d32a133e4f5a) diff --git a/examples/video_write tasks/src/main.rs b/examples/video_write tasks/src/main.rs new file mode 100644 index 00000000..5adf2582 --- /dev/null +++ b/examples/video_write tasks/src/main.rs @@ -0,0 +1,127 @@ +use clap::Parser; +use std::{path::Path, sync::Arc}; +use tokio::signal; +use tokio::sync::Mutex; + +use kornia::{ + image::{Image, ImageSize}, + io::stream::{video::VideoWriterCodec, V4L2CameraConfig, VideoWriter}, +}; + +#[derive(Parser)] +struct Args { + #[arg(short, long, default_value = "0")] + camera_id: u32, + + #[arg(short, long, default_value = "30")] + fps: i32, + + #[arg(short, long)] + duration: Option, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let args = Args::parse(); + + // start the recording stream + let rec = rerun::RecordingStreamBuilder::new("Kornia Video Write App").spawn()?; + + // allocate the image buffers + let frame_size = ImageSize { + width: 640, + height: 480, + }; + + // create a webcam capture object with camera id 0 + // and force the image size to 640x480 + let webcam = V4L2CameraConfig::new() + .with_camera_id(args.camera_id) + .with_fps(args.fps as u32) + .with_size(frame_size) + .build()?; + + // start the video writer + let video_writer = VideoWriter::new( + Path::new("output.mp4"), + VideoWriterCodec::H264, + args.fps, + frame_size, + )?; + let video_writer = Arc::new(Mutex::new(video_writer)); + video_writer.lock().await.start()?; + + // Create a channel to send frames to the video writer + let (tx, rx) = tokio::sync::mpsc::channel::>>>(32); + let rx = Arc::new(Mutex::new(rx)); + + // Spawn a task to read frames from the camera and send them to the video writer + let video_writer_task = tokio::spawn({ + let rx = rx.clone(); + let video_writer = video_writer.clone(); + async move { + while let Some(img) = rx.lock().await.recv().await { + // lock the image and write it to the video writer + let img = img.lock().await; + video_writer + .lock() + .await + .write(&img) + .expect("Failed to write image to video writer"); + } + Ok::<_, Box>(()) + } + }); + + // Visualization thread + let visualization_task = tokio::spawn({ + let rec = rec.clone(); + let rx = rx.clone(); + async move { + while let Some(img) = rx.lock().await.recv().await { + // lock the image and log it + let img = img.lock().await; + rec.log_static( + "image", + &rerun::Image::from_elements( + img.as_slice(), + img.size().into(), + rerun::ColorModel::RGB, + ), + )?; + } + Ok::<_, Box>(()) + } + }); + + // start grabbing frames from the camera + let capture = webcam.run_with_termination( + |img| { + let tx = tx.clone(); + async move { + // send the image to the video writer and the visualization + tx.send(Arc::new(Mutex::new(img))).await?; + Ok(()) + } + }, + async { + signal::ctrl_c().await.expect("Failed to listen for Ctrl+C"); + println!("👋 Finished recording. Closing app."); + }, + ); + + tokio::select! { + _ = capture => (), + _ = video_writer_task => (), + _ = visualization_task => (), + _ = signal::ctrl_c() => (), + } + + video_writer + .lock() + .await + .stop() + .expect("Failed to stop video writer"); + + Ok(()) +} diff --git a/examples/video_write/src/main.rs b/examples/video_write/src/main.rs index 8ff993d0..e505238a 100644 --- a/examples/video_write/src/main.rs +++ b/examples/video_write/src/main.rs @@ -1,20 +1,11 @@ use clap::Parser; -use std::{ - path::Path, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, Mutex, - }, -}; +use std::{path::Path, sync::Arc}; use tokio::signal; +use tokio::sync::Mutex; use kornia::{ - image::{ops, Image, ImageSize}, - imgproc, - io::{ - fps_counter::FpsCounter, - stream::{StreamCaptureError, V4L2CameraConfig, VideoWriter}, - }, + image::ImageSize, + io::stream::{video::VideoWriterCodec, V4L2CameraConfig, VideoWriter}, }; #[derive(Parser)] @@ -23,7 +14,7 @@ struct Args { camera_id: u32, #[arg(short, long, default_value = "30")] - fps: u32, + fps: i32, #[arg(short, long)] duration: Option, @@ -34,7 +25,7 @@ async fn main() -> Result<(), Box> { let args = Args::parse(); // start the recording stream - let rec = rerun::RecordingStreamBuilder::new("Kornia Webcapture App").spawn()?; + let rec = rerun::RecordingStreamBuilder::new("Kornia Video Write App").spawn()?; // allocate the image buffers let frame_size = ImageSize { @@ -46,15 +37,19 @@ async fn main() -> Result<(), Box> { // and force the image size to 640x480 let webcam = V4L2CameraConfig::new() .with_camera_id(args.camera_id) - .with_fps(args.fps) + .with_fps(args.fps as u32) .with_size(frame_size) .build()?; // start the video writer - let mut video_writer = VideoWriter::new(Path::new("output.mp4"), args.fps as f32, frame_size)?; - video_writer.start()?; - + let video_writer = VideoWriter::new( + Path::new("output.mp4"), + VideoWriterCodec::H264, + args.fps, + frame_size, + )?; let video_writer = Arc::new(Mutex::new(video_writer)); + video_writer.lock().await.start()?; // start grabbing frames from the camera webcam @@ -64,19 +59,17 @@ async fn main() -> Result<(), Box> { let video_writer = video_writer.clone(); async move { // write the image to the video writer - video_writer.lock().unwrap().write(img)?; - //println!("Wrote frame"); + video_writer.lock().await.write(&img)?; // log the image - //rec.log_static( - // "image", - // &rerun::Image::from_elements( - // img.as_slice(), - // img.size().into(), - // rerun::ColorModel::RGB, - // ), - //)?; - + rec.log_static( + "image", + &rerun::Image::from_elements( + img.as_slice(), + img.size().into(), + rerun::ColorModel::RGB, + ), + )?; Ok(()) } }, @@ -87,8 +80,11 @@ async fn main() -> Result<(), Box> { ) .await?; - // stop the video writer - video_writer.lock().unwrap().stop()?; + video_writer + .lock() + .await + .stop() + .expect("Failed to stop video writer"); Ok(()) } From 80ad61ab25a4ee5d6a1566e8ab447fb0faa594ec Mon Sep 17 00:00:00 2001 From: edgar Date: Wed, 18 Sep 2024 21:01:32 +0200 Subject: [PATCH 3/7] fix clippy --- crates/kornia-imgproc/benches/bench_flip.rs | 30 ++++++++++++++------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/crates/kornia-imgproc/benches/bench_flip.rs b/crates/kornia-imgproc/benches/bench_flip.rs index 96a3bc71..813c41bc 100644 --- a/crates/kornia-imgproc/benches/bench_flip.rs +++ b/crates/kornia-imgproc/benches/bench_flip.rs @@ -86,8 +86,11 @@ fn bench_flip(c: &mut Criterion) { BenchmarkId::new("par_par_slicecopy", ¶meter_string), &(&image_f32, &output), |b, i| { - let (src, mut dst) = (i.0.clone(), i.1.clone()); - b.iter(|| black_box(par_par_slicecopy(&src, &mut dst))) + let (src, mut dst) = (i.0, i.1.clone()); + b.iter(|| { + par_par_slicecopy(black_box(src), black_box(&mut dst)); + black_box(()) + }) }, ); @@ -95,8 +98,11 @@ fn bench_flip(c: &mut Criterion) { BenchmarkId::new("par_loop_loop", ¶meter_string), &(&image_f32, &output), |b, i| { - let (src, mut dst) = (i.0.clone(), i.1.clone()); - b.iter(|| black_box(par_loop_loop(&src, &mut dst))) + let (src, mut dst) = (i.0, i.1.clone()); + b.iter(|| { + par_loop_loop(black_box(src), black_box(&mut dst)); + black_box(()) + }) }, ); @@ -104,8 +110,11 @@ fn bench_flip(c: &mut Criterion) { BenchmarkId::new("par_loop_slicecopy", ¶meter_string), &(&image_f32, &output), |b, i| { - let (src, mut dst) = (i.0.clone(), i.1.clone()); - b.iter(|| black_box(par_loop_slicecopy(&src, &mut dst))) + let (src, mut dst) = (i.0, i.1.clone()); + b.iter(|| { + par_loop_slicecopy(black_box(src), black_box(&mut dst)); + black_box(()) + }) }, ); @@ -113,8 +122,11 @@ fn bench_flip(c: &mut Criterion) { BenchmarkId::new("par_seq_slicecopy", ¶meter_string), &(&image_f32, &output), |b, i| { - let (src, mut dst) = (i.0.clone(), i.1.clone()); - b.iter(|| black_box(par_seq_slicecopy(&src, &mut dst))) + let (src, mut dst) = (i.0, i.1.clone()); + b.iter(|| { + par_seq_slicecopy(black_box(src), black_box(&mut dst)); + black_box(()) + }) }, ); @@ -123,7 +135,7 @@ fn bench_flip(c: &mut Criterion) { &(&image_f32, &output), |b, i| { let (src, mut dst) = (i.0, i.1.clone()); - b.iter(|| black_box(flip::horizontal_flip(src, &mut dst))) + b.iter(|| flip::horizontal_flip(black_box(src), black_box(&mut dst))) }, ); } From a07d26ab8d9ba5e41f4a6c694111db0cb5a67533 Mon Sep 17 00:00:00 2001 From: edgar Date: Wed, 18 Sep 2024 21:21:25 +0200 Subject: [PATCH 4/7] improve errors --- crates/kornia-io/src/stream/capture.rs | 2 +- crates/kornia-io/src/stream/error.rs | 8 ++++++-- crates/kornia-io/src/stream/video.rs | 22 ++++++++++------------ 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/crates/kornia-io/src/stream/capture.rs b/crates/kornia-io/src/stream/capture.rs index 7a94c4fa..22219b18 100644 --- a/crates/kornia-io/src/stream/capture.rs +++ b/crates/kornia-io/src/stream/capture.rs @@ -161,7 +161,7 @@ impl StreamCapture { fn get_appsink(&self) -> Result { self.pipeline .by_name("sink") - .ok_or_else(|| StreamCaptureError::DowncastAppSinkError)? + .ok_or_else(|| StreamCaptureError::GetElementByNameError)? .dynamic_cast::() .map_err(StreamCaptureError::DowncastPipelineError) } diff --git a/crates/kornia-io/src/stream/error.rs b/crates/kornia-io/src/stream/error.rs index c3f79eff..3333ebd0 100644 --- a/crates/kornia-io/src/stream/error.rs +++ b/crates/kornia-io/src/stream/error.rs @@ -10,8 +10,8 @@ pub enum StreamCaptureError { DowncastPipelineError(gst::Element), /// An error occurred during GStreamer downcast of appsink. - #[error("Failed to downcast appsink")] - DowncastAppSinkError, + #[error("Failed to get an element by name")] + GetElementByNameError, /// An error occurred during GStreamer to get the bus. #[error("Failed to get the bus")] @@ -67,4 +67,8 @@ pub enum StreamCaptureError { /// An error for an invalid configuration. #[error("Invalid configuration: {0}")] InvalidConfig(String), + + /// An error occurred during GStreamer to send end of stream event. + #[error("Error ocurred in the gstreamer flow")] + GstreamerFlowError(#[from] gst::FlowError), } diff --git a/crates/kornia-io/src/stream/video.rs b/crates/kornia-io/src/stream/video.rs index 127917f5..1fbb4e64 100644 --- a/crates/kornia-io/src/stream/video.rs +++ b/crates/kornia-io/src/stream/video.rs @@ -67,7 +67,7 @@ impl VideoWriter { let appsrc = pipeline .by_name("src") - .unwrap() + .ok_or_else(|| StreamCaptureError::GetElementByNameError)? .dynamic_cast::() .map_err(StreamCaptureError::DowncastPipelineError)?; @@ -131,7 +131,7 @@ impl VideoWriter { // Send end of stream to the appsrc self.appsrc .end_of_stream() - .expect("Failed to send end of stream"); + .map_err(StreamCaptureError::GstreamerFlowError)?; // Take the handle and await it // TODO: This is a blocking call, we need to make it non-blocking @@ -158,19 +158,15 @@ impl VideoWriter { /// * `img` - The image to write to the video file. // TODO: support write_async pub fn write(&mut self, img: &Image) -> Result<(), StreamCaptureError> { - let mut buffer = gst::Buffer::with_size(img.storage.len())?; + // TODO: verify is there is a cheaper way to copy the buffer + let mut buffer = gst::Buffer::from_mut_slice(img.as_slice().to_vec()); let pts = gst::ClockTime::from_nseconds(self.counter * 1_000_000_000 / self.fps as u64); let duration = gst::ClockTime::from_nseconds(1_000_000_000 / self.fps as u64); - { - let buffer_ref = buffer.get_mut().expect("Failed to get buffer"); - buffer_ref.set_pts(Some(pts)); - buffer_ref.set_duration(Some(duration)); - - let mut map = buffer_ref.map_writable()?; - map.copy_from_slice(img.as_slice()); - } + let buffer_ref = buffer.get_mut().expect("Failed to get buffer"); + buffer_ref.set_pts(Some(pts)); + buffer_ref.set_duration(Some(duration)); self.counter += 1; @@ -184,6 +180,8 @@ impl VideoWriter { impl Drop for VideoWriter { fn drop(&mut self) { - self.stop().unwrap(); + self.stop().unwrap_or_else(|e| { + eprintln!("Error stopping video writer: {:?}", e); + }); } } From 69e4ca24b19acd739c27d99125035ec392a7ffaf Mon Sep 17 00:00:00 2001 From: edgar Date: Wed, 18 Sep 2024 21:35:24 +0200 Subject: [PATCH 5/7] add fake test --- crates/kornia-io/src/stream/video.rs | 30 ++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/crates/kornia-io/src/stream/video.rs b/crates/kornia-io/src/stream/video.rs index 1fbb4e64..f2b5eaec 100644 --- a/crates/kornia-io/src/stream/video.rs +++ b/crates/kornia-io/src/stream/video.rs @@ -185,3 +185,33 @@ impl Drop for VideoWriter { }); } } + +#[cfg(test)] +mod tests { + use super::{VideoWriter, VideoWriterCodec}; + use kornia_image::{Image, ImageSize}; + + #[test] + #[ignore = "TODO: fix this test as there's a race condition in the gstreamer flow"] + fn video_writer() -> Result<(), Box> { + let tmp_dir = tempfile::tempdir()?; + std::fs::create_dir_all(tmp_dir.path())?; + + let file_path = tmp_dir.path().join("test.mp4"); + + let size = ImageSize { + width: 6, + height: 4, + }; + let mut writer = VideoWriter::new(&file_path, VideoWriterCodec::H264, 30, size)?; + writer.start()?; + + let img = Image::new(size, vec![0; size.width * size.height * 3])?; + writer.write(&img)?; + writer.stop()?; + + assert!(file_path.exists(), "File does not exist: {:?}", file_path); + + Ok(()) + } +} From fe28f335704a13eb9821fd027e7becd9546ab9c4 Mon Sep 17 00:00:00 2001 From: edgar Date: Wed, 18 Sep 2024 21:48:58 +0200 Subject: [PATCH 6/7] improve example output --- crates/kornia-io/src/stream/video.rs | 4 +++- examples/video_write/README.md | 8 ++++---- examples/video_write/src/main.rs | 20 ++++++++++---------- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/crates/kornia-io/src/stream/video.rs b/crates/kornia-io/src/stream/video.rs index f2b5eaec..116bd232 100644 --- a/crates/kornia-io/src/stream/video.rs +++ b/crates/kornia-io/src/stream/video.rs @@ -32,7 +32,7 @@ impl VideoWriter { /// * `fps` - The frames per second of the video. /// * `size` - The size of the video. pub fn new( - path: &Path, + path: impl AsRef, codec: VideoWriterCodec, fps: i32, size: ImageSize, @@ -50,6 +50,8 @@ impl VideoWriter { } }; + let path = path.as_ref().to_owned(); + let pipeline_str = format!( "appsrc name=src ! \ videoconvert ! video/x-raw,format=I420 ! \ diff --git a/examples/video_write/README.md b/examples/video_write/README.md index d762593e..eaf3c218 100644 --- a/examples/video_write/README.md +++ b/examples/video_write/README.md @@ -1,21 +1,21 @@ -An example showing how to use the webcam with the `kornia::io` module with the ability to cancel the feed after a certain amount of time. This example will display the webcam feed in a [`rerun`](https://github.com/rerun-io/rerun) window. +An example showing how to write a video file using the `kornia::io` module along with the webcam capture example. Visualizes the webcam feed in a [`rerun`](https://github.com/rerun-io/rerun) window. NOTE: This example requires the gstremer backend to be enabled. To enable the gstreamer backend, use the `gstreamer` feature flag when building the `kornia` crate and its dependencies. ```bash -Usage: webcam [OPTIONS] +Usage: video_write [OPTIONS] --output Options: + -o, --output -c, --camera-id [default: 0] -f, --fps [default: 30] - -d, --duration -h, --help Print help ``` Example: ```bash -cargo run --bin webcam --release -- --camera-id 0 --duration 5 --fps 30 +cargo run --bin video_write --release -- --output ~/output.mp4 ``` ![Screenshot from 2024-08-28 18-33-56](https://github.com/user-attachments/assets/783619e4-4867-48bc-b7d2-d32a133e4f5a) diff --git a/examples/video_write/src/main.rs b/examples/video_write/src/main.rs index e505238a..a7f5db20 100644 --- a/examples/video_write/src/main.rs +++ b/examples/video_write/src/main.rs @@ -1,5 +1,5 @@ use clap::Parser; -use std::{path::Path, sync::Arc}; +use std::{path::PathBuf, sync::Arc}; use tokio::signal; use tokio::sync::Mutex; @@ -10,20 +10,25 @@ use kornia::{ #[derive(Parser)] struct Args { + #[arg(short, long)] + output: PathBuf, + #[arg(short, long, default_value = "0")] camera_id: u32, #[arg(short, long, default_value = "30")] fps: i32, - - #[arg(short, long)] - duration: Option, } #[tokio::main] async fn main() -> Result<(), Box> { let args = Args::parse(); + // Ensure the output path ends with .mp4 + if args.output.extension().and_then(|ext| ext.to_str()) != Some("mp4") { + return Err("Output file must have a .mp4 extension".into()); + } + // start the recording stream let rec = rerun::RecordingStreamBuilder::new("Kornia Video Write App").spawn()?; @@ -42,12 +47,7 @@ async fn main() -> Result<(), Box> { .build()?; // start the video writer - let video_writer = VideoWriter::new( - Path::new("output.mp4"), - VideoWriterCodec::H264, - args.fps, - frame_size, - )?; + let video_writer = VideoWriter::new(args.output, VideoWriterCodec::H264, args.fps, frame_size)?; let video_writer = Arc::new(Mutex::new(video_writer)); video_writer.lock().await.start()?; From 0e04083dd791a8ad5f35ac44234d90d626a305a6 Mon Sep 17 00:00:00 2001 From: edgar Date: Wed, 18 Sep 2024 21:52:30 +0200 Subject: [PATCH 7/7] added readme for the video_write_tasks example --- examples/video_write tasks/README.md | 11 +++++------ examples/video_write tasks/src/main.rs | 17 ++++++++++------- examples/video_write/README.md | 2 -- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/examples/video_write tasks/README.md b/examples/video_write tasks/README.md index d762593e..b80929e1 100644 --- a/examples/video_write tasks/README.md +++ b/examples/video_write tasks/README.md @@ -1,21 +1,20 @@ -An example showing how to use the webcam with the `kornia::io` module with the ability to cancel the feed after a certain amount of time. This example will display the webcam feed in a [`rerun`](https://github.com/rerun-io/rerun) window. +Example showing how to write a video using different background tasks. NOTE: This example requires the gstremer backend to be enabled. To enable the gstreamer backend, use the `gstreamer` feature flag when building the `kornia` crate and its dependencies. ```bash -Usage: webcam [OPTIONS] +Usage: video_write_tasks [OPTIONS] --output Options: + -o, --output -c, --camera-id [default: 0] -f, --fps [default: 30] -d, --duration - -h, --help Print help + -h, --help Print help Print help ``` Example: ```bash -cargo run --bin webcam --release -- --camera-id 0 --duration 5 --fps 30 +cargo run --bin video_write_tasks --release -- --output output.mp4 ``` - -![Screenshot from 2024-08-28 18-33-56](https://github.com/user-attachments/assets/783619e4-4867-48bc-b7d2-d32a133e4f5a) diff --git a/examples/video_write tasks/src/main.rs b/examples/video_write tasks/src/main.rs index 5adf2582..02ef76ac 100644 --- a/examples/video_write tasks/src/main.rs +++ b/examples/video_write tasks/src/main.rs @@ -1,5 +1,5 @@ use clap::Parser; -use std::{path::Path, sync::Arc}; +use std::{path::PathBuf, sync::Arc}; use tokio::signal; use tokio::sync::Mutex; @@ -10,6 +10,9 @@ use kornia::{ #[derive(Parser)] struct Args { + #[arg(short, long)] + output: PathBuf, + #[arg(short, long, default_value = "0")] camera_id: u32, @@ -24,6 +27,11 @@ struct Args { async fn main() -> Result<(), Box> { let args = Args::parse(); + // Ensure the output path ends with .mp4 + if args.output.extension().and_then(|ext| ext.to_str()) != Some("mp4") { + return Err("Output file must have a .mp4 extension".into()); + } + // start the recording stream let rec = rerun::RecordingStreamBuilder::new("Kornia Video Write App").spawn()?; @@ -42,12 +50,7 @@ async fn main() -> Result<(), Box> { .build()?; // start the video writer - let video_writer = VideoWriter::new( - Path::new("output.mp4"), - VideoWriterCodec::H264, - args.fps, - frame_size, - )?; + let video_writer = VideoWriter::new(args.output, VideoWriterCodec::H264, args.fps, frame_size)?; let video_writer = Arc::new(Mutex::new(video_writer)); video_writer.lock().await.start()?; diff --git a/examples/video_write/README.md b/examples/video_write/README.md index eaf3c218..35027082 100644 --- a/examples/video_write/README.md +++ b/examples/video_write/README.md @@ -17,5 +17,3 @@ Example: ```bash cargo run --bin video_write --release -- --output ~/output.mp4 ``` - -![Screenshot from 2024-08-28 18-33-56](https://github.com/user-attachments/assets/783619e4-4867-48bc-b7d2-d32a133e4f5a)