ImageClassifier

ImageClassifier.swift

import UIKit

import CoreML

protocol ImageClassifierLogic {
  var version: String? { get }
  func prediction(image: UIImage?) -> [String: Double]?
}

final class ImageClassifier: ImageClassifierLogic {

  var version: String? {
    guard #available(iOS 13.0, *) else { return nil }
    guard let imageClassifierModel = try? KarrotMLModel(configuration: MLModelConfiguration()) else { return nil }
    return imageClassifierModel.model.modelDescription.metadata[MLModelMetadataKey.versionString] as? String
  }

  func prediction(image: UIImage?) -> [String: Double]? {
    guard #available(iOS 13.0, *) else { return nil }
    guard let cgImage = image?.cgImage else { return nil }
    guard let imageClassifierModel = try? KarrotMLModel(configuration: MLModelConfiguration()) else { return nil }
    let result = try? imageClassifierModel.prediction(input: KarrotMLModelInput(imageWith: cgImage))
    return result?.scores
  }
}