[nihav-encoder.git] / src / imgseq.rs

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::<usize>() {
                        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<T: Sized+ToString>(&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<NAPacket> {
        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<NAValue> { None }
}

pub struct ImgSeqDemuxerCreator<'a> {
    name:       &'a str,
    start:      Option<u64>,
    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<ImgSeqDemuxer> {
        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<NAValue> {
        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<File>, pgmyuv: bool) -> DemuxerResult<NAVideoInfo> {
    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<usize> {
    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)
}
Commit	Line	Data
90683bc3 KS	1	use nihav_core::frame::*;
	2	use nihav_core::demuxers::*;
	3	use std::fs::File;
	4	use std::io::BufReader;
	5	use std::io::Read;
	6
	7	struct TemplateName {
	8	prefix: String,
	9	pad_size: usize,
	10	suffix: String,
	11	single: bool,
	12	}
	13
	14	trait Deescape {
	15	fn deescape(&mut self);
	16	}
	17
	18	impl Deescape for String {
	19	fn deescape(&mut self) {
	20	while let Some(idx) = self.find("%%") {
	21	self.remove(idx + 1);
	22	}
	23	}
	24	}
	25
	26	impl TemplateName {
	27	fn new(name: &str) -> Self {
	28	let mut off = 0;
	29	let mut tmpl_start = 0;
	30	let mut tmpl_end = 0;
	31	let mut pad_size = 0;
	32	'parse_loop:
	33	while let Some(idx) = name[off..].find('%') {
	34	let idx = idx + off;
	35	if idx + 1 == name.len() {
	36	break;
	37	}
	38	if name[idx + 1..].starts_with('%') { // escape, skip it
	39	off += 1;
	40	}
	41	if name[idx + 1..].starts_with('0') {
	42	if let Some(end_idx) = name[idx + 2..].find('d') {
	43	if let Ok(val) = name[idx + 2..][..end_idx].parse::<usize>() {
	44	if val <= 32 {
	45	tmpl_start = idx;
	46	pad_size = val;
	47	tmpl_end = idx + 2 + end_idx + 1;
	48	}
	49	}
	50	break 'parse_loop;
	51	}
	52	}
	53	if name[idx + 1..].starts_with('d') {
	54	tmpl_start = idx;
	55	tmpl_end = idx + 2;
	56	break;
	57	}
	58	off += idx;
	59	}
	60
	61	if tmpl_end == 0 {
	62	let mut prefix = name.to_owned();
	63	prefix.deescape();
	64	Self {
65	prefix,
66	pad_size: 0,
67	suffix: String::new(),
68	single: true,
69	}
70	} else {
71	let mut prefix = name[..tmpl_start].to_string();
72	prefix.deescape();
73	let mut suffix = name[tmpl_end..].to_string();
74	suffix.deescape();
75	Self {
76	prefix, suffix, pad_size,
77	single: false,
78	}
79	}
80	}
81	fn format<T: Sized+ToString>(&self, id: T) -> String {
82	let mut number = id.to_string();
83	while number.len() < self.pad_size {
84	number.insert(0, '0');
85	}
86	let mut fname = String::with_capacity(self.prefix.len() + number.len() + self.suffix.len());
87	fname.push_str(&self.prefix);
88	fname.push_str(&number);
89	fname.push_str(&self.suffix);
90	fname
91	}
92	}
93
94	pub struct ImgSeqDemuxer {
95	pub stream: NAStreamRef,
96	pub sm: StreamManager,
97	cur_frame: u64,
98	template: TemplateName,
99	pgmyuv: bool,
100	}
101
102	impl ImgSeqDemuxer {
103	fn new(stream: NAStreamRef, cur_frame: u64, template: TemplateName, pgmyuv: bool) -> Self {
104	let mut sm = StreamManager::new();
105	sm.add_stream_ref(stream.clone());
106	Self {
107	stream, sm, cur_frame, template, pgmyuv,
108	}
109	}
110	pub fn seek(&mut self, time: NATimePoint) -> DemuxerResult<()> {
111	self.cur_frame = match time {
112	NATimePoint::None => return Ok(()),
113	NATimePoint::Milliseconds(ms) => NATimeInfo::time_to_ts(ms, 1000, self.stream.tb_num, self.stream.tb_den),
114	NATimePoint::PTS(pts) => pts,
115	};
116	Ok(())
117	}
118	pub fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
119	if self.cur_frame > 0 && self.template.single {
120	return Err(DemuxerError::EOF);
121	}
122	let fname = self.template.format(self.cur_frame);
123	if let Ok(file) = File::open(fname.as_str()) {
124	let mut file = BufReader::new(file);
125	let vinfo = read_pnm_header(&mut file, self.pgmyuv)?;
126	let pkt_size = if vinfo.format.model.is_yuv() && vinfo.format.components == 3 {
127	vinfo.width * (vinfo.height * 3 / 2) * if vinfo.format.get_max_depth() > 8 { 2 } else { 1 }
128	} else {
129	vinfo.width * vinfo.height * usize::from(vinfo.format.components)
130	};
131	let mut buf = vec![0; pkt_size];
132	file.read_exact(&mut buf).map_err(\|_\| DemuxerError::IOError)?;
133	let ts = NATimeInfo::new(Some(self.cur_frame), None, None, self.stream.tb_num, self.stream.tb_den);
134	let pkt = NAPacket::new(self.stream.clone(), ts, true, buf);
135	self.cur_frame += 1;
136	Ok(pkt)
137	} else {
138	Err(DemuxerError::EOF)
139	}
140	}
141	}
142
143	impl NAOptionHandler for ImgSeqDemuxer {
144	fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
145	fn set_options(&mut self, _options: &[NAOption]) {}
146	fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
147	}
148
149	pub struct ImgSeqDemuxerCreator<'a> {
150	name: &'a str,
151	start: Option<u64>,
152	pgmyuv: bool,
153	tb_num: u32,
154	tb_den: u32,
155	}
156
157	impl<'a> ImgSeqDemuxerCreator<'a> {
158	pub fn new(name: &'a str) -> Self {
159	Self {
160	name,
161	start: None,
162	pgmyuv: false,
163	tb_num: 1,
164	tb_den: 25,
165	}
166	}
167	pub fn open(&mut self) -> DemuxerResult<ImgSeqDemuxer> {
168	let template = TemplateName::new(self.name);
169
170	let fname = template.format(self.start.unwrap_or(0));
171
172	if let Ok(file) = File::open(fname.as_str()) {
173	let mut file = BufReader::new(file);
174	let vinfo = read_pnm_header(&mut file, self.pgmyuv)?;
175	let cinfo = NACodecInfo::new(if vinfo.format.model.is_yuv() { "rawvideo" } else { "rawvideo-ms" }, NACodecTypeInfo::Video(vinfo), None);
176	let stream = NAStream::new(StreamType::Video, 0, cinfo, self.tb_num, self.tb_den, 0).into_ref();
177	return Ok(ImgSeqDemuxer::new(stream, self.start.unwrap_or(0), template, self.pgmyuv));
178	}
179
180	// if start is not given, try also starting from one
181	if self.start.is_none() {
182	let new_start = 1;
183	let fname = template.format(new_start);
184	if let Ok(file) = File::open(fname.as_str()) {
185	let mut file = BufReader::new(file);
186	let vinfo = read_pnm_header(&mut file, self.pgmyuv)?;
187	let cinfo = NACodecInfo::new(if vinfo.format.model.is_yuv() { "rawvideo" } else { "rawvideo-ms" }, NACodecTypeInfo::Video(vinfo), None);
188	let stream = NAStream::new(StreamType::Video, 0, cinfo, self.tb_num, self.tb_den, 0).into_ref();
189	return Ok(ImgSeqDemuxer::new(stream, new_start, template, self.pgmyuv));
190	}
191	}
192
193	Err(DemuxerError::NoSuchInput)
194	}
195	}
196
197	const IMGSEQ_OPTIONS: &[NAOptionDefinition] = &[
198	NAOptionDefinition {
199	name: "start", description: "start frame number",
200	opt_type: NAOptionDefinitionType::Int(Some(0), None) },
201	NAOptionDefinition {
202	name: "pgmyuv", description: "Input is in PGMYUV format",
203	opt_type: NAOptionDefinitionType::Bool },
204	NAOptionDefinition {
205	name: "tb_num", description: "timebase numerator",
206	opt_type: NAOptionDefinitionType::Int(Some(1), Some(1000000)) },
207	NAOptionDefinition {
208	name: "tb_den", description: "timebase denominator",
209	opt_type: NAOptionDefinitionType::Int(Some(1), Some(1000000)) },
210	];
211
212	impl<'a> NAOptionHandler for ImgSeqDemuxerCreator<'a> {
213	fn get_supported_options(&self) -> &[NAOptionDefinition] { IMGSEQ_OPTIONS }
214	fn set_options(&mut self, options: &[NAOption]) {
215	for option in options.iter() {
216	for opt_def in IMGSEQ_OPTIONS.iter() {
217	if opt_def.check(option).is_ok() {
218	match (option.name, &option.value) {
219	("start", NAValue::Int(intval)) => {
220	self.start = Some(*intval as u64);
221	},
222	("pgmyuv", NAValue::Bool(bval)) => {
223	self.pgmyuv = *bval;
224	},
225	("tb_num", NAValue::Int(intval)) => {
226	self.tb_num = *intval as u32;
227	},
228	("tb_den", NAValue::Int(intval)) => {
229	self.tb_den = *intval as u32;
230	},
231	_ => {},
232	}
233	}
234	}
235	}
236	}
237	fn query_option_value(&self, name: &str) -> Option<NAValue> {
238	match name {
239	"start" => Some(NAValue::Int(self.start.unwrap_or(0) as i64)),
240	"pgmyuv" => Some(NAValue::Bool(self.pgmyuv)),
241	"tb_num" => Some(NAValue::Int(self.tb_num as i64)),
242	"tb_den" => Some(NAValue::Int(self.tb_den as i64)),
243	_ => None,
244	}
245	}
246	}
247
248	fn read_pnm_header(file: &mut BufReader<File>, pgmyuv: bool) -> DemuxerResult<NAVideoInfo> {
249	let mut fr = FileReader::new_read(file);
250	let mut br = ByteReader::new(&mut fr);
251
252	let mut magic = [0; 2];
253	br.read_buf(&mut magic)?;
254	if magic[0] != b'P' { return Err(DemuxerError::InvalidData); }
255	match magic[1] {
256	b'4' \| // PBM, PBM ASCII
257	b'1' => return Err(DemuxerError::NotImplemented),
258	b'5' => { // PGM
259	},
260	b'2' => return Err(DemuxerError::NotImplemented), // PGM ASCII
261	b'6' => { // PPM
262	},
263	b'3' => return Err(DemuxerError::NotImplemented), // PPM ASCII
264	_ => return Err(DemuxerError::InvalidData),
265	};
266	if br.read_byte()? != b'\n' { return Err(DemuxerError::InvalidData); }
267	let w = read_number(&mut br)?;
268	let h = read_number(&mut br)?;
269	let maxval = if matches!(magic[1], b'4' \| b'1') { 1 } else { read_number(&mut br)? };
270	if maxval > 65535 \|\| (maxval & (maxval + 1)) != 0 { return Err(DemuxerError::InvalidData); }
271	let bits = maxval.count_ones() as u8;
272
273	let mut vinfo = NAVideoInfo::new(w, h, false, RGB24_FORMAT);
274	match magic[1] {
275	b'5' \| b'2' if !pgmyuv => {
276	vinfo.format = NAPixelFormaton {
277	model: ColorModel::YUV(YUVSubmodel::YUVJ),
278	components: 1,
279	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],
280	elem_size: 1,
281	be: true,
282	alpha: false,
283	palette: false,
284	};
285	vinfo.bits = bits;
286	},
287	b'5' \| b'2' => {
288	if ((w & 1) != 0) \|\| ((h % 3) != 0) { return Err(DemuxerError::InvalidData); }
289	vinfo.format = YUV420_FORMAT;
290	vinfo.height = h * 2 / 3;
291	vinfo.bits = bits * 3 / 2;
292	},
293	b'6' \| b'3' => {
294	vinfo.format = RGB24_FORMAT;
295	vinfo.bits = bits;
296	},
297	_ => unreachable!(),
298	};
299	if bits != 8 {
300	for chr in vinfo.format.comp_info.iter_mut().flatten() {
301	chr.depth = bits;
302	if bits > 8 {
303	chr.next_elem = 2;
304	}
305	}
306	if bits > 8 {
307	vinfo.format.elem_size <<= 1;
308	}
309	}
310
311	Ok(vinfo)
312	}
313
314	fn read_number(br: &mut ByteReader) -> DemuxerResult<usize> {
315	let mut val = 0;
316	loop {
317	let c = br.read_byte()?;
318	match c {
319	b'0'..=b'9' => {
320	if val > 1048576 {
321	return Err(DemuxerError::InvalidData);
322	}
323	val = val * 10 + usize::from(c - b'0');
324	},
325	b' ' \| b'\n' => break,
326	_ => return Err(DemuxerError::InvalidData),
327	};
328	}
329	Ok(val)
330	}