iOS应用中使用Swift、UIKit与CoreML调用图像分类模型预测结果异常的问题排查

结合你描述的情况——自己训练的模型在Xcode Preview里正常，但Swift调用结果完全错误，替换官方MobileNet模型却能正常工作——核心问题大概率出在模型输入格式/预处理逻辑不匹配上，毕竟你的代码框架是没问题的。下面是具体的排查步骤和修复建议：

1. 核对模型的输入规范 #

首先打开你的myImageClassifier.mlmodel文件，重点查看以下几个参数：

• 输入图像尺寸：模型训练时要求的固定尺寸（比如常见的224x224、299x299），你的代码目前是直接用原图像尺寸转换，这会导致输入尺寸和模型要求不符，预测结果必然错误。
• 像素格式：模型要求的像素格式（比如RGB、BGR、ARGB），你的转换代码用的是kCVPixelFormatType_32ARGB，如果模型训练时用的是其他格式，通道顺序不匹配就会导致识别失效。
• 预处理规则：训练时是否对图像做了归一化（比如把像素值从0-255缩放到0-1或-1到1）、是否需要翻转图像等，这些细节必须和你的代码处理逻辑完全一致。

2. 修复图像转换代码的关键问题 #

你的buffer(from:)函数缺少了图像缩放的步骤，并且可能存在颜色通道或坐标系的偏差。这里给出修改后的代码，适配大多数CoreML模型的要求：

// 先添加一个将图像缩放到指定尺寸的工具函数
private func scaledImage(_ image: UIImage, to targetSize: CGSize) -> UIImage {
    UIGraphicsBeginImageContextWithOptions(targetSize, false, UIScreen.main.scale)
    image.draw(in: CGRect(origin: .zero, size: targetSize))
    let scaledImage = UIGraphicsGetImageFromCurrentImageContext()!
    UIGraphicsEndImageContext()
    return scaledImage
}

// 修改转换函数，先缩放图像到模型要求的尺寸，再转换为CVPixelBuffer
func buffer(from image: UIImage, targetSize: CGSize) -> CVPixelBuffer? {
    let scaledImage = self.scaledImage(image, to: targetSize)
    let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, 
                 kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
    
    var pixelBuffer: CVPixelBuffer?
    let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                     Int(targetSize.width),
                                     Int(targetSize.height),
                                     kCVPixelFormatType_32ARGB,
                                     attrs,
                                     &pixelBuffer)
    
    guard status == kCVReturnSuccess else { return nil }
    
    CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
    let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)
    let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
    
    // 注意：如果模型要求BGR格式，需要修改bitmapInfo和颜色空间
    let context = CGContext(data: pixelData,
                            width: Int(targetSize.width),
                            height: Int(targetSize.height),
                            bitsPerComponent: 8,
                            bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!),
                            space: rgbColorSpace,
                            bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
    
    // 坐标系翻转：UIKit的图像坐标系和CoreML常用的坐标系可能上下颠倒，根据模型训练时的设置调整
    context?.translateBy(x: 0, y: targetSize.height)
    context?.scaleBy(x: 1.0, y: -1.0)
    
    UIGraphicsPushContext(context!)
    scaledImage.draw(in: CGRect(x: 0, y: 0, width: targetSize.width, height: targetSize.height))
    UIGraphicsPopContext()
    
    CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
    return pixelBuffer
}

调用时，传入从mlmodel文件中查到的目标尺寸：

// 替换成你模型实际要求的尺寸，比如224x224
let modelInputSize = CGSize(width: 224, height: 224)
let pixelImage = buffer(from: (image ?? UIImage(named: "imagePlaceholder"))!, targetSize: modelInputSize)

3. 验证转换后的图像是否符合要求 #

你可以把转换后的CVPixelBuffer再转回UIImage，检查是否和原图像的内容、颜色一致：

func image(from pixelBuffer: CVPixelBuffer) -> UIImage? {
    let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
    let context = CIContext(options: [.useSoftwareRenderer: false])
    guard let cgImage = context.createCGImage(ciImage, from: ciImage.extent) else { return nil }
    return UIImage(cgImage: cgImage)
}

// 调用验证
if let pixelImage = pixelImage, let convertedImage = image(from: pixelImage) {
    print("转换后的图像尺寸：\(convertedImage.size)")
    // 可以把convertedImage显示到一个临时UIImageView里查看颜色是否正常
}

4. 检查模型导出时的设置 #

• 确认导出CoreML模型时，是否正确选择了图像输入类型，并且类别标签（Class Labels）的映射和训练时一致。
• 如果训练时用的是PyTorch/TensorFlow，导出CoreML时是否正确处理了输入输出的格式转换（比如TensorFlow的NHWC格式转CoreML的NCHW）。

最后总结 #

因为官方MobileNet模型能正常工作，说明你的代码框架没有问题，问题就出在你的模型输入要求和图像转换逻辑不匹配上。优先核对尺寸和预处理规则，这两个是最常见的坑，调整后应该就能解决预测结果异常的问题。

内容的提问来源于stack exchange，提问作者Sheldon Li