use super::{owned::*, AVResult};
use crate::ffi::{AVCodecID::*, AVFieldOrder::*, AVMediaType::*, AVPixelFormat::*, *};
use std::convert::TryInto;
use std::fmt::Debug;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
pub trait MediaDesc {
fn codec_id(&self) -> AVCodecID {
Default::default()
}
fn as_audio_desc(&self) -> Option<&AudioDesc> {
None
}
fn as_video_desc(&self) -> Option<&VideoDesc> {
None
}
}
impl Debug for &dyn MediaDesc {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MediaDesc {{ codec_id: {:?} }}", self.codec_id())
}
}
impl Debug for Box<dyn MediaDesc> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MediaDesc {{ codec_id: {:?} }}", self.codec_id())
}
}
pub trait Writer {
fn write_header(&mut self) -> AVResult<()>;
fn write_bytes(
&mut self,
bytes: &[u8],
pts: i64,
duration: i64,
is_key_frame: bool,
stream_index: usize,
) -> AVResult<()>;
fn write_trailer(&mut self) -> AVResult<()>;
fn close(&mut self);
fn flush(&mut self);
fn size(&self) -> u64;
}
impl Debug for &dyn Writer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Writer @ 0x{:p}", self)
}
}
impl Debug for Box<dyn Writer> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Writer @ 0x{:p}", self)
}
}
#[derive(Copy, Clone, Debug, Default)]
pub struct AudioDesc {
pub codec_id: AVCodecID,
pub sample_fmt: AVSampleFormat,
pub bit_rate: i64,
pub sample_rate: usize,
pub channels: usize,
}
impl MediaDesc for AudioDesc {
fn codec_id(&self) -> AVCodecID {
self.codec_id
}
}
impl AudioDesc {
pub fn new() -> Self {
Default::default()
}
}
#[derive(Copy, Clone, Debug, Default)]
pub struct VideoDesc {
pub codec_id: AVCodecID,
pub width: i32,
pub height: i32,
pub bit_rate: i64,
pub time_base: AVRational,
pub gop_size: i32,
pub pix_fmt: AVPixelFormat,
}
impl MediaDesc for VideoDesc {
fn codec_id(&self) -> AVCodecID {
self.codec_id
}
fn as_video_desc(&self) -> Option<&VideoDesc> {
Some(self)
}
}
impl VideoDesc {
pub fn new() -> Self {
Default::default()
}
pub fn with_h264(width: i32, height: i32, bit_rate: i64, time_unit: i32) -> Self {
Self {
codec_id: AV_CODEC_ID_H264,
width,
height,
bit_rate,
time_base: AVRational::with_normalize(time_unit),
gop_size: 12,
pix_fmt: AV_PIX_FMT_YUV420P,
}
}
pub fn with_h265(width: i32, height: i32, bit_rate: i64, time_unit: i32) -> Self {
Self {
codec_id: AV_CODEC_ID_HEVC,
width,
height,
bit_rate,
time_base: AVRational::with_normalize(time_unit),
gop_size: 12,
pix_fmt: AV_PIX_FMT_YUV420P,
}
}
}
#[derive(Debug)]
pub struct Stream {
stream: AVStreamOwned,
in_time_base: AVRational,
}
#[derive(Debug)]
pub struct SimpleWriter {
ctx: AVFormatContextOwned,
format_options: String,
streams: Vec<Stream>,
header_writed: bool,
trailer_writed: bool,
}
impl Drop for SimpleWriter {
fn drop(&mut self) {
self.close();
}
}
impl Writer for SimpleWriter {
fn write_header(&mut self) -> AVResult<()> {
Ok(())
}
fn write_bytes(
&mut self,
bytes: &[u8],
pts: i64,
duration: i64,
is_key_frame: bool,
stream_index: usize,
) -> AVResult<()> {
if !self.header_writed {
self.ctx.write_header(Some(&self.format_options))?;
self.header_writed = true;
}
unsafe {
let stm = self.streams.get(stream_index).unwrap();
let in_time_base = stm.in_time_base;
let out_time_base = stm.stream.time_base;
let mut pkt = AVPacket::default();
let pts = av_rescale_q_rnd(
pts,
in_time_base,
out_time_base,
AVRounding::new().near_inf().pass_min_max(),
);
pkt.pts = pts;
pkt.dts = pts;
pkt.data = bytes.as_ptr() as *mut u8;
pkt.size = bytes.len().try_into()?;
pkt.stream_index = stream_index.try_into()?;
pkt.flags = if is_key_frame { AV_PKT_FLAG_KEY } else { 0 };
pkt.duration = av_rescale_q(duration, in_time_base, out_time_base);
pkt.pos = -1;
self.ctx.write_frame_interleaved(&mut pkt)?;
self.ctx.flush();
Ok(())
}
}
fn write_trailer(&mut self) -> AVResult<()> {
if self.header_writed && !self.trailer_writed {
self.ctx.write_trailer()?;
self.trailer_writed = true;
self.flush();
}
Ok(())
}
fn close(&mut self) {
self.write_trailer().unwrap();
self.ctx.flush();
}
fn flush(&mut self) {
self.ctx.flush();
}
fn size(&self) -> u64 {
self.ctx.size()
}
}
impl SimpleWriter {
pub fn new<P>(
path: P,
descs: &[&dyn MediaDesc],
format: Option<&str>,
format_options: Option<&str>,
) -> AVResult<Self>
where
P: AsRef<Path> + Sized,
{
let mut ctx = AVFormatContextOwned::with_output(path, format, None)?;
let mut streams: Vec<Stream> = vec![];
for desc in descs {
let codec_id = desc.codec_id();
match codec_id {
AV_CODEC_ID_H264 | AV_CODEC_ID_HEVC => {
let desc = desc.as_video_desc().unwrap();
let mut st = ctx.new_stream(codec_id)?;
if let Some(par) = st.codecpar_mut() {
par.codec_type = AVMEDIA_TYPE_VIDEO;
par.codec_id = codec_id;
par.bit_rate = desc.bit_rate;
par.width = desc.width;
par.height = desc.height;
par.field_order = AV_FIELD_UNKNOWN;
par.sample_aspect_ratio = AVRational::new(0, 1);
par.profile = FF_PROFILE_UNKNOWN;
par.level = FF_LEVEL_UNKNOWN;
}
streams.push(Stream {
stream: st,
in_time_base: desc.time_base,
});
}
_ => {}
}
}
Ok(Self {
ctx,
format_options: format_options.unwrap_or("").to_owned(),
streams,
header_writed: false,
trailer_writed: false,
})
}
}
pub type FormatLocationCallback = dyn Fn(usize) -> String;
pub type SplitNotifier = dyn Fn(usize);
#[derive(Default)]
pub struct SplitOptions {
output_path: Option<PathBuf>,
format_location: Option<Box<FormatLocationCallback>>,
before_split: Option<Box<SplitNotifier>>,
after_split: Option<Box<SplitNotifier>>,
max_files: Option<usize>,
max_size_bytes: Option<u64>,
max_size_time: Option<u64>,
max_overhead: Option<f32>,
split_at_keyframe: Option<bool>,
start_index: Option<usize>,
}
impl Debug for SplitOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SplitOptions")
.field("output_path", &self.output_path)
.field("max_files", &self.max_files)
.field("max_size_bytes", &self.max_size_bytes)
.field("max_size_time", &self.max_size_time)
.field("max_overhead", &self.max_overhead)
.field("split_at_keyframe", &self.split_at_keyframe)
.field("start_index", &self.start_index)
.finish()
}
}
pub struct SplitWriter {
medias: Vec<Box<dyn MediaDesc>>,
format: Option<String>,
format_options: Option<String>,
writer: Option<Box<dyn Writer>>,
output_path: PathBuf,
format_location: Option<Box<FormatLocationCallback>>,
before_split: Option<Box<SplitNotifier>>,
after_split: Option<Box<SplitNotifier>>,
max_files: usize,
max_size_bytes: u64,
max_size_time: u64,
max_overhead: f32,
split_at_keyframe: bool,
start_index: usize,
current_index: usize,
start_time: Instant,
started: bool,
need_key_frame: bool,
split_wait_for_key_frame: bool,
}
impl Debug for SplitWriter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "SplitWriter @ 0x{:p}", self)
}
}
impl Writer for SplitWriter {
fn write_header(&mut self) -> AVResult<()> {
if let Some(writer) = &mut self.writer {
writer.write_header()
} else {
Err("The underly writer does not ready".into())
}
}
fn write_bytes(
&mut self,
bytes: &[u8],
pts: i64,
duration: i64,
is_key_frame: bool,
stream_index: usize,
) -> AVResult<()> {
if self.can_split_now(is_key_frame, stream_index) {
self.split_now();
}
if self.writer.is_none() {
let writer = SimpleWriter::new(
self.format_location(self.current_index).to_str().unwrap(),
&self
.medias
.iter()
.map(Deref::deref)
.collect::<Vec<&dyn MediaDesc>>(),
self.format.as_deref(),
self.format_options.as_deref(),
)?;
self.writer = Some(Box::new(writer));
self.start_time = Instant::now();
self.started = true;
}
if let Some(ref mut writer) = self.writer {
writer.write_bytes(bytes, pts, duration, is_key_frame, stream_index)?;
}
Ok(())
}
fn write_trailer(&mut self) -> AVResult<()> {
if let Some(writer) = &mut self.writer {
writer.write_trailer()
} else {
Err("The underly writer does not ready".into())
}
}
fn close(&mut self) {
if let Some(writer) = &mut self.writer {
writer.close();
}
}
fn flush(&mut self) {
if let Some(writer) = &mut self.writer {
writer.flush();
}
}
fn size(&self) -> u64 {
if let Some(writer) = &self.writer {
writer.size()
} else {
0
}
}
}
impl SplitWriter {
pub fn new(
descs: Vec<Box<dyn MediaDesc>>,
format: Option<&str>,
format_options: Option<&str>,
split_options: SplitOptions,
) -> AVResult<Self> {
let mut need_key_frame = false;
for d in descs.iter() {
if d.codec_id().has_gop() {
need_key_frame = true;
}
}
Ok(Self {
medias: descs,
format: format.map(String::from),
format_options: format_options.map(String::from),
writer: None,
output_path: split_options.output_path.unwrap(),
format_location: split_options.format_location,
before_split: split_options.before_split,
after_split: split_options.after_split,
max_files: split_options.max_files.unwrap_or(0),
max_size_bytes: split_options.max_size_bytes.unwrap_or(0),
max_size_time: split_options.max_size_time.unwrap_or(0),
max_overhead: split_options.max_overhead.unwrap_or(0.1f32),
split_at_keyframe: split_options.split_at_keyframe.unwrap_or(true),
start_index: split_options.start_index.unwrap_or(0),
current_index: split_options.start_index.unwrap_or(0),
start_time: Instant::now(),
started: false,
need_key_frame,
split_wait_for_key_frame: false,
})
}
pub(crate) fn is_bytes_overrun(&mut self) -> bool {
let mut exceeded = false;
if let Some(ref writer) = self.writer {
if self.max_size_bytes > 0 && writer.size() >= self.max_size_bytes {
exceeded = true
}
}
exceeded
}
pub(crate) fn is_bytes_overflow(&mut self) -> bool {
let mut exceeded = false;
if let Some(ref writer) = self.writer {
let overhead_bytes = self.max_size_bytes * (self.max_overhead * 100.0) as u64 / 100;
if self.max_size_bytes > 0 && writer.size() >= self.max_size_bytes + overhead_bytes {
exceeded = true
}
}
exceeded
}
pub(crate) fn is_time_overrun(&mut self) -> bool {
self.max_size_time > 0
&& self.start_time.elapsed() >= Duration::from_nanos(self.max_size_time)
}
pub(crate) fn is_time_overflow(&mut self) -> bool {
let overhead_time = self.max_size_time * (self.max_overhead * 100.0) as u64 / 100;
self.max_size_time > 0
&& self.start_time.elapsed() >= Duration::from_nanos(self.max_size_time + overhead_time)
}
pub fn can_split_now(&mut self, is_key_frame: bool, stream_index: usize) -> bool {
let mut split_now: bool = false;
if self.split_wait_for_key_frame {
split_now = self.stream_has_key_frame(stream_index) && is_key_frame;
self.split_wait_for_key_frame = false;
} else {
let overrun = self.is_bytes_overrun() || self.is_time_overrun();
if overrun && self.split_at_keyframe && self.need_key_frame {
self.split_wait_for_key_frame = true;
} else {
split_now = overrun;
}
}
let overflow = self.is_bytes_overflow() || self.is_time_overflow();
split_now || overflow
}
pub fn clean_files(&self) {
if self.max_files > 0 && (self.current_index - self.start_index) >= self.max_files - 1 {
let index = self.current_index - (self.max_files - 1);
if index >= self.start_index {
let old_file = self.format_location(index);
std::fs::remove_file(old_file).unwrap();
}
}
}
pub fn ext_of_format(format: Option<&str>) -> &'static str {
format
.map(|s| match s {
"mp4" => ".mp4",
"mpegts" => ".ts",
_ => "dat",
})
.unwrap_or("dat")
}
pub fn format_location(&self, index: usize) -> PathBuf {
let loc = if let Some(ref cb) = self.format_location {
cb(index)
} else {
format!(
"MED{:06}{}",
index,
Self::ext_of_format(self.format.as_deref())
)
};
let path = self.output_path.join(loc);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).unwrap();
}
path
}
pub fn split_now(&mut self) {
if let Some(ref cb) = self.before_split {
cb(self.current_index);
}
let _ = self.writer.take();
self.clean_files();
self.current_index += 1;
if let Some(ref cb) = self.after_split {
cb(self.current_index);
}
}
pub fn stream_has_key_frame(&self, stream_index: usize) -> bool {
self.medias[stream_index].codec_id().has_gop()
}
}
#[derive(Default)]
pub struct OpenOptions {
medias: Vec<Box<dyn MediaDesc>>,
format: Option<String>,
format_options: Option<String>,
format_location: Option<Box<FormatLocationCallback>>,
before_split: Option<Box<SplitNotifier>>,
after_split: Option<Box<SplitNotifier>>,
max_files: Option<usize>,
max_size_bytes: Option<u64>,
max_size_time: Option<u64>,
max_overhead: Option<f32>,
split_at_keyframe: Option<bool>,
start_index: Option<usize>,
}
impl Debug for OpenOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "OpenOptions @ 0x{:p}", self)
}
}
impl OpenOptions {
pub fn new() -> Self {
Default::default()
}
pub fn media<T>(mut self, media: T) -> Self
where
T: MediaDesc + Sized + 'static,
{
self.medias.push(Box::new(media));
self
}
pub fn format<S>(mut self, format: S) -> Self
where
S: Into<String>,
{
self.format = Some(format.into());
self
}
pub fn format_options<S>(mut self, format_options: S) -> Self
where
S: Into<String>,
{
self.format_options = Some(format_options.into());
self
}
pub fn format_location<F>(mut self, format_location: F) -> Self
where
F: Fn(usize) -> String + 'static,
{
self.format_location = Some(Box::new(format_location));
self
}
pub fn before_split<F>(mut self, before_split: F) -> Self
where
F: Fn(usize) + 'static,
{
self.before_split = Some(Box::new(before_split));
self
}
pub fn after_split<F>(mut self, after_split: F) -> Self
where
F: Fn(usize) + 'static,
{
self.after_split = Some(Box::new(after_split));
self
}
pub fn max_files(mut self, max_files: usize) -> Self {
self.max_files = Some(max_files);
self
}
pub fn max_size_bytes(mut self, max_size_bytes: u64) -> Self {
self.max_size_bytes = Some(max_size_bytes);
self
}
pub fn max_size_time(mut self, max_size_time: u64) -> Self {
self.max_size_time = Some(max_size_time);
self
}
pub fn max_overhead(mut self, max_overhead: f32) -> Self {
self.max_overhead = Some(max_overhead);
self
}
pub fn split_at_keyframe(mut self, split_at_keyframe: bool) -> Self {
self.split_at_keyframe = Some(split_at_keyframe);
self
}
pub fn start_index(mut self, start_index: usize) -> Self {
self.start_index = Some(start_index);
self
}
pub fn open<P>(self, path: P) -> AVResult<Box<dyn Writer>>
where
P: AsRef<Path> + Sized,
{
if self.format_location.is_some() || self.max_files.is_some() {
let split_options = SplitOptions {
output_path: Some(AsRef::<Path>::as_ref(&path).to_path_buf()),
format_location: self.format_location,
before_split: self.before_split,
after_split: self.after_split,
max_files: self.max_files,
max_size_bytes: self.max_size_bytes,
max_size_time: self.max_size_time,
max_overhead: self.max_overhead,
split_at_keyframe: self.split_at_keyframe,
start_index: self.start_index,
};
let writer = SplitWriter::new(
self.medias,
self.format.as_deref(),
self.format_options.as_deref(),
split_options,
)?;
Ok(Box::new(writer))
} else {
let medias: Vec<&dyn MediaDesc> = self.medias.iter().map(Deref::deref).collect();
let writer = SimpleWriter::new(
path,
&medias[..],
self.format.as_deref(),
self.format_options.as_deref(),
)?;
Ok(Box::new(writer))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_simple_writer() {
let a_desc = AudioDesc::new();
let v_desc = VideoDesc::with_h264(352, 288, 4000, 1000000);
let example_bytes = include_bytes!("../../examples/envivio-352x288.264.framed");
for _ in 0..100 {
let mut mp4_writer = SimpleWriter::new(
"/tmp/envivio-352x288.264.mp4",
&[&a_desc, &v_desc],
None,
Some("movflags=frag_keyframe"),
)
.unwrap();
let mut ts_writer = SimpleWriter::new(
"/tmp/envivio-352x288.264.ts",
&[&a_desc, &v_desc],
Some("mpegts"),
Some("mpegts_copyts=1"),
)
.unwrap();
let mut offset: usize = 0;
let mut pts = 0;
while offset + 4 < example_bytes.len() {
let size_bytes = &example_bytes[offset..offset + 4];
let frame_size = i32::from_be_bytes(size_bytes.try_into().unwrap()) as usize;
offset += 4;
let frame_bytes = &example_bytes[offset..offset + frame_size];
offset += frame_size;
mp4_writer
.write_bytes(frame_bytes, pts, 40000, false, 0)
.unwrap();
ts_writer
.write_bytes(frame_bytes, pts, 40000, false, 0)
.unwrap();
pts += 40000;
}
}
}
}