Вопрос относительно UIImage -> CVPixelBuffer -> UIImage преобразование

Question

Я работаю над простым шумоподавляющим PO C в SwiftUI, где я хочу:

1. Загрузка входного изображения
2. Применение модели CoreML (шумоподавления) к входу image
3. Отображение выходного изображения

У меня есть кое-что работающее на основе десятков исходных кодов, которые я нашел в Интернете. Основываясь на том, что я прочитал, модель CoreML (по крайней мере та, которую я использую) принимает CVPixelBuffer и выводит также CVPixelBuffer. Поэтому моя идея заключалась в следующем:

1. Преобразование входного UIImage в CVPixelBuffer
2. Применение модели CoreML к CVPixelBuffer
3. Преобразование недавно созданного CVPixelBuffer в UIImage

(Обратите внимание, что я читал, что с помощью фреймворка Vision можно напрямую ввести CGImage в модель. Я попробую этот подход, как только узнаю, что я Я пытаюсь достичь этого, так как считаю, что это хорошее упражнение.)

Вначале я хотел пропустить шаг (2), чтобы сосредоточиться на проблеме конверсии. В приведенном ниже коде я попытался добиться:

1. Преобразование входного UIImage в CVPixelBuffer
2. Преобразование CVPixelBuffer в UIImage

I ' Я не 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, входные и выходные изображения выглядели одинаково , но не было одинаковых. Я думаю, что входное изображение было связано с картой цветов (ColorSyn c Profile), которая была потеряна во время преобразования. Я предполагал, что должен был использовать cgInput.colorSpace во время создания CGContext, но казалось, что использование CGColorSpace(name: CGColorSpace.sRGB)! работало лучше. Может кто-нибудь объяснить мне это?

Спасибо за вашу помощь.

Answer 1

Вы также можете использовать CGImage объекты с Core ML, но вы должны создать объект MLFeatureValue вручную, а затем поместить его в MLFeatureProvider, чтобы передать его модели. Но это касается только ввода модели, а не вывода.

Другой вариант - использовать код из моего CoreMLHelpers репо.