Add image extension to compilation instructions

setup.py

@@ -171,6 +171,28 @@ def get_extensions():
             )
         )
 
+    try:
+        include_dirs += [os.environ['LIBRARY_INC']]
+    except KeyError:
+        pass
+
+    try:
+        library_dirs = [os.environ['LIBRARY_LIB']]
+    except KeyError:
+        library_dirs = []
+
+    # Image reading extension
+    image_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'image')
+    image_src = glob.glob(os.path.join(image_src_dir, "*.cpp"))
+    ext_modules.append(extension(
+        'torchvision.image',
+        image_src,
+        include_dirs=include_dirs + image_src_dir,
+        library_dirs=library_dirs,
+        define_macros=define_macros,
+        extra_compile_args=extra_compile_args
+    ))
+
     ffmpeg_exe = distutils.spawn.find_executable('ffmpeg')
     has_ffmpeg = ffmpeg_exe is not None
 

test/test_image.py

@@ -35,10 +35,10 @@ def test_decode_png(self):
             size = os.path.getsize(img_path)
             img_lpng = decode_png(torch.from_file(img_path, dtype=torch.uint8, size=size))
             self.assertEqual(img_lpng, img_pil)
-            
+
             self.assertEqual(decode_png(torch.empty()), torch.empty())
             self.assertEqual(decode_png(torch.randint(3, 5, (300,))), torch.empty())
-    
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/io/image.py

@@ -45,4 +45,3 @@ def read_png(path):
         raise ValueError("Expected a non empty file.")
     data = torch.from_file(path, dtype=torch.uint8, size=size)
     return decode_png(data)
-