iris/video/
writer.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
7/// Output format for the video writer.
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub enum OutputFormat {
10    /// Animated GIF
11    Gif,
12    /// PNG image sequence
13    PngSequence,
14    /// JPEG image sequence
15    JpegSequence,
16    /// QOI image sequence
17    QoiSequence,
18}
19
20impl OutputFormat {
21    /// Returns the file extension for this format.
22    #[must_use]
23    pub fn extension(&self) -> &'static str {
24        match self {
25            Self::Gif => "gif",
26            Self::PngSequence => "png",
27            Self::JpegSequence => "jpg",
28            Self::QoiSequence => "qoi",
29        }
30    }
31
32    /// Detects the output format from a file path.
33    #[must_use]
34    pub fn from_path(path: &str) -> Self {
35        let lower = path.to_lowercase();
36        if lower.ends_with(".gif") {
37            Self::Gif
38        } else if lower.ends_with(".qoi") {
39            Self::QoiSequence
40        } else if lower.ends_with(".jpg") || lower.ends_with(".jpeg") {
41            Self::JpegSequence
42        } else {
43            Self::PngSequence
44        }
45    }
46}
47
48/// Configuration for the video writer.
49#[derive(Clone, Debug)]
50pub struct VideoWriteOptions {
51    /// Output format.
52    pub format: OutputFormat,
53    /// Frames per second for the output.
54    pub fps: f64,
55    /// Number of loops for GIF (0 = infinite).
56    pub gif_loops: u32,
57    /// JPEG quality (1-100).
58    pub jpeg_quality: u8,
59}
60
61impl Default for VideoWriteOptions {
62    fn default() -> Self {
63        Self {
64            format: OutputFormat::Gif,
65            fps: 30.0,
66            gif_loops: 0,
67            jpeg_quality: 90,
68        }
69    }
70}
71
72/// A high-level video writer that buffers frames and writes to disk.
73///
74/// Supports GIF output and image sequences. Frames are buffered in
75/// memory and flushed on `finish()` or `Drop`.
76pub struct VideoWriter<B: Backend> {
77    frames: Vec<Frame<B>>,
78    output_path: PathBuf,
79    options: VideoWriteOptions,
80    width: usize,
81    height: usize,
82    finished: bool,
83}
84
85impl<B: Backend> VideoWriter<B> {
86    /// Creates a new video writer for the given output path and dimensions.
87    pub fn create(
88        output_path: impl AsRef<Path>,
89        width: usize,
90        height: usize,
91        options: &VideoWriteOptions,
92    ) -> Result<Self> {
93        let output_path = output_path.as_ref().to_path_buf();
94
95        if let Some(parent) = output_path.parent().filter(|p| !p.exists()) {
96            std::fs::create_dir_all(parent)
97                .map_err(|e| IrisError::Video(format!("Failed to create output directory: {e}")))?;
98        }
99
100        Ok(Self {
101            frames: Vec::new(),
102            output_path,
103            options: options.clone(),
104            width,
105            height,
106            finished: false,
107        })
108    }
109
110    /// Writes a single frame to the buffer.
111    pub fn write_frame(&mut self, frame: &Frame<B>) -> Result<()> {
112        if self.finished {
113            return Err(IrisError::Video("Writer already finished".to_string()));
114        }
115
116        if frame.width() != self.width || frame.height() != self.height {
117            return Err(IrisError::DimensionMismatch {
118                expected: vec![3, self.height, self.width],
119                actual: vec![3, frame.height(), frame.width()],
120            });
121        }
122
123        self.frames.push(frame.clone());
124        Ok(())
125    }
126
127    /// Returns the number of frames currently buffered.
128    #[must_use]
129    pub fn frame_count(&self) -> usize {
130        self.frames.len()
131    }
132
133    /// Returns the current buffer duration.
134    #[must_use]
135    pub fn duration(&self) -> Duration {
136        self.frames
137            .iter()
138            .map(|f| {
139                if f.duration.is_zero() {
140                    Duration::from_secs_f64(1.0 / self.options.fps)
141                } else {
142                    f.duration
143                }
144            })
145            .sum()
146    }
147
148    /// Finishes writing and flushes all frames to disk.
149    pub fn finish(&mut self) -> Result<()> {
150        if self.finished {
151            return Ok(());
152        }
153
154        if self.frames.is_empty() {
155            return Err(IrisError::Video("No frames to write".to_string()));
156        }
157
158        match self.options.format {
159            OutputFormat::Gif => self.write_gif()?,
160            OutputFormat::PngSequence => self.write_png_sequence()?,
161            OutputFormat::JpegSequence => self.write_jpeg_sequence()?,
162            OutputFormat::QoiSequence => self.write_qoi_sequence()?,
163        }
164
165        self.finished = true;
166        Ok(())
167    }
168
169    /// Converts a Burn tensor frame to an `image` crate RGB image.
170    fn frame_to_image(&self, frame: &Frame<B>) -> image::RgbImage {
171        let [c, h, w] = frame.shape();
172        assert!(c == 3, "Only RGB frames supported for output");
173
174        let raw_data: Vec<f32> = frame
175            .image
176            .tensor
177            .to_data()
178            .into_vec()
179            .expect("Tensor data should be valid");
180        let mut img = image::RgbImage::new(w as u32, h as u32);
181
182        let hw = h * w;
183        for y in 0..h {
184            for x in 0..w {
185                let idx = y * w + x;
186                let r = (raw_data[idx] * 255.0).clamp(0.0, 255.0) as u8;
187                let g = (raw_data[hw + idx] * 255.0).clamp(0.0, 255.0) as u8;
188                let b = (raw_data[2 * hw + idx] * 255.0).clamp(0.0, 255.0) as u8;
189                img.put_pixel(x as u32, y as u32, image::Rgb([r, g, b]));
190            }
191        }
192
193        img
194    }
195
196    /// Writes frames as an animated GIF.
197    fn write_gif(&self) -> Result<()> {
198        use image::codecs::gif::{GifEncoder, Repeat};
199
200        let file = std::fs::File::create(&self.output_path)
201            .map_err(|e| IrisError::Video(format!("Failed to create GIF file: {e}")))?;
202
203        let mut encoder = GifEncoder::new(file);
204        encoder
205            .set_repeat(Repeat::Infinite)
206            .map_err(|e| IrisError::Video(format!("Failed to set GIF repeat: {e}")))?;
207
208        let delay =
209            image::Delay::from_saturating_duration(Duration::from_secs_f64(1.0 / self.options.fps));
210
211        let frames: Vec<image::Frame> = self
212            .frames
213            .iter()
214            .map(|f| {
215                let img = self.frame_to_image(f);
216                let dyn_img = image::DynamicImage::ImageRgb8(img);
217                let rgba = dyn_img.to_rgba8();
218                image::Frame::from_parts(rgba, 0, 0, delay)
219            })
220            .collect();
221
222        encoder
223            .encode_frames(frames)
224            .map_err(|e| IrisError::Video(format!("Failed to write GIF: {e}")))?;
225
226        Ok(())
227    }
228
229    /// Writes frames as a PNG image sequence.
230    fn write_png_sequence(&self) -> Result<()> {
231        let base = self
232            .output_path
233            .file_stem()
234            .and_then(|s| s.to_str())
235            .unwrap_or("frame");
236        let parent = self.output_path.parent().unwrap_or(Path::new("."));
237
238        for (i, frame) in self.frames.iter().enumerate() {
239            let img = self.frame_to_image(frame);
240            let path = parent.join(format!("{base}_{i:06}.png"));
241            img.save(&path)
242                .map_err(|e| IrisError::Video(format!("Failed to save PNG frame {i}: {e}")))?;
243        }
244
245        Ok(())
246    }
247
248    /// Writes frames as a JPEG image sequence.
249    fn write_jpeg_sequence(&self) -> Result<()> {
250        use image::codecs::jpeg::JpegEncoder;
251
252        let base = self
253            .output_path
254            .file_stem()
255            .and_then(|s| s.to_str())
256            .unwrap_or("frame");
257        let parent = self.output_path.parent().unwrap_or(Path::new("."));
258
259        for (i, frame) in self.frames.iter().enumerate() {
260            let img = self.frame_to_image(frame);
261            let path = parent.join(format!("{base}_{i:06}.jpg"));
262            let file = std::fs::File::create(&path)
263                .map_err(|e| IrisError::Video(format!("Failed to create JPEG frame {i}: {e}")))?;
264            let encoder = JpegEncoder::new_with_quality(file, self.options.jpeg_quality);
265            img.write_with_encoder(encoder)
266                .map_err(|e| IrisError::Video(format!("Failed to write JPEG frame {i}: {e}")))?;
267        }
268
269        Ok(())
270    }
271
272    /// Writes frames as a QOI image sequence.
273    fn write_qoi_sequence(&self) -> Result<()> {
274        let base = self
275            .output_path
276            .file_stem()
277            .and_then(|s| s.to_str())
278            .unwrap_or("frame");
279        let parent = self.output_path.parent().unwrap_or(Path::new("."));
280
281        for (i, frame) in self.frames.iter().enumerate() {
282            let img = self.frame_to_image(frame);
283            let raw = img.into_raw();
284            let path = parent.join(format!("{base}_{i:06}.qoi"));
285
286            let mut qoi_data = Vec::new();
287            qoi_data.extend_from_slice(b"qoif");
288            qoi_data.extend_from_slice(&(self.width as u32).to_be_bytes());
289            qoi_data.extend_from_slice(&(self.height as u32).to_be_bytes());
290            qoi_data.push(3);
291            qoi_data.push(0);
292
293            let mut offset = 0;
294            while offset < raw.len() {
295                let r = raw[offset];
296                let g = raw[offset + 1];
297                let b = raw[offset + 2];
298                qoi_data.push(0xFF);
299                qoi_data.push(r);
300                qoi_data.push(g);
301                qoi_data.push(b);
302                offset += 3;
303            }
304
305            qoi_data.extend_from_slice(&[0, 0, 0, 0, 0, 0, 0, 1]);
306
307            std::fs::write(&path, &qoi_data)
308                .map_err(|e| IrisError::Video(format!("Failed to save QOI frame {i}: {e}")))?;
309        }
310
311        Ok(())
312    }
313
314    /// Returns whether the writer has been finished.
315    #[must_use]
316    pub fn is_finished(&self) -> bool {
317        self.finished
318    }
319
320    /// Returns a reference to the output path.
321    #[must_use]
322    pub fn output_path(&self) -> &Path {
323        &self.output_path
324    }
325}
326
327impl<B: Backend> Drop for VideoWriter<B> {
328    fn drop(&mut self) {
329        if !self.finished && !self.frames.is_empty() {
330            let _ = self.finish();
331        }
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338    use crate::test_helpers::{TestBackend, test_device};
339    use burn::tensor::TensorData;
340
341    fn make_test_frame(index: usize) -> Frame<TestBackend> {
342        let device = test_device();
343        let data = TensorData::new(vec![0.5f32; 3 * 32 * 32], [3, 32, 32]);
344        let tensor = burn::tensor::Tensor::<TestBackend, 3>::from_data(data, &device);
345        let img = crate::image::Image::new(tensor);
346        Frame::new(img, Duration::from_millis(index as u64 * 33), index)
347            .with_duration(Duration::from_millis(33))
348    }
349
350    #[test]
351    fn test_video_writer_create() {
352        let opts = VideoWriteOptions::default();
353        let writer = VideoWriter::<TestBackend>::create("/tmp/test.gif", 32, 32, &opts);
354        assert!(writer.is_ok());
355        let writer = writer.unwrap();
356        assert_eq!(writer.frame_count(), 0);
357        assert!(!writer.is_finished());
358    }
359
360    #[test]
361    fn test_video_writer_write_frame() {
362        let opts = VideoWriteOptions::default();
363        let mut writer =
364            VideoWriter::<TestBackend>::create("/tmp/test.gif", 32, 32, &opts).unwrap();
365
366        let frame = make_test_frame(0);
367        writer.write_frame(&frame).unwrap();
368        assert_eq!(writer.frame_count(), 1);
369
370        let frame = make_test_frame(1);
371        writer.write_frame(&frame).unwrap();
372        assert_eq!(writer.frame_count(), 2);
373    }
374
375    #[test]
376    fn test_video_writer_dimension_mismatch() {
377        let opts = VideoWriteOptions::default();
378        let mut writer =
379            VideoWriter::<TestBackend>::create("/tmp/test.gif", 64, 64, &opts).unwrap();
380
381        let device = test_device();
382        let data = TensorData::new(vec![0.5f32; 3 * 32 * 32], [3, 32, 32]);
383        let tensor = burn::tensor::Tensor::<TestBackend, 3>::from_data(data, &device);
384        let img = crate::image::Image::new(tensor);
385        let frame = Frame::new(img, Duration::ZERO, 0);
386
387        assert!(writer.write_frame(&frame).is_err());
388    }
389
390    #[test]
391    fn test_video_writer_duration() {
392        let opts = VideoWriteOptions {
393            fps: 30.0,
394            ..Default::default()
395        };
396        let mut writer =
397            VideoWriter::<TestBackend>::create("/tmp/test.gif", 32, 32, &opts).unwrap();
398
399        for i in 0..30 {
400            writer.write_frame(&make_test_frame(i)).unwrap();
401        }
402
403        let dur = writer.duration();
404        assert!((dur.as_secs_f64() - 1.0).abs() < 0.02);
405    }
406
407    #[test]
408    fn test_output_format_detection() {
409        assert_eq!(OutputFormat::from_path("out.gif"), OutputFormat::Gif);
410        assert_eq!(
411            OutputFormat::from_path("out.png"),
412            OutputFormat::PngSequence
413        );
414        assert_eq!(
415            OutputFormat::from_path("out.jpg"),
416            OutputFormat::JpegSequence
417        );
418        assert_eq!(
419            OutputFormat::from_path("out.qoi"),
420            OutputFormat::QoiSequence
421        );
422    }
423}
iris/video/writer.rs

iris/video/
writer.rs
