]> git.nihav.org Git - nihav.git/blame - nihav-ms/src/codecs/msvideo1enc.rs
core/frame: add truncate() for NAAudioBuffer
[nihav.git] / nihav-ms / src / codecs / msvideo1enc.rs
CommitLineData
dab59886
KS
1use nihav_core::codecs::*;
2use nihav_core::io::byteio::*;
3use nihav_codec_support::vq::*;
4
5#[derive(Default,Clone,Copy,PartialEq)]
6struct Pixel16(u16);
7
8impl Pixel16 {
4efceb69 9 fn unpack(self) -> (u8, u8, u8) {
dab59886
KS
10 (((self.0 >> 10) & 0x1F) as u8, ((self.0 >> 5) & 0x1F) as u8, (self.0 & 0x1F) as u8)
11 }
12 fn pack(r: u8, g: u8, b: u8) -> Self {
13 Pixel16{ 0: (u16::from(r) << 10) | (u16::from(g) << 5) | u16::from(b) }
14 }
15}
16impl VQElement for Pixel16 {
17 fn dist(&self, rval: Self) -> u32 {
18 let (r0, g0, b0) = self.unpack();
19 let (r1, g1, b1) = rval.unpack();
20 let rd = i32::from(r0) - i32::from(r1);
21 let gd = i32::from(g0) - i32::from(g1);
22 let bd = i32::from(b0) - i32::from(b1);
23 (rd * rd + gd * gd + bd * bd) as u32
24 }
25 fn min_cw() -> Self { Pixel16(0x0000) }
26 fn max_cw() -> Self { Pixel16(0x7FFF) }
27 fn min(&self, rval: Self) -> Self {
28 let (r0, g0, b0) = self.unpack();
29 let (r1, g1, b1) = rval.unpack();
30 Self::pack(r0.min(r1), g0.min(g1), b0.min(b1))
31 }
32 fn max(&self, rval: Self) -> Self {
33 let (r0, g0, b0) = self.unpack();
34 let (r1, g1, b1) = rval.unpack();
35 Self::pack(r0.max(r1), g0.max(g1), b0.max(b1))
36 }
37 fn num_components() -> usize { 3 }
38 fn sort_by_component(arr: &mut [Self], component: usize) {
39 let mut counts = [0; 32];
40 for pix in arr.iter() {
41 let (r, g, b) = pix.unpack();
42 let idx = match component {
43 0 => r,
44 1 => g,
45 _ => b,
46 } as usize;
47 counts[idx] += 1;
48 }
49 let mut offs = [0; 32];
50 for i in 0..31 {
51 offs[i + 1] = offs[i] + counts[i];
52 }
53 let mut dst = vec![Pixel16(0); arr.len()];
54 for pix in arr.iter() {
55 let (r, g, b) = pix.unpack();
56 let idx = match component {
57 0 => r,
58 1 => g,
59 _ => b,
60 } as usize;
61 dst[offs[idx]] = *pix;
62 offs[idx] += 1;
63 }
64 arr.copy_from_slice(dst.as_slice());
65 }
66 fn max_dist_component(min: &Self, max: &Self) -> usize {
67 let (r0, g0, b0) = max.unpack();
68 let (r1, g1, b1) = min.unpack();
69 let rd = u32::from(r0) - u32::from(r1);
70 let gd = u32::from(g0) - u32::from(g1);
71 let bd = u32::from(b0) - u32::from(b1);
72 if rd > gd && rd > bd {
73 0
74 } else if bd > rd && bd > gd {
75 2
76 } else {
77 1
78 }
79 }
80}
81
82struct Pixel16Sum {
83 rsum: u64,
84 gsum: u64,
85 bsum: u64,
86 count: u64,
87}
88
89impl VQElementSum<Pixel16> for Pixel16Sum {
90 fn zero() -> Self { Pixel16Sum { rsum: 0, gsum: 0, bsum: 0, count: 0 } }
91 fn add(&mut self, rval: Pixel16, count: u64) {
92 let (r, g, b) = rval.unpack();
93 self.rsum += u64::from(r) * count;
94 self.gsum += u64::from(g) * count;
95 self.bsum += u64::from(b) * count;
96 self.count += count;
97 }
98 fn get_centroid(&self) -> Pixel16 {
99 if self.count != 0 {
100 let r = ((self.rsum + self.count / 2) / self.count) as u8;
101 let g = ((self.gsum + self.count / 2) / self.count) as u8;
102 let b = ((self.bsum + self.count / 2) / self.count) as u8;
103 Pixel16::pack(r, g, b)
104 } else {
105 Pixel16(0x0000)
106 }
107 }
108}
109
110#[derive(Default)]
111struct BlockState {
112 fill_dist: u32,
113 fill_val: Pixel16,
114 clr2_dist: u32,
115 clr2_flags: u16,
116 clr2: [Pixel16; 2],
117 clr8_dist: u32,
118 clr8_flags: u16,
119 clr8: [[Pixel16; 2]; 4],
120}
121
122impl BlockState {
123 fn calc_stats(&mut self, buf: &[Pixel16; 16]) {
124 let num_cw = quantise_median_cut::<Pixel16, Pixel16Sum>(buf, &mut self.clr2);
125 if num_cw == 1 {
126 self.fill_val = Pixel16 { 0: buf[0].0 & !0x400 };
127 } else {
128 let mut avg = Pixel16Sum::zero();
129 for pix in buf.iter() {
130 avg.add(*pix, 1);
131 }
132 self.fill_val = Pixel16 { 0: avg.get_centroid().0 & !0x400 };
133 }
134 self.fill_dist = 0;
135 for pix in buf.iter() {
136 self.fill_dist += pix.dist(self.fill_val);
137 }
138 if self.fill_dist == 0 {
139 self.clr2_dist = std::u32::MAX;
140 self.clr8_dist = std::u32::MAX;
141 return;
142 }
143
144 self.clr2_flags = 0u16;
145 if num_cw == 2 {
146 let mut mask = 1;
147 self.clr2_dist = 0;
148 for pix in buf.iter() {
149 let dist0 = pix.dist(self.clr2[0]);
150 let dist1 = pix.dist(self.clr2[1]);
151 if dist0 < dist1 {
152 self.clr2_flags |= mask;
153 self.clr2_dist += dist0;
154 } else {
155 self.clr2_dist += dist1;
156 }
157 mask <<= 1;
158 }
159 if (self.clr2_flags & 0x8000) != 0 {
160 self.clr2_flags = !self.clr2_flags;
161 self.clr2.swap(0, 1);
162 }
163 } else {
164 self.clr2_dist = self.fill_dist;
165 self.clr2 = [self.fill_val; 2];
166 }
167 if self.clr2_dist == 0 {
168 self.clr8_dist = std::u32::MAX;
169 return;
170 }
171
172 self.clr8 = [[Pixel16 { 0: 0}; 2]; 4];
173 self.clr8_flags = 0;
174 self.clr8_dist = 0;
175 let mut mask = 1;
176 for i in 0..4 {
177 let off = (i & 1) * 2 + (i & 2) * 4;
178 let src2 = [buf[off], buf[off + 1], buf[off + 4], buf[off + 5]];
179 let nc = quantise_median_cut::<Pixel16, Pixel16Sum>(&src2, &mut self.clr8[i]);
180 if nc < 2 {
181 self.clr8[i][1] = self.clr8[i][0];
182 }
183 for j in 0..4 {
184 let dist0 = src2[j].dist(self.clr8[i][0]);
185 let dist1 = src2[j].dist(self.clr8[i][1]);
186 if dist0 < dist1 {
187 self.clr8_flags |= mask;
188 self.clr8_dist += dist0;
189 } else {
190 self.clr8_dist += dist1;
191 }
192 mask <<= 1;
193 }
194 }
195 if (self.clr8_flags & 0x8000) != 0 {
196 self.clr8_flags ^= 0xF000;
197 self.clr8[3].swap(0, 1);
198 }
199 }
200 fn put_fill(&self, dst: &mut [u16], dstride: usize) {
201 for line in dst.chunks_mut(dstride) {
202 for i in 0..4 {
203 line[i] = self.fill_val.0;
204 }
205 }
206 }
207 fn put_clr2(&self, dst: &mut [u16], dstride: usize) {
208 for j in 0..4 {
209 for i in 0..4 {
210 if (self.clr2_flags & (1 << (i + j * 4))) == 0 {
211 dst[i + j * dstride] = self.clr2[0].0;
212 } else {
213 dst[i + j * dstride] = self.clr2[1].0;
214 }
215 }
216 }
217 }
218 fn put_clr8(&self, dst: &mut [u16], dstride: usize) {
219 for i in 0..4 {
220 let off = (i & 1) * 2 + (i & 2) * dstride;
221 let cur_flg = (self.clr8_flags >> (i * 4)) & 0xF;
222 dst[off] = self.clr8[i][( !cur_flg & 1) as usize].0;
223 dst[off + 1] = self.clr8[i][((!cur_flg >> 1) & 1) as usize].0;
224 dst[off + dstride] = self.clr8[i][((!cur_flg >> 2) & 1) as usize].0;
225 dst[off + 1 + dstride] = self.clr8[i][((!cur_flg >> 3) & 1) as usize].0;
226 }
227 }
228 fn write_fill(&self, bw: &mut ByteWriter) -> EncoderResult<()> {
229 bw.write_u16le(self.fill_val.0 | 0x8000)?;
230 Ok(())
231 }
232 fn write_clr2(&self, bw: &mut ByteWriter) -> EncoderResult<()> {
233 bw.write_u16le(self.clr2_flags)?;
234 bw.write_u16le(self.clr2[0].0)?;
235 bw.write_u16le(self.clr2[1].0)?;
236 Ok(())
237 }
238 fn write_clr8(&self, bw: &mut ByteWriter) -> EncoderResult<()> {
239 bw.write_u16le(self.clr8_flags)?;
240 bw.write_u16le(self.clr8[0][0].0 | 0x8000)?;
241 bw.write_u16le(self.clr8[0][1].0)?;
242 bw.write_u16le(self.clr8[1][0].0)?;
243 bw.write_u16le(self.clr8[1][1].0)?;
244 bw.write_u16le(self.clr8[2][0].0)?;
245 bw.write_u16le(self.clr8[2][1].0)?;
246 bw.write_u16le(self.clr8[3][0].0)?;
247 bw.write_u16le(self.clr8[3][1].0)?;
248 Ok(())
249 }
250}
251
252struct MSVideo1Encoder {
253 stream: Option<NAStreamRef>,
254 pkt: Option<NAPacket>,
255 pool: NAVideoBufferPool<u16>,
256 lastfrm: Option<NAVideoBufferRef<u16>>,
257 quality: u8,
258 frmcount: u8,
c9c7f648 259 key_int: u8,
dab59886
KS
260}
261
262impl MSVideo1Encoder {
263 fn new() -> Self {
264 Self {
265 stream: None,
266 pkt: None,
267 pool: NAVideoBufferPool::new(2),
268 lastfrm: None,
269 quality: 0,
270 frmcount: 0,
c9c7f648 271 key_int: 25,
dab59886
KS
272 }
273 }
274 fn get_block(src: &[u16], sstride: usize, buf: &mut [Pixel16; 16]) {
275 for (line, dst) in src.chunks(sstride).zip(buf.chunks_mut(4)) {
276 for i in 0..4 {
277 dst[i] = Pixel16 { 0: line[i] };
278 }
279 }
280 }
281 fn write_skips(bw: &mut ByteWriter, skips: usize) -> EncoderResult<()> {
282 bw.write_u16le((skips as u16) | 0x8400)?;
283 Ok(())
284 }
285 fn encode_inter(bw: &mut ByteWriter, cur_frm: &mut NAVideoBuffer<u16>, in_frm: &NAVideoBuffer<u16>, prev_frm: &NAVideoBuffer<u16>, _quality: u8) -> EncoderResult<bool> {
286 let mut is_intra = true;
287 let src = in_frm.get_data();
288 let sstride = in_frm.get_stride(0);
289 let soff = in_frm.get_offset(0);
290 let (w, h) = in_frm.get_dimensions(0);
291 let rsrc = prev_frm.get_data();
292 let rstride = prev_frm.get_stride(0);
293 let roff = prev_frm.get_offset(0);
294 let dstride = cur_frm.get_stride(0);
295 let doff = cur_frm.get_offset(0);
296 let dst = cur_frm.get_data_mut().unwrap();
297 let mut skip_run = 0;
298 for ((sstrip, rstrip), dstrip) in (&src[soff..]).chunks(sstride * 4).take(h / 4).zip((&rsrc[roff..]).chunks(rstride * 4)).zip((&mut dst[doff..]).chunks_mut(dstride * 4)) {
299 for x in (0..w).step_by(4) {
300 let mut buf = [Pixel16::min_cw(); 16];
301 let mut refbuf = [Pixel16::min_cw(); 16];
302 Self::get_block(&sstrip[x..], sstride, &mut buf);
303 Self::get_block(&rstrip[x..], rstride, &mut refbuf);
304
305 let mut skip_dist = 0;
306 for (pix, rpix) in buf.iter().zip(refbuf.iter()) {
307 skip_dist += pix.dist(*rpix);
308 }
309 if skip_dist == 0 {
310 skip_run += 1;
311 is_intra = false;
312 if skip_run == 1023 {
313 Self::write_skips(bw, skip_run)?;
314 skip_run = 0;
315 }
316 continue;
317 }
318
319 let mut bstate = BlockState::default();
320 bstate.calc_stats(&buf);
321
322 let dst = &mut dstrip[x..];
323 if skip_dist <= bstate.fill_dist {
324 skip_run += 1;
325 is_intra = false;
326 if skip_run == 1023 {
327 Self::write_skips(bw, skip_run)?;
328 skip_run = 0;
329 }
330 } else if bstate.fill_dist <= bstate.clr2_dist {
331 bstate.put_fill(dst, dstride);
332 if skip_run != 0 {
333 Self::write_skips(bw, skip_run)?;
334 skip_run = 0;
335 }
336 bstate.write_fill(bw)?;
337 } else if bstate.clr8_dist < bstate.clr2_dist {
338 bstate.put_clr8(dst, dstride);
339 if skip_run != 0 {
340 Self::write_skips(bw, skip_run)?;
341 skip_run = 0;
342 }
343 bstate.write_clr8(bw)?;
344 } else {
345 bstate.put_clr2(dst, dstride);
346 if skip_run != 0 {
347 Self::write_skips(bw, skip_run)?;
348 skip_run = 0;
349 }
350 bstate.write_clr2(bw)?;
351 }
352 }
353 }
354 if skip_run != 0 {
355 Self::write_skips(bw, skip_run)?;
356 }
357 if is_intra {
358 bw.write_u16le(0)?;
359 } //xxx: something for inter?
360 Ok(is_intra)
361 }
362 fn encode_intra(bw: &mut ByteWriter, cur_frm: &mut NAVideoBuffer<u16>, in_frm: &NAVideoBuffer<u16>, _quality: u8) -> EncoderResult<bool> {
363 let src = in_frm.get_data();
364 let sstride = in_frm.get_stride(0);
365 let soff = in_frm.get_offset(0);
366 let (w, h) = in_frm.get_dimensions(0);
367 let dstride = cur_frm.get_stride(0);
368 let doff = cur_frm.get_offset(0);
369 let dst = cur_frm.get_data_mut().unwrap();
370 for (sstrip, dstrip) in (&src[soff..]).chunks(sstride * 4).take(h / 4).zip((&mut dst[doff..]).chunks_mut(dstride * 4)) {
371 for x in (0..w).step_by(4) {
372 let mut buf = [Pixel16::min_cw(); 16];
373 Self::get_block(&sstrip[x..], sstride, &mut buf);
374 let mut bstate = BlockState::default();
375 bstate.calc_stats(&buf);
376
377 let dst = &mut dstrip[x..];
378 if bstate.fill_dist <= bstate.clr2_dist {
379 bstate.put_fill(dst, dstride);
380 bstate.write_fill(bw)?;
381 } else if bstate.clr8_dist < bstate.clr2_dist {
382 bstate.put_clr8(dst, dstride);
383 bstate.write_clr8(bw)?;
384 } else {
385 bstate.put_clr2(dst, dstride);
386 bstate.write_clr2(bw)?;
387 }
388 }
389 }
390 bw.write_u16le(0)?;
391 Ok(true)
392 }
393}
394
395const RGB555_FORMAT: NAPixelFormaton = NAPixelFormaton {
396 model: ColorModel::RGB(RGBSubmodel::RGB), components: 3,
397 comp_info: [
398 Some(NAPixelChromaton{ h_ss: 0, v_ss: 0, packed: true, depth: 5, shift: 10, comp_offs: 0, next_elem: 2 }),
399 Some(NAPixelChromaton{ h_ss: 0, v_ss: 0, packed: true, depth: 5, shift: 5, comp_offs: 0, next_elem: 2 }),
400 Some(NAPixelChromaton{ h_ss: 0, v_ss: 0, packed: true, depth: 5, shift: 0, comp_offs: 0, next_elem: 2 }),
401 None, None],
402 elem_size: 2, be: false, alpha: false, palette: false };
403
404impl NAEncoder for MSVideo1Encoder {
405 fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
406 match encinfo.format {
407 NACodecTypeInfo::None => {
408 let mut ofmt = EncodeParameters::default();
409 ofmt.format = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, true, RGB555_FORMAT));
410 Ok(ofmt)
411 },
4efceb69 412 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
dab59886
KS
413 NACodecTypeInfo::Video(vinfo) => {
414 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, true, RGB555_FORMAT);
d722ffe9 415 let mut ofmt = *encinfo;
dab59886
KS
416 ofmt.format = NACodecTypeInfo::Video(outinfo);
417 Ok(ofmt)
418 }
419 }
420 }
421 fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
422 match encinfo.format {
423 NACodecTypeInfo::None => Err(EncoderError::FormatError),
424 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
425 NACodecTypeInfo::Video(vinfo) => {
426 if vinfo.format != RGB555_FORMAT {
427 return Err(EncoderError::FormatError);
428 }
429 if ((vinfo.width | vinfo.height) & 3) != 0 {
430 return Err(EncoderError::FormatError);
431 }
432
433 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, true, RGB555_FORMAT);
4efceb69 434 let info = NACodecInfo::new("msvideo1", NACodecTypeInfo::Video(out_info), None);
2ff56201
KS
435 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den);
436 stream.set_num(stream_id as usize);
437 let stream = stream.into_ref();
4efceb69 438 if self.pool.prealloc_video(out_info, 2).is_err() {
dab59886
KS
439 return Err(EncoderError::AllocError);
440 }
441
442 self.stream = Some(stream.clone());
443 self.quality = encinfo.quality;
c8db9313 444
dab59886
KS
445 Ok(stream)
446 },
447 }
448 }
449 fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
450 let buf = frm.get_buffer();
451 if let Some(ref vbuf) = buf.get_vbuf16() {
452 let mut cur_frm = self.pool.get_free().unwrap();
453 let mut dbuf = Vec::with_capacity(4);
454 let mut gw = GrowableMemoryWriter::new_write(&mut dbuf);
455 let mut bw = ByteWriter::new(&mut gw);
456 if self.frmcount == 0 {
457 self.lastfrm = None;
458 }
459 let is_intra = if let Some(ref prev_buf) = self.lastfrm {
460 Self::encode_inter(&mut bw, &mut cur_frm, vbuf, prev_buf, self.quality)?
461 } else {
462 Self::encode_intra(&mut bw, &mut cur_frm, vbuf, self.quality)?
463 };
464 self.lastfrm = Some(cur_frm);
465 self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
466 self.frmcount += 1;
c9c7f648 467 if self.frmcount == self.key_int {
dab59886
KS
468 self.frmcount = 0;
469 }
470 Ok(())
471 } else {
472 Err(EncoderError::InvalidParameters)
473 }
474 }
475 fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
476 let mut npkt = None;
477 std::mem::swap(&mut self.pkt, &mut npkt);
478 Ok(npkt)
479 }
480 fn flush(&mut self) -> EncoderResult<()> {
481 self.frmcount = 0;
482 Ok(())
483 }
484}
485
c9c7f648
KS
486const ENCODER_OPTS: &[NAOptionDefinition] = &[
487 NAOptionDefinition {
488 name: KEYFRAME_OPTION, description: KEYFRAME_OPTION_DESC,
489 opt_type: NAOptionDefinitionType::Int(Some(0), Some(128)) },
490];
491
20766f15 492impl NAOptionHandler for MSVideo1Encoder {
c9c7f648
KS
493 fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
494 fn set_options(&mut self, options: &[NAOption]) {
495 for option in options.iter() {
496 for opt_def in ENCODER_OPTS.iter() {
497 if opt_def.check(option).is_ok() {
498 match option.name {
499 KEYFRAME_OPTION => {
500 if let NAValue::Int(intval) = option.value {
501 self.key_int = intval as u8;
502 }
503 },
504 _ => {},
505 };
506 }
507 }
508 }
509 }
510 fn query_option_value(&self, name: &str) -> Option<NAValue> {
511 match name {
512 KEYFRAME_OPTION => Some(NAValue::Int(i64::from(self.key_int))),
513 _ => None,
514 }
515 }
20766f15
KS
516}
517
dab59886
KS
518pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
519 Box::new(MSVideo1Encoder::new())
520}
521
522#[cfg(test)]
523mod test {
524 use nihav_core::codecs::*;
525 use nihav_core::demuxers::*;
526 use nihav_core::muxers::*;
527 use crate::*;
528 use nihav_commonfmt::*;
529 use nihav_codec_support::test::enc_video::*;
530 use super::RGB555_FORMAT;
531
532 #[test]
533 fn test_ms_video1_encoder() {
534 let mut dmx_reg = RegisteredDemuxers::new();
535 generic_register_all_demuxers(&mut dmx_reg);
536 let mut dec_reg = RegisteredDecoders::new();
78fb6560
KS
537 generic_register_all_decoders(&mut dec_reg);
538 ms_register_all_decoders(&mut dec_reg);
dab59886
KS
539 let mut mux_reg = RegisteredMuxers::new();
540 generic_register_all_muxers(&mut mux_reg);
541 let mut enc_reg = RegisteredEncoders::new();
542 ms_register_all_encoders(&mut enc_reg);
543
544 let dec_config = DecoderTestParams {
545 demuxer: "avi",
546 in_name: "assets/Misc/TalkingHead_352x288.avi",
547 stream_type: StreamType::Video,
c75dce78 548 limit: Some(3),
dab59886
KS
549 dmx_reg, dec_reg,
550 };
551 let enc_config = EncoderTestParams {
552 muxer: "avi",
553 enc_name: "msvideo1",
554 out_name: "msvideo1.avi",
555 mux_reg, enc_reg,
556 };
557 let dst_vinfo = NAVideoInfo {
558 width: 0,
559 height: 0,
560 format: RGB555_FORMAT,
561 flipped: true,
30940e74 562 bits: 16,
dab59886
KS
563 };
564 let enc_params = EncodeParameters {
565 format: NACodecTypeInfo::Video(dst_vinfo),
566 quality: 0,
567 bitrate: 0,
568 tb_num: 0,
569 tb_den: 0,
570 flags: 0,
571 };
c75dce78
KS
572 //test_encoding_to_file(&dec_config, &enc_config, enc_params);
573 test_encoding_md5(&dec_config, &enc_config, enc_params,
574 &[0x8f0d3f66, 0xb91a6ee5, 0x59e4933e, 0x59c6fb0b]);
dab59886
KS
575 }
576}