From dbb2cbc9c1274135470608275c1893a201cba731 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Sun, 19 Sep 2021 13:28:36 +0200 Subject: [PATCH] YUV4MPEG demuxer --- nihav-commonfmt/Cargo.toml | 3 +- nihav-commonfmt/src/demuxers/mod.rs | 4 + nihav-commonfmt/src/demuxers/y4m.rs | 215 ++++++++++++++++++++++++++++ nihav-registry/src/detect.rs | 5 + 4 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 nihav-commonfmt/src/demuxers/y4m.rs diff --git a/nihav-commonfmt/Cargo.toml b/nihav-commonfmt/Cargo.toml index b322049..1ac7261 100644 --- a/nihav-commonfmt/Cargo.toml +++ b/nihav-commonfmt/Cargo.toml @@ -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"] diff --git a/nihav-commonfmt/src/demuxers/mod.rs b/nihav-commonfmt/src/demuxers/mod.rs index d602cfc..3edb2bc 100644 --- a/nihav-commonfmt/src/demuxers/mod.rs +++ b/nihav-commonfmt/src/demuxers/mod.rs @@ -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 index 0000000..093d444 --- /dev/null +++ b/nihav-commonfmt/src/demuxers/y4m.rs @@ -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 { + 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 { 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 { + 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::() { + self.width = w; + } + }, + b'H' => { + if let Ok(h) = val.parse::() { + self.height = h; + } + }, + b'F' => { + if let Ok(fden) = val.parse::() { + 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::() { + self.fps_num = fnum; + } + if let Ok(fden) = vals[0].parse::() { + 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 + '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); + } + } +} diff --git a/nihav-registry/src/detect.rs b/nihav-registry/src/detect.rs index b883f02..c51936d 100644 --- a/nihav-registry/src/detect.rs +++ b/nihav-registry/src/detect.rs @@ -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", -- 2.39.5