diff --git a/Cargo.toml b/Cargo.toml index 9afb480..f1867c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ crate-type = ["cdylib"] [dependencies] anyhow = "1.0.79" +byteorder = "1.5.0" numpy = "0.21.0" ogg = "0.9.1" opus = "0.3.0" diff --git a/src/lib.rs b/src/lib.rs index 9647090..1fcff2e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -190,6 +190,26 @@ fn write_wav( Ok(()) } +#[pyfunction] +#[pyo3(signature = (filename, data, sample_rate))] +fn write_opus( + filename: std::path::PathBuf, + data: numpy::PyReadonlyArray1, + sample_rate: u32, +) -> PyResult<()> { + let w = std::fs::File::create(&filename).w_f(filename.as_path())?; + let mut w = std::io::BufWriter::new(w); + let data = data.as_array(); + match data.as_slice() { + None => { + let data = data.to_vec(); + opus::write_ogg(&mut w, data.as_ref(), sample_rate).w_f(filename.as_path())? + } + Some(data) => opus::write_ogg(&mut w, data, sample_rate).w_f(filename.as_path())?, + } + Ok(()) +} + /// Reads the whole content of an ogg/opus encoded file. /// /// This returns a two dimensional array as well as the sample rate. @@ -223,5 +243,6 @@ fn sphn(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(write_wav, m)?)?; m.add_function(wrap_pyfunction!(read_opus, m)?)?; m.add_function(wrap_pyfunction!(read_opus_bytes, m)?)?; + m.add_function(wrap_pyfunction!(write_opus, m)?)?; Ok(()) } diff --git a/src/opus.rs b/src/opus.rs index 6e13520..8ca28ac 100644 --- a/src/opus.rs +++ b/src/opus.rs @@ -1,5 +1,10 @@ use anyhow::Result; +// This must be an allowed value among 120, 240, 480, 960, 1920, and 2880. +// Using a different value would result in a BadArg "invalid argument" error when calling encode. +// https://opus-codec.org/docs/opus_api-1.2/group__opus__encoder.html#ga4ae9905859cd241ef4bb5c59cd5e5309 +const OPUS_ENCODER_FRAME_SIZE: usize = 960; + #[allow(unused)] #[derive(Debug)] struct OpusHeader { @@ -84,3 +89,66 @@ pub fn read_ogg(reader: R) -> Result<(Vec( + w: &mut W, + channels: u8, + sample_rate: u32, +) -> std::io::Result<()> { + use byteorder::WriteBytesExt; + + // https://wiki.xiph.org/OggOpus#ID_Header + w.write_all(b"OpusHead")?; + w.write_u8(1)?; // version + w.write_u8(channels)?; // channel count + w.write_u16::(3840)?; // pre-skip + w.write_u32::(sample_rate)?; // sample-rate in Hz + w.write_i16::(0)?; // output gain Q7.8 in dB + w.write_u8(0)?; // channel map + Ok(()) +} + +fn write_opus_tags(w: &mut W) -> std::io::Result<()> { + use byteorder::WriteBytesExt; + + // https://wiki.xiph.org/OggOpus#Comment_Header + let vendor = "sphn-pyo3"; + w.write_all(b"OpusTags")?; + w.write_u32::(vendor.len() as u32)?; // vendor string length + w.write_all(vendor.as_bytes())?; // vendor string, UTF8 encoded + w.write_u32::(0u32)?; // number of tags + Ok(()) +} + +pub fn write_ogg(w: &mut W, pcm: &[f32], sample_rate: u32) -> Result<()> { + let mut pw = ogg::PacketWriter::new(w); + + // Write the opus headers and tags + let mut head = Vec::new(); + write_opus_header(&mut head, 1, sample_rate)?; + pw.write_packet(head, 42, ogg::PacketWriteEndInfo::EndPage, 0)?; + let mut tags = Vec::new(); + write_opus_tags(&mut tags)?; + pw.write_packet(tags, 42, ogg::PacketWriteEndInfo::EndPage, 0)?; + + // Write the actual pcm data + let mut encoder = + opus::Encoder::new(sample_rate, opus::Channels::Mono, opus::Application::Voip)?; + let mut out_pcm_buf = vec![0u8; 50_000]; + + let mut total_data = 0; + let n_frames = pcm.len() / OPUS_ENCODER_FRAME_SIZE; + for (frame_idx, pcm) in pcm.chunks_exact(OPUS_ENCODER_FRAME_SIZE).enumerate() { + total_data += pcm.len() as u64; + let size = encoder.encode_float(pcm, &mut out_pcm_buf)?; + let msg = out_pcm_buf[..size].to_vec(); + let inf = if frame_idx + 1 == n_frames { + ogg::PacketWriteEndInfo::EndPage + } else { + ogg::PacketWriteEndInfo::NormalPacket + }; + pw.write_packet(msg, 42, inf, total_data)?; + } + + Ok(()) +}