forked from mazzzystar/Queryable
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathImgEncoder.swift
More file actions
69 lines (57 loc) · 2.14 KB
/
ImgEncoder.swift
File metadata and controls
69 lines (57 loc) · 2.14 KB
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
//
// ImgEncoder.swift
// TestEncoder
//
// Created by Ke Fang on 2022/12/08.
//
import Foundation
import CoreML
import UIKit
public struct ImgEncoder {
var model: MLModel
init(resourcesAt baseURL: URL,
configuration config: MLModelConfiguration = .init()
) throws {
let imgEncoderURL = baseURL.appending(path:"ImageEncoder_float32.mlmodelc")
let imgEncoderModel = try MLModel(contentsOf: imgEncoderURL, configuration: config)
self.model = imgEncoderModel
}
public func computeImgEmbedding(img: UIImage) async throws -> MLShapedArray<Float32> {
let imgEmbedding = try await self.encode(image: img)
return imgEmbedding
}
// public init(model: MLModel) {
// self.model = model
// }
/// Prediction queue
let queue = DispatchQueue(label: "imgencoder.predict")
public func encode(image: UIImage) async throws -> MLShapedArray<Float32> {
do {
guard let resizedImage = try image.resizeImageTo(size:CGSize(width: 224, height: 224)) else {
throw ImageEncodingError.resizeError
}
guard let buffer = resizedImage.convertToBuffer() else {
throw ImageEncodingError.bufferConversionError
}
guard let inputFeatures = try? MLDictionaryFeatureProvider(dictionary: ["colorImage": buffer]) else {
throw ImageEncodingError.featureProviderError
}
let result = try queue.sync { try model.prediction(from: inputFeatures) }
guard let embeddingFeature = result.featureValue(for: "embOutput"),
let multiArray = embeddingFeature.multiArrayValue else {
throw ImageEncodingError.predictionError
}
return MLShapedArray<Float32>(converting: multiArray)
} catch {
print("Error in encoding: \(error)")
throw error
}
}
}
// Define the custom errors
enum ImageEncodingError: Error {
case resizeError
case bufferConversionError
case featureProviderError
case predictionError
}