codec_support/test: add remuxing tests
[nihav.git] / nihav-codec-support / src / test / enc_video.rs
CommitLineData
0b257d9f
KS
1use std::fs::File;
2use nihav_core::frame::*;
3use nihav_core::codecs::*;
4use nihav_core::demuxers::*;
5use nihav_core::muxers::*;
6use nihav_core::scale::*;
7use nihav_core::soundcvt::*;
f9304550 8use super::md5::MD5;
0b257d9f
KS
9
10pub struct DecoderTestParams {
11 pub demuxer: &'static str,
12 pub in_name: &'static str,
13 pub limit: Option<u64>,
14 pub stream_type: StreamType,
15 pub dmx_reg: RegisteredDemuxers,
16 pub dec_reg: RegisteredDecoders,
17}
18
19pub struct EncoderTestParams {
20 pub muxer: &'static str,
21 pub enc_name: &'static str,
22 pub out_name: &'static str,
23 pub mux_reg: RegisteredMuxers,
24 pub enc_reg: RegisteredEncoders,
25}
26
f9304550
KS
27pub fn test_remuxing(dec_config: &DecoderTestParams, enc_config: &EncoderTestParams) {
28 let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
29 let mut file = File::open(dec_config.in_name).unwrap();
30 let mut fr = FileReader::new_read(&mut file);
31 let mut br = ByteReader::new(&mut fr);
32 let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
33
34 let mux_f = enc_config.mux_reg.find_muxer(enc_config.muxer).unwrap();
35 let out_name = "assets/test_out/".to_owned() + enc_config.out_name;
36 let file = File::create(&out_name).unwrap();
37 let mut fw = FileWriter::new_write(file);
38 let mut bw = ByteWriter::new(&mut fw);
39 let mut out_sm = StreamManager::new();
40 let mux_caps = mux_f.get_capabilities();
41 let mut stream_map: Vec<Option<usize>> = Vec::new();
42 let mut has_video = false;
43 let mut has_audio = false;
44 for stream in dmx.get_streams() {
45 let mut copy_stream = false;
46 match mux_caps {
47 MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::OnlyVideo => {
48 copy_stream = stream.get_media_type() == StreamType::Video;
49 has_video = true;
50 },
51 MuxerCapabilities::SingleAudio(_) | MuxerCapabilities::OnlyAudio => {
52 copy_stream = stream.get_media_type() == StreamType::Audio;
53 has_audio = true;
54 },
55 MuxerCapabilities::SingleVideoAndAudio(_, _) => {
56 if stream.get_media_type() == StreamType::Video {
57 copy_stream = !has_video;
58 has_video = true;
59 }
60 if stream.get_media_type() == StreamType::Audio {
61 copy_stream = !has_audio;
62 has_audio = true;
63 }
64 },
65 MuxerCapabilities::Universal => {
66 if stream.get_media_type() == StreamType::Video {
67 copy_stream = true;
68 has_video = true;
69 }
70 if stream.get_media_type() == StreamType::Audio {
71 copy_stream = true;
72 has_audio = true;
73 }
74 }
75 };
76 if copy_stream {
77 let streamno = out_sm.add_stream(NAStream::clone(&stream)).unwrap();
78 stream_map.push(Some(streamno));
79 match mux_caps {
80 MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::SingleAudio(_) => break,
81 _ => {},
82 };
83 } else {
84 stream_map.push(None);
85 }
86 }
87 assert!(out_sm.get_num_streams() > 0);
88 let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
89
90 loop {
91 let pktres = dmx.get_frame();
92 if let Err(e) = pktres {
93 if e == DemuxerError::EOF { break; }
94 panic!("error");
95 }
96 let mut pkt = pktres.unwrap();
97 println!("Got {}", pkt);
98 if let Some(new_id) = stream_map[pkt.get_stream().id as usize] {
99 pkt.reassign(mux.get_stream(new_id).unwrap(), pkt.get_time_information());
100 mux.mux_frame(pkt).unwrap();
101 }
102 }
103
104 mux.end().unwrap();
105}
106
107pub fn test_remuxing_md5(dec_config: &DecoderTestParams, muxer: &str, mux_reg: &RegisteredMuxers, md5_hash: [u32; 4]) {
108 let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
109 let mut file = File::open(dec_config.in_name).unwrap();
110 let mut fr = FileReader::new_read(&mut file);
111 let mut br = ByteReader::new(&mut fr);
112 let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
113
114 let mux_f = mux_reg.find_muxer(muxer).unwrap();
115
116 let mut dst = Vec::with_capacity(1048576);
117 let mut gw = GrowableMemoryWriter::new_write(&mut dst);
118 let mut bw = ByteWriter::new(&mut gw);
119 let mut out_sm = StreamManager::new();
120 let mux_caps = mux_f.get_capabilities();
121 let mut stream_map: Vec<Option<usize>> = Vec::new();
122 let mut has_video = false;
123 let mut has_audio = false;
124 for stream in dmx.get_streams() {
125 let mut copy_stream = false;
126 match mux_caps {
127 MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::OnlyVideo => {
128 copy_stream = stream.get_media_type() == StreamType::Video;
129 has_video = true;
130 },
131 MuxerCapabilities::SingleAudio(_) | MuxerCapabilities::OnlyAudio => {
132 copy_stream = stream.get_media_type() == StreamType::Audio;
133 has_audio = true;
134 },
135 MuxerCapabilities::SingleVideoAndAudio(_, _) => {
136 if stream.get_media_type() == StreamType::Video {
137 copy_stream = !has_video;
138 has_video = true;
139 }
140 if stream.get_media_type() == StreamType::Audio {
141 copy_stream = !has_audio;
142 has_audio = true;
143 }
144 },
145 MuxerCapabilities::Universal => {
146 if stream.get_media_type() == StreamType::Video {
147 copy_stream = true;
148 has_video = true;
149 }
150 if stream.get_media_type() == StreamType::Audio {
151 copy_stream = true;
152 has_audio = true;
153 }
154 }
155 };
156 if copy_stream {
157 let streamno = out_sm.add_stream(NAStream::clone(&stream)).unwrap();
158 stream_map.push(Some(streamno));
159 match mux_caps {
160 MuxerCapabilities::SingleVideo(_) | MuxerCapabilities::SingleAudio(_) => break,
161 _ => {},
162 };
163 } else {
164 stream_map.push(None);
165 }
166 }
167 assert!(out_sm.get_num_streams() > 0);
168 let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
169
170 loop {
171 let pktres = dmx.get_frame();
172 if let Err(e) = pktres {
173 if e == DemuxerError::EOF { break; }
174 panic!("error");
175 }
176 let mut pkt = pktres.unwrap();
177 println!("Got {}", pkt);
178 if let Some(new_id) = stream_map[pkt.get_stream().id as usize] {
179 pkt.reassign(mux.get_stream(new_id).unwrap(), pkt.get_time_information());
180 mux.mux_frame(pkt).unwrap();
181 }
182 }
183
184 mux.end().unwrap();
185
186 let mut hash = [0; 4];
187 MD5::calculate_hash(dst.as_slice(), &mut hash);
188 println!("output hash {:08x}{:08x}{:08x}{:08x}", hash[0], hash[1], hash[2], hash[3]);
189 assert_eq!(hash, md5_hash);
190}
191
0b257d9f
KS
192pub fn test_encoding_to_file(dec_config: &DecoderTestParams, enc_config: &EncoderTestParams, mut enc_params: EncodeParameters) {
193 let dmx_f = dec_config.dmx_reg.find_demuxer(dec_config.demuxer).unwrap();
194 let mut file = File::open(dec_config.in_name).unwrap();
195 let mut fr = FileReader::new_read(&mut file);
196 let mut br = ByteReader::new(&mut fr);
197 let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
198
199 let in_stream = dmx.get_streams().find(|str| str.get_media_type() == dec_config.stream_type).unwrap();
200 let in_stream_id = in_stream.id;
201 let decfunc = dec_config.dec_reg.find_decoder(in_stream.get_info().get_name()).unwrap();
202 let mut dec = (decfunc)();
203 let mut dsupp = Box::new(NADecoderSupport::new());
204 dec.init(&mut dsupp, in_stream.get_info()).unwrap();
205
206 let mut out_sm = StreamManager::new();
207 enc_params.tb_num = in_stream.tb_num;
208 enc_params.tb_den = in_stream.tb_den;
209
210 if let (NACodecTypeInfo::Video(ref mut vinfo), Some(ref_vinfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_video_info()) {
211 if vinfo.width == 0 {
212 vinfo.width = ref_vinfo.width;
213 vinfo.height = ref_vinfo.height;
214 }
215 }
216 let mut dst_chmap = NAChannelMap::new();
217 if let (NACodecTypeInfo::Audio(ref mut ainfo), Some(ref_ainfo)) = (&mut enc_params.format, in_stream.get_info().get_properties().get_audio_info()) {
218 if ainfo.sample_rate == 0 {
219 ainfo.sample_rate = ref_ainfo.sample_rate;
220 }
221 if ainfo.channels == 0 {
222 ainfo.channels = ref_ainfo.channels;
223 }
224 match ainfo.channels {
225 1 => {
226 dst_chmap.add_channel(NAChannelType::C);
227 },
228 2 => {
229 dst_chmap.add_channel(NAChannelType::L);
230 dst_chmap.add_channel(NAChannelType::R);
231 },
232 _ => panic!("cannot guess channel map"),
233 }
234 }
235
236 let encfunc = enc_config.enc_reg.find_encoder(enc_config.enc_name).unwrap();
237 let mut encoder = (encfunc)();
238 let out_str = encoder.init(0, enc_params).unwrap();
239 out_sm.add_stream(NAStream::clone(&out_str));
240
241 let mux_f = enc_config.mux_reg.find_muxer(enc_config.muxer).unwrap();
242 let out_name = "assets/test_out/".to_owned() + enc_config.out_name;
243 let file = File::create(&out_name).unwrap();
244 let mut fw = FileWriter::new_write(file);
245 let mut bw = ByteWriter::new(&mut fw);
246 let mut mux = create_muxer(mux_f, out_sm, &mut bw).unwrap();
247
248 let (mut ifmt, dst_vinfo) = if let NACodecTypeInfo::Video(vinfo) = enc_params.format {
249 (ScaleInfo { fmt: vinfo.format, width: vinfo.width, height: vinfo.height },
250 vinfo)
251 } else {
252 (ScaleInfo { fmt: YUV420_FORMAT, width: 2, height: 2 },
253 NAVideoInfo { width: 2, height: 2, format: YUV420_FORMAT, flipped: false })
254 };
255 let ofmt = ifmt;
256 let mut scaler = NAScale::new(ifmt, ofmt).unwrap();
257 let mut cvt_buf = alloc_video_buffer(dst_vinfo, 2).unwrap();
258 loop {
259 let pktres = dmx.get_frame();
260 if let Err(e) = pktres {
261 if e == DemuxerError::EOF { break; }
262 panic!("decoding error");
263 }
264 let pkt = pktres.unwrap();
265 if pkt.get_stream().id != in_stream_id { continue; }
266 let frm = dec.decode(&mut dsupp, &pkt).unwrap();
267 let buf = frm.get_buffer();
268 let cfrm = if let NACodecTypeInfo::Video(_) = enc_params.format {
269 let cur_ifmt = get_scale_fmt_from_pic(&buf);
270 if cur_ifmt != ifmt {
271 ifmt = cur_ifmt;
272 scaler = NAScale::new(ifmt, ofmt).unwrap();
273 }
274 scaler.convert(&buf, &mut cvt_buf).unwrap();
275 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf.clone())
276 } else if let NACodecTypeInfo::Audio(ref dst_ainfo) = enc_params.format {
277 let cvt_buf = convert_audio_frame(&buf, dst_ainfo, &dst_chmap).unwrap();
278 NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cvt_buf)
279 } else {
280 panic!("unexpected format");
281 };
282 encoder.encode(&cfrm).unwrap();
283 while let Ok(Some(pkt)) = encoder.get_packet() {
284 mux.mux_frame(pkt).unwrap();
285 }
286 if let Some(maxts) = dec_config.limit {
287 if frm.get_pts().unwrap_or(0) >= maxts {
288 break;
289 }
290 }
291 }
292 encoder.flush().unwrap();
293 while let Ok(Some(pkt)) = encoder.get_packet() {
294 mux.mux_frame(pkt).unwrap();
295 }
296 mux.end().unwrap();
297}