nihav-codec-support/src/test/enc_video.rs

   1 //! Routines for testing encoders and muxers.
   2 use std::fs::File;
   3 use nihav_core::frame::*;
   4 use nihav_core::codecs::*;
   5 use nihav_core::demuxers::*;
   6 use nihav_core::muxers::*;
   7 use nihav_core::scale::*;
   8 use nihav_core::soundcvt::*;
   9 use super::md5::MD5;
  10
  11 /// Parameters for the source used in the test.
  12 pub struct DecoderTestParams {
  13     /// Demuxer name e.g. `"mov"`.
  14     pub demuxer:        &'static str,
  15     /// Input file name.
  16     pub in_name:        &'static str,
  17     /// Timestamp for last decoded frame.
  18     pub limit:          Option<u64>,
  19     /// Desired input stream type (that will be decoded and fed to the encoder).
  20     pub stream_type:    StreamType,
  21     /// Registered demuxers.
  22     pub dmx_reg:        RegisteredDemuxers,
  23     /// Registered decoders.
  24     pub dec_reg:        RegisteredDecoders,
  25 }
  26
  27 /// Parameters for the encoding test output.
  28 pub struct EncoderTestParams {
  29     /// Muxer name e.g. `"avi"`.
  30     pub muxer:          &'static str,
  31     /// Encoder name.
  32     pub enc_name:       &'static str,
  33     /// Output file name.
  34     pub out_name:       &'static str,
  35     /// Registered muxers.
  36     pub mux_reg:        RegisteredMuxers,
  37     /// Registered encoders.
  38     pub enc_reg:        RegisteredEncoders,
  39 }
  40
  41 /// Tests muxer by making it mux raw streams from the input file.
  42 ///
  43 /// Streams not fitting the output profile (e.g. a video stream or a second audio stream for WAV muxer) will be ignored.
  44 pub fn test_remuxing(dec_config: &DecoderTestParams, enc_config: &EncoderTestParams) {
  45     let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
  46     let mut file = File::open(dec_config.in_name).unwrap();
  47     let mut fr = FileReader::new_read(&mut file);
  48     let mut br = ByteReader::new(&mut fr);
  49     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
  50
  51     let mux_f = enc_config.mux_reg.find_muxer(enc_config.muxer).unwrap();
  52     let out_name = "assets/test_out/".to_owned() + enc_config.out_name;
  53     let file = File::create(&out_name).unwrap();
  54     let mut fw = FileWriter::new_write(file);
  55     let mut bw = ByteWriter::new(&mut fw);
  56     let mut out_sm = StreamManager::new();
  57     let mux_caps = mux_f.get_capabilities();
  58     let mut stream_map: Vec<Option<usize>> = Vec::new();
  59     let mut has_video = false;
  60     let mut has_audio = false;
  61     for stream in dmx.get_streams() {
  62         let mut copy_stream = false;
  63         match mux_caps {
  64             MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::OnlyVideo => {
  65                 copy_stream = stream.get_media_type() == StreamType::Video;
  66                 has_video = true;
  67             },
  68             MuxerCapabilities::SingleAudio(_) | MuxerCapabilities::OnlyAudio => {
  69                 copy_stream = stream.get_media_type() == StreamType::Audio;
  70                 has_audio = true;
  71             },
  72             MuxerCapabilities::SingleVideoAndAudio(_, _) => {
  73                 if stream.get_media_type() == StreamType::Video {
  74                     copy_stream = !has_video;
  75                     has_video = true;
  76                 }
  77                 if stream.get_media_type() == StreamType::Audio {
  78                     copy_stream = !has_audio;
  79                     has_audio = true;
  80                 }
  81             },
  82             MuxerCapabilities::Universal => {
  83                 if stream.get_media_type() == StreamType::Video {
  84                     copy_stream = true;
  85                     has_video = true;
  86                 }
  87                 if stream.get_media_type() == StreamType::Audio {
  88                     copy_stream = true;
  89                     has_audio = true;
  90                 }
  91             }
  92         };
  93         if copy_stream {
  94             let streamno = out_sm.add_stream(NAStream::clone(&stream)).unwrap();
  95             stream_map.push(Some(streamno));
  96             match mux_caps {
  97                 MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::SingleAudio(_) => break,
  98                 _ => {},
  99             };
 100         } else {
 101             stream_map.push(None);
 102         }
 103     }
 104     assert!(out_sm.get_num_streams() > 0);
 105     let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
 106
 107     loop {
 108         let pktres = dmx.get_frame();
 109         if let Err(e) = pktres {
 110             if e == DemuxerError::EOF { break; }
 111             panic!("error");
 112         }
 113         let mut pkt = pktres.unwrap();
 114         println!("Got {}", pkt);
 115         if let Some(new_id) = stream_map[pkt.get_stream().id as usize] {
 116             pkt.reassign(mux.get_stream(new_id).unwrap(), pkt.get_time_information());
 117             mux.mux_frame(pkt).unwrap();
 118         }
 119     }
 120
 121     mux.end().unwrap();
 122 }
 123
 124 /// Tests muxer by making it mux raw streams from the input file and comparing MD5 hash of the result to the provided one.
 125 ///
 126 /// Streams not fitting the output profile (e.g. a video stream or a second audio stream for WAV muxer) will be ignored.
 127 pub fn test_remuxing_md5(dec_config: &DecoderTestParams, muxer: &str, mux_reg: &RegisteredMuxers, md5_hash: [u32; 4]) {
 128     let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
 129     let mut file = File::open(dec_config.in_name).unwrap();
 130     let mut fr = FileReader::new_read(&mut file);
 131     let mut br = ByteReader::new(&mut fr);
 132     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 133
 134     let mux_f = mux_reg.find_muxer(muxer).unwrap();
 135
 136     let mut dst = Vec::with_capacity(1 << 10);
 137     let mut gw = GrowableMemoryWriter::new_write(&mut dst);
 138     let mut bw = ByteWriter::new(&mut gw);
 139     let mut out_sm = StreamManager::new();
 140     let mux_caps = mux_f.get_capabilities();
 141     let mut stream_map: Vec<Option<usize>> = Vec::new();
 142     let mut has_video = false;
 143     let mut has_audio = false;
 144     for stream in dmx.get_streams() {
 145         let mut copy_stream = false;
 146         match mux_caps {
 147             MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::OnlyVideo => {
 148                 copy_stream = stream.get_media_type() == StreamType::Video;
 149                 has_video = true;
 150             },
 151             MuxerCapabilities::SingleAudio(_) | MuxerCapabilities::OnlyAudio => {
 152                 copy_stream = stream.get_media_type() == StreamType::Audio;
 153                 has_audio = true;
 154             },
 155             MuxerCapabilities::SingleVideoAndAudio(_, _) => {
 156                 if stream.get_media_type() == StreamType::Video {
 157                     copy_stream = !has_video;
 158                     has_video = true;
 159                 }
 160                 if stream.get_media_type() == StreamType::Audio {
 161                     copy_stream = !has_audio;
 162                     has_audio = true;
 163                 }
 164             },
 165             MuxerCapabilities::Universal => {
 166                 if stream.get_media_type() == StreamType::Video {
 167                     copy_stream = true;
 168                     has_video = true;
 169                 }
 170                 if stream.get_media_type() == StreamType::Audio {
 171                     copy_stream = true;
 172                     has_audio = true;
 173                 }
 174             }
 175         };
 176         if copy_stream {
 177             let streamno = out_sm.add_stream(NAStream::clone(&stream)).unwrap();
 178             stream_map.push(Some(streamno));
 179             match mux_caps {
 180                 MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::SingleAudio(_) => break,
 181                 _ => {},
 182             };
 183         } else {
 184             stream_map.push(None);
 185         }
 186     }
 187     assert!(out_sm.get_num_streams() > 0);
 188     let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
 189
 190     loop {
 191         let pktres = dmx.get_frame();
 192         if let Err(e) = pktres {
 193             if e == DemuxerError::EOF { break; }
 194             panic!("error");
 195         }
 196         let mut pkt = pktres.unwrap();
 197         println!("Got {}", pkt);
 198         if let Some(new_id) = stream_map[pkt.get_stream().id as usize] {
 199             pkt.reassign(mux.get_stream(new_id).unwrap(), pkt.get_time_information());
 200             mux.mux_frame(pkt).unwrap();
 201         }
 202     }
 203
 204     mux.end().unwrap();
 205
 206     let mut hash = [0; 4];
 207     MD5::calculate_hash(dst.as_slice(), &mut hash);
 208     println!("output hash {:08x}{:08x}{:08x}{:08x}", hash[0], hash[1], hash[2], hash[3]);
 209     assert_eq!(hash, md5_hash);
 210 }
 211
 212 /// Tests an encoder by decoding a stream from input file, feeding it to the encoder and muxing the result into output file.
 213 pub fn test_encoding_to_file(dec_config: &DecoderTestParams, enc_config: &EncoderTestParams, mut enc_params: EncodeParameters, enc_options: &[NAOption]) {
 214     let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
 215     let mut file = File::open(dec_config.in_name).unwrap();
 216     let mut fr = FileReader::new_read(&mut file);
 217     let mut br = ByteReader::new(&mut fr);
 218     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 219
 220     let in_stream = dmx.get_streams().find(|strm| strm.get_media_type() == dec_config.stream_type).unwrap();
 221     let in_stream_id = in_stream.id;
 222     let decfunc = dec_config.dec_reg.find_decoder(in_stream.get_info().get_name()).unwrap();
 223     let mut dec = (decfunc)();
 224     let mut dsupp = Box::new(NADecoderSupport::new());
 225     dec.init(&mut dsupp, in_stream.get_info()).unwrap();
 226
 227     let mut out_sm = StreamManager::new();
 228     enc_params.tb_num = in_stream.tb_num;
 229     enc_params.tb_den = in_stream.tb_den;
 230
 231     if let (NACodecTypeInfo::Video(ref mut vinfo), Some(ref_vinfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_video_info()) {
 232         if vinfo.width == 0 {
 233             vinfo.width  = ref_vinfo.width;
 234             vinfo.height = ref_vinfo.height;
 235         }
 236     }
 237     let mut dst_chmap = NAChannelMap::new();
 238     if let (NACodecTypeInfo::Audio(ref mut ainfo), Some(ref_ainfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_audio_info()) {
 239         if ainfo.sample_rate == 0 {
 240             ainfo.sample_rate = ref_ainfo.sample_rate;
 241         }
 242         if ainfo.channels == 0 {
 243             ainfo.channels = ref_ainfo.channels;
 244         }
 245         match ainfo.channels {
 246             1 => {
 247                 dst_chmap.add_channel(NAChannelType::C);
 248             },
 249             2 => {
 250                 dst_chmap.add_channel(NAChannelType::L);
 251                 dst_chmap.add_channel(NAChannelType::R);
 252             },
 253             _ => panic!("cannot guess channel map"),
 254         }
 255     }
 256
 257     let encfunc = enc_config.enc_reg.find_encoder(enc_config.enc_name).unwrap();
 258     let mut encoder = (encfunc)();
 259     encoder.set_options(enc_options);
 260     let out_str = encoder.init(0, enc_params).unwrap();
 261     out_sm.add_stream(NAStream::clone(&out_str));
 262
 263     let mux_f = enc_config.mux_reg.find_muxer(enc_config.muxer).unwrap();
 264     let out_name = "assets/test_out/".to_owned() + enc_config.out_name;
 265     let file = File::create(&out_name).unwrap();
 266     let mut fw = FileWriter::new_write(file);
 267     let mut bw = ByteWriter::new(&mut fw);
 268     let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
 269
 270     let (mut ifmt, dst_vinfo) = if let NACodecTypeInfo::Video(vinfo) = enc_params.format {
 271             (ScaleInfo { fmt: vinfo.format, width: vinfo.width, height: vinfo.height },
 272              vinfo)
 273         } else {
 274             (ScaleInfo { fmt: YUV420_FORMAT, width: 2, height: 2 },
 275              NAVideoInfo { width: 2, height: 2, format: YUV420_FORMAT, flipped: false, bits: 12 })
 276         };
 277     let ofmt = ifmt;
 278     let mut scaler = NAScale::new(ifmt, ofmt).unwrap();
 279     let mut cvt_buf = alloc_video_buffer(dst_vinfo, 2).unwrap();
 280     loop {
 281         let pktres = dmx.get_frame();
 282         if let Err(e) = pktres {
 283             if e == DemuxerError::EOF { break; }
 284             panic!("decoding error");
 285         }
 286         let pkt = pktres.unwrap();
 287         if pkt.get_stream().id != in_stream_id { continue; }
 288         let frm = dec.decode(&mut dsupp, &pkt).unwrap();
 289         let buf = frm.get_buffer();
 290         let cfrm = if let NACodecTypeInfo::Video(_) = enc_params.format {
 291                 let cur_ifmt = get_scale_fmt_from_pic(&buf);
 292                 if cur_ifmt != ifmt {
 293                     ifmt = cur_ifmt;
 294                     scaler = NAScale::new(ifmt, ofmt).unwrap();
 295                 }
 296                 scaler.convert(&buf, &mut cvt_buf).unwrap();
 297                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf.clone())
 298             } else if let NACodecTypeInfo::Audio(ref dst_ainfo) = enc_params.format {
 299                 let cvt_buf = convert_audio_frame(&buf, dst_ainfo, &dst_chmap).unwrap();
 300                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf)
 301             } else {
 302                 panic!("unexpected format");
 303             };
 304         encoder.encode(&cfrm).unwrap();
 305         while let Ok(Some(pkt)) = encoder.get_packet() {
 306             mux.mux_frame(pkt).unwrap();
 307         }
 308         if let Some(maxts) = dec_config.limit {
 309             if frm.get_pts().unwrap_or(0) >= maxts {
 310                 break;
 311             }
 312         }
 313     }
 314     encoder.flush().unwrap();
 315     while let Ok(Some(pkt)) = encoder.get_packet() {
 316         mux.mux_frame(pkt).unwrap();
 317     }
 318     mux.end().unwrap();
 319 }
 320
 321 /// Tests an encoder by decoding a stream from input file, feeding it to the encoder and calculating the hash of codec information and packet data.
 322 pub fn test_encoding_md5(dec_config: &DecoderTestParams, enc_config: &EncoderTestParams, mut enc_params: EncodeParameters, enc_options: &[NAOption], ref_hash: &[u32; 4]) {
 323     let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
 324     let mut file = File::open(dec_config.in_name).unwrap();
 325     let mut fr = FileReader::new_read(&mut file);
 326     let mut br = ByteReader::new(&mut fr);
 327     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 328
 329     let in_stream = dmx.get_streams().find(|strm| strm.get_media_type() == dec_config.stream_type).unwrap();
 330     let in_stream_id = in_stream.id;
 331     let decfunc = dec_config.dec_reg.find_decoder(in_stream.get_info().get_name()).unwrap();
 332     let mut dec = (decfunc)();
 333     let mut dsupp = Box::new(NADecoderSupport::new());
 334     dec.init(&mut dsupp, in_stream.get_info()).unwrap();
 335
 336     enc_params.tb_num = in_stream.tb_num;
 337     enc_params.tb_den = in_stream.tb_den;
 338
 339     if let (NACodecTypeInfo::Video(ref mut vinfo), Some(ref_vinfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_video_info()) {
 340         if vinfo.width == 0 {
 341             vinfo.width  = ref_vinfo.width;
 342             vinfo.height = ref_vinfo.height;
 343         }
 344     }
 345     let mut dst_chmap = NAChannelMap::new();
 346     if let (NACodecTypeInfo::Audio(ref mut ainfo), Some(ref_ainfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_audio_info()) {
 347         if ainfo.sample_rate == 0 {
 348             ainfo.sample_rate = ref_ainfo.sample_rate;
 349         }
 350         if ainfo.channels == 0 {
 351             ainfo.channels = ref_ainfo.channels;
 352         }
 353         match ainfo.channels {
 354             1 => {
 355                 dst_chmap.add_channel(NAChannelType::C);
 356             },
 357             2 => {
 358                 dst_chmap.add_channel(NAChannelType::L);
 359                 dst_chmap.add_channel(NAChannelType::R);
 360             },
 361             _ => panic!("cannot guess channel map"),
 362         }
 363     }
 364
 365     let encfunc = enc_config.enc_reg.find_encoder(enc_config.enc_name).unwrap();
 366     let mut encoder = (encfunc)();
 367     encoder.set_options(enc_options);
 368     let out_str = encoder.init(0, enc_params).unwrap();
 369
 370     let mut md5 = MD5::new();
 371     let info = out_str.get_info();
 372     md5.update_hash(info.get_name().as_bytes());
 373     match info.get_properties() {
 374         NACodecTypeInfo::Video(ref vinfo) => {
 375             let mut hdr = [0u8; 10];
 376             hdr[0] = (vinfo.width  >> 24) as u8;
 377             hdr[1] = (vinfo.width  >> 16) as u8;
 378             hdr[2] = (vinfo.width  >>  8) as u8;
 379             hdr[3] =  vinfo.width         as u8;
 380             hdr[4] = (vinfo.height >> 24) as u8;
 381             hdr[5] = (vinfo.height >> 16) as u8;
 382             hdr[6] = (vinfo.height >>  8) as u8;
 383             hdr[7] =  vinfo.height        as u8;
 384             hdr[8] = vinfo.flipped as u8;
 385             hdr[9] = vinfo.bits;
 386             md5.update_hash(&hdr);
 387         },
 388         NACodecTypeInfo::Audio(ref ainfo) => {
 389             let mut hdr = [0u8; 10];
 390             hdr[0] = (ainfo.sample_rate >> 24) as u8;
 391             hdr[1] = (ainfo.sample_rate >> 16) as u8;
 392             hdr[2] = (ainfo.sample_rate >>  8) as u8;
 393             hdr[3] =  ainfo.sample_rate        as u8;
 394             hdr[4] = ainfo.channels;
 395             hdr[5] = ainfo.format.bits;
 396             hdr[6] = (ainfo.block_len >> 24) as u8;
 397             hdr[7] = (ainfo.block_len >> 16) as u8;
 398             hdr[8] = (ainfo.block_len >>  8) as u8;
 399             hdr[9] =  ainfo.block_len        as u8;
 400             md5.update_hash(&hdr);
 401         },
 402         _ => {},
 403     };
 404     if let Some(ref buf) = info.get_extradata() {
 405         md5.update_hash(buf.as_slice());
 406     }
 407
 408     let (mut ifmt, dst_vinfo) = if let NACodecTypeInfo::Video(vinfo) = enc_params.format {
 409             (ScaleInfo { fmt: vinfo.format, width: vinfo.width, height: vinfo.height },
 410              vinfo)
 411         } else {
 412             (ScaleInfo { fmt: YUV420_FORMAT, width: 2, height: 2 },
 413              NAVideoInfo { width: 2, height: 2, format: YUV420_FORMAT, flipped: false, bits: 12 })
 414         };
 415     let ofmt = ifmt;
 416     let mut scaler = NAScale::new(ifmt, ofmt).unwrap();
 417     let mut cvt_buf = alloc_video_buffer(dst_vinfo, 2).unwrap();
 418     loop {
 419         let pktres = dmx.get_frame();
 420         if let Err(e) = pktres {
 421             if e == DemuxerError::EOF { break; }
 422             panic!("decoding error");
 423         }
 424         let pkt = pktres.unwrap();
 425         if pkt.get_stream().id != in_stream_id { continue; }
 426         let frm = dec.decode(&mut dsupp, &pkt).unwrap();
 427         let buf = frm.get_buffer();
 428         let cfrm = if let NACodecTypeInfo::Video(_) = enc_params.format {
 429                 let cur_ifmt = get_scale_fmt_from_pic(&buf);
 430                 if cur_ifmt != ifmt {
 431                     ifmt = cur_ifmt;
 432                     scaler = NAScale::new(ifmt, ofmt).unwrap();
 433                 }
 434                 scaler.convert(&buf, &mut cvt_buf).unwrap();
 435                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf.clone())
 436             } else if let NACodecTypeInfo::Audio(ref dst_ainfo) = enc_params.format {
 437                 let cvt_buf = convert_audio_frame(&buf, dst_ainfo, &dst_chmap).unwrap();
 438                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf)
 439             } else {
 440                 panic!("unexpected format");
 441             };
 442         encoder.encode(&cfrm).unwrap();
 443         while let Ok(Some(pkt)) = encoder.get_packet() {
 444             md5.update_hash(pkt.get_buffer().as_slice());
 445         }
 446         if let Some(maxts) = dec_config.limit {
 447             if frm.get_pts().unwrap_or(0) >= maxts {
 448                 break;
 449             }
 450         }
 451     }
 452     encoder.flush().unwrap();
 453     while let Ok(Some(pkt)) = encoder.get_packet() {
 454         md5.update_hash(pkt.get_buffer().as_slice());
 455     }
 456
 457     let mut hash = [0; 4];
 458     md5.finish();
 459     md5.get_hash(&mut hash);
 460     println!("encode hash {}", md5);
 461     assert_eq!(&hash, ref_hash);
 462 }