setup.py

from setuptools import setup, find_packages, Extension
from torch.utils import cpp_extension

# plugins_ext = Extension(
#    name='trt_pose.plugins',
#    sources=['trt_pose/find_peaks.cpp'],
#    include_dirs=cpp_extension.include_paths(),
#    language='c++'
# )

setup(
    name="trt_pose",
    version="0.0.1",
    description="Pose detection accelerated by NVIDIA TensorRT",
    packages=find_packages(),
    ext_package="trt_pose",
    ext_modules=[
        cpp_extension.CppExtension(
            "plugins",
            [
                "trt_pose/parse/find_peaks.cpp",
                "trt_pose/parse/paf_score_graph.cpp",
                "trt_pose/parse/refine_peaks.cpp",
                "trt_pose/parse/munkres.cpp",
                "trt_pose/parse/connect_parts.cpp",
                "trt_pose/plugins.cpp",
                "trt_pose/train/generate_cmap.cpp",
                "trt_pose/train/generate_paf.cpp",
            ],
        )
    ],
    cmdclass={"build_ext": cpp_extension.BuildExtension},
    install_requires=[],
)