511 lines
21 KiB
Python
511 lines
21 KiB
Python
from enum import Enum
|
|
from typing import List, Union
|
|
from warnings import warn
|
|
|
|
import torch
|
|
|
|
from ..extension import _load_library
|
|
from ..utils import _log_api_usage_once
|
|
|
|
|
|
try:
|
|
_load_library("image")
|
|
except (ImportError, OSError) as e:
|
|
warn(
|
|
f"Failed to load image Python extension: '{e}'"
|
|
f"If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. "
|
|
f"Otherwise, there might be something wrong with your environment. "
|
|
f"Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?"
|
|
)
|
|
|
|
|
|
class ImageReadMode(Enum):
|
|
"""Allow automatic conversion to RGB, RGBA, etc while decoding.
|
|
|
|
.. note::
|
|
|
|
You don't need to use this struct, you can just pass strings to all
|
|
``mode`` parameters, e.g. ``mode="RGB"``.
|
|
|
|
The different available modes are the following.
|
|
|
|
- UNCHANGED: loads the image as-is
|
|
- RGB: converts to RGB
|
|
- RGBA: converts to RGB with transparency (also aliased as RGB_ALPHA)
|
|
- GRAY: converts to grayscale
|
|
- GRAY_ALPHA: converts to grayscale with transparency
|
|
|
|
.. note::
|
|
|
|
Some decoders won't support all possible values, e.g. GRAY and
|
|
GRAY_ALPHA are only supported for PNG and JPEG images.
|
|
"""
|
|
|
|
UNCHANGED = 0
|
|
GRAY = 1
|
|
GRAY_ALPHA = 2
|
|
RGB = 3
|
|
RGB_ALPHA = 4
|
|
RGBA = RGB_ALPHA # Alias for convenience
|
|
|
|
|
|
def read_file(path: str) -> torch.Tensor:
|
|
"""
|
|
Return the bytes contents of a file as a uint8 1D Tensor.
|
|
|
|
Args:
|
|
path (str or ``pathlib.Path``): the path to the file to be read
|
|
|
|
Returns:
|
|
data (Tensor)
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(read_file)
|
|
data = torch.ops.image.read_file(str(path))
|
|
return data
|
|
|
|
|
|
def write_file(filename: str, data: torch.Tensor) -> None:
|
|
"""
|
|
Write the content of an uint8 1D tensor to a file.
|
|
|
|
Args:
|
|
filename (str or ``pathlib.Path``): the path to the file to be written
|
|
data (Tensor): the contents to be written to the output file
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(write_file)
|
|
torch.ops.image.write_file(str(filename), data)
|
|
|
|
|
|
def decode_png(
|
|
input: torch.Tensor,
|
|
mode: ImageReadMode = ImageReadMode.UNCHANGED,
|
|
apply_exif_orientation: bool = False,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Decodes a PNG image into a 3 dimensional RGB or grayscale Tensor.
|
|
|
|
The values of the output tensor are in uint8 in [0, 255] for most cases. If
|
|
the image is a 16-bit png, then the output tensor is uint16 in [0, 65535]
|
|
(supported from torchvision ``0.21``). Since uint16 support is limited in
|
|
pytorch, we recommend calling
|
|
:func:`torchvision.transforms.v2.functional.to_dtype()` with ``scale=True``
|
|
after this function to convert the decoded image into a uint8 or float
|
|
tensor.
|
|
|
|
Args:
|
|
input (Tensor[1]): a one dimensional uint8 tensor containing
|
|
the raw bytes of the PNG image.
|
|
mode (str or ImageReadMode): The mode to convert the image to, e.g. "RGB".
|
|
Default is "UNCHANGED". See :class:`~torchvision.io.ImageReadMode`
|
|
for available modes.
|
|
apply_exif_orientation (bool): apply EXIF orientation transformation to the output tensor.
|
|
Default: False.
|
|
|
|
Returns:
|
|
output (Tensor[image_channels, image_height, image_width])
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(decode_png)
|
|
if isinstance(mode, str):
|
|
mode = ImageReadMode[mode.upper()]
|
|
output = torch.ops.image.decode_png(input, mode.value, apply_exif_orientation)
|
|
return output
|
|
|
|
|
|
def encode_png(input: torch.Tensor, compression_level: int = 6) -> torch.Tensor:
|
|
"""
|
|
Takes an input tensor in CHW layout and returns a buffer with the contents
|
|
of its corresponding PNG file.
|
|
|
|
Args:
|
|
input (Tensor[channels, image_height, image_width]): int8 image tensor of
|
|
``c`` channels, where ``c`` must 3 or 1.
|
|
compression_level (int): Compression factor for the resulting file, it must be a number
|
|
between 0 and 9. Default: 6
|
|
|
|
Returns:
|
|
Tensor[1]: A one dimensional int8 tensor that contains the raw bytes of the
|
|
PNG file.
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(encode_png)
|
|
output = torch.ops.image.encode_png(input, compression_level)
|
|
return output
|
|
|
|
|
|
def write_png(input: torch.Tensor, filename: str, compression_level: int = 6):
|
|
"""
|
|
Takes an input tensor in CHW layout (or HW in the case of grayscale images)
|
|
and saves it in a PNG file.
|
|
|
|
Args:
|
|
input (Tensor[channels, image_height, image_width]): int8 image tensor of
|
|
``c`` channels, where ``c`` must be 1 or 3.
|
|
filename (str or ``pathlib.Path``): Path to save the image.
|
|
compression_level (int): Compression factor for the resulting file, it must be a number
|
|
between 0 and 9. Default: 6
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(write_png)
|
|
output = encode_png(input, compression_level)
|
|
write_file(filename, output)
|
|
|
|
|
|
def decode_jpeg(
|
|
input: Union[torch.Tensor, List[torch.Tensor]],
|
|
mode: ImageReadMode = ImageReadMode.UNCHANGED,
|
|
device: Union[str, torch.device] = "cpu",
|
|
apply_exif_orientation: bool = False,
|
|
) -> Union[torch.Tensor, List[torch.Tensor]]:
|
|
"""Decode JPEG image(s) into 3D RGB or grayscale Tensor(s), on CPU or CUDA.
|
|
|
|
The values of the output tensor are uint8 between 0 and 255.
|
|
|
|
.. note::
|
|
When using a CUDA device, passing a list of tensors is more efficient than repeated individual calls to ``decode_jpeg``.
|
|
When using CPU the performance is equivalent.
|
|
The CUDA version of this function has explicitly been designed with thread-safety in mind.
|
|
This function does not return partial results in case of an error.
|
|
|
|
Args:
|
|
input (Tensor[1] or list[Tensor[1]]): a (list of) one dimensional uint8 tensor(s) containing
|
|
the raw bytes of the JPEG image. The tensor(s) must be on CPU,
|
|
regardless of the ``device`` parameter.
|
|
mode (str or ImageReadMode): The mode to convert the image to, e.g. "RGB".
|
|
Default is "UNCHANGED". See :class:`~torchvision.io.ImageReadMode`
|
|
for available modes.
|
|
device (str or torch.device): The device on which the decoded image will
|
|
be stored. If a cuda device is specified, the image will be decoded
|
|
with `nvjpeg <https://developer.nvidia.com/nvjpeg>`_. This is only
|
|
supported for CUDA version >= 10.1
|
|
|
|
.. betastatus:: device parameter
|
|
|
|
.. 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] or list[Tensor[image_channels, image_height, image_width]]):
|
|
The values of the output tensor(s) are uint8 between 0 and 255.
|
|
``output.device`` will be set to the specified ``device``
|
|
|
|
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(decode_jpeg)
|
|
if isinstance(device, str):
|
|
device = torch.device(device)
|
|
if isinstance(mode, str):
|
|
mode = ImageReadMode[mode.upper()]
|
|
|
|
if isinstance(input, list):
|
|
if len(input) == 0:
|
|
raise ValueError("Input list must contain at least one element")
|
|
if not all(isinstance(t, torch.Tensor) for t in input):
|
|
raise ValueError("All elements of the input list must be tensors.")
|
|
if not all(t.device.type == "cpu" for t in input):
|
|
raise ValueError("Input list must contain tensors on CPU.")
|
|
if device.type == "cuda":
|
|
return torch.ops.image.decode_jpegs_cuda(input, mode.value, device)
|
|
else:
|
|
return [torch.ops.image.decode_jpeg(img, mode.value, apply_exif_orientation) for img in input]
|
|
|
|
else: # input is tensor
|
|
if input.device.type != "cpu":
|
|
raise ValueError("Input tensor must be a CPU tensor")
|
|
if device.type == "cuda":
|
|
return torch.ops.image.decode_jpegs_cuda([input], mode.value, device)[0]
|
|
else:
|
|
return torch.ops.image.decode_jpeg(input, mode.value, apply_exif_orientation)
|
|
|
|
|
|
def encode_jpeg(
|
|
input: Union[torch.Tensor, List[torch.Tensor]], quality: int = 75
|
|
) -> Union[torch.Tensor, List[torch.Tensor]]:
|
|
"""Encode RGB tensor(s) into raw encoded jpeg bytes, on CPU or CUDA.
|
|
|
|
.. note::
|
|
Passing a list of CUDA tensors is more efficient than repeated individual calls to ``encode_jpeg``.
|
|
For CPU tensors the performance is equivalent.
|
|
|
|
Args:
|
|
input (Tensor[channels, image_height, image_width] or List[Tensor[channels, image_height, image_width]]):
|
|
(list of) uint8 image tensor(s) of ``c`` channels, where ``c`` must be 1 or 3
|
|
quality (int): Quality of the resulting JPEG file(s). Must be a number between
|
|
1 and 100. Default: 75
|
|
|
|
Returns:
|
|
output (Tensor[1] or list[Tensor[1]]): A (list of) one dimensional uint8 tensor(s) that contain the raw bytes of the JPEG file.
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(encode_jpeg)
|
|
if quality < 1 or quality > 100:
|
|
raise ValueError("Image quality should be a positive number between 1 and 100")
|
|
if isinstance(input, list):
|
|
if not input:
|
|
raise ValueError("encode_jpeg requires at least one input tensor when a list is passed")
|
|
if input[0].device.type == "cuda":
|
|
return torch.ops.image.encode_jpegs_cuda(input, quality)
|
|
else:
|
|
return [torch.ops.image.encode_jpeg(image, quality) for image in input]
|
|
else: # single input tensor
|
|
if input.device.type == "cuda":
|
|
return torch.ops.image.encode_jpegs_cuda([input], quality)[0]
|
|
else:
|
|
return torch.ops.image.encode_jpeg(input, quality)
|
|
|
|
|
|
def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75):
|
|
"""
|
|
Takes an input tensor in CHW layout and saves it in a JPEG file.
|
|
|
|
Args:
|
|
input (Tensor[channels, image_height, image_width]): int8 image tensor of ``c``
|
|
channels, where ``c`` must be 1 or 3.
|
|
filename (str or ``pathlib.Path``): Path to save the image.
|
|
quality (int): Quality of the resulting JPEG file, it must be a number
|
|
between 1 and 100. Default: 75
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(write_jpeg)
|
|
output = encode_jpeg(input, quality)
|
|
assert isinstance(output, torch.Tensor) # Needed for torchscript
|
|
write_file(filename, output)
|
|
|
|
|
|
def decode_image(
|
|
input: Union[torch.Tensor, str],
|
|
mode: ImageReadMode = ImageReadMode.UNCHANGED,
|
|
apply_exif_orientation: bool = False,
|
|
) -> torch.Tensor:
|
|
"""Decode an image into a uint8 tensor, from a path or from raw encoded bytes.
|
|
|
|
Currently supported image formats are jpeg, png, gif and webp.
|
|
|
|
The values of the output tensor are in uint8 in [0, 255] for most cases.
|
|
|
|
If the image is a 16-bit png, then the output tensor is uint16 in [0, 65535]
|
|
(supported from torchvision ``0.21``). Since uint16 support is limited in
|
|
pytorch, we recommend calling
|
|
:func:`torchvision.transforms.v2.functional.to_dtype()` with ``scale=True``
|
|
after this function to convert the decoded image into a uint8 or float
|
|
tensor.
|
|
|
|
.. note::
|
|
|
|
``decode_image()`` doesn't work yet on AVIF or HEIC images. For these
|
|
formats, directly call :func:`~torchvision.io.decode_avif` or
|
|
:func:`~torchvision.io.decode_heic`.
|
|
|
|
Args:
|
|
input (Tensor or str or ``pathlib.Path``): The image to decode. If a
|
|
tensor is passed, it must be one dimensional uint8 tensor containing
|
|
the raw bytes of the image. Otherwise, this must be a path to the image file.
|
|
mode (str or ImageReadMode): The mode to convert the image to, e.g. "RGB".
|
|
Default is "UNCHANGED". See :class:`~torchvision.io.ImageReadMode`
|
|
for available modes.
|
|
apply_exif_orientation (bool): apply EXIF orientation transformation to the output tensor.
|
|
Only applies to JPEG and PNG images. Default: False.
|
|
|
|
Returns:
|
|
output (Tensor[image_channels, image_height, image_width])
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(decode_image)
|
|
if not isinstance(input, torch.Tensor):
|
|
input = read_file(str(input))
|
|
if isinstance(mode, str):
|
|
mode = ImageReadMode[mode.upper()]
|
|
output = torch.ops.image.decode_image(input, mode.value, apply_exif_orientation)
|
|
return output
|
|
|
|
|
|
def read_image(
|
|
path: str,
|
|
mode: ImageReadMode = ImageReadMode.UNCHANGED,
|
|
apply_exif_orientation: bool = False,
|
|
) -> torch.Tensor:
|
|
"""[OBSOLETE] Use :func:`~torchvision.io.decode_image` instead."""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(read_image)
|
|
data = read_file(path)
|
|
return decode_image(data, mode, apply_exif_orientation=apply_exif_orientation)
|
|
|
|
|
|
def decode_gif(input: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Decode a GIF image into a 3 or 4 dimensional RGB Tensor.
|
|
|
|
The values of the output tensor are uint8 between 0 and 255.
|
|
The output tensor has shape ``(C, H, W)`` if there is only one image in the
|
|
GIF, and ``(N, C, H, W)`` if there are ``N`` images.
|
|
|
|
Args:
|
|
input (Tensor[1]): a one dimensional contiguous uint8 tensor containing
|
|
the raw bytes of the GIF image.
|
|
|
|
Returns:
|
|
output (Tensor[image_channels, image_height, image_width] or Tensor[num_images, image_channels, image_height, image_width])
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(decode_gif)
|
|
return torch.ops.image.decode_gif(input)
|
|
|
|
|
|
def decode_webp(
|
|
input: torch.Tensor,
|
|
mode: ImageReadMode = ImageReadMode.UNCHANGED,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Decode a WEBP image into a 3 dimensional RGB[A] Tensor.
|
|
|
|
The values of the output tensor are uint8 between 0 and 255.
|
|
|
|
Args:
|
|
input (Tensor[1]): a one dimensional contiguous uint8 tensor containing
|
|
the raw bytes of the WEBP image.
|
|
mode (str or ImageReadMode): The mode to convert the image to, e.g. "RGB".
|
|
Default is "UNCHANGED". See :class:`~torchvision.io.ImageReadMode`
|
|
for available modes.
|
|
|
|
Returns:
|
|
Decoded image (Tensor[image_channels, image_height, image_width])
|
|
"""
|
|
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
|
|
_log_api_usage_once(decode_webp)
|
|
if isinstance(mode, str):
|
|
mode = ImageReadMode[mode.upper()]
|
|
return torch.ops.image.decode_webp(input, mode.value)
|
|
|
|
|
|
# TODO_AVIF_HEIC: Better support for torchscript. Scripting decode_avif of
|
|
# decode_heic currently fails, mainly because of the logic
|
|
# _load_extra_decoders_once() (using global variables, try/except statements,
|
|
# etc.).
|
|
# The ops (torch.ops.extra_decoders_ns.decode_*) are otherwise torchscript-able,
|
|
# and users who need torchscript can always just wrap those.
|
|
|
|
# TODO_AVIF_HEIC: decode_image() should work for those. The key technical issue
|
|
# we have here is that the format detection logic of decode_image() is
|
|
# implemented in torchvision, and torchvision has zero knowledge of
|
|
# torchvision-extra-decoders, so we cannot call the AVIF/HEIC C++ decoders
|
|
# (those in torchvision-extra-decoders) from there.
|
|
# A trivial check that could be done within torchvision would be to check the
|
|
# file extension, if a path was passed. We could also just implement the
|
|
# AVIF/HEIC detection logic in Python as a fallback, if the file detection
|
|
# didn't find any format. In any case: properly determining whether a file is
|
|
# HEIC is far from trivial, and relying on libmagic would probably be best
|
|
|
|
|
|
_EXTRA_DECODERS_ALREADY_LOADED = False
|
|
|
|
|
|
def _load_extra_decoders_once():
|
|
global _EXTRA_DECODERS_ALREADY_LOADED
|
|
if _EXTRA_DECODERS_ALREADY_LOADED:
|
|
return
|
|
|
|
try:
|
|
import torchvision_extra_decoders
|
|
|
|
# torchvision-extra-decoders only supports linux for now. BUT, users on
|
|
# e.g. MacOS can still install it: they will get the pure-python
|
|
# 0.0.0.dev version:
|
|
# https://pypi.org/project/torchvision-extra-decoders/0.0.0.dev0, which
|
|
# is a dummy version that was created to reserve the namespace on PyPI.
|
|
# We have to check that expose_extra_decoders() exists for those users,
|
|
# so we can properly error on non-Linux archs.
|
|
assert hasattr(torchvision_extra_decoders, "expose_extra_decoders")
|
|
except (AssertionError, ImportError) as e:
|
|
raise RuntimeError(
|
|
"In order to enable the AVIF and HEIC decoding capabilities of "
|
|
"torchvision, you need to `pip install torchvision-extra-decoders`. "
|
|
"Just install the package, you don't need to update your code. "
|
|
"This is only supported on Linux, and this feature is still in BETA stage. "
|
|
"Please let us know of any issue: https://github.com/pytorch/vision/issues/new/choose. "
|
|
"Note that `torchvision-extra-decoders` is released under the LGPL license. "
|
|
) from e
|
|
|
|
# This will expose torch.ops.extra_decoders_ns.decode_avif and torch.ops.extra_decoders_ns.decode_heic
|
|
torchvision_extra_decoders.expose_extra_decoders()
|
|
|
|
_EXTRA_DECODERS_ALREADY_LOADED = True
|
|
|
|
|
|
def decode_avif(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
|
|
"""Decode an AVIF image into a 3 dimensional RGB[A] Tensor.
|
|
|
|
.. warning::
|
|
In order to enable the AVIF decoding capabilities of torchvision, you
|
|
first need to run ``pip install torchvision-extra-decoders``. Just
|
|
install the package, you don't need to update your code. This is only
|
|
supported on Linux, and this feature is still in BETA stage. Please let
|
|
us know of any issue:
|
|
https://github.com/pytorch/vision/issues/new/choose. Note that
|
|
`torchvision-extra-decoders
|
|
<https://github.com/pytorch-labs/torchvision-extra-decoders/>`_ is
|
|
released under the LGPL license.
|
|
|
|
The values of the output tensor are in uint8 in [0, 255] for most images. If
|
|
the image has a bit-depth of more than 8, then the output tensor is uint16
|
|
in [0, 65535]. Since uint16 support is limited in pytorch, we recommend
|
|
calling :func:`torchvision.transforms.v2.functional.to_dtype()` with
|
|
``scale=True`` after this function to convert the decoded image into a uint8
|
|
or float tensor.
|
|
|
|
Args:
|
|
input (Tensor[1]): a one dimensional contiguous uint8 tensor containing
|
|
the raw bytes of the AVIF image.
|
|
mode (str or ImageReadMode): The mode to convert the image to, e.g. "RGB".
|
|
Default is "UNCHANGED". See :class:`~torchvision.io.ImageReadMode`
|
|
for available modes.
|
|
|
|
Returns:
|
|
Decoded image (Tensor[image_channels, image_height, image_width])
|
|
"""
|
|
_load_extra_decoders_once()
|
|
if input.dtype != torch.uint8:
|
|
raise RuntimeError(f"Input tensor must have uint8 data type, got {input.dtype}")
|
|
return torch.ops.extra_decoders_ns.decode_avif(input, mode.value)
|
|
|
|
|
|
def decode_heic(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
|
|
"""Decode an HEIC image into a 3 dimensional RGB[A] Tensor.
|
|
|
|
.. warning::
|
|
In order to enable the AVIF decoding capabilities of torchvision, you
|
|
first need to run ``pip install torchvision-extra-decoders``. Just
|
|
install the package, you don't need to update your code. This is only
|
|
supported on Linux, and this feature is still in BETA stage. Please let
|
|
us know of any issue:
|
|
https://github.com/pytorch/vision/issues/new/choose. Note that
|
|
`torchvision-extra-decoders
|
|
<https://github.com/pytorch-labs/torchvision-extra-decoders/>`_ is
|
|
released under the LGPL license.
|
|
|
|
The values of the output tensor are in uint8 in [0, 255] for most images. If
|
|
the image has a bit-depth of more than 8, then the output tensor is uint16
|
|
in [0, 65535]. Since uint16 support is limited in pytorch, we recommend
|
|
calling :func:`torchvision.transforms.v2.functional.to_dtype()` with
|
|
``scale=True`` after this function to convert the decoded image into a uint8
|
|
or float tensor.
|
|
|
|
Args:
|
|
input (Tensor[1]): a one dimensional contiguous uint8 tensor containing
|
|
the raw bytes of the HEIC image.
|
|
mode (str or ImageReadMode): The mode to convert the image to, e.g. "RGB".
|
|
Default is "UNCHANGED". See :class:`~torchvision.io.ImageReadMode`
|
|
for available modes.
|
|
|
|
Returns:
|
|
Decoded image (Tensor[image_channels, image_height, image_width])
|
|
"""
|
|
_load_extra_decoders_once()
|
|
if input.dtype != torch.uint8:
|
|
raise RuntimeError(f"Input tensor must have uint8 data type, got {input.dtype}")
|
|
return torch.ops.extra_decoders_ns.decode_heic(input, mode.value)
|