2020-08-11 17:32:21 +00:00
|
|
|
// Copyright (C) 2020 Philippe Normand <philn@igalia.com>
|
|
|
|
// Copyright (C) 2020 Natanael Mojica <neithanmo@gmail.com>
|
|
|
|
//
|
2022-01-15 18:40:12 +00:00
|
|
|
// This Source Code Form is subject to the terms of the Mozilla Public License, v2.0.
|
|
|
|
// If a copy of the MPL was not distributed with this file, You can obtain one at
|
|
|
|
// <https://mozilla.org/MPL/2.0/>.
|
|
|
|
//
|
|
|
|
// SPDX-License-Identifier: MPL-2.0
|
2020-08-11 17:32:21 +00:00
|
|
|
|
2021-06-03 18:20:54 +00:00
|
|
|
use gst::glib;
|
2020-08-11 17:32:21 +00:00
|
|
|
use gst::prelude::*;
|
|
|
|
use gst::subclass::prelude::*;
|
|
|
|
use gst_base::subclass::base_transform::BaseTransformImplExt;
|
|
|
|
use gst_base::subclass::base_transform::GenerateOutputSuccess;
|
|
|
|
use gst_base::subclass::prelude::*;
|
2021-03-26 16:35:35 +00:00
|
|
|
|
2020-08-11 17:32:21 +00:00
|
|
|
use nnnoiseless::DenoiseState;
|
2021-03-26 16:35:35 +00:00
|
|
|
|
|
|
|
use byte_slice_cast::*;
|
2020-08-11 17:32:21 +00:00
|
|
|
|
2020-11-22 11:16:10 +00:00
|
|
|
use once_cell::sync::Lazy;
|
2021-03-26 16:35:35 +00:00
|
|
|
|
|
|
|
use atomic_refcell::AtomicRefCell;
|
|
|
|
|
2020-11-22 11:16:10 +00:00
|
|
|
static CAT: Lazy<gst::DebugCategory> = Lazy::new(|| {
|
|
|
|
gst::DebugCategory::new(
|
2020-08-11 17:32:21 +00:00
|
|
|
"audiornnoise",
|
|
|
|
gst::DebugColorFlags::empty(),
|
|
|
|
Some("Rust Audio Denoise Filter"),
|
2020-11-22 11:16:10 +00:00
|
|
|
)
|
|
|
|
});
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
const FRAME_SIZE: usize = DenoiseState::FRAME_SIZE;
|
|
|
|
|
|
|
|
struct ChannelDenoiser {
|
2020-08-19 15:00:39 +00:00
|
|
|
denoiser: Box<DenoiseState<'static>>,
|
2020-08-11 17:32:21 +00:00
|
|
|
frame_chunk: Box<[f32; FRAME_SIZE]>,
|
|
|
|
out_chunk: Box<[f32; FRAME_SIZE]>,
|
|
|
|
}
|
|
|
|
|
|
|
|
struct State {
|
|
|
|
in_info: gst_audio::AudioInfo,
|
|
|
|
denoisers: Vec<ChannelDenoiser>,
|
|
|
|
adapter: gst_base::UniqueAdapter,
|
|
|
|
}
|
|
|
|
|
2021-03-07 16:22:24 +00:00
|
|
|
#[derive(Default)]
|
2020-11-14 16:52:56 +00:00
|
|
|
pub struct AudioRNNoise {
|
2021-03-26 16:35:35 +00:00
|
|
|
state: AtomicRefCell<Option<State>>,
|
2020-08-11 17:32:21 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl State {
|
|
|
|
// The following three functions are copied from the csound filter.
|
2021-05-25 14:37:48 +00:00
|
|
|
fn buffer_duration(&self, buffer_size: u64) -> Option<gst::ClockTime> {
|
2020-08-11 17:32:21 +00:00
|
|
|
let samples = buffer_size / self.in_info.bpf() as u64;
|
|
|
|
self.samples_to_time(samples)
|
|
|
|
}
|
|
|
|
|
2021-05-25 14:37:48 +00:00
|
|
|
fn samples_to_time(&self, samples: u64) -> Option<gst::ClockTime> {
|
|
|
|
samples
|
|
|
|
.mul_div_round(*gst::ClockTime::SECOND, self.in_info.rate() as u64)
|
|
|
|
.map(gst::ClockTime::from_nseconds)
|
2020-08-11 17:32:21 +00:00
|
|
|
}
|
|
|
|
|
2021-05-25 14:37:48 +00:00
|
|
|
fn current_pts(&self) -> Option<gst::ClockTime> {
|
2020-08-11 17:32:21 +00:00
|
|
|
// get the last seen pts and the amount of bytes
|
|
|
|
// since then
|
|
|
|
let (prev_pts, distance) = self.adapter.prev_pts();
|
|
|
|
|
|
|
|
// Use the distance to get the amount of samples
|
|
|
|
// and with it calculate the time-offset which
|
|
|
|
// can be added to the prev_pts to get the
|
|
|
|
// pts at the beginning of the adapter.
|
|
|
|
let samples = distance / self.in_info.bpf() as u64;
|
2021-05-25 14:37:48 +00:00
|
|
|
prev_pts
|
2021-10-09 10:17:05 +00:00
|
|
|
.opt_checked_add(self.samples_to_time(samples))
|
|
|
|
.ok()
|
|
|
|
.flatten()
|
2020-08-11 17:32:21 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
fn needs_more_data(&self) -> bool {
|
|
|
|
self.adapter.available() < (FRAME_SIZE * self.in_info.bpf() as usize)
|
|
|
|
}
|
|
|
|
|
|
|
|
fn process(&mut self, input_plane: &[f32], output_plane: &mut [f32]) {
|
|
|
|
let channels = self.in_info.channels() as usize;
|
|
|
|
let size = FRAME_SIZE * channels;
|
|
|
|
|
|
|
|
for (out_frame, in_frame) in output_plane.chunks_mut(size).zip(input_plane.chunks(size)) {
|
|
|
|
for (index, item) in in_frame.iter().enumerate() {
|
|
|
|
let channel_index = index % channels;
|
|
|
|
let channel_denoiser = &mut self.denoisers[channel_index];
|
|
|
|
let pos = index / channels;
|
|
|
|
channel_denoiser.frame_chunk[pos] = *item;
|
|
|
|
}
|
|
|
|
|
|
|
|
for i in (in_frame.len() / channels)..(size / channels) {
|
|
|
|
for c in 0..channels {
|
|
|
|
let channel_denoiser = &mut self.denoisers[c];
|
|
|
|
channel_denoiser.frame_chunk[i] = 0.0;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// FIXME: The first chunks coming out of the denoisers contains some
|
|
|
|
// fade-in artifacts. We might want to discard those.
|
|
|
|
for channel_denoiser in &mut self.denoisers {
|
|
|
|
channel_denoiser.denoiser.process_frame(
|
|
|
|
&mut channel_denoiser.out_chunk[..],
|
|
|
|
&channel_denoiser.frame_chunk[..],
|
|
|
|
);
|
|
|
|
}
|
|
|
|
|
|
|
|
for (index, item) in out_frame.iter_mut().enumerate() {
|
|
|
|
let channel_index = index % channels;
|
|
|
|
let channel_denoiser = &self.denoisers[channel_index];
|
|
|
|
let pos = index / channels;
|
|
|
|
*item = channel_denoiser.out_chunk[pos];
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl AudioRNNoise {
|
2020-11-14 16:52:56 +00:00
|
|
|
fn drain(&self, element: &super::AudioRNNoise) -> Result<gst::FlowSuccess, gst::FlowError> {
|
2021-03-26 16:35:35 +00:00
|
|
|
let mut state_lock = self.state.borrow_mut();
|
2020-08-11 17:32:21 +00:00
|
|
|
let state = state_lock.as_mut().unwrap();
|
|
|
|
|
|
|
|
let available = state.adapter.available();
|
|
|
|
if available == 0 {
|
|
|
|
return Ok(gst::FlowSuccess::Ok);
|
|
|
|
}
|
|
|
|
|
|
|
|
let mut buffer = gst::Buffer::with_size(available).map_err(|e| {
|
2022-02-21 17:43:46 +00:00
|
|
|
gst::error!(
|
2020-08-11 17:32:21 +00:00
|
|
|
CAT,
|
|
|
|
obj: element,
|
|
|
|
"Failed to allocate buffer at EOS {:?}",
|
|
|
|
e
|
|
|
|
);
|
|
|
|
gst::FlowError::Flushing
|
|
|
|
})?;
|
|
|
|
|
|
|
|
let duration = state.buffer_duration(available as _);
|
2021-04-12 12:49:54 +00:00
|
|
|
let pts = state.current_pts();
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
{
|
|
|
|
let ibuffer = state.adapter.take_buffer(available).unwrap();
|
|
|
|
let in_map = ibuffer.map_readable().map_err(|_| gst::FlowError::Error)?;
|
|
|
|
let in_data = in_map.as_slice_of::<f32>().unwrap();
|
|
|
|
|
|
|
|
let buffer = buffer.get_mut().unwrap();
|
|
|
|
buffer.set_duration(duration);
|
|
|
|
buffer.set_pts(pts);
|
|
|
|
|
|
|
|
let mut out_map = buffer.map_writable().map_err(|_| gst::FlowError::Error)?;
|
2021-11-30 14:31:50 +00:00
|
|
|
let out_data = out_map.as_mut_slice_of::<f32>().unwrap();
|
2020-08-11 17:32:21 +00:00
|
|
|
|
2021-11-30 14:31:50 +00:00
|
|
|
state.process(in_data, out_data);
|
2020-08-11 17:32:21 +00:00
|
|
|
}
|
|
|
|
|
2021-04-20 12:58:11 +00:00
|
|
|
let srcpad = element.static_pad("src").unwrap();
|
2020-08-11 17:32:21 +00:00
|
|
|
srcpad.push(buffer)
|
|
|
|
}
|
|
|
|
|
|
|
|
fn generate_output(
|
|
|
|
&self,
|
2020-11-14 16:52:56 +00:00
|
|
|
_element: &super::AudioRNNoise,
|
2020-08-11 17:32:21 +00:00
|
|
|
state: &mut State,
|
|
|
|
) -> Result<GenerateOutputSuccess, gst::FlowError> {
|
|
|
|
let available = state.adapter.available();
|
|
|
|
let bpf = state.in_info.bpf() as usize;
|
|
|
|
let output_size = available - (available % (FRAME_SIZE * bpf));
|
|
|
|
let duration = state.buffer_duration(output_size as _);
|
2021-04-12 12:49:54 +00:00
|
|
|
let pts = state.current_pts();
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
let mut buffer = gst::Buffer::with_size(output_size).map_err(|_| gst::FlowError::Error)?;
|
|
|
|
|
|
|
|
{
|
|
|
|
let ibuffer = state
|
|
|
|
.adapter
|
|
|
|
.take_buffer(output_size)
|
|
|
|
.map_err(|_| gst::FlowError::Error)?;
|
|
|
|
let in_map = ibuffer.map_readable().map_err(|_| gst::FlowError::Error)?;
|
|
|
|
let in_data = in_map.as_slice_of::<f32>().unwrap();
|
|
|
|
|
|
|
|
let buffer = buffer.get_mut().unwrap();
|
|
|
|
buffer.set_duration(duration);
|
|
|
|
buffer.set_pts(pts);
|
|
|
|
|
|
|
|
let mut out_map = buffer.map_writable().map_err(|_| gst::FlowError::Error)?;
|
2021-11-30 14:31:50 +00:00
|
|
|
let out_data = out_map.as_mut_slice_of::<f32>().unwrap();
|
2020-08-11 17:32:21 +00:00
|
|
|
|
2021-11-30 14:31:50 +00:00
|
|
|
state.process(in_data, out_data);
|
2020-08-11 17:32:21 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
Ok(GenerateOutputSuccess::Buffer(buffer))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-03-07 16:22:24 +00:00
|
|
|
#[glib::object_subclass]
|
2020-08-11 17:32:21 +00:00
|
|
|
impl ObjectSubclass for AudioRNNoise {
|
|
|
|
const NAME: &'static str = "AudioRNNoise";
|
2020-11-14 16:52:56 +00:00
|
|
|
type Type = super::AudioRNNoise;
|
2020-08-11 17:32:21 +00:00
|
|
|
type ParentType = gst_base::BaseTransform;
|
|
|
|
}
|
|
|
|
|
|
|
|
impl ObjectImpl for AudioRNNoise {}
|
2021-01-21 18:21:29 +00:00
|
|
|
|
2021-10-23 08:57:31 +00:00
|
|
|
impl GstObjectImpl for AudioRNNoise {}
|
|
|
|
|
2021-01-21 18:21:29 +00:00
|
|
|
impl ElementImpl for AudioRNNoise {
|
|
|
|
fn metadata() -> Option<&'static gst::subclass::ElementMetadata> {
|
|
|
|
static ELEMENT_METADATA: Lazy<gst::subclass::ElementMetadata> = Lazy::new(|| {
|
|
|
|
gst::subclass::ElementMetadata::new(
|
|
|
|
"Audio denoise",
|
|
|
|
"Filter/Effect/Audio",
|
|
|
|
"Removes noise from an audio stream",
|
|
|
|
"Philippe Normand <philn@igalia.com>",
|
|
|
|
)
|
|
|
|
});
|
|
|
|
|
|
|
|
Some(&*ELEMENT_METADATA)
|
|
|
|
}
|
|
|
|
|
|
|
|
fn pad_templates() -> &'static [gst::PadTemplate] {
|
|
|
|
static PAD_TEMPLATES: Lazy<Vec<gst::PadTemplate>> = Lazy::new(|| {
|
2021-11-06 07:34:10 +00:00
|
|
|
let caps = gst::Caps::builder("audio/x-raw")
|
|
|
|
.field("format", gst_audio::AUDIO_FORMAT_F32.to_str())
|
|
|
|
.field("rate", 48000)
|
|
|
|
.field("channels", gst::IntRange::new(1, std::i32::MAX))
|
|
|
|
.field("layout", "interleaved")
|
|
|
|
.build();
|
2021-01-21 18:21:29 +00:00
|
|
|
let src_pad_template = gst::PadTemplate::new(
|
|
|
|
"src",
|
|
|
|
gst::PadDirection::Src,
|
|
|
|
gst::PadPresence::Always,
|
|
|
|
&caps,
|
|
|
|
)
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
let sink_pad_template = gst::PadTemplate::new(
|
|
|
|
"sink",
|
|
|
|
gst::PadDirection::Sink,
|
|
|
|
gst::PadPresence::Always,
|
|
|
|
&caps,
|
|
|
|
)
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
vec![src_pad_template, sink_pad_template]
|
|
|
|
});
|
|
|
|
|
|
|
|
PAD_TEMPLATES.as_ref()
|
|
|
|
}
|
|
|
|
}
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
impl BaseTransformImpl for AudioRNNoise {
|
2021-01-21 18:21:29 +00:00
|
|
|
const MODE: gst_base::subclass::BaseTransformMode =
|
|
|
|
gst_base::subclass::BaseTransformMode::NeverInPlace;
|
|
|
|
const PASSTHROUGH_ON_SAME_CAPS: bool = false;
|
|
|
|
const TRANSFORM_IP_ON_PASSTHROUGH: bool = false;
|
|
|
|
|
2020-08-11 17:32:21 +00:00
|
|
|
fn set_caps(
|
|
|
|
&self,
|
2020-11-14 16:52:56 +00:00
|
|
|
element: &Self::Type,
|
2020-08-11 17:32:21 +00:00
|
|
|
incaps: &gst::Caps,
|
|
|
|
outcaps: &gst::Caps,
|
|
|
|
) -> Result<(), gst::LoggableError> {
|
|
|
|
// Flush previous state
|
2021-03-26 16:35:35 +00:00
|
|
|
if self.state.borrow_mut().is_some() {
|
2020-08-11 17:32:21 +00:00
|
|
|
self.drain(element).map_err(|e| {
|
2020-12-20 18:43:45 +00:00
|
|
|
gst::loggable_error!(CAT, "Error flusing previous state data {:?}", e)
|
2020-08-11 17:32:21 +00:00
|
|
|
})?;
|
|
|
|
}
|
|
|
|
if incaps != outcaps {
|
2020-12-20 18:43:45 +00:00
|
|
|
return Err(gst::loggable_error!(
|
2020-08-11 17:32:21 +00:00
|
|
|
CAT,
|
|
|
|
"Input and output caps are not the same"
|
|
|
|
));
|
|
|
|
}
|
|
|
|
|
2022-02-21 17:43:46 +00:00
|
|
|
gst::debug!(CAT, obj: element, "Set caps to {}", incaps);
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
let in_info = gst_audio::AudioInfo::from_caps(incaps)
|
2020-12-20 18:43:45 +00:00
|
|
|
.map_err(|e| gst::loggable_error!(CAT, "Failed to parse input caps {:?}", e))?;
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
let mut denoisers = vec![];
|
|
|
|
for _i in 0..in_info.channels() {
|
|
|
|
denoisers.push(ChannelDenoiser {
|
|
|
|
denoiser: DenoiseState::new(),
|
|
|
|
frame_chunk: Box::new([0.0; FRAME_SIZE]),
|
|
|
|
out_chunk: Box::new([0.0; FRAME_SIZE]),
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2021-03-26 16:35:35 +00:00
|
|
|
let mut state_lock = self.state.borrow_mut();
|
2020-08-11 17:32:21 +00:00
|
|
|
*state_lock = Some(State {
|
|
|
|
in_info,
|
|
|
|
denoisers,
|
|
|
|
adapter: gst_base::UniqueAdapter::new(),
|
|
|
|
});
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
fn generate_output(
|
|
|
|
&self,
|
2020-11-14 16:52:56 +00:00
|
|
|
element: &Self::Type,
|
2020-08-11 17:32:21 +00:00
|
|
|
) -> Result<GenerateOutputSuccess, gst::FlowError> {
|
|
|
|
// Check if there are enough data in the queued buffer and adapter,
|
|
|
|
// if it is not the case, just notify the parent class to not generate
|
|
|
|
// an output
|
|
|
|
if let Some(buffer) = self.take_queued_buffer() {
|
2021-04-12 12:49:54 +00:00
|
|
|
if buffer.flags() == gst::BufferFlags::DISCONT {
|
2020-08-11 17:32:21 +00:00
|
|
|
self.drain(element)?;
|
|
|
|
}
|
|
|
|
|
2021-03-26 16:35:35 +00:00
|
|
|
let mut state_guard = self.state.borrow_mut();
|
2020-08-11 17:32:21 +00:00
|
|
|
let state = state_guard.as_mut().ok_or_else(|| {
|
2020-12-20 18:43:45 +00:00
|
|
|
gst::element_error!(
|
2020-08-11 17:32:21 +00:00
|
|
|
element,
|
|
|
|
gst::CoreError::Negotiation,
|
|
|
|
["Can not generate an output without State"]
|
|
|
|
);
|
|
|
|
gst::FlowError::NotNegotiated
|
|
|
|
})?;
|
|
|
|
|
|
|
|
state.adapter.push(buffer);
|
|
|
|
if !state.needs_more_data() {
|
|
|
|
return self.generate_output(element, state);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
Ok(GenerateOutputSuccess::NoOutput)
|
|
|
|
}
|
|
|
|
|
2020-11-14 16:52:56 +00:00
|
|
|
fn sink_event(&self, element: &Self::Type, event: gst::Event) -> bool {
|
2020-08-11 17:32:21 +00:00
|
|
|
use gst::EventView;
|
|
|
|
|
|
|
|
if let EventView::Eos(_) = event.view() {
|
2022-02-21 17:43:46 +00:00
|
|
|
gst::debug!(CAT, obj: element, "Handling EOS");
|
2020-08-11 17:32:21 +00:00
|
|
|
if self.drain(element).is_err() {
|
|
|
|
return false;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
self.parent_sink_event(element, event)
|
|
|
|
}
|
|
|
|
|
|
|
|
fn query(
|
|
|
|
&self,
|
2020-11-14 16:52:56 +00:00
|
|
|
element: &Self::Type,
|
2020-08-11 17:32:21 +00:00
|
|
|
direction: gst::PadDirection,
|
|
|
|
query: &mut gst::QueryRef,
|
|
|
|
) -> bool {
|
|
|
|
if direction == gst::PadDirection::Src {
|
2022-01-19 13:07:45 +00:00
|
|
|
if let gst::QueryViewMut::Latency(q) = query.view_mut() {
|
2021-04-20 12:58:11 +00:00
|
|
|
let sink_pad = element.static_pad("sink").expect("Sink pad not found");
|
2020-08-11 17:32:21 +00:00
|
|
|
let mut upstream_query = gst::query::Latency::new();
|
|
|
|
if sink_pad.peer_query(&mut upstream_query) {
|
2021-04-12 12:49:54 +00:00
|
|
|
let (live, mut min, mut max) = upstream_query.result();
|
2022-02-21 17:43:46 +00:00
|
|
|
gst::debug!(
|
2020-08-11 17:32:21 +00:00
|
|
|
CAT,
|
|
|
|
obj: element,
|
|
|
|
"Peer latency: live {} min {} max {}",
|
|
|
|
live,
|
|
|
|
min,
|
2021-05-25 14:37:48 +00:00
|
|
|
max.display(),
|
2020-08-11 17:32:21 +00:00
|
|
|
);
|
|
|
|
|
|
|
|
min += gst::ClockTime::from_seconds((FRAME_SIZE / 48000) as u64);
|
2021-10-09 10:17:05 +00:00
|
|
|
max = max.opt_add(gst::ClockTime::from_seconds((FRAME_SIZE / 48000) as u64));
|
2020-08-11 17:32:21 +00:00
|
|
|
q.set(live, min, max);
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
BaseTransformImplExt::parent_query(self, element, direction, query)
|
|
|
|
}
|
|
|
|
|
2020-11-14 16:52:56 +00:00
|
|
|
fn stop(&self, _element: &Self::Type) -> Result<(), gst::ErrorMessage> {
|
2020-08-11 17:32:21 +00:00
|
|
|
// Drop state
|
2021-03-26 16:35:35 +00:00
|
|
|
let _ = self.state.borrow_mut().take();
|
2020-08-11 17:32:21 +00:00
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
}
|