|
1
|
|
|
"""PyBM3D packaging and distribution.""" |
|
2
|
|
|
import os |
|
3
|
|
|
from setuptools import setup, Extension |
|
4
|
|
|
from setuptools.command.build_ext import build_ext as _build_ext |
|
5
|
|
|
|
|
6
|
|
|
|
|
7
|
|
|
class CustomBuildExt(_build_ext): |
|
8
|
|
|
"""Custom build extension class.""" |
|
9
|
|
|
|
|
10
|
|
|
def __init__(self, *args, **kwargs): |
|
11
|
|
|
"""Extends default class constructor.""" |
|
12
|
|
|
# _build_ext is not a new-style (object) class, which is required for |
|
13
|
|
|
# super to work |
|
14
|
|
|
_build_ext.__init__(self, *args, **kwargs) |
|
15
|
|
|
self.cuda_config = self.load_cuda_config() |
|
16
|
|
|
|
|
17
|
|
|
def build_extensions(self): |
|
18
|
|
|
"""Extends default build_extensions method. |
|
19
|
|
|
|
|
20
|
|
|
Further customizes the compiler to add C file specific compile |
|
21
|
|
|
arguments and support nvcc compilation of *.cu CUDA files.""" |
|
22
|
|
|
|
|
23
|
|
|
self.customize_compiler_for_c_args_and_nvcc() |
|
24
|
|
|
_build_ext.build_extensions(self) |
|
25
|
|
|
|
|
26
|
|
|
def finalize_options(self): |
|
27
|
|
|
"""Extends default finalize_options method. |
|
28
|
|
|
|
|
29
|
|
|
Injects NumPy`s C include directories into the Cython compilation. This |
|
30
|
|
|
is done after NumPy was installed through the setuptools setup_requires |
|
31
|
|
|
argument which removes NumPy from the necessary preinstalled |
|
32
|
|
|
packages.""" |
|
33
|
|
|
|
|
34
|
|
|
_build_ext.finalize_options(self) |
|
35
|
|
|
# prevent numpy from thinking it is still in its setup process |
|
36
|
|
|
__builtins__.__NUMPY_SETUP__ = False |
|
37
|
|
|
import numpy |
|
38
|
|
|
self.include_dirs.append(numpy.get_include()) |
|
39
|
|
|
|
|
40
|
|
|
def customize_compiler_for_c_args_and_nvcc(self): |
|
41
|
|
|
"""Customize the compiler. |
|
42
|
|
|
|
|
43
|
|
|
The customization adds C file specific compile |
|
44
|
|
|
arguments and support for nvcc compilation of *.cu CUDA files.""" |
|
45
|
|
|
|
|
46
|
|
|
self.compiler.src_extensions.append('.cu') |
|
47
|
|
|
|
|
48
|
|
|
# save references to the default compiler_so and _comple methods |
|
49
|
|
|
default_compiler_so = self.compiler.compiler_so |
|
50
|
|
|
super = self.compiler._compile |
|
51
|
|
|
|
|
52
|
|
|
# now redefine the _compile method. This gets executed for each |
|
53
|
|
|
# object but distutils doesn't have the ability to change compilers |
|
54
|
|
|
# based on source extension: we add it. |
|
55
|
|
|
def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts): |
|
56
|
|
|
if os.path.splitext(src)[1] == '.cu': |
|
57
|
|
|
# use the nvcc for *.cu files |
|
58
|
|
|
self.compiler.set_executable('compiler_so', |
|
59
|
|
|
self.cuda_config['nvcc']) |
|
60
|
|
|
postargs = extra_postargs['nvcc'] |
|
61
|
|
|
else: |
|
62
|
|
|
postargs = extra_postargs['unix'] |
|
63
|
|
|
|
|
64
|
|
|
# add C file specific compile arguments |
|
65
|
|
|
if os.path.splitext(src)[1] == '.c': |
|
66
|
|
|
postargs = postargs + extra_postargs['c_args'] |
|
67
|
|
|
|
|
68
|
|
|
super(obj, src, ext, cc_args, postargs, pp_opts) |
|
69
|
|
|
# reset the default compiler_so |
|
70
|
|
|
self.compiler.compiler_so = default_compiler_so |
|
71
|
|
|
|
|
72
|
|
|
# inject our redefined _compile method into the class |
|
73
|
|
|
self.compiler._compile = _compile |
|
74
|
|
|
|
|
75
|
|
|
@staticmethod |
|
76
|
|
|
def load_cuda_config(): |
|
77
|
|
|
"""Locate the CUDA environment on the system |
|
78
|
|
|
Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64' |
|
79
|
|
|
and values giving the absolute path to each directory. |
|
80
|
|
|
""" |
|
81
|
|
|
def find_in_path(name, path): |
|
82
|
|
|
"""Finds a file by name in a search path.""" |
|
83
|
|
|
for dir in path.split(os.pathsep): |
|
84
|
|
|
binpath = os.path.join(dir, name) |
|
85
|
|
|
if os.path.exists(binpath): |
|
86
|
|
|
return os.path.abspath(binpath) |
|
87
|
|
|
return None |
|
88
|
|
|
|
|
89
|
|
|
# first check if the CUDA_HOME env variable is in use |
|
90
|
|
|
if 'CUDA_HOME' in os.environ: |
|
91
|
|
|
home = os.environ['CUDA_HOME'] |
|
92
|
|
|
nvcc = os.path.join(home, 'bin', 'nvcc') |
|
93
|
|
|
else: |
|
94
|
|
|
# otherwise, search the PATH for NVCC |
|
95
|
|
|
nvcc = find_in_path('nvcc', os.environ['PATH']) |
|
96
|
|
|
if nvcc is None: |
|
97
|
|
|
return {'cuda_available': False} |
|
98
|
|
|
raise EnvironmentError('The nvcc binary could not be located ' |
|
99
|
|
|
'in your $PATH. Either add it to your ' |
|
100
|
|
|
'path, or set $CUDA_HOME') |
|
101
|
|
|
home = os.path.dirname(os.path.dirname(nvcc)) |
|
102
|
|
|
|
|
103
|
|
|
cuda_config = {'home': home, 'nvcc': nvcc, |
|
104
|
|
|
'include': os.path.join(home, 'include'), |
|
105
|
|
|
'lib64': os.path.join(home, 'lib64')} |
|
106
|
|
|
for k, v in cuda_config.items(): |
|
107
|
|
|
if not os.path.exists(v): |
|
108
|
|
|
raise EnvironmentError('The CUDA %s path could not be located ' |
|
109
|
|
|
'in %s' % (k, v)) |
|
110
|
|
|
|
|
111
|
|
|
return cuda_config |
|
112
|
|
|
|
|
113
|
|
|
|
|
114
|
|
|
# optimize to the current CPU and enable warnings |
|
115
|
|
|
extra_compile_args = {'unix': ['-march=native', '-Wall', '-Wextra', ], |
|
116
|
|
|
'c_args': ['-std=c99', ]} |
|
117
|
|
|
libraries = ["png", "tiff", "jpeg", "fftw3", "fftw3f"] |
|
118
|
|
|
ext_modules = [Extension("pybm3d.bm3d", |
|
119
|
|
|
language="c++", |
|
120
|
|
|
sources=["pybm3d/bm3d.pyx", |
|
121
|
|
|
"bm3d_src/iio.c", |
|
122
|
|
|
"bm3d_src/bm3d.cpp", |
|
123
|
|
|
"bm3d_src/lib_transforms.cpp", |
|
124
|
|
|
"bm3d_src/utilities.cpp", ], |
|
125
|
|
|
extra_compile_args=extra_compile_args, |
|
126
|
|
|
libraries=libraries)] |
|
127
|
|
|
|
|
128
|
|
|
setup( |
|
129
|
|
|
name='pybm3d', |
|
130
|
|
|
version='0.2.0', |
|
131
|
|
|
description='Python wrapper around BM3D', |
|
132
|
|
|
author='Eric Jonas', |
|
133
|
|
|
author_email='[email protected]', |
|
134
|
|
|
url='https://github.com/ericmjonas/pybm3d', |
|
135
|
|
|
zip_safe=False, |
|
136
|
|
|
packages=['pybm3d'], |
|
137
|
|
|
ext_modules=ext_modules, |
|
138
|
|
|
cmdclass={'build_ext': CustomBuildExt}, |
|
139
|
|
|
setup_requires=['numpy>=1.13', ], |
|
140
|
|
|
install_requires=[ |
|
141
|
|
|
'setuptools>=18.0', |
|
142
|
|
|
'cython>=0.27', |
|
143
|
|
|
'numpy>=1.13', ], |
|
144
|
|
|
tests_require=[ |
|
145
|
|
|
'pytest', |
|
146
|
|
|
'scikit-image', ], |
|
147
|
|
|
) |
|
148
|
|
|
|
ericmjonas /
pybm3d
