Skip to content

Commit 9d04ed8

Browse files
r-zeniner-zenine
committed
Adds readpng features
1 parent b1ef906 commit 9d04ed8

18 files changed

+165
-73
lines changed

‎.circleci/config.yml

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,9 @@ jobs:
128128
ca-certificates \
129129
curl \
130130
gnupg-agent \
131-
software-properties-common
131+
software-properties-common \
132+
libpng16-16 \
133+
libpng-dev
132134
133135
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
134136
@@ -219,7 +221,8 @@ jobs:
219221
curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
220222
sh conda.sh -b
221223
source $HOME/miniconda3/bin/activate
222-
packaging/build_wheel.sh
224+
conda install -yq libpng
225+
packaging/build_wheel_macos.sh
223226
- store_artifacts:
224227
path: dist
225228
- persist_to_workspace:
@@ -239,6 +242,7 @@ jobs:
239242
sh conda.sh -b
240243
source $HOME/miniconda3/bin/activate
241244
conda install -yq conda-build
245+
conda install -yq libpng
242246
packaging/build_conda.sh
243247
- store_artifacts:
244248
path: /Users/distiller/miniconda3/conda-bld/osx-64

‎.circleci/config.yml.in

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,9 @@ jobs:
128128
ca-certificates \
129129
curl \
130130
gnupg-agent \
131-
software-properties-common
131+
software-properties-common \
132+
libpng16-16 \
133+
libpng-dev
132134

133135
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
134136

@@ -219,7 +221,8 @@ jobs:
219221
curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
220222
sh conda.sh -b
221223
source $HOME/miniconda3/bin/activate
222-
packaging/build_wheel.sh
224+
conda install -yq libpng
225+
packaging/build_wheel_macos.sh
223226
- store_artifacts:
224227
path: dist
225228
- persist_to_workspace:
@@ -239,6 +242,7 @@ jobs:
239242
sh conda.sh -b
240243
source $HOME/miniconda3/bin/activate
241244
conda install -yq conda-build
245+
conda install -yq libpng
242246
packaging/build_conda.sh
243247
- store_artifacts:
244248
path: /Users/distiller/miniconda3/conda-bld/osx-64

‎.travis.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ matrix:
3434

3535
before_install:
3636
- sudo apt-get update
37+
- sudo apt-get install -y libpng16-16 libpng16-dev
3738
- wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
3839
- bash miniconda.sh -b -p $HOME/miniconda
3940
- export PATH="$HOME/miniconda/bin:$PATH"

‎CMakeLists.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ file(GLOB HEADERS torchvision/csrc/vision.h)
1111
file(GLOB MODELS_HEADERS torchvision/csrc/models/*.h)
1212
file(GLOB MODELS_SOURCES torchvision/csrc/models/*.h torchvision/csrc/models/*.cpp)
1313

14-
file(GLOB IMAGE_HEADERS torchvision/csrc/image/*.h)
15-
file(GLOB IMAGE_SOURCES torchvision/csrc/image/*.h torchvision/csrc/image/*.cpp)
14+
file(GLOB IMAGE_HEADERS torchvision/csrc/image.h)
15+
file(GLOB IMAGE_SOURCES torchvision/csrc/cpu/image/*.h torchvision/csrc/cpu/image/*.cpp)
1616

1717
add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${IMAGE_SOURCES})
1818
target_link_libraries(${PROJECT_NAME} PUBLIC "${PNG_LIBRARY}")

‎packaging/build_conda.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,5 @@ export SOURCE_ROOT_DIR="$PWD"
1010
setup_conda_pytorch_constraint
1111
setup_conda_cudatoolkit_constraint
1212
setup_visual_studio_constraint
13+
conda install -yq libpng
1314
conda build $CONDA_CHANNEL_FLAGS -c defaults -c conda-forge --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

‎packaging/build_wheel.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
77
export BUILD_TYPE=wheel
88
setup_env 0.5.0
99
setup_wheel_python
10+
yum install -y libpng libpng-devel
1011
pip_install numpy pyyaml future ninja
1112
# TODO remove after https://github.com/pytorch/pytorch/pull/27282 gets merged
1213
pip_install six

‎packaging/build_wheel_macos.sh

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
#!/bin/bash
2+
set -ex
3+
4+
script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
5+
. "$script_dir/pkg_helpers.bash"
6+
7+
export BUILD_TYPE=wheel
8+
setup_env 0.5.0
9+
setup_wheel_python
10+
pip_install numpy pyyaml future ninja
11+
# TODO remove after https://github.com/pytorch/pytorch/pull/27282 gets merged
12+
pip_install six
13+
setup_pip_pytorch_version
14+
python setup.py clean
15+
IS_WHEEL=1 python setup.py bdist_wheel

‎packaging/torchvision/meta.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ requirements:
1212
host:
1313
- python
1414
- setuptools
15+
- libpng >=1.6.37
1516
{{ environ.get('CONDA_PYTORCH_BUILD_CONSTRAINT') }}
1617
{{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
1718
{{ environ.get('CONDA_CPUONLY_FEATURE') }}

‎packaging/wheel/linux_manywheel.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ rm -rf vision
3232
git clone https://github.com/pytorch/vision
3333

3434
cd /tmp/vision
35-
35+
yum install -y libpng libpng-devel
3636
for PYDIR in "${python_installations[@]}"; do
3737
export PATH=$PYDIR/bin:$OLD_PATH
3838
pip install --upgrade pip

‎setup.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,10 @@ def get_extensions():
8383

8484
main_file = glob.glob(os.path.join(extensions_dir, '*.cpp'))
8585
source_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', '*.cpp'))
86+
source_image_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', 'image', '*.cpp'))
8687
source_cuda = glob.glob(os.path.join(extensions_dir, 'cuda', '*.cu'))
8788

88-
sources = main_file + source_cpu
89+
sources = main_file + source_cpu + source_image_cpu
8990
extension = CppExtension
9091

9192
compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1'
@@ -142,6 +143,7 @@ def get_extensions():
142143
extension(
143144
'torchvision._C',
144145
sources,
146+
libraries=['png'],
145147
include_dirs=include_dirs,
146148
define_macros=define_macros,
147149
extra_compile_args=extra_compile_args,

‎test/test_image.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import os
2+
import unittest
3+
4+
import torch
5+
from PIL import Image
6+
from torchvision.io.image import read_png, decode_png
7+
import numpy as np
8+
9+
IMAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "imagefolder")
10+
11+
12+
def get_png_images(directory):
13+
assert os.path.isdir(directory)
14+
for root, dir, files in os.walk(directory):
15+
for fl in files:
16+
_, ext = os.path.splitext(fl)
17+
if ext == ".png":
18+
yield os.path.join(root, fl)
19+
20+
21+
class ImageTester(unittest.TestCase):
22+
def test_read_png(self):
23+
for img_path in get_png_images(IMAGE_DIR):
24+
img_pil = torch.from_numpy(np.array(Image.open(img_path)))
25+
img_lpng = read_png(img_path)
26+
self.assertTrue(torch.all(img_lpng == img_pil))
27+
28+
def test_decode_png(self):
29+
for img_path in get_png_images(IMAGE_DIR):
30+
img_pil = torch.from_numpy(np.array(Image.open(img_path)))
31+
size = os.path.getsize(img_path)
32+
img_lpng = decode_png(torch.from_file(img_path, dtype=torch.uint8, size=size))
33+
self.assertTrue(torch.all(img_lpng == img_pil))
34+
35+
36+
if __name__ == '__main__':
37+
unittest.main()

‎torchvision/csrc/image/readpng.cpp renamed to ‎torchvision/csrc/cpu/image/readpng_cpu.cpp

Lines changed: 33 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,43 @@
1-
#include "readpng.h"
2-
3-
#include <sxx/sxx.h>
1+
#include "readpng_cpu.h"
42

53
#include <png.h>
4+
#include <setjmp.h>
5+
#include <string>
66

7-
namespace torch {
8-
namespace vision {
9-
namespace image {
10-
namespace impl {
11-
bool is_png(const void* data) {
12-
return png_sig_cmp(png_const_bytep(data), 0, 8) == 0;
13-
}
14-
15-
torch::Tensor readpng(const void* data) {
16-
struct Reader {
17-
png_const_bytep ptr;
18-
} reader;
19-
20-
reader.ptr = png_const_bytep(data) + 8;
21-
auto read_callback =
22-
[](png_structp png_ptr, png_bytep output, png_size_t bytes) {
23-
auto reader = static_cast<Reader*>(png_get_io_ptr(png_ptr));
24-
std::copy(reader->ptr, reader->ptr + bytes, output);
25-
reader->ptr += bytes;
26-
};
27-
7+
torch::Tensor decodePNG(const torch::Tensor& data) {
288
auto png_ptr =
299
png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
30-
if (!png_ptr)
31-
return torch::tensor({0});
32-
10+
TORCH_CHECK(png_ptr, "libpng read structure allocation failed!")
3311
auto info_ptr = png_create_info_struct(png_ptr);
3412
if (!info_ptr) {
3513
png_destroy_read_struct(&png_ptr, nullptr, nullptr);
36-
return torch::tensor({0});
14+
// Seems redundant with the if statement. done here to avoid leaking memory.
15+
TORCH_CHECK(info_ptr, "libpng info structure allocation failed!")
3716
}
3817

18+
auto datap = data.accessor<unsigned char, 1>().data();
19+
3920
if (setjmp(png_jmpbuf(png_ptr)) != 0) {
4021
png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
41-
return torch::tensor({0});
22+
auto errptr = static_cast<png_const_charp>(png_get_error_ptr(png_ptr));
23+
TORCH_CHECK(false, "Internal error.");
4224
}
25+
auto is_png = !png_sig_cmp(datap, 0, 8);
26+
TORCH_CHECK(is_png, "Content is not png!")
4327

44-
png_set_read_fn(png_ptr, &reader, read_callback);
28+
struct Reader {
29+
png_const_bytep ptr;
30+
} reader;
31+
reader.ptr = png_const_bytep(datap) + 8;
32+
33+
auto read_callback =
34+
[](png_structp png_ptr, png_bytep output, png_size_t bytes) {
35+
auto reader = static_cast<Reader*>(png_get_io_ptr(png_ptr));
36+
std::copy(reader->ptr, reader->ptr + bytes, output);
37+
reader->ptr += bytes;
38+
};
4539
png_set_sig_bytes(png_ptr, 8);
40+
png_set_read_fn(png_ptr, &reader, read_callback);
4641
png_read_info(png_ptr, info_ptr);
4742

4843
png_uint_32 width, height;
@@ -58,26 +53,24 @@ torch::Tensor readpng(const void* data) {
5853
nullptr,
5954
nullptr);
6055

61-
if (retval != 1 || color_type != PNG_COLOR_TYPE_RGB) {
56+
if (retval != 1) {
57+
png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
58+
TORCH_CHECK(retval == 1, "Could read image metadata from content.")
59+
}
60+
if (color_type != PNG_COLOR_TYPE_RGB) {
6261
png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
63-
return torch::tensor({0});
62+
TORCH_CHECK(
63+
color_type == PNG_COLOR_TYPE_RGB, "Non RGB images are not supported.")
6464
}
6565

6666
auto tensor =
6767
torch::empty({int64_t(height), int64_t(width), int64_t(3)}, torch::kU8);
68-
auto ptr = tensor.data<uint8_t>();
68+
auto ptr = tensor.accessor<uint8_t, 3>().data();
6969
auto bytes = png_get_rowbytes(png_ptr, info_ptr);
70-
7170
for (decltype(height) i = 0; i < height; ++i) {
7271
png_read_row(png_ptr, ptr, nullptr);
7372
ptr += bytes;
7473
}
75-
7674
png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
7775
return tensor;
7876
}
79-
80-
} // namespace impl
81-
} // namespace image
82-
} // namespace vision
83-
} // namespace torch

‎torchvision/csrc/cpu/image/readpng_cpu.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
#pragma once
2+
3+
#include <torch/torch.h>
4+
#include <string>
5+
6+
torch::Tensor decodePNG(const torch::Tensor& data);

‎torchvision/csrc/image.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
#pragma once
2+
3+
#include "cpu/image/readpng_cpu.h"

‎torchvision/csrc/image/image.h

Lines changed: 0 additions & 6 deletions
This file was deleted.

‎torchvision/csrc/image/readpng.h

Lines changed: 0 additions & 19 deletions
This file was deleted.

‎torchvision/csrc/vision.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include "ROIAlign.h"
1212
#include "ROIPool.h"
1313
#include "empty_tensor_op.h"
14+
#include "image.h"
1415
#include "nms.h"
1516

1617
// If we are in a Windows environment, we need to define
@@ -49,4 +50,5 @@ static auto registry =
4950
.op("torchvision::ps_roi_align", &ps_roi_align)
5051
.op("torchvision::ps_roi_pool", &ps_roi_pool)
5152
.op("torchvision::deform_conv2d", &deform_conv2d)
53+
.op("torchvision::decode_png", &decodePNG)
5254
.op("torchvision::_cuda_version", &_cuda_version);

‎torchvision/io/image.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
import torch
2+
from torch import nn, Tensor
3+
import os
4+
5+
6+
def decode_png(input):
7+
# type: (Tensor) -> Tensor
8+
"""
9+
Decodes a PNG image into a 3 dimensional RGB Tensor.
10+
The values of the output tensor are uint8 between 0 and 255.
11+
12+
Arguments:
13+
input (Tensor[1]): a one dimensional int8 tensor containing
14+
the raw bytes of the PNG image.
15+
16+
Returns:
17+
output (Tensor[image_width, image_height, 3])
18+
"""
19+
if not isinstance(input, torch.Tensor) or len(input) == 0:
20+
raise ValueError("Expected a non empty 1-dimensional tensor.")
21+
22+
if not input.dtype == torch.uint8:
23+
raise ValueError("Expected a torch.uint8 tensor.")
24+
output = torch.ops.torchvision.decode_png(input)
25+
return output
26+
27+
28+
def read_png(path):
29+
# type: (str) -> Tensor
30+
"""
31+
Reads a PNG image into a 3 dimensional RGB Tensor.
32+
The values of the output tensor are uint8 between 0 and 255.
33+
34+
Arguments:
35+
path (str): path of the PNG image.
36+
37+
Returns:
38+
output (Tensor[image_width, image_height, 3])
39+
"""
40+
if not os.path.isfile(path):
41+
raise ValueError("Excepted a valid file path.")
42+
43+
size = os.path.getsize(path)
44+
if size == 0:
45+
raise ValueError("Excepted a non empty file.")
46+
data = torch.from_file(path, dtype=torch.uint8, size=size)
47+
return decode_png(data)

0 commit comments

Comments
 (0)