#![allow(unsafe_code)]
#![allow(unsafe_op_in_unsafe_fn)]
#![allow(clippy::ptr_as_ptr)]
#![allow(clippy::cast_possible_wrap)]
use std::ffi::CString;
use std::path::Path;
use std::ptr;
use ff_format::{PixelFormat, VideoFrame};
use ff_sys::{
AVCodecID, AVCodecID_AV_CODEC_ID_BMP, AVCodecID_AV_CODEC_ID_MJPEG, AVCodecID_AV_CODEC_ID_PNG,
AVCodecID_AV_CODEC_ID_TIFF, AVCodecID_AV_CODEC_ID_WEBP, AVColorRange_AVCOL_RANGE_JPEG,
AVFormatContext, AVPixelFormat, AVPixelFormat_AV_PIX_FMT_BGR24, AVPixelFormat_AV_PIX_FMT_RGB24,
AVPixelFormat_AV_PIX_FMT_YUV420P, AVRational, av_frame_alloc, av_frame_free,
av_interleaved_write_frame, av_packet_alloc, av_packet_free, av_packet_unref, av_write_trailer,
avcodec, avformat, avformat_alloc_output_context2, avformat_free_context, avformat_new_stream,
avformat_write_header, swscale,
};
use crate::EncodeError;
const MAX_PLANES: usize = 8;
pub(super) struct ImageEncodeOptions {
pub(super) width: Option<u32>,
pub(super) height: Option<u32>,
pub(super) quality: Option<u32>,
pub(super) pixel_format: Option<PixelFormat>,
}
struct ImageEncoderInner {
format_ctx: *mut AVFormatContext,
codec_ctx: *mut ff_sys::AVCodecContext,
dst_frame: *mut ff_sys::AVFrame,
packet: *mut ff_sys::AVPacket,
sws_ctx: Option<*mut ff_sys::SwsContext>,
dst_width: u32,
dst_height: u32,
pix_fmt: AVPixelFormat,
}
impl ImageEncoderInner {
unsafe fn open(
path: &Path,
opts: &ImageEncodeOptions,
src: &VideoFrame,
) -> Result<Self, EncodeError> {
let codec_id = codec_from_extension(path)?;
let dst_width = opts.width.unwrap_or_else(|| src.width());
let dst_height = opts.height.unwrap_or_else(|| src.height());
let pix_fmt = opts
.pixel_format
.map_or_else(|| preferred_pix_fmt(codec_id), pixel_format_to_av);
let mut inner = Self {
format_ctx: ptr::null_mut(),
codec_ctx: ptr::null_mut(),
dst_frame: ptr::null_mut(),
packet: ptr::null_mut(),
sws_ctx: None,
dst_width,
dst_height,
pix_fmt,
};
let c_path = CString::new(path.to_str().ok_or_else(|| EncodeError::CannotCreateFile {
path: path.to_path_buf(),
})?)
.map_err(|_| EncodeError::CannotCreateFile {
path: path.to_path_buf(),
})?;
let explicit_fmt = codec_fallback_format(codec_id);
let mut ret = if let Some(fmt) = explicit_fmt {
avformat_alloc_output_context2(
&mut inner.format_ctx,
ptr::null_mut(),
fmt,
c_path.as_ptr(),
)
} else {
avformat_alloc_output_context2(
&mut inner.format_ctx,
ptr::null_mut(),
ptr::null(),
c_path.as_ptr(),
)
};
if ret < 0 || inner.format_ctx.is_null() {
ret = avformat_alloc_output_context2(
&mut inner.format_ctx,
ptr::null_mut(),
ptr::null(),
c_path.as_ptr(),
);
}
if ret < 0 || inner.format_ctx.is_null() {
return Err(EncodeError::Ffmpeg {
code: ret,
message: format!(
"Cannot create output context: {}",
ff_sys::av_error_string(ret)
),
});
}
let stream = avformat_new_stream(inner.format_ctx, ptr::null());
if stream.is_null() {
return Err(EncodeError::Ffmpeg {
code: 0,
message: "Cannot create output stream".to_string(),
});
}
let codec = avcodec::find_encoder(codec_id).ok_or(EncodeError::UnsupportedCodec {
codec: format!("codec_id={codec_id}"),
})?;
inner.codec_ctx = avcodec::alloc_context3(codec).map_err(EncodeError::from_ffmpeg_error)?;
(*inner.codec_ctx).width = dst_width as i32;
(*inner.codec_ctx).height = dst_height as i32;
(*inner.codec_ctx).time_base = AVRational { num: 1, den: 1 };
(*inner.codec_ctx).pix_fmt = pix_fmt;
if codec_id == AVCodecID_AV_CODEC_ID_MJPEG {
(*inner.codec_ctx).color_range = AVColorRange_AVCOL_RANGE_JPEG;
}
if let Some(q) = opts.quality {
apply_quality(inner.codec_ctx, codec_id, q);
}
avcodec::open2(inner.codec_ctx, codec, ptr::null_mut())
.map_err(EncodeError::from_ffmpeg_error)?;
let par = (*stream).codecpar;
(*par).codec_id = codec_id;
(*par).codec_type = ff_sys::AVMediaType_AVMEDIA_TYPE_VIDEO;
(*par).width = (*inner.codec_ctx).width;
(*par).height = (*inner.codec_ctx).height;
(*par).format = pix_fmt;
let io_ctx = avformat::open_output(path, avformat::avio_flags::WRITE)
.map_err(EncodeError::from_ffmpeg_error)?;
(*inner.format_ctx).pb = io_ctx;
let ret = avformat_write_header(inner.format_ctx, ptr::null_mut());
if ret < 0 {
return Err(EncodeError::from_ffmpeg_error(ret));
}
inner.dst_frame = av_frame_alloc();
if inner.dst_frame.is_null() {
return Err(EncodeError::Ffmpeg {
code: 0,
message: "Cannot allocate destination frame".to_string(),
});
}
(*inner.dst_frame).format = pix_fmt;
(*inner.dst_frame).width = dst_width as i32;
(*inner.dst_frame).height = dst_height as i32;
let ret = ff_sys::av_frame_get_buffer(inner.dst_frame, 0);
if ret < 0 {
return Err(EncodeError::from_ffmpeg_error(ret));
}
inner.packet = av_packet_alloc();
if inner.packet.is_null() {
return Err(EncodeError::Ffmpeg {
code: 0,
message: "Cannot allocate packet".to_string(),
});
}
Ok(inner)
}
unsafe fn encode_frame(&mut self, src: &VideoFrame) -> Result<(), EncodeError> {
let src_fmt = pixel_format_to_av(src.format());
let needs_conversion = src_fmt != self.pix_fmt
|| src.width() != self.dst_width
|| src.height() != self.dst_height;
if needs_conversion {
let sws_ctx = swscale::get_context(
src.width() as i32,
src.height() as i32,
src_fmt,
self.dst_width as i32,
self.dst_height as i32,
self.pix_fmt,
swscale::scale_flags::BILINEAR,
)
.map_err(EncodeError::from_ffmpeg_error)?;
self.sws_ctx = Some(sws_ctx);
let mut src_data: [*const u8; MAX_PLANES] = [ptr::null(); MAX_PLANES];
let mut src_linesize: [i32; MAX_PLANES] = [0; MAX_PLANES];
for (i, plane) in src.planes().iter().enumerate() {
if i < MAX_PLANES {
src_data[i] = plane.data().as_ptr();
src_linesize[i] = src.strides()[i] as i32;
}
}
let scale_result = swscale::scale(
sws_ctx,
src_data.as_ptr(),
src_linesize.as_ptr(),
0,
src.height() as i32,
(*self.dst_frame).data.as_mut_ptr().cast_const(),
(*self.dst_frame).linesize.as_mut_ptr(),
);
if let Some(sws) = self.sws_ctx.take() {
swscale::free_context(sws);
}
scale_result.map_err(EncodeError::from_ffmpeg_error)?;
} else {
for (i, plane) in src.planes().iter().enumerate() {
if i >= MAX_PLANES || (*self.dst_frame).data[i].is_null() {
break;
}
let src_stride = src.strides()[i];
let dst_stride = (*self.dst_frame).linesize[i] as usize;
let plane_data = plane.data();
if src_stride == dst_stride {
std::ptr::copy_nonoverlapping(
plane_data.as_ptr(),
(*self.dst_frame).data[i],
plane_data.len(),
);
} else {
let row_bytes = src_stride.min(dst_stride);
let num_rows = plane_data.len() / src_stride;
for row in 0..num_rows {
std::ptr::copy_nonoverlapping(
plane_data[row * src_stride..].as_ptr(),
(*self.dst_frame).data[i].add(row * dst_stride),
row_bytes,
);
}
}
}
}
(*self.dst_frame).pts = 0;
avcodec::send_frame(self.codec_ctx, self.dst_frame)
.map_err(EncodeError::from_ffmpeg_error)?;
self.drain_packets(false)?;
avcodec::send_frame(self.codec_ctx, ptr::null()).map_err(EncodeError::from_ffmpeg_error)?;
self.drain_packets(true)?;
av_write_trailer(self.format_ctx);
avformat::close_output(&mut (*self.format_ctx).pb);
Ok(())
}
unsafe fn drain_packets(&mut self, until_eof: bool) -> Result<(), EncodeError> {
loop {
match avcodec::receive_packet(self.codec_ctx, self.packet) {
Ok(()) => {
(*self.packet).stream_index = 0;
let ret = av_interleaved_write_frame(self.format_ctx, self.packet);
av_packet_unref(self.packet);
if ret < 0 {
return Err(EncodeError::from_ffmpeg_error(ret));
}
}
Err(e) if e == ff_sys::error_codes::EOF => break,
Err(e) if !until_eof && e == ff_sys::error_codes::EAGAIN => break,
Err(e) => return Err(EncodeError::from_ffmpeg_error(e)),
}
}
Ok(())
}
}
impl Drop for ImageEncoderInner {
fn drop(&mut self) {
unsafe {
if !self.dst_frame.is_null() {
av_frame_free(&mut self.dst_frame);
}
if !self.packet.is_null() {
av_packet_free(&mut self.packet);
}
if let Some(sws) = self.sws_ctx.take() {
swscale::free_context(sws);
}
if !self.codec_ctx.is_null() {
avcodec::free_context(&mut self.codec_ctx);
}
if !self.format_ctx.is_null() {
if !(*self.format_ctx).pb.is_null() {
avformat::close_output(&mut (*self.format_ctx).pb);
}
avformat_free_context(self.format_ctx);
self.format_ctx = ptr::null_mut();
}
}
}
}
pub(super) fn codec_from_extension(path: &Path) -> Result<AVCodecID, EncodeError> {
let ext = path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_lowercase();
match ext.as_str() {
"jpg" | "jpeg" => Ok(AVCodecID_AV_CODEC_ID_MJPEG),
"png" => Ok(AVCodecID_AV_CODEC_ID_PNG),
"bmp" => Ok(AVCodecID_AV_CODEC_ID_BMP),
"tif" | "tiff" => Ok(AVCodecID_AV_CODEC_ID_TIFF),
"webp" => Ok(AVCodecID_AV_CODEC_ID_WEBP),
"" => Err(EncodeError::InvalidConfig {
reason: "no file extension".to_string(),
}),
e => Err(EncodeError::UnsupportedCodec {
codec: e.to_string(),
}),
}
}
fn codec_fallback_format(codec_id: AVCodecID) -> Option<*const std::os::raw::c_char> {
if codec_id == AVCodecID_AV_CODEC_ID_MJPEG {
Some(c"mjpeg".as_ptr())
} else if codec_id == AVCodecID_AV_CODEC_ID_PNG {
Some(c"apng".as_ptr())
} else if codec_id == AVCodecID_AV_CODEC_ID_TIFF {
Some(c"tiff".as_ptr())
} else if codec_id == AVCodecID_AV_CODEC_ID_WEBP {
Some(c"webp".as_ptr())
} else {
None
}
}
fn preferred_pix_fmt(codec_id: AVCodecID) -> AVPixelFormat {
match codec_id {
x if x == AVCodecID_AV_CODEC_ID_MJPEG => AVPixelFormat_AV_PIX_FMT_YUV420P,
x if x == AVCodecID_AV_CODEC_ID_PNG => AVPixelFormat_AV_PIX_FMT_RGB24,
x if x == AVCodecID_AV_CODEC_ID_BMP => AVPixelFormat_AV_PIX_FMT_BGR24,
x if x == AVCodecID_AV_CODEC_ID_TIFF => AVPixelFormat_AV_PIX_FMT_RGB24,
x if x == AVCodecID_AV_CODEC_ID_WEBP => AVPixelFormat_AV_PIX_FMT_YUV420P,
_ => AVPixelFormat_AV_PIX_FMT_RGB24,
}
}
fn pixel_format_to_av(fmt: PixelFormat) -> AVPixelFormat {
match fmt {
PixelFormat::Yuv420p => ff_sys::AVPixelFormat_AV_PIX_FMT_YUV420P,
PixelFormat::Yuv422p => ff_sys::AVPixelFormat_AV_PIX_FMT_YUV422P,
PixelFormat::Yuv444p => ff_sys::AVPixelFormat_AV_PIX_FMT_YUV444P,
PixelFormat::Rgb24 => ff_sys::AVPixelFormat_AV_PIX_FMT_RGB24,
PixelFormat::Bgr24 => ff_sys::AVPixelFormat_AV_PIX_FMT_BGR24,
PixelFormat::Rgba => ff_sys::AVPixelFormat_AV_PIX_FMT_RGBA,
PixelFormat::Bgra => ff_sys::AVPixelFormat_AV_PIX_FMT_BGRA,
PixelFormat::Gray8 => ff_sys::AVPixelFormat_AV_PIX_FMT_GRAY8,
PixelFormat::Nv12 => ff_sys::AVPixelFormat_AV_PIX_FMT_NV12,
PixelFormat::Nv21 => ff_sys::AVPixelFormat_AV_PIX_FMT_NV21,
PixelFormat::Yuv420p10le => ff_sys::AVPixelFormat_AV_PIX_FMT_YUV420P10LE,
PixelFormat::P010le => ff_sys::AVPixelFormat_AV_PIX_FMT_P010LE,
_ => ff_sys::AVPixelFormat_AV_PIX_FMT_RGB24,
}
}
unsafe fn apply_quality(codec_ctx: *mut ff_sys::AVCodecContext, codec_id: AVCodecID, quality: u32) {
let q = quality.min(100);
if codec_id == AVCodecID_AV_CODEC_ID_MJPEG {
let qscale = (1 + (100 - q) * 30 / 100) as i32;
(*codec_ctx).qmin = qscale;
(*codec_ctx).qmax = qscale;
log::info!("MJPEG quality applied quality={q} qscale={qscale}");
} else if codec_id == AVCodecID_AV_CODEC_ID_PNG {
let level = q * 9 / 100;
if (*codec_ctx).priv_data.is_null() {
log::warn!("PNG compression_level: priv_data is null, skipping quality={q}");
return;
}
let Ok(key) = CString::new("compression_level") else {
return;
};
let Ok(val) = CString::new(level.to_string()) else {
return;
};
let ret = ff_sys::av_opt_set((*codec_ctx).priv_data, key.as_ptr(), val.as_ptr(), 0);
if ret < 0 {
log::warn!(
"av_opt_set compression_level failed, ignoring \
quality={q} error={}",
ff_sys::av_error_string(ret)
);
} else {
log::info!("PNG compression_level applied quality={q} level={level}");
}
} else if codec_id == AVCodecID_AV_CODEC_ID_WEBP {
if (*codec_ctx).priv_data.is_null() {
log::warn!("WebP quality: priv_data is null, skipping quality={q}");
return;
}
let Ok(key) = CString::new("quality") else {
return;
};
let Ok(val) = CString::new(q.to_string()) else {
return;
};
let ret = ff_sys::av_opt_set((*codec_ctx).priv_data, key.as_ptr(), val.as_ptr(), 0);
if ret < 0 {
log::warn!(
"av_opt_set quality failed for WebP, ignoring \
quality={q} error={}",
ff_sys::av_error_string(ret)
);
} else {
log::info!("WebP quality applied quality={q}");
}
} else {
let fmt_name = if codec_id == AVCodecID_AV_CODEC_ID_BMP {
"bmp"
} else if codec_id == AVCodecID_AV_CODEC_ID_TIFF {
"tiff"
} else {
"this format"
};
log::warn!("quality option has no effect for {fmt_name} images, ignoring quality={q}");
}
}
pub(super) fn encode_image(
path: &Path,
frame: &VideoFrame,
opts: &ImageEncodeOptions,
) -> Result<(), EncodeError> {
unsafe {
ff_sys::ensure_initialized();
let mut inner = ImageEncoderInner::open(path, opts, frame)?;
inner.encode_frame(frame)?;
log::info!(
"Image encoded successfully path={} src={}x{} dst={}x{}",
path.display(),
frame.width(),
frame.height(),
inner.dst_width,
inner.dst_height,
);
Ok(())
} }
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn codec_from_extension_jpeg_should_return_mjpeg() {
let id = codec_from_extension(Path::new("img.jpg")).unwrap();
assert_eq!(id, AVCodecID_AV_CODEC_ID_MJPEG);
}
#[test]
fn codec_from_extension_jpeg_alias_should_return_mjpeg() {
let id = codec_from_extension(Path::new("img.jpeg")).unwrap();
assert_eq!(id, AVCodecID_AV_CODEC_ID_MJPEG);
}
#[test]
fn codec_from_extension_png_should_return_png() {
let id = codec_from_extension(Path::new("img.PNG")).unwrap(); assert_eq!(id, AVCodecID_AV_CODEC_ID_PNG);
}
#[test]
fn codec_from_extension_bmp_should_return_bmp() {
let id = codec_from_extension(Path::new("img.bmp")).unwrap();
assert_eq!(id, AVCodecID_AV_CODEC_ID_BMP);
}
#[test]
fn codec_from_extension_tif_should_return_tiff() {
let id = codec_from_extension(Path::new("img.tif")).unwrap();
assert_eq!(id, AVCodecID_AV_CODEC_ID_TIFF);
}
#[test]
fn codec_from_extension_tiff_should_return_tiff() {
let id = codec_from_extension(Path::new("img.tiff")).unwrap();
assert_eq!(id, AVCodecID_AV_CODEC_ID_TIFF);
}
#[test]
fn codec_from_extension_webp_should_return_webp() {
let id = codec_from_extension(Path::new("img.webp")).unwrap();
assert_eq!(id, AVCodecID_AV_CODEC_ID_WEBP);
}
#[test]
fn codec_from_extension_no_ext_should_return_invalid_config() {
let result = codec_from_extension(Path::new("no_extension"));
assert!(matches!(result, Err(EncodeError::InvalidConfig { .. })));
}
#[test]
fn codec_from_extension_unknown_should_return_unsupported_codec() {
let result = codec_from_extension(Path::new("img.avi"));
assert!(matches!(result, Err(EncodeError::UnsupportedCodec { .. })));
}
#[test]
fn preferred_pix_fmt_mjpeg_should_return_yuv420p() {
assert_eq!(
preferred_pix_fmt(AVCodecID_AV_CODEC_ID_MJPEG),
AVPixelFormat_AV_PIX_FMT_YUV420P
);
}
#[test]
fn preferred_pix_fmt_png_should_return_rgb24() {
assert_eq!(
preferred_pix_fmt(AVCodecID_AV_CODEC_ID_PNG),
AVPixelFormat_AV_PIX_FMT_RGB24
);
}
#[test]
fn preferred_pix_fmt_bmp_should_return_bgr24() {
assert_eq!(
preferred_pix_fmt(AVCodecID_AV_CODEC_ID_BMP),
AVPixelFormat_AV_PIX_FMT_BGR24
);
}
#[test]
fn preferred_pix_fmt_webp_should_return_yuv420p() {
assert_eq!(
preferred_pix_fmt(AVCodecID_AV_CODEC_ID_WEBP),
AVPixelFormat_AV_PIX_FMT_YUV420P
);
}
#[test]
fn pixel_format_to_av_yuv420p_should_match() {
assert_eq!(
pixel_format_to_av(PixelFormat::Yuv420p),
AVPixelFormat_AV_PIX_FMT_YUV420P
);
}
#[test]
fn pixel_format_to_av_rgb24_should_match() {
assert_eq!(
pixel_format_to_av(PixelFormat::Rgb24),
AVPixelFormat_AV_PIX_FMT_RGB24
);
}
#[test]
fn drop_on_uninitialised_inner_should_not_panic() {
let inner = ImageEncoderInner {
format_ctx: ptr::null_mut(),
codec_ctx: ptr::null_mut(),
dst_frame: ptr::null_mut(),
packet: ptr::null_mut(),
sws_ctx: None,
dst_width: 0,
dst_height: 0,
pix_fmt: ff_sys::AVPixelFormat_AV_PIX_FMT_NONE,
};
drop(inner); }
#[test]
fn codec_from_extension_case_insensitive_should_work() {
let _ = codec_from_extension(&PathBuf::from("IMG.JPG")).unwrap();
let _ = codec_from_extension(&PathBuf::from("IMG.BMP")).unwrap();
let _ = codec_from_extension(&PathBuf::from("IMG.WEBP")).unwrap();
}
}