iris/video/
iterator.rs

1use crate::error::{IrisError, Result};
2use crate::video::frame::Frame;
3use burn::tensor::backend::Backend;
4use std::path::{Path, PathBuf};
5use std::time::Duration;
6
7use super::metadata::ContainerFormat;
8
9/// Iterator over frames from a video file.
10pub struct FrameIterator<B: Backend> {
11    frames: Vec<Frame<B>>,
12    current: usize,
13    loop_playback: bool,
14}
15
16impl<B: Backend> FrameIterator<B> {
17    /// Creates a new iterator from a pre-loaded set of frames.
18    #[must_use]
19    pub fn new(frames: Vec<Frame<B>>) -> Self {
20        Self {
21            frames,
22            current: 0,
23            loop_playback: false,
24        }
25    }
26
27    /// Enables loop playback.
28    #[must_use]
29    pub fn with_loop(mut self) -> Self {
30        self.loop_playback = true;
31        self
32    }
33
34    /// Returns the total number of frames.
35    #[must_use]
36    pub fn total_frames(&self) -> usize {
37        self.frames.len()
38    }
39
40    /// Returns the current frame index.
41    #[must_use]
42    pub fn current_index(&self) -> usize {
43        self.current
44    }
45
46    /// Seeks to a specific frame index.
47    pub fn seek(&mut self, index: usize) -> Result<()> {
48        if index >= self.frames.len() {
49            return Err(IrisError::InvalidParameter(format!(
50                "Frame index {} out of range [0, {})",
51                index,
52                self.frames.len()
53            )));
54        }
55        self.current = index;
56        Ok(())
57    }
58
59    /// Seeks to the frame at the given timestamp.
60    pub fn seek_to_time(&mut self, time: Duration, fps: f64) -> Result<()> {
61        let frame_index = (time.as_secs_f64() * fps).round() as usize;
62        self.seek(frame_index)
63    }
64
65    /// Returns the remaining frames count.
66    #[must_use]
67    pub fn remaining(&self) -> usize {
68        self.frames.len().saturating_sub(self.current)
69    }
70
71    /// Returns the total duration of all frames.
72    #[must_use]
73    pub fn total_duration(&self) -> Duration {
74        self.frames
75            .iter()
76            .map(|f| {
77                if f.duration.is_zero() {
78                    Duration::from_secs_f64(1.0 / 30.0)
79                } else {
80                    f.duration
81                }
82            })
83            .sum()
84    }
85}
86
87impl<B: Backend> Iterator for FrameIterator<B> {
88    type Item = Frame<B>;
89
90    fn next(&mut self) -> Option<Self::Item> {
91        if self.current >= self.frames.len() {
92            if self.loop_playback && !self.frames.is_empty() {
93                self.current = 0;
94            } else {
95                return None;
96            }
97        }
98        let frame = self.frames[self.current].clone();
99        self.current += 1;
100        Some(frame)
101    }
102
103    fn size_hint(&self) -> (usize, Option<usize>) {
104        let remaining = self.remaining();
105        (remaining, Some(remaining))
106    }
107}
108
109impl<B: Backend> ExactSizeIterator for FrameIterator<B> {}
110
111/// Loads all frames from an animated GIF file.
112pub fn load_animated_image<B: Backend>(path: &Path, device: &B::Device) -> Result<Vec<Frame<B>>> {
113    use image::AnimationDecoder;
114    use image::codecs::gif::GifDecoder;
115
116    let path_str = path.to_string_lossy().to_string();
117    let format = ContainerFormat::from_path(&path_str);
118
119    let file = std::fs::File::open(path)
120        .map_err(|e| IrisError::Video(format!("Failed to open video file: {e}")))?;
121
122    let reader = std::io::BufReader::new(file);
123
124    match format {
125        ContainerFormat::Gif => {
126            let decoder = GifDecoder::new(reader)
127                .map_err(|e| IrisError::Video(format!("Failed to decode GIF: {e}")))?;
128
129            let frames_iter = decoder.into_frames();
130            let mut frames = Vec::new();
131
132            for (i, result) in frames_iter.enumerate() {
133                let frame = result
134                    .map_err(|e| IrisError::Video(format!("Failed to read GIF frame {i}: {e}")))?;
135
136                let delay = frame.delay();
137                let (numer, denom) = delay.numer_denom_ms();
138                let duration_ms = if denom > 0 {
139                    numer as u64 * 1000 / denom as u64
140                } else {
141                    33
142                };
143                let duration = Duration::from_millis(duration_ms);
144
145                let img = frame.buffer();
146                let (w, h) = img.dimensions();
147                let raw: Vec<f32> = img
148                    .pixels()
149                    .flat_map(|p| {
150                        let [r, g, b, _a] = p.0;
151                        [r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0]
152                    })
153                    .collect();
154
155                let tensor = burn::tensor::Tensor::<B, 3>::from_data(
156                    burn::tensor::TensorData::new(raw, [3, h as usize, w as usize]),
157                    device,
158                );
159                let image = crate::image::Image::new(tensor);
160
161                let pts = Duration::from_secs_f64(i as f64 * duration.as_secs_f64());
162                let frame = Frame::new(image, pts, i)
163                    .with_duration(duration)
164                    .with_keyframe(i == 0 || i % 30 == 0);
165
166                frames.push(frame);
167            }
168
169            Ok(frames)
170        }
171        ContainerFormat::Png => {
172            use image::codecs::png::PngDecoder;
173
174            let file2 = std::fs::File::open(path)
175                .map_err(|e| IrisError::Video(format!("Failed to open PNG file: {e}")))?;
176            let reader2 = std::io::BufReader::new(file2);
177
178            let decoder = PngDecoder::new(reader2)
179                .map_err(|e| IrisError::Video(format!("Failed to decode PNG: {e}")))?;
180
181            let apng = decoder
182                .apng()
183                .map_err(|e| IrisError::Video(format!("Failed to parse APNG: {e}")))?;
184
185            let frames_iter = apng.into_frames();
186            let mut frames = Vec::new();
187
188            for (i, result) in frames_iter.enumerate() {
189                let frame = result
190                    .map_err(|e| IrisError::Video(format!("Failed to read APNG frame {i}: {e}")))?;
191
192                let delay = frame.delay();
193                let (numer, denom) = delay.numer_denom_ms();
194                let duration_ms = if denom > 0 {
195                    numer as u64 * 1000 / denom as u64
196                } else {
197                    33
198                };
199                let duration = Duration::from_millis(duration_ms);
200
201                let img = frame.buffer();
202                let (w, h) = img.dimensions();
203                let raw: Vec<f32> = img
204                    .pixels()
205                    .flat_map(|p| {
206                        let [r, g, b, _a] = p.0;
207                        [r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0]
208                    })
209                    .collect();
210
211                let tensor = burn::tensor::Tensor::<B, 3>::from_data(
212                    burn::tensor::TensorData::new(raw, [3, h as usize, w as usize]),
213                    device,
214                );
215                let image = crate::image::Image::new(tensor);
216
217                let pts = Duration::from_secs_f64(i as f64 * duration.as_secs_f64());
218                let frame = Frame::new(image, pts, i)
219                    .with_duration(duration)
220                    .with_keyframe(i == 0 || i % 30 == 0);
221
222                frames.push(frame);
223            }
224
225            Ok(frames)
226        }
227        _ => {
228            // For unsupported animated formats, try as single image
229            let img = image::open(path)
230                .map_err(|e| IrisError::Video(format!("Failed to open image: {e}")))?;
231
232            let rgb = img.to_rgb8();
233            let (w, h) = rgb.dimensions();
234            let raw: Vec<f32> = rgb
235                .pixels()
236                .flat_map(|p| {
237                    let [r, g, b] = p.0;
238                    [r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0]
239                })
240                .collect();
241
242            let tensor = burn::tensor::Tensor::<B, 3>::from_data(
243                burn::tensor::TensorData::new(raw, [3, h as usize, w as usize]),
244                device,
245            );
246            let image = crate::image::Image::new(tensor);
247            Ok(vec![Frame::new(image, Duration::ZERO, 0)])
248        }
249    }
250}
251
252/// Loads numbered image frames from a directory.
253pub fn load_image_sequence<B: Backend>(
254    dir: &Path,
255    pattern: &str,
256    device: &B::Device,
257    fps: f64,
258) -> Result<Vec<Frame<B>>> {
259    let mut entries: Vec<PathBuf> = std::fs::read_dir(dir)
260        .map_err(|e| IrisError::Video(format!("Failed to read directory: {e}")))?
261        .filter_map(|e| e.ok())
262        .map(|e| e.path())
263        .filter(|p| {
264            p.file_name()
265                .and_then(|n| n.to_str())
266                .map(|n| n.contains(pattern))
267                .unwrap_or(false)
268        })
269        .filter(|p| {
270            p.extension().is_some_and(|ext| {
271                matches!(
272                    ext.to_str(),
273                    Some("png" | "jpg" | "jpeg" | "bmp" | "tiff" | "webp")
274                )
275            })
276        })
277        .collect();
278
279    entries.sort();
280
281    let mut frames = Vec::with_capacity(entries.len());
282    for (i, entry) in entries.iter().enumerate() {
283        let img = image::open(entry)
284            .map_err(|e| IrisError::Video(format!("Failed to open frame {}: {e}", i + 1)))?;
285
286        let rgb = img.to_rgb8();
287        let (w, h) = rgb.dimensions();
288        let raw: Vec<f32> = rgb
289            .pixels()
290            .flat_map(|p| {
291                let [r, g, b] = p.0;
292                [r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0]
293            })
294            .collect();
295
296        let tensor = burn::tensor::Tensor::<B, 3>::from_data(
297            burn::tensor::TensorData::new(raw, [3, h as usize, w as usize]),
298            device,
299        );
300        let image = crate::image::Image::new(tensor);
301        let pts = Duration::from_secs_f64(i as f64 / fps);
302        let frame = Frame::new(image, pts, i)
303            .with_duration(Duration::from_secs_f64(1.0 / fps))
304            .with_keyframe(i == 0 || i % 30 == 0);
305
306        frames.push(frame);
307    }
308
309    Ok(frames)
310}
311
312/// Trait for extending `Frame` with convenience methods.
313pub trait FrameExt {
314    /// Sets the keyframe flag.
315    fn with_keyframe(self, is_keyframe: bool) -> Self;
316}
317
318impl<B: Backend> FrameExt for Frame<B> {
319    fn with_keyframe(mut self, is_keyframe: bool) -> Self {
320        self.is_keyframe = is_keyframe;
321        self
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use crate::test_helpers::{TestBackend, test_device};
329    use burn::tensor::TensorData;
330
331    fn make_frame(index: usize) -> Frame<TestBackend> {
332        let device = test_device();
333        let data = TensorData::new(vec![0.5f32; 3 * 32 * 32], [3, 32, 32]);
334        let tensor = burn::tensor::Tensor::<TestBackend, 3>::from_data(data, &device);
335        let img = crate::image::Image::new(tensor);
336        Frame::new(img, Duration::from_millis(index as u64 * 33), index)
337            .with_duration(Duration::from_millis(33))
338    }
339
340    #[test]
341    fn test_frame_iterator_sequential() {
342        let frames: Vec<_> = (0..5).map(make_frame).collect();
343        let mut iter = FrameIterator::new(frames);
344        assert_eq!(iter.total_frames(), 5);
345        assert_eq!(iter.remaining(), 5);
346
347        let first = iter.next().unwrap();
348        assert_eq!(first.index, 0);
349        assert_eq!(iter.current_index(), 1);
350        assert_eq!(iter.remaining(), 4);
351    }
352
353    #[test]
354    fn test_frame_iterator_exact_size() {
355        let frames: Vec<_> = (0..10).map(make_frame).collect();
356        let iter = FrameIterator::new(frames);
357        assert_eq!(iter.len(), 10);
358    }
359
360    #[test]
361    fn test_frame_iterator_seek() {
362        let frames: Vec<_> = (0..5).map(make_frame).collect();
363        let mut iter = FrameIterator::new(frames);
364        iter.seek(3).unwrap();
365        assert_eq!(iter.current_index(), 3);
366        let frame = iter.next().unwrap();
367        assert_eq!(frame.index, 3);
368    }
369
370    #[test]
371    fn test_frame_iterator_seek_out_of_bounds() {
372        let frames: Vec<_> = (0..5).map(make_frame).collect();
373        let mut iter = FrameIterator::new(frames);
374        assert!(iter.seek(10).is_err());
375    }
376
377    #[test]
378    fn test_frame_iterator_loop() {
379        let frames: Vec<_> = (0..3).map(make_frame).collect();
380        let mut iter = FrameIterator::new(frames).with_loop();
381        for _ in 0..7 {
382            assert!(iter.next().is_some());
383        }
384    }
385
386    #[test]
387    fn test_frame_iterator_total_duration() {
388        let frames: Vec<_> = (0..30).map(make_frame).collect();
389        let iter = FrameIterator::new(frames);
390        let dur = iter.total_duration();
391        assert!((dur.as_secs_f64() - 1.0).abs() < 0.02);
392    }
393
394    #[test]
395    fn test_frame_iterator_seek_to_time() {
396        let frames: Vec<_> = (0..60).map(make_frame).collect();
397        let mut iter = FrameIterator::new(frames);
398        iter.seek_to_time(Duration::from_secs(1), 30.0).unwrap();
399        assert_eq!(iter.current_index(), 30);
400    }
401}
iris/video/iterator.rs

iris/video/
iterator.rs
