video/png: Add PNG encoder element

It can encode raw video formats like Gray8/16, RGB and RGBA and
uses the PNG crate which is a decoding and encoding library in pure Rust

6 files changed, 632 insertions, 0 deletions

diff --git a/video/rspng/Cargo.toml b/video/rspng/Cargo.toml
new file mode 100644
index 000000000..8039145ca
--- /dev/null
+++ b/video/rspng/Cargo.toml
@@ -0,0 +1,26 @@
+[package]
+name = "gst-plugin-rspng"
+version = "0.1.0"
+authors = ["Natanael Mojica <neithanmo@gmail.com>"]
+repository = "https://gitlab.freedesktop.org/gstreamer/gst-plugins-rs"
+license = "MIT/Apache-2.0"
+edition = "2018"
+description = "An PNG encoder/decoder written in pure Rust"
+
+[dependencies]
+glib = { git = "https://github.com/gtk-rs/glib" }
+gst = { package = "gstreamer", git = "https://gitlab.freedesktop.org/gstreamer/gstreamer-rs" }
+gst_video = { package = "gstreamer-video", git = "https://gitlab.freedesktop.org/gstreamer/gstreamer-rs" }
+gst_check = { package = "gstreamer-check", git = "https://gitlab.freedesktop.org/gstreamer/gstreamer-rs" }
+png = "0.16.3"
+once_cell = "1"
+parking_lot = "0.10.2"
+atomic_refcell = "0.1"
+
+[lib]
+name = "gstrspng"
+crate-type = ["cdylib", "rlib"]
+path = "src/lib.rs"
+
+[build-dependencies]
+gst-plugin-version-helper = { path="../../version-helper" }
diff --git a/video/rspng/build.rs b/video/rspng/build.rs
new file mode 100644
index 000000000..17be1215e
--- /dev/null
+++ b/video/rspng/build.rs
@@ -0,0 +1,3 @@
+fn main() {
+    gst_plugin_version_helper::get_info()
+}
diff --git a/video/rspng/examples/pngenc.rs b/video/rspng/examples/pngenc.rs
new file mode 100644
index 000000000..23deb8b25
--- /dev/null
+++ b/video/rspng/examples/pngenc.rs
@@ -0,0 +1,45 @@
+// Copyright (C) 2020 Natanael Mojica <neithanmo@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use gst::prelude::*;
+
+const ENCODE_PIPELINE: &str = "videotestsrc is-live=false num-buffers=1 ! videoconvert ! video/x-raw, format=RGB, width=160, height=120 ! 
+    rspngenc compression-level=2 filter=4 ! filesink location=frame.png";
+
+fn main() {
+    gst::init().unwrap();
+    gstrspng::plugin_register_static().expect("Failed to register gif plugin");
+
+    let pipeline = gst::parse_launch(ENCODE_PIPELINE).unwrap();
+    let bus = pipeline.get_bus().unwrap();
+
+    pipeline
+        .set_state(gst::State::Playing)
+        .expect("Failed to set pipeline state to playing");
+
+    for msg in bus.iter_timed(gst::CLOCK_TIME_NONE) {
+        use gst::MessageView;
+
+        match msg.view() {
+            MessageView::Eos(..) => break,
+            MessageView::Error(err) => {
+                println!(
+                    "Error from {:?}: {} ({:?})",
+                    err.get_src().map(|s| s.get_path_string()),
+                    err.get_error(),
+                    err.get_debug()
+                );
+                break;
+            }
+            _ => (),
+        }
+    }
+    pipeline
+        .set_state(gst::State::Null)
+        .expect("Failed to set pipeline state to null");
+}
diff --git a/video/rspng/src/lib.rs b/video/rspng/src/lib.rs
new file mode 100644
index 000000000..826f87abd
--- /dev/null
+++ b/video/rspng/src/lib.rs
@@ -0,0 +1,28 @@
+// Copyright (C) 2020 Natanael Mojica <neithanmo@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use gst::gst_plugin_define;
+
+mod pngenc;
+
+fn plugin_init(plugin: &gst::Plugin) -> Result<(), glib::BoolError> {
+    pngenc::register(plugin)?;
+    Ok(())
+}
+
+gst_plugin_define!(
+    rspng,
+    env!("CARGO_PKG_DESCRIPTION"),
+    plugin_init,
+    concat!(env!("CARGO_PKG_VERSION"), "-", env!("COMMIT_ID")),
+    "MIT/X11",
+    env!("CARGO_PKG_NAME"),
+    env!("CARGO_PKG_NAME"),
+    env!("CARGO_PKG_REPOSITORY"),
+    env!("BUILD_REL_DATE")
+);
diff --git a/video/rspng/src/pngenc.rs b/video/rspng/src/pngenc.rs
new file mode 100644
index 000000000..e3e946593
--- /dev/null
+++ b/video/rspng/src/pngenc.rs
@@ -0,0 +1,440 @@
+// Copyright (C) 2020 Natanael Mojica <neithanmo@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use std::{io, io::Write, sync::Arc};
+
+use glib::subclass;
+use glib::subclass::prelude::*;
+use glib::{glib_object_impl, glib_object_subclass, gobject_sys, GEnum};
+
+use gst::prelude::*;
+use gst::subclass::prelude::*;
+use gst::{gst_debug, gst_element_error, gst_error, gst_loggable_error};
+use gst_video::prelude::*;
+use gst_video::subclass::prelude::*;
+
+use atomic_refcell::AtomicRefCell;
+use once_cell::sync::Lazy;
+use parking_lot::Mutex;
+
+const DEFAULT_COMPRESSION_LEVEL: CompressionLevel = CompressionLevel::Default;
+const DEFAULT_FILTER_TYPE: FilterType = FilterType::NoFilter;
+
+#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Clone, Copy, GEnum)]
+#[repr(u32)]
+#[genum(type_name = "GstRsPngCompressionLevel")]
+pub(crate) enum CompressionLevel {
+    #[genum(name = "Default: Use the default compression level.", nick = "default")]
+    Default,
+    #[genum(name = "Fast: A fast compression algorithm.", nick = "fast")]
+    Fast,
+    #[genum(
+        name = "Best: Uses the algorithm with the best results.",
+        nick = "best"
+    )]
+    Best,
+    #[genum(name = "Huffman: Huffman compression.", nick = "huffman")]
+    Huffman,
+    #[genum(name = "Rle: Rle compression.", nick = "rle")]
+    Rle,
+}
+
+#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Clone, Copy, GEnum)]
+#[repr(u32)]
+#[genum(type_name = "GstRsPngFilterType")]
+pub(crate) enum FilterType {
+    #[genum(
+        name = "NoFilter: No filtering applied to the output.",
+        nick = "nofilter"
+    )]
+    NoFilter,
+    #[genum(name = "Sub: filter applied to each pixel.", nick = "sub")]
+    Sub,
+    #[genum(name = "Up: Up filter similar to Sub.", nick = "up")]
+    Up,
+    #[genum(
+        name = "Avg: The Average filter uses the average of the two neighboring pixels.",
+        nick = "avg"
+    )]
+    Avg,
+    #[genum(
+        name = "Paeth: The Paeth filter computes a simple linear function of the three neighboring pixels.",
+        nick = "paeth"
+    )]
+    Paeth,
+}
+
+impl From<CompressionLevel> for png::Compression {
+    fn from(value: CompressionLevel) -> Self {
+        match value {
+            CompressionLevel::Default => png::Compression::Default,
+            CompressionLevel::Fast => png::Compression::Fast,
+            CompressionLevel::Best => png::Compression::Best,
+            CompressionLevel::Huffman => png::Compression::Huffman,
+            CompressionLevel::Rle => png::Compression::Rle,
+        }
+    }
+}
+
+impl From<FilterType> for png::FilterType {
+    fn from(value: FilterType) -> Self {
+        match value {
+            FilterType::NoFilter => png::FilterType::NoFilter,
+            FilterType::Sub => png::FilterType::Sub,
+            FilterType::Up => png::FilterType::Up,
+            FilterType::Avg => png::FilterType::Avg,
+            FilterType::Paeth => png::FilterType::Paeth,
+        }
+    }
+}
+
+static CAT: Lazy<gst::DebugCategory> = Lazy::new(|| {
+    gst::DebugCategory::new(
+        "rspngenc",
+        gst::DebugColorFlags::empty(),
+        Some("PNG encoder"),
+    )
+});
+
+// Inner buffer where the result of frame encoding  is written
+// before relay them downstream
+struct CacheBuffer {
+    buffer: AtomicRefCell<Vec<u8>>,
+}
+
+impl CacheBuffer {
+    pub fn new() -> Self {
+        Self {
+            buffer: AtomicRefCell::new(Vec::new()),
+        }
+    }
+
+    pub fn clear(&self) {
+        self.buffer.borrow_mut().clear();
+    }
+
+    pub fn write(&self, buf: &[u8]) {
+        let mut buffer = self.buffer.borrow_mut();
+        buffer.extend_from_slice(buf);
+    }
+
+    pub fn consume(&self) -> Vec<u8> {
+        let mut buffer = self.buffer.borrow_mut();
+        std::mem::replace(&mut *buffer, Vec::new())
+    }
+}
+// The Encoder requires a Writer, so we use here and intermediate structure
+// for caching encoded frames
+struct CacheWriter {
+    cache: Arc<CacheBuffer>,
+}
+
+impl CacheWriter {
+    pub fn new(cache: Arc<CacheBuffer>) -> Self {
+        Self { cache }
+    }
+}
+
+impl Write for CacheWriter {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.cache.write(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+struct Settings {
+    compression: CompressionLevel,
+    filter: FilterType,
+}
+
+impl Default for Settings {
+    fn default() -> Self {
+        Settings {
+            compression: DEFAULT_COMPRESSION_LEVEL,
+            filter: DEFAULT_FILTER_TYPE,
+        }
+    }
+}
+
+static PROPERTIES: [subclass::Property; 2] = [
+    subclass::Property("compression-level", |name| {
+        glib::ParamSpec::enum_(
+            name,
+            "Compression level",
+            "Selects the compression algorithm to use",
+            CompressionLevel::static_type(),
+            DEFAULT_COMPRESSION_LEVEL as i32,
+            glib::ParamFlags::READWRITE,
+        )
+    }),
+    subclass::Property("filter", |name| {
+        glib::ParamSpec::enum_(
+            name,
+            "Filter",
+            "Selects the filter type to applied",
+            FilterType::static_type(),
+            DEFAULT_FILTER_TYPE as i32,
+            glib::ParamFlags::READWRITE,
+        )
+    }),
+];
+
+struct State {
+    video_info: gst_video::VideoInfo,
+    cache: Arc<CacheBuffer>,
+    writer: Option<png::Writer<CacheWriter>>,
+}
+
+impl State {
+    fn new(video_info: gst_video::VideoInfo) -> Self {
+        let cache = Arc::new(CacheBuffer::new());
+        Self {
+            video_info,
+            cache,
+            writer: None,
+        }
+    }
+
+    fn reset(&mut self, settings: Settings) -> Result<(), gst::LoggableError> {
+        // clear the cache
+        self.cache.clear();
+        let width = self.video_info.width();
+        let height = self.video_info.height();
+        let mut encoder = png::Encoder::new(CacheWriter::new(self.cache.clone()), width, height);
+        let color = match self.video_info.format() {
+            gst_video::VideoFormat::Gray8 | gst_video::VideoFormat::Gray16Be => {
+                png::ColorType::Grayscale
+            }
+            gst_video::VideoFormat::Rgb => png::ColorType::RGB,
+            gst_video::VideoFormat::Rgba => png::ColorType::RGBA,
+            _ => {
+                gst_error!(CAT, "format is not supported yet");
+                unreachable!()
+            }
+        };
+        let depth = if self.video_info.format() == gst_video::VideoFormat::Gray16Be {
+            png::BitDepth::Sixteen
+        } else {
+            png::BitDepth::Eight
+        };
+
+        encoder.set_color(color);
+        encoder.set_depth(depth);
+        encoder.set_compression(png::Compression::from(settings.compression));
+        encoder.set_filter(png::FilterType::from(settings.filter));
+        // Write the header for this video format into our inner buffer
+        let writer = encoder.write_header().map_err(|e| {
+            gst_loggable_error!(CAT, "Failed to create encoder error: {}", e.to_string())
+        })?;
+        self.writer = Some(writer);
+        Ok(())
+    }
+
+    fn write_data(&mut self, data: &[u8]) -> Result<(), png::EncodingError> {
+        if let Some(writer) = self.writer.as_mut() {
+            writer.write_image_data(data)
+        } else {
+            unreachable!()
+        }
+    }
+}
+
+struct PngEncoder {
+    state: Mutex<Option<State>>,
+    settings: Mutex<Settings>,
+}
+
+impl ObjectSubclass for PngEncoder {
+    const NAME: &'static str = "PngEncoder";
+    type ParentType = gst_video::VideoEncoder;
+    type Instance = gst::subclass::ElementInstanceStruct<Self>;
+    type Class = subclass::simple::ClassStruct<Self>;
+
+    glib_object_subclass!();
+
+    fn new() -> Self {
+        Self {
+            state: Mutex::new(None),
+            settings: Mutex::new(Default::default()),
+        }
+    }
+
+    fn class_init(klass: &mut subclass::simple::ClassStruct<Self>) {
+        klass.set_metadata(
+            "PNG encoder",
+            "Encoder/Video",
+            "PNG encoder",
+            "Natanael Mojica <neithanmo@gmail>",
+        );
+
+        let sink_caps = gst::Caps::new_simple(
+            "video/x-raw",
+            &[
+                (
+                    "format",
+                    &gst::List::new(&[
+                        &gst_video::VideoFormat::Gray8.to_str(),
+                        &gst_video::VideoFormat::Gray16Be.to_str(),
+                        &gst_video::VideoFormat::Rgb.to_str(),
+                        &gst_video::VideoFormat::Rgba.to_str(),
+                    ]),
+                ),
+                ("width", &gst::IntRange::<i32>::new(1, std::i32::MAX)),
+                ("height", &gst::IntRange::<i32>::new(1, std::i32::MAX)),
+                (
+                    "framerate",
+                    &gst::FractionRange::new(
+                        gst::Fraction::new(1, 1),
+                        gst::Fraction::new(std::i32::MAX, 1),
+                    ),
+                ),
+            ],
+        );
+        let sink_pad_template = gst::PadTemplate::new(
+            "sink",
+            gst::PadDirection::Sink,
+            gst::PadPresence::Always,
+            &sink_caps,
+        )
+        .unwrap();
+        klass.add_pad_template(sink_pad_template);
+
+        let src_caps = gst::Caps::new_simple("image/png", &[]);
+        let src_pad_template = gst::PadTemplate::new(
+            "src",
+            gst::PadDirection::Src,
+            gst::PadPresence::Always,
+            &src_caps,
+        )
+        .unwrap();
+        klass.add_pad_template(src_pad_template);
+        klass.install_properties(&PROPERTIES);
+    }
+}
+
+impl ObjectImpl for PngEncoder {
+    glib_object_impl!();
+
+    fn set_property(&self, _obj: &glib::Object, id: usize, value: &glib::Value) {
+        let prop = &PROPERTIES[id];
+
+        match *prop {
+            subclass::Property("compression-level", ..) => {
+                let mut settings = self.settings.lock();
+                settings.compression = value
+                    .get_some::<CompressionLevel>()
+                    .expect("type checked upstream");
+            }
+            subclass::Property("filter", ..) => {
+                let mut settings = self.settings.lock();
+                settings.filter = value
+                    .get_some::<FilterType>()
+                    .expect("type checked upstream");
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    fn get_property(&self, _obj: &glib::Object, id: usize) -> Result<glib::Value, ()> {
+        let prop = &PROPERTIES[id];
+
+        match *prop {
+            subclass::Property("compression-level", ..) => {
+                let settings = self.settings.lock();
+                Ok(settings.compression.to_value())
+            }
+            subclass::Property("filter", ..) => {
+                let settings = self.settings.lock();
+                Ok(settings.filter.to_value())
+            }
+            _ => unimplemented!(),
+        }
+    }
+}
+
+impl ElementImpl for PngEncoder {}
+
+impl VideoEncoderImpl for PngEncoder {
+    fn stop(&self, _element: &gst_video::VideoEncoder) -> Result<(), gst::ErrorMessage> {
+        *self.state.lock() = None;
+        Ok(())
+    }
+
+    fn set_format(
+        &self,
+        element: &gst_video::VideoEncoder,
+        state: &gst_video::VideoCodecState<'static, gst_video::video_codec_state::Readable>,
+    ) -> Result<(), gst::LoggableError> {
+        let video_info = state.get_info();
+        gst_debug!(CAT, obj: element, "Setting format {:?}", video_info);
+        {
+            let settings = self.settings.lock();
+            let mut state = State::new(video_info);
+            state.reset(*settings)?;
+            *self.state.lock() = Some(state);
+        }
+
+        let output_state = element
+            .set_output_state(gst::Caps::new_simple("image/png", &[]), Some(state))
+            .map_err(|_| gst_loggable_error!(CAT, "Failed to set output state"))?;
+        element
+            .negotiate(output_state)
+            .map_err(|_| gst_loggable_error!(CAT, "Failed to negotiate"))
+    }
+
+    fn handle_frame(
+        &self,
+        element: &gst_video::VideoEncoder,
+        mut frame: gst_video::VideoCodecFrame,
+    ) -> Result<gst::FlowSuccess, gst::FlowError> {
+        let mut state_guard = self.state.lock();
+        let state = state_guard.as_mut().ok_or(gst::FlowError::NotNegotiated)?;
+
+        gst_debug!(
+            CAT,
+            obj: element,
+            "Sending frame {}",
+            frame.get_system_frame_number()
+        );
+        {
+            let input_buffer = frame
+                .get_input_buffer()
+                .expect("frame without input buffer");
+
+            let input_map = input_buffer.map_readable().unwrap();
+            let data = input_map.as_slice();
+            state.write_data(data).map_err(|e| {
+                gst_element_error!(element, gst::CoreError::Failed, [&e.to_string()]);
+                gst::FlowError::Error
+            })?;
+        }
+
+        let buffer = state.cache.consume();
+        drop(state_guard);
+
+        let output_buffer = gst::Buffer::from_mut_slice(buffer);
+        // There are no such incremental frames in the png format
+        frame.set_flags(gst_video::VideoCodecFrameFlags::SYNC_POINT);
+        frame.set_output_buffer(output_buffer);
+        element.finish_frame(Some(frame))
+    }
+}
+
+pub fn register(plugin: &gst::Plugin) -> Result<(), glib::BoolError> {
+    gst::Element::register(
+        Some(plugin),
+        "rspngenc",
+        gst::Rank::Primary,
+        PngEncoder::get_type(),
+    )
+}
diff --git a/video/rspng/tests/pngenc.rs b/video/rspng/tests/pngenc.rs
new file mode 100644
index 000000000..09af25a01
--- /dev/null
+++ b/video/rspng/tests/pngenc.rs
@@ -0,0 +1,90 @@
+// Copyright (C) 2020 Natanael Mojica <neithanmo@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+fn init() {
+    use std::sync::Once;
+    static INIT: Once = Once::new();
+
+    INIT.call_once(|| {
+        gst::init().unwrap();
+        gstrspng::plugin_register_static().expect("Failed to register rspng plugin");
+    });
+}
+
+#[test]
+fn test_png_encode_gray() {
+    init();
+
+    let video_info = gst_video::VideoInfo::new(gst_video::VideoFormat::Gray8, 160, 120)
+        .fps((30, 1))
+        .build()
+        .unwrap();
+    test_png_encode(&video_info);
+}
+
+#[test]
+fn test_png_encode_gray16() {
+    init();
+
+    let video_info = gst_video::VideoInfo::new(gst_video::VideoFormat::Gray16Be, 160, 120)
+        .fps((30, 1))
+        .build()
+        .unwrap();
+    test_png_encode(&video_info);
+}
+
+#[test]
+fn test_png_encode_rgb() {
+    init();
+
+    let video_info = gst_video::VideoInfo::new(gst_video::VideoFormat::Rgb, 160, 120)
+        .fps((30, 1))
+        .build()
+        .unwrap();
+    test_png_encode(&video_info);
+}
+
+#[test]
+fn test_png_encode_rgba() {
+    init();
+
+    let video_info = gst_video::VideoInfo::new(gst_video::VideoFormat::Rgba, 160, 120)
+        .fps((30, 1))
+        .build()
+        .unwrap();
+    test_png_encode(&video_info);
+}
+
+fn test_png_encode(video_info: &gst_video::VideoInfo) {
+    let mut h = gst_check::Harness::new("rspngenc");
+    h.set_src_caps(video_info.to_caps().unwrap());
+    h.play();
+
+    for pts in 0..5 {
+        let buffer = {
+            let mut buffer = gst::Buffer::with_size(video_info.size()).unwrap();
+            {
+                let buffer = buffer.get_mut().unwrap();
+                buffer.set_pts(gst::ClockTime::from_seconds(pts));
+            }
+            let mut vframe =
+                gst_video::VideoFrame::from_buffer_writable(buffer, &video_info).unwrap();
+            for v in vframe.plane_data_mut(0).unwrap() {
+                *v = 128;
+            }
+            vframe.into_buffer()
+        };
+        h.push(buffer.clone()).unwrap();
+    }
+    h.push_event(gst::Event::new_eos().build());
+
+    (0..5).for_each(|_| {
+        let buffer = h.pull().unwrap();
+        assert!(!buffer.get_flags().contains(gst::BufferFlags::DELTA_UNIT))
+    });
+}