File size: 1,515 Bytes
1f71164 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | import CoreML
import Foundation
import ImageIO
guard CommandLine.arguments.count == 3 else {
fatalError("Usage: encode model.mlpackage image.jpg (image must already be 448×448 RGB)")
}
let package = URL(fileURLWithPath: CommandLine.arguments[1])
let imageURL = URL(fileURLWithPath: CommandLine.arguments[2])
let compiled = try MLModel.compileModel(at: package)
let configuration = MLModelConfiguration()
configuration.computeUnits = .all
let model = try MLModel(contentsOf: compiled, configuration: configuration)
guard let source = CGImageSourceCreateWithURL(imageURL as CFURL, nil),
let image = CGImageSourceCreateImageAtIndex(source, 0, nil),
let constraint = model.modelDescription.inputDescriptionsByName["image"]?.imageConstraint,
image.width == constraint.pixelsWide, image.height == constraint.pixelsHigh else {
fatalError("Provide an orientation-corrected RGB image resized to the model's input dimensions")
}
let input = try MLFeatureValue(cgImage: image, constraint: constraint, options: nil)
let output = try model.prediction(from: MLDictionaryFeatureProvider(dictionary: ["image": input]))
guard let array = output.featureValue(for: "embedding")?.multiArrayValue,
array.count == 768, array.dataType == .float32 else {
fatalError("Unexpected model output")
}
let vector = Array(UnsafeBufferPointer(start: array.dataPointer.assumingMemoryBound(to: Float.self), count: array.count))
print("\(vector.count) values, L2 norm \(sqrt(vector.reduce(0) { $0 + $1 * $1 }))")
|