1use bytes::Bytes;
9
10use super::backend::{self, Backend};
11use crate::frame::{I420, Surface};
12use crate::{Error, Size};
13
14#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
22#[non_exhaustive]
23pub enum Codec {
24 #[default]
27 H264,
28 H265,
30}
31
32#[derive(Clone, Debug, Default, PartialEq, Eq)]
35#[non_exhaustive]
36pub enum Kind {
37 #[default]
40 Auto,
41 Hardware,
43 Software,
45 Named(String),
48}
49
50#[derive(Clone, Debug)]
56#[non_exhaustive]
57pub struct Config {
58 pub width: u32,
59 pub height: u32,
60 pub framerate: u32,
61 pub bitrate: Option<u64>,
64 pub gop: u32,
67 pub codec: Codec,
69 pub kind: Kind,
70}
71
72impl Config {
73 pub fn new(width: u32, height: u32, framerate: u32) -> Self {
76 Self {
77 width,
78 height,
79 framerate,
80 bitrate: None,
81 gop: framerate.saturating_mul(2).max(1),
83 codec: Codec::default(),
84 kind: Kind::Auto,
85 }
86 }
87
88 pub fn size(&self) -> Size {
90 Size::new(self.width, self.height)
91 }
92
93 pub(crate) fn resolved_bitrate(&self) -> u64 {
95 self.bitrate.unwrap_or_else(|| {
96 ((self.size().pixels() * self.framerate as u64) as f64 * 0.07) as u64
99 })
100 }
101}
102
103pub struct Encoder {
107 backend: Box<dyn Backend>,
108 codec: Codec,
109 size: Size,
110 bitrate: u64,
111}
112
113impl Encoder {
114 pub fn new(config: &Config) -> Result<Self, Error> {
116 if config.framerate == 0 {
120 return Err(Error::InvalidFramerate(0));
121 }
122 let size = config.size();
124 size.validate("encoder")?;
125
126 let backend = backend::open(config)?;
127 Ok(Self {
128 backend,
129 codec: config.codec,
130 size,
131 bitrate: config.resolved_bitrate(),
132 })
133 }
134
135 pub fn name(&self) -> &str {
137 self.backend.name()
138 }
139
140 pub fn size(&self) -> Size {
142 self.size
143 }
144
145 pub fn bitrate(&self) -> u64 {
149 self.bitrate
150 }
151
152 pub fn set_bitrate(&mut self, bitrate: u64) -> Result<(), Error> {
167 if bitrate == self.bitrate {
168 return Ok(());
169 }
170 self.backend.set_bitrate(bitrate)?;
171 self.bitrate = bitrate;
174 Ok(())
175 }
176
177 pub fn codec(&self) -> Codec {
180 self.codec
181 }
182
183 pub fn encode_rgba(&mut self, rgba: &[u8], size: Size, keyframe: bool) -> Result<Vec<Bytes>, Error> {
190 self.check_frame(size, rgba.len(), size.pixels() as usize * 4, "RGBA")?;
194
195 let frame = Surface::I420(I420::from_rgba(rgba, size.width * 4, size.width, size.height)?);
196 self.encode(&frame, keyframe)
197 }
198
199 pub fn encode_i420(&mut self, i420: &[u8], size: Size, keyframe: bool) -> Result<Vec<Bytes>, Error> {
211 self.check_frame(size, i420.len(), I420::len(size.width, size.height), "I420")?;
212
213 let frame = Surface::I420(I420 {
214 width: size.width,
215 height: size.height,
216 data: i420.to_vec(),
217 });
218 self.encode(&frame, keyframe)
219 }
220
221 fn check_frame(&self, size: Size, got: usize, expected: usize, what: &str) -> Result<(), Error> {
229 if size != self.size {
230 return Err(Error::Codec(anyhow::anyhow!(
231 "frame {size} does not match encoder {}",
232 self.size
233 )));
234 }
235 if got != expected {
236 return Err(Error::Codec(anyhow::anyhow!(
237 "{what} buffer is {got} bytes, expected {expected} for {size}"
238 )));
239 }
240 Ok(())
241 }
242
243 pub fn encode(&mut self, frame: &Surface, keyframe: bool) -> Result<Vec<Bytes>, Error> {
253 let size = Size::new(frame.width(), frame.height());
254 if size != self.size {
255 return Err(Error::Codec(anyhow::anyhow!(
256 "frame {size} does not match encoder {}",
257 self.size
258 )));
259 }
260 self.backend.encode(frame, keyframe)
261 }
262
263 pub fn finish(mut self) -> Result<Vec<Bytes>, Error> {
268 self.backend.finish()
269 }
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275
276 fn gray_rgba(width: u32, height: u32) -> Vec<u8> {
278 vec![0x80u8; width as usize * height as usize * 4]
279 }
280
281 #[test]
282 fn software_encoder_emits_annexb() {
283 let config = Config {
284 kind: Kind::Software,
285 ..Config::new(320, 240, 30)
286 };
287 let mut encoder = Encoder::new(&config).expect("openh264 is vendored, always available");
288 assert_eq!(encoder.name(), "openh264");
289
290 let frame = gray_rgba(320, 240);
291 let mut packets = Vec::new();
292 for i in 0..30 {
293 packets.extend(encoder.encode_rgba(&frame, Size::new(320, 240), i == 0).unwrap());
294 }
295 packets.extend(encoder.finish().unwrap());
296
297 assert!(!packets.is_empty(), "encoder produced no packets");
298
299 let first = &packets[0];
302 let has_start_code = first.starts_with(&[0, 0, 0, 1]) || first.starts_with(&[0, 0, 1]);
303 assert!(
304 has_start_code,
305 "first packet is not Annex-B: {:02x?}",
306 &first[..first.len().min(8)]
307 );
308 }
309
310 #[test]
311 fn encode_rgba_emits_annexb() {
312 let config = Config {
313 kind: Kind::Software,
314 ..Config::new(320, 240, 30)
315 };
316 let mut encoder = Encoder::new(&config).unwrap();
317
318 let rgba = gray_rgba(320, 240);
319 let mut packets = encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
320 packets.extend(encoder.finish().unwrap());
321 assert!(!packets.is_empty());
322 assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
323 }
324
325 #[test]
326 fn encode_i420_emits_annexb() {
327 let config = Config {
328 kind: Kind::Software,
329 ..Config::new(320, 240, 30)
330 };
331 let mut encoder = Encoder::new(&config).unwrap();
332
333 let data = vec![0x80u8; I420::len(320, 240)];
335 let mut packets = encoder.encode_i420(&data, Size::new(320, 240), true).unwrap();
336 packets.extend(encoder.finish().unwrap());
337 assert!(!packets.is_empty());
338 assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
339 }
340
341 #[test]
344 fn encode_i420_rejects_wrong_size() {
345 let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
346 return;
347 };
348 assert!(matches!(
349 encoder.encode_i420(&[0u8; 16], Size::new(320, 240), false),
350 Err(Error::Codec(_))
351 ));
352 }
353
354 #[test]
355 fn encode_rgba_rejects_short_buffer() {
356 let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
357 return;
358 };
359 assert!(matches!(
361 encoder.encode_rgba(&[0u8; 16], Size::new(320, 240), false),
362 Err(Error::Codec(_))
363 ));
364 }
365
366 #[test]
369 fn encode_rgba_rejects_dimension_mismatch() {
370 let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
371 return;
372 };
373 let rgba = gray_rgba(640, 480);
374 assert!(matches!(
375 encoder.encode_rgba(&rgba, Size::new(640, 480), false),
376 Err(Error::Codec(_))
377 ));
378 }
379
380 #[test]
383 fn encode_i420_rejects_oversized_buffer() {
384 let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
385 return;
386 };
387 let data = vec![0x80u8; I420::len(640, 480)];
388 assert!(matches!(
389 encoder.encode_i420(&data, Size::new(640, 480), false),
390 Err(Error::Codec(_))
391 ));
392 }
393
394 #[test]
398 fn encode_rejects_transposed_frame() {
399 let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
400 return;
401 };
402
403 let rgba = gray_rgba(240, 320);
404 assert_eq!(rgba.len(), gray_rgba(320, 240).len(), "the byte counts must collide");
405 assert!(matches!(
406 encoder.encode_rgba(&rgba, Size::new(240, 320), false),
407 Err(Error::Codec(_))
408 ));
409
410 let i420 = vec![0x80u8; I420::len(240, 320)];
411 assert_eq!(i420.len(), I420::len(320, 240), "the byte counts must collide");
412 assert!(matches!(
413 encoder.encode_i420(&i420, Size::new(240, 320), false),
414 Err(Error::Codec(_))
415 ));
416 }
417
418 #[test]
419 fn new_rejects_zero_framerate() {
420 let config = Config::new(320, 240, 0);
423 assert!(matches!(Encoder::new(&config), Err(Error::InvalidFramerate(0))));
424 }
425
426 #[test]
427 fn unknown_named_encoder_errors() {
428 let config = Config {
429 kind: Kind::Named("definitely_not_a_codec".into()),
430 ..Config::new(320, 240, 30)
431 };
432 assert!(matches!(Encoder::new(&config), Err(Error::NoEncoder(_))));
433 }
434
435 #[cfg(target_os = "macos")]
439 #[test]
440 fn videotoolbox_emits_annexb_keyframe() {
441 let config = Config {
442 kind: Kind::Named("videotoolbox".into()),
443 ..Config::new(320, 240, 30)
444 };
445 let mut encoder = Encoder::new(&config).expect("videotoolbox is available on macOS");
446 assert_eq!(encoder.name(), "videotoolbox");
447
448 let frame = gray_rgba(320, 240);
449 let mut packets = Vec::new();
450 for i in 0..10 {
451 packets.extend(encoder.encode_rgba(&frame, Size::new(320, 240), i == 0).unwrap());
452 }
453 packets.extend(encoder.finish().unwrap());
454
455 assert!(!packets.is_empty(), "encoder produced no packets");
456 let first = &packets[0];
457 assert!(
458 first.starts_with(&[0, 0, 0, 1]) || first.starts_with(&[0, 0, 1]),
459 "first packet is not Annex-B"
460 );
461
462 let types = nal_types(first);
465 assert!(types.contains(&7), "no SPS in first packet: {types:?}");
466 assert!(types.contains(&8), "no PPS in first packet: {types:?}");
467 assert!(types.contains(&5), "first packet is not an IDR: {types:?}");
468 }
469
470 #[cfg(target_os = "macos")]
474 #[test]
475 fn videotoolbox_emits_annexb_keyframe_h265() {
476 let config = Config {
477 codec: Codec::H265,
478 kind: Kind::Named("videotoolbox".into()),
479 ..Config::new(320, 240, 30)
480 };
481 let mut encoder = Encoder::new(&config).expect("videotoolbox HEVC is available on macOS");
482 assert_eq!(encoder.name(), "videotoolbox");
483 assert_eq!(encoder.codec(), Codec::H265);
484
485 let frame = gray_rgba(320, 240);
486 let mut packets = Vec::new();
487 for i in 0..10 {
488 packets.extend(encoder.encode_rgba(&frame, Size::new(320, 240), i == 0).unwrap());
489 }
490 packets.extend(encoder.finish().unwrap());
491
492 assert!(!packets.is_empty(), "encoder produced no packets");
493 let first = &packets[0];
494 assert!(
495 first.starts_with(&[0, 0, 0, 1]) || first.starts_with(&[0, 0, 1]),
496 "first packet is not Annex-B"
497 );
498
499 let types = hevc_nal_types(first);
502 assert!(types.contains(&32), "no VPS in first packet: {types:?}");
503 assert!(types.contains(&33), "no SPS in first packet: {types:?}");
504 assert!(types.contains(&34), "no PPS in first packet: {types:?}");
505 assert!(
506 types.iter().any(|t| (16..=23).contains(t)),
507 "first packet is not an IRAP: {types:?}"
508 );
509 }
510
511 #[cfg(target_os = "macos")]
513 fn hevc_nal_types(annexb: &[u8]) -> Vec<u8> {
514 let mut types = Vec::new();
515 let mut i = 0;
516 while i + 3 < annexb.len() {
517 if annexb[i..i + 3] == [0, 0, 1] {
518 types.push((annexb[i + 3] >> 1) & 0x3f);
519 i += 3;
520 } else {
521 i += 1;
522 }
523 }
524 types
525 }
526
527 #[cfg(target_os = "macos")]
530 #[test]
531 fn videotoolbox_encodes_surface_zero_copy() {
532 let config = Config {
533 kind: Kind::Named("videotoolbox".into()),
534 ..Config::new(320, 240, 30)
535 };
536 let mut encoder = Encoder::new(&config).unwrap();
537
538 let mut packets = Vec::new();
539 for i in 0..10 {
540 let frame = Surface::PixelBuffer(nv12_surface(320, 240));
541 packets.extend(encoder.encode(&frame, i == 0).unwrap());
542 }
543 packets.extend(encoder.finish().unwrap());
544
545 assert!(!packets.is_empty());
546 let types = nal_types(&packets[0]);
547 assert!(
548 types.contains(&7) && types.contains(&8) && types.contains(&5),
549 "no IDR: {types:?}"
550 );
551 }
552
553 #[cfg(target_os = "macos")]
556 #[test]
557 fn openh264_downloads_surface() {
558 let config = Config {
559 kind: Kind::Software,
560 ..Config::new(320, 240, 30)
561 };
562 let mut encoder = Encoder::new(&config).unwrap();
563
564 let frame = Surface::PixelBuffer(nv12_surface(320, 240));
565 let mut packets = encoder.encode(&frame, true).unwrap();
566 packets.extend(encoder.finish().unwrap());
567
568 assert!(!packets.is_empty());
569 assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
570 }
571
572 #[cfg(target_os = "macos")]
575 fn nv12_surface(width: u32, height: u32) -> crate::frame::macos::PixelBuffer {
576 use std::ptr::{self, NonNull};
577
578 use objc2_core_foundation::CFRetained;
579 use objc2_core_video::{
580 CVPixelBuffer, CVPixelBufferCreate, CVPixelBufferGetBaseAddressOfPlane, CVPixelBufferGetBytesPerRowOfPlane,
581 CVPixelBufferLockBaseAddress, CVPixelBufferLockFlags, CVPixelBufferUnlockBaseAddress,
582 kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
583 };
584
585 let mut raw: *mut CVPixelBuffer = ptr::null_mut();
586 let status = unsafe {
587 CVPixelBufferCreate(
588 None,
589 width as usize,
590 height as usize,
591 kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
592 None,
593 NonNull::new(&mut raw).unwrap(),
594 )
595 };
596 assert_eq!(status, 0, "CVPixelBufferCreate failed");
597 let buffer = unsafe { CFRetained::from_raw(NonNull::new(raw).unwrap()) };
598
599 let flags = CVPixelBufferLockFlags(0);
600 assert_eq!(unsafe { CVPixelBufferLockBaseAddress(&buffer, flags) }, 0);
601 for (plane, rows) in [(0usize, height as usize), (1usize, height as usize / 2)] {
602 let base = CVPixelBufferGetBaseAddressOfPlane(&buffer, plane) as *mut u8;
603 let stride = CVPixelBufferGetBytesPerRowOfPlane(&buffer, plane);
604 unsafe { ptr::write_bytes(base, 128, stride * rows) };
605 }
606 unsafe { CVPixelBufferUnlockBaseAddress(&buffer, flags) };
607
608 crate::frame::macos::PixelBuffer::new(buffer, width, height)
609 }
610
611 #[cfg(any(target_os = "macos", target_os = "windows"))]
614 fn nal_types(annexb: &[u8]) -> Vec<u8> {
615 let mut types = Vec::new();
616 let mut i = 0;
617 while i + 3 < annexb.len() {
618 if annexb[i..i + 3] == [0, 0, 1] {
619 types.push(annexb[i + 3] & 0x1f);
620 i += 3;
621 } else {
622 i += 1;
623 }
624 }
625 types
626 }
627
628 #[cfg(target_os = "windows")]
632 #[test]
633 #[ignore]
634 fn mediafoundation_cpu_rgba() {
635 let config = Config {
636 kind: Kind::Named("mediafoundation".into()),
637 ..Config::new(640, 480, 30)
638 };
639 let mut encoder = Encoder::new(&config).expect("hardware H.264 encoder available");
640 assert_eq!(encoder.name(), "mediafoundation");
641
642 let frame = gray_rgba(640, 480);
643 let mut packets = Vec::new();
644 for i in 0..30 {
645 packets.extend(encoder.encode_rgba(&frame, Size::new(640, 480), i == 0).unwrap());
646 }
647 packets.extend(encoder.finish().unwrap());
648
649 assert!(!packets.is_empty(), "encoder produced no packets");
650 let types = nal_types(&packets[0]);
651 assert!(types.contains(&7), "no SPS in first packet: {types:?}");
652 assert!(types.contains(&8), "no PPS in first packet: {types:?}");
653 assert!(types.contains(&5), "first packet is not an IDR: {types:?}");
654 }
655
656 #[cfg(target_os = "windows")]
660 #[tokio::test]
661 #[ignore]
662 async fn mediafoundation_camera_texture() {
663 let mut camera = crate::capture::open(&crate::capture::Config::default())
664 .await
665 .expect("open default camera");
666 let (w, h) = (camera.width(), camera.height());
667
668 let config = Config {
669 kind: Kind::Named("mediafoundation".into()),
670 ..Config::new(w, h, camera.framerate().unwrap_or(30))
671 };
672 let mut encoder = Encoder::new(&config).expect("hardware H.264 encoder available");
673
674 let mut packets = Vec::new();
675 let mut textures = 0;
676 for i in 0..30 {
677 let frame = camera.read().await.expect("frame, not end of stream");
678 if matches!(frame, Surface::Texture(_)) {
679 textures += 1;
680 }
681 packets.extend(encoder.encode(&frame, i == 0).unwrap());
682 }
683 packets.extend(encoder.finish().unwrap());
684
685 assert!(textures > 0, "capture never produced a GPU texture");
688 assert!(!packets.is_empty(), "encoder produced no packets");
689 let types = nal_types(&packets[0]);
690 assert!(
691 types.contains(&7) && types.contains(&8) && types.contains(&5),
692 "no IDR: {types:?}"
693 );
694 }
695
696 #[test]
700 fn set_bitrate_retunes_software_encoder() {
701 let config = Config {
702 kind: Kind::Software,
703 ..Config::new(320, 240, 30)
704 };
705 let mut encoder = Encoder::new(&config).unwrap();
706 let rgba = gray_rgba(320, 240);
707
708 let opened = encoder.bitrate();
709 assert_eq!(opened, config.resolved_bitrate());
710
711 encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
713
714 let halved = opened / 2;
715 encoder.set_bitrate(halved).unwrap();
716 assert_eq!(encoder.bitrate(), halved);
717
718 let packets = encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
720 assert!(!packets.is_empty(), "encoder produced nothing after a retune");
721 assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
722 }
723
724 #[test]
728 fn set_bitrate_before_the_first_frame_is_deferred() {
729 let config = Config {
730 kind: Kind::Software,
731 ..Config::new(320, 240, 30)
732 };
733 let mut encoder = Encoder::new(&config).unwrap();
734
735 let halved = encoder.bitrate() / 2;
736 encoder.set_bitrate(halved).expect("a retune before the first frame");
737 assert_eq!(encoder.bitrate(), halved);
738
739 let rgba = gray_rgba(320, 240);
741 let packets = encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
742 assert!(!packets.is_empty());
743
744 encoder.set_bitrate(halved / 2).unwrap();
746 assert!(encoder.encode_rgba(&rgba, Size::new(320, 240), false).is_ok());
747 }
748
749 #[test]
752 fn set_bitrate_to_current_is_a_noop() {
753 let config = Config {
754 kind: Kind::Software,
755 ..Config::new(320, 240, 30)
756 };
757 let mut encoder = Encoder::new(&config).unwrap();
758
759 let opened = encoder.bitrate();
760 encoder.set_bitrate(opened).unwrap();
761 assert_eq!(encoder.bitrate(), opened);
762 }
763
764 #[test]
765 fn default_bitrate_scales_with_resolution() {
766 let small = Config::new(320, 240, 30).resolved_bitrate();
767 let large = Config::new(1920, 1080, 30).resolved_bitrate();
768 assert!(large > small);
769 assert!(small > 0);
770 }
771}