目录
https://github.com/pjl1995/CTracker/blob/master/model.py
调用:
from lib.cython_nms.cpu_nms import soft_nms as cython_soft_nms_wrapper
demo:
import numpy as np
import cpu_nms
# from lib.cython_nms import cpu_nms import cython_soft_nms_wrapper
x = np.ones([11,10], dtype =np.float32)
aa= cpu_nms.soft_nms(x,0.7)
print(aa)
编译python库:
python setup_cpu.py build_ext --inplace
python编译通用库:用vs的native命令窗口
python setup.py build_ext --inplace
有时报错,需要用vs的native命令窗口
(2)打开vs2019(或者vs2017)专用的命令行窗口
因为我是64位的系统,因此选择的x64 native tools command prompt for vs2019
from setuptools import setup, Extension
import numpy as np
from Cython.Build import cythonize
from Cython.Distutils import build_ext
# Obtain the numpy include directory. This logic works across numpy versions.
try:
numpy_include = np.get_include()
except AttributeError:
numpy_include = np.get_numpy_include()
# extensions
ext_args = dict(
include_dirs=[numpy_include],
language='c++',
)
ext_modules = [
Extension(
"bbox",
["bbox.pyx"],
**ext_args
),
Extension(
"cpu_nms",
["cpu_nms.pyx"],
**ext_args
),
]
setup(
name='frcnn_cython',
ext_modules=cythonize(ext_modules),
# inject our custom trigger
cmdclass={'build_ext': build_ext},
)
pytorch nms
import torch
def nms(boxes, scores, overlap=0.7, top_k=200):
"""
输入:
boxes: 存储一个图片的所有预测框。[num_positive,4].
scores:置信度。如果为多分类则需要将nms函数套在一个循环内。[num_positive].
overlap: nms抑制时iou的阈值.
top_k: 先选取置信度前top_k个框再进行nms.
返回:
nms后剩余预测框的索引.
"""
keep = scores.new(scores.size(0)).zero_().long()
# 保存留下来的box的索引 [num_positive]
# 函数new(): 构建一个有相同数据类型的tensor
# 如果输入box为空则返回空Tensor
if boxes.numel() == 0:
return keep
x1 = boxes[:, 0] # x1 坐标
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
area = torch.mul(x2 - x1, y2 - y1) # 并行化计算所有框的面积
v, idx = scores.sort(0) # 升序排序
idx = idx[-top_k:] # 前top-k的索引,从小到大
xx1 = boxes.new()
yy1 = boxes.new()
xx2 = boxes.new() # new() 无参数,创建 相同类型的空值;
yy2 = boxes.new()
w = boxes.new()
h = boxes.new()
count = 0
while idx.numel() > 0:
i = idx[-1] # 目前最大score对应的索引 # 选取得分最大的框索引;
keep[count] = i # 存储在keep中
count += 1
if idx.size(0) == 1: # 跳出循环条件:box被筛选完了
break
idx = idx[:-1] # 去掉最后一个
# 剩下boxes的信息存储在xx,yy中
torch.index_select(x1, 0, idx, out=xx1) # 从x1中再维度0选取索引为idx 数据 输出到xx1中;
torch.index_select(y1, 0, idx, out=yy1) # torch.index_select() # 从tensor中按指定维度和索引 取值;
torch.index_select(x2, 0, idx, out=xx2)
torch.index_select(y2, 0, idx, out=yy2)
# 计算当前最大置信框与其他剩余框的交集,不知道clamp的同学确实容易被误导
xx1 = torch.clamp(xx1, min=x1[i]) # max(x1,xx1) # x1 y1 的最大值
yy1 = torch.clamp(yy1, min=y1[i]) # max(y1,yy1)
xx2 = torch.clamp(xx2, max=x2[i]) # min(x2,xx2) # x2 x3 最小值;
yy2 = torch.clamp(yy2, max=y2[i]) # min(y2,yy2)
w.resize_as_(xx2)
h.resize_as_(yy2)
w = xx2 - xx1 # w=min(x2,xx2)−max(x1,xx1)
h = yy2 - yy1 # h=min(y2,yy2)−max(y1,yy1)
w = torch.clamp(w, min=0.0) # max(w,0)
h = torch.clamp(h, min=0.0) # max(h,0)
inter = w * h
# 计算当前最大置信框与其他剩余框的IOU
# IoU = i / (area(a) + area(b) - i)
rem_areas = torch.index_select(area, 0, idx) # 剩余的框的面积
union = rem_areas + area[i] - inter # 并集
IoU = inter / union # 计算iou
# 选出IoU <= overlap的boxes(注意le函数的使用)
idx = idx[IoU.le(overlap)] # le: 小于等于 返回的bool , 去除大于overlap的值;
return keep, count