]> git.nihav.org Git - nihav.git/commitdiff
YUV4MPEG demuxer
authorKostya Shishkov <kostya.shishkov@gmail.com>
Sun, 19 Sep 2021 11:28:36 +0000 (13:28 +0200)
committerKostya Shishkov <kostya.shishkov@gmail.com>
Thu, 30 Sep 2021 16:13:05 +0000 (18:13 +0200)
nihav-commonfmt/Cargo.toml
nihav-commonfmt/src/demuxers/mod.rs
nihav-commonfmt/src/demuxers/y4m.rs [new file with mode: 0644]
nihav-registry/src/detect.rs

index b322049d522ad2e56552222c19203ddaecc5e7f1..1ac726122854aa6171ab40468ec2ddeed3190882 100644 (file)
@@ -23,10 +23,11 @@ decoders = []
 demuxers = []
 encoders = []
 muxers = []
-all_demuxers = ["demuxer_avi", "demuxer_mov", "demuxer_wav"]
+all_demuxers = ["demuxer_avi", "demuxer_mov", "demuxer_wav", "demuxer_y4m"]
 demuxer_avi = ["demuxers"]
 demuxer_mov = ["demuxers"]
 demuxer_wav = ["demuxers"]
+demuxer_y4m = ["demuxers"]
 all_muxers = ["muxer_avi", "muxer_wav"]
 muxer_avi = ["muxers"]
 muxer_wav = ["muxers"]
index d602cfca5288043d6e1b964a32f9210a9e129951..3edb2bc50c1ed3e87b6aaa6602e9b7f646fe9c65 100644 (file)
@@ -14,6 +14,8 @@ mod avi;
 mod mov;
 #[cfg(feature="demuxer_wav")]
 mod wav;
+#[cfg(feature="demuxer_y4m")]
+mod y4m;
 
 const DEMUXERS: &[&dyn DemuxerCreator] = &[
 #[cfg(feature="demuxer_avi")]
@@ -22,6 +24,8 @@ const DEMUXERS: &[&dyn DemuxerCreator] = &[
     &mov::MOVDemuxerCreator {},
 #[cfg(feature="demuxer_wav")]
     &wav::WAVDemuxerCreator {},
+#[cfg(feature="demuxer_y4m")]
+    &y4m::Y4MDemuxerCreator {},
 ];
 
 /// Registers all available demuxers provided by this crate.
diff --git a/nihav-commonfmt/src/demuxers/y4m.rs b/nihav-commonfmt/src/demuxers/y4m.rs
new file mode 100644 (file)
index 0000000..093d444
--- /dev/null
@@ -0,0 +1,215 @@
+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);
+        }
+    }
+}
index b883f02112cf5c139041a0f76b2f2be81f0ea2a9..c51936d3a3c4b22526f4d0317c25802fe3c76834 100644 (file)
@@ -230,6 +230,11 @@ const DETECTORS: &[DetectConditions] = &[
                                                                 &CC::Str(b"moov")),
                                                                 &CC::Str(b"ftyp")) }],
     },
+    DetectConditions {
+        demux_name: "yuv4mpeg",
+        extensions: ".y4m",
+        conditions: &[CheckItem{offs: 0, cond: &CC::Str(b"YUV4MPEG2 ") }],
+    },
     DetectConditions {
         demux_name: "fcmp",
         extensions: ".cmp",