From 3cc1e1556852eb520a3757aaaeb049dbbdeb889e Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Fri, 3 Apr 2026 13:40:04 -0230 Subject: [PATCH 1/3] simplify for baked-in nms --- Sources/Roboflow/Classes/Roboflow.swift | 3 + .../core/RFYoloLiteObjectDetectionModel.swift | 117 ++++++++++++++++++ .../RoboflowTests/ObjectDetectionTests.swift | 86 ++++++++++++- Tests/RoboflowTests/TestUtils.swift | 14 +++ 4 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift diff --git a/Sources/Roboflow/Classes/Roboflow.swift b/Sources/Roboflow/Classes/Roboflow.swift index fb25624..d02835e 100644 --- a/Sources/Roboflow/Classes/Roboflow.swift +++ b/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() } diff --git a/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift b/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift new file mode 100644 index 0000000..d8327cc --- /dev/null +++ b/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift @@ -0,0 +1,117 @@ +// +// 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 = "" + var confidence: Float = result.confidence + + if let topLabel = result.labels.first { + label = topLabel.identifier + // Use per-class confidence (0-1) instead of observation-level (can be >1) + confidence = topLabel.confidence + + 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) + } + } +} diff --git a/Tests/RoboflowTests/ObjectDetectionTests.swift b/Tests/RoboflowTests/ObjectDetectionTests.swift index 18a1bdc..77e2cb1 100644 --- a/Tests/RoboflowTests/ObjectDetectionTests.swift +++ b/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 { diff --git a/Tests/RoboflowTests/TestUtils.swift b/Tests/RoboflowTests/TestUtils.swift index 785edb0..382150b 100644 --- a/Tests/RoboflowTests/TestUtils.swift +++ b/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: 3) + + 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) From f15b6926ce5e8d6d872e78235454d8799557cd09 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Fri, 3 Apr 2026 17:25:05 -0230 Subject: [PATCH 2/3] fix confidence calculation --- .../Classes/core/RFYoloLiteObjectDetectionModel.swift | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift b/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift index d8327cc..fbdc36e 100644 --- a/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift +++ b/Sources/Roboflow/Classes/core/RFYoloLiteObjectDetectionModel.swift @@ -86,12 +86,14 @@ public class RFYoloLiteObjectDetectionModel: RFObjectDetectionModel { 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 - // Use per-class confidence (0-1) instead of observation-level (can be >1) - confidence = topLabel.confidence + confidence = min(result.confidence * topLabel.confidence, 1.0) if let intValue = Int(label), !classes.contains(label), intValue < classes.count { label = classes[intValue] From 2e6a3d62fa9ffefc3fad5563566bf9a295d17c26 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Fri, 3 Apr 2026 17:35:29 -0230 Subject: [PATCH 3/3] update test hard hatmodel version --- Tests/RoboflowTests/TestUtils.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tests/RoboflowTests/TestUtils.swift b/Tests/RoboflowTests/TestUtils.swift index 382150b..6c1129d 100644 --- a/Tests/RoboflowTests/TestUtils.swift +++ b/Tests/RoboflowTests/TestUtils.swift @@ -149,7 +149,7 @@ public class TestUtils { // 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: 3) + 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)")