Skip to main content

yscv_video/
video_io.rs

1use std::collections::HashMap;
2use std::io::{Read, Write};
3use std::path::Path;
4
5use super::error::VideoError;
6use super::frame::Rgb8Frame;
7
8/// Video container metadata.
9#[derive(Debug, Clone)]
10pub struct VideoMeta {
11    pub width: u32,
12    pub height: u32,
13    pub frame_count: u32,
14    pub fps: f32,
15    pub properties: HashMap<String, String>,
16}
17
18/// Reads raw RGB8 frames from a simple uncompressed video file.
19///
20/// Format: [8 bytes: magic "RCVVIDEO"] [4 bytes: width LE] [4 bytes: height LE] [4 bytes: frame_count LE] [4 bytes: fps as f32 LE bits] [frames: width*height*3 bytes each].
21pub struct RawVideoReader {
22    pub meta: VideoMeta,
23    data: Vec<u8>,
24    frame_offset: usize,
25    current_frame: u32,
26}
27
28const MAGIC: &[u8; 8] = b"RCVVIDEO";
29
30impl RawVideoReader {
31    /// Opens a raw video file for reading.
32    pub fn open(path: &Path) -> Result<Self, VideoError> {
33        let mut file = std::fs::File::open(path)
34            .map_err(|e| VideoError::Source(format!("{}: {e}", path.display())))?;
35        let mut data = Vec::new();
36        file.read_to_end(&mut data)
37            .map_err(|e| VideoError::Source(e.to_string()))?;
38
39        if data.len() < 24 || &data[..8] != MAGIC {
40            return Err(VideoError::Source("invalid raw video file header".into()));
41        }
42
43        let width = u32::from_le_bytes([data[8], data[9], data[10], data[11]]);
44        let height = u32::from_le_bytes([data[12], data[13], data[14], data[15]]);
45        let frame_count = u32::from_le_bytes([data[16], data[17], data[18], data[19]]);
46        let fps = f32::from_le_bytes([data[20], data[21], data[22], data[23]]);
47
48        Ok(Self {
49            meta: VideoMeta {
50                width,
51                height,
52                frame_count,
53                fps,
54                properties: HashMap::new(),
55            },
56            data,
57            frame_offset: 24,
58            current_frame: 0,
59        })
60    }
61
62    /// Reads the next frame, if available.
63    pub fn next_frame(&mut self) -> Option<Rgb8Frame> {
64        if self.current_frame >= self.meta.frame_count {
65            return None;
66        }
67        let frame_size = self.meta.width as usize * self.meta.height as usize * 3;
68        let start = self.frame_offset + self.current_frame as usize * frame_size;
69        let end = start + frame_size;
70        if end > self.data.len() {
71            return None;
72        }
73        self.current_frame += 1;
74        Rgb8Frame::from_bytes(
75            self.current_frame as u64 - 1,
76            0,
77            self.meta.width as usize,
78            self.meta.height as usize,
79            bytes::Bytes::copy_from_slice(&self.data[start..end]),
80        )
81        .ok()
82    }
83
84    /// Resets to the beginning.
85    pub fn seek_start(&mut self) {
86        self.current_frame = 0;
87    }
88
89    /// Returns the frame count.
90    pub fn frame_count(&self) -> u32 {
91        self.meta.frame_count
92    }
93}
94
95/// Writes raw RGB8 frames to a simple uncompressed video file.
96pub struct RawVideoWriter {
97    width: u32,
98    height: u32,
99    fps: f32,
100    frames: Vec<Vec<u8>>,
101}
102
103impl RawVideoWriter {
104    pub fn new(width: u32, height: u32, fps: f32) -> Self {
105        Self {
106            width,
107            height,
108            fps,
109            frames: Vec::new(),
110        }
111    }
112
113    /// Appends an RGB8 frame (must be width*height*3 bytes).
114    pub fn push_frame(&mut self, rgb8_data: &[u8]) -> Result<(), VideoError> {
115        let expected = self.width as usize * self.height as usize * 3;
116        if rgb8_data.len() != expected {
117            return Err(VideoError::Source(format!(
118                "frame size mismatch: expected {expected}, got {}",
119                rgb8_data.len()
120            )));
121        }
122        self.frames.push(rgb8_data.to_vec());
123        Ok(())
124    }
125
126    /// Writes the video to a file.
127    pub fn save(&self, path: &Path) -> Result<(), VideoError> {
128        let mut file = std::fs::File::create(path)
129            .map_err(|e| VideoError::Source(format!("{}: {e}", path.display())))?;
130
131        let wr = |f: &mut std::fs::File, d: &[u8]| -> Result<(), VideoError> {
132            f.write_all(d)
133                .map_err(|e| VideoError::Source(e.to_string()))
134        };
135        wr(&mut file, MAGIC)?;
136        wr(&mut file, &self.width.to_le_bytes())?;
137        wr(&mut file, &self.height.to_le_bytes())?;
138        wr(&mut file, &(self.frames.len() as u32).to_le_bytes())?;
139        wr(&mut file, &self.fps.to_le_bytes())?;
140
141        for frame in &self.frames {
142            wr(&mut file, frame)?;
143        }
144
145        Ok(())
146    }
147
148    pub fn frame_count(&self) -> usize {
149        self.frames.len()
150    }
151}
152
153/// Image sequence reader: reads numbered image files as a video stream.
154///
155/// Pattern example: `frames/frame_%04d.png` -> reads `frame_0000.png`, `frame_0001.png`, etc.
156pub struct ImageSequenceReader {
157    pub width: usize,
158    pub height: usize,
159    paths: Vec<std::path::PathBuf>,
160    current: usize,
161}
162
163impl ImageSequenceReader {
164    /// Creates a reader from a sorted list of image file paths.
165    pub fn from_paths(paths: Vec<std::path::PathBuf>) -> Self {
166        Self {
167            width: 0,
168            height: 0,
169            paths,
170            current: 0,
171        }
172    }
173
174    /// Returns the total number of frames.
175    pub fn frame_count(&self) -> usize {
176        self.paths.len()
177    }
178
179    /// Resets to the beginning.
180    pub fn seek_start(&mut self) {
181        self.current = 0;
182    }
183
184    /// Returns the next image path without loading.
185    pub fn next_path(&mut self) -> Option<&Path> {
186        if self.current >= self.paths.len() {
187            return None;
188        }
189        let path = &self.paths[self.current];
190        self.current += 1;
191        Some(path)
192    }
193}
194
195// ---------------------------------------------------------------------------
196// MP4 / H.264 Video Reader
197// ---------------------------------------------------------------------------
198
199/// Reads H.264-encoded MP4 video files and decodes frames to RGB8.
200///
201/// Combines the MP4 box parser, Annex B NAL extraction, and H.264 decoder
202/// into a single end-to-end reader.
203///
204/// ```ignore
205/// let mut reader = Mp4VideoReader::open("input.mp4")?;
206/// while let Some(frame) = reader.next_frame()? {
207///     // frame.rgb8_data is Vec<u8>, frame.width / frame.height
208/// }
209/// ```
210pub struct Mp4VideoReader {
211    decoder: super::h264_decoder::H264Decoder,
212    nal_units: Vec<super::codec::NalUnit>,
213    current_nal: usize,
214}
215
216impl Mp4VideoReader {
217    /// Open an MP4 file containing H.264 video.
218    ///
219    /// Reads the entire file, parses MP4 boxes to find the `mdat` box,
220    /// and extracts NAL units from the raw data.
221    pub fn open(path: &Path) -> Result<Self, VideoError> {
222        let data = std::fs::read(path)
223            .map_err(|e| VideoError::Source(format!("{}: {e}", path.display())))?;
224
225        // Parse MP4 boxes to find mdat (media data)
226        let boxes = super::codec::parse_mp4_boxes(&data)?;
227        let mdat = boxes
228            .iter()
229            .find(|b| b.type_str() == "mdat")
230            .ok_or_else(|| VideoError::ContainerParse("no mdat box found in MP4".into()))?;
231
232        let mdat_start = (mdat.offset + mdat.header_size as u64) as usize;
233        let mdat_end = (mdat.offset + mdat.size) as usize;
234        if mdat_end > data.len() {
235            return Err(VideoError::ContainerParse(
236                "mdat box extends past EOF".into(),
237            ));
238        }
239
240        let mdat_data = &data[mdat_start..mdat_end];
241
242        // Try Annex B format first (start codes 0x000001 / 0x00000001)
243        let mut nal_units = super::codec::parse_annex_b(mdat_data);
244
245        // If no NAL units found with Annex B, try length-prefixed (AVCC) format
246        if nal_units.is_empty() {
247            nal_units = parse_avcc_nals(mdat_data);
248        }
249
250        // Also check for SPS/PPS in moov/stsd boxes (common in MP4)
251        if let Some(moov) = boxes.iter().find(|b| b.type_str() == "moov") {
252            let moov_start = (moov.offset + moov.header_size as u64) as usize;
253            let moov_end = (moov.offset + moov.size) as usize;
254            if moov_end <= data.len() {
255                let moov_data = &data[moov_start..moov_end];
256                let extra_nals = super::codec::parse_annex_b(moov_data);
257                // Prepend parameter set NALs before video NALs
258                let mut combined = extra_nals;
259                combined.extend(nal_units);
260                nal_units = combined;
261            }
262        }
263
264        if nal_units.is_empty() {
265            return Err(VideoError::ContainerParse(
266                "no NAL units found in MP4 mdat".into(),
267            ));
268        }
269
270        Ok(Self {
271            decoder: super::h264_decoder::H264Decoder::new(),
272            nal_units,
273            current_nal: 0,
274        })
275    }
276
277    /// Decode the next frame. Returns `None` when all NAL units are consumed.
278    pub fn next_frame(&mut self) -> Result<Option<super::codec::DecodedFrame>, VideoError> {
279        while self.current_nal < self.nal_units.len() {
280            let nal = &self.nal_units[self.current_nal];
281            self.current_nal += 1;
282            if let Some(frame) = self.decoder.process_nal(nal)? {
283                return Ok(Some(frame));
284            }
285        }
286        Ok(None)
287    }
288
289    /// Reset to the beginning (re-decode from first NAL).
290    pub fn seek_start(&mut self) {
291        self.current_nal = 0;
292        self.decoder = super::h264_decoder::H264Decoder::new();
293    }
294
295    /// Total number of NAL units found.
296    pub fn nal_count(&self) -> usize {
297        self.nal_units.len()
298    }
299}
300
301/// Parse length-prefixed NAL units (AVCC format, common in MP4).
302/// Each NAL is preceded by a 4-byte big-endian length.
303fn parse_avcc_nals(data: &[u8]) -> Vec<super::codec::NalUnit> {
304    let mut units = Vec::new();
305    let mut i = 0;
306    while i + 4 <= data.len() {
307        let len = u32::from_be_bytes([data[i], data[i + 1], data[i + 2], data[i + 3]]) as usize;
308        i += 4;
309        if len == 0 || i + len > data.len() {
310            break;
311        }
312        let header = data[i];
313        let nal_type = super::codec::NalUnitType::from_byte(header & 0x1F);
314        let nal_ref_idc = (header >> 5) & 3;
315        units.push(super::codec::NalUnit {
316            nal_type,
317            nal_ref_idc,
318            data: data[i..i + len].to_vec(),
319        });
320        i += len;
321    }
322    units
323}