]> git.nihav.org Git - nihav.git/blob - nihav-itu/src/codecs/h264/test/raw_demux.rs
replace vec.truncate(0) with vec.clear()
[nihav.git] / nihav-itu / src / codecs / h264 / test / raw_demux.rs
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<()> {
13 dst.clear();
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);
127 self.frame_buf.clear();
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