Update rust_highlight/src/lib.rs

rust_highlight/src/lib.rs

@@ -1,60 +1,86 @@
-use pyo3::prelude::*;
 use std::process::Command;
+use std::path::Path;
 
-#[pyfunction]
-fn render_video(
+pub fn render_video(
     id: usize,
     image_path: &str,
     audio_path: &str,
     duration: f64,
     words: Vec<(String, (u32, u32, u32, u32))>,
-) -> PyResult<String> {
-    let clip = format!("clip{}.mp4", id);
+    output_dir: &str,
+) -> Result<String, String> {
+    // Validate paths
+    if !Path::new(image_path).exists() {
+        return Err(format!("Image not found: {}", image_path));
+    }
+    if !Path::new(audio_path).exists() {
+        return Err(format!("Audio not found: {}", audio_path));
+    }
+
+    let output_path = format!("{}/clip{}.mp4", output_dir, id);
     let duration_str = duration.to_string();
+    let n_words = words.len() as f64;
+    let highlight_duration = duration / n_words;
 
-    // --- FFmpeg drawbox filters for each word ---
-    let mut vf = String::new();
-    let n = words.len() as f64;
+    // Build complex filter graph
+    let mut filter_complex = String::new();
+    let mut last_output = "0:v".to_string();
 
     for (i, (_, (x, y, w, h))) in words.iter().enumerate() {
-        let start = (i as f64) * (duration / n);
-        let end = start + (duration / n);
-        vf.push_str(&format!(
-            "drawbox=x={}:y={}:w={}:h={}:[email protected]:enable='between(t,{:.3},{:.3})',",
-            x, y, w, h, start, end
+        let start = i as f64 * highlight_duration;
+        let end = start + highlight_duration;
+
+        // Create highlight box
+        filter_complex.push_str(&format!(
+            "[{last_out}]drawbox=x={x}:y={y}:w={w}:h={h}:[email protected]:t=fill:enable='between(t,{start:.3},{end:.3})'[v{i}];",
+            last_out = last_output,
+            x = x,
+            y = y,
+            w = w,
+            h = h,
+            start = start,
+            end = end,
+            i = i
         ));
+        last_output = format!("[v{i}]");
     }
 
-    // Trim trailing comma and add pixel format fix
-    let vf_with_format = format!("{},format=yuv420p", vf.trim_end_matches(','));
+    // Remove trailing semicolon
+    filter_complex.pop();
 
-    // --- FFmpeg command ---
+    // FFmpeg command with optimized settings
     let status = Command::new("ffmpeg")
         .args(&[
-            "-y",                    // Overwrite output
-            "-loop", "1",            // Loop image
-            "-i", image_path,        // Input image
-            "-i", audio_path,        // Input audio
-            "-vf", &vf_with_format,  // Video filter with highlights
-            "-t", &duration_str,     // Duration
-            "-c:v", "libx264",       // Video codec
-            "-pix_fmt", "yuv420p",   // Pixel format (for compatibility)
-            "-c:a", "aac",           // Audio codec
-            &clip,                   // Output file
+            "-y", // Overwrite output
+            "-hwaccel", "auto", // Enable hardware acceleration
+            "-loop", "1",
+            "-i", image_path,
+            "-i", audio_path,
+            "-filter_complex", &filter_complex,
+            "-map", &last_output,
+            "-map", "1:a",
+            "-c:v", "libx264",
+            "-preset", "fast", // Balance between speed and compression
+            "-crf", "18", // High quality (lower = better quality)
+            "-pix_fmt", "yuv420p",
+            "-movflags", "+faststart", // For web streaming
+            "-c:a", "aac",
+            "-b:a", "192k", // Higher audio bitrate
+            "-r", "60", // Higher framerate for smoother animation
+            "-t", &duration_str,
+            "-vsync", "2", // Frame rate conversion method
+            &output_path,
         ])
-        .status()
-        .expect("Failed to run ffmpeg");
-
-    // Return result
-    if status.success() {
-        Ok(clip)
-    } else {
-        Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-            "FFmpeg failed",
-        ))
+        .status();
+
+    match status {
+        Ok(exit_status) if exit_status.success() => Ok(output_path),
+        Ok(_) => Err("FFmpeg command failed".to_string()),
+        Err(e) => Err(format!("Failed to execute FFmpeg: {}", e)),
     }
 }
 
+
 #[pymodule]
 fn rust_highlight(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(render_video, m)?)?;