- Notifications
You must be signed in to change notification settings - Fork2
PyTorch implementation of soft-nms
License
NotificationsYou must be signed in to change notification settings
MrParosk/soft_nms
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
Implementation of the soft-nms algorithm described in the paper:Soft-NMS -- Improving Object Detection With One Line of Code
The algorithm is implemented in PyTorch's C++ frontend for better performance.
Make sure that you have installed PyTorch, version 1.7 or higher. Install the package by
pip install git+https://github.com/MrParosk/soft_nms.git
Note that if you are using Windows, you need MSVC installed.
importtorchfrompt_soft_nmsimportbatched_soft_nms,soft_nmssigma=0.5score_threshold=0.1boxes=torch.tensor([[20,20,40,40], [10,10,20,20], [20,20,35,35]],device="cpu",dtype=torch.float)scores=torch.tensor([0.5,0.9,0.11],device="cpu",dtype=torch.float)updated_scores,keep=soft_nms(boxes,scores,sigma,score_threshold)# updated_scores=tensor([0.9000, 0.5000]), keep=tensor([1, 0])# With batched_soft_nms, the soft-nms will be applied per batch, which is specified with indiciesindicies=torch.tensor([0,0,1],device="cpu")keep_batch=batched_soft_nms(boxes,scores,indicies,sigma,score_threshold)# keep_batch=tensor([1, 0, 2])