setup.py

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 >= [2, 4], "Requires PyTorch >= 2.4"


def get_extension():
    this_dir = path.dirname(path.abspath(__file__))
    extensions_dir = path.join(this_dir, "pt_seq_nms", "csrc")
    sources = glob.glob(path.join(extensions_dir, "*.cpp"))

    include_dirs = [extensions_dir]
    extensions = [CppExtension("seq_nms", sources, include_dirs=include_dirs)]

    return extensions


setup(
    name="pt_seq_nms",
    version="0.1",
    packages=("pt_seq_nms",),
    description="PyTorch implementation of the seq-nms algorithm",
    python_requires=">=3.10",
    install_requires=[],
    ext_modules=get_extension(),
    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)