Kieran/ffmpeg-rs-raw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: stream groups

Browse Source
This commit is contained in:
kieran 2024-11-11 14:10:03 +00:00
parent b66a8db0c6
commit 363ad4f55c
No known key found for this signature in database
GPG Key ID: DE71CEB3925BE941
6 changed files with 154 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
examples/main.rs
View File

@ -43,7 +43,7 @@ unsafe fn decode_input(demuxer: Demuxer, info: DemuxerInfo) {
let mut decoder = Decoder::new();
decoder.enable_hw_decoder_any();
for ref stream in info.channels {
for ref stream in info.streams {
decoder
.setup_decoder(stream, None)
.expect("decoder setup failed");
@ -66,7 +66,7 @@ unsafe fn loop_decoder(mut demuxer: Demuxer, mut decoder: Decoder) {
av_packet_free(&mut pkt);
continue;
}
if let Ok(frames) = decoder.decode_pkt(pkt, stream) {
if let Ok(frames) = decoder.decode_pkt(pkt) {
for mut frame in frames {
// do nothing but decode entire stream
if media_type == AVMediaType::AVMEDIA_TYPE_VIDEO {

28
src/decode.rs
View File

@ -1,4 +1,4 @@
use crate::{bail_ffmpeg, options_to_dict, rstr, StreamInfoChannel};
use crate::{bail_ffmpeg, options_to_dict, rstr, StreamInfo};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
@ -117,7 +117,7 @@ impl Decoder {
/// Set up a decoder for a given channel
pub fn setup_decoder(
&mut self,
channel: &StreamInfoChannel,
channel: &StreamInfo,
options: Option<HashMap<String, String>>,
) -> Result<&mut DecoderCodecContext, Error> {
unsafe { self.setup_decoder_for_stream(channel.stream, options) }
@ -241,11 +241,7 @@ impl Decoder {
pub unsafe fn flush(&mut self) -> Result<Vec<*mut AVFrame>, Error> {
let mut pkgs = Vec::new();
for ctx in self.codecs.values_mut() {
pkgs.extend(Self::decode_pkt_internal(
ctx.context,
ptr::null_mut(),
ctx.stream,
)?);
pkgs.extend(Self::decode_pkt_internal(ctx.context, ptr::null_mut())?);
}
Ok(pkgs)
}
@ -253,7 +249,6 @@ impl Decoder {
pub unsafe fn decode_pkt_internal(
ctx: *mut AVCodecContext,
pkt: *mut AVPacket,
stream: *mut AVStream,
) -> Result<Vec<*mut AVFrame>, Error> {
let mut ret = avcodec_send_packet(ctx, pkt);
bail_ffmpeg!(ret, "Failed to decode packet");
@ -275,23 +270,12 @@ impl Decoder {
Ok(pkgs)
}
pub unsafe fn decode_pkt(
&mut self,
pkt: *mut AVPacket,
stream: *mut AVStream,
) -> Result<Vec<*mut AVFrame>, Error> {
pub unsafe fn decode_pkt(&mut self, pkt: *mut AVPacket) -> Result<Vec<*mut AVFrame>, Error> {
if pkt.is_null() {
return self.flush();
}
let stream_index = (*pkt).stream_index;
assert_eq!(
stream_index,
(*stream).index,
"Passed stream reference does not match stream_index of packet"
);
if let Some(ctx) = self.codecs.get_mut(&stream_index) {
Self::decode_pkt_internal(ctx.context, pkt, stream)
if let Some(ctx) = self.codecs.get_mut(&(*pkt).stream_index) {
Self::decode_pkt_internal(ctx.context, pkt)
} else {
Ok(vec![])
}

89
src/demux.rs
View File

@ -1,7 +1,9 @@
use crate::{bail_ffmpeg, cstr};
use crate::{DemuxerInfo, StreamChannelType, StreamInfoChannel};
use crate::{bail_ffmpeg, cstr, rstr, StreamGroupInfo, StreamGroupType};
use crate::{DemuxerInfo, StreamInfo, StreamType};
use anyhow::{bail, Error, Result};
use ffmpeg_sys_the_third::AVStreamGroupParamsType::AV_STREAM_GROUP_PARAMS_TILE_GRID;
use ffmpeg_sys_the_third::*;
use log::warn;
use slimbox::{slimbox_unsize, SlimBox, SlimMut};
use std::collections::HashMap;
use std::io::Read;
@ -114,43 +116,73 @@ impl Demuxer {
return Err(Error::msg("Could not find stream info"));
}
let mut channel_infos = vec![];
let mut streams = vec![];
let mut stream_groups = vec![];
for n in 0..(*self.ctx).nb_streams as usize {
let stream = *(*self.ctx).streams.add(n);
let mut n_stream = 0;
for n in 0..(*self.ctx).nb_stream_groups as usize {
let group = *(*self.ctx).stream_groups.add(n);
n_stream += (*group).nb_streams as usize;
match (*group).type_ {
AV_STREAM_GROUP_PARAMS_TILE_GRID => {
let tg = (*group).params.tile_grid;
let codec_par = (*(*(*group).streams.add(0))).codecpar;
stream_groups.push(StreamGroupInfo {
index: (*group).index as usize,
group_type: StreamGroupType::TileGrid {
tiles: (*tg).nb_tiles as usize,
width: (*tg).width as usize,
height: (*tg).height as usize,
format: (*codec_par).format as isize,
codec: (*codec_par).codec_id as isize,
},
group,
});
}
t => {
warn!(
"Unsupported stream group type {}, skipping.",
rstr!(avformat_stream_group_name(t))
);
}
}
}
while n_stream < (*self.ctx).nb_streams as usize {
let stream = *(*self.ctx).streams.add(n_stream);
n_stream += 1;
match (*(*stream).codecpar).codec_type {
AVMediaType::AVMEDIA_TYPE_VIDEO => {
channel_infos.push(StreamInfoChannel {
streams.push(StreamInfo {
stream,
index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize,
channel_type: StreamChannelType::Video,
codec: (*(*stream).codecpar).codec_id as isize,
stream_type: StreamType::Video,
width: (*(*stream).codecpar).width as usize,
height: (*(*stream).codecpar).height as usize,
fps: av_q2d((*stream).avg_frame_rate) as f32,
format: (*(*stream).codecpar).format as usize,
format: (*(*stream).codecpar).format as isize,
sample_rate: 0,
});
}
AVMediaType::AVMEDIA_TYPE_AUDIO => {
channel_infos.push(StreamInfoChannel {
streams.push(StreamInfo {
stream,
index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize,
channel_type: StreamChannelType::Audio,
codec: (*(*stream).codecpar).codec_id as isize,
stream_type: StreamType::Audio,
width: (*(*stream).codecpar).width as usize,
height: (*(*stream).codecpar).height as usize,
fps: 0.0,
format: (*(*stream).codecpar).format as usize,
format: (*(*stream).codecpar).format as isize,
sample_rate: (*(*stream).codecpar).sample_rate as usize,
});
}
AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
channel_infos.push(StreamInfoChannel {
streams.push(StreamInfo {
stream,
index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize,
channel_type: StreamChannelType::Subtitle,
codec: (*(*stream).codecpar).codec_id as isize,
stream_type: StreamType::Subtitle,
width: 0,
height: 0,
fps: 0.0,
@ -167,7 +199,8 @@ impl Demuxer {
let info = DemuxerInfo {
duration: (*self.ctx).duration as f32 / AV_TIME_BASE as f32,
bitrate: (*self.ctx).bit_rate as usize,
channels: channel_infos,
streams,
groups: stream_groups,
};
Ok(info)
}
@ -200,3 +233,25 @@ impl Drop for Demuxer {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[ignore] // WIP
fn test_stream_groups() -> Result<()> {
unsafe {
let mut demux =
Demuxer::new("https://trac.ffmpeg.org/raw-attachment/ticket/11170/IMG_4765.HEIC")?;
let probe = demux.probe_input()?;
assert_eq!(1, probe.streams.len());
assert_eq!(1, probe.groups.len());
assert!(matches!(
probe.groups[0].group_type,
StreamGroupType::TileGrid { .. }
));
}
Ok(())
}
}

6
src/muxer.rs
View File

@ -3,8 +3,8 @@ use anyhow::{bail, Result};
use ffmpeg_sys_the_third::{
av_dump_format, av_interleaved_write_frame, av_packet_rescale_ts, av_write_trailer,
avcodec_parameters_copy, avcodec_parameters_from_context, avformat_alloc_output_context2,
avformat_free_context, avformat_new_stream, avformat_write_header, avio_flush, avio_open,
AVFormatContext, AVPacket, AVStream, AVFMT_GLOBALHEADER, AVFMT_NOFILE, AVIO_FLAG_WRITE,
avformat_free_context, avformat_new_stream, avformat_write_header, avio_open, AVFormatContext,
AVPacket, AVStream, AVFMT_GLOBALHEADER, AVFMT_NOFILE, AVIO_FLAG_WRITE,
AV_CODEC_FLAG_GLOBAL_HEADER,
};
use std::collections::HashMap;
@ -185,7 +185,7 @@ mod tests {
.with_stream_encoder(&encoder)?;
muxer.open()?;
let mut pts = 0;
for z in 0..100 {
for _z in 0..100 {
(*frame).pts = pts;
for pkt in encoder.encode_frame(frame)? {
muxer.write_packet(pkt)?;

93
src/stream_info.rs
View File

@ -1,6 +1,7 @@
use crate::{format_time, rstr};
use ffmpeg_sys_the_third::{
av_get_pix_fmt_name, av_get_sample_fmt_name, avcodec_get_name, AVMediaType, AVStream,
AVStreamGroup,
};
use std::fmt::{Display, Formatter};
use std::intrinsics::transmute;
@ -9,14 +10,15 @@ use std::intrinsics::transmute;
pub struct DemuxerInfo {
pub bitrate: usize,
pub duration: f32,
pub channels: Vec<StreamInfoChannel>,
pub streams: Vec<StreamInfo>,
pub groups: Vec<StreamGroupInfo>,
}
impl DemuxerInfo {
pub fn best_stream(&self, t: StreamChannelType) -> Option<&StreamInfoChannel> {
self.channels
pub fn best_stream(&self, t: StreamType) -> Option<&StreamInfo> {
self.streams
.iter()
.filter(|a| a.channel_type == t)
.filter(|a| a.stream_type == t)
.reduce(|acc, channel| {
if channel.best_metric() > acc.best_metric() {
channel
@ -26,16 +28,16 @@ impl DemuxerInfo {
})
}
pub fn best_video(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Video)
pub fn best_video(&self) -> Option<&StreamInfo> {
self.best_stream(StreamType::Video)
}
pub fn best_audio(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Audio)
pub fn best_audio(&self) -> Option<&StreamInfo> {
self.best_stream(StreamType::Audio)
}
pub fn best_subtitle(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Subtitle)
pub fn best_subtitle(&self) -> Option<&StreamInfo> {
self.best_stream(StreamType::Subtitle)
}
pub unsafe fn is_best_stream(&self, stream: *mut AVStream) -> bool {
@ -71,7 +73,7 @@ impl Display for DemuxerInfo {
format_time(self.duration),
bitrate_str
)?;
for c in &self.channels {
for c in &self.streams {
write!(f, "\n {}", c)?;
}
Ok(())
@ -79,62 +81,61 @@ impl Display for DemuxerInfo {
}
#[derive(Clone, Debug, PartialEq)]
pub enum StreamChannelType {
pub enum StreamType {
Video,
Audio,
Subtitle,
}
impl Display for StreamChannelType {
impl Display for StreamType {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}",
match self {
StreamChannelType::Video => "video",
StreamChannelType::Audio => "audio",
StreamChannelType::Subtitle => "subtitle",
StreamType::Video => "video",
StreamType::Audio => "audio",
StreamType::Subtitle => "subtitle",
}
)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct StreamInfoChannel {
pub struct StreamInfo {
pub index: usize,
pub channel_type: StreamChannelType,
pub codec: usize,
pub stream_type: StreamType,
pub codec: isize,
pub format: isize,
pub width: usize,
pub height: usize,
pub fps: f32,
pub sample_rate: usize,
pub format: usize,
// private stream pointer
pub(crate) stream: *mut AVStream,
}
unsafe impl Send for StreamInfoChannel {}
unsafe impl Sync for StreamInfoChannel {}
impl StreamInfoChannel {
impl StreamInfo {
pub fn best_metric(&self) -> f32 {
match self.channel_type {
StreamChannelType::Video => self.width as f32 * self.height as f32 * self.fps,
StreamChannelType::Audio => self.sample_rate as f32,
StreamChannelType::Subtitle => 999. - self.index as f32,
match self.stream_type {
StreamType::Video => self.width as f32 * self.height as f32 * self.fps,
StreamType::Audio => self.sample_rate as f32,
StreamType::Subtitle => 999. - self.index as f32,
}
}
}
impl Display for StreamInfoChannel {
impl Display for StreamInfo {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let codec_name = unsafe { rstr!(avcodec_get_name(transmute(self.codec as i32))) };
match self.channel_type {
StreamChannelType::Video => write!(
match self.stream_type {
StreamType::Video => write!(
f,
"{} #{}: codec={},size={}x{},fps={:.3},pix_fmt={}",
self.channel_type,
self.stream_type,
self.index,
codec_name,
self.width,
@ -142,10 +143,10 @@ impl Display for StreamInfoChannel {
self.fps,
unsafe { rstr!(av_get_pix_fmt_name(transmute(self.format as libc::c_int))) },
),
StreamChannelType::Audio => write!(
StreamType::Audio => write!(
f,
"{} #{}: codec={},format={},sample_rate={}",
self.channel_type,
self.stream_type,
self.index,
codec_name,
unsafe {
@ -155,11 +156,31 @@ impl Display for StreamInfoChannel {
},
self.sample_rate,
),
StreamChannelType::Subtitle => write!(
StreamType::Subtitle => write!(
f,
"{} #{}: codec={}",
self.channel_type, self.index, codec_name
self.stream_type, self.index, codec_name
),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum StreamGroupType {
TileGrid {
tiles: usize,
width: usize,
height: usize,
codec: isize,
format: isize,
},
}
#[derive(Clone, Debug, PartialEq)]
pub struct StreamGroupInfo {
pub index: usize,
pub group_type: StreamGroupType,
// private pointer
pub(crate) group: *mut AVStreamGroup,
}

23
src/transcode.rs
View File

@ -1,4 +1,4 @@
use crate::{Decoder, Demuxer, DemuxerInfo, Encoder, Muxer, Scaler, StreamInfoChannel};
use crate::{Decoder, Demuxer, DemuxerInfo, Encoder, Muxer, Scaler, StreamInfo};
use anyhow::Result;
use ffmpeg_sys_the_third::{av_frame_free, av_packet_free};
use std::collections::HashMap;
@ -39,7 +39,7 @@ impl Transcoder {
/// a pre-configured output encoder
pub unsafe fn transcode_stream(
&mut self,
in_stream: &StreamInfoChannel,
in_stream: &StreamInfo,
encoder_out: Encoder,
) -> Result<()> {
let src_index = in_stream.index as i32;
@ -47,7 +47,8 @@ impl Transcoder {
let out_ctx = encoder_out.codec_context();
if in_stream.width != (*out_ctx).width as usize
|| in_stream.height != (*out_ctx).height as usize
|| in_stream.format != (*out_ctx).pix_fmt as usize {
|| in_stream.format != (*out_ctx).pix_fmt as isize
{
// Setup scaler if the size/format is different from what the codec expects
self.scalers.insert(src_index, Scaler::new());
}
@ -60,7 +61,7 @@ impl Transcoder {
}
/// Copy a stream from the input to the output
pub unsafe fn copy_stream(&mut self, in_stream: StreamInfoChannel) -> Result<()> {
pub unsafe fn copy_stream(&mut self, in_stream: StreamInfo) -> Result<()> {
let dst_stream = self.muxer.add_copy_stream(in_stream.stream)?;
self.copy_stream
.insert(in_stream.index as i32, (*dst_stream).index);
@ -84,12 +85,16 @@ impl Transcoder {
let src_index = (*stream).index;
// check if encoded stream
if let Some(enc) = self.encoders.get_mut(&src_index) {
for mut frame in self.decoder.decode_pkt(pkt, stream)? {
for mut frame in self.decoder.decode_pkt(pkt)? {
// scale frame before sending to encoder
let mut frame = if let Some(mut sws) = self.scalers.get_mut(&src_index) {
let mut frame = if let Some(sws) = self.scalers.get_mut(&src_index) {
let enc_ctx = enc.codec_context();
let new_frame = sws.process_frame(frame, (*enc_ctx).width as u16, (*enc_ctx).height as u16, (*enc_ctx).pix_fmt)?;
let new_frame = sws.process_frame(
frame,
(*enc_ctx).width as u16,
(*enc_ctx).height as u16,
(*enc_ctx).pix_fmt,
)?;
av_frame_free(&mut frame);
new_frame
} else {
@ -135,7 +140,7 @@ mod tests {
let mut transcoder =
Transcoder::new("test_output/test.mp4", "test_output/test_transcode.mkv")?;
let info = transcoder.prepare()?;
for c in info.channels {
for c in info.streams {
transcoder.copy_stream(c)?;
}
transcoder.run()?;