1use std::fs::File;
2use std::io;
3use std::io::{Read, Write, Seek, SeekFrom};
4use std::io::{Error, ErrorKind};
5use super::Bitstream;
6
7pub struct Bin {
8 pub input: String,
9 pub output: String,
10 pub frame_num: usize,
11 pub h265: bool,
12}
13
14impl Bin {
15 pub fn helper() {
16 println!("Usage: rsplit bin input.bin output frame_num h264|h265")
17 }
18
19 pub fn new(args: &[String]) -> Result<Bin, &'static str> {
20 let l = args.len() as usize;
21 if l < 6 {
22 return Err("too less arguments for rsplit bin mode");
23 }
24
25 let input = args[2].clone();
26 let output = args[3].clone();
27 let frame_num_opt = args[4].clone().parse::<usize>();
28 let frame_num = match frame_num_opt {
29 Ok(frame_num) => frame_num,
30 Err(_) => {
31 return Err("can't parse frame_num as usize");
32 }
33 };
34 let h265 = match args[5].to_lowercase().as_ref() {
35 "h265" => true,
36 "h264" => false,
37 _ => {
38 return Err("only support h264 and h265");
39 }
40 };
41
42 Ok(Bin {
43 input: input,
44 output: output,
45 frame_num: frame_num,
46 h265: h265,
47 })
48 }
49
50 pub fn run(&self) -> io::Result<()> {
51 println!("rsplit H26{} {} into {}",
52 4 + (self.h265 as i32),
53 self.input,
54 self.output);
55 let mut fi = try!(File::open(self.input.clone()));
56
57 let mut bk_container: Vec<u8> = Vec::new();
58 let mut bs_container: Vec<Bitstream> = Vec::new();
59 let mut pre_frame_no = 0;
60 let mut cur_frame_no = 0;
61 let mut bak_byte_pos = 0;
62 loop {
63 let (eof, opt) = if self.h265 {
64 self.find_h265_nal_units(&mut fi)
65 }else{
66 self.find_h264_nal_units(&mut fi)
67 };
68
69 let mut bs = match opt {
70 Ok(bs) => bs,
71 Err(_) => {
72 break;
73 }
74 };
75 bs.buf_size = bs.frame_location[bs.nal_size] - bs.frame_location[0];
76
77 if bs.idr_flag {
78 print!("IDR");
79 } else {
80 print!(".");
81 }
82
83 if bs.idr_flag && cur_frame_no - pre_frame_no >= (self.frame_num as i32) {
84 if let Err(e) = self.write_to_file(&mut pre_frame_no,
85 cur_frame_no,
86 bak_byte_pos,
87 &mut bk_container,
88 &mut bs_container) {
89 return Err(e);
90 }
91 bs_container.clear();
92 bak_byte_pos = bk_container.len();
93 }
94
95 if self.h265 {
96 for i in 0..bs.nal_size {
97 if bs.frame_header[i]==32 || bs.frame_header[i]==33 || bs.frame_header[i]==34 {
99 for j in 0..bs.frame_location[i + 1] - bs.frame_location[i] {
100 bk_container.push(bs.frame_data[(bs.frame_location[i] + j) as usize]);
101 }
102 }
103 }
104 }else{
105 for i in 0..bs.nal_size {
106 if bs.frame_header[i]==7 || bs.frame_header[i]==8 {
108 for j in 0..bs.frame_location[i + 1] - bs.frame_location[i] {
109 bk_container.push(bs.frame_data[(bs.frame_location[i] + j) as usize]);
110 }
111 }
112 }
113 }
114
115 bs_container.push(bs);
116
117 if eof {
118 break;
119 }
120
121 cur_frame_no += 1;
122 }
123
124 if let Err(e) = self.write_to_file(&mut pre_frame_no,
125 cur_frame_no,
126 bak_byte_pos,
127 &mut bk_container,
128 &mut bs_container) {
129 return Err(e);
130 }
131
132 Ok(())
133 }
134
135 fn find_h265_nal_units(&self, fp_bs: &mut File) -> (bool, io::Result<Bitstream>) {
136 const START_CODE_SIZE: i32 = 3;
137 const MAX_NAL_UNITS_PER_BS: usize = 600;
138
139 let mut bs = Bitstream {
140 frame_header: vec![0u8; MAX_NAL_UNITS_PER_BS+1],
141 frame_location: vec![0u32; MAX_NAL_UNITS_PER_BS+1],
142 nal_size: 0,
143 frame_data: vec![0u8; 0],
144 buf_size: 0,
145 idr_flag: false,
146 };
147
148 let mut buf = [0u8; 1];
149 let mut frame_data_size = 0;
150 let mut num_nal_units = 0;
151 let mut zeros = 0;
152 let mut pic_found_flag = false;
153 let mut bs_size_since_last_slice = START_CODE_SIZE;
154 let mut num_nal_units_since_last_slice = 0;
155 let mut last_slice_flag = true;
156
157 loop {
158 match fp_bs.read(&mut buf) {
159 Ok(n) => {
160 if n != 1 {
161 bs.frame_location[num_nal_units] = frame_data_size;
162 bs.nal_size = num_nal_units;
163 return (true, Ok(bs));
164 }
165 }
166 Err(e) => {
167 return (false, Err(e));
168 }
169 };
170 bs.frame_data.push(buf[0]);
171 frame_data_size += 1;
172 if !last_slice_flag {
173 bs_size_since_last_slice += 1;
174 }
175
176 match buf[0] {
177 0 => {
178 zeros += 1;
179 }
180
181 1 => {
182 if zeros > 1 {
183 bs.frame_location[num_nal_units] = frame_data_size - zeros - 1;
185
186 match fp_bs.read(&mut buf) {
187 Ok(n) => {
188 if n != 1 {
189 return (true, Ok(bs));
190 }
191 }
192 Err(e) => {
193 return (false, Err(e));
194 }
195 };
196
197 let nal_unit_type = (buf[0] & 0x7E) >> 1;
198
199 bs.frame_header[num_nal_units] = nal_unit_type;
200
201 if nal_unit_type <= 23 {
202 match fp_bs.read(&mut buf) {
204 Ok(n) => {
205 if n != 1 {
206 return (true, Ok(bs));
207 }
208 }
209 Err(e) => {
210 return (false, Err(e));
211 }
212 };
213 match fp_bs.read(&mut buf) {
214 Ok(n) => {
215 if n != 1 {
216 return (true, Ok(bs));
217 }
218 }
219 Err(e) => {
220 return (false, Err(e));
221 }
222 };
223
224 let first_slice_in_pic_flag = (buf[0] >> 7) != 0;
225
226 if first_slice_in_pic_flag {
227 if pic_found_flag {
228 fp_bs.seek(SeekFrom::Current(-3 -
229 (bs_size_since_last_slice as i64)))
230 .unwrap();
231 bs.nal_size = num_nal_units - num_nal_units_since_last_slice;
232
233 return (false, Ok(bs));
234 } else {
235 fp_bs.seek(SeekFrom::Current(-3)).unwrap();
236 num_nal_units += 1;
237 zeros = 0;
238 pic_found_flag = true;
239 bs.idr_flag = nal_unit_type==19 || nal_unit_type == 20 ;
240 }
241 } else {
242 fp_bs.seek(SeekFrom::Current(-3)).unwrap();
243 num_nal_units += 1;
244 zeros = 0;
245 }
246
247 last_slice_flag = true;
248 bs_size_since_last_slice = START_CODE_SIZE;
249 num_nal_units_since_last_slice = 0;
250 } else {
251 if nal_unit_type == 40 {
252 last_slice_flag = true;
254 bs_size_since_last_slice = START_CODE_SIZE;
255 num_nal_units_since_last_slice = 0;
256 } else {
257 last_slice_flag = false;
258 num_nal_units_since_last_slice += 1;
259 }
260
261 fp_bs.seek(SeekFrom::Current(-1)).unwrap();
262 num_nal_units += 1;
263 zeros = 0;
264 }
265 } else {
266 zeros = 0;
267 }
268 }
269
270 _ => {
271 zeros = 0;
272 }
273 }
274 }
275 }
276
277 fn find_h264_nal_units(&self, fp_bs: &mut File) -> (bool, io::Result<Bitstream>) {
278 const START_CODE_SIZE: i32 = 3;
279 const MAX_NAL_UNITS_PER_BS: usize = 600;
280
281 let mut bs = Bitstream {
282 frame_header: vec![0u8; MAX_NAL_UNITS_PER_BS+1],
283 frame_location: vec![0u32; MAX_NAL_UNITS_PER_BS+1],
284 nal_size: 0,
285 frame_data: vec![0u8; 0],
286 buf_size: 0,
287 idr_flag: false,
288 };
289
290 let mut buf = [0u8; 1];
291 let mut frame_data_size = 0;
292 let mut num_nal_units = 0;
293 let mut zeros = 0;
294 let mut pic_found_flag = false;
295 let mut bs_size_since_last_slice = START_CODE_SIZE;
296 let mut num_nal_units_since_last_slice = 0;
297 let mut last_slice_flag = true;
298
299 loop {
300 match fp_bs.read(&mut buf) {
301 Ok(n) => {
302 if n != 1 {
303 bs.frame_location[num_nal_units] = frame_data_size;
304 bs.nal_size = num_nal_units;
305 return (true, Ok(bs));
306 }
307 }
308 Err(e) => {
309 return (false, Err(e));
310 }
311 };
312 bs.frame_data.push(buf[0]);
313 frame_data_size += 1;
314 if !last_slice_flag {
315 bs_size_since_last_slice += 1;
316 }
317
318 match buf[0] {
319 0 => {
320 zeros += 1;
321 }
322
323 1 => {
324 if zeros > 1 {
325 bs.frame_location[num_nal_units] = frame_data_size - zeros - 1;
327
328 match fp_bs.read(&mut buf) {
329 Ok(n) => {
330 if n != 1 {
331 return (true, Ok(bs));
332 }
333 }
334 Err(e) => {
335 return (false, Err(e));
336 }
337 };
338
339 let nal_unit_type = buf[0] & 0x1F;
340
341 bs.frame_header[num_nal_units] = nal_unit_type;
342
343 if nal_unit_type <= 5 {
344 if nal_unit_type == 1 || nal_unit_type == 2 ||nal_unit_type == 5 {
346 match fp_bs.read(&mut buf) {
347 Ok(n) => {
348 if n != 1 {
349 return (true, Ok(bs));
350 }
351 }
352 Err(e) => {
353 return (false, Err(e));
354 }
355 };
356
357 let first_slice_in_pic_flag = (buf[0] >> 7) != 0;
358
359 if first_slice_in_pic_flag {
360 if pic_found_flag {
361 fp_bs.seek(SeekFrom::Current(-2 - (bs_size_since_last_slice as i64))).unwrap();
362 bs.nal_size = num_nal_units - num_nal_units_since_last_slice;
363
364 return (false, Ok(bs));
365 } else {
366 fp_bs.seek(SeekFrom::Current(-2)).unwrap();
367 num_nal_units += 1;
368 zeros = 0;
369 pic_found_flag = true;
370 bs.idr_flag = nal_unit_type==5;
371 }
372 } else {
373 fp_bs.seek(SeekFrom::Current(-2)).unwrap();
374 num_nal_units += 1;
375 zeros = 0;
376 }
377 }else{
378 fp_bs.seek(SeekFrom::Current(-1)).unwrap();
379 num_nal_units += 1;
380 zeros = 0;
381 }
382
383 last_slice_flag = true;
384 bs_size_since_last_slice = START_CODE_SIZE;
385 num_nal_units_since_last_slice = 0;
386 } else {
387 last_slice_flag = false;
388 num_nal_units_since_last_slice += 1;
389
390 fp_bs.seek(SeekFrom::Current(-1)).unwrap();
391 num_nal_units += 1;
392 zeros = 0;
393 }
394 } else {
395 zeros = 0;
396 }
397 }
398
399 _ => {
400 zeros = 0;
401 }
402 }
403 }
404 }
405
406 fn write_to_file(&self,
407 pre_frame_no: &mut i32,
408 cur_frame_no: i32,
409 bak_byte_pos: usize,
410 bk_container: &mut [u8],
411 bs_container: &mut [Bitstream])
412 -> io::Result<()> {
413 let output_bin = format!("{}_{:04}_{:04}.bin",
414 self.output,
415 *pre_frame_no,
416 cur_frame_no - 1);
417
418 println!("\nFrames[{:04}-{:04}] => {}\n",
419 *pre_frame_no,
420 cur_frame_no - 1,
421 output_bin);
422
423 let mut fo = try!(File::create(output_bin));
424 {
425 let bytes_write = fo.write(&bk_container[0..bak_byte_pos]).unwrap();
426 if bytes_write != bak_byte_pos {
427 return Err(Error::new(ErrorKind::Other, "bytes write is not expected ..."));
428 }
429 }
430 for b in bs_container {
431 let bytes_write = fo.write(&b.frame_data[0..(b.buf_size as usize)]).unwrap();
432 if bytes_write != (b.buf_size as usize) {
433 return Err(Error::new(ErrorKind::Other, "bytes write is not expected ..."));
434 }
435 }
436
437 *pre_frame_no = cur_frame_no;
438
439 Ok(())
440 }
441}