Код, который я предоставляю ниже, кажется, работает нормально. Это может быть не оптимально, но по крайней мере это работает, что является хорошим первым шагом. Основная трудность заключалась в том, чтобы понять, как работает MLFeatureProvider
:
class SigmaProvider: MLFeatureProvider {
let sigma: Double
init(sigma: Double) {
self.sigma = sigma
}
var featureNames: Set<String> {
get {
return ["my_sigma"]
}
}
func featureValue(for featureName: String) -> MLFeatureValue? {
if (featureName == "my_sigma") {
let array = try! MLMultiArray(shape: [1], dataType: .float32)
array[0] = NSNumber(value: self.sigma)
return MLFeatureValue(multiArray: array)
}
return nil
}
}
Этот класс определяет вход с именем my_sigma
и типом MLMultiArray
. Ниже приведен код, который применяет модель CoreML (с именем model
) к входному UIImage:
// Input CGImage
guard let cgImage = uiImage.cgImage else {
return nil
}
// Load and setup the Vision model
guard let visionModel = try? VNCoreMLModel(for: model.model) else {
fatalError("Cannot load Vision ML model")
}
visionModel.inputImageFeatureName = "my_input"
visionModel.featureProvider = SigmaProvider(sigma: self.sigma)
// Create the request
let request = VNCoreMLRequest(model: visionModel)
request.usesCPUOnly = false
// Handler
let handler = VNImageRequestHandler(cgImage: cgImage)
// Process
do {
try handler.perform([request])
}
catch {
print("Error while performing Vision request: \(error)")
return nil
}
// Get the result
guard let results = request.results else {
print("No request results")
return nil
}
// Convert the resulting CVPixelBuffer into a UIImage
for case let resultingImage as VNPixelBufferObservation in results {
if resultingImage.featureName == "my_output" {
let ciOutput = CIImage(cvPixelBuffer: resultingImage.pixelBuffer)
UIGraphicsBeginImageContextWithOptions(uiImage.size, true, 1.0)
UIImage(ciImage: ciOutput).draw(at: CGPoint(x: 0, y: 0))
let output = UIGraphicsGetImageFromCurrentImageContext()
UIGraphicsEndImageContext()
return output
}
}
Обратите внимание, что первый аргумент моей модели, my_input
, не требуется в SigmaProvider
учебный класс. Это указывается с помощью свойства inputImageFeatureName
. Наконец, я получаю доступ к my_output
, просматривая все результаты.
Надеюсь, это кому-нибудь поможет.