ImageProcessor.swift

import Metal
import MetalKit

class ImageProcessor: NSObject {
    
    //ハードウェアとしてのGPUを抽象化したプロトコル
    lazy var device: MTLDevice! = MTLCreateSystemDefaultDevice()
    //コマンドバッファの実行順を管理するキュー
    var commandQueue: MTLCommandQueue!
    //シェーダプログラムの関数名
    var funcName: String = "duotone"

    var outTexture: MTLTexture!
    
    var pipelineState: MTLComputePipelineState!
    let threadGroupCount = MTLSizeMake(16, 16, 1)
    
    /*===================
     初回処理
     ===================*/
    public func Setup() {
        let defaultLibrary = device.makeDefaultLibrary()!
        if let target = defaultLibrary.makeFunction(name: funcName) {
            commandQueue = device.makeCommandQueue()
            do {
                pipelineState = try device.makeComputePipelineState(function: target)
            } catch {
                fatalError("Impossible to setup MTL")
            }
        }
    }
    
    /*===================
     シェーダーの実行
     ===================*/
    public func Run(_ image:UIImage) -> UIImage{
        //GPUで実行されるコマンドを格納するコンテナ
        let buffer = commandQueue.makeCommandBuffer()
        //コマンドを作成し、コマンドバッファに追加する（エンコード）
        let encoder = buffer?.makeComputeCommandEncoder()
        encoder?.setComputePipelineState(pipelineState)

        //出力用テクスチャ作成
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: MTLPixelFormat.rgba8Unorm, width:1920, height: 1080, mipmapped: false)
        textureDescriptor.usage = [.shaderRead, .shaderWrite]
        outTexture = self.device.makeTexture(descriptor: textureDescriptor)
        
        //コマンドバッファにデータを追加
        //textures
        encoder?.setTexture(outTexture, index: 0)
        encoder?.setTexture(mtlTexture(from: image), index: 1)
        encoder?.dispatchThreadgroups( MTLSizeMake(
            Int(ceil(image.size.width / CGFloat(self.threadGroupCount.width))),
            Int(ceil(image.size.height / CGFloat(self.threadGroupCount.height))),
            1), threadsPerThreadgroup: threadGroupCount)
        encoder?.endEncoding()
        //コマンドを実行
        buffer?.commit()
        //完了まで待つ
        buffer?.waitUntilCompleted()
        //出力を返す
        return self.image(from: self.outTexture)
    }
    
    /*=========================
     UIImage -> MTLTexture
     =========================*/
    func mtlTexture(from image: UIImage) -> MTLTexture {
        //CGImage変換時に向きがおかしくならないように
        UIGraphicsBeginImageContext(image.size);
        image.draw(in: CGRect(x:0, y:0, width:image.size.width, height:image.size.height))
        let orientationImage = UIGraphicsGetImageFromCurrentImageContext();
        UIGraphicsEndImageContext();
        //CGImageに変換
        guard let cgImage = orientationImage?.cgImage else {
            fatalError("Can't open image \(image)")
        }
        //MTKTextureLoaderを使用してCGImageをMTLTextureに変換
        let textureLoader = MTKTextureLoader(device: self.device)
        do {
            let tex = try textureLoader.newTexture(cgImage: cgImage, options: nil)
            let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: tex.pixelFormat, width: tex.width, height: tex.height, mipmapped: false)
            textureDescriptor.usage = [.shaderRead, .shaderWrite]
            return tex
        }
        catch {
            fatalError("Can't load texture")
        }
    }
    
    /*=========================
     MTLTexture -> UIImage
     =========================*/
    func image(from mtlTexture: MTLTexture) -> UIImage {
        //画像サイズ
        let w = mtlTexture.width
        let h = mtlTexture.height
        let bytesPerPixel: Int = 4
        let imageByteCount = w * h * bytesPerPixel
        let bytesPerRow = w * bytesPerPixel
        var src = [UInt8](repeating: 0, count: Int(imageByteCount))
        //CGImageに変換
        let region = MTLRegionMake2D(0, 0, w, h)
        mtlTexture.getBytes(&src, bytesPerRow: bytesPerRow, from: region, mipmapLevel: 0)
        let bitmapInfo = CGBitmapInfo(rawValue: (CGBitmapInfo.byteOrder32Big.rawValue | CGImageAlphaInfo.premultipliedLast.rawValue))
        let colorSpace = CGColorSpaceCreateDeviceRGB()
        let bitsPerComponent = 8
        let context = CGContext(data: &src,
                                width: w,
                                height: h,
                                bitsPerComponent: bitsPerComponent,
                                bytesPerRow: bytesPerRow,
                                space: colorSpace,
                                bitmapInfo: bitmapInfo.rawValue)
        let cgImage = context?.makeImage()
        //UIImageに変換して返す
        let image = UIImage(cgImage: cgImage!)
        return image
    }
    
}

// Singleton✌️
extension ImageProcessor {
    class var Shared : ImageProcessor {
        struct Static { static let instance : ImageProcessor = ImageProcessor() }
        return Static.instance
    }
}