]> git.nihav.org Git - nihav.git/blob - nihav-duck/src/codecs/truemotionrtenc.rs
add MPEG-4 ASP decoder
[nihav.git] / nihav-duck / src / codecs / truemotionrtenc.rs
1 use nihav_core::codecs::*;
2 use nihav_core::io::byteio::*;
3 use nihav_core::io::bitwriter::*;
4
5 const TMRT_DELTA_TAB: [&[i16]; 3] = [
6 &[ 5, -7, 36, -36 ],
7 &[ 2, -3, 8, -8, 18, -18, 36, -36 ],
8 &[ 1, -1, 2, -3, 8, -8, 18, -18, 36, -36, 54, -54, 96, -96, 144, -144 ]
9 ];
10
11 const FIRST_NODE: u8 = 255;
12 const ERR_MAX: u32 = std::u32::MAX;
13
14 #[derive(Clone, Copy, Default)]
15 struct TrellisNode {
16 err: u32,
17 hpred: i16,
18 idx: u8,
19 }
20
21 struct TMRTEncoder {
22 stream: Option<NAStreamRef>,
23 pkt: Option<NAPacket>,
24 dbits: u8,
25 hscale: bool,
26 do_trellis: bool,
27 top_line: Vec<u8>,
28 trellis: Vec<TrellisNode>,
29 indices: Vec<u8>,
30 }
31
32 fn find_delta(cur_delta: i16, delta_tab: &[i16]) -> (i16, usize) {
33 let mut idx = 0;
34 let mut best_diff = 512;
35 let mut ndelta = delta_tab[0];
36 for (i, &delta) in delta_tab.iter().enumerate() {
37 let ddiff = (delta - cur_delta).abs();
38 if ddiff < best_diff {
39 idx = i;
40 best_diff = ddiff;
41 ndelta = delta;
42 }
43 }
44 (ndelta, idx)
45 }
46
47 impl TMRTEncoder {
48 fn new() -> Self {
49 Self {
50 stream: None,
51 pkt: None,
52 dbits: 4,
53 hscale: false,
54 do_trellis: false,
55 top_line: Vec::new(),
56 trellis: Vec::new(),
57 indices: Vec::new(),
58 }
59 }
60 fn encode_plane(&mut self, bw: &mut BitWriter, in_frm: &NAVideoBuffer<u8>, plane_no: usize) -> EncoderResult<()> {
61 let (width, height) = in_frm.get_dimensions(plane_no);
62
63 let stride = in_frm.get_stride(plane_no);
64 let offset = in_frm.get_offset(plane_no);
65 let src = in_frm.get_data();
66
67 let delta_tab = TMRT_DELTA_TAB[(self.dbits - 2) as usize];
68 self.top_line.clear();
69 self.top_line.resize(width, if plane_no == 0 { 0 } else { 0x80 });
70
71 let step = if self.hscale { 2 } else { 1 };
72 for line in src[offset..].chunks(stride).take(height) {
73 let mut hor_pred = 0;
74 for (&cur, pred) in line[..width].iter().zip(self.top_line.iter_mut()).step_by(step) {
75 let cur = i16::from(cur);
76 let cur_delta = cur - i16::from(*pred) - hor_pred;
77
78 let (ndelta, idx) = find_delta(cur_delta, delta_tab);
79
80 bw.write(idx as u32, self.dbits);
81 hor_pred += ndelta;
82 *pred = (i16::from(*pred) + hor_pred).max(0).min(255) as u8;
83 }
84 }
85
86 Ok(())
87 }
88
89 fn encode_plane_trellis(&mut self, bw: &mut BitWriter, in_frm: &NAVideoBuffer<u8>, plane_no: usize) -> EncoderResult<()> {
90 let (width, height) = in_frm.get_dimensions(plane_no);
91
92 let stride = in_frm.get_stride(plane_no);
93 let offset = in_frm.get_offset(plane_no);
94 let src = in_frm.get_data();
95
96 let delta_tab = TMRT_DELTA_TAB[(self.dbits - 2) as usize];
97 self.top_line.clear();
98 self.top_line.resize(width, if plane_no == 0 { 0 } else { 0x80 });
99
100 let trellis_size = delta_tab.len();
101 self.trellis.resize(trellis_size * (width + 1), TrellisNode::default());
102 self.indices.resize(width, 0);
103 for node in self.trellis[..trellis_size].iter_mut() {
104 node.idx = FIRST_NODE;
105 }
106
107 let step = if self.hscale { 2 } else { 1 };
108 for line in src[offset..].chunks(stride).take(height) {
109 let mut tsplit = trellis_size;
110 for (&cur, &pred) in line[..width].iter().zip(self.top_line.iter()).step_by(step) {
111 let (tprev, tcur) = self.trellis.split_at_mut(tsplit);
112 let hist = &tprev[tprev.len() - trellis_size..];
113
114 let pix_val = i32::from(cur);
115 let top_val = i16::from(pred);
116 for (dst, &delta) in tcur.iter_mut().zip(delta_tab.iter()) {
117 dst.err = ERR_MAX;
118 for (idx, src) in hist.iter().enumerate() {
119 if src.err == ERR_MAX {
120 continue;
121 }
122 let nval = i32::from((top_val + src.hpred + delta).max(0).min(255));
123 let new_err = src.err + (((nval - pix_val) * (nval - pix_val)) as u32);
124 if new_err < dst.err {
125 dst.err = new_err;
126 dst.idx = idx as u8;
127 dst.hpred = src.hpred + delta;
128 }
129 }
130 }
131
132 tsplit += trellis_size;
133 }
134
135 tsplit -= trellis_size;
136 let mut best_idx = 0;
137 let mut best_err = self.trellis[tsplit].err;
138
139 for (idx, node) in self.trellis[tsplit..].iter().take(trellis_size).enumerate() {
140 if node.err < best_err {
141 best_idx = idx as u8;
142 best_err = node.err;
143 }
144 }
145
146 let mut cur_idx = best_idx;
147 for dst in self.indices[..width / step].iter_mut().rev() {
148 *dst = cur_idx;
149 cur_idx = self.trellis[tsplit + (cur_idx as usize)].idx;
150 tsplit -= trellis_size;
151 }
152
153 let mut hor_pred = 0;
154 for (pred, &idx) in self.top_line.iter_mut().step_by(step).zip(self.indices.iter()) {
155 bw.write(u32::from(idx), self.dbits);
156 hor_pred += delta_tab[usize::from(idx)];
157 *pred = (i16::from(*pred) + hor_pred).max(0).min(255) as u8;
158 }
159 }
160
161 Ok(())
162 }
163 }
164
165 impl NAEncoder for TMRTEncoder {
166 fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
167 match encinfo.format {
168 NACodecTypeInfo::None => {
169 Ok(EncodeParameters {
170 format: NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, true, YUV410_FORMAT)),
171 ..Default::default()
172 })
173 },
174 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
175 NACodecTypeInfo::Video(vinfo) => {
176 let pix_fmt = YUV410_FORMAT;
177 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, false, pix_fmt);
178 let mut ofmt = *encinfo;
179 ofmt.format = NACodecTypeInfo::Video(outinfo);
180 Ok(ofmt)
181 }
182 }
183 }
184 fn get_capabilities(&self) -> u64 { ENC_CAPS_PARAMCHANGE }
185 fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
186 match encinfo.format {
187 NACodecTypeInfo::None => Err(EncoderError::FormatError),
188 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
189 NACodecTypeInfo::Video(vinfo) => {
190 if vinfo.format != YUV410_FORMAT {
191 return Err(EncoderError::FormatError);
192 }
193 if ((vinfo.width | vinfo.height) & 3) != 0 {
194 return Err(EncoderError::FormatError);
195 }
196 if (vinfo.width | vinfo.height) >= (1 << 10) {
197 return Err(EncoderError::FormatError);
198 }
199
200 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, false, vinfo.format);
201 let info = NACodecInfo::new("truemotionrt", NACodecTypeInfo::Video(out_info), None);
202 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
203 stream.set_num(stream_id as usize);
204 let stream = stream.into_ref();
205
206 self.stream = Some(stream.clone());
207
208 Ok(stream)
209 },
210 }
211 }
212 fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
213 let buf = frm.get_buffer();
214 if let Some(ref vbuf) = buf.get_vbuf() {
215 let mut dbuf = Vec::with_capacity(4);
216 let mut gw = GrowableMemoryWriter::new_write(&mut dbuf);
217 let mut bw = ByteWriter::new(&mut gw);
218
219 bw.write_byte(0)?; // header size
220 bw.write_byte(17)?;
221 bw.write_byte(self.dbits)?;
222 bw.write_byte(5)?;
223 bw.write_byte(self.hscale as u8)?;
224 bw.write_byte(0)?; // pad
225 let (width, height) = vbuf.get_dimensions(0);
226 bw.write_u16le(height as u16)?;
227 bw.write_u16le(width as u16)?;
228 bw.write_u16le(0)?; // pad
229 bw.write_u32le(0)?; // full frame_size
230 bw.write_u16le(0)?; // pad
231
232 let hdr_size = bw.tell() as usize;
233
234 bw.write_u32le(0)?; // data size
235
236 let mut bw = BitWriter::new(dbuf, BitWriterMode::LE);
237 if !self.do_trellis {
238 for plane in 0..3 {
239 self.encode_plane(&mut bw, vbuf, plane)?;
240 }
241 } else {
242 for plane in 0..3 {
243 self.encode_plane_trellis(&mut bw, vbuf, plane)?;
244 }
245 }
246 dbuf = bw.end();
247
248 let frame_size = dbuf.len() as u32;
249 write_u32le(&mut dbuf[12..], frame_size)?;
250 write_u32le(&mut dbuf[hdr_size..], frame_size - (hdr_size as u32) - 4)?;
251
252 dbuf[0] = ((hdr_size | 0x80) as u8).rotate_right(3);
253 for i in (1..hdr_size).rev() {
254 dbuf[i] ^= dbuf[i + 1];
255 }
256
257 self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, true, dbuf));
258 Ok(())
259 } else {
260 Err(EncoderError::InvalidParameters)
261 }
262 }
263 fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
264 let mut npkt = None;
265 std::mem::swap(&mut self.pkt, &mut npkt);
266 Ok(npkt)
267 }
268 fn flush(&mut self) -> EncoderResult<()> {
269 Ok(())
270 }
271 }
272
273 const ENCODER_OPTS: &[NAOptionDefinition] = &[
274 NAOptionDefinition {
275 name: "bits", description: "Bits per delta",
276 opt_type: NAOptionDefinitionType::Int(Some(2), Some(4)) },
277 NAOptionDefinition {
278 name: "hscale", description: "Horizontal scaling mode",
279 opt_type: NAOptionDefinitionType::Bool },
280 NAOptionDefinition {
281 name: "trellis", description: "Trellis search for optimal deltas",
282 opt_type: NAOptionDefinitionType::Bool },
283 ];
284
285 impl NAOptionHandler for TMRTEncoder {
286 fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
287 fn set_options(&mut self, options: &[NAOption]) {
288 for option in options.iter() {
289 for opt_def in ENCODER_OPTS.iter() {
290 if opt_def.check(option).is_ok() {
291 match option.name {
292 "bits" => {
293 if let NAValue::Int(val) = option.value {
294 self.dbits = val as u8;
295 }
296 },
297 "hscale" => {
298 if let NAValue::Bool(val) = option.value {
299 self.hscale = val;
300 }
301 },
302 "trellis" => {
303 if let NAValue::Bool(val) = option.value {
304 self.do_trellis = val;
305 }
306 },
307 _ => {},
308 };
309 }
310 }
311 }
312 }
313 fn query_option_value(&self, name: &str) -> Option<NAValue> {
314 match name {
315 "bits" => Some(NAValue::Int(i64::from(self.dbits))),
316 "hscale" => Some(NAValue::Bool(self.hscale)),
317 "trellis" => Some(NAValue::Bool(self.do_trellis)),
318 _ => None,
319 }
320 }
321 }
322
323 pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
324 Box::new(TMRTEncoder::new())
325 }
326
327 #[cfg(test)]
328 mod test {
329 use nihav_core::codecs::*;
330 use nihav_core::demuxers::*;
331 use nihav_core::muxers::*;
332 use crate::*;
333 use nihav_commonfmt::*;
334 use nihav_codec_support::test::enc_video::*;
335
336 fn encode_test(name: &'static str, enc_options: &[NAOption], hash: &[u32; 4]) {
337 let mut dmx_reg = RegisteredDemuxers::new();
338 generic_register_all_demuxers(&mut dmx_reg);
339 let mut dec_reg = RegisteredDecoders::new();
340 duck_register_all_decoders(&mut dec_reg);
341 let mut mux_reg = RegisteredMuxers::new();
342 generic_register_all_muxers(&mut mux_reg);
343 let mut enc_reg = RegisteredEncoders::new();
344 duck_register_all_encoders(&mut enc_reg);
345
346 // sample from private collection
347 let dec_config = DecoderTestParams {
348 demuxer: "avi",
349 in_name: "assets/Duck/tr20_low.avi",
350 stream_type: StreamType::Video,
351 limit: Some(1),
352 dmx_reg, dec_reg,
353 };
354 let enc_config = EncoderTestParams {
355 muxer: "avi",
356 enc_name: "truemotionrt",
357 out_name: name,
358 mux_reg, enc_reg,
359 };
360 let dst_vinfo = NAVideoInfo {
361 width: 0,
362 height: 0,
363 format: YUV410_FORMAT,
364 flipped: false,
365 bits: 9,
366 };
367 let enc_params = EncodeParameters {
368 format: NACodecTypeInfo::Video(dst_vinfo),
369 quality: 0,
370 bitrate: 0,
371 tb_num: 0,
372 tb_den: 0,
373 flags: 0,
374 };
375 //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
376 test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options, hash);
377 }
378 #[test]
379 fn test_truemotionrt_encoder_2bit() {
380 let enc_options = &[
381 NAOption { name: "bits", value: NAValue::Int(2) },
382 ];
383 encode_test("tmrt-2bit.avi", enc_options, &[0x2c2a5ae3, 0xde1646e4, 0xf76bb219, 0xd09602fa]);
384 }
385 #[test]
386 fn test_truemotionrt_encoder_3bit() {
387 let enc_options = &[
388 NAOption { name: "bits", value: NAValue::Int(3) },
389 ];
390 encode_test("tmrt-3bit.avi", enc_options, &[0x36cf8f48, 0x3e8ff2ce, 0x6f3822cf, 0xf7fbf19d]);
391 }
392 #[test]
393 fn test_truemotionrt_encoder_4bit() {
394 let enc_options = &[
395 NAOption { name: "bits", value: NAValue::Int(4) },
396 ];
397 encode_test("tmrt-4bit.avi", enc_options, &[0xa5a7fbe3, 0x7bac0b2b, 0x2af6f97f, 0xa65cd1fc]);
398 }
399 #[test]
400 fn test_truemotionrt_encoder_hscale() {
401 let enc_options = &[
402 NAOption { name: "bits", value: NAValue::Int(3) },
403 NAOption { name: "hscale", value: NAValue::Bool(true) },
404 ];
405 encode_test("tmrt-hscale.avi", enc_options, &[0xc17afa21, 0x5bdf49c9, 0x57997840, 0xfc2f17b6]);
406 }
407 #[test]
408 fn test_truemotionrt_encoder_trellis() {
409 let enc_options = &[
410 NAOption { name: "bits", value: NAValue::Int(3) },
411 NAOption { name: "hscale", value: NAValue::Bool(true) },
412 NAOption { name: "trellis", value: NAValue::Bool(true) },
413 ];
414 encode_test("tmrt-trellis.avi", enc_options, &[0x3586b450, 0x6ea1ed31, 0xe14c0c7d, 0x0886bc4f]);
415 }
416 }