关于 UIImage -> CVPixelBuffer -> UIImage 转换的问题

Question

我正在 SwiftUI 中开发一个简单的去噪 POC，我想这样做：

加载输入图像
对输入图像应用 CoreML 模型（去噪）
显示输出图像

我有一些东西基于我在网上找到的几十个源代码。根据我读过的内容，CoreML 模型（至少是我正在使用的模型）接受 CVPixelBuffer 并输出 CVPixelBuffer。所以我的想法是执行以下操作：

将输入UI图像转换为 CVPixelBuffer
将 CoreML 模型应用于 CVPixelBuffer
将新创建的 CVPixelBuffer 转换为 UIImage

（请注意，我读过使用 Vision 框架，可以将 CGImage 直接输入到模型中。一旦我熟悉我要在这里实现的目标，我就会尝试这种方法因为我认为这是一个很好的练习。）

一开始，我想跳过步骤 (2) 以专注于转换问题。我试图在下面的代码中实现的是：

将输入UI图像转换为 CVPixelBuffer
将 CVPixelBuffer 转换为 UIImage

我不是 Swift 或 Objective-C 开发人员，所以我很确定我至少犯了一些错误。我发现这段代码非常复杂，我想知道是否有更好/更简单的方法来做同样的事情？

func convert(input: UIImage) -> UIImage? {

    // Input CGImage
    guard let cgInput = input.cgImage else {
        return nil
    }

    // Image size
    let width = cgInput.width
    let height = cgInput.height
    let region = CGRect(x: 0, y: 0, width: width, height: height)

    // Attributes needed to create the CVPixelBuffer
    let attributes = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
                      kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue]

    // Create the input CVPixelBuffer
    var pbInput:CVPixelBuffer? = nil
    let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                     width,
                                     height,
                                     kCVPixelFormatType_32ARGB,
                                     attributes as CFDictionary,
                                     &pbInput)

    // Sanity check
    if status != kCVReturnSuccess {
        return nil
    }

    // Fill the input CVPixelBuffer with the content of the input CGImage
    CVPixelBufferLockBaseAddress(pbInput!, CVPixelBufferLockFlags(rawValue: 0))
    guard let context = CGContext(data: CVPixelBufferGetBaseAddress(pbInput!),
                                  width: width,
                                  height: height,
                                  bitsPerComponent: cgInput.bitsPerComponent,
                                  bytesPerRow: cgInput.bytesPerRow,
                                  space: cgInput.colorSpace!,
                                  bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue) else {
                                    return nil
    }
    context.draw(cgInput, in: region)
    CVPixelBufferUnlockBaseAddress(pbInput!, CVPixelBufferLockFlags(rawValue: 0))

    // Create the output CGImage
    let ciOutput = CIImage(cvPixelBuffer: pbInput!)
    let temporaryContext = CIContext(options: nil)
    guard let cgOutput = temporaryContext.createCGImage(ciOutput, from: region) else {
        return nil
    }

    // Create and return the output UIImage
    return UIImage(cgImage: cgOutput)
}

当我在我的 SwiftUI 项目中使用此代码时，输入和输出图像看起来相同，但并不相同。我认为输入图像有一个与之关联的色彩图（ColorSync 配置文件），但在转换过程中丢失了。我假设我应该在 CGContext 创建期间使用 cgInput.colorSpace，但似乎使用 CGColorSpace(name: CGColorSpace.sRGB)! 效果更好。有人可以给我解释一下吗？

感谢您的帮助。

Answer 1

您也可以在 Core ML 中使用 CGImage 对象，但是您必须手动创建 MLFeatureValue 对象，然后将其放入 MLFeatureProvider 中以将其提供给模型.但这只处理模型输入，而不处理输出。

另一种选择是使用我的 CoreMLHelpers 存储库中的代码。

关于 UIImage -> CVPixelBuffer -> UIImage 转换的问题

Question regarding UIImage -> CVPixelBuffer -> UIImage conversion

uiimage

swift

cvpixelbuffer

coreml