DepthToJETConverter.swift

/*
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
	}

}