iosswiftclasscoremlmlmodel

Create a model using Create ML and find classes in metaData


I am using Create ML to create a .mlmodel file using a Image Classifier project. Create ML, creates this file for 2 classes.

After dragging the created ML Model file , and adding it to compile sources, below code is able to find the file and create the model.

But the meta data does not show classes, any reason why? How would I know the classes in the mlmodel file?

let error: NSError! = nil
        
        guard let modelURL = Bundle.main.url(forResource: "ObjectDetector", withExtension: "mlmodelc") else {
            return 
        }
        do {
            let model = try MLModel(contentsOf: modelURL)
            let visionModel = try VNCoreMLModel(for: model)
            let metaData =  model.modelDescription.metadata[.creatorDefinedKey] as! [String:String]
             // ERROR: It  could not find metaData["classes"]
            let allClasses = metaData["classes"]!.components(separatedBy: ",")


            let objectRecognition = VNCoreMLRequest(model: visionModel, completionHandler: { (request, error) in
                DispatchQueue.main.async(execute: {
                    // perform all the UI updates on the main queue
                    if let results = request.results {
                        self.drawVisionRequestResults(results)
                    }
                })
            })
            self.requests = [objectRecognition]
        } catch let error as NSError {
            print("Model loading went wrong: \(error)")
        }

Please note I am not getting "Model class has not been generated yet.". https://stackoverflow.com/questions/462476…

Clicking on ML Model file shows Automatically generated Swift model class.


Solution

  • The class data is part of the mlmodel file but not in the metadata. You can add it to the metadata if you want, but you'll have to write a short Python script for that.

    As of iOS 14, there is MLModelDescription.classLabels that also lets you access the list of class names.