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