Custom Models with Create ML

In Lesson 1, you laid the groundwork for the MoodTracker app by implementing the basic UI for emotion detection. In this demo, you’ll go through the process of integrating the Core ML model into your MoodTracker app to perform emotion detection on images. This includes setting up a view model, configuring the classifier, and updating the user interface to display the results.

In the starter folder, you’ll find the MoodTracker app as you left it in Lesson 1, and you’ll find the Create ML project containing the three classifiers you created in Lesson 2. First, you’ll extract the .mlmodel file from the Create ML project. The CreateML project contains data sources saved as absolute links, which occurs when you drag and drop them. To work with this model, you’ll need to update the paths in Create ML to match your local setup.

Open the EmotionsImageClassifier project and, in the Model Sources section, choose the second classifier you configured, which has the best accuracy among the model sources. Then, open the Output tab. Next, press the Get button to export the model. When saving the file, name it EmotionsImageClassifier to ensure it matches the instructions. Now, you have the model ready to use in your project in the Core ML extension.

Now, open the MoodTracker app. In this demo, you’ll introduce more functionality to the EmotionDetectionView. That’s why you’ll create a view model for this view to handle logic and functionalities in it. Create a new folder named ViewModel, then add a new Swift file named EmotionDetectionViewModel. This view model will hold only a property for the image and the reset method to reset the image.

import SwiftUI
import Combine

class EmotionDetectionViewModel: ObservableObject {
  @Published var image: UIImage?

  func reset() {
    DispatchQueue.main.async {
      self.image = nil
    }
  }
}

Now, open the EmotionDetectionView and replace the image property with viewModel property of our newly created EmotionDetectionViewModel type. Then, replace each representation of $image with $viewModel.image. Also, change reset to viewModel.reset.

@StateObject private var viewModel = EmotionDetectionViewModel()

Build and run your app. Press the Start Emotion Detection button, then select an image and ensure that the image appears as it did previously. Then, press the Select Another Image button to verify that the reset functionality works as expected. Now, you’re ready to start integrating your model into the app.

Drag and drop the EmotionsImageClassifier.mlmodel file into the Helper folder. Next, create a swift file in the same folder and name it EmotionClassifier. In this file, you’ll create the classifier as you learned in the Instruction section of this lesson.

You’ll load the Core ML model in the initializer. Then, the classify method will convert the UIImage to CIImage. Next, you’ll create a VNCoreMLRequest with the model. After that, you’ll handle the classification results and find the top one. Finally, you’ll create a handler and perform this request in a background thread as a best practice, as you learned in the previous section. If you need more details about any of these steps in the classifier, you can refer back to the Instruction section to review them.

import SwiftUI
import Vision
import CoreML

class EmotionClassifier {
  private let model: VNCoreMLModel

  init() {
    // 1. Load the Core ML model
    let configuration = MLModelConfiguration()
    guard let mlModel = try? EmotionsImageClassifier(configuration: configuration).model else {
      fatalError("Failed to load model")
    }
    self.model = try! VNCoreMLModel(for: mlModel)
  }

  func classify(image: UIImage, completion: @escaping (String?, Float?) -> Void) {
    // 2. Convert UIImage to CIImage
    guard let ciImage = CIImage(image: image) else {
      completion(nil, nil)
      return
    }

    // 3. Create a VNCoreMLRequest with the model
    let request = VNCoreMLRequest(model: model) { request, error in
      if let error = error {
        print("Error during classification: \(error.localizedDescription)")
        completion(nil, nil)
        return
      }

      // 4. Handle the classification results
      guard let results = request.results as? [VNClassificationObservation] else {
        print("No results found")
        completion(nil, nil)
        return
      }

      // 5. Find the top result based on confidence
      let topResult = results.max(by: { a, b in a.confidence < b.confidence })
      guard let bestResult = topResult else {
        print("No top result found")
        completion(nil, nil)
        return
      }

      // 6. Pass the top result to the completion handler
      completion(bestResult.identifier, bestResult.confidence)
    }

    // 7. Create a VNImageRequestHandler
    let handler = VNImageRequestHandler(ciImage: ciImage)

    // 8. Perform the request on a background thread
    DispatchQueue.global(qos: .userInteractive).async {
      do {
        try handler.perform([request])
      } catch {
        print("Failed to perform classification: \(error.localizedDescription)")
        completion(nil, nil)
      }
    }
  }
}

Next, open the EmotionDetectionViewModel. Add the classifier property and another two properties to hold the final emotion after classification and the accuracy of this emotion.

@Published var emotion: String?
@Published var accuracy: String?

private let classifier = EmotionClassifier()

Now, create a classifyImage method that resizes the image before classification as you learned is a best practice. Then, use the new EmotionClassifier class to classify the image and update both the emotion and accuracy properties accordingly after completing the classification.

func classifyImage() {
  if let image = self.image {
    // Resize the image before classification
    let resizedImage = resizeImage(image)
    DispatchQueue.global(qos: .userInteractive).async {
      self.classifier.classify(image: resizedImage ?? image) { [weak self] emotion, confidence in
        // Update the published properties on the main thread
        DispatchQueue.main.async {
          self?.emotion = emotion ?? "Unknown"
          self?.accuracy = String(format: "%.2f%%", (confidence ?? 0) * 100.0)
        }
      }
    }
  }
}

private func resizeImage(_ image: UIImage) -> UIImage? {
  UIGraphicsBeginImageContext(CGSize(width: 224, height: 224))
  image.draw(in: CGRect(x: 0, y: 0, width: 224, height: 224))
  let resizedImage = UIGraphicsGetImageFromCurrentImageContext()
  UIGraphicsEndImageContext()
  return resizedImage
}

Next, inside the reset method, reset both the emotion and accuracy properties to ensure that the EmotionResultView is hidden after resetting. Now, it’s time to apply this in your view to allow users to classify the chosen image and display the classification results.

self.emotion = nil
self.accuracy = nil

Create a view in the Views folder and name it EmotionResultView. This view will show users the dominant emotion for the chosen image with its accuracy. Create a simple VStack with two Text views to display the detected emotion and its accuracy.

import SwiftUI

struct EmotionResultView: View {
  let emotion: String
  let accuracy: String

  var body: some View {
    VStack(spacing: 5) {
      Text("Detected Emotion: \(emotion)")
        .font(.title2)
        .padding(.bottom)
      Text("Accuracy: \(accuracy)")
        .font(.subheadline)
        .foregroundColor(.secondary)
    }
    .padding()
    .background(Color.blue.opacity(0.1))
    .cornerRadius(10)
    .shadow(radius: 10)
  }
}

#Preview {
  EmotionResultView(emotion: "Happy", accuracy: "100%")
}

Next, open EmotionDetectionView. Then, add this EmotionResultView inside the body in case there are valid values for the emotion and accuracy properties in the view model.

if let emotion = viewModel.emotion, let accuracy = viewModel.accuracy {
    EmotionResultView(emotion: emotion, accuracy: accuracy)
}

You’re only one step away from finishing. You need to give the user the ability to classify the image after choosing it. Open ActionButtonsView and then add a Button inside the if condition to classify the image. Add the property needed to classify the image in the var classifyImage: () -> Void view. Make sure to fix the preview for this view after the last change.

Button(action: classifyImage) {
  Text("Detect Emotion")
    .font(.headline)
    .padding()
    .frame(maxWidth: .infinity)
    .background(Color.blue)
    .foregroundColor(.white)
    .cornerRadius(10)
}
.padding(.horizontal)

Finally, open the EmotionDetectionView again and add the implementation of this newly added property into the ActionButtonsView in the body.

ActionButtonsView(image: $viewModel.image, reset: viewModel.reset, classifyImage: viewModel.classifyImage)

Build and run the app on a real device to either import an image from your gallery or take a picture of a happy or sad face. Choose an image as you did previously. Notice that you have a new button now named Detect Emotion to classify this image. Press it and notice the results view that appears, showing whether this emotion is more happy or sad with the accuracy. If you try this process on the simulator like here, you’ll get incorrect results. That’s why it’s essential to test it on a real device to obtain accurate data. Try different images with different emotions and notice that the model might make some mistakes in identification, which is acceptable.

Congratulations! You did a great job implementing the MoodTracker app to detect the dominant emotion from an image. Now, you have a real-world app ready to use.