如何将UIImage转换为形状为[1,3,512,512]的4维MLMultiArray？

要把UIImage转成符合CoreML模型要求的4维MLMultiArray，只需要在你现有代码基础上改两处核心逻辑：MLMultiArray的形状定义和内存偏移计算，直接看修改后的完整代码：

func prepareData(){
    // 先把图片缩到512x512，保证和模型输入尺寸匹配
    let targetSize = CGSize(width: 512, height: 512)
    guard let resizedImage = inputImage.resize(to: targetSize),
          let cvBufferInput = resizedImage.pixelBuffer() else {
        return
    }
    
    // 创建4维MLMultiArray：[batchSize, channel, height, width]，这里batchSize固定为1
    guard let mlImg = try? MLMultiArray(shape: [1, 3, NSNumber(value: targetSize.width), NSNumber(value: targetSize.height)], dataType: MLMultiArrayDataType.float32) else {
        return
    }
    
    self.testMLMultiArray(pixelBuffer: cvBufferInput, data: mlImg, height: Int(targetSize.height), width: Int(targetSize.width))
}

func testMLMultiArray(pixelBuffer: CVPixelBuffer, data: MLMultiArray, height: Int, width: Int) {
    CVPixelBufferLockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
    defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0)) } // 确保用完解锁，防止内存泄漏
    
    guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else { return }
    let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
    let buffer = baseAddress.assumingMemoryBound(to: UInt8.self)

    let ptrData = UnsafeMutablePointer<Float>(OpaquePointer(data.dataPointer))
    
    // 计算各维度的内存步长：batch步长=通道数*宽*高；通道步长=宽*高
    let batchStride = 3 * width * height
    let channelStride = width * height

    for y in 0..<height {
        for x in 0..<width {
            // 因为batch固定为0，所以batch偏移是0*batchStride，再叠加通道和坐标的偏移
            ptrData[0 * batchStride + 0 * channelStride + y * width + x] = Float(buffer[y * bytesPerRow + x * 4 + 1])
            ptrData[0 * batchStride + 1 * channelStride + y * width + x] = Float(buffer[y * bytesPerRow + x * 4 + 2])
            ptrData[0 * batchStride + 2 * channelStride + y * width + x] = Float(buffer[y * bytesPerRow + x * 4 + 3])
        }
    }
}

// 给UIImage加个缩放扩展，方便处理尺寸不匹配的情况
extension UIImage {
    func resize(to size: CGSize) -> UIImage? {
        UIGraphicsBeginImageContextWithOptions(size, false, UIScreen.main.scale)
        defer { UIGraphicsEndImageContext() }
        draw(in: CGRect(origin: .zero, size: size))
        return UIGraphicsGetImageFromCurrentImageContext()
    }
}

关键修改说明： #

1. 4维MLMultiArray创建：形状改成[1, 3, 512, 512]，对应CoreML常用的[batch, channel, height, width]格式，这里batch设为1表示单张图片输入。
2. 内存偏移计算：新增batchStride，因为batch索引固定为0，所以实际偏移是0，但保留这个逻辑方便后续扩展多图批量输入；通道和坐标的偏移逻辑和原来一致，只是叠加了batch维度的偏移。
3. 图片尺寸保证：加了缩放逻辑，确保输入图片是512x512，避免因尺寸不匹配导致模型报错。
4. 内存安全优化：用defer语句确保CVPixelBuffer用完后解锁，避免内存泄漏。

内容的提问来源于stack exchange，提问作者Phuong nguyen