使用创建ML创建模型并在metaData中查找类

问题描述

我正在使用Create ML通过图像分类器项目创建.mlmodel文件。 创建ML,为此文件创建2个类。

拖动创建的ML模型文件并将其添加到编译源后,下面的代码可以找到该文件并创建模型。

但是元数据未显示类,为什么?我怎么知道mlmodel文件中的类?

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)")
        }

请注意,我没有得到“尚未生成Model类。” https://stackoverflow.com/questions/462476…

单击ML模型文件显示自动生成的Swift模型类。

解决方法

类数据是mlmodel文件的一部分,但不在元数据中。您可以根据需要将其添加到元数据中,但是必须为此编写一个简短的Python脚本。

从iOS 14开始,有MLModelDescription.classLabels也可以让您访问类名列表。