//
|
// H264Encoder.swift
|
// LiveProject
|
//
|
// Created by 倪路朋 on 6/28/25.
|
//
|
import Foundation
|
import AVFoundation
|
import VideoToolbox
|
|
class H264Encoder {
|
private var session: VTCompressionSession?
|
private let width: Int
|
private let height: Int
|
private let fps: Int
|
private let bitrate: Int
|
|
private let converter : PixelBufferConverter = PixelBufferConverter()
|
|
var onEncoded: ((Data, CMTime, Bool) -> Void)?
|
|
init(width: Int, height: Int, fps: Int, bitrate: Int) {
|
self.width = width
|
self.height = height
|
self.fps = fps
|
self.bitrate = bitrate
|
setupSession()
|
}
|
|
private func setupSession() {
|
let status = VTCompressionSessionCreate(
|
allocator: nil,
|
width: Int32(width),
|
height: Int32(height),
|
codecType: kCMVideoCodecType_H264,
|
encoderSpecification: nil,
|
imageBufferAttributes: nil,
|
compressedDataAllocator: nil,
|
outputCallback: encodeCallback,
|
refcon: UnsafeMutableRawPointer(Unmanaged.passRetained(self).toOpaque()),
|
compressionSessionOut: &session
|
)
|
|
guard status == noErr, let session = session else {
|
print("❌ Failed to create session: \(status)")
|
return
|
}
|
|
VTSessionSetProperty(session, key: kVTCompressionPropertyKey_RealTime, value: kCFBooleanTrue)
|
VTSessionSetProperty(session, key: kVTCompressionPropertyKey_AverageBitRate, value: bitrate as CFTypeRef)
|
let frameInterval = Int(fps)
|
VTSessionSetProperty(session, key: kVTCompressionPropertyKey_MaxKeyFrameInterval, value: frameInterval as CFTypeRef)
|
VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ExpectedFrameRate, value: frameInterval as CFTypeRef)
|
|
VTCompressionSessionPrepareToEncodeFrames(session)
|
}
|
|
func encode(pixelBuffer: CVPixelBuffer, pts: CMTime) {
|
guard let session = session else {
|
print("❌ Session is nil")
|
return
|
}
|
|
let format = CVPixelBufferGetPixelFormatType(pixelBuffer)
|
let supportedFormats: [OSType] = [
|
kCVPixelFormatType_420YpCbCr8BiPlanarFullRange,
|
kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
|
]
|
switch format{
|
case kCVPixelFormatType_32BGRA:
|
print("32BGRA")
|
break;
|
case kCVPixelFormatType_32ARGB:
|
print("32ARGB")
|
break;
|
default:
|
print("????")
|
break;
|
}
|
|
if let buffer = converter.convertBGRAtoNV12(pixelBuffer) {
|
print("converter \(decodeOSType(CVPixelBufferGetPixelFormatType(buffer)))")
|
let timestamp = CMTimeMake(value: Int64(CACurrentMediaTime() * 1000), timescale: 1000)
|
|
var flags = VTEncodeInfoFlags()
|
let status = VTCompressionSessionEncodeFrame(
|
session,
|
imageBuffer:buffer ,
|
presentationTimeStamp: pts,
|
duration: .invalid,
|
frameProperties: nil,
|
sourceFrameRefcon: nil,
|
infoFlagsOut: &flags
|
)
|
if status != noErr {
|
print("❌ Encoding failed: \(status)")
|
}
|
}
|
|
}
|
|
func finish() {
|
guard let session = session else { return }
|
VTCompressionSessionCompleteFrames(session, untilPresentationTimeStamp: .invalid)
|
}
|
|
func invalidate() {
|
guard let session = session else { return }
|
VTCompressionSessionInvalidate(session)
|
self.session = nil
|
}
|
|
func decodeOSType(_ format: OSType) -> String {
|
let characters = [
|
UInt8((format >> 24) & 0xFF),
|
UInt8((format >> 16) & 0xFF),
|
UInt8((format >> 8) & 0xFF),
|
UInt8(format & 0xFF)
|
]
|
return String(bytes: characters, encoding: .ascii) ?? "????"
|
}
|
}
|
|
// MARK: - VideoToolbox Callback
|
|
private func encodeCallback(
|
outputCallbackRefCon: UnsafeMutableRawPointer?,
|
sourceFrameRefCon: UnsafeMutableRawPointer?,
|
status: OSStatus,
|
infoFlags: VTEncodeInfoFlags,
|
sampleBuffer: CMSampleBuffer?
|
) {
|
guard
|
status == noErr,
|
let sampleBuffer = sampleBuffer,
|
CMSampleBufferDataIsReady(sampleBuffer),
|
let ref = outputCallbackRefCon
|
else { return }
|
|
let encoder = Unmanaged<H264Encoder>.fromOpaque(ref).takeUnretainedValue()
|
|
guard let blockBuffer = CMSampleBufferGetDataBuffer(sampleBuffer) else { return }
|
|
var length = 0
|
var dataPointer: UnsafeMutablePointer<Int8>?
|
guard CMBlockBufferGetDataPointer(blockBuffer, atOffset: 0, lengthAtOffsetOut: nil, totalLengthOut: &length, dataPointerOut: &dataPointer) == noErr else {
|
return
|
}
|
|
let data = Data(bytes: dataPointer!, count: length)
|
|
var isKeyframe = true
|
if let attachments = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, createIfNecessary: false),
|
let dict = CFArrayGetValueAtIndex(attachments, 0) {
|
let d = unsafeBitCast(dict, to: CFDictionary.self) as NSDictionary
|
if let notSync = d[kCMSampleAttachmentKey_NotSync] as? Bool {
|
isKeyframe = !notSync
|
}
|
}
|
|
let pts = CMSampleBufferGetPresentationTimeStamp(sampleBuffer)
|
encoder.onEncoded?(data, pts, isKeyframe)
|
|
// ⚠️ 这里必须 release 与 passRetained 配对
|
Unmanaged<H264Encoder>.fromOpaque(ref).release()
|
}
|
