|
| 1 | +package com.visionutils |
| 2 | + |
| 3 | +import kotlin.math.max |
| 4 | +import kotlin.math.min |
| 5 | + |
| 6 | +/** |
| 7 | + * Result classes for bounding box operations |
| 8 | + */ |
| 9 | +data class ScaleBoxesResult( |
| 10 | + val boxes: List<List<Double>>, |
| 11 | + val format: String, |
| 12 | + val processingTimeMs: Double |
| 13 | +) |
| 14 | + |
| 15 | +data class ClipBoxesResult( |
| 16 | + val boxes: List<List<Double>>, |
| 17 | + val format: String, |
| 18 | + val removedCount: Int, |
| 19 | + val processingTimeMs: Double |
| 20 | +) |
| 21 | + |
| 22 | +data class NMSResult( |
| 23 | + val indices: List<Int>, |
| 24 | + val detections: List<Map<String, Any>>, |
| 25 | + val suppressedCount: Int, |
| 26 | + val processingTimeMs: Double |
| 27 | +) |
| 28 | + |
| 29 | +/** |
| 30 | + * Utility class for bounding box operations |
| 31 | + */ |
| 32 | +object BoundingBoxUtilsAndroid { |
| 33 | + |
| 34 | + /** |
| 35 | + * Convert boxes between formats (xyxy, xywh, cxcywh) |
| 36 | + */ |
| 37 | + fun convertBoxFormat( |
| 38 | + boxes: List<List<Double>>, |
| 39 | + sourceFormat: String, |
| 40 | + targetFormat: String |
| 41 | + ): List<List<Double>> { |
| 42 | + if (sourceFormat == targetFormat) { |
| 43 | + return boxes |
| 44 | + } |
| 45 | + |
| 46 | + return boxes.map { box -> |
| 47 | + if (box.size != 4) return@map box |
| 48 | + |
| 49 | + // First convert to xyxy |
| 50 | + val xyxy = when (sourceFormat) { |
| 51 | + "xywh" -> listOf(box[0], box[1], box[0] + box[2], box[1] + box[3]) |
| 52 | + "cxcywh" -> { |
| 53 | + val halfW = box[2] / 2.0 |
| 54 | + val halfH = box[3] / 2.0 |
| 55 | + listOf(box[0] - halfW, box[1] - halfH, box[0] + halfW, box[1] + halfH) |
| 56 | + } |
| 57 | + else -> box // xyxy |
| 58 | + } |
| 59 | + |
| 60 | + // Then convert from xyxy to target |
| 61 | + when (targetFormat) { |
| 62 | + "xywh" -> listOf(xyxy[0], xyxy[1], xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]) |
| 63 | + "cxcywh" -> { |
| 64 | + val w = xyxy[2] - xyxy[0] |
| 65 | + val h = xyxy[3] - xyxy[1] |
| 66 | + listOf(xyxy[0] + w / 2.0, xyxy[1] + h / 2.0, w, h) |
| 67 | + } |
| 68 | + else -> xyxy // xyxy |
| 69 | + } |
| 70 | + } |
| 71 | + } |
| 72 | + |
| 73 | + /** |
| 74 | + * Scale boxes from one image size to another |
| 75 | + */ |
| 76 | + fun scaleBoxes( |
| 77 | + boxes: List<List<Double>>, |
| 78 | + sourceWidth: Double, |
| 79 | + sourceHeight: Double, |
| 80 | + targetWidth: Double, |
| 81 | + targetHeight: Double, |
| 82 | + format: String, |
| 83 | + clip: Boolean = true |
| 84 | + ): ScaleBoxesResult { |
| 85 | + val startTime = System.nanoTime() |
| 86 | + |
| 87 | + val scaleX = targetWidth / sourceWidth |
| 88 | + val scaleY = targetHeight / sourceHeight |
| 89 | + |
| 90 | + // Convert to xyxy for scaling |
| 91 | + val xyxyBoxes = if (format != "xyxy") { |
| 92 | + convertBoxFormat(boxes, format, "xyxy") |
| 93 | + } else { |
| 94 | + boxes |
| 95 | + } |
| 96 | + |
| 97 | + var scaledBoxes = xyxyBoxes.map { box -> |
| 98 | + if (box.size != 4) return@map box |
| 99 | + listOf(box[0] * scaleX, box[1] * scaleY, box[2] * scaleX, box[3] * scaleY) |
| 100 | + } |
| 101 | + |
| 102 | + // Clip if requested |
| 103 | + if (clip) { |
| 104 | + scaledBoxes = scaledBoxes.map { box -> |
| 105 | + if (box.size != 4) return@map box |
| 106 | + listOf( |
| 107 | + max(0.0, min(targetWidth, box[0])), |
| 108 | + max(0.0, min(targetHeight, box[1])), |
| 109 | + max(0.0, min(targetWidth, box[2])), |
| 110 | + max(0.0, min(targetHeight, box[3])) |
| 111 | + ) |
| 112 | + } |
| 113 | + } |
| 114 | + |
| 115 | + // Convert back to original format if needed |
| 116 | + val finalBoxes = if (format != "xyxy") { |
| 117 | + convertBoxFormat(scaledBoxes, "xyxy", format) |
| 118 | + } else { |
| 119 | + scaledBoxes |
| 120 | + } |
| 121 | + |
| 122 | + val endTime = System.nanoTime() |
| 123 | + val processingTimeMs = (endTime - startTime) / 1_000_000.0 |
| 124 | + |
| 125 | + return ScaleBoxesResult( |
| 126 | + boxes = finalBoxes, |
| 127 | + format = format, |
| 128 | + processingTimeMs = processingTimeMs |
| 129 | + ) |
| 130 | + } |
| 131 | + |
| 132 | + /** |
| 133 | + * Clip boxes to image boundaries |
| 134 | + */ |
| 135 | + fun clipBoxes( |
| 136 | + boxes: List<List<Double>>, |
| 137 | + width: Double, |
| 138 | + height: Double, |
| 139 | + format: String, |
| 140 | + removeInvalid: Boolean = false |
| 141 | + ): ClipBoxesResult { |
| 142 | + val startTime = System.nanoTime() |
| 143 | + |
| 144 | + // Convert to xyxy for clipping |
| 145 | + val xyxyBoxes = if (format != "xyxy") { |
| 146 | + convertBoxFormat(boxes, format, "xyxy") |
| 147 | + } else { |
| 148 | + boxes |
| 149 | + } |
| 150 | + |
| 151 | + var clippedBoxes = xyxyBoxes.map { box -> |
| 152 | + if (box.size != 4) return@map box |
| 153 | + listOf( |
| 154 | + max(0.0, min(width, box[0])), |
| 155 | + max(0.0, min(height, box[1])), |
| 156 | + max(0.0, min(width, box[2])), |
| 157 | + max(0.0, min(height, box[3])) |
| 158 | + ) |
| 159 | + } |
| 160 | + |
| 161 | + var removedCount = 0 |
| 162 | + if (removeInvalid) { |
| 163 | + val validBoxes = clippedBoxes.filter { box -> |
| 164 | + if (box.size != 4) return@filter false |
| 165 | + val boxWidth = box[2] - box[0] |
| 166 | + val boxHeight = box[3] - box[1] |
| 167 | + boxWidth > 0 && boxHeight > 0 |
| 168 | + } |
| 169 | + removedCount = clippedBoxes.size - validBoxes.size |
| 170 | + clippedBoxes = validBoxes |
| 171 | + } |
| 172 | + |
| 173 | + // Convert back to original format if needed |
| 174 | + val finalBoxes = if (format != "xyxy") { |
| 175 | + convertBoxFormat(clippedBoxes, "xyxy", format) |
| 176 | + } else { |
| 177 | + clippedBoxes |
| 178 | + } |
| 179 | + |
| 180 | + val endTime = System.nanoTime() |
| 181 | + val processingTimeMs = (endTime - startTime) / 1_000_000.0 |
| 182 | + |
| 183 | + return ClipBoxesResult( |
| 184 | + boxes = finalBoxes, |
| 185 | + format = format, |
| 186 | + removedCount = removedCount, |
| 187 | + processingTimeMs = processingTimeMs |
| 188 | + ) |
| 189 | + } |
| 190 | + |
| 191 | + /** |
| 192 | + * Calculate IoU between two boxes |
| 193 | + */ |
| 194 | + fun calculateIoU( |
| 195 | + box1: List<Double>, |
| 196 | + box2: List<Double>, |
| 197 | + format: String |
| 198 | + ): Map<String, Double> { |
| 199 | + val startTime = System.nanoTime() |
| 200 | + |
| 201 | + require(box1.size == 4 && box2.size == 4) { "Boxes must have 4 elements" } |
| 202 | + |
| 203 | + // Convert to xyxy |
| 204 | + val xyxy1 = if (format != "xyxy") { |
| 205 | + convertBoxFormat(listOf(box1), format, "xyxy")[0] |
| 206 | + } else { |
| 207 | + box1 |
| 208 | + } |
| 209 | + val xyxy2 = if (format != "xyxy") { |
| 210 | + convertBoxFormat(listOf(box2), format, "xyxy")[0] |
| 211 | + } else { |
| 212 | + box2 |
| 213 | + } |
| 214 | + |
| 215 | + // Calculate intersection |
| 216 | + val x1 = max(xyxy1[0], xyxy2[0]) |
| 217 | + val y1 = max(xyxy1[1], xyxy2[1]) |
| 218 | + val x2 = min(xyxy1[2], xyxy2[2]) |
| 219 | + val y2 = min(xyxy1[3], xyxy2[3]) |
| 220 | + |
| 221 | + val intersectionWidth = max(0.0, x2 - x1) |
| 222 | + val intersectionHeight = max(0.0, y2 - y1) |
| 223 | + val intersectionArea = intersectionWidth * intersectionHeight |
| 224 | + |
| 225 | + // Calculate areas |
| 226 | + val area1 = (xyxy1[2] - xyxy1[0]) * (xyxy1[3] - xyxy1[1]) |
| 227 | + val area2 = (xyxy2[2] - xyxy2[0]) * (xyxy2[3] - xyxy2[1]) |
| 228 | + val unionArea = area1 + area2 - intersectionArea |
| 229 | + |
| 230 | + val iou = if (unionArea > 0) intersectionArea / unionArea else 0.0 |
| 231 | + |
| 232 | + val endTime = System.nanoTime() |
| 233 | + val processingTimeMs = (endTime - startTime) / 1_000_000.0 |
| 234 | + |
| 235 | + return mapOf( |
| 236 | + "iou" to iou, |
| 237 | + "intersection" to intersectionArea, |
| 238 | + "union" to unionArea, |
| 239 | + "processingTimeMs" to processingTimeMs |
| 240 | + ) |
| 241 | + } |
| 242 | + |
| 243 | + /** |
| 244 | + * Apply Non-Maximum Suppression |
| 245 | + */ |
| 246 | + fun nonMaxSuppression( |
| 247 | + detections: List<Map<String, Any>>, |
| 248 | + iouThreshold: Double, |
| 249 | + scoreThreshold: Double, |
| 250 | + maxDetections: Int?, |
| 251 | + format: String |
| 252 | + ): NMSResult { |
| 253 | + val startTime = System.nanoTime() |
| 254 | + |
| 255 | + // Extract boxes and scores from detections |
| 256 | + val boxes = detections.mapNotNull { det -> |
| 257 | + @Suppress("UNCHECKED_CAST") |
| 258 | + det["box"] as? List<Double> |
| 259 | + } |
| 260 | + val scores = detections.mapNotNull { det -> |
| 261 | + (det["score"] as? Number)?.toDouble() |
| 262 | + } |
| 263 | + |
| 264 | + require(boxes.size == scores.size) { "Boxes and scores must have same length" } |
| 265 | + |
| 266 | + val maxDets = maxDetections ?: 100 |
| 267 | + |
| 268 | + // Convert to xyxy for IoU calculation |
| 269 | + val xyxyBoxes = if (format != "xyxy") { |
| 270 | + convertBoxFormat(boxes, format, "xyxy") |
| 271 | + } else { |
| 272 | + boxes |
| 273 | + } |
| 274 | + |
| 275 | + // Filter by score threshold and sort by score descending |
| 276 | + data class Candidate(val index: Int, val box: List<Double>, val score: Double, val detection: Map<String, Any>) |
| 277 | + |
| 278 | + val candidates = detections.indices |
| 279 | + .filter { scores[it] >= scoreThreshold } |
| 280 | + .map { Candidate(it, xyxyBoxes[it], scores[it], detections[it]) } |
| 281 | + .sortedByDescending { it.score } |
| 282 | + .toMutableList() |
| 283 | + |
| 284 | + val keepIndices = mutableListOf<Int>() |
| 285 | + val keepDetections = mutableListOf<Map<String, Any>>() |
| 286 | + val suppressed = mutableSetOf<Int>() |
| 287 | + |
| 288 | + for (candidate in candidates) { |
| 289 | + if (candidate.index in suppressed) continue |
| 290 | + if (keepIndices.size >= maxDets) break |
| 291 | + |
| 292 | + keepIndices.add(candidate.index) |
| 293 | + // Preserve original detection with all fields (box, score, classIndex, label) |
| 294 | + keepDetections.add(candidate.detection) |
| 295 | + |
| 296 | + // Suppress overlapping boxes |
| 297 | + for (other in candidates) { |
| 298 | + if (other.index !in suppressed && other.index != candidate.index) { |
| 299 | + val iouResult = calculateIoU(candidate.box, other.box, "xyxy") |
| 300 | + if ((iouResult["iou"] ?: 0.0) > iouThreshold) { |
| 301 | + suppressed.add(other.index) |
| 302 | + } |
| 303 | + } |
| 304 | + } |
| 305 | + } |
| 306 | + |
| 307 | + val endTime = System.nanoTime() |
| 308 | + val processingTimeMs = (endTime - startTime) / 1_000_000.0 |
| 309 | + |
| 310 | + return NMSResult( |
| 311 | + indices = keepIndices, |
| 312 | + detections = keepDetections, |
| 313 | + suppressedCount = detections.size - keepIndices.size, |
| 314 | + processingTimeMs = processingTimeMs |
| 315 | + ) |
| 316 | + } |
| 317 | +} |
0 commit comments