Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save createwithswift/bb778e1f8807eff0fddf74e499c022b7 to your computer and use it in GitHub Desktop.
Save createwithswift/bb778e1f8807eff0fddf74e499c022b7 to your computer and use it in GitHub Desktop.

This is sample code to use Core ML models in Swift Playgrounds. For this code snippet, the MobileNetV2 image classification model and YOLOv3 object detection model are used. The Core ML model files and sample projects for usage in Xcode are provided by Apple. The models can be used to classify images or detect dominant objects in a camera frame or image. You can replace these placeholder models with you own custom models to leverage the power of machine learning inside Swift Playgrounds.

For Core ML models to be used in Playground, combined binary files or the model as well as corresponding model classes are needed. The procedure to obtain them is described below. The compiled binary file needs to be places in the Resources folder or the Playground page, the model class can conveniently be places in a Swift file in the Sources folder of the Playground page.

ImageClassificationInPlayground.swift

The Image Classification Playground Page is initialising an UIImage that should be in the Resources folder of the Playground page and the Core ML model instance of MobileNetV2. The image is then resized and converted to a CVPixelBuffer as the required input type for the Core ML model. Finally, the CVPixelBuffer is run through the Core ML model to get a prediction. The resulting classification from the model is printed on the console.

ObjectDetectionInPlayground.swift

The Object Detection Playground Page is initialising an UIImage that should be in the Resources folder of the Playground page and a VNCoreMLModel that uses a Core ML based model as an input partner. For this the Core ML model instance of YOLOv3 is used. Then, a VNCoreMLRequests is created to be used with the VNCoreMLModel. The request is printing the detected objects on the console. The request itself is handled by a VNImageRequestHandler that uses a CVPixelBuffer as a required input. To perform the request the image is resized and converted to a CVPixelBuffer.

UIImage+Extension.swift

With this Extension for UIImage, you can resize an UIImage to a provided size with the resizeImageTo function. This is useful as many machine learning models requires images to be provided in specific dimensions. MobileNetV2 expects square images with 224 x 224 resolution. YOLOv3 expects a resolution of 416 x 416. The convertToBuffer function creates a CVPixelBuffer from an UIImage which is the data type Core ML models such as MobileNetV2 or YOLOv3 are expecting as an input parameter.

For a more detailed elaboration on how this extension works, check out the snippet UIImage+Resize+CVPixelBuffer.

Core ML model binary and corresponding source file

In order to work in Playground, the .mlmodelc compiled Core ML model binary and a corresponding Swift source file containing stub classes for the model are required.

The .mlmodelc is generated when the model (.mlmodel) is added to an Xcode Project once the project is compiled and build. It can be found in the ~/Library/Developer/Xcode/DerivedData where all projects are build by default. In the corresponding projects folder it is located in /Build/Products/Debug-iphonesimulator/NameOfTheApp.app/NameOfTheMachineLearningModerl.mlmodelc. You can access the .app bundle content by control clicking on the .app file and selecting Show Package Content. Copy the .mlmodelc file to the Resources Folder of your Playground page.

createwithswift.com-derived-data

The Swift source file can be found by accessing the Model Class when selecting the .mlmodel file in either Playground or any other Xcode project. The source code can be added in a .swift file to the Playground Sources folder. The swift file should have the exact same name of the model, in this case it would be MobileNetV2.swift or YOLOv3.swift.

createwithswift.com-core-ml-model-class

By default the Swift auto-generated stub sources have internal protection levels. All classes, their properties, functions and initializers in MobileNetV2.swift and YOLOv3.swifthave to be made public for Swift Playground to access the source, as any source that is not part of a page needs to have public accessibility to be usable from a Playground page.

Using Core ML models in Xcode Projects

The example code provided here is intended for use inside Swift Playgrounds. To use Core ML models inside Xcode projects for iPhone, iPad, Apple Watch or Mac, consider the various sample projects provided by Apple, such as

/**
* ImageClassificationInPlayground.swift
* Created by Moritz Philip Recke for Create with Swift on 26 March 2021.
* Reviewed by Tiago Gomes Pereira for Create with Swift on 14 April 2021.
*
* Using a MobileNetV2 Image Classification Core ML model in Swift Playground
*/
import UIKit
import CoreML
/**
* Preparing the image to be used by the model.
*/
// 1. Load the image from the 'Resources' folder.
let newImage = UIImage(named: "lemon.jpg")
// 2. Resize the image to the required input dimension of the Core ML model
// Method from UIImage+Extension.swift
let newSize = CGSize(width: 224, height: 224)
guard let resizedImage = newImage?.resizeImageTo(size: newSize) else {
fatalError("⚠️ The image could not be found or resized.")
}
// 3. Convert the resized image to CVPixelBuffer as it is the required input
// type of the Core ML model. Method from UIImage+Extension.swift
guard let convertedImage = resizedImage.convertToBuffer() else {
fatalError("⚠️ The image could not be converted to CVPixelBugger")
}
/**
* Using the ML model to classify the image.
*/
// 1. Create the ML model instance from the model class in the 'Sources' folder
let mlModel = MobileNetV2()
// 2. Get the prediction output
guard let prediction = try? mlModel.prediction(image: convertedImage) else {
fatalError("⚠️ The model could not return a prediction")
}
/**
* Checking the results of the prediction
*/
let mostLikelyImageCategory = prediction.classLabel
let probabilityOfEachCategory = prediction.classLabelProbs
var highestProbability: Double {
let probabilty = probabilityOfEachCategory[mostLikelyImageCategory] ?? 0.0
let roundedProbability = (probabilty * 100).rounded(.toNearestOrEven)
return roundedProbability
}
print("\(mostLikelyImageCategory): \(highestProbability)%")
/**
* ObjectDetectionInPlayground.swift
* Created by Moritz Philip Recke for Create with Swift on 26 March 2021.
* Reviewed by Tiago Gomes Pereira for Create with Swift on 14 April 2021.
*
* Using a YOLOv3 Object Detection Core ML model with the Vision Framework in Swift Playground.
*/
import UIKit
import CoreML
import Vision
/**
* Preparing the Core ML model container and
* creating the image analysis request.
*/
// 1. Create a container for a Core ML model used with Vision requests.
var visionModel: VNCoreMLModel = {
do {
let modelToBeUsed = YOLOv3().model
return try VNCoreMLModel(for: modelToBeUsed)
} catch {
fatalError("⚠️ Failed to create VNCoreMLModel: \(error)")
}
}()
// 2. Create an image analysis request that uses a Core ML model to process images.
var visionRequest: VNCoreMLRequest = {
// 1. Create the completion handler for the analysis request
var requestCompletionHandler: VNRequestCompletionHandler = { request, error in
if let results = request.results as? [VNRecognizedObjectObservation] {
/**
* In here you do something with the request results.
* In this example we will just list the detected objects.
*/
let detectionConfidenceThreshold: Float = 0.99
for result in results {
let resultDetectionConfidence = result.labels.first?.confidence ?? 0
if resultDetectionConfidence >= detectionConfidenceThreshold {
let detectedObject = result.labels.first?.identifier ?? "Nothing"
let detectedObjectConfidence = result.labels.first?.confidence ?? 0
print("\(detectedObject) detected with \(detectedObjectConfidence) confidence")
} else {
print("The result does not match the confidence threshold.")
}
}
} else {
print("Error while getting the request results.")
}
}
// 2. Create the request with the model container and completion handler
let request = VNCoreMLRequest(model: visionModel,
completionHandler: requestCompletionHandler)
// 3. Inform the Vision algorithm how to scale the input image
request.imageCropAndScaleOption = .scaleFill
return request
}()
/**
* Preparing the image to be used by the model.
*/
// 1. Load the image from the 'Resources' folder.
let newImage = UIImage(named: "lemon.jpg")
// 2. Resize the image to the required input dimension of the Core ML model
// Method from UIImage+Extension.swift
let newSize = CGSize(width: 416, height: 416)
guard let resizedImage = newImage?.resizeImageTo(size: newSize) else {
fatalError("⚠️ The image could not be found or resized.")
}
// 3. Convert the resized image to CVPixelBuffer as it is the required input
// type of the Core ML model. Method from UIImage+Extension.swift
guard let convertedImage = resizedImage.convertToBuffer() else {
fatalError("⚠️ The image could not be converted to CVPixelBugger")
}
/**
* Performing an analysis request over the prepared image.
*/
// 1. Create the handler, which will perform requests on a single image
let handler = VNImageRequestHandler(cvPixelBuffer: convertedImage)
// 2. Performs the image analysis request on the image.
do {
try handler.perform([visionRequest])
} catch {
print("Failed to perform the Vision request: \(error)")
}
// UIImage+Extension.swift
//
// Created by Moritz Philip Recke for Create with Swift on 26 March 2021.
//
import Foundation
import UIKit
extension UIImage {
public func resizeImageTo(size: CGSize) -> UIImage? {
UIGraphicsBeginImageContextWithOptions(size, false, 0.0)
self.draw(in: CGRect(origin: CGPoint.zero, size: size))
let resizedImage = UIGraphicsGetImageFromCurrentImageContext()!
UIGraphicsEndImageContext()
return resizedImage
}
public func convertToBuffer() -> CVPixelBuffer? {
let attributes = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
var pixelBuffer: CVPixelBuffer?
let status = CVPixelBufferCreate(kCFAllocatorDefault, Int(self.size.width), Int(self.size.height), kCVPixelFormatType_32ARGB, attributes, &pixelBuffer)
guard (status == kCVReturnSuccess) else {
return nil
}
CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)
let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
let context = CGContext(data: pixelData, width: Int(self.size.width), height: Int(self.size.height), bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
context?.translateBy(x: 0, y: self.size.height)
context?.scaleBy(x: 1.0, y: -1.0)
UIGraphicsPushContext(context!)
self.draw(in: CGRect(x: 0, y: 0, width: self.size.width, height: self.size.height))
UIGraphicsPopContext()
CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
return pixelBuffer
}
}

MIT License

Copyright (c) 2021 Create with Swift

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment