lumen_ffmpeg/encode/
video.rs

1use std::ptr;
2
3use crate::ffi::{self, AvFrame, AvPacket, sys};
4use crate::gpu::GpuBackend;
5use crate::video::{CpuVideoFrame, EncodeMode, PixelFormat, VideoCodec};
6use crate::{FfmpegError, Result};
7#[cfg(target_os = "linux")]
8use sys::SWS_BILINEAR;
9#[cfg(not(target_os = "linux"))]
10use sys::SwsFlags::SWS_BILINEAR;
11use sys::{
12    AVMediaType::AVMEDIA_TYPE_VIDEO,
13    AVPixelFormat::{AV_PIX_FMT_CUDA, AV_PIX_FMT_VIDEOTOOLBOX, AV_PIX_FMT_YUV420P},
14};
15
16use super::codec::find_encoder;
17use super::common::refresh_stream_time_base;
18use super::hw::{AvBufferRef, create_encoder_hw_contexts};
19use super::output::OutputContext;
20use super::telemetry::GpuEncodeTelemetry;
21
22#[derive(Debug, Clone)]
23pub struct VideoEncoderConfig {
24    pub width: u32,
25    pub height: u32,
26    pub fps: u32,
27    pub codec: VideoCodec,
28    pub encoder_name: Option<String>,
29    pub mode: EncodeMode,
30    pub bit_rate: i64,
31}
32
33impl VideoEncoderConfig {
34    pub fn cpu_rgba(width: u32, height: u32, fps: u32, codec: VideoCodec) -> Self {
35        Self {
36            width,
37            height,
38            fps,
39            codec,
40            encoder_name: None,
41            mode: EncodeMode::CpuUpload,
42            bit_rate: 8_000_000,
43        }
44    }
45
46    pub fn h264_videotoolbox(width: u32, height: u32, fps: u32) -> Self {
47        Self {
48            encoder_name: Some("h264_videotoolbox".to_string()),
49            ..Self::cpu_rgba(width, height, fps, VideoCodec::H264)
50        }
51    }
52}
53
54pub struct VideoEncoder {
55    stream_index: usize,
56    stream_time_base: sys::AVRational,
57    pub(in crate::encode) context: *mut sys::AVCodecContext,
58    frame: AvFrame,
59    scaler: *mut sys::SwsContext,
60    #[allow(dead_code)]
61    pub(in crate::encode) hw_device: Option<AvBufferRef>,
62    pub(in crate::encode) hw_frames: Option<AvBufferRef>,
63    pub(in crate::encode) next_pts: i64,
64    pub(in crate::encode) mode: EncodeMode,
65    pub(in crate::encode) gpu_telemetry: GpuEncodeTelemetry,
66}
67
68unsafe impl Send for VideoEncoder {}
69
70impl VideoEncoder {
71    pub fn create(output: &mut OutputContext, config: VideoEncoderConfig) -> Result<Self> {
72        if config.width == 0 || config.height == 0 || config.fps == 0 {
73            return Err(FfmpegError::new(
74                "VideoEncoder::create",
75                "width, height, and fps must be greater than zero",
76            ));
77        }
78        let codec = find_encoder(&config)?;
79        if codec.is_null() {
80            return Err(FfmpegError::new(
81                "avcodec_find_encoder",
82                "requested encoder is unavailable",
83            )
84            .with_codec(config.codec));
85        }
86        let stream = unsafe { sys::avformat_new_stream(output.ptr, ptr::null()) };
87        if stream.is_null() {
88            return Err(FfmpegError::new(
89                "avformat_new_stream",
90                "failed to allocate output stream",
91            ));
92        }
93        let context = unsafe { sys::avcodec_alloc_context3(codec) };
94        if context.is_null() {
95            return Err(FfmpegError::new(
96                "avcodec_alloc_context3",
97                "failed to allocate encoder context",
98            ));
99        }
100
101        let time_base = sys::AVRational {
102            num: 1,
103            den: config.fps as i32,
104        };
105        let hw_contexts = create_encoder_hw_contexts(&config)?;
106
107        unsafe {
108            (*context).codec_id = config.codec.to_av_codec_id();
109            (*context).codec_type = AVMEDIA_TYPE_VIDEO;
110            (*context).width = config.width as i32;
111            (*context).height = config.height as i32;
112            (*context).time_base = time_base;
113            (*context).framerate = sys::AVRational {
114                num: config.fps as i32,
115                den: 1,
116            };
117            (*context).pix_fmt = match config.mode {
118                EncodeMode::GpuTexture(GpuBackend::Cuda) => AV_PIX_FMT_CUDA,
119                EncodeMode::GpuTexture(GpuBackend::Metal) => AV_PIX_FMT_VIDEOTOOLBOX,
120                EncodeMode::GpuTexture(GpuBackend::Vulkan) | EncodeMode::CpuUpload => {
121                    AV_PIX_FMT_YUV420P
122                }
123            };
124            (*context).bit_rate = config.bit_rate;
125            (*context).gop_size = config.fps as i32 * 2;
126            if let Some(frames) = hw_contexts
127                .as_ref()
128                .and_then(|contexts| contexts.frames.as_ref())
129            {
130                (*context).hw_frames_ctx = sys::av_buffer_ref(frames.ptr);
131                if (*context).hw_frames_ctx.is_null() {
132                    return Err(FfmpegError::new(
133                        "av_buffer_ref",
134                        "failed to reference encoder hardware frames context",
135                    )
136                    .with_backend(GpuBackend::Cuda));
137                }
138            }
139            if ((*(*output.ptr).oformat).flags & sys::AVFMT_GLOBALHEADER) != 0 {
140                (*context).flags |= sys::AV_CODEC_FLAG_GLOBAL_HEADER as i32;
141            }
142            ffi::check(
143                sys::avcodec_open2(context, codec, ptr::null_mut()),
144                "avcodec_open2",
145            )?;
146            ffi::check(
147                sys::avcodec_parameters_from_context((*stream).codecpar, context),
148                "avcodec_parameters_from_context",
149            )?;
150            (*stream).time_base = time_base;
151        }
152
153        let mut frame = AvFrame::new()?;
154        let scaler = if config.mode == EncodeMode::CpuUpload {
155            unsafe {
156                (*frame.as_mut_ptr()).format = AV_PIX_FMT_YUV420P as i32;
157                (*frame.as_mut_ptr()).width = config.width as i32;
158                (*frame.as_mut_ptr()).height = config.height as i32;
159                ffi::check(
160                    sys::av_frame_get_buffer(frame.as_mut_ptr(), 32),
161                    "av_frame_get_buffer",
162                )?;
163            }
164
165            let scaler = unsafe {
166                sys::sws_getContext(
167                    config.width as i32,
168                    config.height as i32,
169                    PixelFormat::Rgba8.to_av_pixel_format(),
170                    config.width as i32,
171                    config.height as i32,
172                    AV_PIX_FMT_YUV420P,
173                    SWS_BILINEAR as i32,
174                    ptr::null_mut(),
175                    ptr::null_mut(),
176                    ptr::null(),
177                )
178            };
179            if scaler.is_null() {
180                return Err(FfmpegError::new(
181                    "sws_getContext",
182                    "failed to create encoder color conversion context",
183                ));
184            }
185            scaler
186        } else {
187            ptr::null_mut()
188        };
189
190        Ok(Self {
191            stream_index: unsafe { (*stream).index as usize },
192            stream_time_base: time_base,
193            context,
194            frame,
195            scaler,
196            hw_device: hw_contexts
197                .as_ref()
198                .and_then(|contexts| contexts.device.clone_ref()),
199            hw_frames: hw_contexts.and_then(|contexts| contexts.frames),
200            next_pts: 0,
201            mode: config.mode,
202            gpu_telemetry: GpuEncodeTelemetry::default(),
203        })
204    }
205
206    pub fn gpu_telemetry(&self) -> &GpuEncodeTelemetry {
207        &self.gpu_telemetry
208    }
209
210    pub(in crate::encode) fn refresh_stream_time_base(
211        &mut self,
212        output: &OutputContext,
213    ) -> Result<()> {
214        self.stream_time_base = refresh_stream_time_base(
215            output,
216            self.stream_index,
217            "VideoEncoder::refresh_stream_time_base",
218        )?;
219        Ok(())
220    }
221
222    pub(in crate::encode) fn send_cpu_frame(
223        &mut self,
224        output: &mut OutputContext,
225        frame: &CpuVideoFrame,
226    ) -> Result<()> {
227        if let EncodeMode::GpuTexture(backend) = self.mode {
228            return Err(FfmpegError::new(
229                "VideoEncoder::send_cpu_frame",
230                "hardware texture encoders consume GPU inputs; create a CPU upload encoder to send CPU bytes",
231            )
232            .with_backend(backend));
233        }
234        unsafe {
235            ffi::check(
236                sys::av_frame_make_writable(self.frame.as_mut_ptr()),
237                "av_frame_make_writable",
238            )?;
239        }
240        let src_data = [frame.data.as_ptr(), ptr::null(), ptr::null(), ptr::null()];
241        let src_stride = [frame.stride as i32, 0, 0, 0];
242        unsafe {
243            sys::sws_scale(
244                self.scaler,
245                src_data.as_ptr(),
246                src_stride.as_ptr(),
247                0,
248                frame.height as i32,
249                (*self.frame.as_mut_ptr()).data.as_mut_ptr(),
250                (*self.frame.as_mut_ptr()).linesize.as_mut_ptr(),
251            );
252            (*self.frame.as_mut_ptr()).pts = frame.pts.unwrap_or(self.next_pts);
253        }
254        self.next_pts = self.next_pts.saturating_add(1);
255        self.send_frame(output, self.frame.as_ptr())
256    }
257
258    pub(in crate::encode) fn flush(&mut self, output: &mut OutputContext) -> Result<()> {
259        self.send_frame(output, ptr::null())
260    }
261
262    pub(in crate::encode) fn send_frame(
263        &mut self,
264        output: &mut OutputContext,
265        frame: *const sys::AVFrame,
266    ) -> Result<()> {
267        unsafe {
268            ffi::check(
269                sys::avcodec_send_frame(self.context, frame),
270                "avcodec_send_frame",
271            )?;
272        }
273        loop {
274            let mut packet = AvPacket::new()?;
275            let result = unsafe { sys::avcodec_receive_packet(self.context, packet.as_mut_ptr()) };
276            if result == sys::AVERROR(libc::EAGAIN) || result == sys::AVERROR_EOF {
277                break;
278            }
279            if result < 0 {
280                return Err(ffi::error_from_code("avcodec_receive_packet", result));
281            }
282            unsafe {
283                (*packet.as_mut_ptr()).stream_index = self.stream_index as i32;
284                if (*packet.as_mut_ptr()).duration == 0 {
285                    (*packet.as_mut_ptr()).duration = 1;
286                }
287                sys::av_packet_rescale_ts(
288                    packet.as_mut_ptr(),
289                    (*self.context).time_base,
290                    self.stream_time_base,
291                );
292                ffi::check(
293                    sys::av_interleaved_write_frame(output.ptr, packet.as_mut_ptr()),
294                    "av_interleaved_write_frame",
295                )
296                .map_err(|error| error.with_path(output.path().to_string()))?;
297            }
298        }
299        Ok(())
300    }
301}
302
303impl Drop for VideoEncoder {
304    fn drop(&mut self) {
305        unsafe {
306            if !self.scaler.is_null() {
307                sys::sws_freeContext(self.scaler);
308            }
309            sys::avcodec_free_context(&mut self.context);
310        }
311    }
312}
lumen_ffmpeg/encode/video.rs

lumen_ffmpeg/encode/
video.rs
