1 // Copyright (C) 2019 Sebastian Dröge <sebastian@centricular.com>
2 //
3 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6 // option. This file may not be copied, modified, or distributed
7 // except according to those terms.
8 
9 use gst::glib;
10 use gst::gst_debug;
11 use gst::subclass::prelude::*;
12 use gst_video::prelude::*;
13 use gst_video::subclass::prelude::*;
14 use once_cell::sync::Lazy;
15 use rav1e::color;
16 use rav1e::config;
17 use rav1e::data;
18 use std::sync::Mutex;
19 
20 const DEFAULT_SPEED_PRESET: u32 = 5;
21 const DEFAULT_LOW_LATENCY: bool = false;
22 const DEFAULT_MIN_KEY_FRAME_INTERVAL: u64 = 12;
23 const DEFAULT_MAX_KEY_FRAME_INTERVAL: u64 = 240;
24 const DEFAULT_BITRATE: i32 = 0;
25 const DEFAULT_QUANTIZER: usize = 100;
26 const DEFAULT_TILE_COLS: usize = 0;
27 const DEFAULT_TILE_ROWS: usize = 0;
28 const DEFAULT_TILES: usize = 0;
29 const DEFAULT_THREADS: usize = 0;
30 
31 #[derive(Debug, Clone, Copy)]
32 struct Settings {
33     speed_preset: u32,
34     low_latency: bool,
35     min_key_frame_interval: u64,
36     max_key_frame_interval: u64,
37     bitrate: i32,
38     quantizer: usize,
39     tile_cols: usize,
40     tile_rows: usize,
41     tiles: usize,
42     threads: usize,
43 }
44 
45 impl Default for Settings {
default() -> Self46     fn default() -> Self {
47         Settings {
48             speed_preset: DEFAULT_SPEED_PRESET,
49             low_latency: DEFAULT_LOW_LATENCY,
50             min_key_frame_interval: DEFAULT_MIN_KEY_FRAME_INTERVAL,
51             max_key_frame_interval: DEFAULT_MAX_KEY_FRAME_INTERVAL,
52             bitrate: DEFAULT_BITRATE,
53             quantizer: DEFAULT_QUANTIZER,
54             tile_cols: DEFAULT_TILE_COLS,
55             tile_rows: DEFAULT_TILE_ROWS,
56             tiles: DEFAULT_TILES,
57             threads: DEFAULT_THREADS,
58         }
59     }
60 }
61 
62 enum Context {
63     Eight(rav1e::Context<u8>),
64     Sixteen(rav1e::Context<u16>),
65 }
66 
67 impl Context {
receive_packet( &mut self, ) -> Result<(data::FrameType, u64, u32, Vec<u8>), data::EncoderStatus>68     fn receive_packet(
69         &mut self,
70     ) -> Result<(data::FrameType, u64, u32, Vec<u8>), data::EncoderStatus> {
71         match self {
72             Context::Eight(ref mut context) => context.receive_packet().map(|packet| {
73                 (
74                     packet.frame_type,
75                     packet.input_frameno,
76                     *packet
77                         .opaque
78                         .and_then(|o| o.downcast::<u32>().ok())
79                         .expect("no frame number"),
80                     packet.data,
81                 )
82             }),
83             Context::Sixteen(ref mut context) => context.receive_packet().map(|packet| {
84                 (
85                     packet.frame_type,
86                     packet.input_frameno,
87                     *packet
88                         .opaque
89                         .and_then(|o| o.downcast::<u32>().ok())
90                         .expect("no frame number"),
91                     packet.data,
92                 )
93             }),
94         }
95     }
96 
send_frame( &mut self, in_frame: Option<(u32, &gst_video::VideoFrameRef<&gst::BufferRef>)>, force_keyframe: bool, ) -> Result<(), data::EncoderStatus>97     fn send_frame(
98         &mut self,
99         in_frame: Option<(u32, &gst_video::VideoFrameRef<&gst::BufferRef>)>,
100         force_keyframe: bool,
101     ) -> Result<(), data::EncoderStatus> {
102         match self {
103             Context::Eight(ref mut context) => {
104                 if let Some((frame_number, in_frame)) = in_frame {
105                     let mut enc_frame = context.new_frame();
106                     enc_frame.planes[0].copy_from_raw_u8(
107                         in_frame.plane_data(0).unwrap(),
108                         in_frame.plane_stride()[0] as usize,
109                         1,
110                     );
111 
112                     if in_frame.n_planes() > 1 {
113                         enc_frame.planes[1].copy_from_raw_u8(
114                             in_frame.plane_data(1).unwrap(),
115                             in_frame.plane_stride()[1] as usize,
116                             1,
117                         );
118                         enc_frame.planes[2].copy_from_raw_u8(
119                             in_frame.plane_data(2).unwrap(),
120                             in_frame.plane_stride()[2] as usize,
121                             1,
122                         );
123                     }
124 
125                     context.send_frame((
126                         enc_frame,
127                         Some(rav1e::data::FrameParameters {
128                             frame_type_override: if force_keyframe {
129                                 rav1e::prelude::FrameTypeOverride::Key
130                             } else {
131                                 rav1e::prelude::FrameTypeOverride::No
132                             },
133                             opaque: Some(Box::new(frame_number)),
134                         }),
135                     ))
136                 } else {
137                     context.send_frame(None)
138                 }
139             }
140             Context::Sixteen(ref mut context) => {
141                 if let Some((frame_number, in_frame)) = in_frame {
142                     let mut enc_frame = context.new_frame();
143                     enc_frame.planes[0].copy_from_raw_u8(
144                         in_frame.plane_data(0).unwrap(),
145                         in_frame.plane_stride()[0] as usize,
146                         2,
147                     );
148 
149                     if in_frame.n_planes() > 1 {
150                         enc_frame.planes[1].copy_from_raw_u8(
151                             in_frame.plane_data(1).unwrap(),
152                             in_frame.plane_stride()[1] as usize,
153                             2,
154                         );
155                         enc_frame.planes[2].copy_from_raw_u8(
156                             in_frame.plane_data(2).unwrap(),
157                             in_frame.plane_stride()[2] as usize,
158                             2,
159                         );
160                     }
161 
162                     context.send_frame((
163                         enc_frame,
164                         Some(rav1e::data::FrameParameters {
165                             frame_type_override: if force_keyframe {
166                                 rav1e::prelude::FrameTypeOverride::Key
167                             } else {
168                                 rav1e::prelude::FrameTypeOverride::No
169                             },
170                             opaque: Some(Box::new(frame_number)),
171                         }),
172                     ))
173                 } else {
174                     context.send_frame(None)
175                 }
176             }
177         }
178     }
179 
flush(&mut self)180     fn flush(&mut self) {
181         match self {
182             Context::Eight(ref mut context) => context.flush(),
183             Context::Sixteen(ref mut context) => context.flush(),
184         }
185     }
186 }
187 
188 struct State {
189     context: Context,
190     video_info: gst_video::VideoInfo,
191 }
192 
193 #[derive(Default)]
194 pub struct Rav1Enc {
195     state: Mutex<Option<State>>,
196     settings: Mutex<Settings>,
197 }
198 
199 static CAT: Lazy<gst::DebugCategory> = Lazy::new(|| {
200     gst::DebugCategory::new(
201         "rav1enc",
202         gst::DebugColorFlags::empty(),
203         Some("rav1e AV1 encoder"),
204     )
205 });
206 
207 #[glib::object_subclass]
208 impl ObjectSubclass for Rav1Enc {
209     const NAME: &'static str = "Rav1Enc";
210     type Type = super::Rav1Enc;
211     type ParentType = gst_video::VideoEncoder;
212 }
213 
214 impl ObjectImpl for Rav1Enc {
properties() -> &'static [glib::ParamSpec]215     fn properties() -> &'static [glib::ParamSpec] {
216         static PROPERTIES: Lazy<Vec<glib::ParamSpec>> = Lazy::new(|| {
217             vec![
218                 glib::ParamSpec::new_uint(
219                     "speed-preset",
220                     "Speed Preset",
221                     "Speed preset (10 fastest, 0 slowest)",
222                     0,
223                     10,
224                     DEFAULT_SPEED_PRESET,
225                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
226                 ),
227                 glib::ParamSpec::new_boolean(
228                     "low-latency",
229                     "Low Latency",
230                     "Low Latency",
231                     DEFAULT_LOW_LATENCY,
232                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
233                 ),
234                 glib::ParamSpec::new_uint64(
235                     "min-key-frame-interval",
236                     "Min Key Frame Interval",
237                     "Min Key Frame Interval",
238                     0,
239                     std::u64::MAX,
240                     DEFAULT_MIN_KEY_FRAME_INTERVAL,
241                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
242                 ),
243                 glib::ParamSpec::new_uint64(
244                     "max-key-frame-interval",
245                     "Max Key Frame Interval",
246                     "Max Key Frame Interval",
247                     0,
248                     std::u64::MAX,
249                     DEFAULT_MAX_KEY_FRAME_INTERVAL,
250                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
251                 ),
252                 glib::ParamSpec::new_int(
253                     "bitrate",
254                     "Bitrate",
255                     "Bitrate",
256                     0,
257                     std::i32::MAX,
258                     DEFAULT_BITRATE,
259                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
260                 ),
261                 glib::ParamSpec::new_uint(
262                     "quantizer",
263                     "Quantizer",
264                     "Quantizer",
265                     0,
266                     std::u32::MAX,
267                     DEFAULT_QUANTIZER as u32,
268                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
269                 ),
270                 glib::ParamSpec::new_uint(
271                     "tile-cols",
272                     "Tile Cols",
273                     "Tile Cols",
274                     0,
275                     std::u32::MAX,
276                     DEFAULT_TILE_COLS as u32,
277                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
278                 ),
279                 glib::ParamSpec::new_uint(
280                     "tile-rows",
281                     "Tile Rows",
282                     "Tile Rows",
283                     0,
284                     std::u32::MAX,
285                     DEFAULT_TILE_ROWS as u32,
286                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
287                 ),
288                 glib::ParamSpec::new_uint(
289                     "tiles",
290                     "Tiles",
291                     "Tiles",
292                     0,
293                     std::u32::MAX,
294                     DEFAULT_TILES as u32,
295                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
296                 ),
297                 glib::ParamSpec::new_uint(
298                     "threads",
299                     "Threads",
300                     "Threads",
301                     0,
302                     std::u32::MAX,
303                     DEFAULT_THREADS as u32,
304                     glib::ParamFlags::READWRITE | gst::PARAM_FLAG_MUTABLE_READY,
305                 ),
306             ]
307         });
308 
309         PROPERTIES.as_ref()
310     }
311 
set_property( &self, _obj: &Self::Type, _id: usize, value: &glib::Value, pspec: &glib::ParamSpec, )312     fn set_property(
313         &self,
314         _obj: &Self::Type,
315         _id: usize,
316         value: &glib::Value,
317         pspec: &glib::ParamSpec,
318     ) {
319         match pspec.name() {
320             "speed-preset" => {
321                 let mut settings = self.settings.lock().unwrap();
322                 settings.speed_preset = value.get().expect("type checked upstream");
323             }
324             "low-latency" => {
325                 let mut settings = self.settings.lock().unwrap();
326                 settings.low_latency = value.get().expect("type checked upstream");
327             }
328             "min-key-frame-interval" => {
329                 let mut settings = self.settings.lock().unwrap();
330                 settings.min_key_frame_interval = value.get().expect("type checked upstream");
331             }
332             "max-key-frame-interval" => {
333                 let mut settings = self.settings.lock().unwrap();
334                 settings.max_key_frame_interval = value.get().expect("type checked upstream");
335             }
336             "bitrate" => {
337                 let mut settings = self.settings.lock().unwrap();
338                 settings.bitrate = value.get().expect("type checked upstream");
339             }
340             "quantizer" => {
341                 let mut settings = self.settings.lock().unwrap();
342                 settings.quantizer = value.get::<u32>().expect("type checked upstream") as usize;
343             }
344             "tile-cols" => {
345                 let mut settings = self.settings.lock().unwrap();
346                 settings.tile_cols = value.get::<u32>().expect("type checked upstream") as usize;
347             }
348             "tile-rows" => {
349                 let mut settings = self.settings.lock().unwrap();
350                 settings.tile_rows = value.get::<u32>().expect("type checked upstream") as usize;
351             }
352             "tiles" => {
353                 let mut settings = self.settings.lock().unwrap();
354                 settings.tiles = value.get::<u32>().expect("type checked upstream") as usize;
355             }
356             "threads" => {
357                 let mut settings = self.settings.lock().unwrap();
358                 settings.threads = value.get::<u32>().expect("type checked upstream") as usize;
359             }
360             _ => unimplemented!(),
361         }
362     }
363 
property(&self, _obj: &Self::Type, _id: usize, pspec: &glib::ParamSpec) -> glib::Value364     fn property(&self, _obj: &Self::Type, _id: usize, pspec: &glib::ParamSpec) -> glib::Value {
365         match pspec.name() {
366             "speed-preset" => {
367                 let settings = self.settings.lock().unwrap();
368                 settings.speed_preset.to_value()
369             }
370             "low-latency" => {
371                 let settings = self.settings.lock().unwrap();
372                 settings.low_latency.to_value()
373             }
374             "min-key-frame-interval" => {
375                 let settings = self.settings.lock().unwrap();
376                 settings.min_key_frame_interval.to_value()
377             }
378             "max-key-frame-interval" => {
379                 let settings = self.settings.lock().unwrap();
380                 settings.max_key_frame_interval.to_value()
381             }
382             "bitrate" => {
383                 let settings = self.settings.lock().unwrap();
384                 settings.bitrate.to_value()
385             }
386             "quantizer" => {
387                 let settings = self.settings.lock().unwrap();
388                 (settings.quantizer as u32).to_value()
389             }
390             "tile-cols" => {
391                 let settings = self.settings.lock().unwrap();
392                 (settings.tile_cols as u32).to_value()
393             }
394             "tile-rows" => {
395                 let settings = self.settings.lock().unwrap();
396                 (settings.tile_rows as u32).to_value()
397             }
398             "tiles" => {
399                 let settings = self.settings.lock().unwrap();
400                 (settings.tiles as u32).to_value()
401             }
402             "threads" => {
403                 let settings = self.settings.lock().unwrap();
404                 (settings.threads as u32).to_value()
405             }
406             _ => unimplemented!(),
407         }
408     }
409 }
410 
411 impl ElementImpl for Rav1Enc {
metadata() -> Option<&'static gst::subclass::ElementMetadata>412     fn metadata() -> Option<&'static gst::subclass::ElementMetadata> {
413         static ELEMENT_METADATA: Lazy<gst::subclass::ElementMetadata> = Lazy::new(|| {
414             gst::subclass::ElementMetadata::new(
415                 "rav1e AV1 encoder",
416                 "Encoder/Video",
417                 "rav1e AV1 encoder",
418                 "Sebastian Dröge <sebastian@centricular.com>",
419             )
420         });
421 
422         Some(&*ELEMENT_METADATA)
423     }
424 
pad_templates() -> &'static [gst::PadTemplate]425     fn pad_templates() -> &'static [gst::PadTemplate] {
426         static PAD_TEMPLATES: Lazy<Vec<gst::PadTemplate>> = Lazy::new(|| {
427             let sink_caps = gst::Caps::new_simple(
428                 "video/x-raw",
429                 &[
430                     (
431                         "format",
432                         &gst::List::new(&[
433                             &gst_video::VideoFormat::I420.to_str(),
434                             &gst_video::VideoFormat::Y42b.to_str(),
435                             &gst_video::VideoFormat::Y444.to_str(),
436                             &gst_video::VideoFormat::I42010le.to_str(),
437                             &gst_video::VideoFormat::I42210le.to_str(),
438                             &gst_video::VideoFormat::Y44410le.to_str(),
439                             &gst_video::VideoFormat::I42012le.to_str(),
440                             &gst_video::VideoFormat::I42212le.to_str(),
441                             &gst_video::VideoFormat::Y44412le.to_str(),
442                             // &gst_video::VideoFormat::Gray8.to_str(),
443                         ]),
444                     ),
445                     ("width", &gst::IntRange::<i32>::new(1, std::i32::MAX)),
446                     ("height", &gst::IntRange::<i32>::new(1, std::i32::MAX)),
447                     (
448                         "framerate",
449                         &gst::FractionRange::new(
450                             gst::Fraction::new(0, 1),
451                             gst::Fraction::new(std::i32::MAX, 1),
452                         ),
453                     ),
454                 ],
455             );
456             let sink_pad_template = gst::PadTemplate::new(
457                 "sink",
458                 gst::PadDirection::Sink,
459                 gst::PadPresence::Always,
460                 &sink_caps,
461             )
462             .unwrap();
463 
464             let src_caps = gst::Caps::new_simple("video/x-av1", &[]);
465             let src_pad_template = gst::PadTemplate::new(
466                 "src",
467                 gst::PadDirection::Src,
468                 gst::PadPresence::Always,
469                 &src_caps,
470             )
471             .unwrap();
472 
473             vec![src_pad_template, sink_pad_template]
474         });
475 
476         PAD_TEMPLATES.as_ref()
477     }
478 }
479 
480 impl VideoEncoderImpl for Rav1Enc {
stop(&self, _element: &Self::Type) -> Result<(), gst::ErrorMessage>481     fn stop(&self, _element: &Self::Type) -> Result<(), gst::ErrorMessage> {
482         *self.state.lock().unwrap() = None;
483 
484         Ok(())
485     }
486 
487     // For the colorimetry mapping below
488     #[allow(clippy::wildcard_in_or_patterns)]
set_format( &self, element: &Self::Type, state: &gst_video::VideoCodecState<'static, gst_video::video_codec_state::Readable>, ) -> Result<(), gst::LoggableError>489     fn set_format(
490         &self,
491         element: &Self::Type,
492         state: &gst_video::VideoCodecState<'static, gst_video::video_codec_state::Readable>,
493     ) -> Result<(), gst::LoggableError> {
494         self.finish(element)
495             .map_err(|_| gst::loggable_error!(CAT, "Failed to drain"))?;
496 
497         let video_info = state.info();
498         gst_debug!(CAT, obj: element, "Setting format {:?}", video_info);
499 
500         let settings = self.settings.lock().unwrap();
501 
502         // TODO: More properties, HDR information
503         let cfg = config::Config::new()
504             .with_encoder_config(config::EncoderConfig {
505                 width: video_info.width() as usize,
506                 height: video_info.height() as usize,
507                 bit_depth: video_info.format_info().depth()[0] as usize,
508                 chroma_sampling: match video_info.format() {
509                     gst_video::VideoFormat::I420
510                     | gst_video::VideoFormat::I42010le
511                     | gst_video::VideoFormat::I42012le => color::ChromaSampling::Cs420,
512                     gst_video::VideoFormat::Y42b
513                     | gst_video::VideoFormat::I42210le
514                     | gst_video::VideoFormat::I42212le => color::ChromaSampling::Cs422,
515                     gst_video::VideoFormat::Y444
516                     | gst_video::VideoFormat::Y44410le
517                     | gst_video::VideoFormat::Y44412le => color::ChromaSampling::Cs444,
518                     // gst_video::VideoFormat::Gray8 => color::ChromaSampling::Cs400,
519                     _ => unreachable!(),
520                 },
521                 chroma_sample_position: match video_info.chroma_site() {
522                     gst_video::VideoChromaSite::H_COSITED => color::ChromaSamplePosition::Vertical,
523                     gst_video::VideoChromaSite::COSITED => color::ChromaSamplePosition::Colocated,
524                     _ => color::ChromaSamplePosition::Unknown,
525                 },
526                 pixel_range: match video_info.colorimetry().range() {
527                     gst_video::VideoColorRange::Range0_255 => color::PixelRange::Full,
528                     _ => color::PixelRange::Limited,
529                 },
530                 color_description: {
531                     let matrix = match video_info.colorimetry().matrix() {
532                         gst_video::VideoColorMatrix::Rgb => color::MatrixCoefficients::Identity,
533                         gst_video::VideoColorMatrix::Fcc => color::MatrixCoefficients::FCC,
534                         gst_video::VideoColorMatrix::Bt709 => color::MatrixCoefficients::BT709,
535                         gst_video::VideoColorMatrix::Bt601 => color::MatrixCoefficients::BT601,
536                         gst_video::VideoColorMatrix::Smpte240m => {
537                             color::MatrixCoefficients::SMPTE240
538                         }
539                         gst_video::VideoColorMatrix::Bt2020 => color::MatrixCoefficients::BT2020NCL,
540                         _ => color::MatrixCoefficients::Unspecified,
541                     };
542                     let transfer = match video_info.colorimetry().transfer() {
543                         gst_video::VideoTransferFunction::Gamma10 => {
544                             color::TransferCharacteristics::Linear
545                         }
546                         gst_video::VideoTransferFunction::Bt709 => {
547                             color::TransferCharacteristics::BT709
548                         }
549                         gst_video::VideoTransferFunction::Smpte240m => {
550                             color::TransferCharacteristics::SMPTE240
551                         }
552                         gst_video::VideoTransferFunction::Srgb => {
553                             color::TransferCharacteristics::SRGB
554                         }
555                         gst_video::VideoTransferFunction::Log100 => {
556                             color::TransferCharacteristics::Log100
557                         }
558                         gst_video::VideoTransferFunction::Log316 => {
559                             color::TransferCharacteristics::Log100Sqrt10
560                         }
561                         gst_video::VideoTransferFunction::Bt202012 => {
562                             color::TransferCharacteristics::BT2020_12Bit
563                         }
564                         gst_video::VideoTransferFunction::Gamma18
565                         | gst_video::VideoTransferFunction::Gamma20
566                         | gst_video::VideoTransferFunction::Gamma22
567                         | gst_video::VideoTransferFunction::Gamma28
568                         | gst_video::VideoTransferFunction::Adobergb
569                         | _ => color::TransferCharacteristics::Unspecified,
570                     };
571                     let primaries = match video_info.colorimetry().primaries() {
572                         gst_video::VideoColorPrimaries::Bt709 => color::ColorPrimaries::BT709,
573                         gst_video::VideoColorPrimaries::Bt470m => color::ColorPrimaries::BT470M,
574                         gst_video::VideoColorPrimaries::Bt470bg => color::ColorPrimaries::BT470BG,
575                         gst_video::VideoColorPrimaries::Smpte170m => color::ColorPrimaries::BT601,
576                         gst_video::VideoColorPrimaries::Smpte240m => {
577                             color::ColorPrimaries::SMPTE240
578                         }
579                         gst_video::VideoColorPrimaries::Film => color::ColorPrimaries::GenericFilm,
580                         gst_video::VideoColorPrimaries::Bt2020 => color::ColorPrimaries::BT2020,
581                         gst_video::VideoColorPrimaries::Adobergb | _ => {
582                             color::ColorPrimaries::Unspecified
583                         }
584                     };
585 
586                     Some(color::ColorDescription {
587                         color_primaries: primaries,
588                         transfer_characteristics: transfer,
589                         matrix_coefficients: matrix,
590                     })
591                 },
592                 speed_settings: config::SpeedSettings::from_preset(settings.speed_preset as usize),
593                 time_base: if video_info.fps() != gst::Fraction::new(0, 1) {
594                     data::Rational {
595                         num: *video_info.fps().numer() as u64,
596                         den: *video_info.fps().denom() as u64,
597                     }
598                 } else {
599                     data::Rational { num: 30, den: 1 }
600                 },
601                 low_latency: settings.low_latency,
602                 min_key_frame_interval: settings.min_key_frame_interval,
603                 max_key_frame_interval: settings.max_key_frame_interval,
604                 bitrate: settings.bitrate,
605                 quantizer: settings.quantizer,
606                 tile_cols: settings.tile_cols,
607                 tile_rows: settings.tile_rows,
608                 tiles: settings.tiles,
609                 ..Default::default()
610             })
611             .with_threads(settings.threads);
612         // TODO: RateControlConfig
613 
614         *self.state.lock().unwrap() = Some(State {
615             context: if video_info.format_info().depth()[0] > 8 {
616                 Context::Sixteen(cfg.new_context().map_err(|err| {
617                     gst::loggable_error!(CAT, "Failed to create context: {:?}", err)
618                 })?)
619             } else {
620                 Context::Eight(cfg.new_context().map_err(|err| {
621                     gst::loggable_error!(CAT, "Failed to create context: {:?}", err)
622                 })?)
623             },
624             video_info,
625         });
626 
627         let output_state = element
628             .set_output_state(gst::Caps::new_simple("video/x-av1", &[]), Some(state))
629             .map_err(|_| gst::loggable_error!(CAT, "Failed to set output state"))?;
630         element
631             .negotiate(output_state)
632             .map_err(|_| gst::loggable_error!(CAT, "Failed to negotiate"))?;
633 
634         self.parent_set_format(element, state)
635     }
636 
flush(&self, element: &Self::Type) -> bool637     fn flush(&self, element: &Self::Type) -> bool {
638         gst_debug!(CAT, obj: element, "Flushing");
639 
640         let mut state_guard = self.state.lock().unwrap();
641         if let Some(ref mut state) = *state_guard {
642             state.context.flush();
643             while let Ok(_) | Err(data::EncoderStatus::Encoded) = state.context.receive_packet() {
644                 gst_debug!(CAT, obj: element, "Dropping packet on flush",);
645             }
646         }
647 
648         true
649     }
650 
finish(&self, element: &Self::Type) -> Result<gst::FlowSuccess, gst::FlowError>651     fn finish(&self, element: &Self::Type) -> Result<gst::FlowSuccess, gst::FlowError> {
652         gst_debug!(CAT, obj: element, "Finishing");
653 
654         let mut state_guard = self.state.lock().unwrap();
655         if let Some(ref mut state) = *state_guard {
656             if let Err(data::EncoderStatus::Failure) = state.context.send_frame(None, false) {
657                 return Err(gst::FlowError::Error);
658             }
659             state.context.flush();
660             self.output_frames(element, state)?;
661         }
662 
663         Ok(gst::FlowSuccess::Ok)
664     }
665 
handle_frame( &self, element: &Self::Type, frame: gst_video::VideoCodecFrame, ) -> Result<gst::FlowSuccess, gst::FlowError>666     fn handle_frame(
667         &self,
668         element: &Self::Type,
669         frame: gst_video::VideoCodecFrame,
670     ) -> Result<gst::FlowSuccess, gst::FlowError> {
671         let mut state_guard = self.state.lock().unwrap();
672         let state = state_guard.as_mut().ok_or(gst::FlowError::NotNegotiated)?;
673 
674         self.output_frames(element, state)?;
675 
676         gst_debug!(
677             CAT,
678             obj: element,
679             "Sending frame {}",
680             frame.system_frame_number()
681         );
682 
683         let input_buffer = frame.input_buffer().expect("frame without input buffer");
684 
685         let in_frame =
686             gst_video::VideoFrameRef::from_buffer_ref_readable(&*input_buffer, &state.video_info)
687                 .map_err(|_| {
688                 gst::element_error!(
689                     element,
690                     gst::CoreError::Failed,
691                     ["Failed to map output buffer readable"]
692                 );
693                 gst::FlowError::Error
694             })?;
695 
696         match state.context.send_frame(
697             Some((frame.system_frame_number(), &in_frame)),
698             frame
699                 .flags()
700                 .contains(gst_video::VideoCodecFrameFlags::FORCE_KEYFRAME),
701         ) {
702             Ok(_) => {
703                 gst_debug!(
704                     CAT,
705                     obj: element,
706                     "Sent frame {}",
707                     frame.system_frame_number()
708                 );
709             }
710             Err(data::EncoderStatus::Failure) => {
711                 gst::element_error!(element, gst::CoreError::Failed, ["Failed to send frame"]);
712                 return Err(gst::FlowError::Error);
713             }
714             Err(_) => (),
715         }
716 
717         self.output_frames(element, state)
718     }
719 }
720 
721 impl Rav1Enc {
output_frames( &self, element: &super::Rav1Enc, state: &mut State, ) -> Result<gst::FlowSuccess, gst::FlowError>722     fn output_frames(
723         &self,
724         element: &super::Rav1Enc,
725         state: &mut State,
726     ) -> Result<gst::FlowSuccess, gst::FlowError> {
727         loop {
728             match state.context.receive_packet() {
729                 Ok((packet_type, packet_number, frame_number, packet_data)) => {
730                     gst_debug!(
731                         CAT,
732                         obj: element,
733                         "Received packet {} of size {}, frame type {:?}",
734                         packet_number,
735                         packet_data.len(),
736                         packet_type
737                     );
738 
739                     let mut frame = element.frame(frame_number as i32).expect("frame not found");
740 
741                     if packet_type == data::FrameType::KEY {
742                         frame.set_flags(gst_video::VideoCodecFrameFlags::SYNC_POINT);
743                     }
744                     let output_buffer = gst::Buffer::from_mut_slice(packet_data);
745                     frame.set_output_buffer(output_buffer);
746                     element.finish_frame(Some(frame))?;
747                 }
748                 Err(data::EncoderStatus::Encoded) => {
749                     gst_debug!(CAT, obj: element, "Encoded but not output frame yet",);
750                 }
751                 Err(data::EncoderStatus::Failure) => {
752                     gst::element_error!(
753                         element,
754                         gst::CoreError::Failed,
755                         ["Failed to receive frame"]
756                     );
757                     return Err(gst::FlowError::Error);
758                 }
759                 Err(err) => {
760                     gst_debug!(
761                         CAT,
762                         obj: element,
763                         "Soft error when receiving frame: {:?}",
764                         err
765                     );
766                     return Ok(gst::FlowSuccess::Ok);
767                 }
768             }
769         }
770     }
771 }
772