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