From 120d8d8c4921e1798c91a23b0092c7e90d338777 Mon Sep 17 00:00:00 2001 From: Kai Vogelgesang Date: Mon, 25 Oct 2021 23:58:05 +0200 Subject: [PATCH] Add beat detection script --- beat_detection/.gitignore | 14 ++ beat_detection/Cargo.toml | 13 ++ beat_detection/src/capture.rs | 95 +++++++++++++ beat_detection/src/dsp.rs | 71 ++++++++++ beat_detection/src/main.rs | 258 ++++++++++++++++++++++++++++++++++ 5 files changed, 451 insertions(+) create mode 100644 beat_detection/.gitignore create mode 100644 beat_detection/Cargo.toml create mode 100644 beat_detection/src/capture.rs create mode 100644 beat_detection/src/dsp.rs create mode 100644 beat_detection/src/main.rs diff --git a/beat_detection/.gitignore b/beat_detection/.gitignore new file mode 100644 index 0000000..6985cf1 --- /dev/null +++ b/beat_detection/.gitignore @@ -0,0 +1,14 @@ +# Generated by Cargo +# will have compiled files and executables +debug/ +target/ + +# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries +# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html +Cargo.lock + +# These are backup files generated by rustfmt +**/*.rs.bk + +# MSVC Windows builds of rustc generate these, which store debugging information +*.pdb diff --git a/beat_detection/Cargo.toml b/beat_detection/Cargo.toml new file mode 100644 index 0000000..967b475 --- /dev/null +++ b/beat_detection/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "beat_detection" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +pulse = { version = "2.23.1", package = "libpulse-binding" } +psimple = { version = "2.23.0", package = "libpulse-simple-binding" } +anyhow = "1.0.44" +sdl2 = "0.35.1" +ringbuffer = "0.8.3" diff --git a/beat_detection/src/capture.rs b/beat_detection/src/capture.rs new file mode 100644 index 0000000..14e321b --- /dev/null +++ b/beat_detection/src/capture.rs @@ -0,0 +1,95 @@ +use anyhow::{anyhow, Result}; +use psimple::Simple; +use pulse::context::{Context, FlagSet as ContextFlagSet}; +use pulse::mainloop::standard::{IterateResult, Mainloop}; +use pulse::sample::{Format, Spec}; +use pulse::stream::Direction; +use std::cell::RefCell; +use std::ops::Deref; +use std::rc::Rc; + +/* + Some manual poking around in PulseAudio to get the name of the default sink +*/ + +fn poll_mainloop(mainloop: &Rc>) -> Result<()> { + match mainloop.borrow_mut().iterate(true) { + IterateResult::Quit(_) | IterateResult::Err(_) => { + return Err(anyhow!("Iterate state was not success")); + } + IterateResult::Success(_) => { + return Ok(()); + } + } +} + +fn get_pulse_default_sink() -> Result { + let mainloop = Rc::new(RefCell::new( + Mainloop::new().ok_or(anyhow!("Failed to create mainloop"))?, + )); + let ctx = Rc::new(RefCell::new( + Context::new(mainloop.borrow().deref(), "gib_default_sink") + .ok_or(anyhow!("Failed to create context"))?, + )); + + ctx.borrow_mut() + .connect(None, ContextFlagSet::NOFLAGS, None)?; + + // Wait for context to be ready + loop { + poll_mainloop(&mainloop)?; + + match ctx.borrow().get_state() { + pulse::context::State::Ready => { + break; + } + pulse::context::State::Failed | pulse::context::State::Terminated => { + return Err(anyhow!("Context was in failed/terminated state")); + } + _ => {} + } + } + + let result = Rc::new(RefCell::new(None)); + let cb_result_ref = result.clone(); + + ctx.borrow().introspect().get_server_info(move |info| { + *cb_result_ref.borrow_mut() = if let Some(ref sink_name) = info.default_sink_name { + Some(Ok(sink_name.to_string())) + } else { + Some(Err(())) + } + }); + + loop { + if let Some(result) = result.borrow().deref() { + return result + .to_owned() + .map_err(|_| anyhow!("Default sink name was empty")); + } + + poll_mainloop(&mainloop)?; + } +} + +/* + Get a PASimple instance which reads from the default sink monitor +*/ + +pub fn get_audio_reader(spec: &Spec) -> Result { + let mut default_sink_name = get_pulse_default_sink()?; + default_sink_name.push_str(".monitor"); + + let simple = Simple::new( + None, + "piano_thingy", + Direction::Record, + Some(&default_sink_name), + "sample_yoinker", + &spec, + None, + None, + )?; + + Ok(simple) +} diff --git a/beat_detection/src/dsp.rs b/beat_detection/src/dsp.rs new file mode 100644 index 0000000..1c46c30 --- /dev/null +++ b/beat_detection/src/dsp.rs @@ -0,0 +1,71 @@ +/* + Taken from Till's magic Arduino Sketch +*/ +pub trait ZTransformFilter { + fn process(&mut self, sample: f32) -> f32; +} + +// 20 - 200Hz Single Pole Bandpass IIR Filter +#[derive(Default)] +pub struct BassFilter { + xv: [f32; 3], + yv: [f32; 3], +} + +impl ZTransformFilter for BassFilter { + fn process(&mut self, sample: f32) -> f32 { + self.xv[0] = self.xv[1]; + self.xv[1] = self.xv[2]; + self.xv[2] = sample / 3.0f32; + + self.yv[0] = self.yv[1]; + self.yv[1] = self.yv[2]; + self.yv[2] = (self.xv[2] - self.xv[0]) + + (-0.7960060012f32 * self.yv[0]) + + (1.7903124146f32 * self.yv[1]); + + self.yv[2] + } +} + +// 10Hz Single Pole Lowpass IIR Filter +#[derive(Default)] +pub struct EnvelopeFilter { + xv: [f32; 2], + yv: [f32; 2], +} + +impl ZTransformFilter for EnvelopeFilter { + fn process(&mut self, sample: f32) -> f32 { + self.xv[0] = self.xv[1]; + self.xv[1] = sample / 50.0f32; + + self.yv[0] = self.yv[1]; + self.yv[1] = (self.xv[0] + self.xv[1]) + (0.9875119299f32 * self.yv[0]); + + self.yv[1] + } +} + +// 1.7 - 3.0Hz Single Pole Bandpass IIR Filter +#[derive(Default)] +pub struct BeatFilter { + xv: [f32; 3], + yv: [f32; 3], +} + +impl ZTransformFilter for BeatFilter { + fn process(&mut self, sample: f32) -> f32 { + self.xv[0] = self.xv[1]; + self.xv[1] = self.xv[2]; + self.xv[2] = sample / 2.7f32; + + self.yv[0] = self.yv[1]; + self.yv[1] = self.yv[2]; + self.yv[2] = (self.xv[2] - self.xv[0]) + + (-0.7169861741f32 * self.yv[0]) + + (1.4453653501f32 * self.yv[1]); + + self.yv[2] + } +} diff --git a/beat_detection/src/main.rs b/beat_detection/src/main.rs new file mode 100644 index 0000000..4c3aab1 --- /dev/null +++ b/beat_detection/src/main.rs @@ -0,0 +1,258 @@ +mod capture; +mod dsp; + +use anyhow::Result; +use pulse::sample::{Format, Spec}; +use ringbuffer::{ConstGenericRingBuffer, RingBuffer, RingBufferExt, RingBufferWrite}; +use sdl2::{ + event::Event, + keyboard::Keycode, + pixels::Color, + rect::{Point, Rect}, +}; + +use crate::dsp::ZTransformFilter; + +const SAMPLE_RATE: usize = 5000; +const FPS: usize = 50; + +const BUFFER_SIZE: usize = SAMPLE_RATE / FPS; + +const POINT_COUNT: usize = SAMPLE_RATE / 200; +const POINT_BUFFER_SIZE: usize = POINT_COUNT.next_power_of_two(); + +const MIN_MAX_SAMPLE_COUNT: usize = 10; + +fn main() -> Result<()> { + let spec = Spec { + format: Format::F32le, + rate: SAMPLE_RATE as u32, + channels: 1, + }; + assert!(spec.is_valid()); + + let reader = capture::get_audio_reader(&spec)?; + let mut buffer = [0u8; 4 * BUFFER_SIZE]; + + let mut bass_filter = dsp::BassFilter::default(); + let mut envelope_filter = dsp::EnvelopeFilter::default(); + let mut beat_filter = dsp::BeatFilter::default(); + let mut j = 0; + + let mut threshold = 1.5f32; + + let mut brightness = 0f32; + + let mut beat_toggle = false; + + // sdl + + let sdl = sdl2::init().unwrap(); + let sdl_video = sdl.video().unwrap(); + + let window = sdl_video + .window("Beat Detectinator", 800, 600) + .resizable() + .build()?; + let mut canvas = window.into_canvas().build()?; + let texture_creator = canvas.texture_creator(); + let texture_size = 10 * POINT_COUNT as u32; + let mut texture = texture_creator.create_texture_target( + canvas.default_pixel_format(), + texture_size, + texture_size, + )?; + + let mut point_buf = ConstGenericRingBuffer::::new(); + point_buf.fill_default(); + + let mut event_pump = sdl.event_pump().unwrap(); + + 'running: loop { + for event in event_pump.poll_iter() { + match event { + Event::Quit { .. } + | Event::KeyDown { + keycode: Some(Keycode::Escape), + .. + } => break 'running, + Event::KeyDown { + keycode: Some(Keycode::Up), + .. + } => { + threshold += 0.01f32; + println!("threshold: {:.2}", threshold); + } + Event::KeyDown { + keycode: Some(Keycode::Down), + .. + } => { + threshold -= 0.01f32; + println!("threshold: {:.2}", threshold); + } + Event::KeyDown { + keycode: Some(k), .. + } => { + println!("{}", k) + } + _ => {} + } + } + + reader.read(&mut buffer)?; + + for i in 0..BUFFER_SIZE { + let mut float_bytes = [0u8; 4]; + float_bytes.copy_from_slice(&buffer[4 * i..4 * i + 4]); + + j += 1; + let sample = f32::from_le_bytes(float_bytes); + let mut value = bass_filter.process(sample); + + if value < 0f32 { + value = -value; + } + + let envelope = envelope_filter.process(value); + + if j == 200 { + let beat = beat_filter.process(envelope); + point_buf.push(beat); + + canvas.with_texture_canvas(&mut texture, |canvas| { + let min_y = 0f32; + let max_y = 3f32; + + let get_y = |y: &f32| { + + let mut y = y.clone(); + + if y <= 0f32 { + y = 0f32; + } + + y = (1f32 + y).log2(); + + let mut y = (y - min_y) / (max_y - min_y); + + ((1f32 - y) * texture_size as f32) as u32 + }; + + // background + + let prev_brightness = brightness; + brightness = if beat > threshold { + 1f32 + } else { + 0.75f32 * brightness + }; + + if brightness - prev_brightness > 0.5 { + println!("beat. {}", if beat_toggle {"-"} else {"|"}); + beat_toggle ^= true; + } + + let v = (255f32 * brightness) as u8; + + canvas.set_draw_color(Color::RGB(v, v, v)); + canvas.clear(); + + // zero + + canvas.set_draw_color(Color::RGB(0, 128, 0)); + let y = get_y(&0f32); + + canvas + .draw_line( + Point::new(0, y as i32), + Point::new(texture_size as i32, y as i32), + ) + .unwrap(); + + // min / max lines + + canvas.set_draw_color(Color::RGB(255, 0, 0)); + + let min_beat = point_buf + .iter() + .skip(point_buf.capacity() - MIN_MAX_SAMPLE_COUNT) + .reduce(|a, b| if a < b { a } else { b }) + .unwrap(); + + let x = texture_size - MIN_MAX_SAMPLE_COUNT as u32 * 10; + + let y = get_y(min_beat); + + canvas + .draw_line( + Point::new(x as i32, y as i32), + Point::new(texture_size as i32, y as i32), + ) + .unwrap(); + + let max_beat = point_buf + .iter() + .skip(point_buf.capacity() - MIN_MAX_SAMPLE_COUNT) + .reduce(|a, b| if a > b { a } else { b }) + .unwrap(); + + let y = get_y(max_beat); + + canvas + .draw_line( + Point::new(x as i32, y as i32), + Point::new(texture_size as i32, y as i32), + ) + .unwrap(); + + // threshhold line + + canvas.set_draw_color(Color::RGB(0, 0, 255)); + let y = get_y(&threshold); + + canvas + .draw_line( + Point::new(0, y as i32), + Point::new(texture_size as i32, y as i32), + ) + .unwrap(); + + // values + + canvas.set_draw_color(Color::RGB(0, 255, 0)); + + for (i, beat) in point_buf + .iter() + .skip(point_buf.capacity() - POINT_COUNT) + .enumerate() + { + let x = 10 * i; + let y = get_y(beat); + + canvas + .draw_rect(Rect::new((x + 1) as i32, (y - 1) as i32, 8, 3)) + .unwrap(); + } + })?; + + j = 0; + } + } + + canvas.set_draw_color(Color::RGB(0, 0, 0)); + canvas.clear(); + + let (w, h) = canvas.window().drawable_size(); + + let target_rect = if w > h { + Rect::new(((w - h) / 2) as i32, 0, h, h) + } else { + Rect::new(0, ((h - w) / 2) as i32, w, w) + }; + + canvas.copy(&texture, None, Some(target_rect)).unwrap(); + canvas.present(); + } + + Ok(()) +}