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

Question

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

1. 加载输入图像
2. 将CoreML模型(去噪)应用于输入图像
3. 显示输出图像

我有一些基于我在网上找到的几十个源代码的东西。根据我所读到的，一个CoreML模型(至少我正在使用的模型)接受一个CVPixelBuffer，并输出一个CVPixelBuffer。我的想法是做以下几件事：

1. 将输入UIImage转换为CVPixelBuffer
2. 将CoreML模型应用于CVPixelBuffer
3. 将新创建的CVPixelBuffer转换为UIImage

(请注意，我已经读过，使用Vision框架，您可以直接将CGImage输入到模型中。一旦我熟悉了我在这里试图达到的目标，我就会尝试这种方法，因为我认为这是一个很好的练习。)

首先，我想跳过第(2)步，专注于转换问题。在下面的代码中，我试图实现的是：

1. 将输入UIImage转换为CVPixelBuffer
2. 将CVPixelBuffer转换为UIImage

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

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配置文件)，它在转换过程中丢失了。我以为我应该在创建cgInput.colorSpace的过程中使用CGContext，但似乎使用CGColorSpace(name: CGColorSpace.sRGB)!更好。谁能给我解释一下吗？

谢谢你的帮助。

Answer 1

您也可以在Core中使用CGImage对象，但是您必须手工创建MLFeatureValue对象，然后将其放入MLFeatureProvider中以将其交给模型。但这只涉及到模型输入，而不是输出。

另一种选择是使用我的CoreMLHelpers回购的代码。