From d645fcbb154bbbac2f6343eeaaedcb99b7959a0d Mon Sep 17 00:00:00 2001 From: Hailey Somerville Date: Wed, 22 Jul 2020 18:51:00 +1000 Subject: [PATCH] implement letterbox scaling --- codec/src/avc/encode.rs | 6 +- codec/src/ffmpeg.rs | 4 +- codec/src/ffmpeg/frame.rs | 153 ++++++++++++++++++++++++++++++++----- codec/src/ffmpeg/pixfmt.rs | 138 +++++++++++++++++++++++++++++++++ codec/src/ffmpeg/scale.rs | 18 ++--- src/rtmp/packet.rs | 6 -- src/video/encode.rs | 36 ++++++++- 7 files changed, 315 insertions(+), 46 deletions(-) create mode 100644 codec/src/ffmpeg/pixfmt.rs diff --git a/codec/src/avc/encode.rs b/codec/src/avc/encode.rs index 5611c0c..918e8e8 100644 --- a/codec/src/avc/encode.rs +++ b/codec/src/avc/encode.rs @@ -8,7 +8,7 @@ use bytes::Bytes; use ffmpeg_dev::sys as ff; use crate::avc::{bitstream, nal, AvcError, DecoderConfigurationRecord}; -use crate::ffmpeg::{AvCodecContext, AvFrame, AvError, AvDict, AvPacket}; +use crate::ffmpeg::{AvCodecContext, AvFrame, AvError, AvDict, AvPacket, PixelFormat}; #[derive(Debug)] pub struct AvcEncoder { @@ -17,7 +17,7 @@ pub struct AvcEncoder { pub struct AvcParams { pub time_base: usize, - pub pixel_format: ff::AVPixelFormat, + pub pixel_format: PixelFormat, pub color_space: ff::AVColorSpace, pub picture_width: usize, pub picture_height: usize, @@ -126,7 +126,7 @@ impl AvcEncoder { avctx.width = params.picture_width.try_into().expect("picture_width too large"); avctx.height = params.picture_height.try_into().expect("picture_height too large"); avctx.colorspace = params.color_space; - avctx.pix_fmt = params.pixel_format; + avctx.pix_fmt = params.pixel_format.into_raw(); avctx.time_base.num = 1; avctx.time_base.den = params.time_base as c_int; avctx.flags |= ff::AV_CODEC_FLAG_GLOBAL_HEADER as i32; diff --git a/codec/src/ffmpeg.rs b/codec/src/ffmpeg.rs index 90ab369..172852e 100644 --- a/codec/src/ffmpeg.rs +++ b/codec/src/ffmpeg.rs @@ -8,11 +8,13 @@ use sys as ff; mod frame; mod packet; +mod pixfmt; mod scale; -pub use frame::{AvFrame, PictureSettings}; +pub use frame::{AvFrame, PictureSettings, PictureData, PictureDataMut}; pub use packet::AvPacket; pub use scale::SwsContext; +pub use pixfmt::{PixelFormat, PixFmtDescriptor, PlaneInfo, ColorFormat}; #[derive(Debug)] pub struct AvCodecContext { diff --git a/codec/src/ffmpeg/frame.rs b/codec/src/ffmpeg/frame.rs index 2a45f7b..8444dda 100644 --- a/codec/src/ffmpeg/frame.rs +++ b/codec/src/ffmpeg/frame.rs @@ -1,10 +1,11 @@ use std::convert::TryInto; -use std::ptr; +use std::marker::PhantomData; use std::os::raw::c_int; +use std::ptr; use ffmpeg_dev::sys as ff; -use crate::ffmpeg::AvError; +use crate::ffmpeg::{AvError, PixelFormat, PixFmtDescriptor, ColorFormat}; #[derive(Debug)] pub struct AvFrame { @@ -32,7 +33,7 @@ impl AvFrame { let underlying = frame.as_underlying_mut(); underlying.width = settings.width.try_into().expect("width too large"); underlying.height = settings.height.try_into().expect("height too large"); - underlying.format = settings.pixel_format; + underlying.format = settings.pixel_format.into_raw(); unsafe { ff::av_frame_get_buffer(frame.as_mut_ptr(), 0); @@ -77,8 +78,8 @@ impl AvFrame { self.coded_height() - underlying.crop_top - underlying.crop_bottom } - pub fn pixel_format(&self) -> ff::AVPixelFormat { - self.as_underlying().format + pub fn pixel_format(&self) -> PixelFormat { + unsafe { PixelFormat::from_raw(self.as_underlying().format) } } pub fn is_key_frame(&self) -> bool { @@ -131,16 +132,18 @@ impl AvFrame { } } - pub fn frame_data(&self) -> FrameData { + pub fn frame_data(&self) -> PictureData { let underlying = self.as_underlying(); - FrameData { - data: &underlying.data, - stride: &underlying.linesize, + PictureData { + picture: self.picture_settings(), + data: underlying.data, + stride: underlying.linesize, + _phantom: PhantomData, } } - pub fn frame_data_mut(&mut self) -> FrameDataMut { + pub fn frame_data_mut(&mut self) -> PictureDataMut { unsafe { let rc = ff::av_frame_make_writable(self.ptr); @@ -149,23 +152,131 @@ impl AvFrame { } } + let picture = self.picture_settings(); let underlying = self.as_underlying_mut(); - FrameDataMut { - data: &mut underlying.data, - stride: &mut underlying.linesize, + PictureDataMut { + picture: picture, + data: underlying.data, + stride: underlying.linesize, + _phantom: PhantomData, + } + } + + fn subframe(&self, x: usize, y: usize, w: usize, h: usize) -> (PictureSettings, PlanarData, PlanarStride) { + let right = x.checked_add(w).expect("x + w overflow"); + let bottom = y.checked_add(h).expect("y + h overflow"); + + if x > self.coded_width() || right > self.coded_width() { + panic!("horizontal section out of bounds (x: {:?}, w: {:?}, picture width: {:?})", + x, w, self.coded_width()); + } + + if y > self.coded_height() || bottom > self.coded_height() { + panic!("vertical section out of bounds (y: {:?}, h: {:?}, picture height: {:?})", + y, h, self.coded_height()); + } + + // scale x and y to align to log2_chroma boundary + let pixdesc = self.pixel_format().descriptor(); + + let x = pixdesc.align_horizontal(x); + let y = pixdesc.align_vertical(y); + + let w = pixdesc.align_horizontal(right) - x; + let h = pixdesc.align_vertical(bottom) - y; + + let picture = PictureSettings { + width: w, + height: h, + pixel_format: self.pixel_format(), + }; + + let mut data = [ptr::null_mut(); 8]; + let underlying = self.as_underlying(); + + // TODO - this should work just fine for non-planar pixfmts as long as + // we don't mutate data - only assign into it from underlying.data + for (idx, component) in pixdesc.components().iter().enumerate() { + let is_chroma = match pixdesc.color() { + ColorFormat::Yuv => idx > 0, + _ => false, + }; + + let plane = component.plane(); + + let x_off = if is_chroma { + x >> pixdesc.log2_chroma_w() + } else { + x + }; + + let y_off = if is_chroma { + y >> pixdesc.log2_chroma_h() + } else { + y + }; + + data[plane] = unsafe { + underlying.data[plane] + .add(x_off * component.step()) + .add(y_off * underlying.linesize[plane] as usize) + }; + } + + (picture, data, underlying.linesize) + } + + pub fn subframe_data(&self, x: usize, y: usize, w: usize, h: usize) -> PictureData { + let (picture, data, stride) = self.subframe(x, y, w, h); + + PictureData { + picture, + data, + stride, + _phantom: PhantomData, + } + } + + pub fn subframe_data_mut(&self, x: usize, y: usize, w: usize, h: usize) -> PictureDataMut { + let (picture, data, stride) = self.subframe(x, y, w, h); + + PictureDataMut { + picture, + data, + stride, + _phantom: PhantomData, } } } -pub struct FrameData<'a> { - pub data: &'a [*mut u8; ff::AV_NUM_DATA_POINTERS as usize], - pub stride: &'a [c_int; ff::AV_NUM_DATA_POINTERS as usize], +type PlanarData = [*mut u8; ff::AV_NUM_DATA_POINTERS as usize]; +type PlanarStride = [c_int; ff::AV_NUM_DATA_POINTERS as usize]; + +pub struct PictureData<'a> { + pub(in crate::ffmpeg) picture: PictureSettings, + pub(in crate::ffmpeg) data: PlanarData, + pub(in crate::ffmpeg) stride: PlanarStride, + _phantom: PhantomData<&'a AvFrame>, } -pub struct FrameDataMut<'a> { - pub data: &'a mut [*mut u8; ff::AV_NUM_DATA_POINTERS as usize], - pub stride: &'a mut [c_int; ff::AV_NUM_DATA_POINTERS as usize], +impl<'a> PictureData<'a> { + pub fn picture_settings(&self) -> &PictureSettings { + &self.picture + } +} + +pub struct PictureDataMut<'a> { + pub(in crate::ffmpeg) picture: PictureSettings, + pub(in crate::ffmpeg) data: PlanarData, + pub(in crate::ffmpeg) stride: PlanarStride, + _phantom: PhantomData<&'a mut AvFrame>, +} + +impl<'a> PictureDataMut<'a> { + pub fn picture_settings(&self) -> &PictureSettings { + &self.picture + } } impl Clone for AvFrame { @@ -192,7 +303,7 @@ impl Drop for AvFrame { pub struct PictureSettings { pub width: usize, pub height: usize, - pub pixel_format: ff::AVPixelFormat, + pub pixel_format: PixelFormat, } impl PictureSettings { @@ -200,7 +311,7 @@ impl PictureSettings { PictureSettings { width, height, - pixel_format: ff::AVPixelFormat_AV_PIX_FMT_YUV420P, + pixel_format: PixelFormat::yuv420p(), } } } diff --git a/codec/src/ffmpeg/pixfmt.rs b/codec/src/ffmpeg/pixfmt.rs new file mode 100644 index 0000000..419bc1b --- /dev/null +++ b/codec/src/ffmpeg/pixfmt.rs @@ -0,0 +1,138 @@ +use std::convert::TryInto; +use std::ffi::CStr; +use std::fmt::{self, Debug}; +use std::slice; + +use ffmpeg_dev::sys as ff; + +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct PixelFormat(ff::AVPixelFormat); + +impl PixelFormat { + pub fn yuv420p() -> Self { + PixelFormat(ff::AVPixelFormat_AV_PIX_FMT_YUV420P) + } + + pub unsafe fn from_raw(pixfmt: ff::AVPixelFormat) -> Self { + PixelFormat(pixfmt) + } + + pub fn into_raw(self) -> ff::AVPixelFormat { + self.0 + } + + pub fn name(&self) -> &'static str { + unsafe { + let ptr = ff::av_get_pix_fmt_name(self.0); + CStr::from_ptr(ptr).to_str().expect("CStr::to_str") + } + } + + pub fn descriptor(&self) -> PixFmtDescriptor { + PixFmtDescriptor { + desc: unsafe { &*ff::av_pix_fmt_desc_get(self.0) }, + } + } +} + +impl Debug for PixelFormat { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PixelFormat({:?}; {:?})", + self.0, + self.name()) + } +} + +#[derive(Debug)] +pub struct PixFmtDescriptor { + desc: &'static ff::AVPixFmtDescriptor, +} + +impl PixFmtDescriptor { + pub fn components(&self) -> &[PlaneInfo] { + unsafe { + let ptr = self.desc.comp.as_ptr() as *const PlaneInfo; + let len = self.desc.nb_components.into(); + slice::from_raw_parts(ptr, len) + } + } + + pub fn planar(&self) -> bool { + (self.desc.flags & ff::AV_PIX_FMT_FLAG_PLANAR as u64) != 0 + } + + pub fn rgb(&self) -> bool { + (self.desc.flags & ff::AV_PIX_FMT_FLAG_RGB as u64) != 0 + } + + pub fn color(&self) -> ColorFormat { + let flags = self.desc.flags; + + if (flags & ff::AV_PIX_FMT_FLAG_RGB as u64) != 0 { + ColorFormat::Rgb + } else if (flags & ff::AV_PIX_FMT_FLAG_HWACCEL as u64) != 0 { + ColorFormat::Hwaccel + } else if (flags & ff::AV_PIX_FMT_FLAG_PAL as u64) != 0 { + ColorFormat::Palette + } else if (flags & ff::AV_PIX_FMT_FLAG_PSEUDOPAL as u64) != 0 { + ColorFormat::PseudoPalette + } else { + ColorFormat::Yuv + } + } + + /// Amount to shift the luma (Y) width right to find the chroma (U, V) width + pub fn log2_chroma_w(&self) -> usize { + self.desc.log2_chroma_w as usize + } + + /// Amount to shift the luma (Y) height right to find the chroma (U, V) height + pub fn log2_chroma_h(&self) -> usize { + self.desc.log2_chroma_h as usize + } + + pub fn align_horizontal(&self, value: usize) -> usize { + value & (usize::max_value() << self.log2_chroma_w()) + } + + pub fn align_vertical(&self, value: usize) -> usize { + value & (usize::max_value() << self.log2_chroma_h()) + } +} + +#[derive(Debug, Clone, Copy)] +pub enum ColorFormat { + Yuv, + Rgb, + Hwaccel, + Palette, + PseudoPalette, +} + +#[repr(transparent)] +#[derive(Debug)] +pub struct PlaneInfo { + comp: ff::AVComponentDescriptor, +} + +impl PlaneInfo { + pub fn plane(&self) -> usize { + self.comp.plane.try_into().unwrap() + } + + pub fn step(&self) -> usize { + self.comp.step.try_into().unwrap() + } + + pub fn offset(&self) -> usize { + self.comp.offset.try_into().unwrap() + } + + pub fn shift(&self) -> usize { + self.comp.shift.try_into().unwrap() + } + + pub fn depth(&self) -> usize { + self.comp.depth.try_into().unwrap() + } +} diff --git a/codec/src/ffmpeg/scale.rs b/codec/src/ffmpeg/scale.rs index bfcce6d..de9005f 100644 --- a/codec/src/ffmpeg/scale.rs +++ b/codec/src/ffmpeg/scale.rs @@ -1,10 +1,9 @@ use std::convert::TryInto; -use std::os::raw::c_int; use std::ptr; use ffmpeg_dev::sys as ff; -use crate::ffmpeg::{AvFrame, PictureSettings}; +use crate::ffmpeg::{AvFrame, PictureSettings, PictureData, PictureDataMut}; #[derive(Debug)] pub struct SwsContext { @@ -22,8 +21,8 @@ impl SwsContext { let ptr = unsafe { ff::sws_getContext( - input_width, input_height, input.pixel_format, - output_width, output_height, output.pixel_format, + input_width, input_height, input.pixel_format.into_raw(), + output_width, output_height, output.pixel_format.into_raw(), ff::SWS_BICUBIC as i32, ptr::null_mut(), ptr::null_mut(), ptr::null(), ) }; @@ -47,21 +46,18 @@ impl SwsContext { &self.output } - pub fn process(&mut self, input: &AvFrame, output: &mut AvFrame) { + pub fn process(&mut self, input: &PictureData, output: &mut PictureDataMut) { let input_settings = input.picture_settings(); let output_settings = output.picture_settings(); - if input_settings != self.input { + if input_settings != &self.input { panic!("wrong picture settings for input frame: {:?}; expected: {:?}", input_settings, self.input); } - if output_settings != self.output { - panic!("wrong picture settings for output frame: {:?}; expected: {:?}", input_settings, self.input); + if output_settings != &self.output { + panic!("wrong picture settings for output frame: {:?}; expected: {:?}", output_settings, self.output); } - let input = input.frame_data(); - let output = output.frame_data_mut(); - let input_data = input.data.as_ptr() as *const *const _; let input_stride = input.stride.as_ptr(); diff --git a/src/rtmp/packet.rs b/src/rtmp/packet.rs index 02735ac..6eaebed 100644 --- a/src/rtmp/packet.rs +++ b/src/rtmp/packet.rs @@ -63,12 +63,6 @@ pub enum VideoFrameType { VideoInfoFrame, } -impl VideoFrameType { - pub fn is_key_frame(&self) -> bool { - *self == VideoFrameType::KeyFrame || *self == VideoFrameType::GeneratedKeyFrame - } -} - #[derive(Debug)] pub enum VideoPacketError { Eof, diff --git a/src/video/encode.rs b/src/video/encode.rs index 50e39e9..66105d6 100644 --- a/src/video/encode.rs +++ b/src/video/encode.rs @@ -1,8 +1,10 @@ +use std::cmp; use std::collections::VecDeque; use std::convert::TryInto; use bytes::Bytes; use fdk_aac::enc as aac; +use num_rational::Ratio; use mixlab_codec::avc::DecoderConfigurationRecord; use mixlab_codec::avc::encode::{AvcEncoder, AvcParams, Preset, Tune, RateControl}; @@ -337,15 +339,41 @@ impl DynamicScaler { } } + let width_ratio = Ratio::::new(output_picture.width, input_picture.width); + let height_ratio = Ratio::::new(output_picture.height, input_picture.height); + let scale_factor = cmp::min(width_ratio, height_ratio); + + let pixdesc = output_picture.pixel_format.descriptor(); + + let scaled_width = pixdesc.align_horizontal( + (scale_factor * input_picture.width).to_integer()); + + let scaled_height = pixdesc.align_vertical( + (scale_factor * input_picture.height).to_integer()); + + let scaled_picture = PictureSettings { + width: scaled_width, + height: scaled_height, + pixel_format: output_picture.pixel_format, + }; + + let scaled_x = pixdesc.align_horizontal((output_picture.width - scaled_width) / 2); + let scaled_y = pixdesc.align_vertical((output_picture.height - scaled_height) / 2); + let scale = self.scale.get_or_insert_with(|| { eprintln!("new dynamic rescaler from: {:?}", input_picture); eprintln!(" to: {:?}", output_picture); - SwsContext::new(input_picture, output_picture) + SwsContext::new(input_picture, scaled_picture) }); - scale.process(frame, &mut self.frame); - self.frame.set_presentation_timestamp(frame.presentation_timestamp()); - // self.frame.copy_props_from(frame); + scale.process( + &frame.frame_data(), + &mut self.frame.subframe_data_mut( + scaled_x, scaled_y, + scaled_width, scaled_height, + ), + ); + self.frame.copy_props_from(frame); &mut self.frame }