Last active
February 4, 2019 20:10
-
-
Save piemonte/089c257d7e82e3cce19c4230f5cc6452 to your computer and use it in GitHub Desktop.
CustomMetalContext
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| // | |
| // 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)") | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment