1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374 |
- # vim: expandtab:ts=4:sw=4
- import numpy as np
- import cv2
- def non_max_suppression(boxes, classes, max_bbox_overlap, scores=None):
- """Suppress overlapping detections.
- Original code from [1]_ has been adapted to include confidence score.
- .. [1] http://www.pyimagesearch.com/2015/02/16/
- faster-non-maximum-suppression-python/
- Examples
- --------
- >>> boxes = [d.roi for d in detections]
- >>> classes = [d.classes for d in detections]
- >>> scores = [d.confidence for d in detections]
- >>> indices = non_max_suppression(boxes, max_bbox_overlap, scores)
- >>> detections = [detections[i] for i in indices]
- Parameters
- ----------
- boxes : ndarray
- Array of ROIs (x, y, width, height).
- max_bbox_overlap : float
- ROIs that overlap more than this values are suppressed.
- scores : Optional[array_like]
- Detector confidence score.
- Returns
- -------
- List[int]
- Returns indices of detections that have survived non-maxima suppression.
- """
- if len(boxes) == 0:
- return []
- boxes = boxes.astype(np.float)
- pick = []
- x1 = boxes[:, 0]
- y1 = boxes[:, 1]
- x2 = boxes[:, 2] + boxes[:, 0]
- y2 = boxes[:, 3] + boxes[:, 1]
- area = (x2 - x1 + 1) * (y2 - y1 + 1)
- if scores is not None:
- idxs = np.argsort(scores)
- else:
- idxs = np.argsort(y2)
- while len(idxs) > 0:
- last = len(idxs) - 1
- i = idxs[last]
- pick.append(i)
- xx1 = np.maximum(x1[i], x1[idxs[:last]])
- yy1 = np.maximum(y1[i], y1[idxs[:last]])
- xx2 = np.minimum(x2[i], x2[idxs[:last]])
- yy2 = np.minimum(y2[i], y2[idxs[:last]])
- w = np.maximum(0, xx2 - xx1 + 1)
- h = np.maximum(0, yy2 - yy1 + 1)
- overlap = (w * h) / area[idxs[:last]]
- idxs = np.delete(
- idxs, np.concatenate(
- ([last], np.where(overlap > max_bbox_overlap)[0])))
- return pick
|