近两年可以说是视频行业的风口,前有直播,现在有短视频。对短视频应用来说滤镜几乎是必不可少的功能,在 iOS 上大家通常都选择基于 GPUImage 这个库来实现。不过最近我在学习苹果在 WWDC 2014 推出的 Metal 框架,因此这篇博客我将通过 “山寨” Shadertoy 上一个 WebGL 程序来简单介绍如何使用 Metal 框架处理视频。

本文 demo 使用 Swift 4 和 Xcode 9 beta 4 开发,并且目前 Metal 应用只能在真机上运行调试。

初始设置

打开 Xcode 新建一个 Single View App 项目并命名为 VideoProcessWithMetal

由于这个 demo 项目将用到 MetalKitAVFoundationCoreVideo 这几个框架,因此先在 ViewController.swift 顶部导入这几个框架:

import AVFoundation
import MetalKit
import CoreVideo

打开 Main.StoryboardView Controllerview 类型设为 MTKView

按住 Ctrl 拖动连线到 ViewController.swift 中并命名为 mtkView,然后在 ViewController.swift 中添加一个 property:

let device = MTLCreateSystemDefaultDevice()!

这里,device 表示手机的 GPU,然后增加一个 setupView 方法:

private func setupView() {
    mtkView.device = device  // 绑定 GPU
    mtkView.framebufferOnly = false // 读写 currentDrawable.texture 内存
    mtkView.delegate = self
    mtkView.colorPixelFormat = .bgra8Unorm
    mtkView.contentMode = .scaleAspectFit
    mtkView.isPaused = true
}

这里设置 colorPixelFormat.bgra8Unorm(目前 iOS 仅支持 .bgra8Unorm.bgra8Unorm_srgb)。由于接下来我将手动处理摄像头得到的数据并显示到屏幕上,因此将它的 isPaused 置为 true 来防止 MTKView 自动调用 draw 方法。

接着,在 ViewController.swift 里添加下面的 property:

var computePipelineState: MTLComputePipelineState?

并增加一个 initializeComputePipeline 方法:

private func initializeComputePipeline() {
    let library = device.makeDefaultLibrary()
    let shader = library?.makeFunction(name: "separateRGB")
    computePipelineState = try! device.makeComputePipelineState(function: shader!)
}

这里,获取一个在 app 编译时同时编译好的 Metal 代码库方法 separateRGB 并用它创建一个 compute pipelinecompute pipeline 主要用来维护渲染计算过程中的各种状态。具体的 separateRGB 代码将在后面添加。

同时,在 ViewController.swift 里添加下面这个 property:

lazy var commandQueue: MTLCommandQueue = {
    return self.device.makeCommandQueue()!
}()

所有的渲染计算指令都将通过这个 command queue 发送到 GPU,这可以保证渲染指令的先后顺序。

然后,在 ViewController.swift 里添加下面两个 property:

var sourceTexture: MTLTexture?
var textureCache: CVMetalTextureCache?

并增加一个 createTextureCache 方法:

private func createTextureCache() {
    CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, device, nil, &textureCache)
}

这个方法将初始化一个 CVMetalTextCache 对象,它的用法类似 CVOpenGLESTextureCacheRef,在下面的代码里将用它把摄像头捕获的图像数据转成 MTLTexture 对象并传给 sourceTexture

最后,添加摄像头相关的属性:

let captureSession = AVCaptureSession()
let sampleBufferCallbackQueue = DispatchQueue(label: "video.process.metal")
var writer: AVAssetWriter!
var writerInput: AVAssetWriterInput!
var adaptor: AVAssetWriterInputPixelBufferAdaptor!
var isRecording: Bool = false
    
var beginTime = CACurrentMediaTime()  // app 开始运行的时间
var lastSampleTime: CMTime = kCMTimeZero // 上一个摄像头采样的时间
var processedPixelBuffer: CVPixelBuffer?

其中,processedPixelBuffer 是由 adaptor 创建的 CVPixelBufferPool 对象根据 adaptorsourcePixelBufferAttributes 创建出来的一个临时对象,用于存储处理后的图像数据,具体使用将在下面的代码里说明。

接着,增加一个 configCaptureSession 方法:

private func configCaptureSession() {
    captureSession.beginConfiguration()
        
    // preset
    captureSession.sessionPreset = AVCaptureSession.Preset.hd1280x720
        
    // video input
    guard let camera = AVCaptureDevice.default(for: .video) else {
        return
    }
        
    do {
        let videoInput = try AVCaptureDeviceInput(device: camera)
        if captureSession.canAddInput(videoInput) {
            captureSession.addInput(videoInput)
        }
    } catch {
        return
    }
        
    // video output
    let videoOutput = AVCaptureVideoDataOutput()
    videoOutput.videoSettings = [kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA]
    videoOutput.setSampleBufferDelegate(self, queue: sampleBufferCallbackQueue)
    if captureSession.canAddOutput(videoOutput) {
        captureSession.addOutput(videoOutput)
    }
        
    if let connection = videoOutput.connection(with: .video) {
        if connection.isVideoOrientationSupported {
            connection.videoOrientation = AVCaptureVideoOrientation(rawValue: UIApplication.shared.statusBarOrientation.rawValue)!
        }
    }
        
    // writer
    do {
        writer = try! AVAssetWriter(outputURL: URL.randomUrl(), fileType: .mov)
            
        let videoCompressionProperties = [
            AVVideoAverageBitRateKey: 6000000
        ]
            
        let videoSettings: [String : Any] = [
            AVVideoCodecKey: AVVideoCodecH264,
            AVVideoWidthKey: 720,
            AVVideoHeightKey: 1280,
            AVVideoCompressionPropertiesKey: videoCompressionProperties
        ]
            
        writerInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoSettings)
        writerInput.expectsMediaDataInRealTime = true
        writerInput.transform = .identity
        let sourcePixelBufferAttributes = [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
            kCVPixelBufferWidthKey as String: 720,
            kCVPixelBufferHeightKey as String: 1280
        ]
        adaptor = AVAssetWriterInputPixelBufferAdaptor(assetWriterInput: writerInput, sourcePixelBufferAttributes: sourcePixelBufferAttributes)
            
        if writer.canAdd(writerInput) {
            writer.add(writerInput)
        }
    }
        
    captureSession.commitConfiguration()
}

extension URL {
    static public func randomUrl() -> URL {
        let path = NSTemporaryDirectory() + "/" + String.random(length: 5) + ".mov"
        return URL(fileURLWithPath: path)
    }
    
    public func saveToAlbum() {
        UISaveVideoAtPathToSavedPhotosAlbum(self.path, nil, nil, nil)
    }
}

extension String {
    static func random(length: Int = 20) -> String {
        let base = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
        var randomString: String = ""
        
        for _ in 0..<length {
            let randomValue = arc4random_uniform(UInt32(base.characters.count))
            randomString += "\(base[base.index(base.startIndex, offsetBy: Int(randomValue))])"
        }
        return randomString
    }
}

这里,为了简化 demo 只添加了视频相关的输入输出,并使用 1280x720 的分辨率。可以注意到 videoOutput.videoSettingssourcePixelBufferAttributes 里都有这么一个键值对:

kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA

这是因为在接下来的 Metal 代码里将使用 BGRA 颜色空间来处理。同时,在 videoOutput.videoSettings 里有这个设置:

AVVideoCodecKey: AVVideoCodecH264

这在最后保存到文件里时将使用 adaptorBGRA 自动转换为 yuv 并使用 H.264 编码格式。

最后,在 viewDidLoad 方法中调用这些初始化方法:

override func viewDidLoad() {
    super.viewDidLoad()

    setupView()
    initializeComputePipeline()
    createTextureCache()
    configCaptureSession()
}

摄像头捕获图像的处理

在上面的代码中我们通过

videoOutput.setSampleBufferDelegate(self, queue: sampleBufferCallbackQueue)

videoOutputdelegate 设置为 ViewController,因此,为了处理数据,需要实现这个代理中的方法:

extension ViewController: AVCaptureVideoDataOutputSampleBufferDelegate {
    
    func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
        guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
            return
        }
        
        var cvmTexture: CVMetalTexture?
        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, textureCache!, pixelBuffer, nil, mtkView.colorPixelFormat, width, height, 0, &cvmTexture)
        if let cvmTexture = cvmTexture, let texture = CVMetalTextureGetTexture(cvmTexture) {
            sourceTexture = texture
        }
        
        lastSampleTime = CMSampleBufferGetPresentationTimeStamp(sampleBuffer)
        
        DispatchQueue.main.sync {
            mtkView.draw()
        }
    }
}

代码很直白,通过 sampleBuffer 获取它里面的 CVPixelBuffer 对象,并通过刚才创建的 textureCache 创建这个 CVPixelBuffer 对象对应的 MTLTexture 对象。需要注意的是,由于 CVMetalTexture 的创建是一件资源消耗较大的操作,因此 textureCache 会在内部创建一个类似缓冲队列的结构来存储 CVMetalTexutre 对象并循环使用。此外,CVPixelBuffer 对象和对应的 CVMetalTexture 对象在底层共享内存。

MTLTexture 对象处理

同理,我们还要实现 MTKView 的代码方法:

extension ViewController: MTKViewDelegate {
    func draw(in view: MTKView) {
        guard let currentDrawable = mtkView.currentDrawable, let texture = sourceTexture else {
            return
        }
        
        let commandBuffer = commandQueue.makeCommandBuffer()
        let computeCommandEncoder = commandBuffer?.makeComputeCommandEncoder()
        computeCommandEncoder?.setComputePipelineState(computePipelineState!)
        computeCommandEncoder?.setTexture(texture, index: 0)
        computeCommandEncoder?.setTexture(currentDrawable.texture, index: 1)
        
        var diff = Float(CACurrentMediaTime() - beginTime)
        computeCommandEncoder?.setBytes(&diff, length: MemoryLayout<Float>.size, index: 0)
        computeCommandEncoder?.dispatchThreadgroups(texture.threadGroups(pipeline: computePipelineState!), threadsPerThreadgroup: texture.threadGroupCount(pipeline: computePipelineState!))
        computeCommandEncoder?.endEncoding()
        
        commandBuffer?.present(currentDrawable)
        commandBuffer?.commit()
    }
    
    func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
        
    }
}

extension MTLTexture {
    func threadGroupCount(pipeline: MTLComputePipelineState) -> MTLSize {
        return MTLSizeMake(pipeline.threadExecutionWidth,
                           pipeline.maxTotalThreadsPerThreadgroup / pipeline.threadExecutionWidth,
                           1)
    }
    
    func threadGroups(pipeline: MTLComputePipelineState) -> MTLSize {
        let groupCount = threadGroupCount(pipeline: pipeline)
        return MTLSizeMake((self.width + groupCount.width - 1) / groupCount.width, (self.height + groupCount.height - 1) / groupCount.height, 1)
    }
}

这里,通过上面创建的 commandQueue 对象创建了一个 MTLCommandBuffer 对象 commandBuffer,它相当于一个指令缓冲区,具体的指令将被编码进入这个对象中并通过 commandQueue 发送到 GPU 执行。至于具体的指令由 MTLCommandEncoder 对象表示,这里我将使用它的一个子类 MTLComputeCommandEncoder,通过设置它的 pipelineState 它可以知道在执行时将使用哪个渲染函数。

computeCommandEncoder?.setTexture(texture, index: 0)
computeCommandEncoder?.setTexture(currentDrawable.texture, index: 1)
        
var diff = Float(CACurrentMediaTime() - beginTime)
computeCommandEncoder?.setBytes(&diff, length: MemoryLayout<Float>.size, index: 0)

这几个方法的调用将设置渲染函数所需要的几个参数。那么渲染函数具体是什么样呢?我们新建一个 Metal 文件命名为 SeparateRGB.metal,打开这个文件并添加如下代码:

#include <metal_stdlib>
using namespace metal;

float hash(float);
float noise(float3);

float hash(float n) {
    return fract(sin(n) * 43758.5453);
}

float noise(float3 x) {
    float3 p = floor(x);
    float3 f = fract(x);
    
    f = f * f * (3.0 - 2.0 * f);
    
    float n = p.x + p.y * 57.0 + 113.0 * p.z;
    
    float res = mix(mix(mix(hash(n + 0.0), hash(n + 1.0), f.x),
                        mix(hash(n + 57.0), hash(n + 58.0), f.x), f.y),
                    mix(mix(hash(n + 113.0), hash(n + 114.0), f.x),
                        mix(hash(n + 170.0), hash(n + 171.0), f.x), f.y), f.z);
    return res;
}

kernel void separateRGB(texture2d<float, access::read> inTexture [[ texture(0) ]],
                        texture2d<float, access::write> outTexture [[ texture(1) ]],
                        device const float *time [[ buffer(0) ]],
                        uint2 gid [[ thread_position_in_grid ]]) {
    float2 uv = float2(gid);
    uv.x /= inTexture.get_width();
    uv.y /= inTexture.get_height();
    
    float iGlobalTime = *time;
    float blurx = noise(float3(iGlobalTime * 10.0, 0.0, 0.0)) * 2.0 - 1.0;
    float offsetx = blurx * 0.025;
    
    float blury = noise(float3(iGlobalTime * 10.0, 1.0, 0.0)) * 2.0 - 1.0;
    float offsety = blury * 0.01;
    
    float2 ruv = uv + float2(offsetx, offsety);
    float2 guv = uv + float2(-offsetx, -offsety);
    float2 buv = uv + float2(0.00, 0.0);
    
    float r = inTexture.read(uint2(ruv * float2(inTexture.get_width(), inTexture.get_height()))).r;
    float g = inTexture.read(uint2(guv * float2(inTexture.get_width(), inTexture.get_height()))).g;
    float b = inTexture.read(uint2(buv * float2(inTexture.get_width(), inTexture.get_height()))).b;
    
    outTexture.write(float4(r, g, b, 1.0), gid);
}

这段代码是我从这里获取,由于它是用 WebGL 所写,并不能直接使用在 Metal 代码中,因此我用 Metal shading language 将这段代码重写为上面的 Metal 代码。这段代码里最重要的就是 separateRGB 这个方法了,它接受 4 个参数,前两个表示输入和输出的 MTLTexture 对象,第三个表示一个时间参数,最后一个是内置的参数,在运行的时候系统将自动设置这个参数的值,因此上面的代码里我们只设置了前三个参数的值。Metal shading language 基于 c++ 14,具体的语法可以参考官方文档

运行显示

大部分工作已经做好,为了简化 demo,我们在 viewDidAppear 里打开摄像头:

override func viewDidAppear(_ animated: Bool) {
    super.viewDidAppear(animated)
        
    sampleBufferCallbackQueue.async {
        self.captureSession.startRunning()
    }
}

点击运行,如果不出意外的话,it worked!

录制到文件

然而上面的工作只是将处理后的图像显示出来,大部分情况下我们还是需要将渲染的结果保存到视频文件里,在 draw(in view: MTKView) 方法里添加如下代码:

...
computeCommandEncoder?.endEncoding()
        
if self.isRecording && adaptor.assetWriterInput.isReadyForMoreMediaData {
    if processedPixelBuffer == nil {
        CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, adaptor.pixelBufferPool!, &processedPixelBuffer)
    }
            
    guard let p = processedPixelBuffer else {
        return
    }
            
    CVPixelBufferLockBaseAddress(p, CVPixelBufferLockFlags(rawValue: 0))
            
    let outputTexture = currentDrawable.texture
    let region = MTLRegionMake2D(0, 0, 720, 1280)
            
    let buffer = CVPixelBufferGetBaseAddress(p)
    let bytesPerRow = 4 * region.size.width
    outputTexture.getBytes(buffer!, bytesPerRow: bytesPerRow, from: region, mipmapLevel: 0)
            
    adaptor.append(p, withPresentationTime: lastSampleTime)
            
    CVPixelBufferUnlockBaseAddress(p, CVPixelBufferLockFlags(rawValue: 0))
}

这里通过 MTKView.currentDrawable 获取到渲染后的 MTLTexture 对象,并将它的字节数据写入到上面创建的 processedPixelBuffer 对象里,最后调用 adaptorappend(_ :, withPresentationTime:) 方法将这个 CVPixelBuffer 对象存储到视频文件中。这里,图像的数据格式将由 adaptor 根据设置从 BGRA 自动转换为 yuv。我的 demo 跑在一台 iPhone 7 上面,容易知道,currentDrawable 的尺寸大小为 750x1334,那么为什么这里 region 尺寸设置为 720, 1080 呢?由于我也是刚刚学习图形编程,在查了大量资料之后依然没有找到准确答案,因此我只能猜测:GPU 在将渲染后的 MTLTexture 显示到屏幕上之前会根据 drawableSize 这个值有一个 ”resampling“ 的过程。如果读者里有熟悉这方面的还望不吝赐教:)。

接下来,打开 Main.Storyboard 并添加一个 UIButton 到视图上:

并连线到 ViewController.swift 里创建一个点击方法:

@IBAction func captureAction(_ sender: Any) {
    let button = sender as! UIButton
        
    if button.title(for: .normal) == "Start" {
        sampleBufferCallbackQueue.sync {
            if self.writer.startWriting() {
                self.writer.startSession(atSourceTime: lastSampleTime)
            }
                
            self.isRecording = true
        }
            
        button.setTitle("Stop", for: .normal)
    } else {
        sampleBufferCallbackQueue.sync {
            self.isRecording = false
            self.writer.finishWriting {
                self.writer.outputURL.saveToAlbum()
            }
        }
            
        button.setTitle("Start", for: .normal)
    }
}

代码很简单,点击按钮后开始录制,注意到这里:

self.writer.startSession(atSourceTime: lastSampleTime)

并没有使用 kCMTimeZero 作为这个方法的参数而是使用 lastSampleTime,这可以简化时间转换方面的操作。

点击运行,就可以将图像内容录制下来了!

总结

至此,我通过简单的 demo 展示了 Metal 框架的使用,由于这只是一个 demo,因此代码里有许多地方并没有进行可能的错误处理,如果你想要将它使用到实际代码里还需要补充上这些错误处理。如果读者里发现我有使用错误的地方,还望告知,谢谢:)

最终的 demo 代码可以在这里获取。

参考