Merge pull request #1 from alesgenova/restructure

Restructure the code into various modules and update dependencies

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "pitch-detection"
 description = "A collection of algorithms to determine the pitch of a sound sample."
-version = "0.1.0"
+version = "0.2.0"
 authors = ["Alessandro Genova <ales.genova@gmail.com>"]
 edition = "2018"
 license = "MIT/Apache-2.0"
@@ -12,6 +12,6 @@ categories = ["algorithms", "multimedia::audio", "no-std"]
 readme = "README.md"
 
 [dependencies]
-rustfft = { version = "3.0.0", default-features = false }
+rustfft = { version = "4.0.0", default-features = false }
 num-traits = { version = "0.2", default-features = false }
-num-complex = { version = "0.2", default-features = false }
+num-complex = { version = "0.3", default-features = false }

README.md

@@ -2,14 +2,16 @@
 
 ## Usage
 ```rust
-use pitch_detection::{AutocorrelationDetector, McLeodDetector};
+use pitch_detection::detector::mcleod::McLeodDetector;
+use pitch_detection::detector::autocorrelation::AutocorrelationDetector;
 
 const SAMPLE_RATE : usize = 44100;
 const SIZE : usize = 1024;
 const PADDING : usize = SIZE / 2;
 const POWER_THRESHOLD : f64 = 5.0;
 const CLARITY_THRESHOLD : f64 = 0.7;
 
+// Signal coming from some source (microphone, generated, etc...)
 let signal = vec![0.0; SIZE];
 let mut detector = McLeodDetector::new(SIZE, PADDING);
 

src/detector/autocorrelation.rs

@@ -0,0 +1,47 @@
+use crate::float::Float;
+use crate::detector::PitchDetector;
+use crate::detector::internals::DetectorInternals;
+use crate::detector::internals::Pitch;
+use crate::detector::internals::autocorrelation;
+use crate::detector::internals::get_power_level;
+use crate::detector::internals::pitch_from_peaks;
+use crate::utils::peak::PeakCorrection;
+
+
+pub struct AutocorrelationDetector<T>
+    where T : Float
+{
+    internals: DetectorInternals<T>
+}
+
+impl<T> AutocorrelationDetector<T>
+    where T: Float
+{
+    pub fn new(size: usize, padding: usize) -> Self {
+        let internals = DetectorInternals::new(1, 2, size, padding);
+        AutocorrelationDetector {
+            internals
+        }
+    }
+}
+
+impl<T> PitchDetector<T> for AutocorrelationDetector<T>
+    where T : Float
+{
+    fn get_pitch(&mut self, signal: &[T], sample_rate: usize, power_threshold: T, clarity_threshold: T) -> Option<Pitch<T>> {
+        assert_eq!(signal.len(), self.internals.size);
+
+        if get_power_level(signal) < power_threshold {
+            return None;
+        }
+
+        let (signal_complex, rest) = self.internals.complex_buffers.split_first_mut().unwrap();
+        let (scratch, _) = rest.split_first_mut().unwrap();
+        let (autocorr, _) = self.internals.real_buffers.split_first_mut().unwrap();
+
+        autocorrelation(signal, signal_complex, scratch, autocorr);
+        let clarity_threshold = clarity_threshold * autocorr[0];
+
+        pitch_from_peaks(autocorr, sample_rate, clarity_threshold, PeakCorrection::None)
+    }
+}

src/detector/internals.rs

@@ -0,0 +1,146 @@
+use rustfft::FFTplanner;
+use num_complex::Complex;
+
+use crate::float::Float;
+use crate::utils::buffer::ComplexComponent;
+use crate::utils::buffer::copy_real_to_complex;
+use crate::utils::buffer::copy_complex_to_real;
+use crate::utils::buffer::new_real_buffer;
+use crate::utils::buffer::new_complex_buffer;
+use crate::utils::peak::PeakCorrection;
+use crate::utils::peak::detect_peaks;
+use crate::utils::peak::choose_peak;
+use crate::utils::peak::correct_peak;
+
+pub struct Pitch<T>
+    where T: Float
+{
+    pub frequency: T,
+    pub clarity: T
+}
+
+pub struct DetectorInternals<T>
+    where T: Float
+{
+    pub size: usize,
+    pub padding: usize,
+    pub real_buffers: Vec<Vec<T>>,
+    pub complex_buffers: Vec<Vec<Complex<T>>>
+}
+
+impl<T> DetectorInternals<T>
+    where T : Float
+{
+    pub fn new(n_real_buffers: usize, n_complex_buffers: usize, size: usize, padding: usize) -> Self {
+        let mut real_buffers: Vec<Vec<T>> = Vec::new();
+        let mut complex_buffers: Vec<Vec<Complex<T>>> = Vec::new();
+
+        for _i in 0..n_real_buffers {
+            let v = new_real_buffer(size + padding);
+            real_buffers.push(v);
+        }
+
+        for _i in 0..n_complex_buffers {
+            let v = new_complex_buffer(size + padding);
+            complex_buffers.push(v);
+        }
+
+        DetectorInternals {
+            size,
+            padding,
+            real_buffers,
+            complex_buffers
+        }
+    }
+}
+
+pub fn autocorrelation<T>(signal: &[T], signal_complex: &mut [Complex<T>], scratch: &mut [Complex<T>], result: &mut [T])
+    where T : Float
+{
+    copy_real_to_complex(signal, signal_complex, ComplexComponent::Re);
+    let mut planner = FFTplanner::new(false);
+    let fft = planner.plan_fft(signal_complex.len());
+    fft.process(signal_complex, scratch);
+    for i in 0..scratch.len() {
+        scratch[i].re = scratch[i].re * scratch[i].re + scratch[i].im * scratch[i].im;
+        scratch[i].im = T::zero();
+    }
+    let mut planner = FFTplanner::new(true);
+    let inv_fft = planner.plan_fft(signal_complex.len());
+    inv_fft.process(scratch, signal_complex);
+    copy_complex_to_real(signal_complex, result, ComplexComponent::Re);
+}
+
+pub fn pitch_from_peaks<T>(input: &[T], sample_rate: usize, clarity_threshold: T, correction: PeakCorrection) -> Option<Pitch<T>>
+    where T: Float
+{
+    let peaks = detect_peaks(input);
+    let chosen_peak = choose_peak(&peaks, clarity_threshold);
+    let chosen_peak = match chosen_peak {
+        Some(peak) => {
+                Some(correct_peak(peak, input, correction))
+        },
+        None => {
+            None
+        }
+    };
+
+    let pitch = match chosen_peak {
+        Some(peak) => {
+            let frequency = T::from_usize(sample_rate).unwrap() / peak.0;
+            let clarity = peak.1 / input[0];
+            Some(Pitch { frequency, clarity })
+        },
+        None => {
+            None
+        }
+    };
+    pitch
+}
+
+pub fn get_power_level<T>(signal: &[T]) -> T
+    where T : Float
+{
+    let mut power = T::zero();
+    for i in 0..signal.len() {
+        power = power + signal[i] * signal[i];
+    }
+    power
+}
+
+fn m_of_tau<T>(signal: &[T], signal_square_sum: Option<T>, result: &mut [T])
+    where T : Float
+{
+    assert!(result.len() >= signal.len());
+
+    let signal_square_sum = match signal_square_sum {
+        Some(val) => val,
+        None => {
+            let mut val = T::zero();
+            for i in 0..signal.len() {
+                val = val + signal[i] * signal[i];
+            }
+            val
+        }
+    };
+
+    result[0] = T::from_usize(2).unwrap() * signal_square_sum;
+    for i in 1..signal.len() {
+        result[i] = result[i - 1] - signal[i - 1] * signal[i - 1];
+    }
+
+    // Signal has no padding, but result does
+    for i in signal.len()..result.len() {
+        result[i] = result[i - 1];
+    }
+}
+
+pub fn normalized_square_difference<T>(signal: &[T], scratch0: &mut [Complex<T>], scratch1: &mut [Complex<T>], scratch2: &mut [T], result: &mut [T])
+    where T : Float
+{
+    autocorrelation(signal, scratch0, scratch1, result);
+    m_of_tau(signal, Some(result[0]), scratch2);
+    for i in 0..result.len() {
+        result[i] = T::from_usize(2).unwrap() * result[i] / scratch2[i];
+    }
+}

src/detector/mcleod.rs

@@ -0,0 +1,46 @@
+use crate::float::Float;
+use crate::detector::PitchDetector;
+use crate::detector::internals::DetectorInternals;
+use crate::detector::internals::Pitch;
+use crate::detector::internals::normalized_square_difference;
+use crate::detector::internals::get_power_level;
+use crate::detector::internals::pitch_from_peaks;
+use crate::utils::peak::PeakCorrection;
+
+pub struct McLeodDetector<T>
+    where T : Float
+{
+    internals: DetectorInternals<T>
+}
+
+impl<T> McLeodDetector<T>
+    where T: Float
+{
+    pub fn new(size: usize, padding: usize) -> Self {
+        let internals = DetectorInternals::new(2, 2, size, padding);
+        McLeodDetector {
+            internals
+        }
+    }
+}
+
+impl<T> PitchDetector<T> for McLeodDetector<T>
+    where T : Float
+{
+    fn get_pitch(&mut self, signal: &[T], sample_rate: usize, power_threshold: T, clarity_threshold: T) -> Option<Pitch<T>> {
+        assert_eq!(signal.len(), self.internals.size);
+
+        if get_power_level(signal) < power_threshold {
+            return None;
+        }
+
+        let (signal_complex, rest) = self.internals.complex_buffers.split_first_mut().unwrap();
+        let (scratch0, _) = rest.split_first_mut().unwrap();
+        let (scratch1, rest) = self.internals.real_buffers.split_first_mut().unwrap();
+        let (nsdf, _) = rest.split_first_mut().unwrap();
+
+        normalized_square_difference(signal, signal_complex, scratch0, scratch1, nsdf);
+
+        pitch_from_peaks(nsdf, sample_rate, clarity_threshold, PeakCorrection::Quadratic)
+    }
+}

src/detector/mod.rs

@@ -0,0 +1,12 @@
+use crate::float::Float;
+use crate::detector::internals::Pitch;
+
+pub mod internals;
+pub mod autocorrelation;
+pub mod mcleod;
+
+pub trait PitchDetector<T>
+    where T : Float
+{
+    fn get_pitch(&mut self, signal: &[T], sample_rate: usize, power_threshold: T, clarity_threshold: T) -> Option<Pitch<T>>;
+}

src/float/mod.rs

@@ -0,0 +1,7 @@
+use num_traits::float::{FloatCore as NumFloatCore};
+use rustfft::FFTnum;
+
+pub trait Float : NumFloatCore + FFTnum {}
+
+impl Float for f64 {}
+impl Float for f32 {}