py-torchvision: add variant to set image backend (#18500)

1 files changed, 22 insertions, 16 deletions

diff --git a/var/spack/repos/builtin/packages/py-torchvision/package.py b/var/spack/repos/builtin/packages/py-torchvision/package.py
index 9cedfa25a0..385a40b209 100644
--- a/var/spack/repos/builtin/packages/py-torchvision/package.py
+++ b/var/spack/repos/builtin/packages/py-torchvision/package.py
@@ -3,8 +3,6 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-from spack import *
-
 
 class PyTorchvision(PythonPackage):
     """The torchvision package consists of popular datasets, model
@@ -12,6 +10,7 @@ class PyTorchvision(PythonPackage):
 
     homepage = "https://github.com/pytorch/vision"
     url      = "https://github.com/pytorch/vision/archive/v0.7.0.tar.gz"
+    git      = "https://github.com/pytorch/vision.git"
 
     maintainers = ['adamjstewart']
     import_modules = [
@@ -21,6 +20,7 @@ class PyTorchvision(PythonPackage):
         'torchvision.models.detection'
     ]
 
+    version('master', branch='master')
     version('0.7.0', sha256='fa0a6f44a50451115d1499b3f2aa597e0092a07afce1068750260fa7dd2c85cb')
     version('0.6.1', sha256='8173680a976c833640ecbd0d7e6f0a11047bf8833433e2147180efc905e48656')
     version('0.6.0', sha256='02de11b3abe6882de4032ce86dab9c7794cbc84369b44d04e667486580f0f1f7')
@@ -30,7 +30,11 @@ class PyTorchvision(PythonPackage):
     version('0.4.0', sha256='c270d74e568bad4559fed4544f6dd1e22e2eb1c60b088e04a5bd5787c4150589')
     version('0.3.0', sha256='c205f0618c268c6ed2f8abb869ef6eb83e5339c1336c243ad321a2f2a85195f0')
 
-    # See README.rst
+    # https://github.com/pytorch/vision#image-backend
+    variant('backend', default='pil', description='Image backend',
+            values=('pil', 'accimage', 'png', 'jpeg'), multi=False)
+
+    # https://github.com/pytorch/vision#installation
     depends_on('python@3.6:', when='@0.7:', type=('build', 'run'))
     depends_on('python@3.5:', when='@0.6.0:0.6.999', type=('build', 'run'))
     depends_on('python@2.7:2.8,3.5:3.8', when='@0.5.0', type=('build', 'run'))
@@ -41,7 +45,7 @@ class PyTorchvision(PythonPackage):
     depends_on('py-numpy', type=('build', 'run'))
     depends_on('py-six', when='@:0.5', type=('build', 'run'))
 
-    # See README.rst
+    # https://github.com/pytorch/vision#installation
     depends_on('py-torch@1.6.0', when='@0.7.0', type=('build', 'link', 'run'))
     depends_on('py-torch@1.5.1', when='@0.6.1', type=('build', 'link', 'run'))
     depends_on('py-torch@1.5.0', when='@0.6.0', type=('build', 'link', 'run'))
@@ -52,16 +56,12 @@ class PyTorchvision(PythonPackage):
     depends_on('py-torch@1.1.0', when='@0.3.0', type=('build', 'link', 'run'))
     depends_on('py-torch@:1.0.1', when='@0.2.2', type=('build', 'link', 'run'))
 
-    # TODO: Torchvision supports the following backends:
-    # * pillow
-    # * pillow-simd
-    # * accimage
-    # * libpng
-    # * libjpeg
-
     # https://github.com/pytorch/vision/issues/1712
-    depends_on('pil@4.1.1:6', when='@:0.4', type=('build', 'run'))
-    depends_on('pil@4.1.1:',  when='@0.5:', type=('build', 'run'))
+    depends_on('pil@4.1.1:6', when='@:0.4 backend=pil', type=('build', 'run'))
+    depends_on('pil@4.1.1:',  when='@0.5: backend=pil', type=('build', 'run'))
+    depends_on('py-accimage', when='backend=accimage', type=('build', 'run'))
+    depends_on('libpng', when='backend=png')
+    depends_on('jpeg', when='backend=jpeg')
 
     # Many of the datasets require additional dependencies to use.
     # These can be installed after the fact.
@@ -69,6 +69,9 @@ class PyTorchvision(PythonPackage):
 
     depends_on('ffmpeg@3.1:', when='@0.4.2:')
 
+    conflicts('backend=png', when='@:0.7')
+    conflicts('backend=jpeg', when='@:0.7')
+
     def setup_build_environment(self, env):
         include = []
         library = []
@@ -78,9 +81,12 @@ class PyTorchvision(PythonPackage):
             library.extend(query.libs.directories)
 
         # README says to use TORCHVISION_INCLUDE and TORCHVISION_LIBRARY,
-        # but these do not work. Build uses a mix of Spack's compiler wrapper
-        # and the actual compiler, so this is needed to get parts of the build
-        # working. See https://github.com/pytorch/vision/issues/2591
+        # but these do not work for older releases. Build uses a mix of
+        # Spack's compiler wrapper and the actual compiler, so this is
+        # needed to get parts of the build working.
+        # See https://github.com/pytorch/vision/issues/2591
+        env.set('TORCHVISION_INCLUDE', ':'.join(include))
+        env.set('TORCHVISION_LIBRARY', ':'.join(library))
         env.set('CPATH', ':'.join(include))
         env.set('LIBRARY_PATH', ':'.join(library))