Question regarding UIImage -> CVPixelBuffer -> UIImage conversion

Question

I am working on a simple denoising POC in SwiftUI where I want to:

Load an input image
Apply a CoreML model (denoising) to the input image
Display the output image

I have something working based on dozens of source codes I found online. Based on what I've read, a CoreML model (at least the one I'm using) accepts a CVPixelBuffer and outputs also a CVPixelBuffer. So my idea was to do the following:

Convert the input UIImage into a CVPixelBuffer
Apply the CoreML model to the CVPixelBuffer
Convert the newly created CVPixelBuffer into a UIImage

(Note that I've read that using the Vision framework, one can input a CGImage directly into the model. I'll try this approach as soon as I'm familiar with what I'm trying to achieve here as I think it is a good exercise.)

As a start, I wanted to skip the step (2) to focus on the conversion problem. What I tried to achieve in the code bellow is:

Convert the input UIImage into a CVPixelBuffer
Convert the CVPixelBuffer into a UIImage

I'm not a Swift or an Objective-C developer, so I'm pretty sure that I've made at least a few mistakes. I found this code quite complex and I was wondering if there was a better / simpler way to do the same thing?

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)
}

When I used this code in my SwiftUI project, input and output images looked the same, but there were not identical. I think the input image had a colormap (ColorSync Profile) associated to it that have been lost during the conversion. I assumed I was supposed to use cgInput.colorSpace during the CGContext creation, but it seemed that using CGColorSpace(name: CGColorSpace.sRGB)! was working better. Can somebody please explain that to me?

Thanks for your help.

score 0 · Accepted Answer · answered Feb 19 '20 at 21:38

0

You can also use CGImage objects with Core ML, but you have to create the MLFeatureValue object by hand and then put it into an MLFeatureProvider to give it to the model. But that only takes care of the model input, not the output.

Another option is to use the code from my CoreMLHelpers repo.

answered Feb 19 '20 at 21:38

Matthijs Hollemans

7,706
2
16
23

Thanks for your answer. I'll give it a try as it might simplify my code, especially the `CGImage` to `CVPixelBuffer` conversion. But in any case I have to deal with the CVPixelBuffer to UIImage conversion. – Vincent Garcia Feb 20 '20 at 16:10
That is pretty simple: `let image = UIImage(ciImage: CIImage(cvPixelBuffer: pixelBuffer))`. – Matthijs Hollemans Feb 20 '20 at 16:51
For some reason `let image = UIImage(ciImage: CIImage(cvPixelBuffer: pixelBuffer))` doesn't work. I had to use a `CIContext` and then use `createCGImage` to make it work. But I'm really not confident with the code I have. This is why I created this post. – Vincent Garcia Feb 20 '20 at 18:39
Sorry for the lack of precision. By "not working" I meant that my output image is totally white. I must have done something wrong I guess. – Vincent Garcia Feb 21 '20 at 08:13
You mean you always get a totally white image? In that case, I'd try running the model from Python to see what you get there. – Matthijs Hollemans Feb 21 '20 at 11:21
I'm not even using the model here. If I replace in the code above the last part where I convert the CVPixelBuffer into a UIImage using what you proposed, the output image is white. But it is fine with the original code. – Vincent Garcia Feb 21 '20 at 14:51

Question regarding UIImage -> CVPixelBuffer -> UIImage conversion

1 Answers1