Skip to content

Commit 44f4d81

Browse files
vfdev-5vfdev-5
committed
Addressed PR review
1 parent 4401c2e commit 44f4d81

File tree

4 files changed

+125
-49
lines changed

4 files changed

+125
-49
lines changed

‎test/test_image.py

+19
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,25 @@ def test_decode_jpeg_with_exif_orientation(tmpdir, orientation):
119119
torch.testing.assert_close(expected, output)
120120

121121

122+
@pytest.mark.parametrize("size", [65533, 1, 7, 10, 23, 33])
123+
def test_invalid_exif(tmpdir, size):
124+
# Inspired from a PIL test:
125+
# https://github.com/python-pillow/Pillow/blob/8f63748e50378424628155994efd7e0739a4d1d1/Tests/test_file_jpeg.py#L299
126+
fp = os.path.join(tmpdir, "invalid_exif.jpg")
127+
t = torch.randint(0, 256, size=(3, 256, 257), dtype=torch.uint8)
128+
im = F.to_pil_image(t)
129+
im.save(fp, "JPEG", exif=b"1" * size)
130+
131+
data = read_file(fp)
132+
output = decode_image(data, apply_exif_orientation=True)
133+
134+
pimg = Image.open(fp)
135+
pimg = ImageOps.exif_transpose(pimg)
136+
137+
expected = F.pil_to_tensor(pimg)
138+
torch.testing.assert_close(expected, output)
139+
140+
122141
def test_decode_jpeg_errors():
123142
with pytest.raises(RuntimeError, match="Expected a non empty 1-dimensional tensor"):
124143
decode_jpeg(torch.empty((100, 1), dtype=torch.uint8))

‎torchvision/csrc/io/image/cpu/decode_image.cpp

-3
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,6 @@ torch::Tensor decode_image(
2727
if (memcmp(jpeg_signature, datap, 3) == 0) {
2828
return decode_jpeg(data, mode, apply_exif_orientation);
2929
} else if (memcmp(png_signature, datap, 4) == 0) {
30-
TORCH_CHECK(
31-
!apply_exif_orientation,
32-
"Unsupported option apply_exif_orientation=true for PNG")
3330
return decode_png(data, mode);
3431
} else {
3532
TORCH_CHECK(

‎torchvision/csrc/io/image/cpu/decode_jpeg.cpp

+2-45
Original file line numberDiff line numberDiff line change
@@ -198,52 +198,9 @@ torch::Tensor decode_jpeg(
198198
jpeg_calc_output_dimensions(&cinfo);
199199
}
200200

201-
int exif_orientation = 0;
201+
int exif_orientation = -1;
202202
if (apply_exif_orientation) {
203-
// Check for Exif marker APP1
204-
jpeg_saved_marker_ptr exif_marker = 0;
205-
jpeg_saved_marker_ptr cmarker = cinfo.marker_list;
206-
while (cmarker && exif_marker == 0) {
207-
if (cmarker->marker == APP1) {
208-
exif_marker = cmarker;
209-
}
210-
cmarker = cmarker->next;
211-
}
212-
213-
if (exif_marker) {
214-
// Code below is inspired from OpenCV
215-
// https://github.com/opencv/opencv/blob/097891e311fae1d8354eb092a0fd0171e630d78c/modules/modules/imgcodecs/src/exif.cpp
216-
217-
// Bytes from Exif size field to the first TIFF header
218-
constexpr size_t start_offset = 6;
219-
if (exif_marker->data_length > start_offset) {
220-
auto* exif_data_ptr = exif_marker->data + start_offset;
221-
auto size = exif_marker->data_length - start_offset;
222-
std::vector<unsigned char> exif_data_vec(
223-
exif_data_ptr, exif_data_ptr + size);
224-
225-
auto endianness = get_endianness(exif_data_vec);
226-
227-
// Checking whether Tag Mark (0x002A) correspond to one contained in the
228-
// Jpeg file
229-
uint16_t tag_mark = get_uint16(exif_data_vec, endianness, 2);
230-
if (tag_mark == REQ_EXIF_TAG_MARK) {
231-
auto offset = get_uint32(exif_data_vec, endianness, 4);
232-
size_t num_entry = get_uint16(exif_data_vec, endianness, offset);
233-
offset += 2; // go to start of tag fields
234-
constexpr size_t tiff_field_size = 12;
235-
for (size_t entry = 0; entry < num_entry; entry++) {
236-
// Here we just search for orientation tag and parse it
237-
auto tag_num = get_uint16(exif_data_vec, endianness, offset);
238-
if (tag_num == ORIENTATION_EXIF_TAG) {
239-
exif_orientation =
240-
get_uint16(exif_data_vec, endianness, offset + 8);
241-
}
242-
offset += tiff_field_size;
243-
}
244-
}
245-
}
246-
}
203+
exif_orientation = fetch_exif_orientation(&cinfo);
247204
}
248205

249206
jpeg_start_decompress(&cinfo);

‎torchvision/csrc/io/image/cpu/exif.h

+104-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,55 @@
1+
/*M///////////////////////////////////////////////////////////////////////////////////////
2+
//
3+
// IMPORTANT: READ BEFORE DOWNLOADING, COPYING, INSTALLING OR USING.
4+
//
5+
// By downloading, copying, installing or using the software you agree to this
6+
license.
7+
// If you do not agree to this license, do not download, install,
8+
// copy or use the software.
9+
//
10+
//
11+
// License Agreement
12+
// For Open Source Computer Vision Library
13+
//
14+
// Copyright (C) 2000-2008, Intel Corporation, all rights reserved.
15+
// Copyright (C) 2009, Willow Garage Inc., all rights reserved.
16+
// Third party copyrights are property of their respective owners.
17+
//
18+
// Redistribution and use in source and binary forms, with or without
19+
modification,
20+
// are permitted provided that the following conditions are met:
21+
//
22+
// * Redistribution's of source code must retain the above copyright notice,
23+
// this list of conditions and the following disclaimer.
24+
//
25+
// * Redistribution's in binary form must reproduce the above copyright
26+
notice,
27+
// this list of conditions and the following disclaimer in the documentation
28+
// and/or other materials provided with the distribution.
29+
//
30+
// * The name of the copyright holders may not be used to endorse or promote
31+
products
32+
// derived from this software without specific prior written permission.
33+
//
34+
// This software is provided by the copyright holders and contributors "as is"
35+
and
36+
// any express or implied warranties, including, but not limited to, the implied
37+
// warranties of merchantability and fitness for a particular purpose are
38+
disclaimed.
39+
// In no event shall the Intel Corporation or contributors be liable for any
40+
direct,
41+
// indirect, incidental, special, exemplary, or consequential damages
42+
// (including, but not limited to, procurement of substitute goods or services;
43+
// loss of use, data, or profits; or business interruption) however caused
44+
// and on any theory of liability, whether in contract, strict liability,
45+
// or tort (including negligence or otherwise) arising in any way out of
46+
// the use of this software, even if advised of the possibility of such damage.
47+
//
48+
//M*/
149
#pragma once
50+
#include <jpeglib.h>
251
#include <torch/types.h>
52+
#include <vector>
353

454
namespace vision {
555
namespace image {
@@ -13,7 +63,7 @@ constexpr uint16_t ORIENTATION_EXIF_TAG = 0x0112;
1363
constexpr uint16_t INCORRECT_TAG = -1;
1464

1565
// Functions in this module are taken from OpenCV
16-
// https://github.com/opencv/opencv/blob/097891e311fae1d8354eb092a0fd0171e630d78c/modules/modules/imgcodecs/src/exif.cpp
66+
// https://github.com/opencv/opencv/blob/097891e311fae1d8354eb092a0fd0171e630d78c/modules/imgcodecs/src/exif.cpp
1767
inline uint16_t get_endianness(const std::vector<unsigned char>& exif_data) {
1868
if ((exif_data.size() < 1) ||
1969
(exif_data.size() > 1 && exif_data[0] != exif_data[1])) {
@@ -58,6 +108,59 @@ inline uint32_t get_uint32(
58108
(exif_data[offset + 2] << 8) + exif_data[offset + 3];
59109
}
60110

111+
inline int fetch_exif_orientation(j_decompress_ptr cinfo) {
112+
int exif_orientation = -1;
113+
// Check for Exif marker APP1
114+
jpeg_saved_marker_ptr exif_marker = 0;
115+
jpeg_saved_marker_ptr cmarker = cinfo->marker_list;
116+
while (cmarker && exif_marker == 0) {
117+
if (cmarker->marker == APP1) {
118+
exif_marker = cmarker;
119+
}
120+
cmarker = cmarker->next;
121+
}
122+
123+
if (exif_marker) {
124+
// Code below is inspired from OpenCV
125+
// https://github.com/opencv/opencv/blob/097891e311fae1d8354eb092a0fd0171e630d78c/modules/imgcodecs/src/exif.cpp
126+
127+
// Bytes from Exif size field to the first TIFF header
128+
constexpr size_t start_offset = 6;
129+
if (exif_marker->data_length > start_offset) {
130+
auto* exif_data_ptr = exif_marker->data + start_offset;
131+
auto size = exif_marker->data_length - start_offset;
132+
std::vector<unsigned char> exif_data_vec(
133+
exif_data_ptr, exif_data_ptr + size);
134+
135+
auto endianness = get_endianness(exif_data_vec);
136+
137+
// Checking whether Tag Mark (0x002A) correspond to one contained in the
138+
// Jpeg file
139+
uint16_t tag_mark = get_uint16(exif_data_vec, endianness, 2);
140+
if (tag_mark == REQ_EXIF_TAG_MARK) {
141+
auto offset = get_uint32(exif_data_vec, endianness, 4);
142+
size_t num_entry = get_uint16(exif_data_vec, endianness, offset);
143+
offset += 2; // go to start of tag fields
144+
constexpr size_t tiff_field_size = 12;
145+
for (size_t entry = 0; entry < num_entry; entry++) {
146+
// Here we just search for orientation tag and parse it
147+
auto tag_num = get_uint16(exif_data_vec, endianness, offset);
148+
if (tag_num == INCORRECT_TAG) {
149+
break;
150+
}
151+
if (tag_num == ORIENTATION_EXIF_TAG) {
152+
exif_orientation =
153+
get_uint16(exif_data_vec, endianness, offset + 8);
154+
break;
155+
}
156+
offset += tiff_field_size;
157+
}
158+
}
159+
}
160+
}
161+
return exif_orientation;
162+
}
163+
61164
constexpr uint16_t IMAGE_ORIENTATION_TL = 1; // normal orientation
62165
constexpr uint16_t IMAGE_ORIENTATION_TR = 2; // needs horizontal flip
63166
constexpr uint16_t IMAGE_ORIENTATION_BR = 3; // needs 180 rotation

0 commit comments

Comments
 (0)