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(1048576);
 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) {
 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(|str| str.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     let out_str = encoder.init(0, enc_params).unwrap();
 260     out_sm.add_stream(NAStream::clone(&out_str));
 261
 262     let mux_f = enc_config.mux_reg.find_muxer(enc_config.muxer).unwrap();
 263     let out_name = "assets/test_out/".to_owned() + enc_config.out_name;
 264     let file = File::create(&out_name).unwrap();
 265     let mut fw = FileWriter::new_write(file);
 266     let mut bw = ByteWriter::new(&mut fw);
 267     let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
 268
 269     let (mut ifmt, dst_vinfo) = if let NACodecTypeInfo::Video(vinfo) = enc_params.format {
 270             (ScaleInfo { fmt: vinfo.format, width: vinfo.width, height: vinfo.height },
 271              vinfo)
 272         } else {
 273             (ScaleInfo { fmt: YUV420_FORMAT, width: 2, height: 2 },
 274              NAVideoInfo { width: 2, height: 2, format: YUV420_FORMAT, flipped: false, bits: 12 })
 275         };
 276     let ofmt = ifmt;
 277     let mut scaler = NAScale::new(ifmt, ofmt).unwrap();
 278     let mut cvt_buf = alloc_video_buffer(dst_vinfo, 2).unwrap();
 279     loop {
 280         let pktres = dmx.get_frame();
 281         if let Err(e) = pktres {
 282             if e == DemuxerError::EOF { break; }
 283             panic!("decoding error");
 284         }
 285         let pkt = pktres.unwrap();
 286         if pkt.get_stream().id != in_stream_id { continue; }
 287         let frm = dec.decode(&mut dsupp, &pkt).unwrap();
 288         let buf = frm.get_buffer();
 289         let cfrm = if let NACodecTypeInfo::Video(_) = enc_params.format {
 290                 let cur_ifmt = get_scale_fmt_from_pic(&buf);
 291                 if cur_ifmt != ifmt {
 292                     ifmt = cur_ifmt;
 293                     scaler = NAScale::new(ifmt, ofmt).unwrap();
 294                 }
 295                 scaler.convert(&buf, &mut cvt_buf).unwrap();
 296                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf.clone())
 297             } else if let NACodecTypeInfo::Audio(ref dst_ainfo) = enc_params.format {
 298                 let cvt_buf = convert_audio_frame(&buf, dst_ainfo, &dst_chmap).unwrap();
 299                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf)
 300             } else {
 301                 panic!("unexpected format");
 302             };
 303         encoder.encode(&cfrm).unwrap();
 304         while let Ok(Some(pkt)) = encoder.get_packet() {
 305             mux.mux_frame(pkt).unwrap();
 306         }
 307         if let Some(maxts) = dec_config.limit {
 308             if frm.get_pts().unwrap_or(0) >= maxts {
 309                 break;
 310             }
 311         }
 312     }
 313     encoder.flush().unwrap();
 314     while let Ok(Some(pkt)) = encoder.get_packet() {
 315         mux.mux_frame(pkt).unwrap();
 316     }
 317     mux.end().unwrap();
 318 }