1use anyhow::Result;
2use core_foundation::base::{OSStatus, TCFType};
3use media::{
4 core_media::{CMSampleBufferRef, CMSampleTimingInfo, CMVideoCodecType},
5 core_video::CVImageBuffer,
6 video_toolbox::{VTCompressionSession, VTEncodeInfoFlags},
7};
8use std::ffi::c_void;
9
10pub struct CompressionSession<F> {
11 session: VTCompressionSession,
12 output_callback: Box<F>,
13}
14
15impl<F: 'static + Send + FnMut(OSStatus, VTEncodeInfoFlags, CMSampleBufferRef)>
16 CompressionSession<F>
17{
18 pub fn new(width: usize, height: usize, codec: CMVideoCodecType, callback: F) -> Result<Self> {
19 let callback = Box::new(callback);
20 let session = VTCompressionSession::new(
21 width,
22 height,
23 codec,
24 Some(Self::output_callback),
25 callback.as_ref() as *const _ as *const c_void,
26 )?;
27 Ok(Self {
28 session,
29 output_callback: callback,
30 })
31 }
32
33 pub fn encode_frame(&self, buffer: &CVImageBuffer, timing: CMSampleTimingInfo) -> Result<()> {
34 self.session.encode_frame(
35 buffer.as_concrete_TypeRef(),
36 timing.presentationTimeStamp,
37 timing.duration,
38 )
39 }
40
41 extern "C" fn output_callback(
42 output_callback_ref_con: *mut c_void,
43 _: *mut c_void,
44 status: OSStatus,
45 flags: VTEncodeInfoFlags,
46 sample_buffer: CMSampleBufferRef,
47 ) {
48 let callback = unsafe { &mut *(output_callback_ref_con as *mut F) };
49 callback(status, flags, sample_buffer);
50 }
51}
52
53// unsafe extern "C" fn output(
54// output_callback_ref_con: *mut c_void,
55// source_frame_ref_con: *mut c_void,
56// status: OSStatus,
57// info_flags: VTEncodeInfoFlags,
58// sample_buffer: CMSampleBufferRef,
59// ) {
60// if status != 0 {
61// println!("error encoding frame, code: {}", status);
62// return;
63// }
64// let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
65
66// let mut is_iframe = false;
67// let attachments = sample_buffer.attachments();
68// if let Some(attachments) = attachments.first() {
69// is_iframe = attachments
70// .find(bindings::kCMSampleAttachmentKey_NotSync as CFStringRef)
71// .map_or(true, |not_sync| {
72// CFBooleanGetValue(*not_sync as CFBooleanRef)
73// });
74// }
75
76// const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
77// if is_iframe {
78// let format_description = sample_buffer.format_description();
79// for ix in 0..format_description.h264_parameter_set_count() {
80// let parameter_set = format_description.h264_parameter_set_at_index(ix);
81// stream.extend(START_CODE);
82// stream.extend(parameter_set);
83// }
84// }
85
86// println!("YO!");
87// }
88
89// static void videoFrameFinishedEncoding(void *outputCallbackRefCon,
90// void *sourceFrameRefCon,
91// OSStatus status,
92// VTEncodeInfoFlags infoFlags,
93// CMSampleBufferRef sampleBuffer) {
94// // Check if there were any errors encoding
95// if (status != noErr) {
96// NSLog(@"Error encoding video, err=%lld", (int64_t)status);
97// return;
98// }
99
100// // In this example we will use a NSMutableData object to store the
101// // elementary stream.
102// NSMutableData *elementaryStream = [NSMutableData data];
103
104// // Find out if the sample buffer contains an I-Frame.
105// // If so we will write the SPS and PPS NAL units to the elementary stream.
106// BOOL isIFrame = NO;
107// CFArrayRef attachmentsArray = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, 0);
108// if (CFArrayGetCount(attachmentsArray)) {
109// CFBooleanRef notSync;
110// CFDictionaryRef dict = CFArrayGetValueAtIndex(attachmentsArray, 0);
111// BOOL keyExists = CFDictionaryGetValueIfPresent(dict,
112// kCMSampleAttachmentKey_NotSync,
113// (const void **)¬Sync);
114// // An I-Frame is a sync frame
115// isIFrame = !keyExists || !CFBooleanGetValue(notSync);
116// }
117
118// // This is the start code that we will write to
119// // the elementary stream before every NAL unit
120// static const size_t startCodeLength = 4;
121// static const uint8_t startCode[] = {0x00, 0x00, 0x00, 0x01};
122
123// // Write the SPS and PPS NAL units to the elementary stream before every I-Frame
124// if (isIFrame) {
125// CMFormatDescriptionRef description = CMSampleBufferGetFormatDescription(sampleBuffer);
126
127// // Find out how many parameter sets there are
128// size_t numberOfParameterSets;
129// CMVideoFormatDescriptionGetH264ParameterSetAtIndex(description,
130// 0, NULL, NULL,
131// &numberOfParameterSets,
132// NULL);
133
134// // Write each parameter set to the elementary stream
135// for (int i = 0; i < numberOfParameterSets; i++) {
136// const uint8_t *parameterSetPointer;
137// size_t parameterSetLength;
138// CMVideoFormatDescriptionGetH264ParameterSetAtIndex(description,
139// i,
140// ¶meterSetPointer,
141// ¶meterSetLength,
142// NULL, NULL);
143
144// // Write the parameter set to the elementary stream
145// [elementaryStream appendBytes:startCode length:startCodeLength];
146// [elementaryStream appendBytes:parameterSetPointer length:parameterSetLength];
147// }
148// }
149
150// // Get a pointer to the raw AVCC NAL unit data in the sample buffer
151// size_t blockBufferLength;
152// uint8_t *bufferDataPointer = NULL;
153// CMBlockBufferGetDataPointer(CMSampleBufferGetDataBuffer(sampleBuffer),
154// 0,
155// NULL,
156// &blockBufferLength,
157// (char **)&bufferDataPointer);
158
159// // Loop through all the NAL units in the block buffer
160// // and write them to the elementary stream with
161// // start codes instead of AVCC length headers
162// size_t bufferOffset = 0;
163// static const int AVCCHeaderLength = 4;
164// while (bufferOffset < blockBufferLength - AVCCHeaderLength) {
165// // Read the NAL unit length
166// uint32_t NALUnitLength = 0;
167// memcpy(&NALUnitLength, bufferDataPointer + bufferOffset, AVCCHeaderLength);
168// // Convert the length value from Big-endian to Little-endian
169// NALUnitLength = CFSwapInt32BigToHost(NALUnitLength);
170// // Write start code to the elementary stream
171// [elementaryStream appendBytes:startCode length:startCodeLength];
172// // Write the NAL unit without the AVCC length header to the elementary stream
173// [elementaryStream appendBytes:bufferDataPointer + bufferOffset + AVCCHeaderLength
174// length:NALUnitLength];
175// // Move to the next NAL unit in the block buffer
176// bufferOffset += AVCCHeaderLength + NALUnitLength;
177// }
178// }