Improve video thumbnailing by using ffmpeg instead

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.toml

@@ -4,7 +4,7 @@ readme = "README.md"
 license = "Apache-2.0"
 authors = ["trivernis <trivernis@protonmail.com>"]
 description = "An image thumbnail creation library"
-version = "0.2.1"
+version = "0.2.2"
 edition = "2018"
 repository = "https://github.com/Trivernis/thumbnailer"
 
@@ -14,7 +14,7 @@ repository = "https://github.com/Trivernis/thumbnailer"
 webp = "0.2.0"
 mime = "0.3.16"
 rayon = "1.5.1"
-vid2img = "0.1.0"
+ffmpeg-next = "4.4.0"
 tempfile = "3.2.0"
 
 [dependencies.image]

src/error.rs

@@ -2,7 +2,6 @@ use image::ImageError;
 use mime::Mime;
 use std::fmt::{Debug, Display, Formatter};
 use std::io;
-use vid2img::{CaptureError, StreamError};
 
 pub type ThumbResult<T> = Result<T, ThumbError>;
 
@@ -18,9 +17,7 @@ pub enum ThumbError {
 
     NullVideo,
 
-    CaptureError(vid2img::CaptureError),
-
-    StreamError(vid2img::StreamError),
+    FFMPEG(ffmpeg_next::Error),
 }
 
 impl Display for ThumbError {
@@ -31,12 +28,7 @@ impl Display for ThumbError {
             ThumbError::Decode => write!(f, "failed to decode image"),
             ThumbError::Unsupported(mime) => write!(f, "Unsupported media type {}", mime),
             ThumbError::NullVideo => write!(f, "no video data found in file"),
-            ThumbError::CaptureError(c) => {
-                write!(f, "capture error when creating video thumbnail: {:?}", c)
-            }
-            ThumbError::StreamError(s) => {
-                write!(f, "stream error when creating video thumbnail: {:?}", s)
-            }
+            ThumbError::FFMPEG(e) => write!(f, "ffmpeg error: {}", e),
         }
     }
 }
@@ -46,6 +38,7 @@ impl std::error::Error for ThumbError {
         match self {
             ThumbError::IO(e) => e.source(),
             ThumbError::Image(i) => i.source(),
+            ThumbError::FFMPEG(e) => e.source(),
             _ => None,
         }
     }
@@ -63,14 +56,8 @@ impl From<image::error::ImageError> for ThumbError {
     }
 }
 
-impl From<vid2img::CaptureError> for ThumbError {
-    fn from(e: CaptureError) -> Self {
-        Self::CaptureError(e)
-    }
-}
-
-impl From<vid2img::StreamError> for ThumbError {
-    fn from(s: StreamError) -> Self {
-        Self::StreamError(s)
+impl From<ffmpeg_next::Error> for ThumbError {
+    fn from(e: ffmpeg_next::Error) -> Self {
+        Self::FFMPEG(e)
     }
 }

src/formats/video_format.rs

@@ -1,11 +1,14 @@
 use crate::error::{ThumbError, ThumbResult};
-use image::png::PngDecoder;
-use image::DynamicImage;
+use ffmpeg_next::filter;
+use ffmpeg_next::filter::Graph;
+use ffmpeg_next::frame::Video;
+use ffmpeg_next::media::Type as MediaType;
+use ffmpeg_next::threading::Config;
+use image::{DynamicImage, RgbaImage};
 use mime::Mime;
 use std::fs;
 use std::io::{BufRead, Seek};
 use std::path::PathBuf;
-use vid2img::FileSource;
 
 pub fn get_video_frame<R: BufRead + Seek>(mut reader: R, mime: Mime) -> ThumbResult<DynamicImage> {
     let tempdir = tempfile::tempdir()?;
@@ -25,13 +28,91 @@ pub fn get_video_frame<R: BufRead + Seek>(mut reader: R, mime: Mime) -> ThumbRes
 }
 
 fn extract_frame_from_video(path: &PathBuf) -> ThumbResult<DynamicImage> {
-    let source = FileSource::new(path, (2000, 2000))?;
-    for frame in source.into_iter() {
-        if let Ok(Some(data)) = frame {
-            let decoder = PngDecoder::new(data.as_slice())?;
-            let img = DynamicImage::from_decoder(decoder)?;
-            return Ok(img);
+    let mut input = ffmpeg_next::format::input(path)?;
+    let stream = input
+        .streams()
+        .best(MediaType::Video)
+        .ok_or_else(|| ThumbError::NullVideo)?;
+
+    let mut decoder = stream.codec().decoder().video()?;
+    decoder.set_threading(Config::count(1));
+
+    let args = format!(
+        "width={w}:height={h}:video_size={w}x{h}:pix_fmt={fmt}:time_base={base}",
+        w = decoder.width(),
+        h = decoder.height(),
+        fmt = decoder
+            .format()
+            .descriptor()
+            .ok_or_else(|| ThumbError::NullVideo)?
+            .name(),
+        base = decoder.time_base(),
+    );
+    let mut filter = Graph::new();
+    filter.add(
+        &filter::find("buffer").ok_or_else(|| ThumbError::NullVideo)?,
+        "in",
+        &args,
+    )?;
+    filter.add(
+        &filter::find("buffersink").ok_or_else(|| ThumbError::NullVideo)?,
+        "out",
+        "",
+    )?;
+    filter
+        .output("in", 0)?
+        .input("out", 0)?
+        .parse("format=rgba")?;
+    filter.validate()?;
+
+    let stream_index = stream.index();
+
+    let packets = input
+        .packets()
+        .filter(|(s, _)| s.index() == stream_index)
+        .map(|(_, p)| p);
+
+    let mut frame = Video::empty();
+    let mut output_frame = Video::empty();
+    let mut count = 0;
+
+    for packet in packets {
+        decoder.send_packet(&packet)?;
+        while let Err(ffmpeg_next::Error::DecoderNotFound) = decoder.receive_frame(&mut frame) {}
+
+        if decode_single_frame(&mut filter, &mut frame, &mut output_frame).is_ok() {
+            count += 1;
+        }
+        if count > 2 {
+            // take the second frame because the first one often is just blank
+            break;
         }
     }
-    Err(ThumbError::NullVideo)
+    decoder.send_eof()?;
+    let image = RgbaImage::from_raw(
+        decoder.width(),
+        decoder.height(),
+        output_frame.data(0).to_vec(),
+    )
+    .ok_or_else(|| ThumbError::NullVideo)?;
+    let image = DynamicImage::ImageRgba8(image);
+
+    Ok(image)
+}
+
+fn decode_single_frame(
+    filter: &mut Graph,
+    frame: &mut Video,
+    mut output_frame: &mut Video,
+) -> ThumbResult<()> {
+    filter
+        .get("in")
+        .ok_or_else(|| ThumbError::NullVideo)?
+        .source()
+        .add(&frame)?;
+    let mut ctx = filter.get("out").ok_or_else(|| ThumbError::NullVideo)?;
+    let mut out = ctx.sink();
+    out.frame(&mut output_frame)?;
+
+    Ok(())
 }