diff --git a/test/assets/damaged_jpeg/TensorFlow-LICENSE b/test/assets/damaged_jpeg/TensorFlow-LICENSE new file mode 100644 index 00000000000..c7563fe4e5b --- /dev/null +++ b/test/assets/damaged_jpeg/TensorFlow-LICENSE @@ -0,0 +1,13 @@ + Copyright 2019 The TensorFlow Authors. All rights reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/test/assets/damaged_jpeg/bad_huffman.jpg b/test/assets/damaged_jpeg/bad_huffman.jpg new file mode 100644 index 00000000000..ef5b6f12c55 Binary files /dev/null and b/test/assets/damaged_jpeg/bad_huffman.jpg differ diff --git a/test/assets/damaged_jpeg/corrupt.jpg b/test/assets/damaged_jpeg/corrupt.jpg new file mode 100644 index 00000000000..5e2fe6c56f5 Binary files /dev/null and b/test/assets/damaged_jpeg/corrupt.jpg differ diff --git a/test/assets/damaged_jpeg/corrupt34_2.jpg b/test/assets/damaged_jpeg/corrupt34_2.jpg new file mode 100644 index 00000000000..4211155c455 Binary files /dev/null and b/test/assets/damaged_jpeg/corrupt34_2.jpg differ diff --git a/test/assets/damaged_jpeg/corrupt34_3.jpg b/test/assets/damaged_jpeg/corrupt34_3.jpg new file mode 100644 index 00000000000..c1c2a9d1e1e Binary files /dev/null and b/test/assets/damaged_jpeg/corrupt34_3.jpg differ diff --git a/test/assets/damaged_jpeg/corrupt34_4.jpg b/test/assets/damaged_jpeg/corrupt34_4.jpg new file mode 100644 index 00000000000..b8e7308ba00 Binary files /dev/null and b/test/assets/damaged_jpeg/corrupt34_4.jpg differ diff --git a/test/test_image.py b/test/test_image.py index 09a591d19ef..0bf3daf5528 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -1,4 +1,5 @@ import os +import glob import unittest import sys @@ -10,11 +11,15 @@ IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") IMAGE_DIR = os.path.join(IMAGE_ROOT, "fakedata", "imagefolder") +DAMAGED_JPEG = os.path.join(IMAGE_ROOT, 'damaged_jpeg') def get_images(directory, img_ext): assert os.path.isdir(directory) for root, _, files in os.walk(directory): + if os.path.basename(root) == 'damaged_jpeg': + continue + for fl in files: _, ext = os.path.splitext(fl) if ext == img_ext: @@ -44,6 +49,21 @@ def test_decode_jpeg(self): with self.assertRaises(RuntimeError): decode_jpeg(torch.empty((100), dtype=torch.uint8)) + def test_damaged_images(self): + # Test image with bad Huffman encoding (should not raise) + bad_huff = os.path.join(DAMAGED_JPEG, 'bad_huffman.jpg') + try: + _ = read_jpeg(bad_huff) + except RuntimeError: + self.assertTrue(False) + + # Truncated images should raise an exception + truncated_images = glob.glob( + os.path.join(DAMAGED_JPEG, 'corrupt*.jpg')) + for image_path in truncated_images: + with self.assertRaises(RuntimeError): + read_jpeg(image_path) + def test_read_png(self): # Check across .png for img_path in get_images(IMAGE_DIR, ".png"): diff --git a/torchvision/csrc/cpu/image/readjpeg_cpu.cpp b/torchvision/csrc/cpu/image/readjpeg_cpu.cpp index 4954b2c2474..b3e3d2ffa5a 100644 --- a/torchvision/csrc/cpu/image/readjpeg_cpu.cpp +++ b/torchvision/csrc/cpu/image/readjpeg_cpu.cpp @@ -48,7 +48,10 @@ static void torch_jpeg_init_source(j_decompress_ptr cinfo) {} static boolean torch_jpeg_fill_input_buffer(j_decompress_ptr cinfo) { torch_jpeg_mgr* src = (torch_jpeg_mgr*)cinfo->src; - // No more data. Probably an incomplete image; just output EOI. + // No more data. Probably an incomplete image; Raise exception. + torch_jpeg_error_ptr myerr = (torch_jpeg_error_ptr)cinfo->err; + strcpy(jpegLastErrorMsg, "Image is incomplete or truncated"); + longjmp(myerr->setjmp_buffer, 1); src->pub.next_input_byte = EOI_BUFFER; src->pub.bytes_in_buffer = 1; return TRUE;