Created
August 7, 2017 15:16
-
-
Save mkj-is/aef9f19e6e8a1d2bf2ed54b6eef24d75 to your computer and use it in GitHub Desktop.
Example use of image classification using Vision framework
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Foundation | |
import FuntastyKit | |
import PromiseKit | |
import Photos | |
import Vision | |
import QuartzCore | |
protocol ImageClassificationService: Service { | |
func classify(asset: PHAsset) -> Promise<String> | |
} | |
final class VisionImageClassificationService: ImageClassificationService { | |
private let model: VNCoreMLModel = try! VNCoreMLModel(for: Resnet50().model) | |
private let modelInputResolution: CGFloat = 224 | |
private var modelInputSize: CGSize { | |
return CGSize(width: modelInputResolution, height: modelInputResolution) | |
} | |
func classify(asset: PHAsset) -> Promise<String> { | |
return Promise { fullfill, reject in | |
PHImageManager.default().requestImage(for: asset, targetSize: modelInputSize, contentMode: .aspectFill, options: options) { image, info in | |
let request = VNCoreMLRequest(model: self.model) { request, error in | |
if let error = error { | |
reject(error) | |
} else { | |
let observations = request.results as? [VNClassificationObservation] ?? [] | |
fullfill(observations.first?.identifier ?? "Picture") | |
} | |
} | |
let handler = VNImageRequestHandler(cgImage: image!.cgImage!) | |
try! handler.perform([request]) | |
} | |
} | |
} | |
private var options: PHImageRequestOptions { | |
let options = PHImageRequestOptions() | |
options.isSynchronous = true | |
options.deliveryMode = .fastFormat | |
options.resizeMode = .exact | |
options.normalizedCropRect = CGRect(origin: .zero, size: modelInputSize) | |
return options | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment