//
|
// MetalPixelConverter.swift
|
// LiveProject
|
//
|
// Created by 倪路朋 on 7/1/25.
|
//
|
|
import Metal
|
import MetalKit
|
import CoreVideo
|
|
final class MetalPixelConverter {
|
|
// MARK: - Properties
|
private let device: MTLDevice
|
private var textureCache: CVMetalTextureCache?
|
private var commandQueue: MTLCommandQueue?
|
private var computePipeline: MTLComputePipelineState?
|
|
// MARK: - Initialization
|
init?(metalDevice: MTLDevice? = MTLCreateSystemDefaultDevice()) {
|
guard let device = metalDevice else {
|
print("⚠️ Metal 不可用")
|
return nil
|
}
|
self.device = device
|
|
// 1. 创建纹理缓存
|
guard CVMetalTextureCacheCreate(
|
kCFAllocatorDefault,
|
nil,
|
device,
|
nil,
|
&textureCache
|
) == kCVReturnSuccess else {
|
print("❌ 创建 Metal 纹理缓存失败")
|
return nil
|
}
|
|
// 2. 创建命令队列
|
self.commandQueue = device.makeCommandQueue()
|
|
// 3. 加载着色器
|
do {
|
self.computePipeline = try makeComputePipeline(device: device)
|
} catch {
|
print("❌ 加载着色器失败: \(error)")
|
return nil
|
}
|
}
|
|
// MARK: - Public Methods
|
/// 将 BGRA CVPixelBuffer 转换为 NV12 CVPixelBuffer
|
func convertBGRAtoNV12(
|
_ inputBuffer: CVPixelBuffer,
|
completion: @escaping (Result<CVPixelBuffer, Error>) -> Void
|
) {
|
// 0. 验证输入格式
|
guard CVPixelBufferGetPixelFormatType(inputBuffer) == kCVPixelFormatType_32BGRA else {
|
completion(.failure(ConversionError.invalidInputFormat))
|
return
|
}
|
|
// 1. 创建输出 NV12 Buffer
|
let (width, height) = (CVPixelBufferGetWidth(inputBuffer), CVPixelBufferGetHeight(inputBuffer))
|
guard let outputBuffer = createNV12PixelBuffer(width: width, height: height) else {
|
completion(.failure(ConversionError.outputBufferCreationFailed))
|
return
|
}
|
|
// 2. 异步处理(避免阻塞主线程)
|
DispatchQueue.global(qos: .userInitiated).async { [weak self] in
|
guard let self = self else { return }
|
|
do {
|
let result = try self.performConversion(
|
inputBuffer: inputBuffer,
|
outputBuffer: outputBuffer
|
)
|
DispatchQueue.main.async {
|
completion(.success(result))
|
}
|
} catch {
|
DispatchQueue.main.async {
|
completion(.failure(error))
|
}
|
}
|
}
|
}
|
|
// MARK: - Private Methods
|
private func makeComputePipeline(device: MTLDevice) throws -> MTLComputePipelineState {
|
// 1. 获取默认 Metal 库(需在项目中添加 .metal 文件)
|
let library = try device.makeDefaultLibrary(bundle: Bundle(for: Self.self))
|
|
// 2. 加载着色器函数
|
guard let kernelFunction = library.makeFunction(name: "bgraToNV12") else {
|
throw ConversionError.shaderNotFound
|
}
|
|
// 3. 创建计算管线
|
return try device.makeComputePipelineState(function: kernelFunction)
|
}
|
|
private func createNV12PixelBuffer(width: Int, height: Int) -> CVPixelBuffer? {
|
var pixelBuffer: CVPixelBuffer?
|
let status = CVPixelBufferCreate(
|
kCFAllocatorDefault,
|
width,
|
height,
|
kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
|
nil,
|
&pixelBuffer
|
)
|
return status == kCVReturnSuccess ? pixelBuffer : nil
|
}
|
|
private func performConversion(
|
inputBuffer: CVPixelBuffer,
|
outputBuffer: CVPixelBuffer
|
) throws -> CVPixelBuffer {
|
// 1. 锁定缓冲区
|
CVPixelBufferLockBaseAddress(inputBuffer, .readOnly)
|
CVPixelBufferLockBaseAddress(outputBuffer, [])
|
defer {
|
CVPixelBufferUnlockBaseAddress(inputBuffer, .readOnly)
|
CVPixelBufferUnlockBaseAddress(outputBuffer, [])
|
}
|
|
// 2. 创建 Metal 纹理
|
guard let textureCache = textureCache,
|
let inputTexture = createMetalTexture(
|
from: inputBuffer,
|
pixelFormat: .bgra8Unorm,
|
textureCache: textureCache
|
),
|
let yTexture = createMetalTexture(
|
from: outputBuffer,
|
planeIndex: 0,
|
pixelFormat: .r8Unorm,
|
textureCache: textureCache
|
),
|
let uvTexture = createMetalTexture(
|
from: outputBuffer,
|
planeIndex: 1,
|
pixelFormat: .rg8Unorm,
|
textureCache: textureCache
|
) else {
|
throw ConversionError.textureCreationFailed
|
}
|
|
// 3. 执行 Metal 计算
|
guard let commandBuffer = commandQueue?.makeCommandBuffer(),
|
let encoder = commandBuffer.makeComputeCommandEncoder() else {
|
throw ConversionError.metalCommandFailed
|
}
|
|
encoder.setComputePipelineState(computePipeline!)
|
encoder.setTexture(inputTexture, index: 0)
|
encoder.setTexture(yTexture, index: 1)
|
encoder.setTexture(uvTexture, index: 2)
|
|
// 4. 调度计算任务
|
let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
|
let threadgroupCount = MTLSize(
|
width: (inputTexture.width + threadgroupSize.width - 1) / threadgroupSize.width,
|
height: (inputTexture.height + threadgroupSize.height - 1) / threadgroupSize.height,
|
depth: 1
|
)
|
encoder.dispatchThreadgroups(threadgroupCount, threadsPerThreadgroup: threadgroupSize)
|
encoder.endEncoding()
|
|
// 5. 提交并等待完成
|
commandBuffer.commit()
|
commandBuffer.waitUntilCompleted()
|
|
return outputBuffer
|
}
|
|
private func createMetalTexture(
|
from pixelBuffer: CVPixelBuffer,
|
planeIndex: Int = 0,
|
pixelFormat: MTLPixelFormat,
|
textureCache: CVMetalTextureCache
|
) -> MTLTexture? {
|
let width = CVPixelBufferGetWidthOfPlane(pixelBuffer, planeIndex)
|
let height = CVPixelBufferGetHeightOfPlane(pixelBuffer, planeIndex)
|
|
var cvMetalTexture: CVMetalTexture?
|
let status = CVMetalTextureCacheCreateTextureFromImage(
|
nil,
|
textureCache,
|
pixelBuffer,
|
nil,
|
pixelFormat,
|
width,
|
height,
|
planeIndex,
|
&cvMetalTexture
|
)
|
|
guard status == kCVReturnSuccess, let texture = cvMetalTexture else {
|
return nil
|
}
|
|
return CVMetalTextureGetTexture(texture)
|
}
|
|
// MARK: - Error Handling
|
enum ConversionError: Error, LocalizedError {
|
case invalidInputFormat
|
case outputBufferCreationFailed
|
case shaderNotFound
|
case textureCreationFailed
|
case metalCommandFailed
|
|
var errorDescription: String? {
|
switch self {
|
case .invalidInputFormat: return "输入格式必须是 BGRA"
|
case .outputBufferCreationFailed: return "无法创建 NV12 输出缓冲区"
|
case .shaderNotFound: return "找不到 Metal 着色器"
|
case .textureCreationFailed: return "无法创建 Metal 纹理"
|
case .metalCommandFailed: return "Metal 命令执行失败"
|
}
|
}
|
}
|
}
|