diff --git a/Cargo.toml b/Cargo.toml index 72ea522..ce1f598 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ readme = "README.md" license = "Apache-2.0" authors = ["trivernis "] 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] diff --git a/src/error.rs b/src/error.rs index 8367325..ad4e59b 100644 --- a/src/error.rs +++ b/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 = Result; @@ -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 for ThumbError { } } -impl From for ThumbError { - fn from(e: CaptureError) -> Self { - Self::CaptureError(e) - } -} - -impl From for ThumbError { - fn from(s: StreamError) -> Self { - Self::StreamError(s) +impl From for ThumbError { + fn from(e: ffmpeg_next::Error) -> Self { + Self::FFMPEG(e) } } diff --git a/src/formats/video_format.rs b/src/formats/video_format.rs index b8e6669..a642f2f 100644 --- a/src/formats/video_format.rs +++ b/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(mut reader: R, mime: Mime) -> ThumbResult { let tempdir = tempfile::tempdir()?; @@ -25,13 +28,91 @@ pub fn get_video_frame(mut reader: R, mime: Mime) -> ThumbRes } fn extract_frame_from_video(path: &PathBuf) -> ThumbResult { - 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(()) }