iris/video/
reader.rs

1use crate::error::{IrisError, Result};
2use crate::video::frame::Frame;
3use crate::video::iterator::FrameIterator;
4use crate::video::metadata::{ContainerFormat, VideoMetadata};
5use burn::tensor::backend::Backend;
6use std::path::{Path, PathBuf};
7use std::time::Duration;
8
9use super::iterator::{load_animated_image, load_image_sequence};
10
11/// Seek mode for random access within a video.
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub enum SeekMode {
14    /// Seek to the nearest keyframe (fast, may not be precise).
15    ByKeyframe,
16    /// Seek to the exact frame (slower, precise).
17    Exact,
18    /// Seek to the nearest frame at the given timestamp.
19    ByTimestamp,
20}
21
22/// Configuration for opening a video source.
23#[derive(Clone, Debug)]
24pub struct VideoOpenOptions {
25    /// Seek mode for random access.
26    pub seek_mode: SeekMode,
27    /// Maximum number of frames to load (0 = unlimited).
28    pub max_frames: usize,
29    /// Whether to preload all frames into memory.
30    pub preload_all: bool,
31    /// Target width to resize frames to (0 = keep original).
32    pub target_width: usize,
33    /// Target height to resize frames to (0 = keep original).
34    pub target_height: usize,
35    /// Frame rate for image sequences (default: 30.0).
36    pub sequence_fps: f64,
37    /// Pattern to match in image sequence filenames.
38    pub sequence_pattern: String,
39}
40
41impl Default for VideoOpenOptions {
42    fn default() -> Self {
43        Self {
44            seek_mode: SeekMode::Exact,
45            max_frames: 0,
46            preload_all: true,
47            target_width: 0,
48            target_height: 0,
49            sequence_fps: 30.0,
50            sequence_pattern: String::new(),
51        }
52    }
53}
54
55/// A high-level video reader that loads frames into memory.
56///
57/// Supports animated images (GIF, APNG), image sequences, and basic
58/// container formats. All frames are cached for fast random access.
59pub struct VideoReader<B: Backend> {
60    frames: Vec<Frame<B>>,
61    metadata: VideoMetadata,
62    source_path: Option<PathBuf>,
63}
64
65impl<B: Backend> VideoReader<B> {
66    /// Opens a video file and loads all frames into memory.
67    pub fn open(path: impl AsRef<Path>, options: &VideoOpenOptions) -> Result<Self> {
68        let path = path.as_ref();
69        let path_str = path.to_string_lossy().to_string();
70        let format = ContainerFormat::from_path(&path_str);
71
72        if !path.exists() {
73            return Err(IrisError::Video(format!(
74                "Video file not found: {}",
75                path.display()
76            )));
77        }
78
79        let file_size = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
80
81        let device = B::Device::default();
82        let mut frames = match format {
83            ContainerFormat::Gif | ContainerFormat::Png | ContainerFormat::WebP => {
84                load_animated_image(path, &device)?
85            }
86            _ => {
87                if path.is_dir() {
88                    load_image_sequence(
89                        path,
90                        &options.sequence_pattern,
91                        &device,
92                        options.sequence_fps,
93                    )?
94                } else {
95                    let img = image::open(path)
96                        .map_err(|e| IrisError::Video(format!("Failed to open image: {e}")))?;
97
98                    let rgb = img.to_rgb8();
99                    let (w, h) = rgb.dimensions();
100                    let raw: Vec<f32> = rgb
101                        .pixels()
102                        .flat_map(|p| {
103                            let [r, g, b] = p.0;
104                            [r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0]
105                        })
106                        .collect();
107
108                    let tensor = burn::tensor::Tensor::<B, 3>::from_data(
109                        burn::tensor::TensorData::new(raw, [3, h as usize, w as usize]),
110                        &device,
111                    );
112                    let image = crate::image::Image::new(tensor);
113                    vec![Frame::new(image, Duration::ZERO, 0)]
114                }
115            }
116        };
117
118        if options.max_frames > 0 && frames.len() > options.max_frames {
119            frames.truncate(options.max_frames);
120        }
121
122        let metadata = if !frames.is_empty() {
123            let w = frames[0].width();
124            let h = frames[0].height();
125            let fps = if frames.len() > 1 {
126                let total_time: f64 = frames.iter().map(|f| f.duration.as_secs_f64()).sum();
127                if total_time > 0.0 {
128                    frames.len() as f64 / total_time
129                } else {
130                    30.0
131                }
132            } else {
133                30.0
134            };
135
136            VideoMetadata {
137                format,
138                duration: frames.iter().map(|f| f.duration).sum(),
139                fps,
140                width: w,
141                height: h,
142                frame_count: frames.len(),
143                video_codec: format!("{format:?}"),
144                pixel_format: super::metadata::PixelFormat::Rgb8,
145                rotation: 0,
146                bit_rate: 0,
147                streams: Vec::new(),
148                has_audio: false,
149                has_subtitles: false,
150                file_size,
151            }
152        } else {
153            return Err(IrisError::Video(
154                "No frames found in video file".to_string(),
155            ));
156        };
157
158        Ok(Self {
159            frames,
160            metadata,
161            source_path: Some(path.to_path_buf()),
162        })
163    }
164
165    /// Creates a reader from pre-loaded frames and metadata.
166    ///
167    /// Useful for in-memory video processing pipelines where frames
168    /// are generated or loaded programmatically.
169    #[must_use]
170    pub fn from_frames(frames: Vec<Frame<B>>, metadata: VideoMetadata) -> Self {
171        Self {
172            frames,
173            metadata,
174            source_path: None,
175        }
176    }
177
178    /// Returns a reference to all loaded frames.
179    #[must_use]
180    pub fn frames(&self) -> &[Frame<B>] {
181        &self.frames
182    }
183
184    /// Returns the total number of loaded frames.
185    #[must_use]
186    pub fn frame_count(&self) -> usize {
187        self.frames.len()
188    }
189
190    /// Returns video metadata.
191    #[must_use]
192    pub fn metadata(&self) -> &VideoMetadata {
193        &self.metadata
194    }
195
196    /// Returns the source file path.
197    #[must_use]
198    pub fn source_path(&self) -> Option<&Path> {
199        self.source_path.as_deref()
200    }
201
202    /// Returns a frame iterator for sequential access.
203    #[must_use]
204    pub fn iter(&self) -> FrameIterator<B> {
205        FrameIterator::new(self.frames.clone())
206    }
207
208    /// Returns a looping frame iterator.
209    #[must_use]
210    pub fn loop_iter(&self) -> FrameIterator<B> {
211        FrameIterator::new(self.frames.clone()).with_loop()
212    }
213
214    /// Gets a specific frame by index.
215    pub fn get_frame(&self, index: usize) -> Result<&Frame<B>> {
216        self.frames.get(index).ok_or_else(|| {
217            IrisError::InvalidParameter(format!(
218                "Frame index {index} out of range [0, {})",
219                self.frames.len()
220            ))
221        })
222    }
223
224    /// Gets frames in a range.
225    pub fn get_range(&self, start: usize, end: usize) -> Result<&[Frame<B>]> {
226        if start >= self.frames.len() || end > self.frames.len() || start >= end {
227            return Err(IrisError::InvalidParameter(format!(
228                "Invalid range [{start}, {end}) for {} frames",
229                self.frames.len()
230            )));
231        }
232        Ok(&self.frames[start..end])
233    }
234
235    /// Seeks to a specific timestamp and returns the frame.
236    pub fn seek_to_time(&self, time: Duration) -> Result<&Frame<B>> {
237        let frame_index = (time.as_secs_f64() * self.metadata.fps).round() as usize;
238        self.get_frame(frame_index)
239    }
240
241    /// Converts frames to a 4D tensor [B, C, H, W] (batch dimension).
242    pub fn to_batch_tensor(&self) -> Result<burn::tensor::Tensor<B, 4>> {
243        if self.frames.is_empty() {
244            return Err(IrisError::Video("No frames to batch".to_string()));
245        }
246
247        let tensors: Vec<burn::tensor::Tensor<B, 4>> = self
248            .frames
249            .iter()
250            .map(|f| f.image.tensor.clone().unsqueeze())
251            .collect();
252
253        Ok(burn::tensor::Tensor::cat(tensors, 0))
254    }
255
256    /// Computes per-frame differences (absolute difference between consecutive frames).
257    pub fn frame_differences(&self) -> Result<Vec<burn::tensor::Tensor<B, 3>>> {
258        if self.frames.len() < 2 {
259            return Err(IrisError::Video("Need at least 2 frames".to_string()));
260        }
261
262        let diffs = self
263            .frames
264            .windows(2)
265            .map(|w| {
266                let diff = w[1].image.tensor.clone() - w[0].image.tensor.clone();
267                diff.abs()
268            })
269            .collect();
270
271        Ok(diffs)
272    }
273
274    /// Computes optical flow magnitude (sum of absolute differences) between consecutive frames.
275    pub fn motion_magnitudes(&self) -> Result<Vec<f32>> {
276        let diffs = self.frame_differences()?;
277        let mut mags = Vec::with_capacity(diffs.len());
278        for d in &diffs {
279            let data: Vec<f32> =
280                d.to_data().convert::<f32>().into_vec().map_err(|e| {
281                    IrisError::Tensor(format!("Failed to convert tensor data: {e}"))
282                })?;
283            let sum: f32 = data.iter().sum();
284            mags.push(sum);
285        }
286        Ok(mags)
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    use crate::test_helpers::{TestBackend, test_device};
294    use burn::tensor::TensorData;
295
296    fn make_test_frame(index: usize) -> Frame<TestBackend> {
297        let device = test_device();
298        let data = TensorData::new(vec![0.5f32; 3 * 32 * 32], [3, 32, 32]);
299        let tensor = burn::tensor::Tensor::<TestBackend, 3>::from_data(data, &device);
300        let img = crate::image::Image::new(tensor);
301        Frame::new(img, Duration::from_millis(index as u64 * 33), index)
302            .with_duration(Duration::from_millis(33))
303    }
304
305    #[test]
306    fn test_video_reader_metadata() {
307        let frames: Vec<_> = (0..30).map(make_test_frame).collect();
308        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 30);
309        let reader = VideoReader {
310            frames,
311            metadata,
312            source_path: None,
313        };
314
315        assert_eq!(reader.frame_count(), 30);
316        assert_eq!(reader.metadata().width, 32);
317        assert!((reader.metadata().fps - 30.0).abs() < 1e-6);
318    }
319
320    #[test]
321    fn test_video_reader_get_frame() {
322        let frames: Vec<_> = (0..5).map(make_test_frame).collect();
323        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 5);
324        let reader = VideoReader {
325            frames,
326            metadata,
327            source_path: None,
328        };
329
330        let frame = reader.get_frame(2).unwrap();
331        assert_eq!(frame.index, 2);
332
333        assert!(reader.get_frame(10).is_err());
334    }
335
336    #[test]
337    fn test_video_reader_get_range() {
338        let frames: Vec<_> = (0..10).map(make_test_frame).collect();
339        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 10);
340        let reader = VideoReader {
341            frames,
342            metadata,
343            source_path: None,
344        };
345
346        let range = reader.get_range(2, 5).unwrap();
347        assert_eq!(range.len(), 3);
348        assert_eq!(range[0].index, 2);
349    }
350
351    #[test]
352    fn test_video_reader_seek_to_time() {
353        let frames: Vec<_> = (0..60).map(make_test_frame).collect();
354        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 60);
355        let reader = VideoReader {
356            frames,
357            metadata,
358            source_path: None,
359        };
360
361        let frame = reader.seek_to_time(Duration::from_secs(1)).unwrap();
362        assert_eq!(frame.index, 30);
363    }
364
365    #[test]
366    fn test_video_reader_to_batch_tensor() {
367        let frames: Vec<_> = (0..3).map(make_test_frame).collect();
368        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 3);
369        let reader = VideoReader {
370            frames,
371            metadata,
372            source_path: None,
373        };
374
375        let batch = reader.to_batch_tensor().unwrap();
376        assert_eq!(batch.dims(), [3, 3, 32, 32]);
377    }
378
379    #[test]
380    fn test_video_reader_frame_differences() {
381        let frames: Vec<_> = (0..5).map(make_test_frame).collect();
382        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 5);
383        let reader = VideoReader {
384            frames,
385            metadata,
386            source_path: None,
387        };
388
389        let diffs = reader.frame_differences().unwrap();
390        assert_eq!(diffs.len(), 4);
391    }
392
393    #[test]
394    fn test_video_reader_motion_magnitudes() {
395        let frames: Vec<_> = (0..5).map(make_test_frame).collect();
396        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 5);
397        let reader = VideoReader {
398            frames,
399            metadata,
400            source_path: None,
401        };
402
403        let mags = reader.motion_magnitudes().unwrap();
404        assert_eq!(mags.len(), 4);
405        assert!(mags.iter().all(|&m| m.abs() < 1e-6));
406    }
407
408    #[test]
409    fn test_video_reader_iter() {
410        let frames: Vec<_> = (0..3).map(make_test_frame).collect();
411        let metadata = VideoMetadata::synthetic(32, 32, 30.0, 3);
412        let reader = VideoReader {
413            frames,
414            metadata,
415            source_path: None,
416        };
417
418        let mut iter = reader.iter();
419        assert_eq!(iter.len(), 3);
420        assert!(iter.next().is_some());
421    }
422}
iris/video/reader.rs

iris/video/
reader.rs
