-
Notifications
You must be signed in to change notification settings - Fork 2
/
setup.py
32 lines (24 loc) · 923 Bytes
/
setup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import glob
from os import path
import torch
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
torch_ver = [int(x) for x in torch.__version__.split(".")[:2]]
assert torch_ver >= [1, 7], "Requires PyTorch >= 1.7"
def get_extension():
this_dir = path.dirname(path.abspath(__file__))
extensions_dir = path.join(this_dir, "pt_soft_nms", "csrc")
sources = glob.glob(path.join(extensions_dir, "*.cpp"))
include_dirs = [extensions_dir]
extensions = [CppExtension("soft_nms", sources, include_dirs=include_dirs)]
return extensions
setup(
name="pt_soft_nms",
version="0.1",
packages=("pt_soft_nms",),
description="PyTorch implementation of the soft-nms algorithm",
python_requires=">=3.8",
install_requires=[],
ext_modules=get_extension(),
cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)