Hands-on Metal: Image Processing using Apple’s GPU framework

The Filter Class

Create a new swift file named Filter. Make necessary imports and create a class named Filter.

import Metal
import MetalKit
class Filter {

Class Members

The first set of variables are related to GPU handling. Please to refer to the previous article for reference. The second set of variables is for image reading and its dimensions. The image dimensions are used across many functions so it makes sense to have it as a property. The last property is how many threads do we launch per block of the GPU. We wish to divide the 2D image into an equally split grid so that each block can be assigned a grid. You can read more here and here for more efficient means of thread grouping.

class Filter {var device: MTLDevice
var defaultLib: MTLLibrary?
var grayscaleShader: MTLFunction?
var commandQueue: MTLCommandQueue?
var commandBuffer: MTLCommandBuffer?
var commandEncoder: MTLComputeCommandEncoder?
var pipelineState: MTLComputePipelineState?
var inputImage: UIImage
var height, width: Int
// most devices have a limit of 512 threads per group
let threadsPerBlock = MTLSize(width: 16, height: 16, depth: 1)

The Constructor

We initiate the necessary GPU handling objects and set the pipeline state. Please to refer to the previous article for reference.

drag and drop any image file into the project directory
init(){self.device = MTLCreateSystemDefaultDevice()!
self.defaultLib = self.device.makeDefaultLibrary()
self.grayscaleShader = self.defaultLib?.makeFunction(name: "black")
self.commandQueue = self.device.makeCommandQueue()
self.commandBuffer = self.commandQueue?.makeCommandBuffer()
self.commandEncoder = self.commandBuffer?.makeComputeCommandEncoder()
if let shader = grayscaleShader { self.pipelineState = try? self.device.makeComputePipelineState(function: shader) } else { fatalError("unable to make compute pipeline") }self.inputImage = UIImage(named: "spidey.jpg")!
self.height = Int(self.inputImage.size.height)
self.width = Int(self.inputImage.size.width)
}

The ‘black’ Shader

#include <metal_stdlib>using namespace metal;kernel void black (
texture2d<float, access::write> outTexture [[texture(0)]],
texture2d<float, access::read> inTexture [[texture(1)]],
uint2 id [[thread_position_in_grid]]) {
float3 val = inTexture.read(id).rgb;
float gray = (val.r + val.g + val.b)/3.0;
float4 out = float4(gray, gray, gray, 1.0);
outTexture.write(out.rgba, id);
}

Helper Functions

There are 4 helper functions to convert images between different formats and 3 to just make the code cleaner. Look into the gist at the bottom for function definitions.

// create CGImage from UIImage
func getCGImage(from uiimg: UIImage) -> CGImage?
// create MTLTexture from CGImage
func getMTLTexture(from cgimg: CGImage) -> MTLTexture
// create CGImage from MTLTexture
func getCGImage(from mtlTexture: MTLTexture) -> CGImage?
// convert CGImage to UIImage
func getUIImage(from cgimg: CGImage) -> UIImage?
// we need an empty texture to write the output value after processing of each pixel
func getEmptyMTLTexture() -> MTLTexture?
// function that uses above helpers to convert the input image to texture
func getInputMTLTexture() -> MTLTexture?
// divides the image dimensions by threads per block to determine number of blocks to launch
func getBlockDimensions() -> MTLSize

ApplyFilter Function

Initially, we unwrap the optionals. Then we set the input and output texture and index them. We set GPU threads to work on the textures. Once commands are encoded the buffer is committed. We await the buffer to finish and convert the output texture back to UIImage to be displayed.

func applyFilter() -> UIImage? {if let encoder = self.commandEncoder, let buffer = self.commandBuffer, let outputTexture = getEmptyMTLTexture(), let inputTexture = getInputMTLTexture() {encoder.setTextures([outputTexture, inputTexture], range: 0..<2)encoder.dispatchThreadgroups(self.getBlockDimensions(), threadsPerThreadgroup: threadsPerBlock)encoder.endEncoding()buffer.commit()
buffer.waitUntilCompleted()
guard let outputImage = getCGImage(from: outputTexture) else { fatalError("Couldn't obtain CGImage from MTLTexture") }return getUIImage(from: outputImage)} else { fatalError("optional unwrapping failed") }}

Display Output

Create an imageView in the storyboard and hook it up to the view controller. In viewDidLoad() instantiate Filter and call its applyFilter. You can now display the result in your image view.

override func viewDidLoad() {super.viewDidLoad()
let filter = Filter()
let resultImage = Filter.applyFilter()
outputImageView.image = resultImage
}

--

--

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store
Avinash

Avinash

Data Science at ShareChat. Ola. IIT Madras.