# custom_fastNMS

**Repository Path**: yang-minghai22/custom-fast-nms

## Basic Information

- **Project Name**: custom_fastNMS
- **Description**: fast Non-maximum suppression
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2023-08-01
- **Last Updated**: 2023-12-22

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# custom4fastNMS

#### 环境依赖
CANN包使用最新社区包 [CANN 社区版下载-昇腾社区](https://www.hiascend.com/zh/software/cann/community)；
确保安装了pytorch+cpu1.11及torch_npu1.11，未安装参考[pytorch: Ascend PyTorch adapter (gitee.com)](https://gitee.com/ascend/pytorch#%E4%BD%BF%E7%94%A8%E6%BA%90%E7%A0%81%E7%BC%96%E8%AF%91%E5%AE%89%E8%A3%85pytorch%E6%A1%86%E6%9E%B6)

#### 编译安装
1. 编译：编译前，确保环境变量已适配source setenv.sh, 将算子代码仓clone到本地，运行
python3 setup.py bdist_wheel
2. 安装：wheel 包在./dist下，运行（生成的wheel包*.wheel）
pip3 install  *.wheel

#### 使用

```
import torch
from fast_nms import fast_nms

# masks, boxes, scores 为前文输出
idxs = torch.arange(boxes.shape[0]).npu()
masks_tmp, score_res, idxs, gather_cnt, _ = fast_nms(boxes.half(), masks.half(), scoes.half(), idxs.half())

gather_cnt = gather_cnt.cpu()[:,0].reshape(8, -1)
boxes_gather_cnt = gather_cnt.view(-1)
boxes_res = torch.Tensor()
idxs = idxs.cpu().int()
for i in range(80):
    take_out_box = torch.index_select(boxes.cpu(), 0, idxs[i][:boxes_gather_cnt[i]])
    boxes_res = torch.cat((boxes_res, take_out_box), dim=0)

masks_gather_cnt = torch.sum(gather_cnt, -1).long()
masks_res = torch.Tensor()
for i in range(8):
    masks_res = torch.cat((masks_res, masks_tmp[i][:gather_cnt[i]]), dim=0)
```