add MS variant of IMA ADPCM encoder
authorKostya Shishkov <kostya.shishkov@gmail.com>
Sun, 8 Aug 2021 12:02:42 +0000 (14:02 +0200)
committerKostya Shishkov <kostya.shishkov@gmail.com>
Sun, 8 Aug 2021 12:02:42 +0000 (14:02 +0200)
nihav-ms/Cargo.toml
nihav-ms/src/codecs/imaadpcmenc.rs [new file with mode: 0644]
nihav-ms/src/codecs/mod.rs

index cf9a79f..ffc6126 100644 (file)
@@ -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 (file)
index 0000000..a64f58f
--- /dev/null
@@ -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<NAStreamRef>,
+    samples:    Vec<i16>,
+    block_len:  usize,
+    block_size: usize,
+    channels:   usize,
+    flush:      bool,
+    srate:      u32,
+    trellis:    bool,
+    nodes:      Vec<[TrellisNode; 16]>,
+    first:      Vec<IMAState>,
+    nibs:       Vec<Vec<u8>>,
+}
+
+const DEFAULT_BLOCK_LEN: usize = 256;
+
+impl IMAADPCMEncoder {
+    fn new() -> Self { Self::default() }
+    fn encode_packet(&mut self) -> EncoderResult<NAPacket> {
+        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<EncodeParameters> {
+        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<NAStreamRef> {
+        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<Option<NAPacket>> {
+        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<NAValue> {
+        match name {
+            "trellis" => Some(NAValue::Bool(self.trellis)),
+            _ => None,
+        }
+    }
+}
+
+pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
+    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]);
+    }
+}
index 752f5e8..1216205 100644 (file)
@@ -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.