From 79eb70eb2b30e67eeb7db8becd6f837320ea498f Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Tue, 3 Oct 2023 17:16:37 +0200 Subject: [PATCH] add option to calculate number of frames in the stream This information may be useful when e.g. transcoding full input file to Bink. --- src/main.rs | 204 ++++++++++++++++++++++++++++++++-------------------- 1 file changed, 125 insertions(+), 79 deletions(-) diff --git a/src/main.rs b/src/main.rs index 73c4c43..8199311 100644 --- a/src/main.rs +++ b/src/main.rs @@ -186,6 +186,9 @@ struct Transcoder { start: NATimePoint, end: NATimePoint, verbose: u8, + + calc_len: bool, + nframes: Vec, } macro_rules! parse_and_apply_options { @@ -533,7 +536,7 @@ impl Transcoder { parse_and_apply_options!(dec, &self.istr_opts[str_idx].dec_opts, name); } } - fn register_output_stream(&mut self, cname: &str, istr: NAStreamRef, out_sm: &mut StreamManager, enc_reg: &RegisteredEncoders) -> RegisterResult { + fn register_output_stream(&mut self, cname: &str, istr: NAStreamRef, iidx: usize, out_sm: &mut StreamManager, enc_reg: &RegisteredEncoders) -> RegisterResult { let out_id = out_sm.get_num_streams() as u32; if let Some(str_idx) = self.istr_opts.iter().position(|el| el.id == (istr.get_num() as u32)) { if self.istr_opts[str_idx].drop { @@ -640,6 +643,10 @@ println!("can't generate default channel map for {} channels", dainfo.channels); let name = format!("output stream {}", out_id); parse_and_apply_options!(encoder, &oopts.enc_opts, name); + if self.calc_len && self.nframes.len() > iidx { + encoder.set_options(&[NAOption{name: "nframes", value: NAValue::Int(self.nframes[iidx] as i64)}]); + } + let ret = encoder.init(out_id, ret_eparams); if ret.is_err() { println!("error initialising encoder"); @@ -730,11 +737,11 @@ println!("can't generate default channel map for {} channels", dainfo.channels); } fn map_single(&mut self, cname: &str, ctype: StreamType, src_sm: &StreamManager, out_sm: &mut StreamManager, enc_reg: &RegisteredEncoders) -> bool { let mut found_stream = false; - for istr in src_sm.iter() { + for (iidx, istr) in src_sm.iter().enumerate() { if istr.get_media_type() != ctype || found_stream { self.encoders.push(OutputMode::Drop); } else { - match self.register_output_stream(cname, istr, out_sm, enc_reg) { + match self.register_output_stream(cname, istr, iidx, out_sm, enc_reg) { RegisterResult::Ok => found_stream = true, RegisterResult::Failed => return false, RegisterResult::Ignored => {}, @@ -756,15 +763,15 @@ println!("can't generate default channel map for {} channels", dainfo.channels); MuxerCapabilities::SingleVideoAndAudio(vname, aname) => { let mut found_vid = false; let mut found_aud = false; - for istr in src_sm.iter() { + for (iidx, istr) in src_sm.iter().enumerate() { if istr.get_media_type() == StreamType::Video && !found_vid && !self.no_video { - match self.register_output_stream(vname, istr, out_sm, enc_reg) { + match self.register_output_stream(vname, istr, iidx, out_sm, enc_reg) { RegisterResult::Ok => found_vid = true, RegisterResult::Failed => return false, RegisterResult::Ignored => {}, }; } else if istr.get_media_type() == StreamType::Audio && !found_aud && !self.no_audio { - match self.register_output_stream(aname, istr, out_sm, enc_reg) { + match self.register_output_stream(aname, istr, iidx, out_sm, enc_reg) { RegisterResult::Ok => found_aud = true, RegisterResult::Failed => return false, RegisterResult::Ignored => {}, @@ -779,9 +786,9 @@ println!("can't generate default channel map for {} channels", dainfo.channels); if self.no_video { return false; } let mut found_vid = false; - for istr in src_sm.iter() { + for (iidx, istr) in src_sm.iter().enumerate() { if istr.get_media_type() == StreamType::Video && !found_vid { - match self.register_output_stream("any", istr, out_sm, enc_reg) { + match self.register_output_stream("any", istr, iidx, out_sm, enc_reg) { RegisterResult::Ok => found_vid = true, RegisterResult::Failed => return false, RegisterResult::Ignored => {}, @@ -796,9 +803,9 @@ println!("can't generate default channel map for {} channels", dainfo.channels); if self.no_audio { return false; } let mut found_aud = false; - for istr in src_sm.iter() { + for (iidx, istr) in src_sm.iter().enumerate() { if istr.get_media_type() == StreamType::Audio && !found_aud { - match self.register_output_stream("any", istr, out_sm, enc_reg) { + match self.register_output_stream("any", istr, iidx, out_sm, enc_reg) { RegisterResult::Ok => found_aud = true, RegisterResult::Failed => return false, RegisterResult::Ignored => {}, @@ -810,13 +817,13 @@ println!("can't generate default channel map for {} channels", dainfo.channels); found_aud }, MuxerCapabilities::Universal => { - for istr in src_sm.iter() { + for (iidx, istr) in src_sm.iter().enumerate() { if (istr.get_media_type() == StreamType::Video && self.no_video) || (istr.get_media_type() == StreamType::Audio && self.no_audio) { self.encoders.push(OutputMode::Drop); continue; } - if self.register_output_stream("any", istr, out_sm, enc_reg) == RegisterResult::Failed { + if self.register_output_stream("any", istr, iidx, out_sm, enc_reg) == RegisterResult::Failed { return false; } } @@ -824,6 +831,81 @@ println!("can't generate default channel map for {} channels", dainfo.channels); }, } } + fn create_demuxers(&mut self, demuxers: &mut Vec<(DemuxerObject, bool)>, full_reg: &FullRegister, print_info: bool) -> bool { + let mut isn_start = 0; + for (i, (iname, ifmt)) in self.input_name.iter().zip( + self.input_fmt.iter()).enumerate() { + match (iname, ifmt.as_ref().map(|s| s.as_str())) { + (Some(name), Some("imgseq")) => { + println!("trying image sequence {}", name); + let mut isdc = ImgSeqDemuxerCreator::new(name.as_str()); + parse_and_apply_options!(isdc, &self.demux_opts[i], "input"); + let isd = if let Ok(ctx) = isdc.open() { + ctx + } else { + println!("failed to create image sequence demuxer!"); + return false; + }; + let dmx = DemuxerObject::create_imgseq(isd); + if print_info { + for i in 0..dmx.get_num_streams() { + let s = dmx.get_stream(i).unwrap(); + let info = s.get_info(); + println!(" stream {}({}) - {} {}", i, i + isn_start, s, info.get_name()); + } + } + isn_start += dmx.get_num_streams(); + demuxers.push((dmx, false)) + }, + (Some(name), _) => { + let res = File::open(name); + if res.is_err() { + println!("error opening input"); + return false; + } + let file = res.unwrap(); + let file = BufReader::new(file); + let mut fr = FileReader::new_read(file); + let mut br = ByteReader::new(&mut fr); + let (is_raw, start, end) = if ifmt.is_none() { + detect_tags(&mut br) + } else { + (false, 0, None) + }; + + let nfr: Box = if start != 0 || end.is_some() { + let file = fr.finish(); + Box::new(BoundedFileReader::new_read(file, start, end).unwrap()) + } else { + Box::new(fr) + }; + let sb = SelfBorrow::new(nfr, |rd| { + unsafe { + ByteReader::new(rd.as_mut().unwrap().as_mut()) + } + }); + + let mut dmx = DemuxerObject::create(sb, full_reg, name, ifmt, is_raw, print_info); + if dmx.is_none() { + println!("cannot find demuxer for '{}'", name); + return false; + } + parse_and_apply_options!(dmx, &self.demux_opts[i], "input"); + if print_info { + for i in 0..dmx.get_num_streams() { + let s = dmx.get_stream(i).unwrap(); + let info = s.get_info(); + println!(" stream {}({}) - {} {}", i, i + isn_start, s, info.get_name()); + } + } + isn_start += dmx.get_num_streams(); + demuxers.push((dmx, false)); + }, + _ => {}, + }; + } + true + } } fn encode_frame(dst_id: u32, encoder: &mut Box, cvt: &mut OutputConvert, frm: NAFrameRef, scale_opts: &[(String, String)]) -> bool { @@ -1088,6 +1170,9 @@ fn main() { return; } }, + "--calc-len" => { + transcoder.calc_len = true; + }, "--verbose" | "-v" => transcoder.verbose = 1, "-vv" => transcoder.verbose = 2, "-v-" => transcoder.verbose = 0, @@ -1149,73 +1234,8 @@ fn main() { } let mut demuxers = Vec::with_capacity(1); - let mut isn_start = 0; - for (i, (iname, ifmt)) in transcoder.input_name.iter().zip( - transcoder.input_fmt.iter()).enumerate() { - match (iname, ifmt.as_ref().map(|s| s.as_str())) { - (Some(name), Some("imgseq")) => { - println!("trying image sequence {}", name); - let mut isdc = ImgSeqDemuxerCreator::new(name.as_str()); - parse_and_apply_options!(isdc, &transcoder.demux_opts[i], "input"); - let isd = if let Ok(ctx) = isdc.open() { - ctx - } else { - println!("failed to create image sequence demuxer!"); - return; - }; - let dmx = DemuxerObject::create_imgseq(isd); - for i in 0..dmx.get_num_streams() { - let s = dmx.get_stream(i).unwrap(); - let info = s.get_info(); - println!(" stream {}({}) - {} {}", i, i + isn_start, s, info.get_name()); - } - isn_start += dmx.get_num_streams(); - demuxers.push((dmx, false)) - }, - (Some(name), _) => { - let res = File::open(name); - if res.is_err() { - println!("error opening input"); - return; - } - let file = res.unwrap(); - let file = BufReader::new(file); - let mut fr = FileReader::new_read(file); - let mut br = ByteReader::new(&mut fr); - let (is_raw, start, end) = if ifmt.is_none() { - detect_tags(&mut br) - } else { - (false, 0, None) - }; - - let nfr: Box = if start != 0 || end.is_some() { - let file = fr.finish(); - Box::new(BoundedFileReader::new_read(file, start, end).unwrap()) - } else { - Box::new(fr) - }; - let sb = SelfBorrow::new(nfr, |rd| { - unsafe { - ByteReader::new(rd.as_mut().unwrap().as_mut()) - } - }); - - let mut dmx = DemuxerObject::create(sb, &full_reg, name, ifmt, is_raw, true); - if dmx.is_none() { - println!("cannot find demuxer for '{}'", name); - return; - } - parse_and_apply_options!(dmx, &transcoder.demux_opts[i], "input"); - for i in 0..dmx.get_num_streams() { - let s = dmx.get_stream(i).unwrap(); - let info = s.get_info(); - println!(" stream {}({}) - {} {}", i, i + isn_start, s, info.get_name()); - } - isn_start += dmx.get_num_streams(); - demuxers.push((dmx, false)); - }, - _ => {}, - }; + if !transcoder.create_demuxers(&mut demuxers, &full_reg, true) { + return; } let duration = demuxers.iter().fold(0u64, |mindur, (dmx, _)| { @@ -1301,6 +1321,32 @@ fn main() { } let mux_creator = ret.unwrap(); + if transcoder.calc_len { + let mut sids = Vec::new(); + transcoder.nframes.clear(); + for (dmx, _) in demuxers.iter_mut() { + let sstart = transcoder.nframes.len(); + let sm = dmx.get_stream_manager(); + sids.clear(); + for stream in sm.iter() { + transcoder.nframes.push(0); + sids.push(stream.get_id()); + } + + while let Ok(pkt) = dmx.get_frame() { + let stream = pkt.get_stream(); + let pos = sstart + sids.iter().position(|&x| x == stream.get_id()).unwrap(); + transcoder.nframes[pos] += 1; + } + } + // this is necessary since not all demuxers allow to seek even back to the start + demuxers.clear(); + if !transcoder.create_demuxers(&mut demuxers, &full_reg, false) { + println!("failed to re-create demuxer(s)"); + return; + } + } + let mux_caps = mux_creator.get_capabilities(); let mut out_sm = StreamManager::new(); if !transcoder.negotiate_stream_map(&ism, mux_caps, &mut out_sm, &full_reg.enc_reg) { -- 2.39.5