Skip to content

Instantly share code, notes, and snippets.

@piemonte
Last active February 4, 2019 20:10
Show Gist options
  • Save piemonte/089c257d7e82e3cce19c4230f5cc6452 to your computer and use it in GitHub Desktop.
Save piemonte/089c257d7e82e3cce19c4230f5cc6452 to your computer and use it in GitHub Desktop.
CustomMetalContext
//
// 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