]>
Commit | Line | Data |
---|---|---|
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("rawvideo", 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("rawvideo", 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 | } |