Tensorflow lite image output different between python and iOS/Android - ios

Question

I convert Keras model to TF lite the output dimension is (1, 256, 256, 1).
the result on python is correct, but when I try to construct the image on ios swift the result is wrong.
Here is the code, that I use to construct an UIImage from a list of output.
// helper function
---------------------------------------
// MARK: - Extensions
extension Data {
init<T>(copyingBufferOf array: [T]) {
self = array.withUnsafeBufferPointer(Data.init)
}
/// Convert a Data instance to Array representation.
func toArray<T>(type: T.Type) -> [T] where T: ExpressibleByIntegerLiteral {
var array = [T](repeating: 0, count: self.count/MemoryLayout<T>.stride)
_ = array.withUnsafeMutableBytes { copyBytes(to: $0) }
return array
}
}
func imageFromSRGBColorArray(pixels: [UInt32], width: Int, height: Int) -> UIImage?
{
guard width > 0 && height > 0 else { return nil }
guard pixels.count == width * height else { return nil }
// Make a mutable copy
var data = pixels
// Convert array of pixels to a CGImage instance.
let cgImage = data.withUnsafeMutableBytes { (ptr) -> CGImage in
let ctx = CGContext(
data: ptr.baseAddress,
width: width,
height: height,
bitsPerComponent: 8,
bytesPerRow: MemoryLayout<UInt32>.size * width,
space: CGColorSpace(name: CGColorSpace.sRGB)!,
bitmapInfo: CGBitmapInfo.byteOrder32Little.rawValue
+ CGImageAlphaInfo.premultipliedFirst.rawValue
)!
return ctx.makeImage()!
}
// Convert the CGImage instance to an UIImage instance.
return UIImage(cgImage: cgImage)
}
let results = outputTensor.data.toArray(type: UInt32.self)
let maskImage = imageFromSRGBColorArray(pixels: results, width: 256, height: 256)
the result I get is completely wrong compared to python.
I think the function imageFromSRGBColorArray is not correct.
can anyone help me to figure out the problem?