nonMaxSuppression.js 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. export function nonMaxSuppression(boxes, scores, maxOutputSize, iouThreshold, scoreThreshold) {
  2. var numBoxes = boxes.shape[0];
  3. var outputSize = Math.min(maxOutputSize, numBoxes);
  4. var candidates = scores
  5. .map(function (score, boxIndex) { return ({ score: score, boxIndex: boxIndex }); })
  6. .filter(function (c) { return c.score > scoreThreshold; })
  7. .sort(function (c1, c2) { return c2.score - c1.score; });
  8. var suppressFunc = function (x) { return x <= iouThreshold ? 1 : 0; };
  9. var selected = [];
  10. candidates.forEach(function (c) {
  11. if (selected.length >= outputSize) {
  12. return;
  13. }
  14. var originalScore = c.score;
  15. for (var j = selected.length - 1; j >= 0; --j) {
  16. var iou = IOU(boxes, c.boxIndex, selected[j]);
  17. if (iou === 0.0) {
  18. continue;
  19. }
  20. c.score *= suppressFunc(iou);
  21. if (c.score <= scoreThreshold) {
  22. break;
  23. }
  24. }
  25. if (originalScore === c.score) {
  26. selected.push(c.boxIndex);
  27. }
  28. });
  29. return selected;
  30. }
  31. function IOU(boxes, i, j) {
  32. var boxesData = boxes.arraySync();
  33. var yminI = Math.min(boxesData[i][0], boxesData[i][2]);
  34. var xminI = Math.min(boxesData[i][1], boxesData[i][3]);
  35. var ymaxI = Math.max(boxesData[i][0], boxesData[i][2]);
  36. var xmaxI = Math.max(boxesData[i][1], boxesData[i][3]);
  37. var yminJ = Math.min(boxesData[j][0], boxesData[j][2]);
  38. var xminJ = Math.min(boxesData[j][1], boxesData[j][3]);
  39. var ymaxJ = Math.max(boxesData[j][0], boxesData[j][2]);
  40. var xmaxJ = Math.max(boxesData[j][1], boxesData[j][3]);
  41. var areaI = (ymaxI - yminI) * (xmaxI - xminI);
  42. var areaJ = (ymaxJ - yminJ) * (xmaxJ - xminJ);
  43. if (areaI <= 0 || areaJ <= 0) {
  44. return 0.0;
  45. }
  46. var intersectionYmin = Math.max(yminI, yminJ);
  47. var intersectionXmin = Math.max(xminI, xminJ);
  48. var intersectionYmax = Math.min(ymaxI, ymaxJ);
  49. var intersectionXmax = Math.min(xmaxI, xmaxJ);
  50. var intersectionArea = Math.max(intersectionYmax - intersectionYmin, 0.0) *
  51. Math.max(intersectionXmax - intersectionXmin, 0.0);
  52. return intersectionArea / (areaI + areaJ - intersectionArea);
  53. }
  54. //# sourceMappingURL=nonMaxSuppression.js.map