Skip to content

Commit

Permalink
Fixed smoke test and set signature with dafault apply_exif_orientatio…
Browse files Browse the repository at this point in the history
…n in the op registration
  • Loading branch information
vfdev-5 committed Feb 23, 2024
1 parent a40be0a commit 4401c2e
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 5 deletions.
2 changes: 1 addition & 1 deletion torchvision/csrc/io/image/cpu/decode_jpeg.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions torchvision/csrc/io/image/image.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
9 changes: 7 additions & 2 deletions torchvision/io/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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])
Expand All @@ -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


Expand Down

0 comments on commit 4401c2e

Please sign in to comment.