CustomMetalContext

CustomMetalContext.swift

//
//  CustomMetalContext.swift
//
//  The MIT License (MIT)
//
//  Copyright (c) 2016-present patrick piemonte (http://patrickpiemonte.com/)
//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.

import Foundation
import AVFoundation
import SceneKit
import Metal
import CoreVideo
import CoreImage

// MARK: - Types

struct CustomYUVConversion {
    var matrix: float3x3 = float3x3()
    var offset: float3 = float3()
}

// MARK: - CustomMetalContext

internal class CustomMetalContext: NSObject {
    
    // MARK: - properties

    internal var _renderNode: SCNNode
    
    internal var _device: MTLDevice?
    internal var _library: MTLLibrary?
    internal var _commandQueue: MTLCommandQueue?
    internal var _renderPassDescriptor: MTLRenderPassDescriptor?

    internal var _program: SCNProgram?
    internal var _material: SCNMaterial?
    
    internal var _lumaTexture: MTLTexture? {
        didSet {
            if let texture = self._lumaTexture {
                self._lumaProperty = SCNMaterialProperty(contents: texture)
                self._lumaProperty?.wrapS = .clampToBorder
                self._lumaProperty?.wrapT = .clampToBorder
                self._lumaProperty?.mipFilter = .none
            } else {
                self._lumaProperty = nil
            }
        }
    }
    fileprivate var _lumaProperty: SCNMaterialProperty?
    internal var _chromaTexture: MTLTexture? {
        didSet {
            if let texture = self._chromaTexture {
                self._chromaProperty = SCNMaterialProperty(contents: texture)
                self._chromaProperty?.wrapS = .clampToBorder
                self._chromaProperty?.wrapT = .clampToBorder
                self._chromaProperty?.mipFilter = .none
            } else {
                self._chromaTexture = nil
            }
        }
    }
    fileprivate var _chromaProperty: SCNMaterialProperty?
    internal var _offscreenTexture: MTLTexture?
    internal var _textureCache: CVMetalTextureCache?
    
    internal var _bufferWidth: Int
    internal var _bufferHeight: Int
    internal var _bufferFormatType: OSType
    internal var _presentationFrame: CGRect
    
    internal var _offscreenRenderer: SCNRenderer?
    internal var _ciContext: CIContext?
    internal var _pixelBufferPool: CVPixelBufferPool?
    
    // MARK: - object lifecycle
    
    convenience init(view: SCNView) {
        self.init()
        self._device = view.device
        self._presentationFrame = view.bounds
        self.setupContext()
    }
    
    override init() {
        self._renderNode = SCNNode()
        self._renderNode.name = "video plane"
        self._bufferWidth = 0
        self._bufferHeight = 0
        self._bufferFormatType = OSType(kCVPixelFormatType_32BGRA)
        self._presentationFrame = UIScreen.main.bounds
        super.init()
    }
    
    deinit {
        self.destroyContext()
    }
    
}

// MARK: - setup

extension CustomMetalContext: MixedRealityViewRenderContext {
    
    // MARK: - properties
    
    internal var renderNode: SCNNode {
        get {
            return self._renderNode
        }
    }
    
    // MARK: - resource lifecycle
    
    internal func setupContext() {

        // Unfortunately, according to the documentation, Metal shaders will only load from the main application bundle when they use SceneKit.
        // So dumb!!!
        
        // setup context library
        //do {
        //    self._library = try self._device?.makeDefaultLibrary(bundle: Bundle(for: CustomMetalContext.self))
        //} catch {
        //    fatalError("failed to load Metal shaders")
        //}
        
        // setup texture cache
        if let device = self._device {
            let error = CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, device, nil, &self._textureCache)
            if error != kCVReturnSuccess {
                print("failed to create texture cache")
            }
        }
    
        self.clearTextureCache()
        
        // Note: i have no plans on adding BGRA shader support

        // setup program
        if self._program == nil {
            self._program = SCNProgram()
            self._program?.delegate = self
            self._program?.vertexFunctionName = "yuvVertexShader"
            self._program?.fragmentFunctionName = "yuvFragmentShader"
        }
        
        // setup geometry
        if self._material == nil {
            
            // Note: for some reason all devices use a 2x scaled plane
            let scale: CGFloat = 2
            
            let plane = SCNPlane(width: scale, height: scale)
            
            self._material = plane.firstMaterial
            self._material?.writesToDepthBuffer = false
            self._material?.readsFromDepthBuffer = true
            
            if let program = self._program {
                self._material?.program = program
                self._renderNode.geometry = plane
            }
        }
            
        self.setupPixelBufferRendering()
    }
    
    internal func destroyContext() {
        self._device = nil
        self._library = nil
        self._commandQueue = nil
        self._renderPassDescriptor = nil
        
        self._program = nil
        self._material = nil
        
        self._lumaTexture = nil
        self._chromaTexture = nil
        self._offscreenTexture = nil
        self._textureCache = nil
        
        self._offscreenRenderer = nil
        self._ciContext = nil
        self._pixelBufferPool = nil
    }
    
    internal func clearTextureCache() {
        if let textureCache = self._textureCache {
            CVMetalTextureCacheFlush(textureCache, 0)
            self._lumaTexture = nil
            self._chromaTexture = nil
            self._offscreenTexture = nil
        }
    }

    internal func setupPixelBufferRendering() {
        if let device = self._device {
            // setup a renderer for a second pass
            if self._offscreenRenderer == nil {
                self._offscreenRenderer = SCNRenderer(device: device, options: nil)
            }

            // setup command buffer
            if self._commandQueue == nil {
                self._commandQueue = device.makeCommandQueue()
            }
            
            if self._renderPassDescriptor == nil {
                self._renderPassDescriptor = MTLRenderPassDescriptor()
            }
            
            // setup a context for conversion
            if self._ciContext == nil {
                let options : [String : AnyObject] = [kCIContextWorkingColorSpace : CGColorSpaceCreateDeviceRGB(),
                                                      kCIContextUseSoftwareRenderer : NSNumber(booleanLiteral: false)]
                self._ciContext = CIContext(mtlDevice: device, options: options)
            }
        }
    }
}

// MARK: - layout

extension CustomMetalContext {
    
    internal func updateLayout() {
    }
    
    internal func willRender(time: TimeInterval) {
    }
    
    internal func render(withImageBuffer imageBuffer: CVImageBuffer, time: TimeInterval) {
        
        // update textures
        
        self.clearTextureCache()
        
        if let textureCache = self._textureCache {
            let isPlanar = CVPixelBufferIsPlanar(imageBuffer)
            let width = isPlanar ? CVPixelBufferGetWidthOfPlane(imageBuffer, 0) : CVPixelBufferGetWidth(imageBuffer)
            let height = isPlanar ? CVPixelBufferGetHeightOfPlane(imageBuffer, 0) : CVPixelBufferGetHeight(imageBuffer)
            
            if self._bufferWidth != width || self._bufferHeight != height {
                self._bufferWidth = width
                self._bufferHeight = height
                self.updateLayout()
            }
            
            self._bufferFormatType = CVPixelBufferGetPixelFormatType(imageBuffer)
            
            switch self._bufferFormatType {
            case kCVPixelFormatType_32BGRA:
                let _ = self.texture(withImageBuffer: imageBuffer, textureCache: textureCache, planeIndex: 0, pixelFormat: .bgra8Unorm)
                debugPrint("unsupported pixel format type")
                break
            case kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange:
                fallthrough
            case kCVPixelFormatType_420YpCbCr8BiPlanarFullRange:
                if let textureY = self.texture(withImageBuffer: imageBuffer, textureCache: textureCache, planeIndex: 0, pixelFormat: .r8Unorm),
                    let textureCbCr = self.texture(withImageBuffer: imageBuffer, textureCache: textureCache, planeIndex: 1, pixelFormat: .rg8Unorm) {
                    self._lumaTexture = textureY
                    self._chromaTexture = textureCbCr
                }
                break
            default:
                debugPrint("unsupported pixel format type")
                break
            }
        }
        
        // setup the offscreen texture now that the size was determined
        if self._offscreenTexture == nil && self._bufferWidth > 0 && self._bufferHeight > 0 {
            let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba8Unorm, width: self._bufferWidth, height: self._bufferHeight, mipmapped: false)
            textureDescriptor.usage = [.shaderRead, .renderTarget]
            self._offscreenTexture = self._device?.makeTexture(descriptor: textureDescriptor)
        }
        
        // bind textures and uniforms
        
        // bind textures to the material for sampling
        self._material?.setValue(self._lumaProperty, forKey: "textureY")
        self._material?.setValue(self._chromaProperty, forKey: "textureCbCr")
        
        // setup color conversion uniform buffer
        // BT.709, which is the standard for HDTV.
        var yuvConversion: CustomYUVConversion = CustomYUVConversion()
        yuvConversion.matrix = float3x3([float3( 1.164,  1.164, 1.164),
                                         float3(   0.0, -0.213, 2.112),
                                         float3( 1.793, -0.533,   0.0)])
        yuvConversion.offset = float3( -(16.0/255.0), -0.5, -0.5)
        let yuvUniformData = Data(bytes: &yuvConversion, count: MemoryLayout<CustomYUVConversion>.stride)
        self._material?.setValue(yuvUniformData, forKey: "yuvConversion")
    }
    
    internal func didRender(scene: SCNScene, pointOfView: SCNNode, time: TimeInterval) {
        if let commandBuffer = self._commandQueue?.makeCommandBuffer(),
            let renderPassDescriptor = self._renderPassDescriptor,
            let offscreenTexture = self._offscreenTexture {

            let viewport = CGRect(x: 0, y: 0, width: self._bufferWidth, height: self._bufferHeight)

            renderPassDescriptor.colorAttachments[0].texture = offscreenTexture
            renderPassDescriptor.colorAttachments[0].loadAction = .clear
            renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColorMake(0, 0, 0, 1.0);
            renderPassDescriptor.colorAttachments[0].storeAction = .store

            self._offscreenRenderer?.scene = scene
            self._offscreenRenderer?.pointOfView = pointOfView
            self._offscreenRenderer?.render(atTime: time, viewport: viewport, commandBuffer: commandBuffer, passDescriptor: renderPassDescriptor)
            
            commandBuffer.commit()
        }
    }
    
}

// MARK: - post

fileprivate let CustomMetalContextMinBufferCount = 3

extension CustomMetalContext {
    
    internal func renderedPixelBuffer() -> CVPixelBuffer? {
        
        // allocate a pool, if necessary

        if self._pixelBufferPool == nil {
            let poolAttributes: [String:AnyObject] = [String(kCVPixelBufferPoolMinimumBufferCountKey): NSNumber(integerLiteral:CustomMetalContextMinBufferCount)]
            
            let pixelBufferAttributes: [String:AnyObject] = [String(kCVPixelBufferPixelFormatTypeKey) : NSNumber(integerLiteral: Int(self._bufferFormatType)),
                                                             String(kCVPixelBufferWidthKey) : NSNumber(value: self._bufferWidth),
                                                             String(kCVPixelBufferHeightKey) : NSNumber(value: self._bufferHeight),
                                                             String(kCVPixelBufferMetalCompatibilityKey) : NSNumber(booleanLiteral: true),
                                                             String(kCVPixelBufferIOSurfacePropertiesKey) : [:] as AnyObject ]

            var pixelBufferPool: CVPixelBufferPool? = nil
            let result = CVPixelBufferPoolCreate(kCFAllocatorDefault, poolAttributes as CFDictionary, pixelBufferAttributes as CFDictionary, &pixelBufferPool)
            if result == kCVReturnSuccess {
                self._pixelBufferPool = pixelBufferPool
            }
        }
        
        // allocate a pixel buffer and render into it
        
        var pixelBuffer: CVPixelBuffer? = nil
        if let pixelBufferPool = self._pixelBufferPool {
            let result = CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, pixelBufferPool, &pixelBuffer)
            
            if result == kCVReturnSuccess {
                if let context = self._ciContext,
                    let pixelBuffer = pixelBuffer,
                    let offscreenTexture = self._offscreenTexture {
                    CVPixelBufferLockBaseAddress(pixelBuffer, CVPixelBufferLockFlags.readOnly)
                    let ciImage = CIImage(mtlTexture: offscreenTexture, options: nil)
                    // Metal's origin is bottom left, so we have to flip and rotate
                    if let orientedImage = ciImage?.oriented(forExifOrientation: 4) {
                        context.render(orientedImage, to: pixelBuffer)
                        // Note: it's possible to apply CoreImage filters here
                    }
                    CVPixelBufferUnlockBaseAddress(pixelBuffer, CVPixelBufferLockFlags.readOnly)
                    return pixelBuffer
                }
            }
            
        }
        
        return nil
        
    }
    
}

// MARK: - private
    
extension CustomMetalContext {
    
    fileprivate func texture(withImageBuffer imageBuffer: CVImageBuffer, textureCache: CVMetalTextureCache, planeIndex: Int = 0, pixelFormat: MTLPixelFormat = .bgra8Unorm) -> MTLTexture? {
        let isPlanar = CVPixelBufferIsPlanar(imageBuffer)
        let width = isPlanar ? CVPixelBufferGetWidthOfPlane(imageBuffer, planeIndex) : CVPixelBufferGetWidth(imageBuffer)
        let height = isPlanar ? CVPixelBufferGetHeightOfPlane(imageBuffer, planeIndex) : CVPixelBufferGetHeight(imageBuffer)
        
        var textureImage: CVMetalTexture? = nil
        let error = CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault,
                                                              textureCache,
                                                              imageBuffer,
                                                              nil,
                                                              pixelFormat,
                                                              width,
                                                              height,
                                                              planeIndex,
                                                              &textureImage)
        if error != kCVReturnSuccess {
            debugPrint("failed to create texture cache")
            return nil
        }
        if let texture = textureImage {
           return CVMetalTextureGetTexture(texture)
        }
        return nil
    }
    
}

// MARK: - SCNProgramDelegate

extension CustomMetalContext: SCNProgramDelegate {
    
    internal func program(_ program: SCNProgram, handleError error: Error) {
        debugPrint("program \(program) error \(error)")
    }
    
}