| 1 | use nihav_core::codecs::*; |
| 2 | use nihav_core::io::byteio::*; |
| 3 | use nihav_codec_support::codecs::imaadpcm::*; |
| 4 | |
| 5 | struct NibbleWriter<'a> { |
| 6 | bw: ByteWriter<'a>, |
| 7 | val: u8, |
| 8 | first: bool, |
| 9 | } |
| 10 | |
| 11 | impl<'a> NibbleWriter<'a> { |
| 12 | fn new(bw: ByteWriter<'a>) -> Self { |
| 13 | Self { bw, val: 0, first: true } |
| 14 | } |
| 15 | fn write(&mut self, nib: u8) -> EncoderResult<()> { |
| 16 | if self.first { |
| 17 | self.val = nib; |
| 18 | self.first = false; |
| 19 | } else { |
| 20 | self.val |= nib << 4; |
| 21 | self.bw.write_byte(self.val)?; |
| 22 | self.first = true; |
| 23 | } |
| 24 | Ok(()) |
| 25 | } |
| 26 | } |
| 27 | |
| 28 | #[derive(Clone,Copy)] |
| 29 | struct TrellisNode { |
| 30 | state: IMAState, |
| 31 | nib: u8, |
| 32 | error: i64, |
| 33 | } |
| 34 | |
| 35 | impl Default for TrellisNode { |
| 36 | fn default() -> Self { |
| 37 | TrellisNode { |
| 38 | state: IMAState::new(), |
| 39 | nib: 0, |
| 40 | error: 0, |
| 41 | } |
| 42 | } |
| 43 | } |
| 44 | |
| 45 | #[derive(Default)] |
| 46 | struct IMAADPCMEncoder { |
| 47 | stream: Option<NAStreamRef>, |
| 48 | samples: Vec<i16>, |
| 49 | block_len: usize, |
| 50 | block_size: usize, |
| 51 | channels: usize, |
| 52 | flush: bool, |
| 53 | srate: u32, |
| 54 | trellis: bool, |
| 55 | nodes: Vec<[TrellisNode; 16]>, |
| 56 | first: Vec<IMAState>, |
| 57 | nibs: Vec<Vec<u8>>, |
| 58 | } |
| 59 | |
| 60 | const DEFAULT_BLOCK_LEN: usize = 256; |
| 61 | |
| 62 | impl IMAADPCMEncoder { |
| 63 | fn new() -> Self { Self::default() } |
| 64 | fn encode_packet(&mut self) -> EncoderResult<NAPacket> { |
| 65 | if self.samples.is_empty() { |
| 66 | return Err(EncoderError::TryAgain); |
| 67 | } |
| 68 | let cur_len = (self.samples.len() / self.channels).min(self.block_len); |
| 69 | if cur_len < self.block_len && !self.flush { |
| 70 | return Err(EncoderError::TryAgain); |
| 71 | } |
| 72 | if cur_len < self.block_len { |
| 73 | self.samples.resize(self.block_len * self.channels, 0); |
| 74 | } |
| 75 | |
| 76 | let mut dbuf = vec![0u8; self.block_size]; |
| 77 | let mut mw = MemoryWriter::new_write(dbuf.as_mut_slice()); |
| 78 | let mut bw = ByteWriter::new(&mut mw); |
| 79 | |
| 80 | for ch in 0..self.channels { |
| 81 | self.first[ch].predictor = i32::from(self.samples[ch]); |
| 82 | self.first[ch].step = Self::calc_step(self.samples[ch], self.samples[ch + self.channels]) as usize; |
| 83 | } |
| 84 | |
| 85 | if !self.trellis { |
| 86 | for ch in 0..self.channels { |
| 87 | bw.write_u16le(self.first[ch].predictor as u16)?; |
| 88 | bw.write_byte(self.first[ch].step as u8)?; |
| 89 | bw.write_byte(0)?; |
| 90 | } |
| 91 | let mut nw = NibbleWriter::new(bw); |
| 92 | for samples in self.samples.chunks(self.channels).take(self.block_len).skip(1) { |
| 93 | for (state, &samp) in self.first.iter_mut().zip(samples.iter()) { |
| 94 | let nib = state.compress_sample(samp); |
| 95 | state.expand_sample(nib); |
| 96 | nw.write(nib)?; |
| 97 | } |
| 98 | } |
| 99 | drop(nw); |
| 100 | } else { |
| 101 | self.nodes.reserve(self.block_len); |
| 102 | self.nibs.resize(self.channels, Vec::new()); |
| 103 | for nibs in self.nibs.iter_mut() { |
| 104 | nibs.resize(self.block_len, 0); |
| 105 | } |
| 106 | let step = self.channels; |
| 107 | let mut state = [TrellisNode::default(); 16]; |
| 108 | for ch in 0..self.channels { |
| 109 | self.nodes.clear(); |
| 110 | for i in 0..16 { |
| 111 | let step = (((self.first[ch].step + i) as isize) - 8).max(0).min(IMA_MAX_STEP as isize) as usize; |
| 112 | state[i].state.predictor = self.first[ch].predictor; |
| 113 | state[i].state.step = step; |
| 114 | state[i].error = 0; |
| 115 | state[i].nib = step as u8; |
| 116 | } |
| 117 | self.nodes.push(state); |
| 118 | let mut sidx = ch + step; |
| 119 | for _i in 1..self.block_len { |
| 120 | let sample = self.samples[sidx]; |
| 121 | |
| 122 | for (nnib, cstate) in state.iter_mut().enumerate() { |
| 123 | for nib in 0..16 { |
| 124 | let pnode = &self.nodes[self.nodes.len() - 1][nib]; |
| 125 | let mut ima = pnode.state; |
| 126 | let nsamp = ima.expand_sample(nnib as u8); |
| 127 | let diff = i64::from(i32::from(sample) - i32::from(nsamp)); |
| 128 | let error = pnode.error + diff * diff; |
| 129 | if (nib == 0) || error < cstate.error { |
| 130 | cstate.state = ima; |
| 131 | cstate.nib = nib as u8; |
| 132 | cstate.error = error; |
| 133 | } |
| 134 | } |
| 135 | } |
| 136 | self.nodes.push(state); |
| 137 | |
| 138 | sidx += step; |
| 139 | } |
| 140 | |
| 141 | let mut idx = 0; |
| 142 | let mut best_err = self.nodes[self.nodes.len() - 1][0].error; |
| 143 | for (i, node) in self.nodes[self.nodes.len() - 1].iter().enumerate() { |
| 144 | if node.error < best_err { |
| 145 | best_err = node.error; |
| 146 | idx = i; |
| 147 | } |
| 148 | } |
| 149 | let mut dst = self.nibs[ch].iter_mut().rev(); |
| 150 | while let Some(nodes) = self.nodes.pop() { |
| 151 | *dst.next().unwrap() = idx as u8; |
| 152 | idx = nodes[idx].nib as usize; |
| 153 | } |
| 154 | self.nibs[ch][0] = idx as u8; |
| 155 | } |
| 156 | for ch in 0..self.channels { |
| 157 | bw.write_u16le(self.first[ch].predictor as u16)?; |
| 158 | bw.write_byte(self.nibs[ch][0])?; |
| 159 | bw.write_byte(0)?; |
| 160 | } |
| 161 | let mut nw = NibbleWriter::new(bw); |
| 162 | for i in 1..self.block_len { |
| 163 | for ch in 0..self.channels { |
| 164 | nw.write(self.nibs[ch][i])?; |
| 165 | } |
| 166 | } |
| 167 | drop(nw); |
| 168 | } |
| 169 | |
| 170 | self.samples.drain(..self.block_len * self.channels); |
| 171 | let ts = NATimeInfo::new(None, None, Some(1), 1, self.srate); |
| 172 | Ok(NAPacket::new(self.stream.clone().unwrap(), ts, true, dbuf)) |
| 173 | } |
| 174 | fn calc_block_size(nsamps: usize, channels: usize) -> usize { |
| 175 | ((nsamps - 1) * channels + 1) / 2 + 4 * channels |
| 176 | } |
| 177 | fn calc_step(samp1: i16, samp2: i16) -> u8 { |
| 178 | let diff = (i32::from(samp1) - i32::from(samp2)).abs(); |
| 179 | for (i, &step) in IMA_STEP_TABLE.iter().enumerate() { |
| 180 | if step >= diff { |
| 181 | return i as u8; |
| 182 | } |
| 183 | } |
| 184 | IMA_MAX_STEP |
| 185 | } |
| 186 | } |
| 187 | |
| 188 | impl NAEncoder for IMAADPCMEncoder { |
| 189 | fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> { |
| 190 | match encinfo.format { |
| 191 | NACodecTypeInfo::None => { |
| 192 | let mut ofmt = EncodeParameters::default(); |
| 193 | ofmt.format = NACodecTypeInfo::Audio(NAAudioInfo::new(0, 1, SND_S16_FORMAT, DEFAULT_BLOCK_LEN)); |
| 194 | Ok(ofmt) |
| 195 | }, |
| 196 | NACodecTypeInfo::Video(_) => Err(EncoderError::FormatError), |
| 197 | NACodecTypeInfo::Audio(ainfo) => { |
| 198 | let mut outinfo = ainfo; |
| 199 | outinfo.channels = outinfo.channels.max(1); |
| 200 | if outinfo.format != SND_S16P_FORMAT && outinfo.format != SND_S16_FORMAT { |
| 201 | outinfo.format = SND_S16_FORMAT; |
| 202 | } |
| 203 | if outinfo.block_len == 0 { |
| 204 | outinfo.block_len = DEFAULT_BLOCK_LEN; |
| 205 | } |
| 206 | if outinfo.block_len < 2 { |
| 207 | outinfo.block_len = 2; |
| 208 | } |
| 209 | if (outinfo.block_len & 7) != 0 { |
| 210 | outinfo.block_len = (outinfo.block_len & 7) + 7; |
| 211 | } |
| 212 | let mut ofmt = *encinfo; |
| 213 | ofmt.format = NACodecTypeInfo::Audio(outinfo); |
| 214 | Ok(ofmt) |
| 215 | } |
| 216 | } |
| 217 | } |
| 218 | fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> { |
| 219 | match encinfo.format { |
| 220 | NACodecTypeInfo::None => Err(EncoderError::FormatError), |
| 221 | NACodecTypeInfo::Video(_) => Err(EncoderError::FormatError), |
| 222 | NACodecTypeInfo::Audio(ainfo) => { |
| 223 | if ainfo.format != SND_S16P_FORMAT && ainfo.format != SND_S16_FORMAT { |
| 224 | return Err(EncoderError::FormatError); |
| 225 | } |
| 226 | if ainfo.block_len < 2 || ((ainfo.block_len * (ainfo.channels as usize)) & 7) != 0 { |
| 227 | return Err(EncoderError::FormatError); |
| 228 | } |
| 229 | self.channels = ainfo.channels as usize; |
| 230 | self.block_len = ainfo.block_len; |
| 231 | self.block_size = Self::calc_block_size(self.block_len, ainfo.channels as usize); |
| 232 | self.first = Vec::with_capacity(self.channels); |
| 233 | for _ in 0..self.channels { |
| 234 | self.first.push(IMAState::new()); |
| 235 | } |
| 236 | |
| 237 | let soniton = NASoniton::new(4, 0); |
| 238 | let out_ainfo = NAAudioInfo::new(ainfo.sample_rate, ainfo.channels, soniton, Self::calc_block_size(self.block_len, self.channels)); |
| 239 | let info = NACodecInfo::new("ima-adpcm-ms", NACodecTypeInfo::Audio(out_ainfo), None); |
| 240 | let mut stream = NAStream::new(StreamType::Audio, stream_id, info, self.block_len as u32, ainfo.sample_rate, 0); |
| 241 | stream.set_num(stream_id as usize); |
| 242 | let stream = stream.into_ref(); |
| 243 | |
| 244 | self.stream = Some(stream.clone()); |
| 245 | self.samples = Vec::with_capacity(self.block_len * self.channels); |
| 246 | self.srate = ainfo.sample_rate; |
| 247 | self.flush = false; |
| 248 | |
| 249 | Ok(stream) |
| 250 | }, |
| 251 | } |
| 252 | } |
| 253 | fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> { |
| 254 | let buf = frm.get_buffer(); |
| 255 | if let Some(ref abuf) = buf.get_abuf_i16() { |
| 256 | let src = abuf.get_data(); |
| 257 | let len = abuf.get_length(); |
| 258 | let ch = abuf.get_chmap().num_channels(); |
| 259 | if abuf.get_step() > 1 || ch == 1 { |
| 260 | self.samples.extend_from_slice(&src[..len * ch]); |
| 261 | } else { |
| 262 | let astride = abuf.get_stride(); |
| 263 | self.samples.reserve(len * ch); |
| 264 | for i in 0..len { |
| 265 | for ch in 0..self.channels { |
| 266 | self.samples.push(src[ch * astride + i]); |
| 267 | } |
| 268 | } |
| 269 | } |
| 270 | Ok(()) |
| 271 | } else { |
| 272 | Err(EncoderError::InvalidParameters) |
| 273 | } |
| 274 | } |
| 275 | fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> { |
| 276 | if let Ok(pkt) = self.encode_packet() { |
| 277 | Ok(Some(pkt)) |
| 278 | } else { |
| 279 | Ok(None) |
| 280 | } |
| 281 | } |
| 282 | fn flush(&mut self) -> EncoderResult<()> { |
| 283 | self.flush = true; |
| 284 | Ok(()) |
| 285 | } |
| 286 | } |
| 287 | |
| 288 | const ENCODER_OPTS: &[NAOptionDefinition] = &[ |
| 289 | NAOptionDefinition { |
| 290 | name: "trellis", description: "Use trellis search", |
| 291 | opt_type: NAOptionDefinitionType::Bool }, |
| 292 | ]; |
| 293 | |
| 294 | impl NAOptionHandler for IMAADPCMEncoder { |
| 295 | fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS } |
| 296 | fn set_options(&mut self, options: &[NAOption]) { |
| 297 | for option in options.iter() { |
| 298 | for opt_def in ENCODER_OPTS.iter() { |
| 299 | if opt_def.check(option).is_ok() { |
| 300 | match option.name { |
| 301 | "trellis" => { |
| 302 | if let NAValue::Bool(val) = option.value { |
| 303 | self.trellis = val; |
| 304 | } |
| 305 | }, |
| 306 | _ => {}, |
| 307 | }; |
| 308 | } |
| 309 | } |
| 310 | } |
| 311 | } |
| 312 | fn query_option_value(&self, name: &str) -> Option<NAValue> { |
| 313 | match name { |
| 314 | "trellis" => Some(NAValue::Bool(self.trellis)), |
| 315 | _ => None, |
| 316 | } |
| 317 | } |
| 318 | } |
| 319 | |
| 320 | pub fn get_encoder() -> Box<dyn NAEncoder + Send> { |
| 321 | Box::new(IMAADPCMEncoder::new()) |
| 322 | } |
| 323 | |
| 324 | #[cfg(test)] |
| 325 | mod test { |
| 326 | use nihav_core::codecs::*; |
| 327 | use nihav_core::demuxers::*; |
| 328 | use nihav_core::muxers::*; |
| 329 | use nihav_codec_support::test::enc_video::*; |
| 330 | use crate::*; |
| 331 | use nihav_commonfmt::*; |
| 332 | |
| 333 | fn test_ima_adpcm_ms_encoder(name: &'static str, trellis: bool, hash: &[u32; 4]) { |
| 334 | let mut dmx_reg = RegisteredDemuxers::new(); |
| 335 | generic_register_all_demuxers(&mut dmx_reg); |
| 336 | let mut dec_reg = RegisteredDecoders::new(); |
| 337 | generic_register_all_decoders(&mut dec_reg); |
| 338 | ms_register_all_decoders(&mut dec_reg); |
| 339 | let mut mux_reg = RegisteredMuxers::new(); |
| 340 | generic_register_all_muxers(&mut mux_reg); |
| 341 | let mut enc_reg = RegisteredEncoders::new(); |
| 342 | ms_register_all_encoders(&mut enc_reg); |
| 343 | |
| 344 | // sample: https://samples.mplayerhq.hu/V-codecs/VP4/ot171_vp40.avi |
| 345 | let dec_config = DecoderTestParams { |
| 346 | demuxer: "avi", |
| 347 | in_name: "assets/Duck/ot171_vp40.avi", |
| 348 | stream_type: StreamType::Audio, |
| 349 | limit: None, |
| 350 | dmx_reg, dec_reg, |
| 351 | }; |
| 352 | let enc_config = EncoderTestParams { |
| 353 | muxer: "wav", |
| 354 | enc_name: "ima-adpcm-ms", |
| 355 | out_name: name, |
| 356 | mux_reg, enc_reg, |
| 357 | }; |
| 358 | let dst_ainfo = NAAudioInfo { |
| 359 | sample_rate: 0, |
| 360 | channels: 0, |
| 361 | format: SND_S16_FORMAT, |
| 362 | block_len: 128, |
| 363 | }; |
| 364 | let enc_params = EncodeParameters { |
| 365 | format: NACodecTypeInfo::Audio(dst_ainfo), |
| 366 | quality: 0, |
| 367 | bitrate: 0, |
| 368 | tb_num: 0, |
| 369 | tb_den: 0, |
| 370 | flags: 0, |
| 371 | }; |
| 372 | let enc_options = &[ |
| 373 | NAOption{name: "trellis", value: NAValue::Bool(trellis)}, |
| 374 | ]; |
| 375 | // test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options); |
| 376 | |
| 377 | test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options, |
| 378 | hash); |
| 379 | } |
| 380 | #[test] |
| 381 | fn test_ima_adpcm_ms_encoder_notrellis() { |
| 382 | test_ima_adpcm_ms_encoder("msimaadpcm-notr.wav", false, |
| 383 | &[0x59909f10, 0xf0420dd2, 0xcfee7ef5, 0x7623caa3]); |
| 384 | } |
| 385 | #[test] |
| 386 | fn test_ima_adpcm_ms_encoder_trellis() { |
| 387 | test_ima_adpcm_ms_encoder("msimaadpcm-tr.wav", true, |
| 388 | &[0x99e373e2, 0x80439b4b, 0xcb4f0b78, 0xeb1e9a51]); |
| 389 | } |
| 390 | } |