path: root/video/rspng/src/pngenc.rs

// 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(),
    )
}