| 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 }) |
| 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 | } |