Add threshold configuration methods to YOLO and YOLOView classes

Sources/YOLO/ObjectDetector.swift

@@ -64,24 +64,22 @@ class ObjectDetector: BasePredictor, @unchecked Sendable {
     if let results = request.results as? [VNRecognizedObjectObservation] {
       var boxes = [Box]()
 
-      for i in 0..<100 {
-        if i < results.count && i < self.numItemsThreshold {
-          let prediction = results[i]
-          let invertedBox = CGRect(
-            x: prediction.boundingBox.minX, y: 1 - prediction.boundingBox.maxY,
-            width: prediction.boundingBox.width, height: prediction.boundingBox.height)
-          let imageRect = VNImageRectForNormalizedRect(
-            invertedBox, Int(inputSize.width), Int(inputSize.height))
-
-          // The labels array is a list of VNClassificationObservation objects,
-          // with the highest scoring class first in the list.
-          let label = prediction.labels[0].identifier
-          let index = self.labels.firstIndex(of: label) ?? 0
-          let confidence = prediction.labels[0].confidence
-          let box = Box(
-            index: index, cls: label, conf: confidence, xywh: imageRect, xywhn: invertedBox)
-          boxes.append(box)
-        }
+      for i in 0..<min(results.count, self.numItemsThreshold) {
+        let prediction = results[i]
+        let invertedBox = CGRect(
+          x: prediction.boundingBox.minX, y: 1 - prediction.boundingBox.maxY,
+          width: prediction.boundingBox.width, height: prediction.boundingBox.height)
+        let imageRect = VNImageRectForNormalizedRect(
+          invertedBox, Int(inputSize.width), Int(inputSize.height))
+
+        // The labels array is a list of VNClassificationObservation objects,
+        // with the highest scoring class first in the list.
+        let label = prediction.labels[0].identifier
+        let index = self.labels.firstIndex(of: label) ?? 0
+        let confidence = prediction.labels[0].confidence
+        let box = Box(
+          index: index, cls: label, conf: confidence, xywh: imageRect, xywhn: invertedBox)
+        boxes.append(box)
       }
 
       // Measure FPS
@@ -125,24 +123,22 @@ class ObjectDetector: BasePredictor, @unchecked Sendable {
     do {
       try requestHandler.perform([request])
       if let results = request.results as? [VNRecognizedObjectObservation] {
-        for i in 0..<100 {
-          if i < results.count && i < self.numItemsThreshold {
-            let prediction = results[i]
-            let invertedBox = CGRect(
-              x: prediction.boundingBox.minX, y: 1 - prediction.boundingBox.maxY,
-              width: prediction.boundingBox.width, height: prediction.boundingBox.height)
-            let imageRect = VNImageRectForNormalizedRect(
-              invertedBox, Int(inputSize.width), Int(inputSize.height))
-
-            // The labels array is a list of VNClassificationObservation objects,
-            // with the highest scoring class first in the list.
-            let label = prediction.labels[0].identifier
-            let index = self.labels.firstIndex(of: label) ?? 0
-            let confidence = prediction.labels[0].confidence
-            let box = Box(
-              index: index, cls: label, conf: confidence, xywh: imageRect, xywhn: invertedBox)
-            boxes.append(box)
-          }
+        for i in 0..<min(results.count, self.numItemsThreshold) {
+          let prediction = results[i]
+          let invertedBox = CGRect(
+            x: prediction.boundingBox.minX, y: 1 - prediction.boundingBox.maxY,
+            width: prediction.boundingBox.width, height: prediction.boundingBox.height)
+          let imageRect = VNImageRectForNormalizedRect(
+            invertedBox, Int(inputSize.width), Int(inputSize.height))
+
+          // The labels array is a list of VNClassificationObservation objects,
+          // with the highest scoring class first in the list.
+          let label = prediction.labels[0].identifier
+          let index = self.labels.firstIndex(of: label) ?? 0
+          let confidence = prediction.labels[0].confidence
+          let box = Box(
+            index: index, cls: label, conf: confidence, xywh: imageRect, xywhn: invertedBox)
+          boxes.append(box)
         }
       }
     } catch {

Sources/YOLO/README.md

@@ -250,7 +250,60 @@ class CameraViewController: UIViewController {
 }
 ```
 
-With just a few lines of code, you can integrate real-time, YOLO-based inference into your application’s camera feed. For more advanced use cases, explore the customization options available for these components.
+With just a few lines of code, you can integrate real-time, YOLO-based inference into your application's camera feed. For more advanced use cases, explore the customization options available for these components.
+
+### Threshold Configuration
+
+Both `YOLO` and `YOLOView`/`YOLOCamera` classes provide methods to configure detection thresholds:
+
+#### YOLO Class Methods
+
+```swift
+// Set individual thresholds
+model.setNumItemsThreshold(100)        // Maximum number of detections (default: 30)
+model.setConfidenceThreshold(0.5)      // Confidence threshold 0.0-1.0 (default: 0.25)
+model.setIouThreshold(0.6)             // IoU threshold for NMS 0.0-1.0 (default: 0.4)
+
+// Or set all thresholds at once
+model.setThresholds(numItems: 100, confidence: 0.5, iou: 0.6)
+
+// Get current threshold values
+let numItems = model.getNumItemsThreshold()       // Returns Int?
+let confidence = model.getConfidenceThreshold()   // Returns Double?
+let iou = model.getIouThreshold()                 // Returns Double?
+```
+
+#### YOLOView/YOLOCamera Methods
+
+```swift
+// For UIKit (YOLOView)
+yoloView.setNumItemsThreshold(50)
+yoloView.setConfidenceThreshold(0.3)
+yoloView.setIouThreshold(0.5)
+
+// Or set all at once
+yoloView.setThresholds(numItems: 50, confidence: 0.3, iou: 0.5)
+
+// Get current values
+let numItems = yoloView.getNumItemsThreshold()    // Returns Int
+let confidence = yoloView.getConfidenceThreshold() // Returns Double
+let iou = yoloView.getIouThreshold()              // Returns Double
+
+// For SwiftUI (YOLOCamera) - set during initialization
+YOLOCamera(
+    modelPathOrName: "yolo11n",
+    task: .detect,
+    cameraPosition: .back
+)
+.onAppear {
+}
+```
+
+**Note:**
+
+- `numItemsThreshold`: Controls the maximum number of detections returned
+- `confidenceThreshold`: Filters out detections below this confidence level
+- `iouThreshold`: Controls Non-Maximum Suppression overlap threshold
 
 ## ⚙️ How to Obtain YOLO Core ML Models
 

Sources/YOLO/YOLO.swift

@@ -69,67 +69,86 @@ public class YOLO {
   private func loadModel(
     from modelURL: URL, task: YOLOTask, completion: ((Result<YOLO, Error>) -> Void)?
   ) {
-    func handleSuccess(predictor: Predictor) {
-      self.predictor = predictor
-      completion?(.success(self))
-    }
-
-    func handleFailure(_ error: Error) {
-      print("Failed to load model with error: \(error)")
-      completion?(.failure(error))
+    let handleResult: (Result<BasePredictor, Error>) -> Void = { result in
+      switch result {
+      case .success(let predictor):
+        self.predictor = predictor
+        completion?(.success(self))
+      case .failure(let error):
+        print("Failed to load model with error: \(error)")
+        completion?(.failure(error))
+      }
     }
 
     switch task {
     case .classify:
-      Classifier.create(
-        unwrappedModelURL: modelURL,
-        completion: { result in
-          switch result {
-          case .success(let predictor): handleSuccess(predictor: predictor)
-          case .failure(let error): handleFailure(error)
-          }
-        })
-
+      Classifier.create(unwrappedModelURL: modelURL, completion: handleResult)
     case .segment:
-      Segmenter.create(
-        unwrappedModelURL: modelURL,
-        completion: { result in
-          switch result {
-          case .success(let predictor): handleSuccess(predictor: predictor)
-          case .failure(let error): handleFailure(error)
-          }
-        })
-
+      Segmenter.create(unwrappedModelURL: modelURL, completion: handleResult)
     case .pose:
-      PoseEstimator.create(
-        unwrappedModelURL: modelURL,
-        completion: { result in
-          switch result {
-          case .success(let predictor): handleSuccess(predictor: predictor)
-          case .failure(let error): handleFailure(error)
-          }
-        })
-
+      PoseEstimator.create(unwrappedModelURL: modelURL, completion: handleResult)
     case .obb:
-      ObbDetector.create(
-        unwrappedModelURL: modelURL,
-        completion: { result in
-          switch result {
-          case .success(let predictor): handleSuccess(predictor: predictor)
-          case .failure(let error): handleFailure(error)
-          }
-        })
-
+      ObbDetector.create(unwrappedModelURL: modelURL, completion: handleResult)
     default:
-      ObjectDetector.create(
-        unwrappedModelURL: modelURL,
-        completion: { result in
-          switch result {
-          case .success(let predictor): handleSuccess(predictor: predictor)
-          case .failure(let error): handleFailure(error)
-          }
-        })
+      ObjectDetector.create(unwrappedModelURL: modelURL, completion: handleResult)
+    }
+  }
+
+  // MARK: - Threshold Configuration Methods
+
+  /// Sets the maximum number of detection items to include in results.
+  /// - Parameter numItems: The maximum number of items to include (default is 30).
+  public func setNumItemsThreshold(_ numItems: Int) {
+    (predictor as? BasePredictor)?.setNumItemsThreshold(numItems: numItems)
+  }
+
+  /// Gets the current maximum number of detection items.
+  /// - Returns: The current threshold value, or nil if not applicable.
+  public func getNumItemsThreshold() -> Int? {
+    (predictor as? BasePredictor)?.numItemsThreshold
+  }
+
+  /// Sets the confidence threshold for filtering results.
+  /// - Parameter confidence: The confidence threshold value (0.0 to 1.0, default is 0.25).
+  public func setConfidenceThreshold(_ confidence: Double) {
+    guard (0.0...1.0).contains(confidence) else {
+      print("Warning: Confidence threshold should be between 0.0 and 1.0")
+      return
+    }
+    (predictor as? BasePredictor)?.setConfidenceThreshold(confidence: confidence)
+  }
+
+  /// Gets the current confidence threshold.
+  /// - Returns: The current confidence threshold value, or nil if not applicable.
+  public func getConfidenceThreshold() -> Double? {
+    (predictor as? BasePredictor)?.confidenceThreshold
+  }
+
+  /// Sets the IoU (Intersection over Union) threshold for non-maximum suppression.
+  /// - Parameter iou: The IoU threshold value (0.0 to 1.0, default is 0.4).
+  public func setIouThreshold(_ iou: Double) {
+    guard (0.0...1.0).contains(iou) else {
+      print("Warning: IoU threshold should be between 0.0 and 1.0")
+      return
     }
+    (predictor as? BasePredictor)?.setIouThreshold(iou: iou)
+  }
+
+  /// Gets the current IoU threshold.
+  /// - Returns: The current IoU threshold value, or nil if not applicable.
+  public func getIouThreshold() -> Double? {
+    (predictor as? BasePredictor)?.iouThreshold
+  }
+
+  /// Sets all thresholds at once.
+  /// - Parameters:
+  ///   - numItems: The maximum number of items to include.
+  ///   - confidence: The confidence threshold value (0.0 to 1.0).
+  ///   - iou: The IoU threshold value (0.0 to 1.0).
+  public func setThresholds(numItems: Int? = nil, confidence: Double? = nil, iou: Double? = nil) {
+    numItems.map { setNumItemsThreshold($0) }
+    confidence.map { setConfidenceThreshold($0) }
+    iou.map { setIouThreshold($0) }
   }
 
   public func callAsFunction(_ uiImage: UIImage, returnAnnotatedImage: Bool = true) -> YOLOResult {

Sources/YOLO/YOLOCamera.swift

@@ -65,17 +65,9 @@ struct YOLOViewRepresentable: UIViewRepresentable {
   let onDetection: ((YOLOResult) -> Void)?
 
   func makeUIView(context: Context) -> YOLOView {
-    let finalModelPathOrName: String
-
-    if let modelURL = modelURL {
-      finalModelPathOrName = modelURL.path
-    } else if let modelPathOrName = modelPathOrName {
-      finalModelPathOrName = modelPathOrName
-    } else {
-      fatalError("Either modelPathOrName or modelURL must be provided")
-    }
-
-    return YOLOView(frame: .zero, modelPathOrName: finalModelPathOrName, task: task)
+    let modelPath = modelURL?.path ?? modelPathOrName ?? ""
+    assert(!modelPath.isEmpty, "Either modelPathOrName or modelURL must be provided")
+    return YOLOView(frame: .zero, modelPathOrName: modelPath, task: task)
   }
 
   func updateUIView(_ uiView: YOLOView, context: Context) {