Non-max Supression for boxes with multiple classes

benihime91: <p>Hi. How do i do <code>non-max-supression</code> for bounding-boxes with mutiple labels in <code>PyTorch</code>?<br>
I was hoping for something that can provide the same functionality as this <code>TensorFlow</code> function :  <a href="https://www.tensorflow.org/api_docs/python/tf/image/combined_non_max_suppression" rel="nofollow noopener"><code>tf.image.combined_non_max_suppression</code></a></p>
<p>Thank You.</p>

Usama_Hasan: <p>Hy <a class="mention" href="/u/benihime91">@benihime91</a> Have you checked <strong>torchvision.ops.nms</strong><br>
<a href="https://pytorch.org/docs/stable/torchvision/ops.html" rel="nofollow noopener">https://pytorch.org/docs/stable/torchvision/ops.html</a></p>

benihime91: <p>Hello <a class="mention" href="/u/usama_hasan">@Usama_Hasan</a><br>
Yes I have …<br>
<code>torchvision.ops.nms</code> performs <code>nms</code> only for given <code>boxes</code> &amp; <code>scores</code> i don’t think it  take into account the <code>classes</code>  of the boxes (correct me if I am wrong)…</p>

Usama_Hasan: <aside class="quote no-group" data-username="benihime91" data-post="3" data-topic="92060">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/benihime91/48/29337_2.png" class="avatar"> benihime91:</div>
<blockquote>
<p><code>classes</code> of the boxes (correct me if I am wrong)</p>
</blockquote>
</aside>
<p>Yes that’s the case.<br>
I guess this will help you in what you’re looking for.</p><aside class="quote quote-modified" data-post="8" data-topic="34210">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/etienne_perot/48/16894_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/batch-non-maximum-suppression-on-the-gpu/34210/8">Batch non-maximum suppression on the GPU</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    hello there, 
using the awesome idea from torchvision “batched_nms”, this following code can decode for several images / several classes at once, it works because batched_nms offsets boxes according to their category, so you never perform a wrong suppression. 
I also tried to accelerate box encoding, if you are interested you can have a peek here: <a href="https://github.com/etienne87/torch_object_rnn/blob/master/core/anchors.py" rel="noopener nofollow ugc">https://github.com/etienne87/torch_object_rnn/blob/master/core/anchors.py</a> 
num_classes = cls_preds.shape[-1] - self.label_offset
num_anchors = box_pre…
  </blockquote>
</aside>


benihime91: <p><a class="mention" href="/u/usama_hasan">@Usama_Hasan</a></p>
<p>I was looking around the sources code of <code>detectron-2</code> where I found this really simple &amp; efficient implementation:<br>
I implemented it like-so:</p>
<pre><code class="lang-auto">def batched_nms(boxes: torch.Tensor, scores: torch.Tensor, idxs: torch.Tensor, iou_threshold: float):
    """
    [
    From :
    (https://github.com/facebookresearch/detectron2/blob/master/detectron2/modeling/roi_heads/fast_rcnn.py#L107-L115)
    ]
    Same as torchvision.ops.boxes.batched_nms, but safer.
    """
    assert boxes.shape[-1] == 4
    if len(boxes) &lt; 40000:
        return torchvision.ops.boxes.batched_nms(boxes, scores, idxs, iou_threshold)

    result_mask = scores.new_zeros(scores.size(), dtype=torch.bool)
    for id in torch.jit.annotate(List[int], torch.unique(idxs).cpu().tolist()):
        mask = (idxs == id).nonzero().view(-1)
        keep = torchvision.ops.nms(boxes[mask], scores[mask], iou_threshold)
        result_mask[mask[keep]] = True
    keep = result_mask.nonzero().view(-1)
    keep = keep[scores[keep].argsort(descending=True)]
    return keep
</code></pre>
<pre><code class="lang-auto">def NonMaxSupression(predictions, score_thres=0.5, iou_thres=0.5):
    predictions[..., :4] = cvt_boxes(predictions[..., :4])
    output_preds = []
    
    for i, pred in enumerate(predictions):
        pred = pred[pred[:, 4] &gt;= score_thres]
        if not pred.size(0):   continue

        boxes = pred[:,:4]
        scores = pred[:, 4] * pred[:, 5:].max(1)[0] 
        _ ,cls_labels = pred[:, 5:].max(1)

        detections = torch.cat((boxes, scores.unsqueeze(dim=1), cls_labels.unsqueeze(dim=1).float()),
                               dim=-1)
        
        scores, cls_labels = scores.view([-1]), cls_labels.type(torch.int32).view([-1])
        
        assert boxes.size(0) == scores.size(0) == cls_labels.size(0)
        
        keep = batched_nms(boxes, scores, cls_labels, iou_thres)
        output_preds.append(detections[keep])
        
    assert len(output_preds) &gt; 0, "[INFO] No objects detected in current Image"
    
    return output_preds
</code></pre>

Etienne_Perot: <aside class="quote no-group" data-username="benihime91" data-post="5" data-topic="92060">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/benihime91/48/29337_2.png" class="avatar"> benihime91:</div>
<blockquote>
<p>where I found this really simple &amp;</p>
</blockquote>
</aside>
<p>I think this is fine, just not really parallel implementation.</p>