使用rust实现rtsp码流截图

中文互联网上的rust示例程序源码还是太稀少，找资料很是麻烦，下面是自己用业余时间开发实现的一个对批量rtsp码流源进行关键帧截图并存盘的rust demo源码记录。

要编译这个源码需要先安装vcpkg，然后用vcpkg install ffmpeg安装最新版本的ffmpeg库，当然了，你要是想vcpkg成功编译安装ffmpeg，vc编译器和windows sdk也是必不可少的，这些对于做rust windows开发的人来说都不是事，还有llvm及clang windows编译器环境也要安装，这都是准备工作。

代码使用了ffmpeg-next库，这个库在ubuntu 22上面使用sudo apt install 的ffmpeg相关libdev包和windows不一样，ubuntu 22里面默认是ffmpeg 4.3，windows平台默认是ffmpeg 7.0.2 ,这就导致了在跨平台编译的时候会出现问题，linux平台获取video decodec解码器和windows平台不一样，代码里面注释掉的内容就是在linux平台编译的时候要使用的函数，如果要在linux平台且使用ffmpeg 4.x版本编译注意打开注释掉的内容。

use ffmpeg_next as ffmpeg;
use tokio;
use std::sync::Arc;
use tokio::sync::Semaphore;
use std::error::Error;
use image::{ImageBuffer, Rgb};
use std::fmt;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use ffmpeg::format::input;
use ffmpeg::software::scaling::{context::Context, flag::Flags};
use ffmpeg::util::frame::video::Video;
use ffmpeg::format::stream::Stream;

#[derive(Debug)]
enum CustomError {
    FfmpegError(ffmpeg::Error),
    ImageError(image::ImageError),
    Other(String),
}

impl fmt::Display for CustomError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CustomError::FfmpegError(e) => write!(f, "FFmpeg error: {}", e),
            CustomError::ImageError(e) => write!(f, "Image error: {}", e),
            CustomError::Other(e) => write!(f, "Other error: {}", e),
        }
    }
}

impl std::error::Error for CustomError {}

impl From<ffmpeg::Error> for CustomError {
    fn from(error: ffmpeg::Error) -> Self {
        CustomError::FfmpegError(error)
    }
}

impl From<image::ImageError> for CustomError {
    fn from(error: image::ImageError) -> Self {
        CustomError::ImageError(error)
    }
}

impl From<&str> for CustomError {
    fn from(error: &str) -> Self {
        CustomError::Other(error.to_string())
    }
}

struct RtspSource {
    url: String,
}


fn get_decoder(input_stream: &Stream) -> Result<ffmpeg::decoder::Video, ffmpeg::Error> {
    let decoder_params = input_stream.parameters();
    let mut ctx = ffmpeg::codec::context::Context::new();
    ctx.set_parameters(decoder_params)?;
    ctx.decoder().video()
}

// #[cfg(not(feature = "ffmpeg_5_0"))]
// fn get_decoder(input_stream: &Stream) -> Result<ffmpeg::decoder::Video, ffmpeg::Error> {
//     input_stream.codec().decoder().video()
// }

async fn capture_frame(source: &RtspSource, frame_counter: Arc<AtomicUsize>) -> Result<(), Box<dyn Error>> {
    let mut ictx = input(&source.url)?;
    let input_stream = ictx
        .streams()
        .best(ffmpeg::media::Type::Video)
        .ok_or("Could not find best video stream")?;
    let video_stream_index = input_stream.index();

    let mut decoder = get_decoder(&input_stream)?;

    let mut scaler = Context::get(
        decoder.format(),
        decoder.width(),
        decoder.height(),
        ffmpeg::format::Pixel::RGB24,
        decoder.width(),
        decoder.height(),
        Flags::BILINEAR,
    )?;

    let mut frame = Video::empty();
    let current_path = std::env::current_dir()?;

    for (stream, packet) in ictx.packets() {
        if stream.index() == video_stream_index && packet.is_key() {
            decoder.send_packet(&packet)?;
            while decoder.receive_frame(&mut frame).is_ok() {
                let mut rgb_frame = Video::empty();
                scaler.run(&frame, &mut rgb_frame)?;

                let buffer = rgb_frame.data(0);
                let width = rgb_frame.width() as u32;
                let height = rgb_frame.height() as u32;
                let img: ImageBuffer<Rgb<u8>, _> =
                    ImageBuffer::from_raw(width, height, buffer.to_owned())
                        .ok_or("Failed to create image buffer")?;

                let index = frame_counter.fetch_add(1, Ordering::SeqCst);
                let file_save_name = format!("captured_frame_{}.jpg", index);
                let save_path: PathBuf = current_path.join("./images/").join(&file_save_name);
                img.save(&save_path)?;
                println!("Frame captured and saved to {}", save_path.display());
                return Ok(());
            }
        }
    }

    Ok(())
}

async fn process_sources(sources: Vec<RtspSource>, max_concurrent: usize) -> Result<(), Box<dyn Error>> {
    let semaphore = Arc::new(Semaphore::new(max_concurrent));
    let frame_counter = Arc::new(AtomicUsize::new(0));

    let mut handles = vec![];

    for source in sources {
        let permit = semaphore.clone().acquire_owned().await?;
        let frame_counter_clone = Arc::clone(&frame_counter);
        let handle = tokio::spawn(async move {
            let result = capture_frame(&source, frame_counter_clone).await;
            match result {
                Ok(_) => println!("Successfully captured frame from {}", source.url),
                Err(e) => eprintln!("Error capturing frame from {}: {}", source.url, e),
            }
            drop(permit);
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.await?;
    }

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    ffmpeg::init()?;

    let mut sources:Vec<RtspSource>=Vec::with_capacity(100);

    for _ in 1..=100 {
        sources.push(RtspSource {
            url: format!("rtsp://你的rtsp源ip地址:8554/stream"),
        });
    }

    let max_concurrent = 20; // Set the maximum number of concurrent captures
    let start_time = tokio::time::Instant::now();
    process_sources(sources, max_concurrent).await?;
    let end_time = tokio::time::Instant::now();
    println!("Time taken to capture frames: {:?}", end_time.duration_since(start_time));

    Ok(())
}

本文发表于https://blog.csdn.net/peihexian，欢迎转载，当博客写完的时候我想到一个问题，那就是其实是不是可以通过调用ffmpeg.exe命令行的方式传参实现截图的抓取，不过在实现上面的算法中我尝试了连上rtsp源头以后立马抓第一帧图像就存盘是不行的，因为没有关键帧数据，第一帧抓到的是乱码，所以代码里面改成了抓关键帧，这样存盘的时候肯定是完整的图像，不知道使用命令行方式传参的方式能不能解决取关键帧的问题。

补充一下Cargo.toml的文件内容:

[package]
name = "ffmpeg-test1"
version = "0.1.0"
edition = "2021"

[dependencies]
ffmpeg-next = { version = "7.0" }
tokio = { version = "1.0", features = ["full"] }
image = "0.25"