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.

Helpmate

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)
    }
}

Ryeland

Weird. It outputs a PNG image with the proper size, but completely white. Is this probably an issue with my kernel code? – Helpmate

Possibly. What happens when you write a solid color (say red) to every pixel in your kernel? – Ryeland

I tried that, and it did nothing - still a blank white background. Is there anything special you have to do to prepare a texture for writing by a compute kernel? – Helpmate

Its usage should be set to .shaderWrite upon creation, but otherwise, not really. Can you confirm its storage mode is set to .managed and that you're waiting on the command buffer containing the blit synchronization command to finish before attempting to read back? – Ryeland

If you set a breakpoint on the getBytes line and QuickLook the texture object, does it have the expected contents? If not, the problem is with how you're dispatching your kernel. If so, the problem is likely in one of my functions (provided you're synchronizing correctly). – Ryeland

Thanks! I just had to set it to .managed and commit the blit encoder. Thanks! Accepted and upvoted. – Helpmate

Just FYI, but you don't need to swizzle BGRA to RGBA, CoreGraphics on iOS uses BGRA as the native pixel format and BGRA provides the best performance. – Fortuity

I am on macOS. So that wouldn't really apply, right? – Helpmate

@Fortuity How does one create a context or image with its components in BGRA order such that it can be written with ImageIO? – Ryeland

See https://mcmap.net/q/1188942/-how-to-determine-and-interpret-the-pixel-format-of-a-cgimage specifically passing kCGBitmapByteOrder32Big along in the bitmap info. – Fortuity

How does this work with other pixel formats? I'm trying to save a display-P3 image from a bgra10_xr_srgb texture, and I wrote a 16-bit swizzle method using vImagePermuteChannels_ARGB16F. The output looks pale, though. I suppose it must be something with the padding in the bgra10 format? I may be interpreting the 10 bits of each component as 16 bits... – Premolar

This is probably better as a separate question. In short, you need to carefully manage your colorspaces so that either the output file has the same colorspace as your texture and enough bits to store the requisite precision, or you need to apply tone-mapping and/or gamut mapping to transform the color data into a lower-precision format before using routines like the ones above. – Ryeland

According to Apple docs, on Apple silicon GPUs the default storage mode of an MTLTexture is shared. That must be why I don't need to do the blit to read the texture correctly in my M1 Mac. However, the default storage mode is managed for Intel-based Macs. I just get a blank texture on my Intel mac without the blit. I suppose we could have an "if" statement to check the storage mode & skip the blit in ARM Macs? – Premolar

The way I normally handle this is to check the hasUnifiedMemory property on the device and set the storage mode and optionally skip the blit based on it. Setting the storage mode explicitly helps document intent, rather than relying on the default. – Ryeland

Answer 2

based on warrenm's answer simplified way to save mtltexture:

func saveTexturePNG(texture:MTLTexture, to url:URL){
    let options = [CIImageOption.colorSpace: CGColorSpaceCreateDeviceRGB(),
               CIContextOption.outputPremultiplied: true,
               CIContextOption.useSoftwareRenderer: false] as! [CIImageOption : Any]
    guard let ciimage = CIImage(mtlTexture: texture, options: options) else {
        print("CIImage not created")
        return
    }
    let flipped = ciimage.transformed(by: CGAffineTransform(scaleX: 1, y: -1))
    guard let cgImage = CIContext().createCGImage(flipped,
                                 from: flipped.extent,
                                 format: CIFormat.RGBA8,
                                 colorSpace: CGColorSpace(name: CGColorSpace.sRGB)!)  else {
        print("CGImage not created")
        return
    }
    if let imageDestination = CGImageDestinationCreateWithURL(url as CFURL, UTType.png.identifier as CFString, 1, nil) {
        CGImageDestinationAddImage(imageDestination, cgImage, nil)
        CGImageDestinationFinalize(imageDestination)
    }

}

Blades