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(); let mut decoder = Decoder::new();
decoder.enable_hw_decoder_any(); decoder.enable_hw_decoder_any();
for ref stream in info.channels { for ref stream in info.streams {
decoder decoder
.setup_decoder(stream, None) .setup_decoder(stream, None)
.expect("decoder setup failed"); .expect("decoder setup failed");
@ -66,7 +66,7 @@ unsafe fn loop_decoder(mut demuxer: Demuxer, mut decoder: Decoder) {
av_packet_free(&mut pkt); av_packet_free(&mut pkt);
continue; continue;
} }
if let Ok(frames) = decoder.decode_pkt(pkt, stream) { if let Ok(frames) = decoder.decode_pkt(pkt) {
for mut frame in frames { for mut frame in frames {
// do nothing but decode entire stream // do nothing but decode entire stream
if media_type == AVMediaType::AVMEDIA_TYPE_VIDEO { 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::hash_map::Entry;
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter}; use std::fmt::{Display, Formatter};
@ -117,7 +117,7 @@ impl Decoder {
/// Set up a decoder for a given channel /// Set up a decoder for a given channel
pub fn setup_decoder( pub fn setup_decoder(
&mut self, &mut self,
channel: &StreamInfoChannel, channel: &StreamInfo,
options: Option<HashMap<String, String>>, options: Option<HashMap<String, String>>,
) -> Result<&mut DecoderCodecContext, Error> { ) -> Result<&mut DecoderCodecContext, Error> {
unsafe { self.setup_decoder_for_stream(channel.stream, options) } 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> { pub unsafe fn flush(&mut self) -> Result<Vec<*mut AVFrame>, Error> {
let mut pkgs = Vec::new(); let mut pkgs = Vec::new();
for ctx in self.codecs.values_mut() { for ctx in self.codecs.values_mut() {
pkgs.extend(Self::decode_pkt_internal( pkgs.extend(Self::decode_pkt_internal(ctx.context, ptr::null_mut())?);
ctx.context,
ptr::null_mut(),
ctx.stream,
)?);
} }
Ok(pkgs) Ok(pkgs)
} }
@ -253,7 +249,6 @@ impl Decoder {
pub unsafe fn decode_pkt_internal( pub unsafe fn decode_pkt_internal(
ctx: *mut AVCodecContext, ctx: *mut AVCodecContext,
pkt: *mut AVPacket, pkt: *mut AVPacket,
stream: *mut AVStream,
) -> Result<Vec<*mut AVFrame>, Error> { ) -> Result<Vec<*mut AVFrame>, Error> {
let mut ret = avcodec_send_packet(ctx, pkt); let mut ret = avcodec_send_packet(ctx, pkt);
bail_ffmpeg!(ret, "Failed to decode packet"); bail_ffmpeg!(ret, "Failed to decode packet");
@ -275,23 +270,12 @@ impl Decoder {
Ok(pkgs) Ok(pkgs)
} }
pub unsafe fn decode_pkt( pub unsafe fn decode_pkt(&mut self, pkt: *mut AVPacket) -> Result<Vec<*mut AVFrame>, Error> {
&mut self,
pkt: *mut AVPacket,
stream: *mut AVStream,
) -> Result<Vec<*mut AVFrame>, Error> {
if pkt.is_null() { if pkt.is_null() {
return self.flush(); return self.flush();
} }
let stream_index = (*pkt).stream_index; if let Some(ctx) = self.codecs.get_mut(&(*pkt).stream_index) {
assert_eq!( Self::decode_pkt_internal(ctx.context, pkt)
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)
} else { } else {
Ok(vec![]) Ok(vec![])
} }

89
src/demux.rs
View File

@ -1,7 +1,9 @@
use crate::{bail_ffmpeg, cstr}; use crate::{bail_ffmpeg, cstr, rstr, StreamGroupInfo, StreamGroupType};
use crate::{DemuxerInfo, StreamChannelType, StreamInfoChannel}; use crate::{DemuxerInfo, StreamInfo, StreamType};
use anyhow::{bail, Error, Result}; use anyhow::{bail, Error, Result};
use ffmpeg_sys_the_third::AVStreamGroupParamsType::AV_STREAM_GROUP_PARAMS_TILE_GRID;
use ffmpeg_sys_the_third::*; use ffmpeg_sys_the_third::*;
use log::warn;
use slimbox::{slimbox_unsize, SlimBox, SlimMut}; use slimbox::{slimbox_unsize, SlimBox, SlimMut};
use std::collections::HashMap; use std::collections::HashMap;
use std::io::Read; use std::io::Read;
@ -114,43 +116,73 @@ impl Demuxer {
return Err(Error::msg("Could not find stream info")); 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 mut n_stream = 0;
let stream = *(*self.ctx).streams.add(n); 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 { match (*(*stream).codecpar).codec_type {
AVMediaType::AVMEDIA_TYPE_VIDEO => { AVMediaType::AVMEDIA_TYPE_VIDEO => {
channel_infos.push(StreamInfoChannel { streams.push(StreamInfo {
stream, stream,
index: (*stream).index as usize, index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize, codec: (*(*stream).codecpar).codec_id as isize,
channel_type: StreamChannelType::Video, stream_type: StreamType::Video,
width: (*(*stream).codecpar).width as usize, width: (*(*stream).codecpar).width as usize,
height: (*(*stream).codecpar).height as usize, height: (*(*stream).codecpar).height as usize,
fps: av_q2d((*stream).avg_frame_rate) as f32, fps: av_q2d((*stream).avg_frame_rate) as f32,
format: (*(*stream).codecpar).format as usize, format: (*(*stream).codecpar).format as isize,
sample_rate: 0, sample_rate: 0,
}); });
} }
AVMediaType::AVMEDIA_TYPE_AUDIO => { AVMediaType::AVMEDIA_TYPE_AUDIO => {
channel_infos.push(StreamInfoChannel { streams.push(StreamInfo {
stream, stream,
index: (*stream).index as usize, index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize, codec: (*(*stream).codecpar).codec_id as isize,
channel_type: StreamChannelType::Audio, stream_type: StreamType::Audio,
width: (*(*stream).codecpar).width as usize, width: (*(*stream).codecpar).width as usize,
height: (*(*stream).codecpar).height as usize, height: (*(*stream).codecpar).height as usize,
fps: 0.0, fps: 0.0,
format: (*(*stream).codecpar).format as usize, format: (*(*stream).codecpar).format as isize,
sample_rate: (*(*stream).codecpar).sample_rate as usize, sample_rate: (*(*stream).codecpar).sample_rate as usize,
}); });
} }
AVMediaType::AVMEDIA_TYPE_SUBTITLE => { AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
channel_infos.push(StreamInfoChannel { streams.push(StreamInfo {
stream, stream,
index: (*stream).index as usize, index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize, codec: (*(*stream).codecpar).codec_id as isize,
channel_type: StreamChannelType::Subtitle, stream_type: StreamType::Subtitle,
width: 0, width: 0,
height: 0, height: 0,
fps: 0.0, fps: 0.0,
@ -167,7 +199,8 @@ impl Demuxer {
let info = DemuxerInfo { let info = DemuxerInfo {
duration: (*self.ctx).duration as f32 / AV_TIME_BASE as f32, duration: (*self.ctx).duration as f32 / AV_TIME_BASE as f32,
bitrate: (*self.ctx).bit_rate as usize, bitrate: (*self.ctx).bit_rate as usize,
channels: channel_infos, streams,
groups: stream_groups,
}; };
Ok(info) 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::{ use ffmpeg_sys_the_third::{
av_dump_format, av_interleaved_write_frame, av_packet_rescale_ts, av_write_trailer, 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, avcodec_parameters_copy, avcodec_parameters_from_context, avformat_alloc_output_context2,
avformat_free_context, avformat_new_stream, avformat_write_header, avio_flush, avio_open, avformat_free_context, avformat_new_stream, avformat_write_header, avio_open, AVFormatContext,
AVFormatContext, AVPacket, AVStream, AVFMT_GLOBALHEADER, AVFMT_NOFILE, AVIO_FLAG_WRITE, AVPacket, AVStream, AVFMT_GLOBALHEADER, AVFMT_NOFILE, AVIO_FLAG_WRITE,
AV_CODEC_FLAG_GLOBAL_HEADER, AV_CODEC_FLAG_GLOBAL_HEADER,
}; };
use std::collections::HashMap; use std::collections::HashMap;
@ -185,7 +185,7 @@ mod tests {
.with_stream_encoder(&encoder)?; .with_stream_encoder(&encoder)?;
muxer.open()?; muxer.open()?;
let mut pts = 0; let mut pts = 0;
for z in 0..100 { for _z in 0..100 {
(*frame).pts = pts; (*frame).pts = pts;
for pkt in encoder.encode_frame(frame)? { for pkt in encoder.encode_frame(frame)? {
muxer.write_packet(pkt)?; muxer.write_packet(pkt)?;

93
src/stream_info.rs
View File

@ -1,6 +1,7 @@
use crate::{format_time, rstr}; use crate::{format_time, rstr};
use ffmpeg_sys_the_third::{ use ffmpeg_sys_the_third::{
av_get_pix_fmt_name, av_get_sample_fmt_name, avcodec_get_name, AVMediaType, AVStream, av_get_pix_fmt_name, av_get_sample_fmt_name, avcodec_get_name, AVMediaType, AVStream,
AVStreamGroup,
}; };
use std::fmt::{Display, Formatter}; use std::fmt::{Display, Formatter};
use std::intrinsics::transmute; use std::intrinsics::transmute;
@ -9,14 +10,15 @@ use std::intrinsics::transmute;
pub struct DemuxerInfo { pub struct DemuxerInfo {
pub bitrate: usize, pub bitrate: usize,
pub duration: f32, pub duration: f32,
pub channels: Vec<StreamInfoChannel>, pub streams: Vec<StreamInfo>,
pub groups: Vec<StreamGroupInfo>,
} }
impl DemuxerInfo { impl DemuxerInfo {
pub fn best_stream(&self, t: StreamChannelType) -> Option<&StreamInfoChannel> { pub fn best_stream(&self, t: StreamType) -> Option<&StreamInfo> {
self.channels self.streams
.iter() .iter()
.filter(|a| a.channel_type == t) .filter(|a| a.stream_type == t)
.reduce(|acc, channel| { .reduce(|acc, channel| {
if channel.best_metric() > acc.best_metric() { if channel.best_metric() > acc.best_metric() {
channel channel
@ -26,16 +28,16 @@ impl DemuxerInfo {
}) })
} }
pub fn best_video(&self) -> Option<&StreamInfoChannel> { pub fn best_video(&self) -> Option<&StreamInfo> {
self.best_stream(StreamChannelType::Video) self.best_stream(StreamType::Video)
} }
pub fn best_audio(&self) -> Option<&StreamInfoChannel> { pub fn best_audio(&self) -> Option<&StreamInfo> {
self.best_stream(StreamChannelType::Audio) self.best_stream(StreamType::Audio)
} }
pub fn best_subtitle(&self) -> Option<&StreamInfoChannel> { pub fn best_subtitle(&self) -> Option<&StreamInfo> {
self.best_stream(StreamChannelType::Subtitle) self.best_stream(StreamType::Subtitle)
} }
pub unsafe fn is_best_stream(&self, stream: *mut AVStream) -> bool { pub unsafe fn is_best_stream(&self, stream: *mut AVStream) -> bool {
@ -71,7 +73,7 @@ impl Display for DemuxerInfo {
format_time(self.duration), format_time(self.duration),
bitrate_str bitrate_str
)?; )?;
for c in &self.channels { for c in &self.streams {
write!(f, "\n {}", c)?; write!(f, "\n {}", c)?;
} }
Ok(()) Ok(())
@ -79,62 +81,61 @@ impl Display for DemuxerInfo {
} }
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
pub enum StreamChannelType { pub enum StreamType {
Video, Video,
Audio, Audio,
Subtitle, Subtitle,
} }
impl Display for StreamChannelType { impl Display for StreamType {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!( write!(
f, f,
"{}", "{}",
match self { match self {
StreamChannelType::Video => "video", StreamType::Video => "video",
StreamChannelType::Audio => "audio", StreamType::Audio => "audio",
StreamChannelType::Subtitle => "subtitle", StreamType::Subtitle => "subtitle",
} }
) )
} }
} }
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
pub struct StreamInfoChannel { pub struct StreamInfo {
pub index: usize, pub index: usize,
pub channel_type: StreamChannelType, pub stream_type: StreamType,
pub codec: usize, pub codec: isize,
pub format: isize,
pub width: usize, pub width: usize,
pub height: usize, pub height: usize,
pub fps: f32, pub fps: f32,
pub sample_rate: usize, pub sample_rate: usize,
pub format: usize,
// private stream pointer // private stream pointer
pub(crate) stream: *mut AVStream, pub(crate) stream: *mut AVStream,
} }
unsafe impl Send for StreamInfoChannel {} impl StreamInfo {
unsafe impl Sync for StreamInfoChannel {}
impl StreamInfoChannel {
pub fn best_metric(&self) -> f32 { pub fn best_metric(&self) -> f32 {
match self.channel_type { match self.stream_type {
StreamChannelType::Video => self.width as f32 * self.height as f32 * self.fps, StreamType::Video => self.width as f32 * self.height as f32 * self.fps,
StreamChannelType::Audio => self.sample_rate as f32, StreamType::Audio => self.sample_rate as f32,
StreamChannelType::Subtitle => 999. - self.index 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 { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let codec_name = unsafe { rstr!(avcodec_get_name(transmute(self.codec as i32))) }; let codec_name = unsafe { rstr!(avcodec_get_name(transmute(self.codec as i32))) };
match self.channel_type { match self.stream_type {
StreamChannelType::Video => write!( StreamType::Video => write!(
f, f,
"{} #{}: codec={},size={}x{},fps={:.3},pix_fmt={}", "{} #{}: codec={},size={}x{},fps={:.3},pix_fmt={}",
self.channel_type, self.stream_type,
self.index, self.index,
codec_name, codec_name,
self.width, self.width,
@ -142,10 +143,10 @@ impl Display for StreamInfoChannel {
self.fps, self.fps,
unsafe { rstr!(av_get_pix_fmt_name(transmute(self.format as libc::c_int))) }, unsafe { rstr!(av_get_pix_fmt_name(transmute(self.format as libc::c_int))) },
), ),
StreamChannelType::Audio => write!( StreamType::Audio => write!(
f, f,
"{} #{}: codec={},format={},sample_rate={}", "{} #{}: codec={},format={},sample_rate={}",
self.channel_type, self.stream_type,
self.index, self.index,
codec_name, codec_name,
unsafe { unsafe {
@ -155,11 +156,31 @@ impl Display for StreamInfoChannel {
}, },
self.sample_rate, self.sample_rate,
), ),
StreamChannelType::Subtitle => write!( StreamType::Subtitle => write!(
f, f,
"{} #{}: codec={}", "{} #{}: 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 anyhow::Result;
use ffmpeg_sys_the_third::{av_frame_free, av_packet_free}; use ffmpeg_sys_the_third::{av_frame_free, av_packet_free};
use std::collections::HashMap; use std::collections::HashMap;
@ -39,7 +39,7 @@ impl Transcoder {
/// a pre-configured output encoder /// a pre-configured output encoder
pub unsafe fn transcode_stream( pub unsafe fn transcode_stream(
&mut self, &mut self,
in_stream: &StreamInfoChannel, in_stream: &StreamInfo,
encoder_out: Encoder, encoder_out: Encoder,
) -> Result<()> { ) -> Result<()> {
let src_index = in_stream.index as i32; let src_index = in_stream.index as i32;
@ -47,7 +47,8 @@ impl Transcoder {
let out_ctx = encoder_out.codec_context(); let out_ctx = encoder_out.codec_context();
if in_stream.width != (*out_ctx).width as usize if in_stream.width != (*out_ctx).width as usize
|| in_stream.height != (*out_ctx).height 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 // Setup scaler if the size/format is different from what the codec expects
self.scalers.insert(src_index, Scaler::new()); self.scalers.insert(src_index, Scaler::new());
} }
@ -60,7 +61,7 @@ impl Transcoder {
} }
/// Copy a stream from the input to the output /// 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)?; let dst_stream = self.muxer.add_copy_stream(in_stream.stream)?;
self.copy_stream self.copy_stream
.insert(in_stream.index as i32, (*dst_stream).index); .insert(in_stream.index as i32, (*dst_stream).index);
@ -84,12 +85,16 @@ impl Transcoder {
let src_index = (*stream).index; let src_index = (*stream).index;
// check if encoded stream // check if encoded stream
if let Some(enc) = self.encoders.get_mut(&src_index) { 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 // 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 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); av_frame_free(&mut frame);
new_frame new_frame
} else { } else {
@ -135,7 +140,7 @@ mod tests {
let mut transcoder = let mut transcoder =
Transcoder::new("test_output/test.mp4", "test_output/test_transcode.mkv")?; Transcoder::new("test_output/test.mp4", "test_output/test_transcode.mkv")?;
let info = transcoder.prepare()?; let info = transcoder.prepare()?;
for c in info.channels { for c in info.streams {
transcoder.copy_stream(c)?; transcoder.copy_stream(c)?;
} }
transcoder.run()?; transcoder.run()?;