|
- #!/usr/bin/env python
-
- import glob
- import os
- from setuptools import find_packages, setup
- import torch
- from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
-
- torch_ver = [int(x) for x in torch.__version__.split(".")[:2]]
- assert torch_ver >= [1, 3], "Requires PyTorch >= 1.3"
-
-
- def get_adet_extensions():
- this_dir = os.path.dirname(os.path.abspath(__file__))
- extensions_dir = os.path.join(this_dir, "adet", "layers", "csrc")
-
- main_source = os.path.join(extensions_dir, "vision.cpp")
- sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"))
- source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu")) + glob.glob(
- os.path.join(extensions_dir, "*.cu")
- )
-
- sources = [main_source] + sources
-
- extension = CppExtension
-
- extra_compile_args = {"cxx": []}
- define_macros = []
-
- if (torch.cuda.is_available() and CUDA_HOME is not None) or os.getenv("FORCE_CUDA", "0") == "1":
- extension = CUDAExtension
- sources += source_cuda
- define_macros += [("WITH_CUDA", None)]
- extra_compile_args["nvcc"] = [
- "-DCUDA_HAS_FP16=1",
- "-D__CUDA_NO_HALF_OPERATORS__",
- "-D__CUDA_NO_HALF_CONVERSIONS__",
- "-D__CUDA_NO_HALF2_OPERATORS__",
- ]
-
- # It's better if pytorch can do this by default ..
- CC = os.environ.get("CC", None)
- if CC is not None:
- extra_compile_args["nvcc"].append("-ccbin={}".format(CC))
-
- sources = [os.path.join(extensions_dir, s) for s in sources]
- include_dirs = [extensions_dir]
- ext_modules = [
- extension(
- "adet._C",
- sources,
- include_dirs=include_dirs,
- define_macros=define_macros,
- extra_compile_args=extra_compile_args,
- )
- ]
-
- return ext_modules
-
-
- def get_deformable_extensions():
- this_dir = os.path.dirname(os.path.abspath(__file__))
- extensions_dir = os.path.join(this_dir, "adet", "modeling", "ops", "src")
-
- main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
- source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
- source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
-
- sources = main_file + source_cpu
- extension = CppExtension
- extra_compile_args = {"cxx": []}
- define_macros = []
-
- if torch.cuda.is_available() and CUDA_HOME is not None:
- extension = CUDAExtension
- sources += source_cuda
- define_macros += [("WITH_CUDA", None)]
- extra_compile_args["nvcc"] = [
- "-DCUDA_HAS_FP16=1",
- "-D__CUDA_NO_HALF_OPERATORS__",
- "-D__CUDA_NO_HALF_CONVERSIONS__",
- "-D__CUDA_NO_HALF2_OPERATORS__",
- ]
- else:
- raise NotImplementedError('Cuda is not availabel')
-
- sources = [os.path.join(extensions_dir, s) for s in sources]
- include_dirs = [extensions_dir]
- ext_modules = [
- extension(
- "MultiScaleDeformableAttention",
- sources,
- include_dirs=include_dirs,
- define_macros=define_macros,
- extra_compile_args=extra_compile_args,
- )
- ]
- return ext_modules
-
-
- setup(
- name="AdelaiDet",
- version="0.2.0",
- author="Adelaide Intelligent Machines",
- url="https://github.com/stanstarks/AdelaiDet",
- description="AdelaiDet is AIM's research "
- "platform for instance-level detection tasks based on Detectron2.",
- packages=find_packages(exclude=("configs", "tests")),
- python_requires=">=3.6",
- install_requires=[
- "setuptools==59.5.0",
- "termcolor>=1.1",
- "Pillow>=6.0",
- "yacs>=0.1.6",
- "tabulate",
- "cloudpickle",
- "matplotlib",
- "tqdm>4.29.0",
- "tensorboard",
- "python-Levenshtein",
- "Polygon3",
- "shapely",
- "kornia==0.6.8",
- "opencv-python",
- "timm"
- ],
- extras_require={"all": ["psutil"]},
- ext_modules=get_adet_extensions() + get_deformable_extensions(),
- cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
- )
|