iOS-bootstrap/Lobe_iOS/PlayViewModel.swift at master · lobe/iOS-bootstrap · GitHub

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
//
//  Copyright © 2020 Microsoft. All rights reserved.
//

import Combine
import SwiftUI
import Vision

enum PlayViewMode {
  case Camera
  case ImagePreview
  case NotLoaded
}

/// View model for the Play View
class PlayViewModel: ObservableObject {
  @Published var classificationLabels: [String]?
  @Published var predictions: [Prediction] = []
  @Published var viewMode: PlayViewMode = PlayViewMode.NotLoaded
  @Published var showImagePicker: Bool = false
  @Published var imageFromPhotoPicker: UIImage?
  @Published var showPredictionView = false
  var captureSessionManager: CaptureSessionManager
  let project: Project
  var imagePredicter: PredictionLayer
  private var disposables = Set<AnyCancellable>()

  init(project: Project) {
    self.project = project
    self.imagePredicter = PredictionLayer(model: project.model)
    self.captureSessionManager = CaptureSessionManager()
    /// Subscribes to two publishers:
    ///     1. `capturedImageOutput` published from `Camera` mode.
    ///     2.  `imageFromPhotoPicker` published from `ImagePreview` mode.
    /// If either of the above publishers emit, we send it's output to the prediction layer for classification results.
    self.$imageFromPhotoPicker
      .merge(with: captureSessionManager.$capturedImageOutput)
      .compactMap { $0 }  // remove non-nill values
      .receive(on: DispatchQueue.global(qos: .userInitiated))
      .sink(receiveValue: { [weak self] image in
        self?.imagePredicter.getPrediction(forImage: image)
      })
      .store(in: &disposables)

    /// Subscribe to classifier results from prediction layer
    self.imagePredicter.$classificationResult
      .receive(on: DispatchQueue.main)
      .sink(receiveValue: {[weak self] classificationResults in
        guard let _classificationResults: [VNClassificationObservation] = classificationResults else {
          self?.classificationLabels = ["Loading..."]
          return
        }

        var predictions: [Prediction] = []
        _classificationResults.forEach { classificationResult in
          predictions.append(Prediction(label: classificationResult.identifier, confidence: classificationResult.confidence))
        }
        self?.predictions = predictions
      })
      .store(in: &disposables)

    /// Update camera session if toggled between view mode.
    self.$viewMode
      .receive(on: DispatchQueue.main)
      .sink(receiveValue: { [weak self] _viewMode in
        if _viewMode == .Camera { self?.captureSessionManager.resetCameraFeed() }
        else { self?.captureSessionManager.tearDown() }
      })
      .store(in: &disposables)

    /// Update if predictions view should appear
    self.$predictions
      .receive(on: DispatchQueue.main)
      .sink(receiveValue: { [weak self] _predictions in
        withAnimation {
          self?.showPredictionView = !_predictions.isEmpty
        }
      })
      .store(in: &disposables)
  }
}