larksuite / rsmpeg

A Rust crate that exposes FFmpeg's power as much as possible.
https://docs.rs/rsmpeg/latest/rsmpeg/
MIT License
644 stars 40 forks source link

how to convert `rsmpeg::avutil::AVFrame` to `opencv::core::Mat`, #194

Open phial3 opened 15 hours ago

phial3 commented 15 hours ago

there's my code:

use anyhow::{Context, Result};
use opencv::{
    core::Mat,
    prelude::*,
};
use rsmpeg::avcodec::{AVCodec, AVCodecContext};
use rsmpeg::avformat::AVFormatContextInput;
use rsmpeg::avutil::AVFrame;
use rsmpeg::ffi;
use rsmpeg::swscale::SwsContext;
use std::ffi::CString;

/// to mat
fn avframe_to_mat(frame: &AVFrame) -> Result<Mat> {
    let width = frame.width as i32;
    let height = frame.height as i32;
    let data = unsafe { std::slice::from_raw_parts(frame.data[0], (width * height * 3) as usize) };

    Ok(Mat::new_rows_cols_with_data(height, width, data).unwrap().try_clone()?)
}

fn main() -> Result<()> {
    // Path to the image file
    let image_path = "assets/test.jpg";

    // Convert the file path to a CString
    let c_image_path = CString::new(image_path).expect("Failed to create CString");

    // Open the input file
    let mut input_context = AVFormatContextInput::open(&c_image_path, None, &mut None)?;

    // Find the first video stream
    let video_stream_index = input_context
        .streams()
        .iter()
        .position(|stream| stream.codecpar().codec_type == ffi::AVMEDIA_TYPE_VIDEO)
        .expect("No video stream found");

    let stream = &input_context.streams()[video_stream_index];
    let codec_id = stream.codecpar().codec_id;
    println!("Found video stream with index: {}, codec_id: {}", video_stream_index, codec_id);

    let codec = AVCodec::find_decoder(codec_id).expect("Failed to find codec");
    let mut codec_context = AVCodecContext::new(&codec);

    // Open the codec context
    codec_context.open(None)?;

    // Set up the SwsContext for converting the frame to RGB
    let mut sws_context = SwsContext::get_context(
        codec_context.width,
        codec_context.height,
        ffi::AV_PIX_FMT_YUV420P,
        codec_context.width,
        codec_context.height,
        ffi::AV_PIX_FMT_RGB24,
        ffi::SWS_BILINEAR,
        None,
        None,
        None,
    ).context("Failed to create SwsContext")?;

    let mut rgb_frame = AVFrame::new();
    rgb_frame.set_format(ffi::AV_PIX_FMT_RGB24);
    rgb_frame.set_width(codec_context.width);
    rgb_frame.set_height(codec_context.height);
    rgb_frame.alloc_buffer()?;

    // Read frames from the file
    loop {
        let packet = input_context.read_packet()?;
        if let Some(packet) = packet {
            if packet.stream_index as i32 == video_stream_index as i32 {
                // Send the packet to the decoder
                codec_context.send_packet(Some(&packet))?;

                println!("Received packet with size: {}, video_stream_index: {}", packet.size, video_stream_index);

                // Receive the frame from the decoder
                while let Ok(f) = codec_context.receive_frame() {
                    // Convert the frame to RGB format
                    sws_context.scale_frame(&f, rgb_frame.width, rgb_frame.height, &mut rgb_frame)?;

                    // Convert the RGB frame to OpenCV Mat
                    let mat = avframe_to_mat(&rgb_frame)?;
                    println!("Converted frame to OpenCV Mat with size: {:?}", mat.size()?);
                }
            }
        }
    }
}

error info:

Found video stream with index: 0, codec_id: 7
[swscaler @ 0x150008000] 0x0 -> 0x0 is invalid scaling dimension
thread 'main' panicked at examples/avframe_to_mat.rs:62:7:
Failed to create SwsContext
ldm0 commented 14 hours ago

There are several issues in your code:

  1. You haven't set any parameters for codec_context, which is why codec_context.width and codec_context.height are both zero. You need to call codec_context.apply_codecpar(&stream.codecpar())?; before opening the codec_context.
  2. The second parameter of scale_frame( should be zero since your image doesn't have stride: https://github.com/FFmpeg/FFmpeg/blob/2eef902d38dded68df7d874bc348aaa42ec87933/libswscale/swscale.h#L235
  3. Make sure to break out of your loop when packet is None.

You can refer to the example here.