Take a snapshot of current screen with Metal in swift

Question

I tried:

let scale = UIScreen.mainScreen().scale        
UIGraphicsBeginImageContextWithOptions(metalLayer.bounds.size, false, scale)
            
// metalLayer.renderInContext(UIGraphicsGetCurrentContext()!)

// self.view.layer ...

metalLayer.presentationLayer()!.renderInContext(UIGraphicsGetCurrentContext()!)
            
let image = UIGraphicsGetImageFromCurrentImageContext()           
UIGraphicsEndImageContext()
UIImageWriteToSavedPhotosAlbum(image, nil, nil, nil)

But the result is an empty screenshot. Any help would be nice!

Please keep in mind that I want to take a snapshot of a CAMetalLayer.

Ipswich

Answer 1

To make a screenshot, you need to get MTLTexture of the frame buffer.

1. If you use MTKView:

let texture = view.currentDrawable!.texture

2. If you don't use MTKView

Here's what I would do - I would have a property which holds last drawable presented to the screen:

let lastDrawableDisplayed: CAMetalDrawable?

And then when you present drawable to the screen, I would update it:

let commandBuffer = commandQueue.commandBuffer()
commandBuffer.addCompletedHandler { buffer in
  self.lastDrawableDisplayed = drawable
}

Now you whenever you need to take a screenshot, you can get a texture like this:

let texture = lastDrawableDisplayed.texture

Ok, now when you have MTLTexture you can convert it to CGImage and then to UIImage or NSImage.

Here's the code for OS X playground (MetalKit.MTLTextureLoader is not available for iOS playgrounds), in which I convert MTLTexture to CGImage

I made a small extension over MTLTexture for this.

import Metal
import MetalKit
import Cocoa

let device = MTLCreateSystemDefaultDevice()!
let textureLoader = MTKTextureLoader(device: device)

let path = "path/to/your/image.jpg"
let data = NSData(contentsOfFile: path)!

let texture = try! textureLoader.newTextureWithData(data, options: nil)

extension MTLTexture {
  
  func bytes() -> UnsafeMutablePointer<Void> {
    let width = self.width
    let height = self.height
    let rowBytes = self.width * 4
    let p = malloc(width * height * 4)
    
    self.getBytes(p, bytesPerRow: rowBytes, fromRegion: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)
    
    return p
  }
  
  func toImage() -> CGImage? {
    let p = bytes()
    
    let pColorSpace = CGColorSpaceCreateDeviceRGB()
    
    let rawBitmapInfo = CGImageAlphaInfo.NoneSkipFirst.rawValue | CGBitmapInfo.ByteOrder32Little.rawValue
    let bitmapInfo:CGBitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)
    
    let selftureSize = self.width * self.height * 4
    let rowBytes = self.width * 4
    let provider = CGDataProviderCreateWithData(nil, p, selftureSize, nil)
    let cgImageRef = CGImageCreate(self.width, self.height, 8, 32, rowBytes, pColorSpace, bitmapInfo, provider, nil, true, CGColorRenderingIntent.RenderingIntentDefault)!
    
    return cgImageRef
  }
}

if let imageRef = texture.toImage() {
  let image = NSImage(CGImage: imageRef, size: NSSize(width: texture.width, height: texture.height))
}

Regolith

at least one } is missing. – Playful

Doesn't capture transparent background. – Minoan

Answer 2

For swift 4.0, Just converting code provided by haawa

let lastDrawableDisplayed = metalView?.currentDrawable?.texture

if let imageRef = lastDrawableDisplayed?.toImage() {
    let uiImage:UIImage = UIImage.init(cgImage: imageRef)
}

extension MTLTexture {

    func bytes() -> UnsafeMutableRawPointer {
        let width = self.width
        let height   = self.height
        let rowBytes = self.width * 4
        let p = malloc(width * height * 4)

        self.getBytes(p!, bytesPerRow: rowBytes, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)

        return p!
    }

    func toImage() -> CGImage? {
        let p = bytes()

        let pColorSpace = CGColorSpaceCreateDeviceRGB()

        let rawBitmapInfo = CGImageAlphaInfo.noneSkipFirst.rawValue | CGBitmapInfo.byteOrder32Little.rawValue
        let bitmapInfo:CGBitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)

        let selftureSize = self.width * self.height * 4
        let rowBytes = self.width * 4
        let releaseMaskImagePixelData: CGDataProviderReleaseDataCallback = { (info: UnsafeMutableRawPointer?, data: UnsafeRawPointer, size: Int) -> () in
            return
        }
        let provider = CGDataProvider(dataInfo: nil, data: p, size: selftureSize, releaseData: releaseMaskImagePixelData)
        let cgImageRef = CGImage(width: self.width, height: self.height, bitsPerComponent: 8, bitsPerPixel: 32, bytesPerRow: rowBytes, space: pColorSpace, bitmapInfo: bitmapInfo, provider: provider!, decode: nil, shouldInterpolate: true, intent: CGColorRenderingIntent.defaultIntent)!

        return cgImageRef
    }
}

Shout

Thank you very much because in my case it is unique way to get image github.com/BradLarson/… #GPUImage3 – Gotama

Doesn't work for me: _validateGetBytes:51: failed assertion `texture must not be a framebufferOnly texture.' – Padlock

Answer 3

I didn't manage to get the accepted answer to work in Swift 4 / Metal 2 with XCode 9.1 on an iPhone 6s. Therefore I used a slightly different approach assuming lastDrawableDisplayed is saved as described in the accepted answer. Quick and dirty and without any exception handling:

let context = CIContext()
let texture = self.lastDrawableDisplayed!.texture
let cImg = CIImage(mtlTexture: texture, options: nil)!
let cgImg = context.createCGImage(cImg, from: cImg.extent)!
let uiImg = UIImage(cgImage: cgImg)

This is based on the documentation on the used CIImage Initializer:

init(mtlTexture:options:) Initializes an image object with data supplied by a Metal texture.

and CIImage Processing which describes how to create a CGImage with the use of CIContext:

CIContext() Create[s] a CIContext object (with default options) [...] context.createCGImage Render[s] the output image to a Core Graphics image that you can display or save to a file.

Hope that helps for anyone using Swift 4.

Edit: Additionally, I have multiple overlaying CAMetalLayer in my project and want to combine them into one single UIImage. Therefore it is needed to have references to the last CAMetalDrawable object of each layer. Before a new layer is added (and therefore used as the provider of nextDrawable()) I simply add the lastDrawableDisplayed to an array [CAMetalDrawable]. When "exporting" the layers I simply write all UIImages subsequently into a bitmap-based graphics context and get the final image with UIGraphicsGetImageFromCurrentImageContext().

Edit: If you are having trouble with orientation, try the following:

let uiImg = UIImage(cgImage: cgImg, scale: 1.0, orientation: UIImageOrientation.downMirrored)

Vitrify

Approach with CIImage works for me when I need to capture snapshot with transparent background! Thank you! – Minoan

this is a brilliant answer! – Pulverulent

Does not work for me either: -[MTLDebugComputeCommandEncoder setTexture:atIndex:]:373: failed assertion frameBufferOnly texture not supported for compute. Needed (with performance hit): metalView.framebufferOnly = false – Padlock

Answer 4

MTLTexture's method toImage needs to release data's memory in release data callback:

let releaseMaskImagePixelData: CGDataProviderReleaseDataCallback = { (info: UnsafeMutableRawPointer?, data: UnsafeRawPointer, size: Int) -> () in
            data.dealloc()
            return
 }           

Elena

Answer 5

There is a memory leak issue due to not free bytes created for CGDataProvider. From @Ankit Kedia answer.

Here is a Swift 5.0 version complete code.

extension MTLTexture {
    func bytes() -> UnsafeMutableRawPointer {
        let width = self.width
        let height = self.height
        let rowBytes = self.width * 4
        let p = malloc(width * height * 4)!
        getBytes(p, bytesPerRow: rowBytes, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)
        return p
    }

    func toImage() -> CGImage? {
        let p = self.bytes()
        let pColorSpace = CGColorSpaceCreateDeviceRGB()
        let rawBitmapInfo = CGImageAlphaInfo.noneSkipFirst.rawValue | CGBitmapInfo.byteOrder32Little.rawValue
        let bitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)
        
        let selftureSize = self.width * self.height * 4
        let rowBytes = self.width * 4
        if let provider = CGDataProvider(dataInfo: nil, data: p, size: selftureSize, releaseData: { _, p, _ in
            p.deallocate()
        }) {
            return CGImage(width: width,
                           height: height,
                           bitsPerComponent: 8,
                           bitsPerPixel: 32,
                           bytesPerRow: rowBytes,
                           space: pColorSpace,
                           bitmapInfo: bitmapInfo,
                           provider: provider,
                           decode: nil,
                           shouldInterpolate: true,
                           intent: CGColorRenderingIntent.defaultIntent)
        }
        return nil
    }
}

Haught

Answer 6

swift 4.2

extension MTLTexture {
    
    func bytes() -> UnsafeMutableRawPointer {
        let width = self.width
        let height   = self.height
        let rowBytes = self.width * 4
        let p = malloc(width * height * 4)
        
        self.getBytes(p!, bytesPerRow: rowBytes, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)
        
        return p!
    }
    
    func toImage() -> CGImage? {
        let p = bytes()
        
        let pColorSpace = CGColorSpaceCreateDeviceRGB()
        
        let rawBitmapInfo = CGImageAlphaInfo.noneSkipFirst.rawValue | CGBitmapInfo.byteOrder32Little.rawValue
        let bitmapInfo:CGBitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)
        
        let selftureSize = self.width * self.height * 4
        let rowBytes = self.width * 4
        let releaseMaskImagePixelData: CGDataProviderReleaseDataCallback = { (info: UnsafeMutableRawPointer?, data: UnsafeRawPointer, size: Int) -> () in
            return
        }
        let provider = CGDataProvider(dataInfo: nil, data: p, size: selftureSize, releaseData: releaseMaskImagePixelData)
        let cgImageRef = CGImage(width: self.width, height: self.height, bitsPerComponent: 8, bitsPerPixel: 32, bytesPerRow: rowBytes, space: pColorSpace, bitmapInfo: bitmapInfo, provider: provider!, decode: nil, shouldInterpolate: true, intent: CGColorRenderingIntent.defaultIntent)!
        
        return cgImageRef
    }
}

Slack

Duplicate of @Ankit's answer – Stelle