Metal shader中核心转化流程如下：
直方图的作用是确保颜色都分配给正在使用的值，最终输出的颜色不代表真实深度，它是当前帧内部相对深度的体现。
首先创建一个Metal文件，建立一个 Compute Shaders
kernel void depthToJET(texture2d<float, access::read> inputTexture [[ texture(0) ]],
texture2d<float, access::write> outputTexture [[ texture(1) ]],
constant JETParams& params [[ buffer(0) ]],
constant float* histogram [[ buffer(1) ]],
constant BGRAPixel *colorTable [[ buffer(2) ]],
uint2 gid [[ thread_position_in_grid ]])
其中 inputTexture 为输入，outputTexture为输出，params 为转换参数，histogram为直方图，colorTable为转换颜色对照表。
struct BGRAPixel {
uchar b;
uchar g;
uchar r;
uchar a;
};
struct JETParams {
int histogramSize;
int binningFactor;
};
kernel void depthToJET(texture2d<float, access::read> inputTexture [[ texture(0) ]],
texture2d<float, access::write> outputTexture [[ texture(1) ]],
constant JETParams& params [[ buffer(0) ]],
constant float* histogram [[ buffer(1) ]],
constant BGRAPixel *colorTable [[ buffer(2) ]],
uint2 gid [[ thread_position_in_grid ]])
{
// Ensure we don't read or write outside of the texture
if ((gid.x >= inputTexture.get_width()) || (gid.y >= inputTexture.get_height())) {
return;
}
// depthDataType is kCVPixelFormatType_DepthFloat16
float depth = inputTexture.read(gid).x;
ushort histIndex = (ushort)(depth * params.binningFactor);
// make sure the value is part of the histogram
if (histIndex >= params.histogramSize) {
return;
}
float colorIndex = histogram[histIndex];
BGRAPixel outputColor = colorTable[(int)colorInd;
outputTexture.write(float4(outputColor.r / 255.0, outputColor.g / 255.0, outputColor.b / 255.0, 1.0), gid);
}
在swift文件中声明和Metal文件一致的数据结构
struct BGRAPixel {
var blue: UInt8 = 0
var green: UInt8 = 0
var red: UInt8 = 0
var alpha: UInt8 = 0
}
struct JETParams {
var histogramSize: Int32 = 1 << 16
var binningFactor: Int32 = 8000
}
颜色对照表
class ColorTable: NSObject {
private var tableBuf: MTLBuffer?
required init (metalDevice: MTLDevice, size: Int) {
self.tableBuf = metalDevice.makeBuffer(length: MemoryLayout<BGRAPixel>.size * size, options: .storageModeShared)
super.init()
self.fillJetTable(size: size)
}
deinit {
}
// second order curve (from HueWave 2010) -- increases saturation at cyan, magenta, and yellow
private func wave(_ pos: Double, phase: Double) -> Double {
let piVal: Double = 4.0 * atan(1.0)
let sinShift: Double = -1.0 / 4.0
// Phase shift sine wave such that sin(2pi * (x+sinShift)) == -1
let xVal: Double = 2.0 * piVal * (pos + sinShift + phase)
let sVal: Double = (sin(xVal) + 1.0) / 2.0
// Normalized sin function
let s2Val: Double = sin(piVal / 2.0 * sVal)
// Flatten top
return s2Val * s2Val
// Symmetrically flattened botton and top
}
private func fillJetTable(size: Int) {
let piVal: Double = 4.0 * atan(1.0)
let rPhase: Double = -1.0 / 4.0
let gPhase: Double = 0.0
let bPhase: Double = +1.0 / 4.0
let table = tableBuf?.contents().bindMemory(to: BGRAPixel.self, capacity: size)
table![0].blue = 0
table![0].green = table![0].blue
table![0].red = table![0].green
// Get pixel info
for idx in 1..<size {
// Get the normalized position
let pos = (Double)(idx) / ((Double)(size) - 1.0)
// Get the current hue value
var red: Double = wave(pos, phase: rPhase)
let green: Double = wave(pos, phase: gPhase)
var blue: Double = wave(pos, phase: bPhase)
// Preserve the jet color table attenuation of red near the start, and blue near the end
// Except instead of making them zero, causing a discontinuity, use an 8th order 1-cos function
if pos < 1.0 / 8.0 {
// Attenuate red channel for 0 < x < 1/8
let xVal: Double = pos * 8.0 * piVal
var attenuation: Double = (cos(xVal) + 1.0) / 2.0
attenuation = 1.0 - pow(attenuation, 0.125)
red *= attenuation
} else if pos > 7.0 / 8.0 {
// Attenuate blue channel for 7/8 < x < 1
let xVal: Double = (1.0 - pos) * 8.0 * piVal
var attenuation: Double = (cos(xVal) + 1.0) / 2.0
attenuation = 1.0 - pow(attenuation, 0.125)
blue *= attenuation
}
table![idx].alpha = (UInt8)(255)
table![idx].red = (UInt8)(255 * red)
table![idx].green = (UInt8)(255 * green)
table![idx].blue = (UInt8)(255 * blue)
}
}
func getColorTable() -> MTLBuffer {
return tableBuf!
}
}
按步骤设置metal
a、在声明中创建Device、CommandQueue
private let metalDevice = MTLCreateSystemDefaultDevice()!
private lazy var commandQueue: MTLCommandQueue? = {
return self.metalDevice.makeCommandQueue()
}()
创建jetParams
private let jetParams = JETParams()
声明metal参数空间及PipelineState
private let jetParamsBuffer: MTLBuffer
private let histogramBuffer: MTLBuffer
private var computePipelineState: MTLComputePipelineState?
声明其他metal所需参数
private(set) var inputFormatDescription: CMFormatDescription?
private(set) var outputFormatDescription: CMFormatDescription?
private var inputTextureFormat: MTLPixelFormat = .invalid
private var outputPixelBufferPool: CVPixelBufferPool!
private let colors = 512
private var textureCache: CVMetalTextureCache!
private var colorBuf: MTLBuffer?
b、
在init()中加载metal文件，创建computePipelineState对象
let defaultLibrary = metalDevice.makeDefaultLibrary()!
let kernelFunction = defaultLibrary.makeFunction(name: "depthToJET")
do {
computePipelineState = try metalDevice.makeComputePipelineState(function: kernelFunction!)
} catch {
fatalError("Unable to create depth converter pipeline state. (\(error))")
}
根据jetParams初始化histBuffer空间
guard let histBuffer = metalDevice.makeBuffer(
length: MemoryLayout<Float>.size * Int(jetParams.histogramSize),
options: .storageModeShared) else {
fatalError("Failed to allocate buffer for histogram")
}
self.histogramBuffer = histBuffer
初始化jetBuffer空间，并写入参数内容
guard let jetBuffer = metalDevice.makeBuffer(length: MemoryLayout<JETParams>.size, options: .storageModeShared) else {
fatalError("Failed to allocate buffer for histogram size")
}
jetBuffer.contents().bindMemory(to: JETParams.self, capacity: 1)
.assign(repeating: self.jetParams, count: 1)
self.jetParamsBuffer = jetBuffer
c、创建allocateOutputBufferPool函数为metal输出分配空间
static private func allocateOutputBufferPool(with formatDescription: CMFormatDescription,
outputRetainedBufferCountHint: Int) -> CVPixelBufferPool? {
let inputDimensions = CMVideoFormatDescriptionGetDimensions(formatDescription)
let outputBufferAttributes: [String: Any] = [
kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
kCVPixelBufferWidthKey as String: Int(inputDimensions.width),
kCVPixelBufferHeightKey as String: Int(inputDimensions.height),
kCVPixelBufferIOSurfacePropertiesKey as String: [:]
]
let poolAttributes = [kCVPixelBufferPoolMinimumBufferCountKey as String: outputRetainedBufferCountHint]
var cvPixelBufferPool: CVPixelBufferPool?
// Create a pixel buffer pool with the same pixel attributes as the input format description
CVPixelBufferPoolCreate(kCFAllocatorDefault, poolAttributes as NSDictionary?, outputBufferAttributes as NSDictionary?, &cvPixelBufferPool)
guard let pixelBufferPool = cvPixelBufferPool else {
assertionFailure("Allocation failure: Could not create pixel buffer pool")
return nil
}
return pixelBufferPool
}
d、创建prepare函数
func prepare(with formatDescription: CMFormatDescription, outputRetainedBufferCountHint: Int) {
...
}
在prepare函数中调用allocateOutputBufferPool
outputPixelBufferPool = DepthToJETConverter.allocateOutputBufferPool(with: formatDescription, outputRetainedBufferCountHint: outputRetainedBufferCountHint)
if outputPixelBufferPool == nil {
return
}
设置inputFormatDescription、outputFormatDescription
var pixelBuffer: CVPixelBuffer?
var pixelBufferFormatDescription: CMFormatDescription?
_ = CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, outputPixelBufferPool!, &pixelBuffer)
if pixelBuffer != nil {
CMVideoFormatDescriptionCreateForImageBuffer(allocator: kCFAllocatorDefault,
imageBuffer: pixelBuffer!,
formatDescriptionOut: &pixelBufferFormatDescription)
}
pixelBuffer = nil
inputFormatDescription = formatDescription
outputFormatDescription = pixelBufferFormatDescription
设置inputTextureFormat
let inputMediaSubType = CMFormatDescriptionGetMediaSubType(formatDescription)
if inputMediaSubType == kCVPixelFormatType_DepthFloat16 {
inputTextureFormat = .r16Float
} else {
assertionFailure("Input format not supported")
}
设置textureCache
var metalTextureCache: CVMetalTextureCache?
if CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, metalDevice, nil, &metalTextureCache) != kCVReturnSuccess {
assertionFailure("Unable to allocate depth converter texture cache")
} else {
textureCache = metalTextureCache
}
设置colorBuf
let colorTable = ColorTable(metalDevice: metalDevice, size: self.colors)
colorBuf = colorTable.getColorTable()
e、在render中进行转换操作
func render(pixelBuffer: CVPixelBuffer) -> CVPixelBuffer? {
... ...
}
建立一个CVPixelBuffer
var newPixelBuffer: CVPixelBuffer?
CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, outputPixelBufferPool!, &newPixelBuffer)
guard let outputPixelBuffer = newPixelBuffer else {
print("Allocation failure: Could not get pixel buffer from pool (\(self.description))")
return nil
}
建立histogramBuffer
let hist = histogramBuffer.contents().bindMemory(to: Float.self, capacity: Int(self.jetParams.histogramSize))
HistogramCalculator.calcHistogram(for: pixelBuffer,
toBuffer: hist,
withSize: self.jetParams.histogramSize,
forColors: Int32(colors),
minDepth: 0.0,
maxDepth: 5.0,
binningFactor: self.jetParams.binningFactor)
建立outputTexture、inputTexture
guard let outputTexture = makeTextureFromCVPixelBuffer(pixelBuffer: outputPixelBuffer, textureFormat: .bgra8Unorm),
let inputTexture = makeTextureFromCVPixelBuffer(pixelBuffer: pixelBuffer, textureFormat: inputTextureFormat) else {
return nil
}
建立并执行commandQueue
guard let commandQueue = commandQueue,
let commandBuffer = commandQueue.makeCommandBuffer(),
let commandEncoder = commandBuffer.makeComputeCommandEncoder() else {
print("Failed to create Metal command queue")
CVMetalTextureCacheFlush(textureCache!, 0)
return nil
}
commandEncoder.label = "Depth to JET"
commandEncoder.setComputePipelineState(computePipelineState!)
commandEncoder.setTexture(inputTexture, index: 0)
commandEncoder.setTexture(outputTexture, index: 1)
commandEncoder.setBuffer(self.jetParamsBuffer, offset: 0, index: 0)
commandEncoder.setBuffer(self.histogramBuffer, offset: 0, index: 1)
commandEncoder.setBuffer(colorBuf, offset: 0, index: 2)
// Set up thread groups as described in https://developer.apple.com/reference/metal/mtlcomputecommandencoder
let width = computePipelineState!.threadExecutionWidth
let height = computePipelineState!.maxTotalThreadsPerThreadgroup / width
let threadsPerThreadgroup = MTLSizeMake(width, height, 1)
let threadgroupsPerGrid = MTLSize(width: (inputTexture.width + width - 1) / width,
height: (inputTexture.height + height - 1) / height,
depth: 1)
commandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
commandEncoder.endEncoding()
commandBuffer.commit()
网友评论