Newer
Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
/*
See LICENSE folder for this sample’s licensing information.
Abstract:
Converts depth values to JET values.
*/
import CoreMedia
import CoreVideo
import 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!
}
}
class DepthToJETConverter: FilterRenderer {
var description: String = "Depth to JET Converter"
var isPrepared = false
private(set) var inputFormatDescription: CMFormatDescription?
private(set) var outputFormatDescription: CMFormatDescription?
private var inputTextureFormat: MTLPixelFormat = .invalid
private var outputPixelBufferPool: CVPixelBufferPool!
private let metalDevice = MTLCreateSystemDefaultDevice()!
private let jetParams = JETParams()
private let colors = 512
private let jetParamsBuffer: MTLBuffer
private let histogramBuffer: MTLBuffer
private var computePipelineState: MTLComputePipelineState?
private lazy var commandQueue: MTLCommandQueue? = {
return self.metalDevice.makeCommandQueue()
}()
private var textureCache: CVMetalTextureCache!
private var colorBuf: MTLBuffer?
required init() {
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))")
}
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
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
}
static private func allocateOutputBufferPool(with formatDescription: CMFormatDescription, outputRetainedBufferCountHint: Int) -> CVPixelBufferPool? {
let inputDimensions = formatDescription.videoDimensions
let outputPixelBufferAttributes: [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?, outputPixelBufferAttributes as NSDictionary?, &cvPixelBufferPool)
guard let pixelBufferPool = cvPixelBufferPool else {
assertionFailure("Allocation failure: Could not create pixel buffer pool")
return nil
}
return pixelBufferPool
}
func prepare(with formatDescription: CMFormatDescription, outputRetainedBufferCountHint: Int) {
reset()
outputPixelBufferPool = DepthToJETConverter.allocateOutputBufferPool(with: formatDescription,
outputRetainedBufferCountHint: outputRetainedBufferCountHint)
if outputPixelBufferPool == nil {
return
}
var pixelBuffer: CVPixelBuffer?
var pixelBufferFormatDescription: CMFormatDescription?
_ = CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, outputPixelBufferPool!, &pixelBuffer)
if let pixelBuffer = pixelBuffer {
CMFormatDescription.createForVideo(allocator: kCFAllocatorDefault, imageBuffer: pixelBuffer, formatDescriptionOut: &pixelBufferFormatDescription)
}
pixelBuffer = nil
inputFormatDescription = formatDescription
outputFormatDescription = pixelBufferFormatDescription
let inputMediaSubType = formatDescription.mediaSubType
if inputMediaSubType == kCVPixelFormatType_DepthFloat16 {
inputTextureFormat = .r16Float
} else {
assertionFailure("Input format not supported")
}
var metalTextureCache: CVMetalTextureCache?
if CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, metalDevice, nil, &metalTextureCache) != kCVReturnSuccess {
assertionFailure("Unable to allocate depth converter texture cache")
} else {
textureCache = metalTextureCache
}
let colorTable = ColorTable(metalDevice: metalDevice, size: self.colors)
colorBuf = colorTable.getColorTable()
isPrepared = true
}
func reset() {
outputPixelBufferPool = nil
outputFormatDescription = nil
inputFormatDescription = nil
textureCache = nil
isPrepared = false
}
// MARK: - Depth to JET Conversion
func render(pixelBuffer: CVPixelBuffer) -> CVPixelBuffer? {
if !isPrepared {
assertionFailure("Invalid state: Not prepared")
return nil
}
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
}
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)
guard let outputTexture = makeTextureFromCVPixelBuffer(pixelBuffer: outputPixelBuffer, textureFormat: .bgra8Unorm),
let inputTexture = makeTextureFromCVPixelBuffer(pixelBuffer: pixelBuffer, textureFormat: inputTextureFormat) else {
return nil
}
// Set up command queue, buffer, and encoder
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()
return outputPixelBuffer
}
func makeTextureFromCVPixelBuffer(pixelBuffer: CVPixelBuffer, textureFormat: MTLPixelFormat) -> MTLTexture? {
let width = CVPixelBufferGetWidth(pixelBuffer)
let height = CVPixelBufferGetHeight(pixelBuffer)
// Create a Metal texture from the image buffer
var cvTextureOut: CVMetalTexture?
CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, textureCache, pixelBuffer, nil, textureFormat, width, height, 0, &cvTextureOut)
guard let cvTexture = cvTextureOut, let texture = CVMetalTextureGetTexture(cvTexture) else {
print("Depth converter failed to create preview texture")
CVMetalTextureCacheFlush(textureCache, 0)
return nil
}
return texture
}
}