--- /dev/null
+use nihav_core::demuxers::*;
+use std::str::FromStr;
+
+struct Y4MDemuxer<'a> {
+ src: &'a mut ByteReader<'a>,
+ width: usize,
+ height: usize,
+ frame_size: usize,
+ hdr_size: u64,
+ fps_num: u32,
+ fps_den: u32,
+ frameno: u64,
+}
+
+impl<'a> DemuxCore<'a> for Y4MDemuxer<'a> {
+ fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
+ let format = self.parse_header()?;
+ seek_index.mode = SeekIndexMode::Automatic;
+
+ let vhdr = NAVideoInfo::new(self.width, self.height, false, format);
+ let vci = NACodecTypeInfo::Video(vhdr);
+ let vinfo = NACodecInfo::new("rawvideo", vci, None);
+ if let None = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, self.fps_num, self.fps_den, 0)) {
+ return Err(DemuxerError::MemoryError);
+ }
+
+ Ok(())
+ }
+
+ fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
+ let mut marker = [0u8; 6];
+ let res = self.src.read_buf(&mut marker);
+ match res {
+ Err(ByteIOError::EOF) => return Err(DemuxerError::EOF),
+ Err(err) => return Err(err.into()),
+ _ => {},
+ };
+ validate!(&marker == b"FRAME\n");
+ let stream = strmgr.get_stream(0).unwrap();
+ let (tb_num, tb_den) = stream.get_timebase();
+ let ts = NATimeInfo::new(Some(self.frameno), None, None, tb_num, tb_den);
+ let pkt = self.src.read_packet(stream, ts, true, self.frame_size)?;
+ self.frameno += 1;
+ Ok(pkt)
+ }
+
+ fn seek(&mut self, time: NATimePoint, _seek_index: &SeekIndex) -> DemuxerResult<()> {
+ let new_fno = match time {
+ NATimePoint::PTS(pts) => {
+ pts
+ },
+ NATimePoint::Milliseconds(ms) => {
+ if (self.fps_num == 0) || (self.fps_den == 0) {
+ return Err(DemuxerError::SeekError);
+ }
+ NATimeInfo::time_to_ts(ms, 1000, self.fps_num, self.fps_den)
+ },
+ NATimePoint::None => return Err(DemuxerError::SeekError),
+ };
+ let pos = self.hdr_size + new_fno * ((self.frame_size + 6) as u64);
+ self.src.seek(SeekFrom::Start(pos))?;
+ self.frameno = new_fno;
+
+ Ok(())
+ }
+ fn get_duration(&self) -> u64 { 0 }
+}
+
+impl<'a> NAOptionHandler for Y4MDemuxer<'a> {
+ fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
+ fn set_options(&mut self, _options: &[NAOption]) { }
+ fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
+}
+
+impl<'a> Y4MDemuxer<'a> {
+ fn new(io: &'a mut ByteReader<'a>) -> Self {
+ Self {
+ src: io,
+ width: 0,
+ height: 0,
+ frame_size: 0,
+ fps_num: 0,
+ fps_den: 0,
+ hdr_size: 0,
+ frameno: 0,
+ }
+ }
+ fn parse_header(&mut self) -> DemuxerResult<NAPixelFormaton> {
+ let mut format = RGB24_FORMAT;
+
+ let mut magic = [0u8; 10];
+ self.src.read_buf(&mut magic)?;
+ validate!(&magic == b"YUV4MPEG2 ");
+ while let Ok((last, tok)) = read_token(&mut self.src) {
+ let (id, val) = tok.split_at(1);
+ validate!(id.len() == 1);
+ match id.bytes().next().unwrap() {
+ b'W' => {
+ if let Ok(w) = val.parse::<usize>() {
+ self.width = w;
+ }
+ },
+ b'H' => {
+ if let Ok(h) = val.parse::<usize>() {
+ self.height = h;
+ }
+ },
+ b'F' => {
+ if let Ok(fden) = val.parse::<u32>() {
+ self.fps_num = 1;
+ self.fps_den = fden;
+ } else {
+ let vals: Vec<&str> = val.split(':').collect();
+ if vals.len() == 2 {
+ if let Ok(fnum) = vals[1].parse::<u32>() {
+ self.fps_num = fnum;
+ }
+ if let Ok(fden) = vals[0].parse::<u32>() {
+ self.fps_den = fden;
+ }
+ }
+ }
+ },
+ b'C' => {
+ let fmt_str = val.as_bytes();
+ validate!(fmt_str.len() >= 3);
+ let mut pix_name: [u8; 7] = *b"yuv000p";
+ validate!(fmt_str[0] == b'4');
+ pix_name[3..6].copy_from_slice(&fmt_str[..3]);
+
+ if let Ok(fmt_name) = std::str::from_utf8(&pix_name) {
+ if let Ok(val) = NAPixelFormaton::from_str(fmt_name) {
+ format = val;
+ }
+ }
+ if format.model.is_yuv() {
+ format.model = ColorModel::YUV(YUVSubmodel::YCbCr);
+ if fmt_str.len() > 3 {
+ let (_, tail) = fmt_str.split_at(3);
+ if tail == b"jpeg" {
+ format.model = ColorModel::YUV(YUVSubmodel::YUVJ);
+ }
+ }
+ }
+ },
+ _ => {},
+ };
+
+ if last {
+ break;
+ }
+ }
+ validate!(self.width > 0 && self.height > 0 && format.model.is_yuv());
+ self.frame_size = 0;
+ for chromaton in format.comp_info.iter() {
+ if let Some(ref chr) = chromaton {
+ self.frame_size += chr.get_data_size(self.width, self.height);
+ }
+ }
+ validate!(self.frame_size > 0);
+
+ Ok(format)
+ }
+}
+
+fn read_token(src: &mut ByteReader) -> DemuxerResult<(bool, String)> {
+ let mut string = String::new();
+ let ws;
+ loop {
+ let b = src.read_byte()?;
+ match b {
+ b' ' | b'\n' => { ws = b; break; },
+ 0..=0x7F => string.push(b as char),
+ _ => return Err(DemuxerError::InvalidData),
+ }
+ }
+
+ Ok((ws == b'\n', string))
+}
+
+pub struct Y4MDemuxerCreator { }
+
+impl DemuxerCreator for Y4MDemuxerCreator {
+ fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
+ Box::new(Y4MDemuxer::new(br))
+ }
+ fn get_name(&self) -> &'static str { "yuv4mpeg" }
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+ use std::fs::File;
+
+ #[test]
+ fn test_y4m_demux() {
+ let mut file = File::open("assets/Misc/test.y4m").unwrap();
+ let mut fr = FileReader::new_read(&mut file);
+ let mut br = ByteReader::new(&mut fr);
+ let mut dmx = Y4MDemuxer::new(&mut br);
+ let mut sm = StreamManager::new();
+ let mut si = SeekIndex::new();
+ dmx.open(&mut sm, &mut si).unwrap();
+
+ loop {
+ let pktres = dmx.get_frame(&mut sm);
+ if let Err(e) = pktres {
+ if e == DemuxerError::EOF { break; }
+ panic!("error");
+ }
+ let pkt = pktres.unwrap();
+ println!("Got {}", pkt);
+ }
+ }
+}
projects / nihav.git / commitdiff