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

   1 use std::fs::File;
   2 use nihav_core::frame::*;
   3 use nihav_core::codecs::*;
   4 use nihav_core::demuxers::*;
   5 use nihav_core::muxers::*;
   6 use nihav_core::scale::*;
   7 use nihav_core::soundcvt::*;
   8
   9 pub struct DecoderTestParams {
  10     pub demuxer:        &'static str,
  11     pub in_name:        &'static str,
  12     pub limit:          Option<u64>,
  13     pub stream_type:    StreamType,
  14     pub dmx_reg:        RegisteredDemuxers,
  15     pub dec_reg:        RegisteredDecoders,
  16 }
  17
  18 pub struct EncoderTestParams {
  19     pub muxer:          &'static str,
  20     pub enc_name:       &'static str,
  21     pub out_name:       &'static str,
  22     pub mux_reg:        RegisteredMuxers,
  23     pub enc_reg:        RegisteredEncoders,
  24 }
  25
  26 pub fn test_encoding_to_file(dec_config: &DecoderTestParams, enc_config: &EncoderTestParams, mut enc_params: EncodeParameters) {
  27     let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
  28     let mut file = File::open(dec_config.in_name).unwrap();
  29     let mut fr = FileReader::new_read(&mut file);
  30     let mut br = ByteReader::new(&mut fr);
  31     let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
  32
  33     let in_stream = dmx.get_streams().find(|str| str.get_media_type() == dec_config.stream_type).unwrap();
  34     let in_stream_id = in_stream.id;
  35     let decfunc = dec_config.dec_reg.find_decoder(in_stream.get_info().get_name()).unwrap();
  36     let mut dec = (decfunc)();
  37     let mut dsupp = Box::new(NADecoderSupport::new());
  38     dec.init(&mut dsupp, in_stream.get_info()).unwrap();
  39
  40     let mut out_sm = StreamManager::new();
  41     enc_params.tb_num = in_stream.tb_num;
  42     enc_params.tb_den = in_stream.tb_den;
  43
  44     if let (NACodecTypeInfo::Video(ref mut vinfo), Some(ref_vinfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_video_info()) {
  45         if vinfo.width == 0 {
  46             vinfo.width  = ref_vinfo.width;
  47             vinfo.height = ref_vinfo.height;
  48         }
  49     }
  50     let mut dst_chmap = NAChannelMap::new();
  51     if let (NACodecTypeInfo::Audio(ref mut ainfo), Some(ref_ainfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_audio_info()) {
  52         if ainfo.sample_rate == 0 {
  53             ainfo.sample_rate = ref_ainfo.sample_rate;
  54         }
  55         if ainfo.channels == 0 {
  56             ainfo.channels = ref_ainfo.channels;
  57         }
  58         match ainfo.channels {
  59             1 => {
  60                 dst_chmap.add_channel(NAChannelType::C);
  61             },
  62             2 => {
  63                 dst_chmap.add_channel(NAChannelType::L);
  64                 dst_chmap.add_channel(NAChannelType::R);
  65             },
  66             _ => panic!("cannot guess channel map"),
  67         }
  68     }
  69
  70     let encfunc = enc_config.enc_reg.find_encoder(enc_config.enc_name).unwrap();
  71     let mut encoder = (encfunc)();
  72     let out_str = encoder.init(0, enc_params).unwrap();
  73     out_sm.add_stream(NAStream::clone(&out_str));
  74
  75     let mux_f = enc_config.mux_reg.find_muxer(enc_config.muxer).unwrap();
  76     let out_name = "assets/test_out/".to_owned() + enc_config.out_name;
  77     let file = File::create(&out_name).unwrap();
  78     let mut fw = FileWriter::new_write(file);
  79     let mut bw = ByteWriter::new(&mut fw);
  80     let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
  81
  82     let (mut ifmt, dst_vinfo) = if let NACodecTypeInfo::Video(vinfo) = enc_params.format {
  83             (ScaleInfo { fmt: vinfo.format, width: vinfo.width, height: vinfo.height },
  84              vinfo)
  85         } else {
  86             (ScaleInfo { fmt: YUV420_FORMAT, width: 2, height: 2 },
  87              NAVideoInfo { width: 2, height: 2, format: YUV420_FORMAT, flipped: false })
  88         };
  89     let ofmt = ifmt;
  90     let mut scaler = NAScale::new(ifmt, ofmt).unwrap();
  91     let mut cvt_buf = alloc_video_buffer(dst_vinfo, 2).unwrap();
  92     loop {
  93         let pktres = dmx.get_frame();
  94         if let Err(e) = pktres {
  95             if e == DemuxerError::EOF { break; }
  96             panic!("decoding error");
  97         }
  98         let pkt = pktres.unwrap();
  99         if pkt.get_stream().id != in_stream_id { continue; }
 100         let frm = dec.decode(&mut dsupp, &pkt).unwrap();
 101         let buf = frm.get_buffer();
 102         let cfrm = if let NACodecTypeInfo::Video(_) = enc_params.format {
 103                 let cur_ifmt = get_scale_fmt_from_pic(&buf);
 104                 if cur_ifmt != ifmt {
 105                     ifmt = cur_ifmt;
 106                     scaler = NAScale::new(ifmt, ofmt).unwrap();
 107                 }
 108                 scaler.convert(&buf, &mut cvt_buf).unwrap();
 109                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf.clone())
 110             } else if let NACodecTypeInfo::Audio(ref dst_ainfo) = enc_params.format {
 111                 let cvt_buf = convert_audio_frame(&buf, dst_ainfo, &dst_chmap).unwrap();
 112                 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf)
 113             } else {
 114                 panic!("unexpected format");
 115             };
 116         encoder.encode(&cfrm).unwrap();
 117         while let Ok(Some(pkt)) = encoder.get_packet() {
 118             mux.mux_frame(pkt).unwrap();
 119         }
 120         if let Some(maxts) = dec_config.limit {
 121             if frm.get_pts().unwrap_or(0) >= maxts {
 122                 break;
 123             }
 124         }
 125     }
 126     encoder.flush().unwrap();
 127     while let Ok(Some(pkt)) = encoder.get_packet() {
 128         mux.mux_frame(pkt).unwrap();
 129     }
 130     mux.end().unwrap();
 131 }