Caps-iOS/Caps/Data/Classifier.swift

import Foundation
import Vision
import CoreML
import UIKit

/// Recognise categories in images
class Classifier: Logger {

    static let userDefaultsKey = "classifier.version"

    let model: VNCoreMLModel

    init(model: VNCoreMLModel) {
        self.model = model
    }
    /**
     Classify an image
     - Parameter image: The image to classify
     - Parameter completion: The callback with the match results
     - Parameter matches: A dictionary with a map from cap id to classifier match
     - Note: This method should not be scheduled on the main thread.
     */
    func recognize(image: CGImage, completion: @escaping (_ matches: [Int: Float]?) -> Void) {
        let image = CIImage(cgImage: image)
        let handler = VNImageRequestHandler(ciImage: image, orientation: .up)
        let request = VNCoreMLRequest(model: model) { request, error in
            let matches = self.process(request: request, error: error)
            completion(matches)
        }
        request.imageCropAndScaleOption = .centerCrop
        do {
            try handler.perform([request])
        } catch {
            self.error("Failed to perform classification: \(error)")
        }
    }

    private func process(request: VNRequest, error: Error?) -> [Int : Float]? {
        if let e = error {
            self.error("Unable to classify image: \(e.localizedDescription)")
            return nil
        }
        guard let result = request.results as? [VNClassificationObservation] else {
            self.error("Invalid classifier result: \(String(describing: request.results))")
            return nil
        }
        let matches = result.reduce(into: [:]) { $0[Int($1.identifier)!] = $1.confidence }

        log("Classifed image with \(matches.count) classes")
        return matches
    }
}