From 4401c2e62794767a0d7f138c4ff54859a9630664 Mon Sep 17 00:00:00 2001 From: vfdev-5 Date: Fri, 23 Feb 2024 18:45:51 +0100 Subject: [PATCH] Fixed smoke test and set signature with dafault apply_exif_orientation in the op registration --- torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 2 +- torchvision/csrc/io/image/image.cpp | 4 ++-- torchvision/io/image.py | 9 +++++++-- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index e8a52c69593..d70f9b024b4 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -212,7 +212,7 @@ torch::Tensor decode_jpeg( if (exif_marker) { // Code below is inspired from OpenCV - // https://github.dev/opencv/opencv/blob/097891e311fae1d8354eb092a0fd0171e630d78c/modules/modules/imgcodecs/src/exif.cpp + // https://github.com/opencv/opencv/blob/097891e311fae1d8354eb092a0fd0171e630d78c/modules/modules/imgcodecs/src/exif.cpp // Bytes from Exif size field to the first TIFF header constexpr size_t start_offset = 6; diff --git a/torchvision/csrc/io/image/image.cpp b/torchvision/csrc/io/image/image.cpp index b5952739a7a..de8979f32bb 100644 --- a/torchvision/csrc/io/image/image.cpp +++ b/torchvision/csrc/io/image/image.cpp @@ -23,11 +23,11 @@ static auto registry = torch::RegisterOperators() .op("image::decode_png", &decode_png) .op("image::encode_png", &encode_png) - .op("image::decode_jpeg", &decode_jpeg) + .op("image::decode_jpeg(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor", &decode_jpeg) .op("image::encode_jpeg", &encode_jpeg) .op("image::read_file", &read_file) .op("image::write_file", &write_file) - .op("image::decode_image", &decode_image) + .op("image::decode_image(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor", &decode_image) .op("image::decode_jpeg_cuda", &decode_jpeg_cuda) .op("image::_jpeg_version", &_jpeg_version) .op("image::_is_compiled_against_turbo", &_is_compiled_against_turbo); diff --git a/torchvision/io/image.py b/torchvision/io/image.py index 8dd61d0d469..de904804078 100644 --- a/torchvision/io/image.py +++ b/torchvision/io/image.py @@ -130,7 +130,10 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6): def decode_jpeg( - input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED, device: str = "cpu" + input: torch.Tensor, + mode: ImageReadMode = ImageReadMode.UNCHANGED, + device: str = "cpu", + apply_exif_orientation: bool = False, ) -> torch.Tensor: """ Decodes a JPEG image into a 3 dimensional RGB or grayscale Tensor. @@ -157,6 +160,8 @@ def decode_jpeg( .. warning:: There is a memory leak in the nvjpeg library for CUDA versions < 11.6. Make sure to rely on CUDA 11.6 or above before using ``device="cuda"``. + apply_exif_orientation (bool): apply EXIF orientation transformation to the output tensor. + Default: False. Only implemented for JPEG format on CPU. Returns: output (Tensor[image_channels, image_height, image_width]) @@ -167,7 +172,7 @@ def decode_jpeg( if device.type == "cuda": output = torch.ops.image.decode_jpeg_cuda(input, mode.value, device) else: - output = torch.ops.image.decode_jpeg(input, mode.value) + output = torch.ops.image.decode_jpeg(input, mode.value, apply_exif_orientation) return output