]>
Commit | Line | Data |
---|---|---|
696e4e20 KS |
1 | use nihav_core::frame::*; |
2 | use nihav_core::io::byteio::*; | |
3 | use nihav_core::demuxers::*; | |
4 | use nihav_core::formats::YUV420_FORMAT; | |
5 | ||
6 | struct RawH264Demuxer<'a> { | |
7 | src: &'a mut ByteReader<'a>, | |
8 | cur_frame: u64, | |
9 | frame_buf: Vec<u8>, | |
10 | } | |
11 | ||
12 | fn read_nal(src: &mut ByteReader, dst: &mut Vec<u8>) -> DemuxerResult<()> { | |
37952415 | 13 | dst.clear(); |
696e4e20 KS |
14 | loop { |
15 | let b = src.read_byte()?; | |
16 | if b == 0 { | |
17 | let b2 = src.read_byte()?; | |
18 | if b2 == 0 { | |
19 | let b3 = src.read_byte()?; | |
20 | if b3 == 0 { | |
21 | while src.read_byte()? != 1 { } | |
22 | break; | |
23 | } else if b3 == 1 { | |
24 | break; | |
25 | } else { | |
26 | dst.push(b); | |
27 | dst.push(b2); | |
28 | dst.push(b3); | |
29 | } | |
30 | } else { | |
31 | dst.push(b); | |
32 | dst.push(b2); | |
33 | } | |
34 | } else { | |
35 | dst.push(b); | |
36 | } | |
37 | } | |
38 | Ok(()) | |
39 | } | |
40 | ||
41 | fn put_nal(dst: &mut Vec<u8>, src: &[u8]) { | |
42 | let len = src.len(); | |
43 | dst.push((len >> 24) as u8); | |
44 | dst.push((len >> 16) as u8); | |
45 | dst.push((len >> 8) as u8); | |
46 | dst.push( len as u8); | |
47 | dst.extend_from_slice(src); | |
48 | } | |
49 | ||
50 | impl<'a> DemuxCore<'a> for RawH264Demuxer<'a> { | |
51 | #[allow(clippy::unreadable_literal)] | |
52 | fn open(&mut self, strmgr: &mut StreamManager, _seek_idx: &mut SeekIndex) -> DemuxerResult<()> { | |
53 | let src = &mut self.src; | |
54 | ||
55 | while src.read_byte()? != 1 { } | |
56 | ||
57 | let mut edata: Vec<u8> = Vec::with_capacity(64); | |
58 | let mut num_sps = 0; | |
59 | let mut sps_buf = Vec::new(); | |
60 | let mut num_pps = 0; | |
61 | let mut pps_buf = Vec::new(); | |
62 | let mut profile = 0; | |
63 | let mut level = 0; | |
64 | ||
65 | let mut nal_buf = Vec::with_capacity(65536); | |
66 | loop { | |
67 | read_nal(src, &mut nal_buf)?; | |
68 | if !nal_buf.is_empty() { | |
69 | let nal_type = nal_buf[0] & 0x1F; | |
70 | match nal_type { | |
71 | 7 => { | |
72 | if nal_buf.len() < 4 { | |
73 | return Err(DemuxerError::InvalidData); | |
74 | } | |
75 | profile = nal_buf[1]; | |
76 | level = nal_buf[3]; | |
77 | sps_buf.push((nal_buf.len() >> 8) as u8); | |
78 | sps_buf.push(nal_buf.len() as u8); | |
79 | sps_buf.extend_from_slice(&nal_buf); | |
80 | num_sps += 1; | |
81 | }, | |
82 | 8 => { | |
83 | pps_buf.push((nal_buf.len() >> 8) as u8); | |
84 | pps_buf.push(nal_buf.len() as u8); | |
85 | pps_buf.extend_from_slice(&nal_buf); | |
86 | num_pps += 1; | |
87 | }, | |
88 | 1 | 5 => { | |
89 | self.frame_buf = nal_buf; | |
90 | break; | |
91 | }, | |
92 | _ => {}, | |
93 | }; | |
94 | } | |
95 | } | |
96 | if num_sps == 0 || num_pps == 0 { | |
97 | return Err(DemuxerError::InvalidData); | |
98 | } | |
99 | edata.extend_from_slice(b"avcC"); | |
100 | edata.push(1); | |
101 | edata.push(profile); | |
102 | edata.push(0); | |
103 | edata.push(level); | |
104 | edata.push(0xFF); | |
105 | edata.push(0xE0 | num_sps); | |
106 | edata.extend_from_slice(&sps_buf); | |
107 | edata.push(num_pps); | |
108 | edata.extend_from_slice(&pps_buf); | |
109 | ||
110 | let width = 16; | |
111 | let height = 16; | |
112 | ||
113 | let vhdr = NAVideoInfo::new(width, height, false, YUV420_FORMAT); | |
114 | let vinfo = NACodecInfo::new("h264", NACodecTypeInfo::Video(vhdr), Some(edata)); | |
115 | if let None = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, 1, 25, 0)) { | |
116 | return Err(DemuxerError::InvalidData); | |
117 | } | |
118 | self.cur_frame = 0; | |
119 | ||
120 | Ok(()) | |
121 | } | |
122 | fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> { | |
123 | let ts = NATimeInfo::new(Some(self.cur_frame), None, None, 1, 25); | |
124 | let mut buf: Vec<u8> = Vec::with_capacity(65536); | |
125 | if !self.frame_buf.is_empty() { | |
126 | put_nal(&mut buf, &self.frame_buf); | |
37952415 | 127 | self.frame_buf.clear(); |
696e4e20 KS |
128 | } |
129 | let strres = strmgr.get_stream(0); | |
130 | if strres.is_none() { | |
131 | return Err(DemuxerError::InvalidData); | |
132 | } | |
133 | let stream = strres.unwrap(); | |
134 | ||
135 | let mut keyframe = false; | |
136 | ||
137 | let mut nal_buf = Vec::with_capacity(65536); | |
138 | let mut is_eof = false; | |
139 | while !is_eof { | |
140 | match read_nal(&mut self.src, &mut nal_buf) { | |
141 | Ok(()) => {}, | |
142 | Err(DemuxerError::IOError) => { is_eof = true; }, | |
143 | Err(err) => return Err(err), | |
144 | }; | |
145 | if !nal_buf.is_empty() { | |
146 | let nal_type = nal_buf[0] & 0x1F; | |
147 | keyframe = nal_type == 5; | |
148 | match nal_type { | |
149 | 1 | 5 => { | |
150 | let first_slice = (nal_buf[1] & 0x80) != 0; | |
151 | if first_slice && !buf.is_empty() { | |
152 | self.frame_buf.extend_from_slice(&nal_buf); | |
153 | break; | |
154 | } else { | |
155 | put_nal(&mut buf, &nal_buf); | |
156 | } | |
157 | }, | |
158 | _ => { | |
159 | //println!("non-slice NAL {} @ {:X}", nal_type, self.src.tell()); | |
160 | put_nal(&mut buf, &nal_buf); | |
161 | }, | |
162 | }; | |
163 | } | |
164 | } | |
165 | ||
166 | if is_eof && buf.is_empty() { | |
167 | return Err(DemuxerError::EOF); | |
168 | } | |
169 | ||
170 | let pkt = NAPacket::new(stream, ts, keyframe, buf); | |
171 | ||
172 | self.cur_frame += 1; | |
173 | ||
174 | Ok(pkt) | |
175 | } | |
176 | fn seek(&mut self, _time: NATimePoint, _seek_idx: &SeekIndex) -> DemuxerResult<()> { | |
177 | Err(DemuxerError::NotImplemented) | |
178 | } | |
179 | fn get_duration(&self) -> u64 { 0 } | |
180 | } | |
181 | ||
182 | impl<'a> NAOptionHandler for RawH264Demuxer<'a> { | |
183 | fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] } | |
184 | fn set_options(&mut self, _options: &[NAOption]) { } | |
185 | fn query_option_value(&self, _name: &str) -> Option<NAValue> { None } | |
186 | } | |
187 | ||
188 | impl<'a> RawH264Demuxer<'a> { | |
189 | fn new(io: &'a mut ByteReader<'a>) -> Self { | |
190 | RawH264Demuxer { | |
191 | src: io, | |
192 | cur_frame: 0, | |
193 | frame_buf: Vec::new(), | |
194 | } | |
195 | } | |
196 | } | |
197 | ||
198 | pub struct RawH264DemuxerCreator { } | |
199 | ||
200 | impl DemuxerCreator for RawH264DemuxerCreator { | |
201 | fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> { | |
202 | Box::new(RawH264Demuxer::new(br)) | |
203 | } | |
204 | fn get_name(&self) -> &'static str { "rawh264" } | |
205 | } | |
206 |