mistralrs_audio/
lib.rs

1//! Audio utilities for `mistral.rs`.
2//!
3//! This crate mirrors `mistralrs-vision` and focuses on audio specific
4//! functionality such as reading audio data, resampling and computing
5//! mel spectrogram features.
6
7use anyhow::Result;
8use symphonia::core::{
9    audio::SampleBuffer, codecs::DecoderOptions, formats::FormatOptions, io::MediaSourceStream,
10    meta::MetadataOptions, probe::Hint,
11};
12
13/// Raw audio input consisting of PCM samples and a sample rate.
14#[derive(Clone, Debug, PartialEq)]
15pub struct AudioInput {
16    pub samples: Vec<f32>,
17    pub sample_rate: u32,
18    pub channels: u16,
19}
20
21impl AudioInput {
22    /// Read a wav file from disk.
23    pub fn read_wav(wav_path: &str) -> Result<Self> {
24        let mut reader = hound::WavReader::open(wav_path)?;
25        let spec = reader.spec();
26        let samples: Vec<f32> = match spec.sample_format {
27            hound::SampleFormat::Float => reader
28                .samples::<f32>()
29                .collect::<std::result::Result<_, _>>()?,
30            hound::SampleFormat::Int => reader
31                .samples::<i16>()
32                .map(|s| s.map(|v| v as f32 / i16::MAX as f32))
33                .collect::<std::result::Result<_, _>>()?,
34        };
35        Ok(Self {
36            samples,
37            sample_rate: spec.sample_rate,
38            channels: spec.channels,
39        })
40    }
41
42    /// Decode audio bytes using `symphonia`.
43    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
44        let cursor = std::io::Cursor::new(bytes.to_vec());
45        let mss = MediaSourceStream::new(Box::new(cursor), Default::default());
46        let hint = Hint::new();
47        let probed = symphonia::default::get_probe().format(
48            &hint,
49            mss,
50            &FormatOptions::default(),
51            &MetadataOptions::default(),
52        )?;
53        let mut format = probed.format;
54        let track = format
55            .default_track()
56            .ok_or_else(|| anyhow::anyhow!("no supported audio tracks"))?;
57        let codec_params = &track.codec_params;
58        let sample_rate = codec_params
59            .sample_rate
60            .ok_or_else(|| anyhow::anyhow!("unknown sample rate"))?;
61        #[allow(clippy::cast_possible_truncation)]
62        let channels = codec_params.channels.map(|c| c.count() as u16).unwrap_or(1);
63        let mut decoder =
64            symphonia::default::get_codecs().make(codec_params, &DecoderOptions::default())?;
65        let mut samples = Vec::new();
66        loop {
67            match format.next_packet() {
68                Ok(packet) => {
69                    let decoded = decoder.decode(&packet)?;
70                    let mut buf =
71                        SampleBuffer::<f32>::new(decoded.capacity() as u64, *decoded.spec());
72                    buf.copy_interleaved_ref(decoded);
73                    samples.extend_from_slice(buf.samples());
74                }
75                Err(symphonia::core::errors::Error::IoError(e))
76                    if e.kind() == std::io::ErrorKind::UnexpectedEof =>
77                {
78                    break;
79                }
80                Err(e) => return Err(e.into()),
81            }
82        }
83        Ok(Self {
84            samples,
85            sample_rate,
86            channels,
87        })
88    }
89
90    /// Convert multi channel audio to mono by averaging channels.
91    pub fn to_mono(&self) -> Vec<f32> {
92        if self.channels <= 1 {
93            return self.samples.clone();
94        }
95        let mut mono = vec![0.0; self.samples.len() / self.channels as usize];
96        for (i, sample) in self.samples.iter().enumerate() {
97            mono[i / self.channels as usize] += *sample;
98        }
99        for s in &mut mono {
100            *s /= self.channels as f32;
101        }
102        mono
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::AudioInput;
109    use hound::{SampleFormat, WavSpec, WavWriter};
110    use std::io::Cursor;
111
112    #[test]
113    fn read_wav_roundtrip() {
114        let spec = WavSpec {
115            channels: 1,
116            sample_rate: 16000,
117            bits_per_sample: 16,
118            sample_format: SampleFormat::Int,
119        };
120        let mut writer = WavWriter::create("/tmp/test.wav", spec).unwrap();
121        for _ in 0..160 {
122            writer.write_sample::<i16>(0).unwrap();
123        }
124        writer.finalize().unwrap();
125        let input = AudioInput::read_wav("/tmp/test.wav").unwrap();
126        assert_eq!(input.samples.len(), 160);
127        assert_eq!(input.sample_rate, 16000);
128        std::fs::remove_file("/tmp/test.wav").unwrap();
129    }
130
131    #[test]
132    fn from_bytes() {
133        let spec = WavSpec {
134            channels: 1,
135            sample_rate: 8000,
136            bits_per_sample: 16,
137            sample_format: SampleFormat::Int,
138        };
139        let mut buffer: Vec<u8> = Vec::new();
140        {
141            let mut writer = WavWriter::new(Cursor::new(&mut buffer), spec).unwrap();
142            for _ in 0..80 {
143                writer.write_sample::<i16>(0).unwrap();
144            }
145            writer.finalize().unwrap();
146        }
147        let input = AudioInput::from_bytes(&buffer).unwrap();
148        assert_eq!(input.samples.len(), 80);
149        assert_eq!(input.sample_rate, 8000);
150    }
151}
mistralrs_audio/lib.rs

mistralrs_audio/
lib.rs
