Skip to main content

memvid_rs/video/
encoder.rs

1//! Video encoding functionality for memvid-rs
2//!
3//! This module provides video encoding using static FFmpeg with H.265 codec
4//! optimized for QR code preservation using exact Python parameters.
5
6use crate::config::VideoConfig;
7use crate::error::{MemvidError, Result};
8use image::DynamicImage;
9use std::path::Path;
10
11/// Video encoder with static FFmpeg support
12pub struct VideoEncoder {
13    config: VideoConfig,
14}
15
16impl Default for VideoEncoder {
17    fn default() -> Self {
18        Self::new(VideoConfig::default())
19    }
20}
21
22impl VideoEncoder {
23    /// Create a new video encoder with custom configuration
24    pub fn new(config: VideoConfig) -> Self {
25        Self { config }
26    }
27
28    /// Encode multiple frames into a video file
29    pub async fn encode_frames(&self, frames: &[DynamicImage], output_path: &str) -> Result<()> {
30        if frames.is_empty() {
31            return Err(MemvidError::Video(
32                "No frames provided for encoding".to_string(),
33            ));
34        }
35
36        let output_path = Path::new(output_path);
37
38        // Create output directory if it doesn't exist
39        if let Some(parent) = output_path.parent() {
40            std::fs::create_dir_all(parent).map_err(MemvidError::Io)?;
41        }
42
43        log::info!(
44            "Encoding {} frames to {} ({}x{} @ {} fps)",
45            frames.len(),
46            output_path.display(),
47            self.config.frame_width,
48            self.config.frame_height,
49            self.config.fps
50        );
51
52        // Initialize FFmpeg
53        ffmpeg_next::init()
54            .map_err(|e| MemvidError::Video(format!("FFmpeg init failed: {}", e)))?;
55
56        // Create output format context
57        let mut output_ctx = ffmpeg_next::format::output(&output_path)
58            .map_err(|e| MemvidError::Video(format!("Failed to create output context: {}", e)))?;
59
60        // Find H.265 encoder (HEVC)
61        let codec = ffmpeg_next::encoder::find(ffmpeg_next::codec::Id::HEVC)
62            .ok_or_else(|| MemvidError::Video("H.265 (HEVC) encoder not found".to_string()))?;
63
64        // Create video stream
65        let mut stream = output_ctx
66            .add_stream(codec)
67            .map_err(|e| MemvidError::Video(format!("Failed to add video stream: {}", e)))?;
68        let stream_index = stream.index();
69
70        // Create encoder context
71        let mut encoder = ffmpeg_next::codec::context::Context::new_with_codec(codec)
72            .encoder()
73            .video()
74            .map_err(|e| MemvidError::Video(format!("Failed to create video encoder: {}", e)))?;
75
76        // Set encoder parameters - use config dimensions, not QR size
77        encoder.set_width(self.config.frame_width);
78        encoder.set_height(self.config.frame_height);
79        encoder.set_format(ffmpeg_next::format::Pixel::YUV420P);
80        encoder.set_time_base(ffmpeg_next::Rational::new(1, self.config.fps as i32));
81        encoder.set_frame_rate(Some(ffmpeg_next::Rational::new(self.config.fps as i32, 1)));
82
83        // Set codec parameters optimized for QR codes (matching Python H.265 implementation exactly)
84        let mut dictionary = ffmpeg_next::Dictionary::new();
85
86        // Use parameters from config.quality_params (matches Python H265_PARAMETERS)
87        for (key, value) in &self.config.quality_params {
88            dictionary.set(key, value);
89        }
90
91        // Open encoder
92        let mut encoder = encoder
93            .open_with(dictionary)
94            .map_err(|e| MemvidError::Video(format!("Failed to open encoder: {}", e)))?;
95
96        // Update stream parameters
97        stream.set_parameters(&encoder);
98
99        // Write header
100        output_ctx
101            .write_header()
102            .map_err(|e| MemvidError::Video(format!("Failed to write header: {}", e)))?;
103
104        // Encode frames with upscaling
105        self.encode_image_frames(frames, &mut encoder, &mut output_ctx, stream_index)
106            .await?;
107
108        // Write trailer
109        output_ctx
110            .write_trailer()
111            .map_err(|e| MemvidError::Video(format!("Failed to write trailer: {}", e)))?;
112
113        log::info!(
114            "Successfully encoded {} frames to {}",
115            frames.len(),
116            output_path.display()
117        );
118        Ok(())
119    }
120
121    /// Encode individual image frames with upscaling for QR preservation
122    async fn encode_image_frames(
123        &self,
124        frames: &[DynamicImage],
125        encoder: &mut ffmpeg_next::encoder::Video,
126        output_ctx: &mut ffmpeg_next::format::context::Output,
127        stream_index: usize,
128    ) -> Result<()> {
129        // Use config dimensions for all frames
130        let target_width = self.config.frame_width;
131        let target_height = self.config.frame_height;
132
133        log::info!(
134            "Upscaling frames from QR size to {}x{} for compression resistance",
135            target_width,
136            target_height
137        );
138
139        for (i, image) in frames.iter().enumerate() {
140            // Always upscale QR codes to target resolution for better compression resistance
141            let upscaled_image = image.resize_exact(
142                target_width,
143                target_height,
144                image::imageops::FilterType::Nearest, // Use nearest neighbor for crisp QR codes
145            );
146
147            // Convert image to RGB format
148            let rgb_image = upscaled_image.to_rgb8();
149            let rgb_data = rgb_image.as_raw();
150
151            // Create video frame
152            let mut frame = ffmpeg_next::frame::Video::new(
153                ffmpeg_next::format::Pixel::RGB24,
154                target_width,
155                target_height,
156            );
157
158            // Copy RGB data to frame
159            frame.data_mut(0)[..rgb_data.len()].copy_from_slice(rgb_data);
160            frame.set_pts(Some((i as f64 / self.config.fps * 1000.0) as i64));
161
162            // Convert RGB to YUV420P for encoding
163            let mut yuv_frame = ffmpeg_next::frame::Video::new(
164                ffmpeg_next::format::Pixel::YUV420P,
165                target_width,
166                target_height,
167            );
168
169            // Set up software scaler for RGB to YUV conversion
170            let mut scaler = ffmpeg_next::software::scaling::Context::get(
171                ffmpeg_next::format::Pixel::RGB24,
172                target_width,
173                target_height,
174                ffmpeg_next::format::Pixel::YUV420P,
175                target_width,
176                target_height,
177                ffmpeg_next::software::scaling::Flags::BILINEAR,
178            )
179            .map_err(|e| MemvidError::Video(format!("Failed to create scaler: {}", e)))?;
180
181            scaler
182                .run(&frame, &mut yuv_frame)
183                .map_err(|e| MemvidError::Video(format!("Failed to scale frame: {}", e)))?;
184
185            yuv_frame.set_pts(frame.pts());
186
187            // Send frame to encoder
188            encoder
189                .send_frame(&yuv_frame)
190                .map_err(|e| MemvidError::Video(format!("Failed to send frame: {}", e)))?;
191
192            // Receive and write packets
193            self.receive_and_write_packets(encoder, output_ctx, stream_index)
194                .await?;
195
196            if (i + 1) % 10 == 0 {
197                log::info!("Encoded {}/{} frames", i + 1, frames.len());
198            }
199        }
200
201        // Flush encoder
202        encoder
203            .send_eof()
204            .map_err(|e| MemvidError::Video(format!("Failed to send EOF: {}", e)))?;
205
206        // Receive remaining packets
207        self.receive_and_write_packets(encoder, output_ctx, stream_index)
208            .await?;
209
210        Ok(())
211    }
212
213    /// Receive packets from encoder and write to output
214    async fn receive_and_write_packets(
215        &self,
216        encoder: &mut ffmpeg_next::encoder::Video,
217        output_ctx: &mut ffmpeg_next::format::context::Output,
218        stream_index: usize,
219    ) -> Result<()> {
220        let mut packet = ffmpeg_next::Packet::empty();
221
222        while encoder.receive_packet(&mut packet).is_ok() {
223            packet.set_stream(stream_index);
224            packet
225                .write_interleaved(output_ctx)
226                .map_err(|e| MemvidError::Video(format!("Failed to write packet: {}", e)))?;
227        }
228
229        Ok(())
230    }
231
232    /// Encode a single image into a one-frame video
233    pub async fn encode_single_image(&self, image: &DynamicImage, output_path: &str) -> Result<()> {
234        self.encode_frames(&[image.clone()], output_path).await
235    }
236
237    /// Get video configuration
238    pub fn config(&self) -> &VideoConfig {
239        &self.config
240    }
241
242    /// Get supported output formats
243    pub fn supported_formats() -> Vec<String> {
244        // Common video formats supported by H.264
245        vec![
246            "mp4".to_string(),
247            "mkv".to_string(),
248            "avi".to_string(),
249            "mov".to_string(),
250        ]
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use crate::qr::encoder::QrEncoder;
258    use tempfile::NamedTempFile;
259
260    #[tokio::test]
261    async fn test_encoder_creation() {
262        let encoder = VideoEncoder::default();
263        assert_eq!(encoder.config.fps, 30.0); // Updated default FPS
264        assert_eq!(encoder.config.frame_width, 256); // QR upscaling target
265        assert_eq!(encoder.config.frame_height, 256); // QR upscaling target
266    }
267
268    #[tokio::test]
269    async fn test_video_config() {
270        let config = VideoConfig::default();
271        assert_eq!(config.fps, 30.0);
272        assert_eq!(config.frame_width, 256);
273        assert_eq!(config.frame_height, 256);
274        assert_eq!(config.codec, "libx265");
275        assert!(config.quality_params.contains_key("crf"));
276        assert!(config.quality_params.contains_key("preset"));
277    }
278
279    #[tokio::test]
280    async fn test_supported_formats() {
281        let formats = VideoEncoder::supported_formats();
282        assert!(formats.contains(&"mp4".to_string()));
283        assert!(formats.contains(&"mkv".to_string()));
284    }
285
286    #[tokio::test]
287    async fn test_encode_single_qr_frame() {
288        // Create a test QR code image
289        let qr_encoder = QrEncoder::default();
290        let qr_frame = qr_encoder.encode_text("Test QR code").unwrap();
291
292        // Create temporary output file
293        let temp_file = NamedTempFile::with_suffix(".mp4").unwrap();
294        let output_path = temp_file.path().to_str().unwrap();
295
296        // Test encoding with QR upscaling
297        let video_encoder = VideoEncoder::default();
298        let result = video_encoder
299            .encode_single_image(&qr_frame.image, output_path)
300            .await;
301
302        // Should succeed or fail gracefully (depending on FFmpeg availability)
303        match result {
304            Ok(_) => {
305                // Verify file was created
306                assert!(temp_file.path().exists());
307                let metadata = std::fs::metadata(temp_file.path()).unwrap();
308                assert!(metadata.len() > 0);
309                log::info!("Video encoding test successful - QR upscaling working");
310            }
311            Err(MemvidError::Video(_)) => {
312                // Expected if FFmpeg not properly configured in test environment
313                log::warn!("Video encoding test skipped - FFmpeg not available");
314            }
315            Err(e) => panic!("Unexpected error: {}", e),
316        }
317    }
318}