Multiple requests support

xzzz9097 · Jul 2, 2017 · ee0f5e3 · ee0f5e3
1 parent b89b142
commit ee0f5e3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 26 deletions.
diff --git a/Machine/ViewController.swift b/Machine/ViewController.swift
@@ -16,8 +16,8 @@ class ViewController: NSViewController,
 
     var requestDelegate = VisionRequestCaptureDelegate.default
 
-    var resnetDelegate  = VisionRequestCaptureDelegate.default
-
+    var visionRequests: [VNRequest] = [ ]
+        
     var captureSession = CaptureSession()
 
     var hideFace = false
@@ -91,35 +91,38 @@ class ViewController: NSViewController,
     }
 
     func loadCaptureSession() {
-        let detectedObjectHandler = VisionRequestResultHandler(delegate: self)
+        prepareFaceRequest()
+        prepareResnetRequest()
 
         requestDelegate.configure(
-            for: VNDetectFaceRectanglesRequest(
-                completionHandler: detectedObjectHandler.requestResultHandler
-            ),
+            for: visionRequests,
             failHandler: { self.resetFaceViews() }
         )
 
-        prepareResnetModel()
-
         captureSession.delegate = requestDelegate
 
         cameraView.layer?.addSublayer(captureSession.previewLayer)
     }
 
-    func prepareResnetModel() {
+    func prepareFaceRequest() {
+        let detectedObjectHandler = VisionRequestResultHandler(delegate: self)
 
+        visionRequests.append(
+            VNDetectFaceRectanglesRequest(
+                completionHandler: detectedObjectHandler.requestResultHandler
+            )
+        )
+    }
+
+    func prepareResnetRequest() {
         guard let resnet = try? VNCoreMLModel(for: Resnet50().model) else {
             fatalError("Failed to load ResNet model")
         }
 
-        resnetDelegate.configure(
-            for: VNCoreMLRequest(model: resnet,
-                                 completionHandler: didReceiveResnetResults),
-            failHandler: { print("Error") }
+        visionRequests.append(
+            VNCoreMLRequest(model: resnet,
+                            completionHandler: didReceiveResnetResults)
         )
-
-        captureSession.delegate = resnetDelegate
     }
 
     override func viewDidLayout() {

diff --git a/Machine/VisionRequestCaptureDelegate.swift b/Machine/VisionRequestCaptureDelegate.swift
@@ -16,36 +16,32 @@ typealias VNRequestSuccesHandler = (VNRequest, Error?) -> ()
 class VisionRequestCaptureDelegate: NSObject,
                                     AVCaptureVideoDataOutputSampleBufferDelegate {
 
-    var request: VNRequest?
+    var requests: [VNRequest]?
 
     var shouldTrack = true
 
     private var failHandler: VNRequestFailHandler?
 
-    private var completionHandler: VNRequestSuccesHandler? {
-        return request?.completionHandler
-    }
-
     static let `default` = VisionRequestCaptureDelegate()
 
     private override init() { }
 
-    init(request: VNRequest,
+    init(requests: [VNRequest],
          failHandler: @escaping VNRequestFailHandler) {
         super.init()
-        configure(for: request, failHandler: failHandler)
+        configure(for: requests, failHandler: failHandler)
     }
 
-    func configure(for request: VNRequest,
+    func configure(for requests: [VNRequest],
                    failHandler: @escaping VNRequestFailHandler) {
-        self.request        = request
+        self.requests       = requests
         self.failHandler    = failHandler
     }
 
     func captureOutput(_ output: AVCaptureOutput,
                        didOutput sampleBuffer: CMSampleBuffer,
                        from connection: AVCaptureConnection) {
-        guard   let request = request,
+        guard   let requests = requests,
                 let failHandler = failHandler else {
                 return
         }
@@ -61,7 +57,7 @@ class VisionRequestCaptureDelegate: NSObject,
                                                         options: [:])
 
         do {
-            try imageRequestHandler.perform([request])
+            try imageRequestHandler.perform(requests)
         } catch {
             print(error)
         }