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

Posted 2023-03-27

【中文标题】关于 UIImage -> CVPixelBuffer -> UIImage 转换的问题

【英文标题】：Question regarding UIImage -> CVPixelBuffer -> UIImage conversion

【发布时间】：2020-06-03 20:40:12

【问题描述】：

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

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

我有一些基于我在网上找到的数十个源代码的工作。根据我所读到的，CoreML 模型（至少我正在使用的模型）接受一个 CVPixelBuffer 并输出一个 CVPixelBuffer。所以我的想法是做以下事情：

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

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

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

将输入的 UIImage 转换为 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)! 效果更好。有人可以向我解释一下吗？

感谢您的帮助。

【参考方案1】：

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

另一种选择是使用我的CoreMLHelpers repo 中的代码。

【讨论】：

感谢您的回答。我会试一试，因为它可能会简化我的代码，尤其是 CGImage 到 CVPixelBuffer 的转换。但无论如何我必须处理 CVPixelBuffer 到 UIImage 的转换。

这很简单：let image = UIImage(ciImage: CIImage(cvPixelBuffer: pixelBuffer)).

由于某种原因let image = UIImage(ciImage: CIImage(cvPixelBuffer: pixelBuffer)) 不起作用。我必须使用CIContext，然后使用createCGImage 使其工作。但我真的对我拥有的代码没有信心。这就是我创建这篇文章的原因。

抱歉不够精确。 “不工作”是指我的输出图像完全是白色的。我想我一定做错了什么。

你的意思是你总是得到一个全白的图像？在这种情况下，我会尝试从 Python 运行模型，看看你会得到什么。