ffmpeg_pipeline/encode/
encoder.rs

1use super::*;
2use ffmpeg_next::{
3    codec::{context::Context, Compliance, Flags as CodecFlags, Id},
4    encoder,
5    ffi::av_opt_set_int,
6    format::{context::Output, Flags as FormatFlags},
7    Packet, Rational,
8};
9
10pub struct Encoder<'o> {
11    index: usize,
12    output: &'o mut Output,
13    encoder: StreamEncoder,
14    in_time_base: Rational,
15    out_time_base: Rational,
16}
17
18impl<'o> Encoder<'o> {
19    pub fn new(
20        output: &'o mut Output,
21        codec_id: Id,
22        codec_params: EncodeParams,
23    ) -> FFmpegResult<Self> {
24        let codec = encoder::find(codec_id).ok_or(FFmpegError::CodecNotFound(codec_id))?;
25        Self::new_with_codec(output, codec, codec_params)
26    }
27
28    /// Creates an encoder by trying implementation names in preference order,
29    /// then falling back to the default encoder for `codec_id`.
30    pub fn new_preferred(
31        output: &'o mut Output,
32        codec_id: Id,
33        preferred_names: &[&str],
34        codec_params: EncodeParams,
35    ) -> FFmpegResult<Self> {
36        let codec = preferred_names
37            .iter()
38            .find_map(|name| encoder::find_by_name(name))
39            .or_else(|| encoder::find(codec_id))
40            .ok_or(FFmpegError::CodecNotFound(codec_id))?;
41        Self::new_with_codec(output, codec, codec_params)
42    }
43
44    fn new_with_codec(
45        output: &'o mut Output,
46        codec: ffmpeg_next::Codec,
47        codec_params: EncodeParams,
48    ) -> FFmpegResult<Self> {
49        let codec_id = codec.id();
50        let output_requires_global_header =
51            output.format().flags().contains(FormatFlags::GLOBAL_HEADER);
52        let mut stream = output.add_stream(codec)?;
53        let mut encoder = Context::from_parameters(stream.parameters())?.encoder();
54        encoder.compliance(Compliance::Experimental);
55
56        let encoder = match codec_params {
57            EncodeParams::Audio {
58                bitrate,
59                channel_layout,
60                compression,
61                global_header,
62                rate,
63                time_base,
64                vbr,
65            } if codec.is_audio() => {
66                let codec = codec.audio()?;
67                let mut encoder = encoder.audio()?;
68                let channel_layout = codec
69                    .channel_layouts()
70                    .map(|cls| cls.best(channel_layout.channels()))
71                    .unwrap_or(channel_layout);
72                let rate = codec
73                    .rates()
74                    .and_then(|rates| {
75                        let rates = rates.collect::<Vec<_>>();
76                        rates
77                            .iter()
78                            .find(|&&r| r == rate)
79                            .or_else(|| rates.iter().min_by_key(|&&r| (r - rate).abs()))
80                            .copied()
81                    })
82                    .unwrap_or(rate);
83                encoder.set_bit_rate(bitrate);
84                encoder.set_channel_layout(channel_layout);
85                encoder.set_compression(compression);
86                encoder.set_format(
87                    codec
88                        .formats()
89                        .and_then(|mut formats| formats.next())
90                        .ok_or(FFmpegError::CodecNotFound(codec_id))?,
91                );
92                encoder.set_rate(rate);
93                encoder.set_time_base(time_base);
94                if global_header || output_requires_global_header {
95                    encoder.set_flags(CodecFlags::GLOBAL_HEADER);
96                }
97                stream.set_time_base(time_base);
98
99                let mut encoder = encoder.open_as(codec)?;
100                if vbr && encoder.id() == Id::OPUS {
101                    unsafe {
102                        match av_opt_set_int(
103                            (*encoder.as_mut_ptr()).priv_data,
104                            c"vbr".as_ptr(),
105                            2,
106                            0,
107                        ) {
108                            0 => Ok(()),
109                            e => Err(ffmpeg_next::Error::from(e)),
110                        }
111                    }?;
112                }
113
114                stream.set_parameters(&encoder);
115                StreamEncoder::Audio(encoder)
116            }
117            EncodeParams::Video {
118                time_base,
119                global_header,
120            } if codec.is_video() => {
121                let mut encoder = encoder.video()?;
122                encoder.set_time_base(time_base);
123                if global_header || output_requires_global_header {
124                    encoder.set_flags(CodecFlags::GLOBAL_HEADER);
125                }
126                stream.set_time_base(time_base);
127
128                let encoder = encoder.open_as(codec)?;
129                stream.set_parameters(&encoder);
130                StreamEncoder::Video(encoder)
131            }
132            _ => {
133                return Err(FFmpegError::CodecNotFound(codec_id));
134            }
135        };
136
137        let in_time_base = codec_params.time_base();
138        let out_time_base = stream.time_base();
139
140        Ok(Self {
141            index: stream.index(),
142            output,
143            encoder,
144            in_time_base,
145            out_time_base,
146        })
147    }
148
149    pub fn get_encoder(&self) -> &StreamEncoder {
150        &self.encoder
151    }
152
153    pub fn set_metadata(&mut self, key: &str, value: &str) -> FFmpegResult {
154        let mut stream = self
155            .output
156            .stream_mut(self.index)
157            .ok_or(FFmpegError::StreamNotFound(self.index))?;
158        let mut metadata = stream.metadata().to_owned();
159        metadata.set(key, value);
160        stream.set_metadata(metadata);
161        Ok(())
162    }
163
164    pub fn write_header(&mut self) -> FFmpegResult {
165        Ok(self.output.write_header()?)
166    }
167
168    pub fn set_size(&mut self, size: FrameSize) {
169        if let StreamEncoder::Video(encoder) = &mut self.encoder {
170            encoder.set_height(size.height as u32);
171            encoder.set_width(size.width as u32);
172        }
173    }
174
175    pub fn send_frame(&mut self, frame: &StreamFrame) -> FFmpegResult {
176        match (&mut self.encoder, frame) {
177            (StreamEncoder::Audio(ref mut encoder), StreamFrame::Audio(frame)) => {
178                encoder.send_frame(frame)?;
179            }
180            (StreamEncoder::Audio(ref mut encoder), StreamFrame::Eof) => {
181                encoder.send_eof()?;
182            }
183            (StreamEncoder::Video(ref mut encoder), StreamFrame::Video(frame)) => {
184                encoder.send_frame(frame)?;
185            }
186            (StreamEncoder::Video(ref mut encoder), StreamFrame::Eof) => {
187                encoder.send_eof()?;
188            }
189            _ => {
190                return Err(FFmpegError::InvalidFrameType(frame.to_string()));
191            }
192        }
193        Ok(())
194    }
195
196    pub fn encode_frame(&mut self) -> FFmpegResult {
197        let mut encoded = Packet::empty();
198
199        match self.encoder {
200            StreamEncoder::Audio(ref mut encoder) => {
201                while encoder.receive_packet(&mut encoded).is_ok() {
202                    encoded.set_stream(self.index);
203                    encoded.rescale_ts(self.in_time_base, self.out_time_base);
204                    encoded.write_interleaved(self.output)?;
205                }
206            }
207            StreamEncoder::Video(ref mut _encoder) => {
208                todo!()
209            }
210        }
211
212        Ok(())
213    }
214}
215
216impl Drop for Encoder<'_> {
217    fn drop(&mut self) {
218        if let Err(e) = self.output.write_trailer() {
219            error!("Error writing trailer: {}", e);
220        }
221    }
222}
ffmpeg_pipeline/encode/encoder.rs

ffmpeg_pipeline/encode/
encoder.rs
