From c7cd83f3edfc537ae8c2dc913e2dd1b62a449734 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Sun, 8 Aug 2021 14:02:42 +0200 Subject: [PATCH] add MS variant of IMA ADPCM encoder --- nihav-ms/Cargo.toml | 3 +- nihav-ms/src/codecs/imaadpcmenc.rs | 389 +++++++++++++++++++++++++++++ nihav-ms/src/codecs/mod.rs | 5 + 3 files changed, 396 insertions(+), 1 deletion(-) create mode 100644 nihav-ms/src/codecs/imaadpcmenc.rs diff --git a/nihav-ms/Cargo.toml b/nihav-ms/Cargo.toml index cf9a79f..ffc6126 100644 --- a/nihav-ms/Cargo.toml +++ b/nihav-ms/Cargo.toml @@ -28,7 +28,8 @@ all_audio_decoders = ["decoder_ima_adpcm_ms", "decoder_ms_adpcm"] decoder_ima_adpcm_ms = ["decoders"] decoder_ms_adpcm = ["decoders"] -all_encoders = ["encoder_msvideo1", "encoder_ms_adpcm"] +all_encoders = ["encoder_msvideo1", "encoder_ms_adpcm", "encoder_ima_adpcm_ms"] encoder_msvideo1 = ["encoders"] encoder_ms_adpcm = ["encoders"] +encoder_ima_adpcm_ms = ["encoders"] encoders = [] diff --git a/nihav-ms/src/codecs/imaadpcmenc.rs b/nihav-ms/src/codecs/imaadpcmenc.rs new file mode 100644 index 0000000..a64f58f --- /dev/null +++ b/nihav-ms/src/codecs/imaadpcmenc.rs @@ -0,0 +1,389 @@ +use nihav_core::codecs::*; +use nihav_core::io::byteio::*; +use nihav_codec_support::codecs::imaadpcm::*; + +struct NibbleWriter<'a> { + bw: ByteWriter<'a>, + val: u8, + first: bool, +} + +impl<'a> NibbleWriter<'a> { + fn new(bw: ByteWriter<'a>) -> Self { + Self { bw, val: 0, first: true } + } + fn write(&mut self, nib: u8) -> EncoderResult<()> { + if self.first { + self.val = nib; + self.first = false; + } else { + self.val |= nib << 4; + self.bw.write_byte(self.val)?; + self.first = true; + } + Ok(()) + } +} + +#[derive(Clone,Copy)] +struct TrellisNode { + state: IMAState, + nib: u8, + error: i64, +} + +impl Default for TrellisNode { + fn default() -> Self { + TrellisNode { + state: IMAState::new(), + nib: 0, + error: 0, + } + } +} + +#[derive(Default)] +struct IMAADPCMEncoder { + stream: Option, + samples: Vec, + block_len: usize, + block_size: usize, + channels: usize, + flush: bool, + srate: u32, + trellis: bool, + nodes: Vec<[TrellisNode; 16]>, + first: Vec, + nibs: Vec>, +} + +const DEFAULT_BLOCK_LEN: usize = 256; + +impl IMAADPCMEncoder { + fn new() -> Self { Self::default() } + fn encode_packet(&mut self) -> EncoderResult { + if self.samples.is_empty() { + return Err(EncoderError::TryAgain); + } + let cur_len = (self.samples.len() / self.channels).min(self.block_len); + if cur_len < self.block_len && !self.flush { + return Err(EncoderError::TryAgain); + } + if cur_len < self.block_len { + self.samples.resize(self.block_len * self.channels, 0); + } + + let mut dbuf = vec![0u8; self.block_size]; + let mut mw = MemoryWriter::new_write(dbuf.as_mut_slice()); + let mut bw = ByteWriter::new(&mut mw); + + for ch in 0..self.channels { + self.first[ch].predictor = i32::from(self.samples[ch]); + self.first[ch].step = Self::calc_step(self.samples[ch], self.samples[ch + self.channels]) as usize; + } + + if !self.trellis { + for ch in 0..self.channels { + bw.write_u16le(self.first[ch].predictor as u16)?; + bw.write_byte(self.first[ch].step as u8)?; + bw.write_byte(0)?; + } + let mut nw = NibbleWriter::new(bw); + for samples in self.samples.chunks(self.channels).take(self.block_len).skip(1) { + for (state, &samp) in self.first.iter_mut().zip(samples.iter()) { + let nib = state.compress_sample(samp); + state.expand_sample(nib); + nw.write(nib)?; + } + } + drop(nw); + } else { + self.nodes.reserve(self.block_len); + self.nibs.resize(self.channels, Vec::new()); + for nibs in self.nibs.iter_mut() { + nibs.resize(self.block_len, 0); + } + let step = self.channels; + let mut state = [TrellisNode::default(); 16]; + for ch in 0..self.channels { + self.nodes.clear(); + for i in 0..16 { + let step = (((self.first[ch].step + i) as isize) - 8).max(0).min(IMA_MAX_STEP as isize) as usize; + state[i].state.predictor = self.first[ch].predictor; + state[i].state.step = step; + state[i].error = 0; + state[i].nib = step as u8; + } + self.nodes.push(state); + let mut sidx = ch + step; + for _i in 1..self.block_len { + let sample = self.samples[sidx]; + + for (nnib, cstate) in state.iter_mut().enumerate() { + for nib in 0..16 { + let pnode = &self.nodes[self.nodes.len() - 1][nib]; + let mut ima = pnode.state; + let nsamp = ima.expand_sample(nnib as u8); + let diff = i64::from(i32::from(sample) - i32::from(nsamp)); + let error = pnode.error + diff * diff; + if (nib == 0) || error < cstate.error { + cstate.state = ima; + cstate.nib = nib as u8; + cstate.error = error; + } + } + } + self.nodes.push(state); + + sidx += step; + } + + let mut idx = 0; + let mut best_err = self.nodes[self.nodes.len() - 1][0].error; + for (i, node) in self.nodes[self.nodes.len() - 1].iter().enumerate() { + if node.error < best_err { + best_err = node.error; + idx = i; + } + } + let mut dst = self.nibs[ch].iter_mut().rev(); + while let Some(nodes) = self.nodes.pop() { + *dst.next().unwrap() = idx as u8; + idx = nodes[idx].nib as usize; + } + self.nibs[ch][0] = idx as u8; + } + for ch in 0..self.channels { + bw.write_u16le(self.first[ch].predictor as u16)?; + bw.write_byte(self.nibs[ch][0])?; + bw.write_byte(0)?; + } + let mut nw = NibbleWriter::new(bw); + for i in 1..self.block_len { + for ch in 0..self.channels { + nw.write(self.nibs[ch][i])?; + } + } + drop(nw); + } + + self.samples.drain(..self.block_len * self.channels); + let ts = NATimeInfo::new(None, None, Some(1), 1, self.srate); + Ok(NAPacket::new(self.stream.clone().unwrap(), ts, true, dbuf)) + } + fn calc_block_size(nsamps: usize, channels: usize) -> usize { + ((nsamps - 1) * channels + 1) / 2 + 4 * channels + } + fn calc_step(samp1: i16, samp2: i16) -> u8 { + let diff = (i32::from(samp1) - i32::from(samp2)).abs(); + for (i, &step) in IMA_STEP_TABLE.iter().enumerate() { + if step >= diff { + return i as u8; + } + } + IMA_MAX_STEP + } +} + +impl NAEncoder for IMAADPCMEncoder { + fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult { + match encinfo.format { + NACodecTypeInfo::None => { + let mut ofmt = EncodeParameters::default(); + ofmt.format = NACodecTypeInfo::Audio(NAAudioInfo::new(0, 1, SND_S16_FORMAT, DEFAULT_BLOCK_LEN)); + Ok(ofmt) + }, + NACodecTypeInfo::Video(_) => Err(EncoderError::FormatError), + NACodecTypeInfo::Audio(ainfo) => { + let mut outinfo = ainfo; + outinfo.channels = outinfo.channels.max(1); + if outinfo.format != SND_S16P_FORMAT && outinfo.format != SND_S16_FORMAT { + outinfo.format = SND_S16_FORMAT; + } + if outinfo.block_len == 0 { + outinfo.block_len = DEFAULT_BLOCK_LEN; + } + if outinfo.block_len < 2 { + outinfo.block_len = 2; + } + if (outinfo.block_len & 7) != 0 { + outinfo.block_len = (outinfo.block_len & 7) + 7; + } + let mut ofmt = *encinfo; + ofmt.format = NACodecTypeInfo::Audio(outinfo); + Ok(ofmt) + } + } + } + fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult { + match encinfo.format { + NACodecTypeInfo::None => Err(EncoderError::FormatError), + NACodecTypeInfo::Video(_) => Err(EncoderError::FormatError), + NACodecTypeInfo::Audio(ainfo) => { + if ainfo.format != SND_S16P_FORMAT && ainfo.format != SND_S16_FORMAT { + return Err(EncoderError::FormatError); + } + if ainfo.block_len < 2 || ((ainfo.block_len * (ainfo.channels as usize)) & 7) != 0 { + return Err(EncoderError::FormatError); + } + self.channels = ainfo.channels as usize; + self.block_len = ainfo.block_len; + self.block_size = Self::calc_block_size(self.block_len, ainfo.channels as usize); + self.first = Vec::with_capacity(self.channels); + for _ in 0..self.channels { + self.first.push(IMAState::new()); + } + + let soniton = NASoniton::new(4, 0); + let out_ainfo = NAAudioInfo::new(ainfo.sample_rate, ainfo.channels, soniton, Self::calc_block_size(self.block_len, self.channels)); + let info = NACodecInfo::new("ima-adpcm-ms", NACodecTypeInfo::Audio(out_ainfo), None); + let mut stream = NAStream::new(StreamType::Audio, stream_id, info, self.block_len as u32, ainfo.sample_rate, 0); + stream.set_num(stream_id as usize); + let stream = stream.into_ref(); + + self.stream = Some(stream.clone()); + self.samples = Vec::with_capacity(self.block_len * self.channels); + self.srate = ainfo.sample_rate; + self.flush = false; + + Ok(stream) + }, + } + } + fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> { + let buf = frm.get_buffer(); + if let Some(ref abuf) = buf.get_abuf_i16() { + let src = abuf.get_data(); + let len = abuf.get_length(); + let ch = abuf.get_chmap().num_channels(); + if abuf.get_step() > 1 || ch == 1 { + self.samples.extend_from_slice(&src[..len * ch]); + } else { + let astride = abuf.get_stride(); + self.samples.reserve(len * ch); + for i in 0..len { + for ch in 0..self.channels { + self.samples.push(src[ch * astride + i]); + } + } + } + Ok(()) + } else { + Err(EncoderError::InvalidParameters) + } + } + fn get_packet(&mut self) -> EncoderResult> { + if let Ok(pkt) = self.encode_packet() { + Ok(Some(pkt)) + } else { + Ok(None) + } + } + fn flush(&mut self) -> EncoderResult<()> { + self.flush = true; + Ok(()) + } +} + +const ENCODER_OPTS: &[NAOptionDefinition] = &[ + NAOptionDefinition { + name: "trellis", description: "Use trellis search", + opt_type: NAOptionDefinitionType::Bool }, +]; + +impl NAOptionHandler for IMAADPCMEncoder { + fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS } + fn set_options(&mut self, options: &[NAOption]) { + for option in options.iter() { + for opt_def in ENCODER_OPTS.iter() { + if opt_def.check(option).is_ok() { + match option.name { + "trellis" => { + if let NAValue::Bool(val) = option.value { + self.trellis = val; + } + }, + _ => {}, + }; + } + } + } + } + fn query_option_value(&self, name: &str) -> Option { + match name { + "trellis" => Some(NAValue::Bool(self.trellis)), + _ => None, + } + } +} + +pub fn get_encoder() -> Box { + Box::new(IMAADPCMEncoder::new()) +} + +#[cfg(test)] +mod test { + use nihav_core::codecs::*; + use nihav_core::demuxers::*; + use nihav_core::muxers::*; + use nihav_codec_support::test::enc_video::*; + use crate::*; + use nihav_commonfmt::*; + + fn test_ima_adpcm_ms_encoder(name: &'static str, trellis: bool, hash: &[u32; 4]) { + let mut dmx_reg = RegisteredDemuxers::new(); + generic_register_all_demuxers(&mut dmx_reg); + let mut dec_reg = RegisteredDecoders::new(); + generic_register_all_decoders(&mut dec_reg); + ms_register_all_decoders(&mut dec_reg); + let mut mux_reg = RegisteredMuxers::new(); + generic_register_all_muxers(&mut mux_reg); + let mut enc_reg = RegisteredEncoders::new(); + ms_register_all_encoders(&mut enc_reg); + + let dec_config = DecoderTestParams { + demuxer: "avi", + in_name: "assets/Duck/ot171_vp40.avi", + stream_type: StreamType::Audio, + limit: None, + dmx_reg, dec_reg, + }; + let enc_config = EncoderTestParams { + muxer: "wav", + enc_name: "ima-adpcm-ms", + out_name: name, + mux_reg, enc_reg, + }; + let dst_ainfo = NAAudioInfo { + sample_rate: 0, + channels: 0, + format: SND_S16_FORMAT, + block_len: 128, + }; + let enc_params = EncodeParameters { + format: NACodecTypeInfo::Audio(dst_ainfo), + quality: 0, + bitrate: 0, + tb_num: 0, + tb_den: 0, + flags: 0, + }; + let enc_options = &[ + NAOption{name: "trellis", value: NAValue::Bool(trellis)}, + ]; + test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options); + + test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options, + hash); + } + #[test] + fn test_ima_adpcm_ms_encoder_notrellis() { + test_ima_adpcm_ms_encoder("msimaadpcm-notr.wav", false, + &[0x59909f10, 0xf0420dd2, 0xcfee7ef5, 0x7623caa3]); + } + #[test] + fn test_ima_adpcm_ms_encoder_trellis() { + test_ima_adpcm_ms_encoder("msimaadpcm-tr.wav", true, + &[0x99e373e2, 0x80439b4b, 0xcb4f0b78, 0xeb1e9a51]); + } +} diff --git a/nihav-ms/src/codecs/mod.rs b/nihav-ms/src/codecs/mod.rs index 752f5e8..1216205 100644 --- a/nihav-ms/src/codecs/mod.rs +++ b/nihav-ms/src/codecs/mod.rs @@ -13,6 +13,9 @@ pub mod msvideo1; #[cfg(feature="decoder_ima_adpcm_ms")] pub mod imaadpcm; +#[cfg(feature="encoder_ima_adpcm_ms")] +pub mod imaadpcmenc; + #[cfg(any(feature="decoder_ms_adpcm", feature="encoder_ms_adpcm"))] pub mod msadpcm; @@ -42,6 +45,8 @@ const MS_ENCODERS: &[EncoderInfo] = &[ EncoderInfo { name: "msvideo1", get_encoder: msvideo1enc::get_encoder }, #[cfg(feature="encoder_ms_adpcm")] EncoderInfo { name: "ms-adpcm", get_encoder: msadpcm::get_encoder }, +#[cfg(feature="encoder_ima_adpcm_ms")] + EncoderInfo { name: "ima-adpcm-ms", get_encoder: imaadpcmenc::get_encoder }, ]; /// Registers all available encoders provided by this crate. -- 2.30.2