| 9 | |
| 10 | |
| 11 | def nms_2d(boxes, overlap_threshold): |
| 12 | x1 = boxes[:, 0] |
| 13 | y1 = boxes[:, 1] |
| 14 | x2 = boxes[:, 2] |
| 15 | y2 = boxes[:, 3] |
| 16 | score = boxes[:, 4] |
| 17 | area = (x2 - x1) * (y2 - y1) |
| 18 | |
| 19 | I = np.argsort(score) |
| 20 | pick = [] |
| 21 | while I.size != 0: |
| 22 | last = I.size |
| 23 | i = I[-1] |
| 24 | pick.append(i) |
| 25 | suppress = [last - 1] |
| 26 | for pos in range(last - 1): |
| 27 | j = I[pos] |
| 28 | xx1 = max(x1[i], x1[j]) |
| 29 | yy1 = max(y1[i], y1[j]) |
| 30 | xx2 = min(x2[i], x2[j]) |
| 31 | yy2 = min(y2[i], y2[j]) |
| 32 | w = xx2 - xx1 |
| 33 | h = yy2 - yy1 |
| 34 | if w > 0 and h > 0: |
| 35 | o = w * h / area[j] |
| 36 | print("Overlap is", o) |
| 37 | if o > overlap_threshold: |
| 38 | suppress.append(pos) |
| 39 | I = np.delete(I, suppress) |
| 40 | return pick |
| 41 | |
| 42 | |
| 43 | def nms_2d_faster(boxes, overlap_threshold, old_type=False): |