add PNM image sequence as possible input source
[nihav-encoder.git] / src / demux.rs
CommitLineData
91a15e39
KS
1use std::io::SeekFrom;
2use nihav_core::codecs::*;
3use nihav_core::demuxers::*;
4use nihav_core::muxers::*;
5use nihav_registry::detect;
6use nihav_core::io::byteio::ByteReader;
7use nihav_allstuff::*;
8use crate::null::*;
90683bc3 9use crate::imgseq::*;
91a15e39
KS
10
11pub struct FullRegister {
12 pub dmx_reg: RegisteredDemuxers,
13 pub rdmx_reg: RegisteredRawDemuxers,
14 pub pkt_reg: RegisteredPacketisers,
15 pub dec_reg: RegisteredDecoders,
16 pub enc_reg: RegisteredEncoders,
17 pub mux_reg: RegisteredMuxers,
18}
19
20impl FullRegister {
21 pub fn new() -> Self {
22 let mut dmx_reg = RegisteredDemuxers::new();
23 nihav_register_all_demuxers(&mut dmx_reg);
24 let mut rdmx_reg = RegisteredRawDemuxers::new();
25 nihav_register_all_raw_demuxers(&mut rdmx_reg);
26 let mut dec_reg = RegisteredDecoders::new();
27 nihav_register_all_decoders(&mut dec_reg);
28 let mut pkt_reg = RegisteredPacketisers::new();
29 nihav_register_all_packetisers(&mut pkt_reg);
30 let mut enc_reg = RegisteredEncoders::new();
31 nihav_register_all_encoders(&mut enc_reg);
32 enc_reg.add_encoder(NULL_ENCODER);
33 let mut mux_reg = RegisteredMuxers::new();
34 nihav_register_all_muxers(&mut mux_reg);
35 mux_reg.add_muxer(NULL_MUXER);
36 Self { dmx_reg, rdmx_reg, pkt_reg, dec_reg, enc_reg, mux_reg }
37 }
38}
39
40pub struct RawStreamCtx<'a> {
41 stream: NAStreamRef,
42 sm: StreamManager,
43 packetiser: Box<dyn NAPacketiser + Send>,
44 br: &'a mut ByteReader<'a>,
206dd66d 45 pts: u64,
91a15e39
KS
46}
47
48impl<'a> RawStreamCtx<'a> {
49 fn new(stream: NAStreamRef, packetiser: Box<dyn NAPacketiser + Send>, br: &'a mut ByteReader<'a>) -> Self {
50 let mut sm = StreamManager::new();
51 sm.add_stream_ref(stream.clone());
206dd66d 52 Self { stream, sm, packetiser, br, pts: 0 }
91a15e39
KS
53 }
54}
55
56pub enum DemuxerObject<'a> {
57 None,
58 Normal(Demuxer<'a>),
59 Raw(RawDemuxer<'a>, Vec<Option<Box<dyn NAPacketiser + Send>>>, bool),
60 RawStream(RawStreamCtx<'a>),
90683bc3 61 ImageSequence(ImgSeqDemuxer),
91a15e39
KS
62}
63
64impl<'a> DemuxerObject<'a> {
65 pub fn create(br: &'a mut ByteReader<'a>, reg: &FullRegister, name: &str, ifmt: &Option<String>, is_raw: bool) -> DemuxerObject<'a> {
66 if !is_raw {
67 let dmx_name = if let Some(ref dname) = ifmt {
68 dname.as_str()
69 } else if let Some((dname, score)) = detect::detect_format(name, br) {
70 println!("detected {} with score {:?}", dname, score);
71 dname
72 } else {
73 ""
74 };
76ea2e41 75 if !dmx_name.is_empty() {
91a15e39
KS
76 println!("trying demuxer {} on {}", dmx_name, name);
77 if let Some(dmx_fact) = reg.dmx_reg.find_demuxer(dmx_name) {
78 br.seek(SeekFrom::Start(0)).unwrap();
79 let dmx = create_demuxer(dmx_fact, br).unwrap();
80 return DemuxerObject::Normal(dmx);
81 }
82 }
83 if ifmt.is_some() {
84 return DemuxerObject::None;
85 }
76ea2e41 86 if !dmx_name.is_empty() {
91a15e39
KS
87 println!("trying raw demuxer {} on {}", dmx_name, name);
88 if let Some(rdmx_fact) = reg.rdmx_reg.find_demuxer(dmx_name) {
89 br.seek(SeekFrom::Start(0)).unwrap();
90 let dmx = create_raw_demuxer(rdmx_fact, br).unwrap();
91 let mut pkts = Vec::new();
92 for stream in dmx.get_streams() {
93 if let Some(pcreate) = reg.pkt_reg.find_packetiser(stream.get_info().get_name()) {
94 let packetiser = (pcreate)();
95 pkts.push(Some(packetiser));
96 } else {
97 pkts.push(None);
98 }
99 }
100 return DemuxerObject::Raw(dmx, pkts, false);
101 }
102 }
103 for rdmx in reg.rdmx_reg.iter() {
104 if rdmx.check_format(br) {
105 println!("detected {} as {}", name, rdmx.get_name());
106 br.seek(SeekFrom::Start(0)).unwrap();
107 let dmx = create_raw_demuxer(*rdmx, br).unwrap();
108 let mut pkts = Vec::new();
109 for stream in dmx.get_streams() {
110 if let Some(pcreate) = reg.pkt_reg.find_packetiser(stream.get_info().get_name()) {
111 let packetiser = (pcreate)();
112 pkts.push(Some(packetiser));
113 } else {
114 pkts.push(None);
115 }
116 }
117 return DemuxerObject::Raw(dmx, pkts, false);
118 }
119 }
120 }
121 br.seek(SeekFrom::Start(0)).unwrap();
122 let mut buf = vec![0; 1048576];
123 let size = br.peek_buf(&mut buf).unwrap();
124 let mut pname = "";
125
126 for pinfo in reg.pkt_reg.iter() {
127 let mut packetiser = (pinfo.get_packetiser)();
128 packetiser.add_data(&buf[..size]);
129 if packetiser.parse_stream(0).is_ok() {
130 pname = pinfo.name;
131 break;
132 }
133 }
76ea2e41 134 if !pname.is_empty() {
91a15e39
KS
135 println!("found raw stream of type {} for {}", pname, name);
136 let pcreate = reg.pkt_reg.find_packetiser(pname).unwrap();
137 let mut packetiser = (pcreate)();
138 packetiser.add_data(&buf[..size]);
139 let stream = packetiser.parse_stream(0).unwrap();
140 packetiser.reset();
141 DemuxerObject::RawStream(RawStreamCtx::new(stream, packetiser, br))
142 } else {
143 DemuxerObject::None
144 }
145 }
90683bc3
KS
146 pub fn create_imgseq(isd: ImgSeqDemuxer) -> Self {
147 DemuxerObject::ImageSequence(isd)
148 }
91a15e39 149 pub fn is_none(&self) -> bool {
76ea2e41 150 matches!(*self, DemuxerObject::None)
91a15e39 151 }
5ec8115f 152 pub fn get_duration(&self) -> u64 {
91a15e39
KS
153 match *self {
154 DemuxerObject::Normal(ref dmx) => dmx.get_duration(),
155 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_duration(),
4dd9047c
KS
156 DemuxerObject::RawStream(ref ctx) => {
157 NATimeInfo::ts_to_time(ctx.stream.duration, 1000, ctx.stream.tb_num, ctx.stream.tb_den)
158 },
91a15e39
KS
159 _ => 0,
160 }
5ec8115f 161 }
91a15e39
KS
162 pub fn get_num_streams(&self) -> usize {
163 match *self {
164 DemuxerObject::None => 0,
165 DemuxerObject::Normal(ref dmx) => dmx.get_num_streams(),
166 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_num_streams(),
167 DemuxerObject::RawStream(_) => 1,
90683bc3 168 DemuxerObject::ImageSequence(_) => 1,
91a15e39
KS
169 }
170 }
171 pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
172 match *self {
173 DemuxerObject::Normal(ref dmx) => dmx.get_stream(idx),
174 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_stream(idx),
175 DemuxerObject::RawStream(ref ctx) if idx == 0 => Some(ctx.stream.clone()),
90683bc3 176 DemuxerObject::ImageSequence(ref ctx) if idx == 0 => Some(ctx.stream.clone()),
91a15e39
KS
177 _ => None,
178 }
179 }
180 pub fn get_streams(&self) -> StreamIter {
181 match *self {
182 DemuxerObject::Normal(ref dmx) => dmx.get_streams(),
183 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_streams(),
184 DemuxerObject::RawStream(ref ctx) => ctx.sm.iter(),
90683bc3 185 DemuxerObject::ImageSequence(ref ctx) => ctx.sm.iter(),
91a15e39
KS
186 _ => unreachable!(),
187 }
188 }
189 pub fn get_stream_manager(&self) -> &StreamManager {
190 match *self {
191 DemuxerObject::Normal(ref dmx) => dmx.get_stream_manager(),
192 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_stream_manager(),
193 DemuxerObject::RawStream(ref ctx) => &ctx.sm,
90683bc3 194 DemuxerObject::ImageSequence(ref ctx) => &ctx.sm,
91a15e39
KS
195 _ => unreachable!(),
196 }
197 }
198 pub fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
199 match *self {
200 DemuxerObject::Normal(ref mut dmx) => dmx.get_frame(),
201 DemuxerObject::Raw(ref mut dmx, ref mut packetisers, ref mut eof) => {
202 loop {
203 let mut has_some = false;
204 for (stream, p) in dmx.get_streams().zip(packetisers.iter_mut()) {
205 if let Some(ref mut pkts) = p {
206 match pkts.get_packet(stream.clone()) {
207 Ok(Some(pkt)) => return Ok(pkt),
208 Ok(None) | Err(DecoderError::ShortData) => {
209 if *eof {
210 *p = None;
211 }
212 },
213 Err(err) => {
214 println!("packetisation error {:?}", err);
215 return Err(DemuxerError::InvalidData);
216 }
217 };
218 has_some |= p.is_some();
219 }
220 }
221 if !has_some {
222 return Err(DemuxerError::EOF);
223 }
224 if let Ok(data) = dmx.get_data() {
225 let id = data.get_stream().get_id();
226 for (i, stream) in dmx.get_streams().enumerate() {
227 if stream.get_id() == id {
228 if let Some(ref mut pkts) = packetisers[i] {
229 pkts.add_data(&data.get_buffer());
230 }
231 break;
232 }
233 }
234 } else {
235 *eof = true;
236 }
237 }
238 },
239 DemuxerObject::RawStream(ref mut ctx) => {
240 let mut buf = [0; 65536];
241 loop {
242 match ctx.packetiser.get_packet(ctx.stream.clone()) {
206dd66d
KS
243 Ok(Some(mut packet)) => {
244 if packet.get_pts().is_none() && packet.get_duration().is_some() {
245 packet.ts.pts = Some(ctx.pts);
246 }
247 ctx.pts += packet.get_duration().unwrap_or(0);
248 return Ok(packet);
249 },
91a15e39
KS
250 Ok(None) => {},
251 Err(DecoderError::ShortData) => {},
252 _ => return Err(DemuxerError::InvalidData),
253 };
254 match ctx.br.read_buf_some(&mut buf) {
255 Ok(size) => {
256 ctx.packetiser.add_data(&buf[..size]);
257 },
258 Err(_) => {
259 match ctx.packetiser.get_packet(ctx.stream.clone()) {
206dd66d
KS
260 Ok(Some(mut packet)) => {
261 if packet.get_pts().is_none() && packet.get_duration().is_some() {
262 packet.ts.pts = Some(ctx.pts);
263 }
264 ctx.pts += packet.get_duration().unwrap_or(0);
265 return Ok(packet);
266 },
91a15e39
KS
267 Ok(None) | Err(DecoderError::ShortData) => return Err(DemuxerError::EOF),
268 _ => return Err(DemuxerError::InvalidData),
269 };
270 },
271 };
272 }
273 },
90683bc3 274 DemuxerObject::ImageSequence(ref mut ctx) => ctx.get_frame(),
91a15e39
KS
275 _ => unreachable!(),
276 }
277 }
278 pub fn seek(&mut self, seek_time: NATimePoint) -> DemuxerResult<()> {
279 match *self {
280 DemuxerObject::Normal(ref mut dmx) => dmx.seek(seek_time),
281 DemuxerObject::Raw(ref mut dmx, _, _) => dmx.seek(seek_time),
90683bc3 282 DemuxerObject::ImageSequence(ref mut ctx) => ctx.seek(seek_time),
91a15e39
KS
283 _ => Err(DemuxerError::NotImplemented),
284 }
285 }
286}
287
288impl<'a> NAOptionHandler for DemuxerObject<'a> {
289 fn get_supported_options(&self) -> &[NAOptionDefinition] {
290 match *self {
291 DemuxerObject::Normal(ref dmx) => dmx.get_supported_options(),
292 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_supported_options(),
90683bc3 293 DemuxerObject::ImageSequence(ref ctx) => ctx.get_supported_options(),
91a15e39
KS
294 _ => &[],
295 }
296 }
297 fn set_options(&mut self, options: &[NAOption]) {
298 match *self {
299 DemuxerObject::Normal(ref mut dmx) => dmx.set_options(options),
300 DemuxerObject::Raw(ref mut dmx, _, _) => dmx.set_options(options),
90683bc3 301 DemuxerObject::ImageSequence(ref mut ctx) => ctx.set_options(options),
91a15e39
KS
302 _ => {},
303 }
304 }
305 fn query_option_value(&self, name: &str) -> Option<NAValue> {
306 match *self {
307 DemuxerObject::Normal(ref dmx) => dmx.query_option_value(name),
308 DemuxerObject::Raw(ref dmx, _, _) => dmx.query_option_value(name),
90683bc3 309 DemuxerObject::ImageSequence(ref ctx) => ctx.query_option_value(name),
91a15e39
KS
310 _ => None,
311 }
312 }
313}
314
315pub fn detect_tags(br: &mut ByteReader) -> (bool, u64, Option<u64>) {
316 let mut is_raw = false;
317 let mut start = 0;
318 let mut end = None;
319
320 // check for ID3v{2-4}
321 let mut buf = [0; 5];
322 br.peek_buf(&mut buf).unwrap();
323 if &buf[0..3] == b"ID3" && buf[3] > 0 && buf[3] < 5 && buf[4] == 0 { //ID3 tag found, must be a raw stream
324 br.read_skip(6).unwrap();
325 let mut size = 0;
326 for _ in 0..4 {
327 let b = br.read_byte().unwrap();
328 if (b & 0x80) != 0 {
329 println!("Invalid ID3 size");
330 break;
331 }
332 size = (size << 7) | u64::from(b);
333 }
334 start = size + 10;
335 is_raw = true;
336 }
337 // check for ID3v1
338 br.seek(SeekFrom::End(-128)).unwrap();
339 let off = br.tell();
340 br.peek_buf(&mut buf[..3]).unwrap();
341 if &buf[0..3] == b"TAG" {
342 end = Some(off);
343 }
344 // check for APETAG
345 let mut buf = [0; 8];
346 if let Some(off) = end {
347 br.seek(SeekFrom::Start(off - 32)).unwrap();
348 } else {
349 br.seek(SeekFrom::End(-32)).unwrap();
350 }
351 let off = br.tell();
352 br.read_buf(&mut buf).unwrap();
353 if &buf == b"APETAGEX" {
354 let ver = br.read_u32le().unwrap();
355 let size = u64::from(br.read_u32le().unwrap());
356 let _items = br.read_u32le().unwrap();
357 let flags = br.read_u32le().unwrap();
358 if ver == 1000 || (flags & 0x80000000) == 0 {
359 end = Some(off - size + 32);
360 } else {
361 end = Some(off - size);
362 }
363 }
364
365 (is_raw, start, end)
366}