Add yololite coreML

Sources/Roboflow/Classes/Roboflow.swift

@@ -50,6 +50,9 @@ public class RoboflowMobile: NSObject {
         if (modelType.contains("detr") || modelType.contains("rfdetr")) {
             return RFDetrObjectDetectionModel()
         }
+        if (modelType.contains("yololite")) {
+            return RFYoloLiteObjectDetectionModel()
+        }
         return RFObjectDetectionModel()
     }
 

Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift

@@ -0,0 +1,119 @@
+//
+//  RFYoloLiteObjectDetectionModel.swift
+//  Roboflow
+//
+//  Created by AI Assistant
+//
+
+import Foundation
+import CoreML
+import Vision
+
+/// Object detection model that uses YOLOLite for inference.
+///
+/// The CoreML model includes baked-in NMS, so VNCoreML returns
+/// VNRecognizedObjectObservation results directly — same path as YOLOv5.
+/// Runtime thresholds (iouThreshold, confidenceThreshold) are adjustable
+/// via the inherited ThresholdProvider.
+///
+/// Uses `.cpuAndGPU` compute units to avoid Neural Engine fp16 underflow.
+public class RFYoloLiteObjectDetectionModel: RFObjectDetectionModel {
+
+    /// Load the retrieved CoreML model for YOLOLite
+    override func loadMLModel(modelPath: URL, colors: [String: String], classes: [String], environment: [String: Any]) -> Error? {
+        self.colors = colors
+        self.classes = classes
+        self.environment = environment
+        self.modelPath = modelPath
+
+        do {
+            if #available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *) {
+                let config = MLModelConfiguration()
+                // Use CPU+GPU only — the Neural Engine ignores compute_precision=FLOAT32
+                // and runs in fp16, which causes underflow in the decode head.
+                config.computeUnits = .cpuAndGPU
+
+                mlModel = try MLModel(contentsOf: modelPath, configuration: config)
+
+                do {
+                    visionModel = try VNCoreMLModel(for: mlModel)
+                    visionModel.featureProvider = thresholdProvider
+                    let request = VNCoreMLRequest(model: visionModel)
+                    request.imageCropAndScaleOption = .scaleFill
+                    coreMLRequest = request
+                } catch {
+                    print("Error to initialize YOLOLite model: \(error)")
+                }
+            } else {
+                return UnsupportedOSError()
+            }
+        } catch {
+            return error
+        }
+        return nil
+    }
+
+    /// Run image through YOLOLite model and return object detection predictions.
+    ///
+    /// CoreML NMS returns all N entries with suppressed ones zeroed.
+    /// We filter to non-zero observations and use the per-class label
+    /// confidence (in [0, 1]) rather than the observation-level confidence
+    /// which VNCoreML computes as a sum of per-class probabilities.
+    public override func detect(pixelBuffer buffer: CVPixelBuffer, completion: @escaping (([RFPrediction]?, Error?) -> Void)) {
+        guard let coreMLRequest = self.coreMLRequest else {
+            completion(nil, NSError(domain: "RFYoloLiteObjectDetectionModel", code: 1,
+                                    userInfo: [NSLocalizedDescriptionKey: "VNCoreML model initialization failed."]))
+            return
+        }
+
+        let handler = VNImageRequestHandler(cvPixelBuffer: buffer)
+        do {
+            try handler.perform([coreMLRequest])
+
+            guard let detectResults = coreMLRequest.results as? [VNRecognizedObjectObservation] else {
+                completion([], nil)
+                return
+            }
+
+            var detections: [RFObjectDetectionPrediction] = []
+            for result in detectResults {
+                guard result.confidence > 0 else { continue }
+
+                let flippedBox = CGRect(x: result.boundingBox.minX,
+                                        y: 1 - result.boundingBox.maxY,
+                                        width: result.boundingBox.width,
+                                        height: result.boundingBox.height)
+                let box = VNImageRectForNormalizedRect(flippedBox, Int(buffer.width()), Int(buffer.height()))
+
+                var label = ""
+                // VNCoreML normalizes per-class confidences to sum to 1.0.
+                // Multiply observation-level confidence (raw sum) by the normalized
+                // top-label share to recover the original obj * max(cls) score.
+                var confidence: Float = result.confidence
+
+                if let topLabel = result.labels.first {
+                    label = topLabel.identifier
+                    confidence = min(result.confidence * topLabel.confidence, 1.0)
+
+                    if let intValue = Int(label), !classes.contains(label), intValue < classes.count {
+                        label = classes[intValue]
+                    }
+                }
+
+                detections.append(RFObjectDetectionPrediction(
+                    x: Float((box.maxX + box.minX) / 2.0),
+                    y: Float((box.maxY + box.minY) / 2.0),
+                    width: Float(box.maxX - box.minX),
+                    height: Float(box.maxY - box.minY),
+                    className: label,
+                    confidence: confidence,
+                    color: hexStringToCGColor(hex: colors[label] ?? "#ff0000"),
+                    box: box
+                ))
+            }
+            completion(detections, nil)
+        } catch {
+            completion(nil, error)
+        }
+    }
+}

Tests/RoboflowTests/ObjectDetectionTests.swift

@@ -148,7 +148,91 @@ final class ObjectDetectionTests: XCTestCase {
     #endif
 
 
-
+
+    // MARK: - YOLOLite Model Tests
+
+    func testLoadYOLOLiteModel() async {
+        guard let model = await TestUtils.loadYOLOLiteModel() else {
+            XCTFail("Failed to load YOLOLite model")
+            return
+        }
+
+        model.configure(threshold: 0.4, overlap: 0.5, maxObjects: 100)
+        XCTAssertNotNil(model, "YOLOLite model should load successfully")
+    }
+
+    func testYOLOLiteInference() async {
+        guard let model = await TestUtils.loadYOLOLiteModel() else {
+            XCTFail("Failed to load YOLOLite model")
+            return
+        }
+
+        model.configure(threshold: 0.4, overlap: 0.5, maxObjects: 100)
+
+        guard let buffer = TestUtils.loadImageAsPixelBuffer(from: "Tests/assets/hard-hat.jpeg") else {
+            XCTFail("Failed to load hard-hat test image")
+            return
+        }
+
+        let (predictions, inferenceError) = await model.detect(pixelBuffer: buffer)
+        XCTAssertNil(inferenceError, "YOLOLite inference failed: \(inferenceError?.localizedDescription ?? "unknown error")")
+        XCTAssertNotNil(predictions, "Predictions should not be nil")
+        XCTAssertGreaterThan(predictions?.count ?? 0, 0, "Should have at least one detection")
+
+        if let predictions = predictions {
+            for prediction in predictions {
+                guard let objPrediction = prediction as? RFObjectDetectionPrediction else {
+                    XCTFail("Prediction should be of type RFObjectDetectionPrediction")
+                    continue
+                }
+
+                XCTAssertFalse(objPrediction.className.isEmpty, "Class name should not be empty")
+                XCTAssertGreaterThanOrEqual(objPrediction.confidence, 0.0, "Confidence should be >= 0")
+                XCTAssertLessThanOrEqual(objPrediction.confidence, 1.0, "Confidence should be <= 1")
+                XCTAssertGreaterThan(objPrediction.width, 0, "Width should be > 0")
+                XCTAssertGreaterThan(objPrediction.height, 0, "Height should be > 0")
+            }
+        }
+    }
+
+    #if canImport(UIKit)
+    func testYOLOLiteUIImageInference() async {
+        guard let model = await TestUtils.loadYOLOLiteModel() else {
+            XCTFail("Failed to load YOLOLite model")
+            return
+        }
+
+        model.configure(threshold: 0.4, overlap: 0.5, maxObjects: 100)
+
+        guard let image = TestUtils.loadUIImage(from: "Tests/assets/hard-hat.jpeg") else {
+            XCTFail("Failed to load hard-hat test image as UIImage")
+            return
+        }
+
+        let (predictions, inferenceError) = await model.detect(image: image)
+
+        XCTAssertNil(inferenceError, "UIImage YOLOLite inference failed: \(inferenceError?.localizedDescription ?? "unknown error")")
+        XCTAssertNotNil(predictions, "Predictions should not be nil")
+
+        if let predictions = predictions {
+            print("YOLOLite UIImage detected \(predictions.count) objects")
+
+            for prediction in predictions {
+                guard let objPrediction = prediction as? RFObjectDetectionPrediction else {
+                    XCTFail("Prediction should be of type RFObjectDetectionPrediction")
+                    continue
+                }
+
+                XCTAssertFalse(objPrediction.className.isEmpty, "Class name should not be empty")
+                XCTAssertGreaterThanOrEqual(objPrediction.confidence, 0.0, "Confidence should be >= 0")
+                XCTAssertLessThanOrEqual(objPrediction.confidence, 1.0, "Confidence should be <= 1")
+
+                print("YOLOLite detected: \(objPrediction.className) with confidence \(objPrediction.confidence)")
+            }
+        }
+    }
+    #endif
+
     // MARK: - RFDetr Model Tests
 
     func testLoadRFDetrModel() async {

Tests/RoboflowTests/TestUtils.swift

@@ -146,6 +146,20 @@ public class TestUtils {
         return model
     }
 
+    // Helper function to load YOLOLite model from API
+    public static func loadYOLOLiteModel() async -> RFModel? {
+        let rf = RoboflowMobile(apiKey: API_KEY)
+        let (model, error, _, _) = await rf.load(model: "hard-hat-workers-ry4vt", modelVersion: 4)
+
+        if let error = error {
+            XCTFail("Failed to load YOLOLite model: \(error.localizedDescription)")
+            return nil
+        }
+
+        XCTAssertNotNil(model, "YOLOLite model should not be nil")
+        return model
+    }
+
     // Helper function to load RFDetr model from API
     public static func loadRFDetrModel() async -> RFModel? {
         let rf = RoboflowMobile(apiKey: API_KEY)