diff --git a/setup.py b/setup.py index d49fddde177..0e5d7d3bcf8 100644 --- a/setup.py +++ b/setup.py @@ -5,6 +5,7 @@ import shutil import sys from setuptools import setup, find_packages +from pkg_resources import get_distribution, DistributionNotFound def read(*names, **kwargs): @@ -15,6 +16,13 @@ def read(*names, **kwargs): return fp.read() +def get_dist(pkgname): + try: + return get_distribution(pkgname) + except DistributionNotFound: + return None + + def find_version(*file_paths): version_file = read(*file_paths) version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", @@ -30,12 +38,15 @@ def find_version(*file_paths): requirements = [ 'numpy', - 'pillow >= 4.1.1', 'six', 'torch', 'tqdm' ] +pillow_ver = ' >= 4.1.1' +pillow_req = 'pillow-simd' if get_dist('pillow-simd') is not None else 'pillow' +requirements.append(pillow_req + pillow_ver) + setup( # Metadata name='torchvision',