nihav-ms/src/codecs/msvideo1enc.rs

   1 use nihav_core::codecs::*;
   2 use nihav_core::io::byteio::*;
   3 use nihav_codec_support::vq::*;
   4
   5 #[derive(Default,Clone,Copy,PartialEq)]
   6 struct Pixel16(u16);
   7
   8 impl Pixel16 {
   9     fn unpack(&self) -> (u8, u8, u8) {
  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 }
  16 impl 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
  82 struct Pixel16Sum {
  83     rsum: u64,
  84     gsum: u64,
  85     bsum: u64,
  86     count: u64,
  87 }
  88
  89 impl 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)]
 111 struct 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
 122 impl 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
 252 struct MSVideo1Encoder {
 253     stream:     Option<NAStreamRef>,
 254     pkt:        Option<NAPacket>,
 255     pool:       NAVideoBufferPool<u16>,
 256     lastfrm:    Option<NAVideoBufferRef<u16>>,
 257     quality:    u8,
 258     frmcount:   u8,
 259 }
 260
 261 impl MSVideo1Encoder {
 262     fn new() -> Self {
 263         Self {
 264             stream:     None,
 265             pkt:        None,
 266             pool:       NAVideoBufferPool::new(2),
 267             lastfrm:    None,
 268             quality:    0,
 269             frmcount:   0,
 270         }
 271     }
 272     fn get_block(src: &[u16], sstride: usize, buf: &mut [Pixel16; 16]) {
 273         for (line, dst) in src.chunks(sstride).zip(buf.chunks_mut(4)) {
 274             for i in 0..4 {
 275                 dst[i] = Pixel16 { 0: line[i] };
 276             }
 277         }
 278     }
 279     fn write_skips(bw: &mut ByteWriter, skips: usize) -> EncoderResult<()> {
 280         bw.write_u16le((skips as u16) | 0x8400)?;
 281         Ok(())
 282     }
 283     fn encode_inter(bw: &mut ByteWriter, cur_frm: &mut NAVideoBuffer<u16>, in_frm: &NAVideoBuffer<u16>, prev_frm: &NAVideoBuffer<u16>, _quality: u8) -> EncoderResult<bool> {
 284         let mut is_intra = true;
 285         let src = in_frm.get_data();
 286         let sstride = in_frm.get_stride(0);
 287         let soff = in_frm.get_offset(0);
 288         let (w, h) = in_frm.get_dimensions(0);
 289         let rsrc = prev_frm.get_data();
 290         let rstride = prev_frm.get_stride(0);
 291         let roff = prev_frm.get_offset(0);
 292         let dstride = cur_frm.get_stride(0);
 293         let doff = cur_frm.get_offset(0);
 294         let dst = cur_frm.get_data_mut().unwrap();
 295         let mut skip_run = 0;
 296         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)) {
 297             for x in (0..w).step_by(4) {
 298                 let mut buf = [Pixel16::min_cw(); 16];
 299                 let mut refbuf = [Pixel16::min_cw(); 16];
 300                 Self::get_block(&sstrip[x..], sstride, &mut buf);
 301                 Self::get_block(&rstrip[x..], rstride, &mut refbuf);
 302
 303                 let mut skip_dist = 0;
 304                 for (pix, rpix) in buf.iter().zip(refbuf.iter()) {
 305                     skip_dist += pix.dist(*rpix);
 306                 }
 307                 if skip_dist == 0 {
 308                     skip_run += 1;
 309                     is_intra = false;
 310                     if skip_run == 1023 {
 311                         Self::write_skips(bw, skip_run)?;
 312                         skip_run = 0;
 313                     }
 314                     continue;
 315                 }
 316
 317                 let mut bstate = BlockState::default();
 318                 bstate.calc_stats(&buf);
 319
 320                 let dst = &mut dstrip[x..];
 321                 if skip_dist <= bstate.fill_dist {
 322                     skip_run += 1;
 323                     is_intra = false;
 324                     if skip_run == 1023 {
 325                         Self::write_skips(bw, skip_run)?;
 326                         skip_run = 0;
 327                     }
 328                 } else if bstate.fill_dist <= bstate.clr2_dist {
 329                     bstate.put_fill(dst, dstride);
 330                     if skip_run != 0 {
 331                         Self::write_skips(bw, skip_run)?;
 332                         skip_run = 0;
 333                     }
 334                     bstate.write_fill(bw)?;
 335                 } else if bstate.clr8_dist < bstate.clr2_dist {
 336                     bstate.put_clr8(dst, dstride);
 337                     if skip_run != 0 {
 338                         Self::write_skips(bw, skip_run)?;
 339                         skip_run = 0;
 340                     }
 341                     bstate.write_clr8(bw)?;
 342                 } else {
 343                     bstate.put_clr2(dst, dstride);
 344                     if skip_run != 0 {
 345                         Self::write_skips(bw, skip_run)?;
 346                         skip_run = 0;
 347                     }
 348                     bstate.write_clr2(bw)?;
 349                 }
 350             }
 351         }
 352         if skip_run != 0 {
 353             Self::write_skips(bw, skip_run)?;
 354         }
 355         if is_intra {
 356             bw.write_u16le(0)?;
 357         } //xxx: something for inter?
 358         Ok(is_intra)
 359     }
 360     fn encode_intra(bw: &mut ByteWriter, cur_frm: &mut NAVideoBuffer<u16>, in_frm: &NAVideoBuffer<u16>, _quality: u8) -> EncoderResult<bool> {
 361         let src = in_frm.get_data();
 362         let sstride = in_frm.get_stride(0);
 363         let soff = in_frm.get_offset(0);
 364         let (w, h) = in_frm.get_dimensions(0);
 365         let dstride = cur_frm.get_stride(0);
 366         let doff = cur_frm.get_offset(0);
 367         let dst = cur_frm.get_data_mut().unwrap();
 368         for (sstrip, dstrip) in (&src[soff..]).chunks(sstride * 4).take(h / 4).zip((&mut dst[doff..]).chunks_mut(dstride * 4)) {
 369             for x in (0..w).step_by(4) {
 370                 let mut buf = [Pixel16::min_cw(); 16];
 371                 Self::get_block(&sstrip[x..], sstride, &mut buf);
 372                 let mut bstate = BlockState::default();
 373                 bstate.calc_stats(&buf);
 374
 375                 let dst = &mut dstrip[x..];
 376                 if bstate.fill_dist <= bstate.clr2_dist {
 377                     bstate.put_fill(dst, dstride);
 378                     bstate.write_fill(bw)?;
 379                 } else if bstate.clr8_dist < bstate.clr2_dist {
 380                     bstate.put_clr8(dst, dstride);
 381                     bstate.write_clr8(bw)?;
 382                 } else {
 383                     bstate.put_clr2(dst, dstride);
 384                     bstate.write_clr2(bw)?;
 385                 }
 386             }
 387         }
 388         bw.write_u16le(0)?;
 389         Ok(true)
 390     }
 391 }
 392
 393 const RGB555_FORMAT: NAPixelFormaton = NAPixelFormaton {
 394         model: ColorModel::RGB(RGBSubmodel::RGB), components: 3,
 395         comp_info: [
 396             Some(NAPixelChromaton{ h_ss: 0, v_ss: 0, packed: true, depth: 5, shift: 10, comp_offs: 0, next_elem: 2 }),
 397             Some(NAPixelChromaton{ h_ss: 0, v_ss: 0, packed: true, depth: 5, shift:  5, comp_offs: 0, next_elem: 2 }),
 398             Some(NAPixelChromaton{ h_ss: 0, v_ss: 0, packed: true, depth: 5, shift:  0, comp_offs: 0, next_elem: 2 }),
 399             None, None],
 400         elem_size: 2, be: false, alpha: false, palette: false };
 401
 402 impl NAEncoder for MSVideo1Encoder {
 403     fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
 404         match encinfo.format {
 405             NACodecTypeInfo::None => {
 406                 let mut ofmt = EncodeParameters::default();
 407                 ofmt.format = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, true, RGB555_FORMAT));
 408                 Ok(ofmt)
 409             },
 410             NACodecTypeInfo::Audio(_) => return Err(EncoderError::FormatError),
 411             NACodecTypeInfo::Video(vinfo) => {
 412                 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, true, RGB555_FORMAT);
 413                 let mut ofmt = *encinfo;
 414                 ofmt.format = NACodecTypeInfo::Video(outinfo);
 415                 Ok(ofmt)
 416             }
 417         }
 418     }
 419     fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
 420         match encinfo.format {
 421             NACodecTypeInfo::None => Err(EncoderError::FormatError),
 422             NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
 423             NACodecTypeInfo::Video(vinfo) => {
 424                 if vinfo.format != RGB555_FORMAT {
 425                     return Err(EncoderError::FormatError);
 426                 }
 427                 if ((vinfo.width | vinfo.height) & 3) != 0 {
 428                     return Err(EncoderError::FormatError);
 429                 }
 430
 431                 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, true, RGB555_FORMAT);
 432                 let info = NACodecInfo::new("msvideo1", NACodecTypeInfo::Video(out_info.clone()), None);
 433                 let stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den).into_ref();
 434                 if let Err(_) = self.pool.prealloc_video(out_info, 2) {
 435                     return Err(EncoderError::AllocError);
 436                 }
 437
 438                 self.stream = Some(stream.clone());
 439                 self.quality = encinfo.quality;
 440
 441                 Ok(stream)
 442             },
 443         }
 444     }
 445     fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
 446         let buf = frm.get_buffer();
 447         if let Some(ref vbuf) = buf.get_vbuf16() {
 448             let mut cur_frm = self.pool.get_free().unwrap();
 449             let mut dbuf = Vec::with_capacity(4);
 450             let mut gw   = GrowableMemoryWriter::new_write(&mut dbuf);
 451             let mut bw   = ByteWriter::new(&mut gw);
 452             if self.frmcount == 0 {
 453                 self.lastfrm = None;
 454             }
 455             let is_intra = if let Some(ref prev_buf) = self.lastfrm {
 456                     Self::encode_inter(&mut bw, &mut cur_frm, vbuf, prev_buf, self.quality)?
 457                 } else {
 458                     Self::encode_intra(&mut bw, &mut cur_frm, vbuf, self.quality)?
 459                 };
 460             self.lastfrm = Some(cur_frm);
 461             self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
 462             self.frmcount += 1;
 463             if self.frmcount == 25 {
 464                 self.frmcount = 0;
 465             }
 466             Ok(())
 467         } else {
 468             Err(EncoderError::InvalidParameters)
 469         }
 470     }
 471     fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
 472         let mut npkt = None;
 473         std::mem::swap(&mut self.pkt, &mut npkt);
 474         Ok(npkt)
 475     }
 476     fn flush(&mut self) -> EncoderResult<()> {
 477         self.frmcount = 0;
 478         Ok(())
 479     }
 480 }
 481
 482 impl NAOptionHandler for MSVideo1Encoder {
 483     fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
 484     fn set_options(&mut self, _options: &[NAOption]) { }
 485     fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
 486 }
 487
 488 pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
 489     Box::new(MSVideo1Encoder::new())
 490 }
 491
 492 #[cfg(test)]
 493 mod test {
 494     use nihav_core::codecs::*;
 495     use nihav_core::demuxers::*;
 496     use nihav_core::muxers::*;
 497     use crate::*;
 498     use nihav_commonfmt::*;
 499     use nihav_codec_support::test::enc_video::*;
 500     use super::RGB555_FORMAT;
 501
 502     #[test]
 503     fn test_ms_video1_encoder() {
 504         let mut dmx_reg = RegisteredDemuxers::new();
 505         generic_register_all_demuxers(&mut dmx_reg);
 506         let mut dec_reg = RegisteredDecoders::new();
 507         generic_register_all_codecs(&mut dec_reg);
 508         ms_register_all_codecs(&mut dec_reg);
 509         let mut mux_reg = RegisteredMuxers::new();
 510         generic_register_all_muxers(&mut mux_reg);
 511         let mut enc_reg = RegisteredEncoders::new();
 512         ms_register_all_encoders(&mut enc_reg);
 513
 514         let dec_config = DecoderTestParams {
 515                 demuxer:        "avi",
 516                 in_name:        "assets/Misc/TalkingHead_352x288.avi",
 517                 stream_type:    StreamType::Video,
 518                 limit:          Some(32),
 519                 dmx_reg, dec_reg,
 520             };
 521         let enc_config = EncoderTestParams {
 522                 muxer:          "avi",
 523                 enc_name:       "msvideo1",
 524                 out_name:       "msvideo1.avi",
 525                 mux_reg, enc_reg,
 526             };
 527         let dst_vinfo = NAVideoInfo {
 528                 width:   0,
 529                 height:  0,
 530                 format:  RGB555_FORMAT,
 531                 flipped: true,
 532             };
 533         let enc_params = EncodeParameters {
 534                 format:  NACodecTypeInfo::Video(dst_vinfo),
 535                 quality: 0,
 536                 bitrate: 0,
 537                 tb_num:  0,
 538                 tb_den:  0,
 539                 flags:   0,
 540             };
 541         test_encoding_to_file(&dec_config, &enc_config, enc_params);
 542     }
 543 }