y4m.rs

   1 use nihav_core::demuxers::*;
   2 use std::str::FromStr;
   3
   4 struct Y4MDemuxer<'a> {
   5     src:        &'a mut ByteReader<'a>,
   6     width:      usize,
   7     height:     usize,
   8     frame_size: usize,
   9     hdr_size:   u64,
  10     fps_num:    u32,
  11     fps_den:    u32,
  12     frameno:    u64,
  13 }
  14
  15 impl<'a> DemuxCore<'a> for Y4MDemuxer<'a> {
  16     fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
  17         let format = self.parse_header()?;
  18         seek_index.mode = SeekIndexMode::Automatic;
  19
  20         let vhdr = NAVideoInfo::new(self.width, self.height, false, format);
  21         let vci = NACodecTypeInfo::Video(vhdr);
  22         let vinfo = NACodecInfo::new("rawvideo", vci, None);
  23         if let None = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, self.fps_num, self.fps_den, 0)) {
  24             return Err(DemuxerError::MemoryError);
  25         }
  26
  27         Ok(())
  28     }
  29
  30     fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
  31         let mut marker = [0u8; 6];
  32         let res                         = self.src.read_buf(&mut marker);
  33         match res {
  34             Err(ByteIOError::EOF) => return Err(DemuxerError::EOF),
  35             Err(err) => return Err(err.into()),
  36             _ => {},
  37         };
  38         validate!(&marker == b"FRAME\n");
  39         let stream = strmgr.get_stream(0).unwrap();
  40         let (tb_num, tb_den) = stream.get_timebase();
  41         let ts = NATimeInfo::new(Some(self.frameno), None, None, tb_num, tb_den);
  42         let pkt = self.src.read_packet(stream, ts, true, self.frame_size)?;
  43         self.frameno += 1;
  44         Ok(pkt)
  45     }
  46
  47     fn seek(&mut self, time: NATimePoint, _seek_index: &SeekIndex) -> DemuxerResult<()> {
  48         let new_fno = match time {
  49                 NATimePoint::PTS(pts) => {
  50                     pts
  51                 },
  52                 NATimePoint::Milliseconds(ms) => {
  53                     if (self.fps_num == 0) || (self.fps_den == 0) {
  54                         return Err(DemuxerError::SeekError);
  55                     }
  56                     NATimeInfo::time_to_ts(ms, 1000, self.fps_num, self.fps_den)
  57                 },
  58                 NATimePoint::None => return Err(DemuxerError::SeekError),
  59             };
  60         let pos = self.hdr_size + new_fno * ((self.frame_size + 6) as u64);
  61         self.src.seek(SeekFrom::Start(pos))?;
  62         self.frameno = new_fno;
  63
  64         Ok(())
  65     }
  66     fn get_duration(&self) -> u64 { 0 }
  67 }
  68
  69 impl<'a> NAOptionHandler for Y4MDemuxer<'a> {
  70     fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
  71     fn set_options(&mut self, _options: &[NAOption]) { }
  72     fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
  73 }
  74
  75 impl<'a> Y4MDemuxer<'a> {
  76     fn new(io: &'a mut ByteReader<'a>) -> Self {
  77         Self {
  78             src:        io,
  79             width:      0,
  80             height:     0,
  81             frame_size: 0,
  82             fps_num:    0,
  83             fps_den:    0,
  84             hdr_size:   0,
  85             frameno:    0,
  86         }
  87     }
  88     fn parse_header(&mut self) -> DemuxerResult<NAPixelFormaton> {
  89         let mut format = RGB24_FORMAT;
  90
  91         let mut magic = [0u8; 10];
  92                                     self.src.read_buf(&mut magic)?;
  93         validate!(&magic == b"YUV4MPEG2 ");
  94         while let Ok((last, tok)) = read_token(&mut self.src) {
  95             let (id, val) = tok.split_at(1);
  96             validate!(id.len() == 1);
  97             match id.bytes().next().unwrap() {
  98                 b'W' => {
  99                     if let Ok(w) = val.parse::<usize>() {
 100                         self.width = w;
 101                     }
 102                 },
 103                 b'H' => {
 104                     if let Ok(h) = val.parse::<usize>() {
 105                         self.height = h;
 106                     }
 107                 },
 108                 b'F' => {
 109                     if let Ok(fden) = val.parse::<u32>() {
 110                         self.fps_num = 1;
 111                         self.fps_den = fden;
 112                     } else {
 113                         let vals: Vec<&str> = val.split(':').collect();
 114                         if vals.len() == 2 {
 115                             if let Ok(fnum) = vals[1].parse::<u32>() {
 116                                 self.fps_num = fnum;
 117                             }
 118                             if let Ok(fden) = vals[0].parse::<u32>() {
 119                                 self.fps_den = fden;
 120                             }
 121                         }
 122                     }
 123                 },
 124                 b'C' => {
 125                     let fmt_str = val.as_bytes();
 126                     validate!(fmt_str.len() >= 3);
 127                     let mut pix_name: [u8; 7] = *b"yuv000p";
 128                     validate!(fmt_str[0] == b'4');
 129                     pix_name[3..6].copy_from_slice(&fmt_str[..3]);
 130
 131                     if let Ok(fmt_name) = std::str::from_utf8(&pix_name) {
 132                         if let Ok(val) = NAPixelFormaton::from_str(fmt_name) {
 133                             format = val;
 134                         }
 135                     }
 136                     if format.model.is_yuv() {
 137                         format.model = ColorModel::YUV(YUVSubmodel::YCbCr);
 138                         if fmt_str.len() > 3 {
 139                             let (_, tail) = fmt_str.split_at(3);
 140                             if tail == b"jpeg" {
 141                                 format.model = ColorModel::YUV(YUVSubmodel::YUVJ);
 142                             }
 143                         }
 144                     }
 145                 },
 146                 _ => {},
 147             };
 148
 149             if last {
 150                 break;
 151             }
 152         }
 153         validate!(self.width > 0 && self.height > 0 && format.model.is_yuv());
 154         self.frame_size = 0;
 155         for chromaton in format.comp_info.iter() {
 156             if let Some(ref chr) = chromaton {
 157                 self.frame_size += chr.get_data_size(self.width, self.height);
 158             }
 159         }
 160         validate!(self.frame_size > 0);
 161
 162         Ok(format)
 163     }
 164 }
 165
 166 fn read_token(src: &mut ByteReader) -> DemuxerResult<(bool, String)> {
 167     let mut string = String::new();
 168     let ws;
 169     loop {
 170         let b                       = src.read_byte()?;
 171         match b {
 172             b' ' | b'\n' => { ws = b; break; },
 173             0..=0x7F => string.push(b as char),
 174             _ => return Err(DemuxerError::InvalidData),
 175         }
 176     }
 177
 178     Ok((ws == b'\n', string))
 179 }
 180
 181 pub struct Y4MDemuxerCreator { }
 182
 183 impl DemuxerCreator for Y4MDemuxerCreator {
 184     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
 185         Box::new(Y4MDemuxer::new(br))
 186     }
 187     fn get_name(&self) -> &'static str { "yuv4mpeg" }
 188 }
 189
 190 #[cfg(test)]
 191 mod test {
 192     use super::*;
 193     use std::fs::File;
 194
 195     #[test]
 196     fn test_y4m_demux() {
 197         let mut file = File::open("assets/Misc/test.y4m").unwrap();
 198         let mut fr = FileReader::new_read(&mut file);
 199         let mut br = ByteReader::new(&mut fr);
 200         let mut dmx = Y4MDemuxer::new(&mut br);
 201         let mut sm = StreamManager::new();
 202         let mut si = SeekIndex::new();
 203         dmx.open(&mut sm, &mut si).unwrap();
 204
 205         loop {
 206             let pktres = dmx.get_frame(&mut sm);
 207             if let Err(e) = pktres {
 208                 if e == DemuxerError::EOF { break; }
 209                 panic!("error");
 210             }
 211             let pkt = pktres.unwrap();
 212             println!("Got {}", pkt);
 213         }
 214     }
 215 }