WIP: Readimage png + jpeg #1632
+393
−30
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,9 @@ | ||
| [submodule "third_party/libpng"] | ||
| path = third_party/libpng | ||
| url = https://github.com/glennrp/libpng | ||
| [submodule "third_party/libjpeg-turbo"] | ||
| path = third_party/libjpeg-turbo | ||
| url = https://github.com/libjpeg-turbo/libjpeg-turbo | ||
| [submodule "third_party/zlib"] | ||
| path = third_party/zlib | ||
| url = https://github.com/madler/zlib |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -14,7 +14,6 @@ | |
| import torch | ||
| from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME | ||
|
|
||
| def read(*names, **kwargs): | ||
| with io.open( | ||
| os.path.join(os.path.dirname(__file__), *names), | ||
|
|
@@ -76,16 +75,51 @@ def write_version_file(): | |
| pillow_req = 'pillow-simd' if get_dist('pillow-simd') is not None else 'pillow' | ||
| requirements.append(pillow_req + pillow_ver) | ||
|
|
||
| third_party_libraries = ['png', 'turbojpeg'] | ||
| third_party_dir = os.path.join(cwd, "third_party") | ||
| third_party_lib_directories = ['libpng', 'libjpeg-turbo'] | ||
| third_party_search_directories = [os.path.join(third_party_dir, directory) for directory in third_party_lib_directories] | ||
|
|
||
| def _build_cmake_dependency(directory, args=""): | ||
| os.chdir(directory) | ||
| if sys.platform == 'win32': | ||
| os.system("cmake.exe --clean .") | ||
| os.system('cmake.exe -G"Visual Studio 16 2019" -DCMAKE_BUILD_TYPE=Release {} .'.format(args) ) | ||
| os.system("cmake.exe --build .") | ||
| else: | ||
| os.system("cmake --clean .") | ||
| os.system("cmake {} . ".format(args)) | ||
| os.system("cmake --build .") | ||
| os.chdir(cwd) | ||
|
|
||
|
|
||
| def build_dependencies(): | ||
| zlib_path = os.path.join(third_party_dir, "zlib") | ||
| if sys.platform == "win32": | ||
| libpng_cmake_options="-DPNG_BUILD_ZLIB=ON -DZLIB_INCLUDE_DIR:PATH={zlib_path} -DZLIB_LIBRARY:FILEPATH={zlib_path}/Debug/zlibd.dll".format(zlib_path=zlib_path) | ||
| elif sys.platform == "darwin": | ||
| libpng_cmake_options="-DPNG_BUILD_ZLIB=ON -DZLIB_INCLUDE_DIR:PATH={zlib_path} -DZLIB_LIBRARY:FILEPATH={zlib_path}/libz.dylib".format(zlib_path=zlib_path) | ||
| else : | ||
| libpng_cmake_options="-DPNG_BUILD_ZLIB=ON -DZLIB_INCLUDE_DIR:PATH={zlib_path} -DZLIB_LIBRARY:FILEPATH={zlib_path}/libz.so".format(zlib_path=zlib_path) | ||
| _build_cmake_dependency("./third_party/zlib") | ||
| # Deactivate SIMD on macOS | ||
| jpeg_turbo_options = "" | ||
| if sys.platform == "darwin": | ||
| jpeg_turbo_options ="-DWITH_SIMD=0" | ||
| _build_cmake_dependency("./third_party/libpng", libpng_cmake_options) | ||
| _build_cmake_dependency("./third_party/libjpeg-turbo", jpeg_turbo_options) | ||
|
|
||
|
|
||
| def get_extensions(): | ||
| this_dir = os.path.dirname(os.path.abspath(__file__)) | ||
| extensions_dir = os.path.join(this_dir, 'torchvision', 'csrc') | ||
|
|
||
| main_file = glob.glob(os.path.join(extensions_dir, '*.cpp')) | ||
| source_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', '*.cpp')) | ||
| source_image_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', 'image', '*.cpp')) | ||
| source_cuda = glob.glob(os.path.join(extensions_dir, 'cuda', '*.cu')) | ||
|
|
||
| sources = main_file + source_cpu + source_image_cpu | ||
| extension = CppExtension | ||
|
|
||
| compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1' | ||
|
|
@@ -142,7 +176,9 @@ def get_extensions(): | |
| extension( | ||
| 'torchvision._C', | ||
| sources, | ||
| libraries= third_party_libraries, | ||
| include_dirs=include_dirs + third_party_search_directories, | ||
| library_dirs=third_party_search_directories, | ||
| define_macros=define_macros, | ||
| extra_compile_args=extra_compile_args, | ||
| ) | ||
|
|
@@ -196,29 +232,32 @@ def run(self): | |
| # It's an old-style class in Python 2.7... | ||
| distutils.command.clean.clean.run(self) | ||
|
|
||
| def build_ext_with_dependencies(self): | ||
| build_dependencies() | ||
| return BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False )(self) | ||
|
|
||
| setup( | ||
| # Metadata | ||
| name=package_name, | ||
| version=version, | ||
| author='PyTorch Core Team', | ||
| author_email='[email protected]', | ||
| url='https://github.com/pytorch/vision', | ||
| description='image and video datasets and models for torch deep learning', | ||
| long_description=readme, | ||
| license='BSD', | ||
| # Package info | ||
| packages=find_packages(exclude=('test',)), | ||
| zip_safe=False, | ||
| install_requires=requirements, | ||
| extras_require={ | ||
| "scipy": ["scipy"], | ||
| }, | ||
| ext_modules=get_extensions(), | ||
| cmdclass={ | ||
| 'build_ext': build_ext_with_dependencies, | ||
| 'clean': clean, | ||
| } | ||
| ) | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| import os | ||
| import unittest | ||
|
|
||
| import torch | ||
| from PIL import Image | ||
| from torchvision.io.image import read_png, decode_png, decode_jpeg, read_jpeg | ||
| import numpy as np | ||
|
|
||
| IMAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "imagefolder") | ||
|
|
||
|
|
||
| def get_images(directory, img_ext): | ||
| assert os.path.isdir(directory) | ||
| for root, dir, files in os.walk(directory): | ||
| for fl in files: | ||
| _, ext = os.path.splitext(fl) | ||
| if ext == img_ext: | ||
| yield os.path.join(root, fl) | ||
|
|
||
|
|
||
| class ImageTester(unittest.TestCase): | ||
| def test_read_png(self): | ||
| for img_path in get_images(IMAGE_DIR, "png"): | ||
| img_pil = torch.from_numpy(np.array(Image.open(img_path))) | ||
| img_lpng = read_png(img_path) | ||
| self.assertTrue(torch.all(img_lpng == img_pil)) | ||
|
There was a problem hiding this comment. It might be better to use assertEqual since it checks for matching dtypes, etc. etc. ; unless you have a specific reason not to? |
||
|
|
||
| def test_decode_png(self): | ||
| for img_path in get_images(IMAGE_DIR, "png"): | ||
| img_pil = torch.from_numpy(np.array(Image.open(img_path))) | ||
| size = os.path.getsize(img_path) | ||
| img_lpng = decode_png(torch.from_file(img_path, dtype=torch.uint8, size=size)) | ||
| self.assertTrue(torch.all(img_lpng == img_pil)) | ||
|
|
||
| def test_read_jpeg(self): | ||
| for img_path in get_images(IMAGE_DIR, "jpg"): | ||
| img_pil = torch.from_numpy(np.array(Image.open(img_path))) | ||
| img_ljpeg = read_jpeg(img_path) | ||
| self.assertTrue(torch.all(img_ljpeg == img_pil)) | ||
|
|
||
| def test_decode_jpeg(self): | ||
| for img_path in get_images(IMAGE_DIR, "jpg"): | ||
| img_pil = torch.from_numpy(np.array(Image.open(img_path))) | ||
| size = os.path.getsize(img_path) | ||
| img_ljpeg = decode_png(torch.from_file(img_path, dtype=torch.uint8, size=size)) | ||
| self.assertTrue(torch.all(img_ljpeg == img_pil)) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| unittest.main() | ||
Submodule libjpeg-turbo
added at
f81833
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is how PyTorch set this up in case you didn't see it. There might be some patterns that could be borrowed from this.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi,
I took a look. I'll grab some code.
Thanks