PyTorch Extension in ROS

edit

Hi!

I'm trying to use a pose detection library with ROS. Since I want to be able to edit the library I want to package it in a ros package. So far so good. The problem is the library is in python but uses c++ extensions. Which do not seem to work when using catkin? Is there a work around?

This is my setup.py file:

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

setup(
    name='trt_pose',
    version='0.0.1',
    description='Pose detection accelerated by NVIDIA TensorRT',
    package_dir={'': 'trt_pose'},
    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=[
    ],
)