TensorFlowLite.Tensor to UUImage - ios

Question

I am new with swift, TFlite and IOS. I succeed to convert, run my model. However at the end, I need to reconstruct an image. My TFlite model return a TFLite.tensor Float32 4d - shape (1, height, width, 3).
let outputTensor: Tensor
outputTensor = try myInterpreter.output(at: 0)
I am looking to make a RGB picture without alpha. In python, it will like this:
Image.fromarray((np.array(outputTensor.data) * 255).astype(np.uint8))
From my understanding the best way will be to make a CVPixelBuffer, apply a CoreOS transformation (for the x255) and finally make the UUImage. I am deeply lost in the IOS doc, it exists many possibilities, does the community has a suggestion ?
++t

Answer 1

From Google example, an extension of UIImage can be coded:
extension UIImage {
convenience init?(data: Data, size: CGSize) {
let width = Int(size.width)
let height = Int(size.height)
let floats = data.toArray(type: Float32.self)
let bufferCapacity = width * height * 4
let unsafePointer = UnsafeMutablePointer<UInt8>.allocate(capacity: bufferCapacity)
let unsafeBuffer = UnsafeMutableBufferPointer<UInt8>(
start: unsafePointer,
count: bufferCapacity)
defer {
unsafePointer.deallocate()
}
for x in 0..<width {
for y in 0..<height {
let floatIndex = (y * width + x) * 3
let index = (y * width + x) * 4
let red = UInt8(floats[floatIndex] * 255)
let green = UInt8(floats[floatIndex + 1] * 255)
let blue = UInt8(floats[floatIndex + 2] * 255)
unsafeBuffer[index] = red
unsafeBuffer[index + 1] = green
unsafeBuffer[index + 2] = blue
unsafeBuffer[index + 3] = 0
}
}
let outData = Data(buffer: unsafeBuffer)
// Construct image from output tensor data
let alphaInfo = CGImageAlphaInfo.noneSkipLast
let bitmapInfo = CGBitmapInfo(rawValue: alphaInfo.rawValue)
.union(.byteOrder32Big)
let colorSpace = CGColorSpaceCreateDeviceRGB()
guard
let imageDataProvider = CGDataProvider(data: outData as CFData),
let cgImage = CGImage(
width: width,
height: height,
bitsPerComponent: 8,
bitsPerPixel: 32,
bytesPerRow: MemoryLayout<UInt8>.size * 4 * Int(size.width),
space: colorSpace,
bitmapInfo: bitmapInfo,
provider: imageDataProvider,
decode: nil,
shouldInterpolate: false,
intent: .defaultIntent
)
else {
return nil
}
self.init(cgImage: cgImage)
}
}
Then the image can be easily constructed from the inference of TFLite.
let outputTensor: Tensor
outputTensor = try decoder.output(at: 0)
image = UIImage(data: outputTensor.data, size: size) ?? UIImage()