-
Notifications
You must be signed in to change notification settings - Fork 3
/
ViewController.swift
136 lines (116 loc) · 5.48 KB
/
ViewController.swift
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
/*
See LICENSE folder for this sample’s licensing information.
Abstract:
View controller for selecting images and applying Vision + Core ML processing.
*/
import UIKit
import CoreML
import Vision
import ImageIO
class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavigationControllerDelegate {
@IBOutlet weak var imageView: UIImageView!
@IBOutlet weak var correctedImageView: UIImageView!
@IBOutlet weak var classificationLabel: UILabel!
@IBAction func takePicture(_ sender: Any) {
let picker = UIImagePickerController()
picker.delegate = self
picker.sourceType = .camera
picker.cameraCaptureMode = .photo
present(picker, animated: true)
}
@IBAction func chooseImage(_ sender: Any) {
// The photo library is the default source, editing not allowed
let picker = UIImagePickerController()
picker.delegate = self
picker.sourceType = .savedPhotosAlbum
present(picker, animated: true)
}
var inputImage: CIImage! // The image to be processed.
func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [String : Any]) {
picker.dismiss(animated: true)
classificationLabel.text = "Analyzing Image…"
correctedImageView.image = nil
guard let uiImage = info[UIImagePickerControllerOriginalImage] as? UIImage
else { fatalError("no image from image picker") }
guard let ciImage = CIImage(image: uiImage)
else { fatalError("can't create CIImage from UIImage") }
let orientation = CGImagePropertyOrientation(uiImage.imageOrientation)
inputImage = ciImage.oriented(forExifOrientation: Int32(orientation.rawValue))
// Show the image in the UI.
imageView.image = uiImage
// Run the rectangle detector, which upon completion runs the ML classifier.
let handler = VNImageRequestHandler(ciImage: ciImage, orientation: CGImagePropertyOrientation(rawValue: UInt32(Int32(orientation.rawValue)))!)
DispatchQueue.global(qos: .userInteractive).async {
do {
try handler.perform([self.rectanglesRequest])
} catch {
print(error)
}
}
}
lazy var classificationRequest: VNCoreMLRequest = {
// Load the ML model through its generated class and create a Vision request for it.
do {
let model = try VNCoreMLModel(for: MNISTClassifier().model)
return VNCoreMLRequest(model: model, completionHandler: self.handleClassification)
} catch {
fatalError("can't load Vision ML model: \(error)")
}
}()
func handleClassification(request: VNRequest, error: Error?) {
guard let observations = request.results as? [VNClassificationObservation]
else { fatalError("unexpected result type from VNCoreMLRequest") }
guard let best = observations.first
else { fatalError("can't get best result") }
DispatchQueue.main.async {
self.classificationLabel.text = "Classification: \"\(best.identifier)\" Confidence: \(best.confidence)"
}
}
lazy var rectanglesRequest: VNDetectRectanglesRequest = {
return VNDetectRectanglesRequest(completionHandler: self.handleRectangles)
}()
func handleRectangles(request: VNRequest, error: Error?) {
guard let observations = request.results as? [VNRectangleObservation]
else { fatalError("unexpected result type from VNDetectRectanglesRequest") }
guard let detectedRectangle = observations.first else {
DispatchQueue.main.async {
self.classificationLabel.text = "No rectangles detected."
}
return
}
let imageSize = inputImage.extent.size
// Verify detected rectangle is valid.
let boundingBox = detectedRectangle.boundingBox.scaled(to: imageSize)
guard inputImage.extent.contains(boundingBox)
else { print("invalid detected rectangle"); return }
// Rectify the detected image and reduce it to inverted grayscale for applying model.
let topLeft = detectedRectangle.topLeft.scaled(to: imageSize)
let topRight = detectedRectangle.topRight.scaled(to: imageSize)
let bottomLeft = detectedRectangle.bottomLeft.scaled(to: imageSize)
let bottomRight = detectedRectangle.bottomRight.scaled(to: imageSize)
let correctedImage = inputImage
.cropped(to: boundingBox)
.applyingFilter("CIPerspectiveCorrection", parameters: [
"inputTopLeft": CIVector(cgPoint: topLeft),
"inputTopRight": CIVector(cgPoint: topRight),
"inputBottomLeft": CIVector(cgPoint: bottomLeft),
"inputBottomRight": CIVector(cgPoint: bottomRight)
])
.applyingFilter("CIColorControls", parameters: [
kCIInputSaturationKey: 0,
kCIInputContrastKey: 32
])
.applyingFilter("CIColorInvert", parameters: [:])
// Show the pre-processed image
DispatchQueue.main.async {
self.correctedImageView.image = UIImage(ciImage: correctedImage)
}
// Run the Core ML MNIST classifier -- results in handleClassification method
let handler = VNImageRequestHandler(ciImage: correctedImage)
do {
try handler.perform([classificationRequest])
} catch {
print(error)
}
}
}