From 90683bc3ac9d2adbe4a6b7239e727a36bd513678 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Sun, 13 Aug 2023 18:04:00 +0200 Subject: [PATCH] add PNM image sequence as possible input source --- src/demux.rs | 14 +++ src/imgseq.rs | 330 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 22 +++- 3 files changed, 365 insertions(+), 1 deletion(-) create mode 100644 src/imgseq.rs diff --git a/src/demux.rs b/src/demux.rs index 0297c2f..61ddbaf 100644 --- a/src/demux.rs +++ b/src/demux.rs @@ -6,6 +6,7 @@ use nihav_registry::detect; use nihav_core::io::byteio::ByteReader; use nihav_allstuff::*; use crate::null::*; +use crate::imgseq::*; pub struct FullRegister { pub dmx_reg: RegisteredDemuxers, @@ -57,6 +58,7 @@ pub enum DemuxerObject<'a> { Normal(Demuxer<'a>), Raw(RawDemuxer<'a>, Vec>>, bool), RawStream(RawStreamCtx<'a>), + ImageSequence(ImgSeqDemuxer), } impl<'a> DemuxerObject<'a> { @@ -141,6 +143,9 @@ impl<'a> DemuxerObject<'a> { DemuxerObject::None } } + pub fn create_imgseq(isd: ImgSeqDemuxer) -> Self { + DemuxerObject::ImageSequence(isd) + } pub fn is_none(&self) -> bool { matches!(*self, DemuxerObject::None) } @@ -160,6 +165,7 @@ impl<'a> DemuxerObject<'a> { DemuxerObject::Normal(ref dmx) => dmx.get_num_streams(), DemuxerObject::Raw(ref dmx, _, _) => dmx.get_num_streams(), DemuxerObject::RawStream(_) => 1, + DemuxerObject::ImageSequence(_) => 1, } } pub fn get_stream(&self, idx: usize) -> Option { @@ -167,6 +173,7 @@ impl<'a> DemuxerObject<'a> { DemuxerObject::Normal(ref dmx) => dmx.get_stream(idx), DemuxerObject::Raw(ref dmx, _, _) => dmx.get_stream(idx), DemuxerObject::RawStream(ref ctx) if idx == 0 => Some(ctx.stream.clone()), + DemuxerObject::ImageSequence(ref ctx) if idx == 0 => Some(ctx.stream.clone()), _ => None, } } @@ -175,6 +182,7 @@ impl<'a> DemuxerObject<'a> { DemuxerObject::Normal(ref dmx) => dmx.get_streams(), DemuxerObject::Raw(ref dmx, _, _) => dmx.get_streams(), DemuxerObject::RawStream(ref ctx) => ctx.sm.iter(), + DemuxerObject::ImageSequence(ref ctx) => ctx.sm.iter(), _ => unreachable!(), } } @@ -183,6 +191,7 @@ impl<'a> DemuxerObject<'a> { DemuxerObject::Normal(ref dmx) => dmx.get_stream_manager(), DemuxerObject::Raw(ref dmx, _, _) => dmx.get_stream_manager(), DemuxerObject::RawStream(ref ctx) => &ctx.sm, + DemuxerObject::ImageSequence(ref ctx) => &ctx.sm, _ => unreachable!(), } } @@ -262,6 +271,7 @@ impl<'a> DemuxerObject<'a> { }; } }, + DemuxerObject::ImageSequence(ref mut ctx) => ctx.get_frame(), _ => unreachable!(), } } @@ -269,6 +279,7 @@ impl<'a> DemuxerObject<'a> { match *self { DemuxerObject::Normal(ref mut dmx) => dmx.seek(seek_time), DemuxerObject::Raw(ref mut dmx, _, _) => dmx.seek(seek_time), + DemuxerObject::ImageSequence(ref mut ctx) => ctx.seek(seek_time), _ => Err(DemuxerError::NotImplemented), } } @@ -279,6 +290,7 @@ impl<'a> NAOptionHandler for DemuxerObject<'a> { match *self { DemuxerObject::Normal(ref dmx) => dmx.get_supported_options(), DemuxerObject::Raw(ref dmx, _, _) => dmx.get_supported_options(), + DemuxerObject::ImageSequence(ref ctx) => ctx.get_supported_options(), _ => &[], } } @@ -286,6 +298,7 @@ impl<'a> NAOptionHandler for DemuxerObject<'a> { match *self { DemuxerObject::Normal(ref mut dmx) => dmx.set_options(options), DemuxerObject::Raw(ref mut dmx, _, _) => dmx.set_options(options), + DemuxerObject::ImageSequence(ref mut ctx) => ctx.set_options(options), _ => {}, } } @@ -293,6 +306,7 @@ impl<'a> NAOptionHandler for DemuxerObject<'a> { match *self { DemuxerObject::Normal(ref dmx) => dmx.query_option_value(name), DemuxerObject::Raw(ref dmx, _, _) => dmx.query_option_value(name), + DemuxerObject::ImageSequence(ref ctx) => ctx.query_option_value(name), _ => None, } } diff --git a/src/imgseq.rs b/src/imgseq.rs new file mode 100644 index 0000000..dad9c81 --- /dev/null +++ b/src/imgseq.rs @@ -0,0 +1,330 @@ +use nihav_core::frame::*; +use nihav_core::demuxers::*; +use std::fs::File; +use std::io::BufReader; +use std::io::Read; + +struct TemplateName { + prefix: String, + pad_size: usize, + suffix: String, + single: bool, +} + +trait Deescape { + fn deescape(&mut self); +} + +impl Deescape for String { + fn deescape(&mut self) { + while let Some(idx) = self.find("%%") { + self.remove(idx + 1); + } + } +} + +impl TemplateName { + fn new(name: &str) -> Self { + let mut off = 0; + let mut tmpl_start = 0; + let mut tmpl_end = 0; + let mut pad_size = 0; + 'parse_loop: + while let Some(idx) = name[off..].find('%') { + let idx = idx + off; + if idx + 1 == name.len() { + break; + } + if name[idx + 1..].starts_with('%') { // escape, skip it + off += 1; + } + if name[idx + 1..].starts_with('0') { + if let Some(end_idx) = name[idx + 2..].find('d') { + if let Ok(val) = name[idx + 2..][..end_idx].parse::() { + if val <= 32 { + tmpl_start = idx; + pad_size = val; + tmpl_end = idx + 2 + end_idx + 1; + } + } + break 'parse_loop; + } + } + if name[idx + 1..].starts_with('d') { + tmpl_start = idx; + tmpl_end = idx + 2; + break; + } + off += idx; + } + + if tmpl_end == 0 { + let mut prefix = name.to_owned(); + prefix.deescape(); + Self { + prefix, + pad_size: 0, + suffix: String::new(), + single: true, + } + } else { + let mut prefix = name[..tmpl_start].to_string(); + prefix.deescape(); + let mut suffix = name[tmpl_end..].to_string(); + suffix.deescape(); + Self { + prefix, suffix, pad_size, + single: false, + } + } + } + fn format(&self, id: T) -> String { + let mut number = id.to_string(); + while number.len() < self.pad_size { + number.insert(0, '0'); + } + let mut fname = String::with_capacity(self.prefix.len() + number.len() + self.suffix.len()); + fname.push_str(&self.prefix); + fname.push_str(&number); + fname.push_str(&self.suffix); + fname + } +} + +pub struct ImgSeqDemuxer { + pub stream: NAStreamRef, + pub sm: StreamManager, + cur_frame: u64, + template: TemplateName, + pgmyuv: bool, +} + +impl ImgSeqDemuxer { + fn new(stream: NAStreamRef, cur_frame: u64, template: TemplateName, pgmyuv: bool) -> Self { + let mut sm = StreamManager::new(); + sm.add_stream_ref(stream.clone()); + Self { + stream, sm, cur_frame, template, pgmyuv, + } + } + pub fn seek(&mut self, time: NATimePoint) -> DemuxerResult<()> { + self.cur_frame = match time { + NATimePoint::None => return Ok(()), + NATimePoint::Milliseconds(ms) => NATimeInfo::time_to_ts(ms, 1000, self.stream.tb_num, self.stream.tb_den), + NATimePoint::PTS(pts) => pts, + }; + Ok(()) + } + pub fn get_frame(&mut self) -> DemuxerResult { + if self.cur_frame > 0 && self.template.single { + return Err(DemuxerError::EOF); + } + let fname = self.template.format(self.cur_frame); + if let Ok(file) = File::open(fname.as_str()) { + let mut file = BufReader::new(file); + let vinfo = read_pnm_header(&mut file, self.pgmyuv)?; + let pkt_size = if vinfo.format.model.is_yuv() && vinfo.format.components == 3 { + vinfo.width * (vinfo.height * 3 / 2) * if vinfo.format.get_max_depth() > 8 { 2 } else { 1 } + } else { + vinfo.width * vinfo.height * usize::from(vinfo.format.components) + }; + let mut buf = vec![0; pkt_size]; + file.read_exact(&mut buf).map_err(|_| DemuxerError::IOError)?; + let ts = NATimeInfo::new(Some(self.cur_frame), None, None, self.stream.tb_num, self.stream.tb_den); + let pkt = NAPacket::new(self.stream.clone(), ts, true, buf); + self.cur_frame += 1; + Ok(pkt) + } else { + Err(DemuxerError::EOF) + } + } +} + +impl NAOptionHandler for ImgSeqDemuxer { + fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] } + fn set_options(&mut self, _options: &[NAOption]) {} + fn query_option_value(&self, _name: &str) -> Option { None } +} + +pub struct ImgSeqDemuxerCreator<'a> { + name: &'a str, + start: Option, + pgmyuv: bool, + tb_num: u32, + tb_den: u32, +} + +impl<'a> ImgSeqDemuxerCreator<'a> { + pub fn new(name: &'a str) -> Self { + Self { + name, + start: None, + pgmyuv: false, + tb_num: 1, + tb_den: 25, + } + } + pub fn open(&mut self) -> DemuxerResult { + let template = TemplateName::new(self.name); + + let fname = template.format(self.start.unwrap_or(0)); + + if let Ok(file) = File::open(fname.as_str()) { + let mut file = BufReader::new(file); + let vinfo = read_pnm_header(&mut file, self.pgmyuv)?; + let cinfo = NACodecInfo::new(if vinfo.format.model.is_yuv() { "rawvideo" } else { "rawvideo-ms" }, NACodecTypeInfo::Video(vinfo), None); + let stream = NAStream::new(StreamType::Video, 0, cinfo, self.tb_num, self.tb_den, 0).into_ref(); + return Ok(ImgSeqDemuxer::new(stream, self.start.unwrap_or(0), template, self.pgmyuv)); + } + + // if start is not given, try also starting from one + if self.start.is_none() { + let new_start = 1; + let fname = template.format(new_start); + if let Ok(file) = File::open(fname.as_str()) { + let mut file = BufReader::new(file); + let vinfo = read_pnm_header(&mut file, self.pgmyuv)?; + let cinfo = NACodecInfo::new(if vinfo.format.model.is_yuv() { "rawvideo" } else { "rawvideo-ms" }, NACodecTypeInfo::Video(vinfo), None); + let stream = NAStream::new(StreamType::Video, 0, cinfo, self.tb_num, self.tb_den, 0).into_ref(); + return Ok(ImgSeqDemuxer::new(stream, new_start, template, self.pgmyuv)); + } + } + + Err(DemuxerError::NoSuchInput) + } +} + +const IMGSEQ_OPTIONS: &[NAOptionDefinition] = &[ + NAOptionDefinition { + name: "start", description: "start frame number", + opt_type: NAOptionDefinitionType::Int(Some(0), None) }, + NAOptionDefinition { + name: "pgmyuv", description: "Input is in PGMYUV format", + opt_type: NAOptionDefinitionType::Bool }, + NAOptionDefinition { + name: "tb_num", description: "timebase numerator", + opt_type: NAOptionDefinitionType::Int(Some(1), Some(1000000)) }, + NAOptionDefinition { + name: "tb_den", description: "timebase denominator", + opt_type: NAOptionDefinitionType::Int(Some(1), Some(1000000)) }, +]; + +impl<'a> NAOptionHandler for ImgSeqDemuxerCreator<'a> { + fn get_supported_options(&self) -> &[NAOptionDefinition] { IMGSEQ_OPTIONS } + fn set_options(&mut self, options: &[NAOption]) { + for option in options.iter() { + for opt_def in IMGSEQ_OPTIONS.iter() { + if opt_def.check(option).is_ok() { + match (option.name, &option.value) { + ("start", NAValue::Int(intval)) => { + self.start = Some(*intval as u64); + }, + ("pgmyuv", NAValue::Bool(bval)) => { + self.pgmyuv = *bval; + }, + ("tb_num", NAValue::Int(intval)) => { + self.tb_num = *intval as u32; + }, + ("tb_den", NAValue::Int(intval)) => { + self.tb_den = *intval as u32; + }, + _ => {}, + } + } + } + } + } + fn query_option_value(&self, name: &str) -> Option { + match name { + "start" => Some(NAValue::Int(self.start.unwrap_or(0) as i64)), + "pgmyuv" => Some(NAValue::Bool(self.pgmyuv)), + "tb_num" => Some(NAValue::Int(self.tb_num as i64)), + "tb_den" => Some(NAValue::Int(self.tb_den as i64)), + _ => None, + } + } +} + +fn read_pnm_header(file: &mut BufReader, pgmyuv: bool) -> DemuxerResult { + let mut fr = FileReader::new_read(file); + let mut br = ByteReader::new(&mut fr); + + let mut magic = [0; 2]; + br.read_buf(&mut magic)?; + if magic[0] != b'P' { return Err(DemuxerError::InvalidData); } + match magic[1] { + b'4' | // PBM, PBM ASCII + b'1' => return Err(DemuxerError::NotImplemented), + b'5' => { // PGM + }, + b'2' => return Err(DemuxerError::NotImplemented), // PGM ASCII + b'6' => { // PPM + }, + b'3' => return Err(DemuxerError::NotImplemented), // PPM ASCII + _ => return Err(DemuxerError::InvalidData), + }; + if br.read_byte()? != b'\n' { return Err(DemuxerError::InvalidData); } + let w = read_number(&mut br)?; + let h = read_number(&mut br)?; + let maxval = if matches!(magic[1], b'4' | b'1') { 1 } else { read_number(&mut br)? }; + if maxval > 65535 || (maxval & (maxval + 1)) != 0 { return Err(DemuxerError::InvalidData); } + let bits = maxval.count_ones() as u8; + + let mut vinfo = NAVideoInfo::new(w, h, false, RGB24_FORMAT); + match magic[1] { + b'5' | b'2' if !pgmyuv => { + vinfo.format = NAPixelFormaton { + model: ColorModel::YUV(YUVSubmodel::YUVJ), + components: 1, + comp_info: [Some(NAPixelChromaton{h_ss: 0, v_ss: 0, packed: false, depth: 8, shift: 0, comp_offs: 0, next_elem: 1}), None, None, None, None], + elem_size: 1, + be: true, + alpha: false, + palette: false, + }; + vinfo.bits = bits; + }, + b'5' | b'2' => { + if ((w & 1) != 0) || ((h % 3) != 0) { return Err(DemuxerError::InvalidData); } + vinfo.format = YUV420_FORMAT; + vinfo.height = h * 2 / 3; + vinfo.bits = bits * 3 / 2; + }, + b'6' | b'3' => { + vinfo.format = RGB24_FORMAT; + vinfo.bits = bits; + }, + _ => unreachable!(), + }; + if bits != 8 { + for chr in vinfo.format.comp_info.iter_mut().flatten() { + chr.depth = bits; + if bits > 8 { + chr.next_elem = 2; + } + } + if bits > 8 { + vinfo.format.elem_size <<= 1; + } + } + + Ok(vinfo) +} + +fn read_number(br: &mut ByteReader) -> DemuxerResult { + let mut val = 0; + loop { + let c = br.read_byte()?; + match c { + b'0'..=b'9' => { + if val > 1048576 { + return Err(DemuxerError::InvalidData); + } + val = val * 10 + usize::from(c - b'0'); + }, + b' ' | b'\n' => break, + _ => return Err(DemuxerError::InvalidData), + }; + } + Ok(val) +} diff --git a/src/main.rs b/src/main.rs index e033dfd..97be1e2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,6 +23,8 @@ use crate::demux::*; mod null; mod acvt; use crate::acvt::*; +mod imgseq; +use crate::imgseq::*; const SND_NO_FORMAT: NASoniton = NASoniton { bits: 0, be: false, packed: false, planar: false, float: false, signed: false }; @@ -1068,6 +1070,20 @@ fn main() { return; } + if let Some("imgseq") = transcoder.input_fmt.as_deref() { + let mut isdc = ImgSeqDemuxerCreator::new(transcoder.input_name.as_str()); + parse_and_apply_options!(isdc, &transcoder.demux_opts, "input"); + let isd = if let Ok(ctx) = isdc.open() { + ctx + } else { + println!("failed to create image sequence demuxer!"); + return; + }; + let mut dmx = DemuxerObject::create_imgseq(isd); + transcode(transcoder, full_reg, &mut dmx); + return; + } + let res = File::open(transcoder.input_name.as_str()); if res.is_err() { println!("error opening input"); @@ -1097,10 +1113,14 @@ fn main() { println!("cannot find demuxer for '{}'", transcoder.input_name.as_str()); return; } + parse_and_apply_options!(dmx, &transcoder.demux_opts, "input"); + transcode(transcoder, full_reg, &mut dmx); +} + +fn transcode(mut transcoder: Transcoder, full_reg: FullRegister, dmx: &mut DemuxerObject) { let duration = dmx.get_duration(); let duration_string = if duration != 0 { format_time(duration) } else { String::new() }; - parse_and_apply_options!(dmx, &transcoder.demux_opts, "input"); for i in 0..dmx.get_num_streams() { let s = dmx.get_stream(i).unwrap(); let info = s.get_info(); -- 2.30.2