| 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 | } |