py_cpu_nms.py 1.0 KB

1234567891011121314151617181920212223242526272829303132333435363738
  1. # --------------------------------------------------------
  2. # Fast R-CNN
  3. # Copyright (c) 2015 Microsoft
  4. # Licensed under The MIT License [see LICENSE for details]
  5. # Written by Ross Girshick
  6. # --------------------------------------------------------
  7. import numpy as np
  8. def py_cpu_nms(dets, thresh):
  9. """Pure Python NMS baseline."""
  10. x1 = dets[:, 0]
  11. y1 = dets[:, 1]
  12. x2 = dets[:, 2]
  13. y2 = dets[:, 3]
  14. scores = dets[:, 4]
  15. areas = (x2 - x1 + 1) * (y2 - y1 + 1)
  16. order = scores.argsort()[::-1]
  17. keep = []
  18. while order.size > 0:
  19. i = order[0]
  20. keep.append(i)
  21. xx1 = np.maximum(x1[i], x1[order[1:]])
  22. yy1 = np.maximum(y1[i], y1[order[1:]])
  23. xx2 = np.minimum(x2[i], x2[order[1:]])
  24. yy2 = np.minimum(y2[i], y2[order[1:]])
  25. w = np.maximum(0.0, xx2 - xx1 + 1)
  26. h = np.maximum(0.0, yy2 - yy1 + 1)
  27. inter = w * h
  28. ovr = inter / (areas[i] + areas[order[1:]] - inter)
  29. inds = np.where(ovr <= thresh)[0]
  30. order = order[inds + 1]
  31. return keep