Significantly improve timelens support

Cargo.lock

@@ -626,21 +626,6 @@ dependencies = [
  "rayon",
 ]
 
-[[package]]
-name = "media-ingestion"
-version = "1.0.0"
-dependencies = [
- "anyhow",
- "ffmpeg_api",
- "fraction",
- "image",
- "media_time",
- "serde",
- "serde_json",
- "structopt",
- "webvtt",
-]
-
 [[package]]
 name = "media_time"
 version = "0.1.0"
@@ -840,6 +825,21 @@ dependencies = [
  "syn 2.0.103",
 ]
 
+[[package]]
+name = "preview-generator"
+version = "1.0.0"
+dependencies = [
+ "anyhow",
+ "ffmpeg_api",
+ "fraction",
+ "image",
+ "media_time",
+ "serde",
+ "serde_json",
+ "structopt",
+ "webvtt",
+]
+
 [[package]]
 name = "proc-macro-error"
 version = "1.0.4"

Cargo.toml

 [package]
-name = "media-ingestion"
+name = "preview-generator"
 version = "1.0.0"
 authors = ["Janne Mareike Koschinski <janne@kuschku.de>"]
 

src/lib.rs

 #![allow(dead_code)]
 
-pub mod spritesheet;
-mod options;
-
 use std::path::Path;
 
 use anyhow::format_err;
 use ffmpeg_api::api::*;
 use ffmpeg_api::enums::*;
-pub use options::ExtractOptions;
+
+pub mod spritesheet;
+pub mod timelens;
 
 #[allow(clippy::too_many_arguments)]
 pub fn extract(
     input_file: &Path,
     output_folder: &Path,
-    options: ExtractOptions,
+    spritesheet_options: Option<spritesheet::SpritesheetOptions>,
+    timelens_options: Option<timelens::TimelensOptions>,
+    discard: AVDiscard,
     scaler: SwsScaler,
     flags: SwsFlags,
 ) -> anyhow::Result<()> {
@@ -34,7 +35,7 @@ pub fn extract(
             false
         })
         .ok_or_else(|| format_err!("Could not find video stream"))?;
-    stream.set_discard(AVDiscard::NonKey);
+    stream.set_discard(discard);
 
     let index = stream.index();
     let time_base = stream.time_base();
@@ -49,16 +50,26 @@ pub fn extract(
         local_codec.name()?
     );
 
-    let frame_count = stream.duration()?.milliseconds() / options.frame_interval.milliseconds();
-    let mut spritesheet_manager = spritesheet::SpritesheetManager::new(
+    let mut spritesheet_manager: Option<spritesheet::SpritesheetManager> = if let Some(options) = spritesheet_options {
+        Some(spritesheet::SpritesheetManager::new(
             options,
-        u32::try_from(frame_count)?,
             output_folder,
             "preview",
-    );
+        )?)
+    } else {
+        None
+    };
     
-    let mut output_frame =
-        AVFrame::new().map_err(|error| format_err!("Could not create output frame: {}", error))?;
+    let mut timelens_manager: Option<timelens::TimelensManager> = if let Some(options) = timelens_options {
+        Some(timelens::TimelensManager::new(
+            options,
+            stream.duration()?,
+            output_folder,
+            "preview",
+        )?)
+    } else {
+        None
+    };
 
     if codec_parameters.codec_type() == AVMediaType::Video {
         let mut codec_context = AVCodecContext::new(&local_codec)
@@ -66,9 +77,9 @@ pub fn extract(
         codec_context.set_parameters(&codec_parameters);
         codec_context.open(&local_codec);
 
-        codec_context.set_skip_loop_filter(AVDiscard::NonKey);
-        codec_context.set_skip_idct(AVDiscard::NonKey);
-        codec_context.set_skip_frame(AVDiscard::NonKey);
+        codec_context.set_skip_loop_filter(discard);
+        codec_context.set_skip_idct(discard);
+        codec_context.set_skip_frame(discard);
 
         let mut packet = AVPacket::new()
             .map_err(|error| format_err!("Could not init temporary packet: {}", error))?;
@@ -93,45 +104,36 @@ pub fn extract(
                         frame.key_frame()
                     );
                     
+                    if let Some(spritesheet_manager) = &mut spritesheet_manager {
                         if spritesheet_manager.fulfils_frame_interval(timestamp) {
                             if !spritesheet_manager.initialized() {
-                            spritesheet_manager
-                                .initialize(frame.width() as u32, frame.height() as u32);
-                            output_frame
-                                .init(
-                                    spritesheet_manager.sprite_width() as i32,
-                                    spritesheet_manager.sprite_height() as i32,
-                                    AVPixelFormat::RGB24,
-                                )
-                                .map_err(|error| {
-                                    format_err!("Could not init output frame: {}", error)
-                                })?;
-                            scale_context
-                                .reinit(&frame, &output_frame, scaler, flags)
-                                .map_err(|error| {
-                                    format_err!("Could not reinit scale context: {}", error)
-                                })?;
+                                spritesheet_manager.initialize(frame.width() as u32, frame.height() as u32)?;
+                            }
+                            spritesheet_manager.add_frame(&mut scale_context, scaler, flags, timestamp, &frame)?;
+                        }
                     }
 
-                        scale_context.scale(&frame, &mut output_frame);
-
-                        spritesheet_manager.add_image(
-                            timestamp,
-                            image::ImageBuffer::from_raw(
-                                output_frame.width() as u32,
-                                output_frame.height() as u32,
-                                output_frame.data(0).to_vec(),
-                            )
-                            .ok_or_else(|| format_err!("Could not process frame"))?,
-                        )?;
+                    if let Some(timelens_manager) = &mut timelens_manager {
+                        if timelens_manager.fulfils_frame_interval(timestamp) {
+                            if !timelens_manager.initialized() {
+                                timelens_manager.initialize()?;
+                            }
+                            timelens_manager.add_frame(&mut scale_context, scaler, flags, timestamp, &frame)?;
+                        }
                     }
                 }
             }
         }
 
+        if let Some(spritesheet_manager) = &mut spritesheet_manager {
             spritesheet_manager.end_frame(duration);
             spritesheet_manager.save()?;
         }
 
+        if let Some(timelens_manager) = &mut timelens_manager {
+            timelens_manager.save()?;
+        }
+    }
+
     Ok(())
 }

src/main.rs

 use std::path::Path;
 
-use ffmpeg_api::enums::{SwsFlags, SwsScaler};
+use ffmpeg_api::enums::{AVDiscard, SwsFlags, SwsScaler};
 use image::ImageFormat as ImageOutputFormat;
 use media_time::MediaTime;
 use structopt::StructOpt;
-use media_ingestion::ExtractOptions;
+
+use preview_generator::spritesheet::SpritesheetOptions;
+use preview_generator::timelens::TimelensOptions;
 
 fn parse_scaler(src: &str) -> Result<SwsScaler, String> {
     match src {
@@ -28,12 +30,8 @@ fn parse_scaler(src: &str) -> Result<SwsScaler, String> {
 struct Options {
     input: String,
     output: String,
-    #[structopt(long = "frame-interval", default_value = "2")]
-    frame_interval: f64,
-    #[structopt(long = "num-horizontal", default_value = "5")]
-    num_horizontal: u32,
-    #[structopt(long = "num-vertical", default_value = "5")]
-    num_vertical: u32,
+    #[structopt(long = "keyframe-interval", default_value = "2")]
+    keyframe_interval: f64, // in seconds
     #[structopt(long = "max-size", default_value = "240")]
     max_size: u32,
     #[structopt(long = "format", default_value = "jpg")]
@@ -46,8 +44,20 @@ struct Options {
     fast_rounding: bool,
     #[structopt(long = "fast-scaling")]
     fast_scaling: bool,
+
     #[structopt(long = "timelens")]
     timelens: bool,
+    #[structopt(long = "timelens-width", default_value = "1000")]
+    timelens_width: u32,
+    #[structopt(long = "timelens-height", default_value = "90")]
+    timelens_height: u32,
+
+    #[structopt(long = "spritesheet")]
+    spritesheet: bool,
+    #[structopt(long = "spritesheet-columns", default_value = "5")]
+    spritesheet_columns: u32,
+    #[structopt(long = "spritesheet-rows", default_value = "5")]
+    spritesheet_rows: u32,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -64,22 +74,37 @@ fn main() -> anyhow::Result<()> {
         flags |= SwsFlags::BIT_EXACT_SCALING;
     }
 
-    if let Err(err) = media_ingestion::extract(
+    if let Err(err) = preview_generator::extract(
         Path::new(&options.input),
         Path::new(&options.output),
-        ExtractOptions {
+        if options.spritesheet {
+            Some(SpritesheetOptions {
                 max_size: options.max_size,
-            num_horizontal: options.num_horizontal,
-            num_vertical: options.num_vertical,
-            frame_interval: MediaTime::from_seconds_f64(options.frame_interval),
+                columns: options.spritesheet_columns,
+                rows: options.spritesheet_rows,
+                frame_interval: MediaTime::from_seconds_f64(options.keyframe_interval),
                 format: match options.format.as_str() {
                     "jpeg" | "jpg" => ImageOutputFormat::Jpeg,
                     "png" => ImageOutputFormat::Png,
                     "bmp" => ImageOutputFormat::Bmp,
                     _ => panic!("Unsupported image format: {}", options.format),
                 },
-            timelens: options.timelens,
+            })
+        } else { None },
+        if options.timelens {
+            Some(TimelensOptions {
+                width: options.timelens_width,
+                height: options.timelens_height,
+                frame_interval: MediaTime::from_seconds_f64(options.keyframe_interval),
+                format: match options.format.as_str() {
+                    "jpeg" | "jpg" => ImageOutputFormat::Jpeg,
+                    "png" => ImageOutputFormat::Png,
+                    "bmp" => ImageOutputFormat::Bmp,
+                    _ => panic!("Unsupported image format: {}", options.format),
                 },
+            })
+        } else { None },
+        AVDiscard::NonKey,
         options.scaler,
         flags,
     ) {

src/options.rs

-use image::ImageFormat as ImageOutputFormat;
-
-pub struct ExtractOptions {
-    pub max_size: u32,
-    pub num_horizontal: u32,
-    pub num_vertical: u32,
-    pub frame_interval: media_time::MediaTime,
-    pub format: ImageOutputFormat,
-    pub timelens: bool,
-}

src/spritesheet.rs

@@ -2,15 +2,21 @@ use std::fs::File;
 use std::io::BufWriter;
 use std::path::PathBuf;
 
-use anyhow::{bail, Error, format_err};
-use image::{DynamicImage, GenericImageView, ImageFormat as ImageOutputFormat, RgbImage};
-use media_time::MediaTime;
+use anyhow::{Error, format_err};
+use ffmpeg_api::api::{AVFrame, SwsContext};
+use ffmpeg_api::enums::{AVPixelFormat, SwsFlags, SwsScaler};
+use image::{DynamicImage, ImageFormat as ImageOutputFormat, RgbImage};
 use webvtt::{WebVTTCue, WebVTTFile};
-use crate::options::ExtractOptions;
+use media_time::MediaTime;
 
-pub enum ImageFormat {
-    Jpeg(i32),
-    Png,
+
+#[derive(Copy, Clone, Debug)]
+pub struct SpritesheetOptions {
+    pub max_size: u32,
+    pub columns: u32,
+    pub rows: u32,
+    pub frame_interval: MediaTime,
+    pub format: ImageOutputFormat,
 }
 
 pub struct SpritesheetManager {
@@ -28,19 +34,18 @@ pub struct SpritesheetManager {
     name: String,
     format: ImageOutputFormat,
     initialized: bool,
-    timelens: bool,
+    buffer: AVFrame,
 }
 
 impl SpritesheetManager {
     pub fn new(
-        options: ExtractOptions,
-        frame_count: u32,
+        options: SpritesheetOptions,
         output_path: impl Into<PathBuf>,
         name: impl AsRef<str>,
-    ) -> SpritesheetManager {
-        SpritesheetManager {
-            num_horizontal: if options.timelens { frame_count } else { options.num_horizontal },
-            num_vertical: if options.timelens { 1 } else { options.num_vertical },
+    ) -> Result<SpritesheetManager, Error> {
+        Ok(SpritesheetManager {
+            num_horizontal: options.columns,
+            num_vertical: options.rows,
             max_side: options.max_size,
             sprite_width: 0,
             sprite_height: 0,
@@ -53,23 +58,29 @@ impl SpritesheetManager {
             name: String::from(name.as_ref()),
             format: options.format,
             initialized: false,
-            timelens: options.timelens,
-        }
+            buffer: AVFrame::new()
+                .map_err(|error| format_err!("Could not create output frame: {}", error))?,
+        })
     }
 
-    pub fn initialize(&mut self, width: u32, height: u32) {
-        if self.timelens {
-            self.sprite_width = 1;
-            self.sprite_height = self.max_side;
-        } else if width >= height {
+    pub fn initialize(&mut self, width: u32, height: u32) -> Result<(), Error> {
+        if width >= height {
             self.sprite_width = self.max_side;
             self.sprite_height = self.sprite_width * height / width;
         } else {
             self.sprite_height = self.max_side;
             self.sprite_width = self.sprite_height * width / height;
         }
+        self.buffer.init(
+            self.sprite_width() as i32,
+            self.sprite_height() as i32,
+            AVPixelFormat::RGB24,
+        ).map_err(|error| {
+            format_err!("Could not init output frame: {}", error)
+        })?;
         self.spritesheet = self.reinit_buffer();
         self.initialized = true;
+        Ok(())
     }
 
     fn reinit_buffer(&self) -> RgbImage {
@@ -119,19 +130,18 @@ impl SpritesheetManager {
     }
 
     pub fn fulfils_frame_interval(&self, timestamp: MediaTime) -> bool {
-        self.current_image == 0 || timestamp - self.last_timestamp > self.frame_interval
+        self.current_image == 0 || (timestamp - self.last_timestamp > self.frame_interval)
     }
 
-    pub fn add_image(&mut self, timestamp: MediaTime, image: RgbImage) -> Result<(), Error> {
-        if image.width() != self.sprite_width || image.height() != self.sprite_height {
-            bail!(
-                "Wrong image size: {}x{}, but expected {}x{}",
-                image.width(),
-                image.height(),
-                self.sprite_width,
-                self.sprite_height
-            )
-        }
+    pub fn add_frame(&mut self, context: &mut SwsContext, scaler: SwsScaler, flags: SwsFlags, timestamp: MediaTime, frame: &AVFrame) -> Result<(), Error> {
+        context.reinit(frame, &self.buffer, scaler, flags)?;
+        context.scale(frame, &mut self.buffer);
+
+        let image = image::ImageBuffer::from_raw(
+            self.buffer.width() as u32,
+            self.buffer.height() as u32,
+            self.buffer.data(0).to_vec(),
+        ).ok_or_else(|| format_err!("Could not process frame"))?;
 
         let x: i64 = self.x(self.current_image).into();
         let y: i64 = self.y(self.current_image).into();
@@ -185,12 +195,8 @@ impl SpritesheetManager {
         let file = File::create(self.output_path.join(&name))
             .map_err(|err| format_err!("Could not create spritesheet {}: {}", &name, err))?;
 
-        let output = if self.timelens {
-            DynamicImage::ImageRgb8(self.spritesheet.view(0, 0, self.current_image, self.sprite_height).to_image())
-        } else {
         let new_buffer = self.reinit_buffer();
-            DynamicImage::ImageRgb8(std::mem::replace(&mut self.spritesheet, new_buffer))
-        };
+        let output = DynamicImage::ImageRgb8(std::mem::replace(&mut self.spritesheet, new_buffer));
 
         output
             .write_to(&mut BufWriter::new(file), self.format)

src/timelens.rs

+use std::fs::File;
+use std::io::{BufWriter, Write};
+use std::path::PathBuf;
+
+use anyhow::{Error, format_err};
+use ffmpeg_api::api::{AVFrame, SwsContext};
+use ffmpeg_api::enums::{AVPixelFormat, SwsFlags, SwsScaler};
+use image::{DynamicImage, GenericImageView, ImageFormat as ImageOutputFormat, RgbImage};
+use media_time::MediaTime;
+
+pub enum ImageFormat {
+    Jpeg(i32),
+    Png,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct TimelensOptions {
+    pub width: u32,
+    pub height: u32,
+    pub frame_interval: MediaTime,
+    pub format: ImageOutputFormat,
+}
+
+pub struct TimelensManager {
+    width: u32,
+    height: u32,
+    sprite_width: u32,
+    sprite_height: u32,
+    spritesheet: RgbImage,
+    current_image: u32,
+    last_timestamp: MediaTime,
+    frame_interval: MediaTime,
+    output_path: PathBuf,
+    name: String,
+    format: ImageOutputFormat,
+    initialized: bool,
+    buffer: AVFrame,
+}
+
+impl TimelensManager {
+    pub fn new(
+        options: TimelensOptions,
+        duration: MediaTime,
+        output_path: impl Into<PathBuf>,
+        name: impl AsRef<str>,
+    ) -> Result<TimelensManager, Error> {
+        let frame_count = duration.milliseconds() as f64 / options.frame_interval.milliseconds() as f64;
+        let frame_count = frame_count.ceil() as u32;
+        Ok(TimelensManager {
+            width: options.width,
+            height: options.height,
+            sprite_width: 1,
+            sprite_height: options.height,
+            spritesheet: RgbImage::new(frame_count, options.height),
+            current_image: 0,
+            last_timestamp: MediaTime::from_millis(0),
+            frame_interval: options.frame_interval,
+            output_path: output_path.into(),
+            name: String::from(name.as_ref()),
+            format: options.format,
+            initialized: false,
+            buffer: AVFrame::new()
+                .map_err(|error| format_err!("Could not create output frame: {}", error))?,
+        })
+    }
+
+    pub fn initialize(&mut self) -> Result<(), Error> {
+        self.buffer.init(
+            self.sprite_width as i32, self.sprite_height as i32,
+            AVPixelFormat::RGB24,
+        ).map_err(|error| {
+            format_err!("Could not init output frame: {}", error)
+        })?;
+        self.initialized = true;
+        Ok(())
+    }
+
+    pub fn initialized(&self) -> bool {
+        self.initialized
+    }
+
+    fn ending(&self) -> String {
+        String::from(match self.format {
+            ImageOutputFormat::Png => "png",
+            ImageOutputFormat::Jpeg => "jpeg",
+            ImageOutputFormat::Bmp => "bmp",
+            _ => panic!("Invalid image format: {:?}", self.format),
+        })
+    }
+
+    pub fn fulfils_frame_interval(&self, timestamp: MediaTime) -> bool {
+        self.current_image == 0 || timestamp - self.last_timestamp >= self.frame_interval
+    }
+
+    pub fn add_frame(&mut self, context: &mut SwsContext, scaler: SwsScaler, flags: SwsFlags, timestamp: MediaTime, frame: &AVFrame) -> Result<(), Error> {
+        context.reinit(frame, &self.buffer, scaler, flags)?;
+        context.scale(frame, &mut self.buffer);
+
+        let image = image::ImageBuffer::from_raw(
+            self.buffer.width() as u32,
+            self.buffer.height() as u32,
+            self.buffer.data(0).to_vec(),
+        ).ok_or_else(|| format_err!("Could not process frame"))?;
+
+        image::imageops::overlay(&mut self.spritesheet, &image, self.current_image.into(), 0);
+
+        self.last_timestamp = timestamp;
+        self.current_image += 1;
+
+        Ok(())
+    }
+
+    pub fn save(&mut self) -> Result<(), Error> {
+        let name = format!(
+            "{}_{}.{}",
+            self.name,
+            "timelens",
+            self.ending(),
+        );
+
+        let file = File::create(self.output_path.join(&name))
+            .map_err(|err| format_err!("Could not create spritesheet {}: {}", &name, err))?;
+
+          let view = self.spritesheet.view(
+            0,
+            0,
+            self.current_image * self.sprite_width,
+            self.sprite_height,
+        ).to_image();
+
+        let mut context = SwsContext::new();
+        let mut source = AVFrame::new()?;
+        source.init(
+            view.width() as i32,
+            view.height() as i32,
+            AVPixelFormat::RGB24,
+        ).map_err(|error| {
+            format_err!("Could not init output frame: {}", error)
+        })?;
+        source.data_mut(0usize).write_all(
+            view.as_raw().as_slice()
+        )?;
+        let mut scaled_result = AVFrame::new()?;
+        scaled_result.init(
+            self.width as i32, self.height as i32,
+            AVPixelFormat::RGB24,
+        ).map_err(|error| {
+            format_err!("Could not init output frame: {}", error)
+        })?;
+        context.reinit(&source, &scaled_result, SwsScaler::Point, SwsFlags::empty())?;
+        context.scale(&source, &mut scaled_result);
+
+        let image = image::ImageBuffer::from_raw(
+            scaled_result.width() as u32,
+            scaled_result.height() as u32,
+            scaled_result.data(0).to_vec(),
+        ).ok_or_else(|| format_err!("Could not process frame"))?;
+
+        DynamicImage::ImageRgb8(image)
+            .write_to(&mut BufWriter::new(file), self.format)
+            .map_err(|err| format_err!("Could not write spritesheet {}: {}", &name, err))?;
+
+        Ok(())
+    }
+}