用Metal做计算（二）YUV格式相机帧渲染
随着CV技术的不断发展以及其在移动端的落地，越来越多的产品集成了特效相机相关功能，如拍摄小视频时的实时滤镜，人脸贴纸，各类瘦脸、瘦腰、大长腿等“邪术”黑科技......
从移动端技术角度去看，特效相机模块的基本结构是：
- 从相机获取图像帧数据
- 将图像 数据喂给特定的算法，获取算法输出
- 利用算法输出对图像帧进行修改（也可以是叠加其他图像，如3D渲染出的虚拟模型）
- 将处理完的图像帧渲染到手机屏幕上
本文仅对1，4（获取相机图像帧数据，直接将图像帧渲染到手机屏幕）做讨论。
iOS上，借助AVCaptureVideoPreviewLayer，可以很容易地把相机图像帧展示在手机屏幕上。但特效相机功能开发中，我们基本不能使用AVCaptureVideoPreviewLayer，一是在某些场景下，我们需要修改图像帧数据本；另外我们无法控制PreviewLayer的渲染时序，做不到相机图像帧的渲染与CV算法执行之间的同步。
所以在本文中，我们从相机帧获取数据，并使用MTKView渲染图像帧。
另外，一些CV算法仅接受灰度图数据作为输入，有些则只接受RGBA数据，所以本文在配置 AVCaptureVideoDataOutput 时将输出的数据格式设置为 kCVPixelFormatType_420YpCbCr8BiPlanarFullRange。
本文涉及到的示例代码地址：camera_preview_with_metal
实现思路
- 配置相机，输出图像帧数据（YUV格式），并以此创建两个MTLTexture对象（一张纹理包含Y通道数据，即灰度图数据，一张纹理包含UV通道数据）
- 借助Metal的Compute Pipeline，用上一步获取的两张纹理，合成出RGBA格式的纹理
- 借助Metal的Render Pipeline 将RGBA格式的纹理渲染的手机屏幕上
配置相机，获取输入纹理对象
配置相机部分的代码如下，这里就不做过多的说明了
session = AVCaptureSession()
guard let device = AVCaptureDevice.default(for: .video) else {
return
}
let deviceInput: AVCaptureDeviceInput
do {
deviceInput = try AVCaptureDeviceInput(device: device)
guard session.canAddInput(deviceInput) else {
return
}
session.addInput(deviceInput)
} catch {
return
}
let dataOutput = AVCaptureVideoDataOutput()
dataOutput.videoSettings = [kCVPixelBufferPixelFormatTypeKey: kCVPixelFormatType_420YpCbCr8BiPlanarFullRange] as [String : Any]
guard session.canAddOutput(dataOutput) else {
return
}
dataOutput.setSampleBufferDelegate(self, queue: sampleBufferQueue)
session.addOutput(dataOutput)
session.sessionPreset = .hd1280x720
在相机数据的回调中，我们需要从CMSampleBuffer创建出两个MTLTexture对象。 这里先对 420YpCbCr8BiPlanarFullRange 数据格式做一个简单的介绍。
我们熟知的RGBA格式的图像数据中，一个pixel占4个byte，每个Byte分别包含R、G、B颜色通道和A透明度通道信息，图像帧总的大小为 countOfBytes = image.width * image.height * 4。
而420YpCbCr8BiPlanarFullRange，图像帧的数据被分割为两个plane，index为0的plane大小与图像分辨率一致(即 countOfBytes = image.width * image.height)，每个byte包含对应像素点的亮度信息。index为1的plane包含了图像帧的色差信息，且其宽高分别为图像宽高的一半（即 countOfBytes = image.width/2 * image.height/2），4个亮度像素公用一个色差像素。
下面是从CMSampleBuffer创建MTLTexture的代码
let imagePixel = CMSampleBufferGetImageBuffer(sampleBuffer)!
let yWidth = CVPixelBufferGetWidthOfPlane(imagePixel, 0)
let yHeight = CVPixelBufferGetHeightOfPlane(imagePixel, 0)
let uvWidth = CVPixelBufferGetWidthOfPlane(imagePixel, 1)
let uvHeight = CVPixelBufferGetHeightOfPlane(imagePixel, 1)
CVPixelBufferLockBaseAddress(imagePixel, CVPixelBufferLockFlags(rawValue: 0))
var yTexture: CVMetalTexture?
var uvTexture: CVMetalTexture?
CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, textureCache!, imagePixel, nil, .r8Unorm, yWidth, yHeight, 0, &yTexture)
CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, textureCache!, imagePixel, nil, .rg8Unorm, uvWidth, uvHeight, 1, &uvTexture)
CVPixelBufferUnlockBaseAddress(imagePixel, CVPixelBufferLockFlags(rawValue: 0))
guard yTexture != nil && uvTexture != nil else {
return
}
// Get MTLTexture instance
luminance = CVMetalTextureGetTexture(yTexture!)
chroma = CVMetalTextureGetTexture(uvTexture!)
需要特别说明的是，我们在获取MTLTexture时，并不是通过手动创建 MTLTexture对象，并调用它的
replace(region:mipmapLevel:withBytes:bytesPerRow:)
来实现的，而是使用了下面的方法。
CVReturn CVMetalTextureCacheCreateTextureFromImage(CFAllocatorRef allocator, CVMetalTextureCacheRef textureCache, CVImageBufferRef sourceImage, CFDictionaryRef textureAttributes, MTLPixelFormat pixelFormat, size_t width, size_t height, size_t planeIndex, CVMetalTextureRef _Nullable *textureOut);
这样做的好处是可以降低CPU的使用率，感兴趣的同学可以对比下两个方案的性能。
利用Metal的ComputePipeline实现YUV到RGB的转换
我们需要在初始化阶段，创建一个MTLComputePipelineState对象，代码如下:
let yuv2rgbFunc = library.makeFunction(name: "yuvToRGB")!
yuv2rgbComputePipeline = try! device.makeComputePipelineState(function: yuv2rgbFunc)
CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, device, nil, &textureCache)
其中，yuvToRGB方法的代码如下:
kernel void yuvToRGB(texture2d<float, access::read> yTexture[[texture(0)]],
texture2d<float, access::read> uvTexture[[texture(1)]],
texture2d<float, access::write> outTexture[[texture(2)]],
constant float3x3 *convertMatrix [[buffer(0)]],
uint2 gid [[thread_position_in_grid]]) {
float4 ySample = yTexture.read(gid);
float4 uvSample = uvTexture.read(gid/2);
float3 yuv;
yuv.x = ySample.r;
yuv.yz = uvSample.rg - float2(0.5);
float3x3 matrix = *convertMatrix;
float3 rgb = matrix * yuv;
outTexture.write(float4(rgb, yuv.x), gid);
}
shader方法包含5个参数，前两个为输入的纹理，第三个为输出纹理，第四个是用于做 YUV到RGB转换的3x3矩阵。转换矩阵的定义：
convertMatrix = float3x3(float3(1.164, 1.164, 1.164),
float3(0, -0.231, 2.112),
float3(1.793, -0.533, 0))
最后一个是grid id。
需要注意的是
float4 uvSample = uvTexture.read(gid/2);
如前文所说，4个y通道数据对应一个uv通道数据，所以去uv数据时读的是gid/2的坐标点。
然后是将计算相关的指令编码到CommandBuffer当中：
guard let computeEncoder = commandBuffer.makeComputeCommandEncoder() else {
return
}
// compute pass
computeEncoder.setComputePipelineState(yuv2rgbComputePipeline)
computeEncoder.setTexture(luminance, index: 0)
computeEncoder.setTexture(chroma, index: 1)
computeEncoder.setTexture(rgbTexture, index: 2)
computeEncoder.setBytes(&convertMatrix, length: MemoryLayout<float3x3>.size, index: 0)
let width = rgbTexture.width
let height = rgbTexture.height
let groupSize = 32
let groupCountW = (width + groupSize) / groupSize - 1
let groupCountH = (height + groupSize) / groupSize - 1
computeEncoder.dispatchThreadgroups(MTLSize(width: groupCountW, height: groupCountH, depth: 1),
threadsPerThreadgroup: MTLSize(width: groupSize, height: groupSize, depth: 1))
computeEncoder.endEncoding()
到这里，我们获得了RGBA格式的图像数据，并保存在rgbTexture对象当中，接下来可以用Metal的RenderPipeline绘制一个全屏幕的四边形，并将rgbTexture附着在四边形上，部分代码如下（RenderPipeline初始化相关的代码可参考：camera_preview_with_metal）
guard let renderPassDesc = view.currentRenderPassDescriptor else {
return
}
guard let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDesc) else {
return
}
renderEncoder.setRenderPipelineState(renderPipelineState)
renderEncoder.setVertexBuffer(vertexBuffer, offset: 0, index: 0)
renderEncoder.setFragmentTexture(rgbTexture, index: 0)
renderEncoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
renderEncoder.endEncoding()
总结
以上就是获取相机YUV格式的数据，转换为RGB格式，并渲染到屏幕上所需的全部代码了，我们可以在此基础上做很多扩展，比如在渲染到手机屏幕之前对rgbTexture再次进行计算，实现各类滤镜效果。
