establish image formar and video codec

crates/kornia-image/src/image.rs

@@ -52,14 +52,6 @@ impl From<ImageSize> for [u32; 2] {
     }
 }
 
-/// The format of the image to write to the video file.
-pub enum ImageFormat {
-    /// MONO 8-bit unsigned integer format.
-    Mono8,
-    /// RGB 8-bit unsigned integer format.
-    Rgb8,
-}
-
 #[derive(Clone)]
 /// Represents an image with pixel data.
 ///

crates/kornia-image/src/lib.rs

@@ -11,4 +11,4 @@ pub mod error;
 pub mod ops;
 
 pub use crate::error::ImageError;
-pub use crate::image::{Image, ImageFormat, ImageSize};
+pub use crate::image::{Image, ImageSize};

crates/kornia-io/src/stream/error.rs

@@ -71,6 +71,10 @@ pub enum StreamCaptureError {
     /// An error occurred during GStreamer to send end of stream event.
     #[error("Error ocurred in the gstreamer flow")]
     GstreamerFlowError(#[from] gst::FlowError),
+
+    /// An error occurred during checking the image format.
+    #[error("Invalid image format: {0}")]
+    InvalidImageFormat(String),
 }
 
 // ensure that can be sent over threads

crates/kornia-io/src/stream/video.rs

@@ -2,21 +2,32 @@ use std::path::Path;
 
 use gst::prelude::*;
 
-use kornia_image::{Image, ImageFormat, ImageSize};
+use kornia_image::{Image, ImageSize};
 
 use super::StreamCaptureError;
 
 /// The codec to use for the video writer.
-pub enum VideoWriterCodec {
+pub enum VideoCodec {
     /// H.264 codec.
     H264,
 }
 
+/// The format of the image to write to the video file.
+///
+/// Usually will be the combination of the image format and the pixel type.
+pub enum ImageFormat {
+    /// 8-bit RGB format.
+    Rgb8,
+    /// 8-bit mono format.
+    Mono8,
+}
+
 /// A struct for writing video files.
 pub struct VideoWriter {
     pipeline: gst::Pipeline,
     appsrc: gst_app::AppSrc,
     fps: i32,
+    format: ImageFormat,
     counter: u64,
     handle: Option<std::thread::JoinHandle<()>>,
 }
@@ -33,7 +44,7 @@ impl VideoWriter {
     /// * `size` - The size of the video.
     pub fn new(
         path: impl AsRef<Path>,
-        codec: VideoWriterCodec,
+        codec: VideoCodec,
         format: ImageFormat,
         fps: i32,
         size: ImageSize,
@@ -43,7 +54,7 @@ impl VideoWriter {
         // TODO: Add support for other codecs
         #[allow(unreachable_patterns)]
         let _codec = match codec {
-            VideoWriterCodec::H264 => "x264enc",
+            VideoCodec::H264 => "x264enc",
             _ => {
                 return Err(StreamCaptureError::InvalidConfig(
                     "Unsupported codec".to_string(),
@@ -52,7 +63,7 @@ impl VideoWriter {
         };
 
         // TODO: Add support for other formats
-        let format = match format {
+        let format_str = match format {
             ImageFormat::Mono8 => "GRAY8",
             ImageFormat::Rgb8 => "RGB",
         };
@@ -83,7 +94,7 @@ impl VideoWriter {
         appsrc.set_format(gst::Format::Time);
 
         let caps = gst::Caps::builder("video/x-raw")
-            .field("format", format)
+            .field("format", format_str)
             .field("width", size.width as i32)
             .field("height", size.height as i32)
             .field("framerate", gst::Fraction::new(fps, 1))
@@ -98,6 +109,7 @@ impl VideoWriter {
             pipeline,
             appsrc,
             fps,
+            format,
             counter: 0,
             handle: None,
         })
@@ -160,13 +172,26 @@ impl VideoWriter {
     /// # Arguments
     ///
     /// * `img` - The image to write to the video file.
-    // TODO: support write_async
+    // TODO: explore supporting write_async
     pub fn write<const C: usize>(&mut self, img: &Image<u8, C>) -> Result<(), StreamCaptureError> {
         // check if the image channels are correct
-        if C != 1 && C != 3 {
-            return Err(StreamCaptureError::InvalidConfig(
-                "Invalid number of channels".to_string(),
-            ));
+        match self.format {
+            ImageFormat::Mono8 => {
+                if C != 1 {
+                    return Err(StreamCaptureError::InvalidImageFormat(format!(
+                        "Invalid number of channels: expected 1, got {}",
+                        C
+                    )));
+                }
+            }
+            ImageFormat::Rgb8 => {
+                if C != 3 {
+                    return Err(StreamCaptureError::InvalidImageFormat(format!(
+                        "Invalid number of channels: expected 3, got {}",
+                        C
+                    )));
+                }
+            }
         }
 
         // TODO: verify is there is a cheaper way to copy the buffer
@@ -199,7 +224,7 @@ impl Drop for VideoWriter {
 
 #[cfg(test)]
 mod tests {
-    use super::{ImageFormat, VideoWriter, VideoWriterCodec};
+    use super::{ImageFormat, VideoCodec, VideoWriter};
     use kornia_image::{Image, ImageSize};
 
     #[ignore = "need gstreamer in CI"]
@@ -215,13 +240,8 @@ mod tests {
             height: 4,
         };
 
-        let mut writer = VideoWriter::new(
-            &file_path,
-            VideoWriterCodec::H264,
-            ImageFormat::Rgb8,
-            30,
-            size,
-        )?;
+        let mut writer =
+            VideoWriter::new(&file_path, VideoCodec::H264, ImageFormat::Rgb8, 30, size)?;
         writer.start()?;
 
         let img = Image::<u8, 3>::new(size, vec![0; size.width * size.height * 3])?;
@@ -246,13 +266,8 @@ mod tests {
             height: 4,
         };
 
-        let mut writer = VideoWriter::new(
-            &file_path,
-            VideoWriterCodec::H264,
-            ImageFormat::Mono8,
-            30,
-            size,
-        )?;
+        let mut writer =
+            VideoWriter::new(&file_path, VideoCodec::H264, ImageFormat::Mono8, 30, size)?;
         writer.start()?;
 
         let img = Image::<u8, 1>::new(size, vec![0; size.width * size.height])?;

examples/video_write tasks/src/main.rs

@@ -4,8 +4,11 @@ use tokio::signal;
 use tokio::sync::Mutex;
 
 use kornia::{
-    image::{Image, ImageFormat, ImageSize},
-    io::stream::{video::VideoWriterCodec, V4L2CameraConfig, VideoWriter},
+    image::{Image, ImageSize},
+    io::stream::{
+        video::{ImageFormat, VideoCodec},
+        V4L2CameraConfig, VideoWriter,
+    },
 };
 
 #[derive(Parser)]
@@ -52,7 +55,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // start the video writer
     let video_writer = VideoWriter::new(
         args.output,
-        VideoWriterCodec::H264,
+        VideoCodec::H264,
         ImageFormat::Rgb8,
         args.fps,
         frame_size,

examples/video_write/src/main.rs

@@ -4,8 +4,11 @@ use tokio::signal;
 use tokio::sync::Mutex;
 
 use kornia::{
-    image::{ImageFormat, ImageSize},
-    io::stream::{video::VideoWriterCodec, V4L2CameraConfig, VideoWriter},
+    image::ImageSize,
+    io::stream::{
+        video::{ImageFormat, VideoCodec},
+        V4L2CameraConfig, VideoWriter,
+    },
 };
 
 #[derive(Parser)]
@@ -52,7 +55,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // start the video writer
     let video_writer = VideoWriter::new(
         args.output,
-        VideoWriterCodec::H264,
+        VideoCodec::H264,
         ImageFormat::Rgb8,
         args.fps,
         frame_size,