ultralytics-inference 0.0.24

// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

//! I/O utilities for saving results including video encoding.

#[cfg(feature = "video")]
use video_rs::{Encoder, Time, encode::Settings as EncoderSettings};

use crate::error::{InferenceError, Result};
use std::path::{Path, PathBuf};

#[cfg(feature = "video")]
use std::sync::Once;

#[cfg(feature = "video")]
static INIT: Once = Once::new();

/// Initialize global video logging configuration.
///
/// ensuring `video-rs` is initialized and `FFmpeg` logs are silenced
/// (only errors are shown). safe to call multiple times.
#[allow(clippy::missing_const_for_fn)]
pub fn init_logging() {
    #[cfg(feature = "video")]
    INIT.call_once(|| {
        if let Err(e) = video_rs::init() {
            eprintln!("Failed to initialize video-rs: {e}");
        }

        video_rs::ffmpeg::log::set_level(video_rs::ffmpeg::log::Level::Error);
    });
}

/// Create `path` and all of its parent directories if they do not already exist.
///
/// This is a no-op when the directory is already present. On failure the
/// underlying I/O error is wrapped with the offending path for context.
///
/// # Errors
///
/// Returns an error if the directory cannot be created.
pub(crate) fn ensure_dir(path: &Path) -> Result<()> {
    std::fs::create_dir_all(path).map_err(|e| {
        std::io::Error::new(
            e.kind(),
            format!("Failed to create directory {}: {e}", path.display()),
        )
        .into()
    })
}

/// Return the next available run directory, e.g. `runs/detect/predict`, then `predict2`, `predict3`, ...
#[must_use]
pub fn find_next_run_dir(base: &str, prefix: &str) -> String {
    let base_path = Path::new(base);
    let first = base_path.join(prefix);
    if !first.exists() {
        return first.to_string_lossy().into_owned();
    }
    for i in 2.. {
        let numbered = base_path.join(format!("{prefix}{i}"));
        if !numbered.exists() {
            return numbered.to_string_lossy().into_owned();
        }
    }
    first.to_string_lossy().into_owned()
}

/// A wrapper around `video-rs` encoder to simplify video saving.
#[cfg(feature = "video")]
pub struct VideoWriter {
    encoder: Encoder,
    frame_duration: Time,
    position: Time,
    width: usize,
    height: usize,
}

#[cfg(feature = "video")]
impl VideoWriter {
    /// Create a new `VideoWriter`.
    ///
    /// # Arguments
    ///
    /// * `path` - Output video path (e.g., "output.mp4").
    /// * `width` - Video width.
    /// * `height` - Video height.
    /// * `fps` - Frames per second.
    ///
    /// # Errors
    ///
    /// Returns an error if the encoder cannot be initialized.
    pub fn new<P: AsRef<Path>>(path: P, width: usize, height: usize, fps: f32) -> Result<Self> {
        let output_path = path.as_ref().to_path_buf();

        // Ensure parent directory exists
        if let Some(parent) = output_path.parent() {
            ensure_dir(parent)?;
        }

        let settings = EncoderSettings::preset_h264_yuv420p(width, height, false);
        let encoder = Encoder::new(output_path.as_path(), settings).map_err(|e| {
            InferenceError::VideoError(format!("Failed to create video encoder: {e}"))
        })?;

        // Calculate frame duration
        // video-rs uses a rational time base.
        // We can approximate by converting seconds to Time.
        let seconds_per_frame = 1.0 / f64::from(fps);
        let frame_duration = Time::from_secs_f64(seconds_per_frame);

        Ok(Self {
            encoder,
            frame_duration,
            position: Time::zero(),
            width,
            height,
        })
    }

    /// Write a frame to the video.
    ///
    /// # Arguments
    ///
    /// * `frame` - Input frame as `DynamicImage`.
    ///
    /// # Errors
    ///
    /// Returns an error if encoding fails or frame dimensions don't match.
    pub fn write_frame(&mut self, frame: &image::DynamicImage) -> Result<()> {
        let img_buffer = frame.to_rgb8();
        let width = img_buffer.width() as usize;
        let height = img_buffer.height() as usize;

        if width != self.width || height != self.height {
            return Err(InferenceError::VideoError(format!(
                "Frame dimensions {}x{} do not match video dimensions {}x{}",
                width, height, self.width, self.height
            )));
        }

        let raw = img_buffer.into_raw();

        #[cfg(feature = "video")]
        let frame_array = ndarray::Array3::from_shape_vec((height, width, 3), raw)
            .map_err(|e| InferenceError::VideoError(e.to_string()))?;

        self.encoder
            .encode(&frame_array, self.position)
            .map_err(|e| InferenceError::VideoError(format!("Failed to encode frame: {e}")))?;

        self.position = self.position.aligned_with(self.frame_duration).add();
        Ok(())
    }

    /// Finish writing the video.
    ///
    /// Calling this explicitly is optional as `drop` will also clean up,
    /// but this allows catching errors.
    /// # Errors
    ///
    /// Returns an error if the encoder fails to finish.
    pub fn finish(mut self) -> Result<()> {
        self.encoder.finish().map_err(|e| {
            InferenceError::VideoError(format!("Failed to finish video encoding: {e}"))
        })
    }
}

/// Helper struct to handle saving inference results to video or disk.
///
/// This consolidates logic for deciding whether to save as a video file
/// or individual frames, and manages the `VideoWriter` state.
pub struct SaveResults {
    save_dir: PathBuf,
    #[cfg(feature = "video")]
    save_frames: bool,
    #[cfg(feature = "video")]
    video_writer: Option<VideoWriter>,
}

impl SaveResults {
    /// Create a new `SaveResults`.
    ///
    /// # Arguments
    ///
    /// * `save_dir` - Directory to save results.
    /// * `save_frames` - If true, force saving individual frames even for video sources.
    #[must_use]
    #[cfg_attr(not(feature = "video"), allow(unused_variables))]
    pub fn new(save_dir: PathBuf, save_frames: bool) -> Self {
        init_logging();

        Self {
            save_dir,
            #[cfg(feature = "video")]
            save_frames,
            #[cfg(feature = "video")]
            video_writer: None,
        }
    }

    /// Save an annotated frame.
    ///
    /// Decides automatically whether to append to a video or save as an image file
    /// based on the source type and configuration.
    ///
    /// # Arguments
    ///
    /// * `is_video` - Whether the source is a video/stream.
    /// * `meta` - Source metadata (path, frame index, fps).
    /// * `annotated` - The annotated image to save.
    ///
    /// # Errors
    ///
    /// Returns an error if saving the image or video frame fails.
    pub fn save(
        &mut self,
        is_video: bool,
        meta: &crate::source::SourceMeta,
        annotated: &image::DynamicImage,
    ) -> Result<()> {
        init_logging();

        #[cfg(feature = "video")]
        let save_as_video = is_video && !self.save_frames;
        #[cfg(not(feature = "video"))]
        let save_as_video = false;

        if save_as_video {
            #[cfg(feature = "video")]
            {
                // Video saving logic
                if self.video_writer.is_none() {
                    let filename = Path::new(&meta.path)
                        .file_name()
                        .unwrap_or_default()
                        .to_string_lossy();

                    let output_name = Path::new(filename.as_ref())
                        .with_extension("mp4")
                        .file_name()
                        .unwrap_or_default()
                        .to_string_lossy()
                        .to_string();

                    let save_path = self.save_dir.join(output_name);
                    let width = annotated.width() as usize;
                    let height = annotated.height() as usize;
                    let fps = meta.fps.unwrap_or(30.0);

                    // Ensure directory exists
                    if let Some(parent) = save_path.parent() {
                        ensure_dir(parent)?;
                    }

                    self.video_writer = Some(VideoWriter::new(save_path, width, height, fps)?);
                }

                if let Some(writer) = &mut self.video_writer {
                    writer.write_frame(annotated)?;
                }
            }
        } else {
            // Image saving logic
            let (save_dir, filename) = if is_video {
                // For video sources, create a subfolder: {video_name}_frames/
                let video_stem = Path::new(&meta.path)
                    .file_stem()
                    .unwrap_or_default()
                    .to_string_lossy();
                let frames_dir = self.save_dir.join(format!("{video_stem}_frames"));
                let frame_num = meta.frame_idx + 1;
                let filename = format!("{video_stem}_{frame_num}.jpg");
                (frames_dir, filename)
            } else {
                let filename = Path::new(&meta.path)
                    .file_name()
                    .unwrap_or_default()
                    .to_string_lossy()
                    .to_string();
                (self.save_dir.clone(), filename)
            };

            let save_path = save_dir.join(filename);

            // Ensure directory exists
            ensure_dir(&save_dir)?;

            annotated
                .save(&save_path)
                .map_err(|e| InferenceError::ImageError(e.to_string()))?;
        }
        Ok(())
    }

    /// Finish any active video writing.
    ///
    /// # Errors
    ///
    /// Returns an error if the video writer fails to finish.
    pub fn finish(self) -> Result<()> {
        #[cfg(feature = "video")]
        if let Some(writer) = self.video_writer {
            writer.finish()?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::source::SourceMeta;

    #[test]
    fn test_ensure_dir_creates_nested() {
        let tmp = tempfile::tempdir().unwrap();
        let nested = tmp.path().join("a").join("b").join("c");
        assert!(!nested.exists());
        ensure_dir(&nested).unwrap();
        assert!(nested.exists());
        // Idempotent: a second call on an existing dir is a no-op success.
        ensure_dir(&nested).unwrap();
        assert!(nested.exists());
    }

    #[test]
    fn test_find_next_run_dir_numbering() {
        let tmp = tempfile::tempdir().unwrap();
        let base = tmp.path().to_string_lossy().into_owned();

        // First call returns the un-suffixed prefix when nothing exists yet.
        let first = find_next_run_dir(&base, "predict");
        assert!(first.ends_with("predict"));

        // Once `predict` exists, the next is `predict2`, then `predict3`, ...
        std::fs::create_dir_all(&first).unwrap();
        let second = find_next_run_dir(&base, "predict");
        assert!(second.ends_with("predict2"));

        std::fs::create_dir_all(&second).unwrap();
        let third = find_next_run_dir(&base, "predict");
        assert!(third.ends_with("predict3"));
    }

    #[test]
    fn test_init_logging_is_idempotent() {
        // Safe to call multiple times.
        init_logging();
        init_logging();
    }

    #[test]
    fn test_save_results_writes_image() {
        let tmp = tempfile::tempdir().unwrap();
        let mut saver = SaveResults::new(tmp.path().to_path_buf(), false);

        let img = image::DynamicImage::new_rgb8(8, 8);
        let meta = SourceMeta {
            path: "frame.jpg".to_string(),
            ..SourceMeta::default()
        };

        saver.save(false, &meta, &img).unwrap();
        assert!(tmp.path().join("frame.jpg").exists());

        // finish() is a clean no-op when no video writer was opened.
        saver.finish().unwrap();
    }

    #[cfg(feature = "video")]
    #[test]
    fn test_video_writer_roundtrip() {
        // Encode a few frames to a real mp4 and confirm a non-empty file is produced.
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("out.mp4");
        let mut writer = VideoWriter::new(&path, 32, 32, 10.0).unwrap();
        for _ in 0..3 {
            writer
                .write_frame(&image::DynamicImage::new_rgb8(32, 32))
                .unwrap();
        }
        // Mismatched frame size is rejected.
        assert!(
            writer
                .write_frame(&image::DynamicImage::new_rgb8(16, 16))
                .is_err()
        );
        writer.finish().unwrap();

        assert!(path.exists());
        assert!(std::fs::metadata(&path).unwrap().len() > 0);
    }

    #[cfg(feature = "video")]
    #[test]
    fn test_save_results_video_source_writes_frame_image_when_save_frames() {
        // With save_frames=true a video source still goes down the image path,
        // landing under a `{stem}_frames/` subdirectory.
        let tmp = tempfile::tempdir().unwrap();
        let mut saver = SaveResults::new(tmp.path().to_path_buf(), true);

        let img = image::DynamicImage::new_rgb8(8, 8);
        let meta = SourceMeta {
            frame_idx: 0,
            path: "clip.mp4".to_string(),
            ..SourceMeta::default()
        };

        saver.save(true, &meta, &img).unwrap();
        assert!(tmp.path().join("clip_frames").join("clip_1.jpg").exists());
        saver.finish().unwrap();
    }

    #[cfg(feature = "video")]
    #[test]
    fn test_save_results_video_branch_writes_mp4() {
        // Video source with save_frames=false routes through the VideoWriter and
        // produces a single .mp4 named after the source.
        let tmp = tempfile::tempdir().unwrap();
        let mut saver = SaveResults::new(tmp.path().to_path_buf(), false);
        let img = image::DynamicImage::new_rgb8(32, 32);
        let meta = SourceMeta {
            frame_idx: 0,
            path: "movie.mp4".to_string(),
            fps: Some(10.0),
            ..SourceMeta::default()
        };
        saver.save(true, &meta, &img).unwrap();
        saver.save(true, &meta, &img).unwrap();
        saver.finish().unwrap();
        assert!(tmp.path().join("movie.mp4").exists());
    }
}