dump-frames.rs

/// mime = "0.3.0"
/// image = "0.24.0"
/// ffmpeg-next = "5.0.2"

extern crate ffmpeg_next as ffmpeg;

use image::{ImageBuffer, Rgb};
use ffmpeg::{
    util::frame::video::Video,
    format::{input, Pixel},
    software::scaling::{context::Context, flag::Flags},
    codec::context::Context as CodecContext,
    media::Type as MediaType,
};
use std::{env, io::Write, path::Path};

type FrameBuffer = ImageBuffer<Rgb<u8>, Vec<u8>>;

fn main() -> Result<(), ffmpeg::Error> {
    ffmpeg::init().unwrap();
    dump_frames()?;
    Ok(())
}

fn dump_frames() -> Result<(), ffmpeg::Error> {
    let ictx = &env::args().nth(1).expect("please specify input video as argument 1");
    let mut ictx = input(ictx).unwrap();

    let path = &env::args().nth(2).expect("please specify output directory as argument 2");
    let path = Path::new(path);
   

    let input = ictx.streams().find(|stream| {
        let codec = CodecContext::from_parameters(stream.parameters()).unwrap();
        codec.medium() == MediaType::Video
    }).ok_or(ffmpeg::Error::StreamNotFound)?;

    let video_stream_index = input.index();

    let context_decoder = CodecContext::from_parameters(input.parameters())?;
    let mut decoder = context_decoder.decoder().video()?;
    let mut scaler = Context::get(
        decoder.format(),
        decoder.width(),
        decoder.height(),
        Pixel::RGB24,
        decoder.width(),
        decoder.height(),
        Flags::BILINEAR,
    )?;

    let mut frame_index = 0;

    let mut receive_and_process_decoded_frames =
        |decoder: &mut ffmpeg::decoder::Video| -> Result<(), ffmpeg::Error> {
            let mut decoded = Video::empty();
            while decoder.receive_frame(&mut decoded).is_ok() {
                let mut rgb_frame = Video::empty();
                scaler.run(&decoded, &mut rgb_frame)?;
                save_file(&rgb_frame, &path, frame_index).unwrap();
                frame_index += 1;
            }
            Ok(())
        };

    for (stream, packet) in ictx.packets() {
        if stream.index() == video_stream_index {
            decoder.send_packet(&packet)?;
            receive_and_process_decoded_frames(&mut decoder)?;
        }
    }
    decoder.send_eof()?;
    receive_and_process_decoded_frames(&mut decoder)?;

    Ok(())
}

fn save_file(frame: &Video, path: &Path, index: usize) -> std::result::Result<(), std::io::Error> {
    let data = frame.data(0);
    let stride = frame.stride(0);
    let byte_width: usize = 3 * frame.width() as usize;
    let height: usize = frame.height() as usize;

    let mut buf: Vec<u8> = vec![];
    buf.reserve((frame.width() * frame.height() * 3) as usize);
    for line in 0..height {
        let begin = line * stride;
        let end = begin + byte_width;
        buf.write_all(&data[begin..end])?;
    }

    let img = FrameBuffer::from_raw(
        frame.width(),
        frame.height(),
        buf,
    ).expect("failed to read buffer");
    
    if !path.exists() {
        std::fs::create_dir_all(path).expect("couldn't create dump directory")
    }
    img.save(path.join(format!("frame-{:05}.jpg", index))).unwrap();
    Ok(())
}