ada

git clone git://archive.git.mtrnord.blog/MTRNord/ada.git
Log | Files | Refs

data.rs (4042B)

      1 use std::mem;
      2 use std::pin::Pin;
      3 use std::sync::{Arc, Mutex, Once};
      4 
      5 use async_std::stream::Stream;
      6 use async_std::task::Context;
      7 use audrey::read::Reader;
      8 use audrey::sample::interpolate::{Converter, Linear};
      9 use audrey::sample::signal::{from_iter, Signal};
     10 use futures::task::Poll;
     11 
     12 use crate::start_time::StartTime;
     13 use crate::SAMPLE_RATE;
     14 use std::io::{Cursor, Seek, SeekFrom};
     15 
     16 #[derive(Clone, Debug)]
     17 pub struct Data {
     18     buffer: Arc<Mutex<Vec<i16>>>,
     19 }
     20 
     21 fn packets_to_wav(data: Vec<i16>) -> Cursor<Vec<u8>> {
     22     println!("Make WAV");
     23     let wavspec = hound::WavSpec {
     24         channels: 1,
     25         sample_rate: 48_000,
     26         bits_per_sample: 16,
     27         sample_format: hound::SampleFormat::Int,
     28     };
     29 
     30     let mew_data = Vec::<u8>::new();
     31     let mut cursor = Cursor::new(mew_data);
     32 
     33     let mut writer = hound::WavWriter::new(&mut cursor, wavspec).unwrap();
     34     let mut i16_writer = writer.get_i16_writer((data.clone().len()) as u32);
     35 
     36     for sample in data.clone().into_iter().step_by(1) {
     37         i16_writer.write_sample(sample);
     38     }
     39 
     40     i16_writer.flush().unwrap();
     41     drop(writer);
     42     cursor.seek(SeekFrom::Start(0)).unwrap();
     43     println!("Return WAV");
     44     cursor
     45 }
     46 
     47 /// Interpolate the wav to the sample rate used by the deepspeech model.
     48 fn interpolate<F>(wav: F) -> Vec<i16>
     49 where
     50     F: std::io::Read,
     51     F: std::io::Seek,
     52 {
     53     println!("interpolate WAV");
     54     let mut reader = Reader::new(wav).unwrap();
     55     let description = reader.description();
     56 
     57     let audio_buffer: Vec<_> = if description.sample_rate() == SAMPLE_RATE {
     58         reader.samples().map(|s| s.unwrap()).collect()
     59     } else {
     60         let interpolator = Linear::new([0i16], [0]);
     61         let conv = Converter::from_hz_to_hz(
     62             from_iter(reader.samples::<i16>().map(|s| [s.unwrap()])),
     63             interpolator,
     64             description.sample_rate() as f64,
     65             SAMPLE_RATE as f64,
     66         );
     67         conv.until_exhausted().map(|v| v[0]).collect()
     68     };
     69     println!("return interpolated WAV");
     70     audio_buffer
     71 }
     72 
     73 impl Stream for Data {
     74     type Item = Vec<i16>;
     75 
     76     fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
     77         let buffer = self.buffer.lock().unwrap();
     78         let data = buffer.clone();
     79 
     80         let start_time_struct = StartTime::singleton();
     81         let ready = start_time_struct.get();
     82         if !data.is_empty() && ready {
     83             let audio_buffer = packets_to_wav(data.clone());
     84             let audio_buffer = interpolate(audio_buffer);
     85             Poll::Ready(Some(audio_buffer))
     86         } else {
     87             // Hack to keep alive
     88             Poll::Ready(Some(Vec::new()))
     89         }
     90     }
     91 }
     92 
     93 impl Data {
     94     pub(crate) fn singleton() -> Data {
     95         // Initialize it to a null value
     96         static mut DATA_SINGLETON: *const Data = 0 as *const Data;
     97         static DATA_ONCE: Once = Once::new();
     98 
     99         unsafe {
    100             DATA_ONCE.call_once(|| {
    101                 // Make it
    102                 let singleton = Data {
    103                     buffer: Arc::new(Mutex::new(Vec::new())),
    104                 };
    105 
    106                 // Put it in the heap so it can outlive this call
    107                 DATA_SINGLETON = mem::transmute(Box::new(singleton));
    108             });
    109 
    110             // Now we give out a copy of the data that is safe to use concurrently.
    111             (*DATA_SINGLETON).clone()
    112         }
    113     }
    114 
    115     pub fn reset(&mut self) {
    116         let mut state = self.buffer.lock().expect("Could not lock mutex");
    117         let data = Vec::new();
    118         mem::replace(&mut *state, data.clone());
    119     }
    120 
    121     pub fn add(&mut self, buf: cpal::UnknownTypeInputBuffer) {
    122         let mut state = self.buffer.lock().expect("Could not lock mutex");
    123         match buf {
    124             cpal::UnknownTypeInputBuffer::I16(buffer) => {
    125                 let mut local_state = state.clone();
    126                 local_state.extend_from_slice(&*buffer);
    127                 mem::replace(&mut *state, local_state.clone());
    128             }
    129             _ => panic!("Unexpected buffer type"),
    130         }
    131     }
    132 }