用Metal做计算(二)YUV格式相机帧渲染

随着CV技术的不断发展以及其在移动端的落地,越来越多的产品集成了特效相机相关功能,如拍摄小视频时的实时滤镜,人脸贴纸,各类瘦脸、瘦腰、大长腿等“邪术”黑科技......
从移动端技术角度去看,特效相机模块的基本结构是:

  1. 从相机获取图像帧数据
  2. 将图像 数据喂给特定的算法,获取算法输出
  3. 利用算法输出对图像帧进行修改(也可以是叠加其他图像,如3D渲染出的虚拟模型)
  4. 将处理完的图像帧渲染到手机屏幕上

本文仅对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再次进行计算,实现各类滤镜效果。

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 202,529评论 5 475
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,015评论 2 379
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 149,409评论 0 335
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,385评论 1 273
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,387评论 5 364
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,466评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,880评论 3 395
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,528评论 0 256
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,727评论 1 295
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,528评论 2 319
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,602评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,302评论 4 318
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,873评论 3 306
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,890评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,132评论 1 259
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 42,777评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,310评论 2 342

推荐阅读更多精彩内容