Merge pull request #2 from LaurentMazare/opus-read

Sketch an ogg/opus reader using libopus.

.github/workflows/rust-ci.yml

@@ -0,0 +1,72 @@
+on: [push, pull_request]
+
+name: Continuous integration
+
+jobs:
+  check:
+    name: Check
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macOS-latest]
+        rust: [stable, nightly]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: ${{ matrix.rust }}
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: check
+
+  test:
+    name: Test Suite
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macOS-latest]
+        rust: [stable, nightly]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: ${{ matrix.rust }}
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+
+  fmt:
+    name: Rustfmt
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          override: true
+      - run: rustup component add rustfmt
+      - uses: actions-rs/cargo@v1
+        with:
+          command: fmt
+          args: --all -- --check
+
+  clippy:
+    name: Clippy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          override: true
+      - run: rustup component add clippy
+      - uses: actions-rs/cargo@v1
+        with:
+          command: clippy
+          args: -- -D warnings

Cargo.toml

@@ -10,6 +10,8 @@ crate-type = ["cdylib"]
 [dependencies]
 anyhow = "1.0.79"
 numpy = "0.21.0"
+ogg = "0.9.1"
+opus = "0.3.0"
 pyo3 = "0.21.0"
 rayon = "1.8.1"
 rubato = "0.15.0"

py_src/sphn/__init__.pyi

@@ -22,6 +22,15 @@ def read(filename, start_sec=None, duration_sec=None, sample_rate=None):
     """
     pass
 
+@staticmethod
+def read_opus(filename):
+    """
+    Reads the whole content of an ogg/opus encoded file.
+
+    This returns a two dimensional array as well as the sample rate.
+    """
+    pass
+
 @staticmethod
 def write_wav(filename, data, sample_rate):
     """

src/lib.rs

@@ -1,4 +1,5 @@
 mod audio;
+mod opus;
 mod wav;
 
 use pyo3::prelude::*;
@@ -185,11 +186,23 @@ fn write_wav(
     Ok(())
 }
 
+/// Reads the whole content of an ogg/opus encoded file.
+///
+/// This returns a two dimensional array as well as the sample rate.
+#[pyfunction]
+#[pyo3(signature = (filename))]
+fn read_opus(filename: std::path::PathBuf, py: Python) -> PyResult<(PyObject, u32)> {
+    let (data, sample_rate) = opus::read_ogg(&filename).w_f(filename.as_path())?;
+    let data = numpy::PyArray2::from_vec2_bound(py, &data)?.into_py(py);
+    Ok((data, sample_rate))
+}
+
 #[pymodule]
 fn sphn(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<FileReader>()?;
     m.add_function(wrap_pyfunction!(durations, m)?)?;
     m.add_function(wrap_pyfunction!(read, m)?)?;
     m.add_function(wrap_pyfunction!(write_wav, m)?)?;
+    m.add_function(wrap_pyfunction!(read_opus, m)?)?;
     Ok(())
 }

src/opus.rs

@@ -0,0 +1,87 @@
+use anyhow::Result;
+
+#[allow(unused)]
+#[derive(Debug)]
+struct OpusHeader {
+    version: u8,
+    channel_count: u8,
+    pre_skip: u16,
+    input_sample_rate: u32,
+    output_gain: i16,
+    mapping_family: u8,
+}
+
+fn parse_opus_header(packet: &[u8]) -> Result<OpusHeader> {
+    if packet.len() < 8 || &packet[0..8] != b"OpusHead" {
+        anyhow::bail!("not a OpusHead packet")
+    }
+    let header = OpusHeader {
+        version: packet[8],
+        channel_count: packet[9],
+        pre_skip: u16::from_le_bytes([packet[10], packet[11]]),
+        input_sample_rate: u32::from_le_bytes([packet[12], packet[13], packet[14], packet[15]]),
+        output_gain: i16::from_le_bytes([packet[16], packet[17]]),
+        mapping_family: packet[18],
+    };
+    Ok(header)
+}
+
+pub fn read_ogg<P: AsRef<std::path::Path>>(p: P) -> Result<(Vec<Vec<f32>>, u32)> {
+    let file = std::fs::File::open(p.as_ref())?;
+    let file = std::io::BufReader::new(file);
+    let mut packet_reader = ogg::PacketReader::new(file);
+    let mut opus_decoder = None;
+    let mut channels = 1;
+    let mut all_data = vec![];
+    while let Some(packet) = packet_reader.read_packet()? {
+        let is_header = packet.data.len() >= 8 && &packet.data[0..8] == b"OpusHead";
+        let is_tags = packet.data.len() >= 8 && &packet.data[0..8] == b"OpusTags";
+        if is_tags {
+            continue;
+        }
+        match (is_header, opus_decoder.as_mut()) {
+            (true, Some(_)) => anyhow::bail!("multiple OpusHead packets"),
+            (true, None) => {
+                let header = parse_opus_header(&packet.data)?;
+                channels = header.channel_count as usize;
+                let channels = match header.channel_count {
+                    1 => opus::Channels::Mono,
+                    2 => opus::Channels::Stereo,
+                    c => anyhow::bail!("unexpected number of channels {c}"),
+                };
+                let od = opus::Decoder::new(header.input_sample_rate, channels)?;
+                opus_decoder = Some(od)
+            }
+            (false, None) => anyhow::bail!("no initial OpusHead"),
+            (false, Some(od)) => {
+                let nb_samples = od.get_nb_samples(&packet.data)?;
+                let prev_len = all_data.len();
+                all_data.resize(prev_len + nb_samples * channels, 0f32);
+                let samples = od.decode_float(
+                    &packet.data,
+                    &mut all_data[prev_len..],
+                    /* Forward Error Correction */ false,
+                )?;
+                all_data.resize(prev_len + samples * channels, 0f32);
+            }
+        }
+    }
+    let sample_rate = match opus_decoder.as_mut() {
+        None => anyhow::bail!("no data"),
+        Some(od) => od.get_sample_rate()?,
+    };
+    let data = match channels {
+        1 => vec![all_data],
+        2 => {
+            let mut c0 = Vec::with_capacity(all_data.len() / 2);
+            let mut c1 = Vec::with_capacity(all_data.len() / 2);
+            for c in all_data.chunks(2) {
+                c0.push(c[0]);
+                c1.push(c[1]);
+            }
+            vec![c0, c1]
+        }
+        c => anyhow::bail!("unexpected number of channels {c}"),
+    };
+    Ok((data, sample_rate))
+}