CODEWITHSUNDEEP
iosswiftpytorchcoremltoolsmlmodel
Phuong nguyen
Phuong nguyen

Reputation: 1

Convert UIImage to MLMultiArray with 4 dimension shape

I have tried to convert UIImage into MLMultiArray to pass as input to CoreML model. It's successful with 3 dimension of shape. But I don't know how to do with 4 dimension of shape. shape like: [1, 3, 512, 512] the input of model Thank.

func prepareData(){

 guard let cvBufferInput = inputImage.pixelBuffer() else {
               
                   return
               }
               
               guard let mlImg = try? MLMultiArray(shape: [3, NSNumber(value: width), NSNumber(value: height)], dataType: MLMultiArrayDataType.float32) else {
                  
                   return
               }
               
               self.testMLMultiArray(pixelBuffer: cvBufferInput, data: mlImg, height: width, width: height)
}

func testMLMultiArray(pixelBuffer: CVPixelBuffer, data: MLMultiArray, height: Int, width: Int) {
        CVPixelBufferLockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer)
        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
        let buffer = baseAddress!.assumingMemoryBound(to: UInt8.self)

        var ptrData = UnsafeMutablePointer<Float>(OpaquePointer(data.dataPointer))
        ptrData = ptrData.advanced(by: 0)

        let cStride = width * height

        for y in 0..<height {
            for x in 0..<width {
                ptrData[y*width + x + cStride * 0] = (Float)(buffer[y*bytesPerRow+x*4+1])
                ptrData[y*width + x + cStride * 1] = (Float)(buffer[y*bytesPerRow+x*4+2])
                ptrData[y*width + x + cStride * 2] = (Float)(buffer[y*bytesPerRow+x*4+3])
            }
        }
    }

Upvotes: 0

Views: 180

Answers (1)

timyau
timyau

Reputation: 152

I think add 1 in the shape might be OK?

guard let mlImg = try? MLMultiArray(shape: [1, 3, NSNumber(value: width), NSNumber(value: height)], dataType: MLMultiArrayDataType.float32)

Upvotes: 0

Related Questions