Swift Metal save bgra8Unorm texture to PNG file

Question

I have a kernel that outputs a texture, and it is a valid MTLTexture object. I want to save it to a png file in the working directory of my project. How should this be done?

The texture format is .bgra8Unorm, and the target output format is PNG. The texture is stored in a MTLTexture object.

EDIT: I am on macOS XCode.

Answer 1

If your app is using Metal on macOS, the first thing you need to do is ensure that your texture data can be read by the CPU. If the texture that's being written by the kernel is in .private storage mode, that means you'll need to blit (copy) from the texture into another texture in .managed mode. If your texture is starting out in .managed storage, you probably need to create a blit command encoder and call synchronize(resource:) on the texture to ensure that its contents on the GPU are reflected on the CPU:

if let blitEncoder = commandBuffer.makeBlitCommandEncoder() {
    blitEncoder.synchronize(resource: outputTexture)
    blitEncoder.endEncoding()
}

Once the command buffer completes (which you can wait on by calling waitUntilCompleted or by adding a completion handler to the command buffer), you're ready to copy the data and create an image:

func makeImage(for texture: MTLTexture) -> CGImage? {
    assert(texture.pixelFormat == .bgra8Unorm)

    let width = texture.width
    let height = texture.height
    let pixelByteCount = 4 * MemoryLayout<UInt8>.size
    let imageBytesPerRow = width * pixelByteCount
    let imageByteCount = imageBytesPerRow * height
    let imageBytes = UnsafeMutableRawPointer.allocate(byteCount: imageByteCount, alignment: pixelByteCount)
    defer {
        imageBytes.deallocate()
    }

    texture.getBytes(imageBytes,
                     bytesPerRow: imageBytesPerRow,
                     from: MTLRegionMake2D(0, 0, width, height),
                     mipmapLevel: 0)

    swizzleBGRA8toRGBA8(imageBytes, width: width, height: height)

    guard let colorSpace = CGColorSpace(name: CGColorSpace.linearSRGB) else { return nil }
    let bitmapInfo = CGImageAlphaInfo.premultipliedLast.rawValue
    guard let bitmapContext = CGContext(data: nil,
                                        width: width,
                                        height: height,
                                        bitsPerComponent: 8,
                                        bytesPerRow: imageBytesPerRow,
                                        space: colorSpace,
                                        bitmapInfo: bitmapInfo) else { return nil }
    bitmapContext.data?.copyMemory(from: imageBytes, byteCount: imageByteCount)
    let image = bitmapContext.makeImage()
    return image
}

You'll notice a call in the middle of this function to a utility function called swizzleBGRA8toRGBA8. This function swaps the bytes in the image buffer so that they're in the RGBA order expected by CoreGraphics. It uses vImage (be sure to import Accelerate) and looks like this:

func swizzleBGRA8toRGBA8(_ bytes: UnsafeMutableRawPointer, width: Int, height: Int) {
    var sourceBuffer = vImage_Buffer(data: bytes,
                                     height: vImagePixelCount(height),
                                     width: vImagePixelCount(width),
                                     rowBytes: width * 4)
    var destBuffer = vImage_Buffer(data: bytes,
                                   height: vImagePixelCount(height),
                                   width: vImagePixelCount(width),
                                   rowBytes: width * 4)
    var swizzleMask: [UInt8] = [ 2, 1, 0, 3 ] // BGRA -> RGBA
    vImagePermuteChannels_ARGB8888(&sourceBuffer, &destBuffer, &swizzleMask, vImage_Flags(kvImageNoFlags))
}

Now we can write a function that enables us to write a texture to a specified URL:

func writeTexture(_ texture: MTLTexture, url: URL) {
    guard let image = makeImage(for: texture) else { return }

    if let imageDestination = CGImageDestinationCreateWithURL(url as CFURL, kUTTypePNG, 1, nil) {
        CGImageDestinationAddImage(imageDestination, image, nil)
        CGImageDestinationFinalize(imageDestination)
    }
}