nihav-core/src/test/dec_video.rs

   1 use std::fs::File;
   2 use std::io::prelude::*;
   3 use crate::frame::*;
   4 use crate::codecs::*;
   5 use crate::demuxers::*;
   6 //use crate::io::byteio::*;
   7 use crate::scale::*;
   8 use super::wavwriter::WavWriter;
   9 use super::md5::MD5;
  10 pub use super::ExpectedTestResult;
  11
  12 const OUTPUT_PREFIX: &str = "assets/test_out";
  13
  14 fn write_pgmyuv(pfx: &str, strno: usize, num: u64, frm: NAFrameRef) {
  15     if let NABufferType::None = frm.get_buffer() { return; }
  16     let name = format!("{}/{}out{:02}_{:06}.pgm", OUTPUT_PREFIX, pfx, strno, num);
  17     let mut ofile = File::create(name).unwrap();
  18     let buf = frm.get_buffer().get_vbuf().unwrap();
  19     let (w, h) = buf.get_dimensions(0);
  20     let (w2, h2) = buf.get_dimensions(1);
  21     let has_alpha = buf.get_info().get_format().has_alpha();
  22     let mut tot_h = h + h2;
  23     if has_alpha {
  24         tot_h += h;
  25     }
  26     if w2 > w/2 {
  27         tot_h += h2;
  28     }
  29     let hdr = format!("P5\n{} {}\n255\n", w, tot_h);
  30     ofile.write_all(hdr.as_bytes()).unwrap();
  31     let dta = buf.get_data();
  32     let ls = buf.get_stride(0);
  33     let mut idx = 0;
  34     let mut idx2 = w;
  35     let is_flipped = buf.get_info().is_flipped();
  36     if is_flipped {
  37         idx  += h * ls;
  38         idx2 += h * ls;
  39     }
  40     for _ in 0..h {
  41         if is_flipped {
  42             idx  -= ls;
  43             idx2 -= ls;
  44         }
  45         let line = &dta[idx..idx2];
  46         ofile.write_all(line).unwrap();
  47         if !is_flipped {
  48             idx  += ls;
  49             idx2 += ls;
  50         }
  51     }
  52     if w2 <= w/2 {
  53         let pad: Vec<u8> = vec![0xFF; (w - w2 * 2) / 2];
  54         let mut base1 = buf.get_offset(1);
  55         let stride1 = buf.get_stride(1);
  56         let mut base2 = buf.get_offset(2);
  57         let stride2 = buf.get_stride(2);
  58         if is_flipped {
  59             base1 += h2 * stride1;
  60             base2 += h2 * stride2;
  61         }
  62         for _ in 0..h2 {
  63             if is_flipped {
  64                 base1 -= stride1;
  65                 base2 -= stride2;
  66             }
  67             let bend1 = base1 + w2;
  68             let line = &dta[base1..bend1];
  69             ofile.write_all(line).unwrap();
  70             ofile.write_all(pad.as_slice()).unwrap();
  71
  72             let bend2 = base2 + w2;
  73             let line = &dta[base2..bend2];
  74             ofile.write_all(line).unwrap();
  75             ofile.write_all(pad.as_slice()).unwrap();
  76
  77             if !is_flipped {
  78                 base1 += stride1;
  79                 base2 += stride2;
  80             }
  81         }
  82     } else {
  83         let pad: Vec<u8> = vec![0xFF; w - w2];
  84         let mut base1 = buf.get_offset(1);
  85         let stride1 = buf.get_stride(1);
  86         if is_flipped {
  87             base1 += h2 * stride1;
  88         }
  89         for _ in 0..h2 {
  90             if is_flipped {
  91                 base1 -= stride1;
  92             }
  93             let bend1 = base1 + w2;
  94             let line = &dta[base1..bend1];
  95             ofile.write_all(line).unwrap();
  96             ofile.write_all(pad.as_slice()).unwrap();
  97             if !is_flipped {
  98                 base1 += stride1;
  99             }
 100         }
 101         let mut base2 = buf.get_offset(2);
 102         let stride2 = buf.get_stride(2);
 103         if is_flipped {
 104             base2 += h2 * stride2;
 105         }
 106         for _ in 0..h2 {
 107             if is_flipped {
 108                 base2 -= stride2;
 109             }
 110             let bend2 = base2 + w2;
 111             let line = &dta[base2..bend2];
 112             ofile.write_all(line).unwrap();
 113             ofile.write_all(pad.as_slice()).unwrap();
 114             if !is_flipped {
 115                 base2 += stride2;
 116             }
 117         }
 118     }
 119     if has_alpha {
 120         let ls = buf.get_stride(3);
 121         let mut idx = buf.get_offset(3);
 122         let mut idx2 = idx + w;
 123         if is_flipped {
 124             idx  += h * ls;
 125             idx2 += h * ls;
 126         }
 127         for _ in 0..h {
 128             if is_flipped {
 129                 idx  -= ls;
 130                 idx2 -= ls;
 131             }
 132             let line = &dta[idx..idx2];
 133             ofile.write_all(line).unwrap();
 134             if !is_flipped {
 135                 idx  += ls;
 136                 idx2 += ls;
 137             }
 138         }
 139     }
 140 }
 141
 142 fn write_palppm(pfx: &str, strno: usize, num: u64, frm: NAFrameRef) {
 143     let name = format!("{}/{}out{:02}_{:06}.ppm", OUTPUT_PREFIX, pfx, strno, num);
 144     let mut ofile = File::create(name).unwrap();
 145     let buf = frm.get_buffer().get_vbuf().unwrap();
 146     let (w, h) = buf.get_dimensions(0);
 147     let paloff = buf.get_offset(1);
 148     let hdr = format!("P6\n{} {}\n255\n", w, h);
 149     ofile.write_all(hdr.as_bytes()).unwrap();
 150     let dta = buf.get_data();
 151     let ls = buf.get_stride(0);
 152     let offs: [usize; 3] = [
 153             buf.get_info().get_format().get_chromaton(0).unwrap().get_offset() as usize,
 154             buf.get_info().get_format().get_chromaton(1).unwrap().get_offset() as usize,
 155             buf.get_info().get_format().get_chromaton(2).unwrap().get_offset() as usize
 156         ];
 157     let mut idx  = 0;
 158     let mut line: Vec<u8> = vec![0; w * 3];
 159     for _ in 0..h {
 160         let src = &dta[idx..(idx+w)];
 161         for x in 0..w {
 162             let pix = src[x] as usize;
 163             line[x * 3 + 0] = dta[paloff + pix * 3 + offs[0]];
 164             line[x * 3 + 1] = dta[paloff + pix * 3 + offs[1]];
 165             line[x * 3 + 2] = dta[paloff + pix * 3 + offs[2]];
 166         }
 167         ofile.write_all(line.as_slice()).unwrap();
 168         idx  += ls;
 169     }
 170 }
 171
 172 fn write_ppm(pfx: &str, strno: usize, num: u64, frm: NAFrameRef) {
 173     let name = format!("{}/{}out{:02}_{:06}.ppm", OUTPUT_PREFIX, pfx, strno, num);
 174     let mut ofile = File::create(name).unwrap();
 175         let info = frm.get_buffer().get_video_info().unwrap();
 176         let mut dpic = alloc_video_buffer(NAVideoInfo::new(info.get_width(), info.get_height(), false, RGB24_FORMAT), 0).unwrap();
 177         let ifmt = ScaleInfo { width: info.get_width(), height: info.get_height(), fmt: info.get_format() };
 178         let ofmt = ScaleInfo { width: info.get_width(), height: info.get_height(), fmt: RGB24_FORMAT };
 179         let mut scaler = NAScale::new(ifmt, ofmt).unwrap();
 180         scaler.convert(&frm.get_buffer(), &mut dpic).unwrap();
 181         let buf = dpic.get_vbuf().unwrap();
 182         let (w, h) = buf.get_dimensions(0);
 183         let hdr = format!("P6\n{} {}\n255\n", w, h);
 184         ofile.write_all(hdr.as_bytes()).unwrap();
 185         let dta = buf.get_data();
 186         let stride = buf.get_stride(0);
 187         for src in dta.chunks(stride) {
 188             ofile.write_all(&src[0..w*3]).unwrap();
 189         }
 190 }
 191
 192 /*fn open_wav_out(pfx: &str, strno: usize) -> WavWriter {
 193     let name = format!("assets/{}out{:02}.wav", pfx, strno);
 194     let mut file = File::create(name).unwrap();
 195     let mut fw = FileWriter::new_write(&mut file);
 196     let mut wr = ByteWriter::new(&mut fw);
 197     WavWriter::new(&mut wr)
 198 }*/
 199
 200 pub fn test_file_decoding(demuxer: &str, name: &str, limit: Option<u64>,
 201                           decode_video: bool, decode_audio: bool,
 202                           video_pfx: Option<&str>,
 203                           dmx_reg: &RegisteredDemuxers, dec_reg: &RegisteredDecoders) {
 204     let dmx_f = dmx_reg.find_demuxer(demuxer).unwrap();
 205     let mut file = File::open(name).unwrap();
 206     let mut fr = FileReader::new_read(&mut file);
 207     let mut br = ByteReader::new(&mut fr);
 208     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 209
 210     let mut decs: Vec<Option<(Box<NADecoderSupport>, Box<dyn NADecoder>)>> = Vec::new();
 211     for i in 0..dmx.get_num_streams() {
 212         let s = dmx.get_stream(i).unwrap();
 213         let info = s.get_info();
 214         let decfunc = dec_reg.find_decoder(info.get_name());
 215         if let Some(df) = decfunc {
 216             if (decode_video && info.is_video()) || (decode_audio && info.is_audio()) {
 217                 let mut dec = (df)();
 218                 let mut dsupp = Box::new(NADecoderSupport::new());
 219                 dec.init(&mut dsupp, info).unwrap();
 220                 decs.push(Some((dsupp, dec)));
 221             } else {
 222                 decs.push(None);
 223             }
 224         } else {
 225             decs.push(None);
 226         }
 227     }
 228
 229     loop {
 230         let pktres = dmx.get_frame();
 231         if let Err(e) = pktres {
 232             if e == DemuxerError::EOF { break; }
 233             panic!("error");
 234         }
 235         let pkt = pktres.unwrap();
 236         let streamno = pkt.get_stream().get_id() as usize;
 237         if let Some((ref mut dsupp, ref mut dec)) = decs[streamno] {
 238             if let (Some(lim), Some(ppts)) = (limit, pkt.get_pts()) {
 239                 if ppts > lim { break; }
 240             }
 241             let frm = dec.decode(dsupp, &pkt).unwrap();
 242             if pkt.get_stream().get_info().is_video() && video_pfx.is_some() && frm.get_frame_type() != FrameType::Skip {
 243                 let pfx = video_pfx.unwrap();
 244                 let pts = if let Some(fpts) = frm.get_pts() { fpts } else { pkt.get_pts().unwrap() };
 245                 let vinfo = frm.get_buffer().get_video_info().unwrap();
 246                 if vinfo.get_format().is_paletted() {
 247                     write_palppm(pfx, streamno, pts, frm);
 248                 } else if vinfo.get_format().get_model().is_yuv() {
 249                     write_pgmyuv(pfx, streamno, pts, frm);
 250                 } else if vinfo.get_format().get_model().is_rgb() {
 251                     write_ppm(pfx, streamno, pts, frm);
 252                 } else {
 253 panic!(" unknown format");
 254                 }
 255             }
 256         }
 257     }
 258 }
 259
 260 pub fn test_decode_audio(demuxer: &str, name: &str, limit: Option<u64>, audio_pfx: Option<&str>,
 261                          dmx_reg: &RegisteredDemuxers, dec_reg: &RegisteredDecoders) {
 262     let dmx_f = dmx_reg.find_demuxer(demuxer).unwrap();
 263     let mut file = File::open(name).unwrap();
 264     let mut fr = FileReader::new_read(&mut file);
 265     let mut br = ByteReader::new(&mut fr);
 266     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 267
 268     let mut decs: Vec<Option<(Box<NADecoderSupport>, Box<dyn NADecoder>)>> = Vec::new();
 269     for i in 0..dmx.get_num_streams() {
 270         let s = dmx.get_stream(i).unwrap();
 271         let info = s.get_info();
 272         let decfunc = dec_reg.find_decoder(info.get_name());
 273         if let Some(df) = decfunc {
 274             if info.is_audio() {
 275                 let mut dec = (df)();
 276                 let mut dsupp = Box::new(NADecoderSupport::new());
 277                 dec.init(&mut dsupp, info).unwrap();
 278                 decs.push(Some((dsupp, dec)));
 279             } else {
 280                 decs.push(None);
 281             }
 282         } else {
 283             decs.push(None);
 284         }
 285     }
 286
 287     if let Some(audio_pfx) = audio_pfx {
 288         let name = format!("{}/{}out.wav", OUTPUT_PREFIX, audio_pfx);
 289         let file = File::create(name).unwrap();
 290         let mut fw = FileWriter::new_write(file);
 291         let mut wr = ByteWriter::new(&mut fw);
 292         let mut wwr = WavWriter::new(&mut wr);
 293         let mut wrote_header = false;
 294
 295         loop {
 296             let pktres = dmx.get_frame();
 297             if let Err(e) = pktres {
 298                 if e == DemuxerError::EOF { break; }
 299                 panic!("error");
 300             }
 301             let pkt = pktres.unwrap();
 302             if limit.is_some() && pkt.get_pts().is_some() && pkt.get_pts().unwrap() > limit.unwrap() {
 303                 break;
 304             }
 305             let streamno = pkt.get_stream().get_id() as usize;
 306             if let Some((ref mut dsupp, ref mut dec)) = decs[streamno] {
 307                 let frm = dec.decode(dsupp, &pkt).unwrap();
 308                 if frm.get_info().is_audio() {
 309                     if !wrote_header {
 310                         wwr.write_header(frm.get_info().as_ref().get_properties().get_audio_info().unwrap()).unwrap();
 311                         wrote_header = true;
 312                     }
 313                     wwr.write_frame(frm.get_buffer()).unwrap();
 314                 }
 315             }
 316         }
 317     } else {
 318         loop {
 319             let pktres = dmx.get_frame();
 320             if let Err(e) = pktres {
 321                 if e == DemuxerError::EOF { break; }
 322                 panic!("error");
 323             }
 324             let pkt = pktres.unwrap();
 325             if limit.is_some() && pkt.get_pts().is_some() && pkt.get_pts().unwrap() > limit.unwrap() {
 326                 break;
 327             }
 328             let streamno = pkt.get_stream().get_id() as usize;
 329             if let Some((ref mut dsupp, ref mut dec)) = decs[streamno] {
 330                 let _ = dec.decode(dsupp, &pkt).unwrap();
 331             }
 332         }
 333     }
 334 }
 335
 336 fn frame_checksum(md5: &mut MD5, frm: NAFrameRef) {
 337     match frm.get_buffer() {
 338         NABufferType::Video(ref vb) => {
 339             md5.update_hash(vb.get_data());
 340         },
 341         NABufferType::Video16(ref vb) => {
 342             let mut samp = [0u8; 2];
 343             let data = vb.get_data();
 344             for el in data.iter() {
 345                 samp[0] = (*el >> 8) as u8;
 346                 samp[1] = (*el >> 0) as u8;
 347                 md5.update_hash(&samp);
 348             }
 349         },
 350         NABufferType::Video32(ref vb) => {
 351             let mut samp = [0u8; 4];
 352             let data = vb.get_data();
 353             for el in data.iter() {
 354                 samp[0] = (*el >> 24) as u8;
 355                 samp[1] = (*el >> 16) as u8;
 356                 samp[2] = (*el >>  8) as u8;
 357                 samp[3] = (*el >>  0) as u8;
 358                 md5.update_hash(&samp);
 359             }
 360         },
 361         NABufferType::VideoPacked(ref vb) => {
 362             md5.update_hash(vb.get_data());
 363         },
 364         NABufferType::AudioU8(ref ab) => {
 365             md5.update_hash(ab.get_data());
 366         },
 367         NABufferType::AudioI16(ref ab) => {
 368             let mut samp = [0u8; 2];
 369             let data = ab.get_data();
 370             for el in data.iter() {
 371                 samp[0] = (*el >> 8) as u8;
 372                 samp[1] = (*el >> 0) as u8;
 373                 md5.update_hash(&samp);
 374             }
 375         },
 376         NABufferType::AudioI32(ref ab) => {
 377             let mut samp = [0u8; 4];
 378             let data = ab.get_data();
 379             for el in data.iter() {
 380                 samp[0] = (*el >> 24) as u8;
 381                 samp[1] = (*el >> 16) as u8;
 382                 samp[2] = (*el >>  8) as u8;
 383                 samp[3] = (*el >>  0) as u8;
 384                 md5.update_hash(&samp);
 385             }
 386         },
 387         NABufferType::AudioF32(ref ab) => {
 388             let mut samp = [0u8; 4];
 389             let data = ab.get_data();
 390             for el in data.iter() {
 391                 let bits = el.to_bits();
 392                 samp[0] = (bits >> 24) as u8;
 393                 samp[1] = (bits >> 16) as u8;
 394                 samp[2] = (bits >>  8) as u8;
 395                 samp[3] = (bits >>  0) as u8;
 396                 md5.update_hash(&samp);
 397             }
 398         },
 399         NABufferType::AudioPacked(ref ab) => {
 400             md5.update_hash(ab.get_data());
 401         },
 402         NABufferType::Data(ref db) => {
 403             md5.update_hash(db.as_ref());
 404         },
 405         NABufferType::None => {},
 406     };
 407 }
 408
 409 pub fn test_decoding(demuxer: &str, dec_name: &str, filename: &str, limit: Option<u64>,
 410                      dmx_reg: &RegisteredDemuxers, dec_reg: &RegisteredDecoders,
 411                      test: ExpectedTestResult) {
 412     let dmx_f = dmx_reg.find_demuxer(demuxer).unwrap();
 413     let mut file = File::open(filename).unwrap();
 414     let mut fr = FileReader::new_read(&mut file);
 415     let mut br = ByteReader::new(&mut fr);
 416     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 417
 418     let mut decs: Vec<Option<(Box<NADecoderSupport>, Box<dyn NADecoder>)>> = Vec::new();
 419     let mut found = false;
 420     for i in 0..dmx.get_num_streams() {
 421         let s = dmx.get_stream(i).unwrap();
 422         let info = s.get_info();
 423 println!("stream {} codec {} / {}", i, info.get_name(), dec_name);
 424         if !found && (info.get_name() == dec_name) {
 425             let decfunc = dec_reg.find_decoder(info.get_name());
 426             if let Some(df) = decfunc {
 427                 let mut dec = (df)();
 428                 let mut dsupp = Box::new(NADecoderSupport::new());
 429                 dec.init(&mut dsupp, info).unwrap();
 430                 decs.push(Some((dsupp, dec)));
 431                 found = true;
 432             } else {
 433                 decs.push(None);
 434             }
 435         } else {
 436             decs.push(None);
 437         }
 438     }
 439
 440     let mut md5 = MD5::new();
 441     let mut frameiter = if let ExpectedTestResult::MD5Frames(ref vec) = test {
 442             Some(vec.iter())
 443         } else {
 444             None
 445         };
 446     loop {
 447         let pktres = dmx.get_frame();
 448         if let Err(e) = pktres {
 449             if e == DemuxerError::EOF { break; }
 450             panic!("error");
 451         }
 452         let pkt = pktres.unwrap();
 453         let streamno = pkt.get_stream().get_id() as usize;
 454         if let Some((ref mut dsupp, ref mut dec)) = decs[streamno] {
 455             if limit.is_some() && pkt.get_pts().is_some() && pkt.get_pts().unwrap() > limit.unwrap() {
 456                 break;
 457             }
 458             let frm = dec.decode(dsupp, &pkt).unwrap();
 459             match &test {
 460                 ExpectedTestResult::Decodes => {},
 461                 ExpectedTestResult::MD5(_) => { frame_checksum(&mut md5, frm); },
 462                 ExpectedTestResult::MD5Frames(_) => {
 463                     md5 = MD5::new();
 464                     frame_checksum(&mut md5, frm);
 465                     md5.finish();
 466                     if let Some(ref mut iter) = frameiter {
 467                         let ret = iter.next();
 468                         if ret.is_none() { break; }
 469                         let ref_hash = ret.unwrap();
 470                         let mut hash = [0u32; 4];
 471                         md5.get_hash(&mut hash);
 472 println!("frame pts {:?} hash {}", pkt.get_pts(), md5);
 473                         assert_eq!(&hash, ref_hash);
 474                     }
 475                 },
 476                 ExpectedTestResult::GenerateMD5Frames => {
 477                     md5 = MD5::new();
 478                     frame_checksum(&mut md5, frm);
 479                     md5.finish();
 480                     let mut hash = [0u32; 4];
 481                     md5.get_hash(&mut hash);
 482 println!("frame pts {:?} hash [0x{:08x}, 0x{:08x}, 0x{:08x}, 0x{:08x}],", pkt.get_pts(), hash[0], hash[1], hash[2], hash[3]);
 483                 },
 484             };
 485         }
 486     }
 487     if let ExpectedTestResult::MD5(ref ref_hash) = test {
 488         md5.finish();
 489         let mut hash = [0u32; 4];
 490         md5.get_hash(&mut hash);
 491 println!("full hash {}", md5);
 492         assert_eq!(&hash, ref_hash);
 493     }
 494     if let ExpectedTestResult::GenerateMD5Frames = test {
 495         panic!("generated hashes");
 496     }
 497 }