preprocessing.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. # vim: expandtab:ts=4:sw=4
  2. import numpy as np
  3. import cv2
  4. def non_max_suppression(boxes, classes, max_bbox_overlap, scores=None):
  5. """Suppress overlapping detections.
  6. Original code from [1]_ has been adapted to include confidence score.
  7. .. [1] http://www.pyimagesearch.com/2015/02/16/
  8. faster-non-maximum-suppression-python/
  9. Examples
  10. --------
  11. >>> boxes = [d.roi for d in detections]
  12. >>> classes = [d.classes for d in detections]
  13. >>> scores = [d.confidence for d in detections]
  14. >>> indices = non_max_suppression(boxes, max_bbox_overlap, scores)
  15. >>> detections = [detections[i] for i in indices]
  16. Parameters
  17. ----------
  18. boxes : ndarray
  19. Array of ROIs (x, y, width, height).
  20. max_bbox_overlap : float
  21. ROIs that overlap more than this values are suppressed.
  22. scores : Optional[array_like]
  23. Detector confidence score.
  24. Returns
  25. -------
  26. List[int]
  27. Returns indices of detections that have survived non-maxima suppression.
  28. """
  29. if len(boxes) == 0:
  30. return []
  31. boxes = boxes.astype(np.float)
  32. pick = []
  33. x1 = boxes[:, 0]
  34. y1 = boxes[:, 1]
  35. x2 = boxes[:, 2] + boxes[:, 0]
  36. y2 = boxes[:, 3] + boxes[:, 1]
  37. area = (x2 - x1 + 1) * (y2 - y1 + 1)
  38. if scores is not None:
  39. idxs = np.argsort(scores)
  40. else:
  41. idxs = np.argsort(y2)
  42. while len(idxs) > 0:
  43. last = len(idxs) - 1
  44. i = idxs[last]
  45. pick.append(i)
  46. xx1 = np.maximum(x1[i], x1[idxs[:last]])
  47. yy1 = np.maximum(y1[i], y1[idxs[:last]])
  48. xx2 = np.minimum(x2[i], x2[idxs[:last]])
  49. yy2 = np.minimum(y2[i], y2[idxs[:last]])
  50. w = np.maximum(0, xx2 - xx1 + 1)
  51. h = np.maximum(0, yy2 - yy1 + 1)
  52. overlap = (w * h) / area[idxs[:last]]
  53. idxs = np.delete(
  54. idxs, np.concatenate(
  55. ([last], np.where(overlap > max_bbox_overlap)[0])))
  56. return pick