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