call batched_nms from torchvision directly
Summary:
torchvision already has the thresholding logic.
The threshold in torchvision is not good, which will be fixed in
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.
-from typing import List
import torch
from torchvision.ops import boxes as box_ops
-from torchvision.ops import nms # BC-compat
+from torchvision.ops import nms # noqa . for compatibility
def batched_nms(
boxes: torch.Tensor, scores: torch.Tensor, idxs: torch.Tensor, iou_threshold: float
):
"""
- Same as torchvision.ops.boxes.batched_nms, but safer.
+ Same as torchvision.ops.boxes.batched_nms, but with float().
"""
assert boxes.shape[-1] == 4
- # TODO may need better strategy.
- # Investigate after having a fully-cuda NMS op.
- if len(boxes) < 40000:
- # fp16 does not have enough range for batched NMS
+ # Note: Torchvision already has a strategy (https://github.com/pytorch/vision/issues/1311)
+ # to decide whether to use coordinate trick or for loop to implement batched_nms. So we
+ # just call it directly.
+ # Fp16 does not have enough range for batched NMS, so adding float().
return box_ops.batched_nms(boxes.float(), 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 = 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
-
# Note: this function (nms_rotated) might be moved into
# torchvision/ops/boxes.py in the future
