Skip to content

Commit 5817355

Browse files
john-rockyUltralyticsAssistantasabri97
authored
Fix: Threshold sliders not working for non-detect YOLO tasks (#329)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Abrish <55057228+asabri97@users.noreply.github.com>
1 parent b4e6253 commit 5817355

11 files changed

Lines changed: 205 additions & 33 deletions

File tree

‎android/src/main/kotlin/com/ultralytics/yolo/ObbDetector.kt‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ class ObbDetector(
3232
private val useGpu: Boolean = true,
3333
private val customOptions: Interpreter.Options? = null
3434
) : BasePredictor() {
35+
36+
private var numItemsThreshold = 30
3537

3638
private val interpreterOptions: Interpreter.Options = (customOptions ?: Interpreter.Options()).apply {
3739
// If no custom options provided, use default threads
@@ -155,6 +157,11 @@ class ObbDetector(
155157
.add(CastOp(DataType.FLOAT32))
156158
.build()
157159
}
160+
161+
override fun setNumItemsThreshold(n: Int) {
162+
numItemsThreshold = n
163+
super.setNumItemsThreshold(n)
164+
}
158165

159166
override fun predict(bitmap: Bitmap, origWidth: Int, origHeight: Int, rotateForCamera: Boolean, isLandscape: Boolean): YOLOResult {
160167
t0 = System.nanoTime()
@@ -198,12 +205,15 @@ class ObbDetector(
198205
confidenceThreshold = CONFIDENCE_THRESHOLD,
199206
iouThreshold = IOU_THRESHOLD
200207
)
208+
209+
// Apply numItemsThreshold limit
210+
val limitedDetections = obbDetections.take(numItemsThreshold)
201211

202-
val annotatedImage = drawOBBsOnBitmap(bitmap, obbDetections)
212+
val annotatedImage = drawOBBsOnBitmap(bitmap, limitedDetections)
203213

204214
return YOLOResult(
205215
origShape = Size(origWidth, origHeight),
206-
obb = obbDetections,
216+
obb = limitedDetections,
207217
annotatedImage = annotatedImage,
208218
speed = t2,
209219
fps = if (t4 > 0) 1.0 / t4 else 0.0,

‎android/src/main/kotlin/com/ultralytics/yolo/PoseEstimator.kt‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ class PoseEstimator(
4646
private const val INPUT_SIZE = 640
4747
}
4848

49+
private var numItemsThreshold = 30
50+
4951
private val boxPool = ObjectPool<Box>(MAX_POOL_SIZE) { Box(0, "", 0f, RectF(), RectF()) }
5052
private val keypointsPool = ObjectPool<Keypoints>(MAX_POOL_SIZE) {
5153
Keypoints(
@@ -240,8 +242,11 @@ class PoseEstimator(
240242
origHeight = origHeight
241243
)
242244

243-
val boxes = rawDetections.map { it.box }
244-
val keypointsList = rawDetections.map { it.keypoints }
245+
// Apply numItemsThreshold limit
246+
val limitedDetections = rawDetections.take(numItemsThreshold)
247+
248+
val boxes = limitedDetections.map { it.box }
249+
val keypointsList = limitedDetections.map { it.keypoints }
245250

246251
// val annotatedImage = drawPoseOnBitmap(bitmap, keypointsList, boxes)
247252

@@ -449,6 +454,11 @@ class PoseEstimator(
449454
iouThreshold = iou.toFloat()
450455
super.setIouThreshold(iou)
451456
}
457+
458+
override fun setNumItemsThreshold(n: Int) {
459+
numItemsThreshold = n
460+
super.setNumItemsThreshold(n)
461+
}
452462

453463
override fun getConfidenceThreshold(): Double {
454464
return confidenceThreshold.toDouble()

‎android/src/main/kotlin/com/ultralytics/yolo/Segmenter.kt‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ class Segmenter(
4141
private var maskH = 0
4242
private var maskW = 0
4343
private var maskC = 0
44+
private var numItemsThreshold = 30
4445

4546
// TFLite Interpreter options
4647
private val interpreterOptions = (customOptions ?: Interpreter.Options()).apply {
@@ -169,6 +170,11 @@ class Segmenter(
169170
.add(CastOp(DataType.FLOAT32))
170171
.build()
171172
}
173+
174+
override fun setNumItemsThreshold(n: Int) {
175+
numItemsThreshold = n
176+
super.setNumItemsThreshold(n)
177+
}
172178

173179
override fun predict(bitmap: Bitmap, origWidth: Int, origHeight: Int, rotateForCamera: Boolean, isLandscape: Boolean): YOLOResult {
174180
t0 = System.nanoTime()
@@ -228,8 +234,11 @@ class Segmenter(
228234
iouThreshold = IOU_THRESHOLD
229235
)
230236

237+
// Apply numItemsThreshold limit
238+
val limitedDetections = rawDetections.take(numItemsThreshold)
239+
231240
val boxes = mutableListOf<Box>()
232-
for ((normRect, cls, score, maskCoeffs) in rawDetections) {
241+
for ((normRect, cls, score, maskCoeffs) in limitedDetections) {
233242
// normRect already contains normalized coordinates (0-1)
234243

235244
// Convert to absolute pixel coordinates for xywh
@@ -246,7 +255,7 @@ class Segmenter(
246255
}
247256

248257
val (combinedMask, probMasks) = generateCombinedMaskImage(
249-
detections = rawDetections,
258+
detections = limitedDetections,
250259
protos = output1[0],
251260
maskW = maskW,
252261
maskH = maskH,

‎android/src/main/kotlin/com/ultralytics/yolo/YOLOPlatformView.kt‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,36 @@ class YOLOPlatformView(
279279
override fun onMethodCall(call: MethodCall, result: MethodChannel.Result) {
280280
try {
281281
when (call.method) {
282+
"setConfidenceThreshold" -> {
283+
val threshold = call.argument<Double>("threshold")
284+
if (threshold != null) {
285+
yoloView.setConfidenceThreshold(threshold)
286+
Log.d(TAG, "Confidence threshold updated: $threshold")
287+
result.success(null)
288+
} else {
289+
result.error("invalid_args", "threshold is required", null)
290+
}
291+
}
292+
"setIoUThreshold", "setIouThreshold" -> {
293+
val threshold = call.argument<Double>("threshold")
294+
if (threshold != null) {
295+
yoloView.setIouThreshold(threshold)
296+
Log.d(TAG, "IoU threshold updated: $threshold")
297+
result.success(null)
298+
} else {
299+
result.error("invalid_args", "threshold is required", null)
300+
}
301+
}
302+
"setNumItemsThreshold" -> {
303+
val numItems = call.argument<Int>("numItems")
304+
if (numItems != null) {
305+
yoloView.setNumItemsThreshold(numItems)
306+
Log.d(TAG, "NumItems threshold updated: $numItems")
307+
result.success(null)
308+
} else {
309+
result.error("invalid_args", "numItems is required", null)
310+
}
311+
}
282312
"setThresholds" -> {
283313
val confidence = call.argument<Double>("confidenceThreshold")
284314
val iou = call.argument<Double>("iouThreshold")
@@ -314,6 +344,51 @@ class YOLOPlatformView(
314344
}
315345
}
316346
}
347+
"switchCamera" -> {
348+
yoloView.switchCamera()
349+
Log.d(TAG, "Camera switched")
350+
result.success(null)
351+
}
352+
"setZoomLevel" -> {
353+
val zoomLevel = call.argument<Double>("zoomLevel")
354+
if (zoomLevel != null) {
355+
yoloView.setZoomLevel(zoomLevel.toFloat())
356+
Log.d(TAG, "Zoom level set to: $zoomLevel")
357+
result.success(null)
358+
} else {
359+
result.error("invalid_args", "zoomLevel is required", null)
360+
}
361+
}
362+
"setStreamingConfig" -> {
363+
// Parse streaming config from arguments
364+
val configMap = call.arguments as? Map<*, *>
365+
if (configMap != null) {
366+
val streamConfig = YOLOStreamConfig(
367+
includeDetections = configMap["includeDetections"] as? Boolean ?: true,
368+
includeClassifications = configMap["includeClassifications"] as? Boolean ?: true,
369+
includeProcessingTimeMs = configMap["includeProcessingTimeMs"] as? Boolean ?: true,
370+
includeFps = configMap["includeFps"] as? Boolean ?: true,
371+
includeMasks = configMap["includeMasks"] as? Boolean ?: false,
372+
includePoses = configMap["includePoses"] as? Boolean ?: false,
373+
includeOBB = configMap["includeOBB"] as? Boolean ?: false,
374+
includeOriginalImage = configMap["includeOriginalImage"] as? Boolean ?: false,
375+
maxFPS = (configMap["maxFPS"] as? Number)?.toInt(),
376+
throttleIntervalMs = (configMap["throttleInterval"] as? Number)?.toInt(),
377+
inferenceFrequency = (configMap["inferenceFrequency"] as? Number)?.toInt(),
378+
skipFrames = (configMap["skipFrames"] as? Number)?.toInt()
379+
)
380+
yoloView.setStreamConfig(streamConfig)
381+
Log.d(TAG, "Streaming config updated")
382+
result.success(null)
383+
} else {
384+
result.error("invalid_args", "Invalid streaming config", null)
385+
}
386+
}
387+
"stop" -> {
388+
yoloView.stop()
389+
Log.d(TAG, "Camera and inference stopped")
390+
result.success(null)
391+
}
317392
"captureFrame" -> {
318393
val imageData = yoloView.captureFrame()
319394
if (imageData != null) {

‎android/src/main/kotlin/com/ultralytics/yolo/YOLOView.kt‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -297,17 +297,17 @@ class YOLOView @JvmOverloads constructor(
297297

298298
fun setConfidenceThreshold(conf: Double) {
299299
confidenceThreshold = conf
300-
(predictor as? ObjectDetector)?.setConfidenceThreshold(conf)
300+
predictor?.setConfidenceThreshold(conf)
301301
}
302302

303303
fun setIouThreshold(iou: Double) {
304304
iouThreshold = iou
305-
(predictor as? ObjectDetector)?.setIouThreshold(iou)
305+
predictor?.setIouThreshold(iou)
306306
}
307307

308308
fun setNumItemsThreshold(n: Int) {
309309
numItemsThreshold = n
310-
(predictor as? ObjectDetector)?.setNumItemsThreshold(n)
310+
predictor?.setNumItemsThreshold(n)
311311
}
312312

313313
fun setZoomLevel(zoomLevel: Float) {
@@ -331,16 +331,19 @@ class YOLOView @JvmOverloads constructor(
331331
Executors.newSingleThreadExecutor().execute {
332332
try {
333333
val newPredictor = when (task) {
334-
YOLOTask.DETECT -> ObjectDetector(context, modelPath, loadLabels(modelPath), useGpu = useGpu).apply {
335-
setConfidenceThreshold(confidenceThreshold)
336-
setIouThreshold(iouThreshold)
337-
setNumItemsThreshold(numItemsThreshold)
338-
}
334+
YOLOTask.DETECT -> ObjectDetector(context, modelPath, loadLabels(modelPath), useGpu = useGpu)
339335
YOLOTask.SEGMENT -> Segmenter(context, modelPath, loadLabels(modelPath), useGpu = useGpu)
340336
YOLOTask.CLASSIFY -> Classifier(context, modelPath, loadLabels(modelPath), useGpu = useGpu)
341337
YOLOTask.POSE -> PoseEstimator(context, modelPath, loadLabels(modelPath), useGpu = useGpu)
342338
YOLOTask.OBB -> ObbDetector(context, modelPath, loadLabels(modelPath), useGpu = useGpu)
343339
}
340+
341+
// Apply thresholds to all predictor types
342+
newPredictor.apply {
343+
setConfidenceThreshold(confidenceThreshold)
344+
setIouThreshold(iouThreshold)
345+
setNumItemsThreshold(numItemsThreshold)
346+
}
344347

345348
post {
346349
this.task = task

‎ios/Classes/Classifier.swift‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -85,10 +85,7 @@ class Classifier: BasePredictor {
8585
var top5: [String] = []
8686
var top5Confs: [Float] = []
8787

88-
var candidateNumber = 5
89-
if observations.count < candidateNumber {
90-
candidateNumber = observations.count
91-
}
88+
var candidateNumber = min(5, observations.count)
9289
if let topObservation = observations.first {
9390
top1 = topObservation.identifier
9491
top1Conf = Float(topObservation.confidence)

‎ios/Classes/ObbDetector.swift‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,33 @@ import Vision
2121
/// Specialized predictor for YOLO models that detect objects using oriented (rotated) bounding boxes.
2222
class ObbDetector: BasePredictor, @unchecked Sendable {
2323

24+
override func setConfidenceThreshold(confidence: Double) {
25+
confidenceThreshold = confidence
26+
}
27+
28+
override func setIouThreshold(iou: Double) {
29+
iouThreshold = iou
30+
}
31+
32+
override func setNumItemsThreshold(numItems: Int) {
33+
numItemsThreshold = numItems
34+
}
35+
2436
override func processObservations(for request: VNRequest, error: Error?) {
2537
if let results = request.results as? [VNCoreMLFeatureValueObservation] {
2638

2739
if let prediction = results.first?.featureValue.multiArrayValue {
2840
let nmsResults = postProcessOBB(
2941
feature: prediction, // your MLMultiArray
30-
confidenceThreshold: 0.25,
31-
iouThreshold: 0.45
42+
confidenceThreshold: Float(confidenceThreshold),
43+
iouThreshold: Float(iouThreshold)
3244
)
3345

3446
var obbResults: [OBBResult] = []
35-
for result in nmsResults {
47+
// Apply numItemsThreshold limit
48+
let limitedResults = Array(nmsResults.prefix(numItemsThreshold))
49+
50+
for result in limitedResults {
3651
let box = result.box
3752
let score = result.score
3853
let clsIdx = labels[result.cls]

‎ios/Classes/PoseEstimater.swift‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,18 @@ import Vision
2222
class PoseEstimater: BasePredictor, @unchecked Sendable {
2323
var colorsForMask: [(red: UInt8, green: UInt8, blue: UInt8)] = []
2424

25+
override func setConfidenceThreshold(confidence: Double) {
26+
confidenceThreshold = confidence
27+
}
28+
29+
override func setIouThreshold(iou: Double) {
30+
iouThreshold = iou
31+
}
32+
33+
override func setNumItemsThreshold(numItems: Int) {
34+
numItemsThreshold = numItems
35+
}
36+
2537
override func processObservations(for request: VNRequest, error: Error?) {
2638
if let results = request.results as? [VNCoreMLFeatureValueObservation] {
2739

@@ -33,7 +45,10 @@ class PoseEstimater: BasePredictor, @unchecked Sendable {
3345
var keypointsList = [Keypoints]()
3446
var boxes = [Box]()
3547

36-
for person in preds {
48+
// Apply numItemsThreshold limit
49+
let limitedPreds = Array(preds.prefix(numItemsThreshold))
50+
51+
for person in limitedPreds {
3752
boxes.append(person.box)
3853
keypointsList.append(person.keypoints)
3954
}
@@ -92,7 +107,10 @@ class PoseEstimater: BasePredictor, @unchecked Sendable {
92107
var keypointsForImage = [[(x: Float, y: Float)]]()
93108
var confsList: [[Float]] = []
94109

95-
for person in preds {
110+
// Apply numItemsThreshold limit
111+
let limitedPreds = Array(preds.prefix(numItemsThreshold))
112+
113+
for person in limitedPreds {
96114
boxes.append(person.box)
97115
keypointsList.append(person.keypoints)
98116
keypointsForImage.append(person.keypoints.xyn)

‎ios/Classes/Segmenter.swift‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,18 @@ import Vision
2222
class Segmenter: BasePredictor, @unchecked Sendable {
2323
var colorsForMask: [(red: UInt8, green: UInt8, blue: UInt8)] = []
2424

25+
override func setConfidenceThreshold(confidence: Double) {
26+
confidenceThreshold = confidence
27+
}
28+
29+
override func setIouThreshold(iou: Double) {
30+
iouThreshold = iou
31+
}
32+
33+
override func setNumItemsThreshold(numItems: Int) {
34+
numItemsThreshold = numItems
35+
}
36+
2537
override func processObservations(for request: VNRequest, error: Error?) {
2638
if let results = request.results as? [VNCoreMLFeatureValueObservation] {
2739
// DispatchQueue.main.async { [self] in
@@ -46,7 +58,10 @@ class Segmenter: BasePredictor, @unchecked Sendable {
4658
var boxes: [Box] = []
4759
var alphas = [CGFloat]()
4860

49-
for p in detectedObjects {
61+
// Apply numItemsThreshold limit
62+
let limitedDetections = Array(detectedObjects.prefix(numItemsThreshold))
63+
64+
for p in limitedDetections {
5065
let box = p.0
5166
let rect = CGRect(
5267
x: box.minX / 640, y: box.minY / 640, width: box.width / 640, height: box.height / 640)
@@ -65,7 +80,7 @@ class Segmenter: BasePredictor, @unchecked Sendable {
6580
DispatchQueue.global(qos: .userInitiated).async {
6681
guard
6782
let procceessedMasks = generateCombinedMaskImage(
68-
detectedObjects: detectedObjects,
83+
detectedObjects: limitedDetections,
6984
protos: masks,
7085
inputWidth: self.modelInputSize.width,
7186
inputHeight: self.modelInputSize.height,

0 commit comments

Comments
 (0)