Apply non-maximum suppression at test time to avoid detecting too many overlapping bounding boxes for a given object. Args: boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. scores: (tensor) The class predscores for the img, Shape:[num_priors]. ov
(boxes, scores, overlap=0.5, top_k=200)
| 173 | # https://github.com/fmassa/object-detection.torch |
| 174 | # Ported to PyTorch by Max deGroot (02/01/2017) |
| 175 | def nms(boxes, scores, overlap=0.5, top_k=200): |
| 176 | """Apply non-maximum suppression at test time to avoid detecting too many |
| 177 | overlapping bounding boxes for a given object. |
| 178 | Args: |
| 179 | boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. |
| 180 | scores: (tensor) The class predscores for the img, Shape:[num_priors]. |
| 181 | overlap: (float) The overlap thresh for suppressing unnecessary boxes. |
| 182 | top_k: (int) The Maximum number of box preds to consider. |
| 183 | Return: |
| 184 | The indices of the kept boxes with respect to num_priors. |
| 185 | """ |
| 186 | |
| 187 | keep = scores.new(scores.size(0)).zero_().long() |
| 188 | if boxes.numel() == 0: |
| 189 | return keep |
| 190 | x1 = boxes[:, 0] |
| 191 | y1 = boxes[:, 1] |
| 192 | x2 = boxes[:, 2] |
| 193 | y2 = boxes[:, 3] |
| 194 | area = torch.mul(x2 - x1, y2 - y1) |
| 195 | v, idx = scores.sort(0) # sort in ascending order |
| 196 | # I = I[v >= 0.01] |
| 197 | idx = idx[-top_k:] # indices of the top-k largest vals |
| 198 | xx1 = boxes.new() |
| 199 | yy1 = boxes.new() |
| 200 | xx2 = boxes.new() |
| 201 | yy2 = boxes.new() |
| 202 | w = boxes.new() |
| 203 | h = boxes.new() |
| 204 | |
| 205 | # keep = torch.Tensor() |
| 206 | count = 0 |
| 207 | while idx.numel() > 0: |
| 208 | i = idx[-1] # index of current largest val |
| 209 | # keep.append(i) |
| 210 | keep[count] = i |
| 211 | count += 1 |
| 212 | if idx.size(0) == 1: |
| 213 | break |
| 214 | idx = idx[:-1] # remove kept element from view |
| 215 | # load bboxes of next highest vals |
| 216 | torch.index_select(x1, 0, idx, out=xx1) |
| 217 | torch.index_select(y1, 0, idx, out=yy1) |
| 218 | torch.index_select(x2, 0, idx, out=xx2) |
| 219 | torch.index_select(y2, 0, idx, out=yy2) |
| 220 | # store element-wise max with next highest score |
| 221 | xx1 = torch.clamp(xx1, min=x1[i]) |
| 222 | yy1 = torch.clamp(yy1, min=y1[i]) |
| 223 | xx2 = torch.clamp(xx2, max=x2[i]) |
| 224 | yy2 = torch.clamp(yy2, max=y2[i]) |
| 225 | w.resize_as_(xx2) |
| 226 | h.resize_as_(yy2) |
| 227 | w = xx2 - xx1 |
| 228 | h = yy2 - yy1 |
| 229 | # check sizes of xx1 and xx2.. after each iteration |
| 230 | w = torch.clamp(w, min=0.0) |
| 231 | h = torch.clamp(h, min=0.0) |
| 232 | inter = w*h |