move files around and update sigmoid/softmax to allow for correct multiclass model predictions

Author: MadeWithStone

.gitignore

@@ -90,3 +90,6 @@ fastlane/test_output
 iOSInjectionProject/
 
 .DS_Store
+
+**/*.mlpackage
+**/*.mlmodel

README.md

@@ -178,7 +178,11 @@ If you've previously installed the Roboflow SDK via Cocoapods, you'll need to up
 The SDK includes a comprehensive test suite that validates model loading and inference functionality. To run the tests:
 
 ```bash
+# for swift only tests
 swift test
+
+# for iOS simulator tests
+xcodebuild test -scheme RoboflowTests -destination 'platform=iOS Simulator,arch=arm64,OS=18.5,name=iPhone 16'
 ```
 
 The test suite includes:

Sources/Roboflow/Classes/Roboflow.swift

@@ -58,7 +58,8 @@ public class RoboflowMobile: NSObject {
             let colors = modelInfo["colors"] as? [String: String],
             let classes = modelInfo["classes"] as? [String],
             let name = modelInfo["name"] as? String,
-            let modelType = modelInfo["modelType"] as? String {
+            let modelType = modelInfo["modelType"] as? String,
+            let environment = modelInfo["environment"] as? [String: Any] {
 
             getConfigDataBackground(modelName: model, modelVersion: modelVersion, apiKey: apiKey, deviceID: deviceID)
 
@@ -69,7 +70,7 @@ public class RoboflowMobile: NSObject {
                                                                 in: .userDomainMask,
                                                                 appropriateFor: nil,
                                                                 create: false)
-                _ = modelObject.loadMLModel(modelPath: documentsURL.appendingPathComponent(modelURL), colors: colors, classes: classes)
+                _ = modelObject.loadMLModel(modelPath: documentsURL.appendingPathComponent(modelURL), colors: colors, classes: classes, environment: environment)
 
                 completion(modelObject, nil, name, modelType)
             } catch {
@@ -78,15 +79,14 @@ public class RoboflowMobile: NSObject {
         } else if retries > 0 {
             clearModelCache(modelName: model, modelVersion: modelVersion)
             retries -= 1
-            getModelData(modelName: model, modelVersion: modelVersion, apiKey: apiKey, deviceID: deviceID) { [self] fetchedModel, error, modelName, modelType, colors, classes in
+            getModelData(modelName: model, modelVersion: modelVersion, apiKey: apiKey, deviceID: deviceID) { [self] fetchedModel, error, modelName, modelType, colors, classes, environment in
                 if let err = error {
                     completion(nil, err, "", "")
                 } else if let fetchedModel = fetchedModel {
                     let modelObject = getModelClass(modelType: modelType)
-                    _ = modelObject.loadMLModel(modelPath: fetchedModel, colors: colors ?? [:], classes: classes ?? [])
+                    _ = modelObject.loadMLModel(modelPath: fetchedModel, colors: colors ?? [:], classes: classes ?? [], environment: environment ?? [:])
                     completion(modelObject, nil, modelName, modelType)
                 } else {
-                    print("No Model Found. Trying Again.")
                     clearAndRetryLoadingModel(model, modelVersion, completion)
                 }
             }
@@ -138,7 +138,6 @@ public class RoboflowMobile: NSObject {
                 completion(dict, nil)
 
             } catch {
-                print(error.localizedDescription)
                 completion(nil, error.localizedDescription)
             }
         }).resume()
@@ -152,10 +151,10 @@ public class RoboflowMobile: NSObject {
 
 
     //Get the model metadata from the Roboflow API
-    private func getModelData(modelName: String, modelVersion: Int, apiKey: String, deviceID: String, completion: @escaping (URL?, Error?, String, String, [String: String]?, [String]?)->()) {
+    private func getModelData(modelName: String, modelVersion: Int, apiKey: String, deviceID: String, completion: @escaping (URL?, Error?, String, String, [String: String]?, [String]?, [String: Any]?)->()) {
         getConfigData(modelName: modelName, modelVersion: modelVersion, apiKey: apiKey, deviceID: deviceID) { data, error in
             if let error = error {
-                completion(nil, error, "", "", nil, nil)
+                completion(nil, error, "", "", nil, nil, nil)
                 return
             }
 
@@ -165,7 +164,7 @@ public class RoboflowMobile: NSObject {
                   let modelType = coreMLDict["modelType"] as? String,
                   let modelURLString = coreMLDict["model"] as? String,
                   let modelURL = URL(string: modelURLString) else {
-                completion(nil, error, "", "", nil, nil)
+                completion(nil, error, "", "", nil, nil, nil)
                 return
             }
 
@@ -188,28 +187,29 @@ public class RoboflowMobile: NSObject {
             //Download the model from the link in the API response
             self.downloadModelFile(modelName: "\(modelName)-\(modelVersion).mlmodel", modelVersion: modelVersion, modelURL: modelURL) { fetchedModel, error in
                 if let error = error {
-                    completion(nil, error, "", "", nil, nil)
+                    completion(nil, error, "", "", nil, nil, nil)
                     return
                 }
 
                 if let fetchedModel = fetchedModel {
-                    _ = self.cacheModelInfo(modelName: modelName, modelVersion: modelVersion, colors: colors ?? [:], classes: classes ?? [], name: name, modelType: modelType, compiledModelURL: fetchedModel)
-                    completion(fetchedModel, nil, name, modelType, colors, classes)
+                    _ = self.cacheModelInfo(modelName: modelName, modelVersion: modelVersion, colors: colors ?? [:], classes: classes ?? [], name: name, modelType: modelType, compiledModelURL: fetchedModel, environment: environmentDict ?? [:])
+                    completion(fetchedModel, nil, name, modelType, colors, classes, environmentDict)
                 } else {
-                    completion(nil, error, "", "", nil, nil)
+                    completion(nil, error, "", "", nil, nil, nil)
                 }
             }
         }
     }
 
 
-    private func cacheModelInfo(modelName: String, modelVersion: Int, colors: [String: String], classes: [String], name: String, modelType: String, compiledModelURL: URL) -> [String: Any]? {
+    private func cacheModelInfo(modelName: String, modelVersion: Int, colors: [String: String], classes: [String], name: String, modelType: String, compiledModelURL: URL, environment: [String: Any]) -> [String: Any]? {
         let modelInfo: [String : Any] = [
             "colors": colors,
             "classes": classes,
             "name": name,
             "modelType": modelType,
-            "compiledModelURL": compiledModelURL.lastPathComponent
+            "compiledModelURL": compiledModelURL.lastPathComponent,
+            "environment": environment
         ]
 
         do {
@@ -227,8 +227,6 @@ public class RoboflowMobile: NSObject {
             if let modelInfoData = UserDefaults.standard.data(forKey: "\(modelName)-\(modelVersion)") {
                 let decodedData = try NSKeyedUnarchiver.unarchivedObject(ofClasses: [NSDictionary.self, NSString.self, NSArray.self], from: modelInfoData) as? [String: Any]
                 return decodedData
-            } else {
-                print("Error: Could not find data for key \(modelName)-\(modelVersion)")
             }
         } catch {
             print("Error unarchiving data: \(error.localizedDescription)")
@@ -257,13 +255,9 @@ public class RoboflowMobile: NSObject {
                         finalModelURL = try self.unzipModelFile(zipURL: finalModelURL)
                     }
 
-                    print("Compiling model at: \(finalModelURL)")
-                    
                     //Compile the downloaded model
                     let compiledModelURL = try MLModel.compileModel(at: finalModelURL)
 
-                    print("Model compiled to: \(compiledModelURL)")
-                    
                     // Ensure Documents directory exists
                     let documentsURL = try FileManager.default.url(for: .documentDirectory,
                                                                     in: .userDomainMask,
@@ -272,13 +266,9 @@ public class RoboflowMobile: NSObject {
 
                     let savedURL = documentsURL.appendingPathComponent("\(modelName)-\(modelVersion).mlmodelc")
 
-                    print("Attempting to move from: \(compiledModelURL)")
-                    print("To: \(savedURL)")
-                    
                     // Check if the compiled model exists
                     guard FileManager.default.fileExists(atPath: compiledModelURL.path) else {
                         let error = NSError(domain: "ModelCompilationError", code: 1, userInfo: [NSLocalizedDescriptionKey: "Compiled model does not exist at: \(compiledModelURL.path)"])
-                        print("Error: \(error.localizedDescription)")
                         completion(nil, error)
                         return
                     }
@@ -291,15 +281,11 @@ public class RoboflowMobile: NSObject {
                     // Move the compiled model
                     do {
                         try FileManager.default.moveItem(at: compiledModelURL, to: savedURL)
-                        print("Successfully moved model to: \(savedURL)")
                     } catch {
-                        print("Move failed with error: \(error.localizedDescription)")
                         // If move fails, try copying instead
                         do {
                             try FileManager.default.copyItem(at: compiledModelURL, to: savedURL)
-                            print("Successfully copied model to: \(savedURL)")
                         } catch {
-                            print("Copy also failed: \(error.localizedDescription)")
                             completion(nil, error)
                             return
                         }

Sources/Roboflow/Classes/ThresholdProvider.swift renamed to Sources/Roboflow/Classes/Utils/ThresholdProvider.swift

@@ -17,60 +17,19 @@ public class RFClassificationModel: RFModel {
     }
 
     //Default model configuration parameters
-    var threshold: Double = 0.5
-    var classes: [String] = []
-    
-    //Configure the parameters for the model
-    public override func configure(threshold: Double, overlap: Double, maxObjects: Float, processingMode: ProcessingMode = .balanced, maxNumberPoints: Int = 500) {
-        self.threshold = threshold
-    }
+    var multiclass: Bool = false
 
     //Load the retrieved CoreML model into an already created RFClassificationModel instance
-    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String]) -> Error? {
-        self.classes = classes
-        do {
-            if #available(macOS 10.14, *) {
-                let config = MLModelConfiguration()
-                if #available(macOS 10.15, *) {
-                    mlModel = try MLModel(contentsOf: modelPath, configuration: config)
-                } else {
-                    // Fallback on earlier versions
-                    return UnsupportedOSError()
-                }
-                visionModel = try VNCoreMLModel(for: mlModel)
-                let request = VNCoreMLRequest(model: visionModel)
-                request.imageCropAndScaleOption = .scaleFill
-                coreMLRequest = request
-            } else {
-                // Fallback on earlier versions
-                return UnsupportedOSError()
-            }
-            
-        } catch {
-            return error
+    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String], environment: [String: Any]) -> Error? {
+        let _ = super.loadMLModel(modelPath: modelPath, colors: colors, classes: classes, environment: environment)
+        if let _ = environment["MULTICLASS"] {
+            self.multiclass = true
         }
-        return nil
-    }
-    
-    //Load a local model file (for manually placed models like ResNet)
-    public func loadLocalModel(modelPath: URL) -> Error? {
         do {
             if #available(macOS 10.14, *) {
                 let config = MLModelConfiguration()
                 if #available(macOS 10.15, *) {
-                    var modelURL = modelPath
-                    
-                    // If the model is .mlpackage, compile it first
-                    if modelPath.pathExtension == "mlpackage" {
-                        do {
-                            let compiledModelURL = try MLModel.compileModel(at: modelPath)
-                            modelURL = compiledModelURL
-                        } catch {
-                            return error
-                        }
-                    }
-                    
-                    mlModel = try MLModel(contentsOf: modelURL, configuration: config)
+                    mlModel = try MLModel(contentsOf: modelPath, configuration: config)
                 } else {
                     // Fallback on earlier versions
                     return UnsupportedOSError()
@@ -141,20 +100,12 @@ public class RFClassificationModel: RFModel {
                 let rawValues = multiArray.dataPointer.bindMemory(to: Float.self, capacity: multiArray.count)
 
                 // Check if values are logits (outside 0-1 range) and need softmax
-                var needsSoftmax = false
-                for i in 0..<multiArray.count {
-                    if rawValues[i] < 0.0 || rawValues[i] > 1.0 {
-                        needsSoftmax = true
-                        break
-                    }
-                }
-                
                 let probabilities: [Float]
-                if needsSoftmax {
+                if !multiclass {
                     // Apply softmax to convert logits to probabilities
                     probabilities = applySoftmax(logits: rawValues, count: multiArray.count)
                 } else {
-                    // Values are already probabilities
+                    // Values are already probabilities (sigmoid applied in model)
                     probabilities = Array(UnsafeBufferPointer(start: rawValues, count: multiArray.count))
                 }
 
@@ -183,7 +134,11 @@ public class RFClassificationModel: RFModel {
                 return pred1.confidence > pred2.confidence
             }
 
-            completion(predictions, nil)
+            if multiclass {
+                completion(predictions, nil)
+            } else {
+                completion(predictions.isEmpty ? [] : [predictions[0]], nil)
+            }
         } catch let error {
             completion(nil, error)
         }

Sources/Roboflow/Classes/UploadResult.swift renamed to Sources/Roboflow/Classes/Utils/UploadResult.swift

@@ -14,21 +14,11 @@ import Accelerate
 
 //Creates an instance of an ML model that's hosted on Roboflow
 public class RFInstanceSegmentationModel: RFObjectDetectionModel {
-    var classes = [String]()
-    var maskProcessingMode: ProcessingMode = .balanced
-    var maskMaxNumberPoints: Int = 500
-    
-    public override func configure(threshold: Double, overlap: Double, maxObjects: Float, processingMode: ProcessingMode = .balanced, maxNumberPoints: Int = 500) {
-        super.configure(threshold: threshold, overlap: overlap, maxObjects: maxObjects, processingMode: processingMode)
-        maskProcessingMode = processingMode
-        maskMaxNumberPoints = maxNumberPoints
-    }
 
 
     //Load the retrieved CoreML model into an already created RFObjectDetectionModel instance
-    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String]) -> Error? {
-        self.colors = colors
-        self.classes = classes
+    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String], environment: [String: Any]) -> Error? {
+        let _ = super.loadMLModel(modelPath: modelPath, colors: colors, classes: classes, environment: environment)
         do {
 
             if #available(iOS 16.0, macOS 13.0, *) {

Sources/Roboflow/Classes/RFClassificationModel.swift renamed to Sources/Roboflow/Classes/core/RFClassificationModel.swift

@@ -26,12 +26,31 @@ public class RFModel: NSObject {
     var mlModel: MLModel!
     var visionModel: VNCoreMLModel!
     var coreMLRequest: VNCoreMLRequest!
-    
+    var environment: [String: Any]!
+    var modelPath: URL!
+    var colors: [String: String]!
+    var classes: [String]!
+    var threshold: Double = 0.5
+    var overlap: Double = 0.4
+    var maxObjects: Float = 20
+    var maskProcessingMode: ProcessingMode = .balanced
+    var maskMaxNumberPoints: Int = 500
+
     //Configure the parameters for the model
-    public func configure(threshold: Double, overlap: Double, maxObjects: Float, processingMode: ProcessingMode = .balanced, maxNumberPoints: Int = 500) {}
+    public func configure(threshold: Double = 0.5, overlap: Double = 0.5, maxObjects: Float = 20, processingMode: ProcessingMode = .balanced, maxNumberPoints: Int = 500) {
+        self.threshold = threshold
+        self.overlap = overlap
+        self.maxObjects = maxObjects
+        self.maskProcessingMode = processingMode
+        self.maskMaxNumberPoints = maxNumberPoints
+    }
 
     //Load the retrieved CoreML model into an already created RFObjectDetectionModel instance
-    func loadMLModel(modelPath: URL, colors: [String: String], classes: [String]) -> Error? {
+    func loadMLModel(modelPath: URL, colors: [String: String], classes: [String], environment: [String: Any]) -> Error? {
+        self.environment = environment
+        self.modelPath = modelPath
+        self.colors = colors
+        self.classes = classes
         return nil
     }
 

Sources/Roboflow/Classes/RFInstanceSegmentationModel.swift renamed to Sources/Roboflow/Classes/core/RFInstanceSegmentationModel.swift

@@ -15,19 +15,12 @@ public class RFObjectDetectionModel: RFModel {
         super.init()
     }
 
-    //Default model configuration parameters
-    var threshold: Double = 0.5
-    var overlap: Double = 0.5
-    var maxObjects: Float = 20.0
-    var colors: [String: String]!
     //Stores the retreived ML model
     var thresholdProvider = ThresholdProvider()
 
     //Configure the parameters for the model
-    public override func configure(threshold: Double, overlap: Double, maxObjects: Float, processingMode: ProcessingMode = .balanced, maxNumberPoints: Int = 500) {
-        self.threshold = threshold
-        self.overlap = overlap
-        self.maxObjects = maxObjects
+    public override func configure(threshold: Double = 0.5, overlap: Double = 0.5, maxObjects: Float = 20, processingMode: ProcessingMode = .balanced, maxNumberPoints: Int = 500) {
+        super.configure(threshold: threshold, overlap: overlap, maxObjects: maxObjects, processingMode: processingMode, maxNumberPoints: maxNumberPoints)
         thresholdProvider.values = ["iouThreshold": MLFeatureValue(double: self.overlap),
                                     "confidenceThreshold": MLFeatureValue(double: self.threshold)]
         if visionModel != nil {
@@ -40,8 +33,8 @@ public class RFObjectDetectionModel: RFModel {
     }
 
     //Load the retrieved CoreML model into an already created RFObjectDetectionModel instance
-    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String]) -> Error? {
-        self.colors = colors
+    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String], environment: [String: Any]) -> Error? {
+        super.loadMLModel(modelPath: modelPath, colors: colors, classes: classes, environment: environment)
         do {
             if #available(macOS 10.14, *) {
                 let config = MLModelConfiguration()