How to convert MLMultiArray to UIImage on Swift?

Question

Im working on CoreML project. I have an image that needs to be stylized using a mlmodel. I'm using CoreMLHelpers for this. Im having an issue getting the output.

    func prediction(input: MLModelProviderInput) throws -> MLModelProviderOutput {
        let outFeatures = try model.prediction(from: input)
        let multiArray = outFeatures.featureValue(for: outputName)!.multiArrayValue!
        
        let outputImage = multiArray.image(min: 0, max: 255, axes: (4, 2, 3)) // does not work
        let result = MLModelProviderOutput(outputImage: outputImage!.pixelBuffer()!, outputName: outputName)
        
        return result
    }

MultiArray has shape of [1, 256, 256, 3]. Am I using the correct axes for this? (4, 2, 3).

CoreMLHelpers on github. My code breaks at line 146 with the console log- "Invalid axes (4, 2, 3) for shape [1, 256, 256, 3]"

https://github.com/hollance/CoreMLHelpers/blob/179ba6239886d2bc789430d6e466c54fddbbb654/CoreMLHelpers/MLMultiArray%2BImage.swift#L146

TIA

Answer 1

Sample:

extension MLMultiArray {

func postProcessImage(size: Int = 256) -> UIImage? {
    let rawPointer = malloc(size*size*3)!
    let bytes = rawPointer.bindMemory(to: UInt8.self, capacity: size*size*3)
    
    let mlArray = self.dataPointer.bindMemory(to: Float32.self, capacity: size*size*3)
    for index in 0..<self.count/(3) {
        bytes[index*3 + 0] = UInt8(max(min(mlArray[index]*255, 255), 0))
        bytes[index*3 + 1] = UInt8(max(min(mlArray[index + size*size]*255, 255), 0))
        bytes[index*3 + 2] = UInt8(max(min(mlArray[index + size*size*2]*255, 255), 0))
    }
    
    let selftureSize = size*size*3
    
    let provider = CGDataProvider(dataInfo: nil, data: rawPointer, size: selftureSize, releaseData: { (_, data, size) in
        data.deallocate()
    })!
   
    let rawBitmapInfo = CGImageAlphaInfo.none.rawValue
    let bitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)
    let pColorSpace = CGColorSpaceCreateDeviceRGB()

    let rowBytesCount = size*3
    let cgImage = CGImage(width: size, height: size, bitsPerComponent: 8, bitsPerPixel: 24, bytesPerRow: rowBytesCount, space: pColorSpace, bitmapInfo: bitmapInfo, provider: provider, decode: nil, shouldInterpolate: true, intent: CGColorRenderingIntent.defaultIntent)!
    let uiImage = UIImage(cgImage: cgImage)
        
    return uiImage
}

}