]> git.nihav.org Git - nihav.git/blame_incremental - nihav-ms/src/codecs/imaadpcmenc.rs
aacsbr: fix high<->low resolution envelope mapping
[nihav.git] / nihav-ms / src / codecs / imaadpcmenc.rs
... / ...
CommitLineData
1use nihav_core::codecs::*;
2use nihav_core::io::byteio::*;
3use nihav_codec_support::codecs::imaadpcm::*;
4
5struct NibbleWriter<'a> {
6 bw: ByteWriter<'a>,
7 val: u8,
8 first: bool,
9}
10
11impl<'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)]
29struct TrellisNode {
30 state: IMAState,
31 nib: u8,
32 error: i64,
33}
34
35impl 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)]
46struct 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
60const DEFAULT_BLOCK_LEN: usize = 256;
61
62impl 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
188impl NAEncoder for IMAADPCMEncoder {
189 fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
190 match encinfo.format {
191 NACodecTypeInfo::None => {
192 Ok(EncodeParameters {
193 format: NACodecTypeInfo::Audio(NAAudioInfo::new(0, 1, SND_S16_FORMAT, DEFAULT_BLOCK_LEN)),
194 ..Default::default() })
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
288const ENCODER_OPTS: &[NAOptionDefinition] = &[
289 NAOptionDefinition {
290 name: "trellis", description: "Use trellis search",
291 opt_type: NAOptionDefinitionType::Bool },
292];
293
294impl 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
320pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
321 Box::new(IMAADPCMEncoder::new())
322}
323
324#[cfg(test)]
325mod 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}