Improvements (#11)

Sources/Upscaling/Extensions/AVVideoCodecType+Extensions.swift

@@ -1,6 +1,6 @@
 import AVFoundation
 
-extension AVVideoCodecType {
+public extension AVVideoCodecType {
     var isProRes: Bool {
         switch self {
         case .proRes422, .proRes4444, .proRes422HQ, .proRes422LT, .proRes422Proxy:

Sources/Upscaling/Extensions/CMFormatDescription+Extensions.swift

@@ -59,4 +59,14 @@ extension CMFormatDescription {
         default: nil
         }
     }
+
+    @available(macOS 14.0, iOS 17.0, *) var hasLeftAndRightEye: Bool {
+        let hasLeftEye = (tagCollections ?? []).contains {
+            $0.contains { $0 == .stereoView(.leftEye) }
+        }
+        let hasRightEye = (tagCollections ?? []).contains {
+            $0.contains { $0 == .stereoView(.rightEye) }
+        }
+        return hasLeftEye && hasRightEye
+    }
 }

Sources/Upscaling/Extensions/CVPixelBuffer+Extensions.swift

@@ -0,0 +1,7 @@
+import CoreVideo
+
+extension CVPixelBuffer {
+    var width: Int { CVPixelBufferGetWidth(self) }
+    var height: Int { CVPixelBufferGetHeight(self) }
+    var size: CGSize { .init(width: width, height: height) }
+}

Sources/Upscaling/Extensions/MTLFXSpatialScalerDescriptor+Extensions.swift

@@ -0,0 +1,21 @@
+import MetalFX
+
+extension MTLFXSpatialScalerDescriptor {
+    var inputSize: CGSize {
+        get {
+            CGSize(width: inputWidth, height: inputHeight)
+        } set {
+            inputWidth = Int(newValue.width)
+            inputHeight = Int(newValue.height)
+        }
+    }
+
+    var outputSize: CGSize {
+        get {
+            CGSize(width: outputWidth, height: outputHeight)
+        } set {
+            outputWidth = Int(newValue.width)
+            outputHeight = Int(newValue.height)
+        }
+    }
+}

Sources/Upscaling/Extensions/MTLTexture+Extensions.swift

@@ -0,0 +1,5 @@
+import Metal
+
+extension MTLTexture {
+    var size: CGSize { .init(width: width, height: height) }
+}

Sources/Upscaling/Upscaler.swift

@@ -0,0 +1,201 @@
+import AVFoundation
+import CoreImage
+import CoreVideo
+import Foundation
+import MetalFX
+
+// MARK: - Upscaler
+
+public final class Upscaler {
+    // MARK: Lifecycle
+
+    public init?(inputSize: CGSize, outputSize: CGSize) {
+        let spatialScalerDescriptor = MTLFXSpatialScalerDescriptor()
+        spatialScalerDescriptor.inputSize = inputSize
+        spatialScalerDescriptor.outputSize = outputSize
+        spatialScalerDescriptor.colorTextureFormat = .bgra8Unorm
+        spatialScalerDescriptor.outputTextureFormat = .bgra8Unorm
+        spatialScalerDescriptor.colorProcessingMode = .perceptual
+        let textureDescriptor = MTLTextureDescriptor()
+        textureDescriptor.width = Int(outputSize.width)
+        textureDescriptor.height = Int(outputSize.height)
+        textureDescriptor.pixelFormat = .bgra8Unorm
+        textureDescriptor.storageMode = .private
+        textureDescriptor.usage = [.renderTarget, .shaderRead]
+        guard let device = MTLCreateSystemDefaultDevice(),
+              let commandQueue = device.makeCommandQueue(),
+              let spatialScaler = spatialScalerDescriptor.makeSpatialScaler(device: device),
+              let intermediateOutputTexture = device.makeTexture(descriptor: textureDescriptor) else { return nil }
+        self.commandQueue = commandQueue
+        self.spatialScaler = spatialScaler
+        self.intermediateOutputTexture = intermediateOutputTexture
+        var textureCache: CVMetalTextureCache?
+        CVMetalTextureCacheCreate(nil, nil, device, nil, &textureCache)
+        guard let textureCache else { return nil }
+        self.textureCache = textureCache
+        var pixelBufferPool: CVPixelBufferPool?
+        CVPixelBufferPoolCreate(nil, nil, [
+            kCVPixelBufferPixelFormatTypeKey: kCVPixelFormatType_32BGRA,
+            kCVPixelBufferMetalCompatibilityKey as String: true,
+            kCVPixelBufferWidthKey: outputSize.width,
+            kCVPixelBufferHeightKey: outputSize.height
+        ] as CFDictionary, &pixelBufferPool)
+        guard let pixelBufferPool else { return nil }
+        self.pixelBufferPool = pixelBufferPool
+    }
+
+    // MARK: Public
+
+    @discardableResult public func upscale(
+        _ pixelBuffer: CVPixelBuffer,
+        pixelBufferPool: CVPixelBufferPool? = nil,
+        outputPixelBuffer: CVPixelBuffer? = nil
+    ) async -> CVPixelBuffer {
+        do {
+            let (commandBuffer, outputPixelBuffer) = try upscaleCommandBuffer(
+                pixelBuffer,
+                pixelBufferPool: pixelBufferPool,
+                outputPixelBuffer: outputPixelBuffer
+            )
+            try await withCheckedThrowingContinuation { continuation in
+                commandBuffer.addCompletedHandler { commandBuffer in
+                    if let error = commandBuffer.error {
+                        continuation.resume(throwing: error)
+                    } else {
+                        continuation.resume()
+                    }
+                }
+                commandBuffer.commit()
+            } as Void
+            return outputPixelBuffer
+        } catch {
+            return pixelBuffer
+        }
+    }
+
+    @discardableResult public func upscale(
+        _ pixelBuffer: CVPixelBuffer,
+        pixelBufferPool: CVPixelBufferPool? = nil,
+        outputPixelBuffer: CVPixelBuffer? = nil
+    ) -> CVPixelBuffer {
+        do {
+            let (commandBuffer, outputPixelBuffer) = try upscaleCommandBuffer(
+                pixelBuffer,
+                pixelBufferPool: pixelBufferPool,
+                outputPixelBuffer: outputPixelBuffer
+            )
+            commandBuffer.commit()
+            commandBuffer.waitUntilCompleted()
+            if commandBuffer.error != nil { return pixelBuffer }
+            return outputPixelBuffer
+        } catch {
+            return pixelBuffer
+        }
+    }
+
+    public func upscale(
+        _ pixelBuffer: CVPixelBuffer,
+        pixelBufferPool: CVPixelBufferPool? = nil,
+        outputPixelBuffer: CVPixelBuffer? = nil,
+        completionHandler: @escaping (CVPixelBuffer) -> Void
+    ) {
+        do {
+            let (commandBuffer, outputPixelBuffer) = try upscaleCommandBuffer(
+                pixelBuffer,
+                pixelBufferPool: pixelBufferPool,
+                outputPixelBuffer: outputPixelBuffer
+            )
+            commandBuffer.addCompletedHandler { commandBuffer in
+                if commandBuffer.error != nil {
+                    completionHandler(pixelBuffer)
+                } else {
+                    completionHandler(outputPixelBuffer)
+                }
+            }
+            commandBuffer.commit()
+        } catch {
+            completionHandler(pixelBuffer)
+        }
+    }
+
+    // MARK: Private
+
+    private let commandQueue: MTLCommandQueue
+    private let spatialScaler: MTLFXSpatialScaler
+    private let intermediateOutputTexture: MTLTexture
+    private let textureCache: CVMetalTextureCache
+    private let pixelBufferPool: CVPixelBufferPool
+
+    private func upscaleCommandBuffer(
+        _ pixelBuffer: CVPixelBuffer,
+        pixelBufferPool: CVPixelBufferPool? = nil,
+        outputPixelBuffer: CVPixelBuffer? = nil
+    ) throws -> (MTLCommandBuffer, CVPixelBuffer) {
+        guard CVPixelBufferGetPixelFormatType(pixelBuffer) == kCVPixelFormatType_32BGRA else {
+            throw Error.unsupportedPixelFormat
+        }
+
+        guard let outputPixelBuffer = outputPixelBuffer ?? {
+            var outputPixelBuffer: CVPixelBuffer?
+            CVPixelBufferPoolCreatePixelBuffer(nil, pixelBufferPool ?? self.pixelBufferPool, &outputPixelBuffer)
+            return outputPixelBuffer
+        }() else { throw Error.couldNotCreatePixelBuffer }
+
+        var colorTexture: CVMetalTexture!
+        var status = CVMetalTextureCacheCreateTextureFromImage(
+            nil,
+            textureCache,
+            pixelBuffer,
+            [:] as CFDictionary,
+            .bgra8Unorm,
+            pixelBuffer.width,
+            pixelBuffer.height,
+            0,
+            &colorTexture
+        )
+        guard status == kCVReturnSuccess,
+              let colorTexture = CVMetalTextureGetTexture(colorTexture) else {
+            throw Error.couldNotCreateMetalTexture
+        }
+
+        var upscaledTexture: CVMetalTexture!
+        status = CVMetalTextureCacheCreateTextureFromImage(
+            nil,
+            textureCache,
+            outputPixelBuffer,
+            [:] as CFDictionary,
+            .bgra8Unorm,
+            outputPixelBuffer.width,
+            outputPixelBuffer.height,
+            0,
+            &upscaledTexture
+        )
+        guard status == kCVReturnSuccess,
+              let upscaledTexture = CVMetalTextureGetTexture(upscaledTexture) else {
+            throw Error.couldNotCreateMetalTexture
+        }
+
+        guard let commandBuffer = commandQueue.makeCommandBuffer() else { throw Error.couldNotMakeCommandBuffer }
+
+        spatialScaler.colorTexture = colorTexture
+        spatialScaler.outputTexture = intermediateOutputTexture
+        spatialScaler.encode(commandBuffer: commandBuffer)
+
+        let blitCommandEncoder = commandBuffer.makeBlitCommandEncoder()
+        blitCommandEncoder?.copy(from: intermediateOutputTexture, to: upscaledTexture)
+        blitCommandEncoder?.endEncoding()
+
+        return (commandBuffer, outputPixelBuffer)
+    }
+}
+
+// MARK: Upscaler.Error
+
+extension Upscaler {
+    enum Error: Swift.Error {
+        case unsupportedPixelFormat
+        case couldNotCreatePixelBuffer
+        case couldNotCreateMetalTexture
+        case couldNotMakeCommandBuffer
+    }
+}