from math import pi
from typing import Optional, Union
import torch
from torch import Tensor
from torch.nn import Module, Parameter
from kornia.color import hsv_to_rgb, rgb_to_grayscale, rgb_to_hsv
from kornia.testing import KORNIA_CHECK, KORNIA_CHECK_IS_COLOR_OR_GRAY, KORNIA_CHECK_IS_TENSOR
from kornia.utils.helpers import _torch_histc_cast
from kornia.utils.image import perform_keep_shape_image, perform_keep_shape_video
def adjust_saturation_raw(image: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Adjust color saturation of an image.
Expecting image to be in hsv format already.
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(isinstance(factor, (float, Tensor)), "Factor should be float or Tensor.")
if isinstance(factor, float):
# TODO: figure out how to create later a tensor without importing torch
factor = torch.as_tensor(factor, device=image.device, dtype=image.dtype)
elif isinstance(factor, Tensor):
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
# unpack the hsv values
h, s, v = torch.chunk(image, chunks=3, dim=-3)
# transform the hue value and appl module
s_out: Tensor = torch.clamp(s * factor, min=0, max=1)
# pack back back the corrected hue
out: Tensor = torch.cat([h, s_out, v], dim=-3)
return out
def adjust_saturation_with_gray_subtraction(image: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Adjust color saturation of an image by blending the image with its grayscaled version.
The image is expected to be an RGB image or a gray image in the range of [0, 1].
If it is an RGB image, returns blending of the image with its grayscaled version.
If it is a gray image, returns the image.
.. note::
this is just a convenience function to have compatibility with Pil
Args:
image: Image/Tensor to be adjusted in the shape of :math:`(*, 3, H, W)`.
factor: How much to adjust the saturation. 0 will give a black
and white image, 1 will give the original image while 2 will enhance the saturation by a factor of 2.
Return:
Adjusted image in the shape of :math:`(*, 3, H, W)`.
Example:
>>> x = torch.ones(1, 3, 3, 3)
>>> adjust_saturation_with_gray_subtraction(x, 2.).shape
torch.Size([1, 3, 3, 3])
>>> x = torch.ones(2, 3, 3, 3)
>>> y = torch.tensor([1., 2.])
>>> adjust_saturation_with_gray_subtraction(x, y).shape
torch.Size([2, 3, 3, 3])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(isinstance(factor, (float, Tensor)), "Factor should be float or Tensor.")
KORNIA_CHECK_IS_COLOR_OR_GRAY(image, "Image should be an RGB or gray image")
if image.shape[-3] == 1:
return image
if isinstance(factor, float):
# TODO: figure out how to create later a tensor without importing torch
factor = torch.as_tensor(factor, device=image.device, dtype=image.dtype)
elif isinstance(factor, Tensor):
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
x_other: Tensor = rgb_to_grayscale(image)
# blend the image with the grayscaled image
x_adjusted: Tensor = (1 - factor) * x_other + factor * image
# clamp the output
out: Tensor = torch.clamp(x_adjusted, 0.0, 1.0)
return out
[docs]def adjust_saturation(image: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Adjust color saturation of an image.
.. image:: _static/img/adjust_saturation.png
The image is expected to be an RGB image in the range of [0, 1].
Args:
image: Image/Tensor to be adjusted in the shape of :math:`(*, 3, H, W)`.
factor: How much to adjust the saturation. 0 will give a black
and white image, 1 will give the original image while 2 will enhance the saturation by a factor of 2.
saturation_mode: The mode to adjust saturation.
Return:
Adjusted image in the shape of :math:`(*, 3, H, W)`.
.. note::
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
image_enhancement.html>`__.
Example:
>>> x = torch.ones(1, 3, 3, 3)
>>> adjust_saturation(x, 2.).shape
torch.Size([1, 3, 3, 3])
>>> x = torch.ones(2, 3, 3, 3)
>>> y = torch.tensor([1., 2.])
>>> adjust_saturation(x, y).shape
torch.Size([2, 3, 3, 3])
"""
# convert the rgb image to hsv
x_hsv: Tensor = rgb_to_hsv(image)
# perform the conversion
x_adjusted: Tensor = adjust_saturation_raw(x_hsv, factor)
# convert back to rgb
out: Tensor = hsv_to_rgb(x_adjusted)
return out
def adjust_hue_raw(image: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Adjust hue of an image.
Expecting image to be in hsv format already.
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(
isinstance(factor, (float, Tensor)),
f"The factor should be a float number or Tensor in the range between" f" [-PI, PI]. Got {type(factor)}",
)
if isinstance(factor, float):
factor = torch.as_tensor(factor)
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
# unpack the hsv values
h, s, v = torch.chunk(image, chunks=3, dim=-3)
# transform the hue value and appl module
divisor: float = 2 * pi
h_out: Tensor = torch.fmod(h + factor, divisor)
# pack back back the corrected hue
out: Tensor = torch.cat([h_out, s, v], dim=-3)
return out
[docs]def adjust_hue(image: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Adjust hue of an image.
.. image:: _static/img/adjust_hue.png
The image is expected to be an RGB image in the range of [0, 1].
Args:
image: Image to be adjusted in the shape of :math:`(*, 3, H, W)`.
factor: How much to shift the hue channel. Should be in [-PI, PI]. PI
and -PI give complete reversal of hue channel in HSV space in positive and negative
direction respectively. 0 means no shift. Therefore, both -PI and PI will give an
image with complementary colors while 0 gives the original image.
Return:
Adjusted image in the shape of :math:`(*, 3, H, W)`.
.. note::
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
image_enhancement.html>`__.
Example:
>>> x = torch.ones(1, 3, 2, 2)
>>> adjust_hue(x, 3.141516).shape
torch.Size([1, 3, 2, 2])
>>> x = torch.ones(2, 3, 3, 3)
>>> y = torch.ones(2) * 3.141516
>>> adjust_hue(x, y).shape
torch.Size([2, 3, 3, 3])
"""
# convert the rgb image to hsv
x_hsv: Tensor = rgb_to_hsv(image)
# perform the conversion
x_adjusted: Tensor = adjust_hue_raw(x_hsv, factor)
# convert back to rgb
out: Tensor = hsv_to_rgb(x_adjusted)
return out
[docs]def adjust_gamma(input: Tensor, gamma: Union[float, Tensor], gain: Union[float, Tensor] = 1.0) -> Tensor:
r"""Perform gamma correction on an image.
.. image:: _static/img/adjust_contrast.png
The input image is expected to be in the range of [0, 1].
Args:
input: Image to be adjusted in the shape of :math:`(*, H, W)`.
gamma: Non negative real number, same as γ\gammaγ in the equation.
gamma larger than 1 make the shadows darker, while gamma smaller than 1 make
dark regions lighter.
gain: The constant multiplier.
Return:
Adjusted image in the shape of :math:`(*, H, W)`.
.. note::
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
image_enhancement.html>`__.
Example:
>>> x = torch.ones(1, 1, 2, 2)
>>> adjust_gamma(x, 1.0, 2.0)
tensor([[[[1., 1.],
[1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y1 = torch.ones(2) * 1.0
>>> y2 = torch.ones(2) * 2.0
>>> adjust_gamma(x, y1, y2).shape
torch.Size([2, 5, 3, 3])
"""
if not isinstance(input, Tensor):
raise TypeError(f"Input type is not a Tensor. Got {type(input)}")
if not isinstance(gamma, (float, Tensor)):
raise TypeError(f"The gamma should be a positive float or Tensor. Got {type(gamma)}")
if not isinstance(gain, (float, Tensor)):
raise TypeError(f"The gain should be a positive float or Tensor. Got {type(gain)}")
if isinstance(gamma, float):
gamma = Tensor([gamma])
if isinstance(gain, float):
gain = Tensor([gain])
gamma = gamma.to(input.device).to(input.dtype)
gain = gain.to(input.device).to(input.dtype)
if (gamma < 0.0).any():
raise ValueError(f"Gamma must be non-negative. Got {gamma}")
if (gain < 0.0).any():
raise ValueError(f"Gain must be non-negative. Got {gain}")
for _ in range(len(input.shape) - len(gamma.shape)):
gamma = torch.unsqueeze(gamma, dim=-1)
for _ in range(len(input.shape) - len(gain.shape)):
gain = torch.unsqueeze(gain, dim=-1)
# Apply the gamma correction
x_adjust: Tensor = gain * torch.pow(input, gamma)
# Truncate between pixel values
out: Tensor = torch.clamp(x_adjust, 0.0, 1.0)
return out
[docs]def adjust_contrast(image: Tensor, factor: Union[float, Tensor], clip_output: bool = True) -> Tensor:
r"""Adjust the contrast of an image tensor.
.. image:: _static/img/adjust_contrast.png
This implementation follows Szeliski's book convention, where contrast is defined as
a `multiplicative` operation directly to raw pixel values. Beware that other frameworks
might use different conventions which can be difficult to reproduce exact results.
The input image and factor is expected to be in the range of [0, 1].
.. tip::
This is not the preferred way to adjust the contrast of an image. Ideally one must
implement :func:`kornia.enhance.adjust_gamma`. More details in the following link:
https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_log_gamma.html#sphx-glr-auto-examples-color-exposure-plot-log-gamma-py
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
factor: Contrast adjust factor per element
in the batch. 0 generates a completely black image, 1 does not modify
the input image while any other non-negative number modify the
brightness by this factor.
clip_output: whether to clip the output image with range of [0, 1].
Return:
Adjusted image in the shape of :math:`(*, H, W)`.
.. note::
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
image_enhancement.html>`__.
Example:
>>> import torch
>>> x = torch.ones(1, 1, 2, 2)
>>> adjust_contrast(x, 0.5)
tensor([[[[0.5000, 0.5000],
[0.5000, 0.5000]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.tensor([0.65, 0.50])
>>> adjust_contrast(x, y).shape
torch.Size([2, 5, 3, 3])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(isinstance(factor, (float, Tensor)), "Factor should be float or Tensor.")
if isinstance(factor, float):
# TODO: figure out how to create later a tensor without importing torch
factor = torch.as_tensor(factor, device=image.device, dtype=image.dtype)
elif isinstance(factor, Tensor):
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
KORNIA_CHECK(any(factor >= 0), f"Contrast factor must be positive. Got {factor}")
# Apply contrast factor to each channel
img_adjust: Tensor = image * factor
# Truncate between pixel values
if clip_output:
img_adjust = img_adjust.clamp(min=0.0, max=1.0)
return img_adjust
[docs]def adjust_contrast_with_mean_subtraction(image: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Adjust the contrast of an image tensor by subtracting the mean over channels.
.. note::
this is just a convenience function to have compatibility with Pil. For exact
definition of image contrast adjustment consider using :func:`kornia.enhance.adjust_gamma`.
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
factor: Contrast adjust factor per element
in the batch. 0 generates a completely black image, 1 does not modify
the input image while any other non-negative number modify the
brightness by this factor.
Return:
Adjusted image in the shape of :math:`(*, H, W)`.
Example:
>>> import torch
>>> x = torch.ones(1, 1, 2, 2)
>>> adjust_contrast_with_mean_subtraction(x, 0.5)
tensor([[[[1., 1.],
[1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.tensor([0.65, 0.50])
>>> adjust_contrast_with_mean_subtraction(x, y).shape
torch.Size([2, 5, 3, 3])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(isinstance(factor, (float, Tensor)), "Factor should be float or Tensor.")
if isinstance(factor, float):
# TODO: figure out how to create later a tensor without importing torch
factor = torch.as_tensor(factor, device=image.device, dtype=image.dtype)
elif isinstance(factor, Tensor):
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
KORNIA_CHECK(any(factor >= 0), f"Contrast factor must be positive. Got {factor}")
if image.shape[-3] == 3:
img_mean = rgb_to_grayscale(image).mean((-2, -1), True)
else:
img_mean = image.mean()
# Apply contrast factor subtracting the mean
img_adjust: Tensor = image * factor + img_mean * (1 - factor)
img_adjust = img_adjust.clamp(min=0.0, max=1.0)
return img_adjust
[docs]def adjust_brightness(image: Tensor, factor: Union[float, Tensor], clip_output=True) -> Tensor:
r"""Adjust the brightness of an image tensor.
.. image:: _static/img/adjust_brightness.png
This implementation follows Szeliski's book convention, where brightness is defined as
an `additive` operation directly to raw pixel and shift its values according the applied
factor and range of the image values. Beware that other framework might use different
conventions which can be difficult to reproduce exact results.
The input image and factor is expected to be in the range of [0, 1].
.. tip::
By applying a large factor might prouce clipping or loss of image detail. We recommenda to
apply small factors to avoid the mentioned issues. Ideally one must implement the adjustment
of image intensity with other techniques suchs as :func:`kornia.enhance.adjust_gamma`. More
details in the following link:
https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_log_gamma.html#sphx-glr-auto-examples-color-exposure-plot-log-gamma-py
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
factor: Brightness adjust factor per element in the batch. It's recommended to
bound the factor by [0, 1]. 0 does not modify the input image while any other
number modify the brightness.
Return:
Adjusted tensor in the shape of :math:`(*, H, W)`.
.. note::
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
image_enhancement.html>`__.
Example:
>>> x = torch.ones(1, 1, 2, 2)
>>> adjust_brightness(x, 1.)
tensor([[[[1., 1.],
[1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.tensor([0.25, 0.50])
>>> adjust_brightness(x, y).shape
torch.Size([2, 5, 3, 3])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(isinstance(factor, (float, Tensor)), "Factor should be float or Tensor.")
# convert factor to a tensor
if isinstance(factor, float):
# TODO: figure out how to create later a tensor without importing torch
factor = torch.as_tensor(factor, device=image.device, dtype=image.dtype)
elif isinstance(factor, Tensor):
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
# shift pixel values
img_adjust: Tensor = image + factor
# truncate between pixel values
if clip_output:
img_adjust = img_adjust.clamp(min=0.0, max=1.0)
return img_adjust
def adjust_brightness_accumulative(image: Tensor, factor: Union[float, Tensor], clip_output=True) -> Tensor:
r"""Adjust the brightness accumulatively of an image tensor.
This implementation follows PIL convention.
The input image and factor is expected to be in the range of [0, 1].
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
factor: Brightness adjust factor per element in the batch. It's recommended to
bound the factor by [0, 1]. 0 does not modify the input image while any other
number modify the brightness.
Return:
Adjusted tensor in the shape of :math:`(*, H, W)`.
Example:
>>> x = torch.ones(1, 1, 2, 2)
>>> adjust_brightness_accumulative(x, 1.)
tensor([[[[1., 1.],
[1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.tensor([0.25, 0.50])
>>> adjust_brightness_accumulative(x, y).shape
torch.Size([2, 5, 3, 3])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
KORNIA_CHECK(isinstance(factor, (float, Tensor)), "Factor should be float or Tensor.")
# convert factor to a tensor
if isinstance(factor, float):
# TODO: figure out how to create later a tensor without importing torch
factor = torch.as_tensor(factor, device=image.device, dtype=image.dtype)
elif isinstance(factor, Tensor):
factor = factor.to(image.device, image.dtype)
# make factor broadcastable
while len(factor.shape) != len(image.shape):
factor = factor[..., None]
# shift pixel values
img_adjust: Tensor = image * factor
# truncate between pixel values
if clip_output:
img_adjust = img_adjust.clamp(min=0.0, max=1.0)
return img_adjust
[docs]def adjust_sigmoid(image: Tensor, cutoff: float = 0.5, gain: float = 10, inv: bool = False) -> Tensor:
"""Adjust sigmoid correction on the input image tensor.
The input image is expected to be in the range of [0, 1].
Reference:
[1]: Gustav J. Braun, "Image Lightness Rescaling Using Sigmoidal Contrast Enhancement Functions",
http://markfairchild.org/PDFs/PAP07.pdf
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
cutoff: The cutoff of sigmoid function.
gain: The multiplier of sigmoid function.
inv: If is set to True the function will return the inverse sigmoid correction.
Returns:
Adjusted tensor in the shape of :math:`(*, H, W)`.
Example:
>>> x = torch.ones(1, 1, 2, 2)
>>> adjust_sigmoid(x, gain=0)
tensor([[[[0.5000, 0.5000],
[0.5000, 0.5000]]]])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
if inv:
img_adjust = 1 - 1 / (1 + (gain * (cutoff - image)).exp())
else:
img_adjust = 1 / (1 + (gain * (cutoff - image)).exp())
return img_adjust
[docs]def adjust_log(image: Tensor, gain: float = 1, inv: bool = False, clip_output: bool = True) -> Tensor:
"""Adjust log correction on the input image tensor.
The input image is expected to be in the range of [0, 1].
Reference:
[1]: http://www.ece.ucsb.edu/Faculty/Manjunath/courses/ece178W03/EnhancePart1.pdf
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
gain: The multiplier of logarithmic function.
inv: If is set to True the function will return the inverse logarithmic correction.
clip_output: Whether to clip the output image with range of [0, 1].
Returns:
Adjusted tensor in the shape of :math:`(*, H, W)`.
Example:
>>> x = torch.zeros(1, 1, 2, 2)
>>> adjust_log(x, inv=True)
tensor([[[[0., 0.],
[0., 0.]]]])
"""
KORNIA_CHECK_IS_TENSOR(image, "Expected shape (*, H, W)")
if inv:
img_adjust = (2**image - 1) * gain
else:
img_adjust = (1 + image).log2() * gain
# truncate between pixel values
if clip_output:
img_adjust = img_adjust.clamp(min=0.0, max=1.0)
return img_adjust
def _solarize(input: Tensor, thresholds: Union[float, Tensor] = 0.5) -> Tensor:
r"""For each pixel in the image, select the pixel if the value is less than the threshold. Otherwise, subtract
1.0 from the pixel.
Args:
input: image or batched images to solarize.
thresholds: solarize thresholds.
If int or one element tensor, input will be solarized across the whole batch.
If 1-d tensor, input will be solarized element-wise, len(thresholds) == len(input).
Returns:
Solarized images.
"""
if not isinstance(input, Tensor):
raise TypeError(f"Input type is not a Tensor. Got {type(input)}")
if not isinstance(thresholds, (float, Tensor)):
raise TypeError(f"The factor should be either a float or Tensor. " f"Got {type(thresholds)}")
if isinstance(thresholds, Tensor) and len(thresholds.shape) != 0:
if not (input.size(0) == len(thresholds) and len(thresholds.shape) == 1):
raise AssertionError(f"thresholds must be a 1-d vector of shape ({input.size(0)},). Got {thresholds}")
# TODO: I am not happy about this line, but no easy to do batch-wise operation
thresholds = thresholds.to(input.device).to(input.dtype)
thresholds = torch.stack([x.expand(*input.shape[-3:]) for x in thresholds])
return torch.where(input < thresholds, input, 1.0 - input)
[docs]def solarize(
input: Tensor, thresholds: Union[float, Tensor] = 0.5, additions: Optional[Union[float, Tensor]] = None
) -> Tensor:
r"""For each pixel in the image less than threshold.
.. image:: _static/img/solarize.png
We add 'addition' amount to it and then clip the pixel value to be between 0 and 1.0.
The value of 'addition' is between -0.5 and 0.5.
Args:
input: image tensor with shapes like :math:`(*, C, H, W)` to solarize.
thresholds: solarize thresholds.
If int or one element tensor, input will be solarized across the whole batch.
If 1-d tensor, input will be solarized element-wise, len(thresholds) == len(input).
additions: between -0.5 and 0.5.
If None, no addition will be performed.
If int or one element tensor, same addition will be added across the whole batch.
If 1-d tensor, additions will be added element-wisely, len(additions) == len(input).
Returns:
The solarized images with shape :math:`(*, C, H, W)`.
Example:
>>> x = torch.rand(1, 4, 3, 3)
>>> out = solarize(x, thresholds=0.5, additions=0.)
>>> out.shape
torch.Size([1, 4, 3, 3])
>>> x = torch.rand(2, 4, 3, 3)
>>> thresholds = torch.tensor([0.8, 0.5])
>>> additions = torch.tensor([-0.25, 0.25])
>>> solarize(x, thresholds, additions).shape
torch.Size([2, 4, 3, 3])
"""
if not isinstance(input, Tensor):
raise TypeError(f"Input type is not a Tensor. Got {type(input)}")
if not isinstance(thresholds, (float, Tensor)):
raise TypeError(f"The factor should be either a float or Tensor. " f"Got {type(thresholds)}")
if isinstance(thresholds, float):
thresholds = torch.as_tensor(thresholds)
if additions is not None:
if not isinstance(additions, (float, Tensor)):
raise TypeError(f"The factor should be either a float or Tensor. " f"Got {type(additions)}")
if isinstance(additions, float):
additions = torch.as_tensor(additions)
if not torch.all((additions < 0.5) * (additions > -0.5)):
raise AssertionError(f"The value of 'addition' is between -0.5 and 0.5. Got {additions}.")
if isinstance(additions, Tensor) and len(additions.shape) != 0:
if not (input.size(0) == len(additions) and len(additions.shape) == 1):
raise AssertionError(f"additions must be a 1-d vector of shape ({input.size(0)},). Got {additions}")
# TODO: I am not happy about this line, but no easy to do batch-wise operation
additions = additions.to(input.device).to(input.dtype)
additions = torch.stack([x.expand(*input.shape[-3:]) for x in additions])
input = input + additions
input = input.clamp(0.0, 1.0)
return _solarize(input, thresholds)
[docs]@perform_keep_shape_image
def posterize(input: Tensor, bits: Union[int, Tensor]) -> Tensor:
r"""Reduce the number of bits for each color channel.
.. image:: _static/img/posterize.png
Non-differentiable function, ``torch.uint8`` involved.
Args:
input: image tensor with shape :math:`(*, C, H, W)` to posterize.
bits: number of high bits. Must be in range [0, 8].
If int or one element tensor, input will be posterized by this bits.
If 1-d tensor, input will be posterized element-wisely, len(bits) == input.shape[-3].
If n-d tensor, input will be posterized element-channel-wisely, bits.shape == input.shape[:len(bits.shape)]
Returns:
Image with reduced color channels with shape :math:`(*, C, H, W)`.
Example:
>>> x = torch.rand(1, 6, 3, 3)
>>> out = posterize(x, bits=8)
>>> torch.testing.assert_allclose(x, out)
>>> x = torch.rand(2, 6, 3, 3)
>>> bits = torch.tensor([4, 2])
>>> posterize(x, bits).shape
torch.Size([2, 6, 3, 3])
"""
if not isinstance(input, Tensor):
raise TypeError(f"Input type is not a Tensor. Got {type(input)}")
if not isinstance(bits, (int, Tensor)):
raise TypeError(f"bits type is not an int or Tensor. Got {type(bits)}")
if isinstance(bits, int):
bits = torch.as_tensor(bits)
# TODO: find a better way to check boundaries on tensors
# if not torch.all((bits >= 0) * (bits <= 8)) and bits.dtype == torch.int:
# raise ValueError(f"bits must be integers within range [0, 8]. Got {bits}.")
# TODO: Make a differentiable version
# Current version:
# Ref: https://github.com/open-mmlab/mmcv/pull/132/files#diff-309c9320c7f71bedffe89a70ccff7f3bR19
# Ref: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py#L222
# Potential approach: implementing kornia.LUT with floating points
# https://github.com/albumentations-team/albumentations/blob/master/albumentations/augmentations/functional.py#L472
def _left_shift(input: Tensor, shift: Tensor):
return ((input * 255).to(torch.uint8) * (2**shift)).to(input.dtype) / 255.0
def _right_shift(input: Tensor, shift: Tensor):
return (input * 255).to(torch.uint8) / (2**shift).to(input.dtype) / 255.0
def _posterize_one(input: Tensor, bits: Tensor):
# Single bits value condition
if bits == 0:
return torch.zeros_like(input)
if bits == 8:
return input.clone()
bits = 8 - bits
return _left_shift(_right_shift(input, bits), bits)
if len(bits.shape) == 0 or (len(bits.shape) == 1 and len(bits) == 1):
return _posterize_one(input, bits)
res = []
if len(bits.shape) == 1:
if bits.shape[0] != input.shape[0]:
raise AssertionError(
f"Batch size must be equal between bits and input. Got {bits.shape[0]}, {input.shape[0]}."
)
for i in range(input.shape[0]):
res.append(_posterize_one(input[i], bits[i]))
return torch.stack(res, dim=0)
if bits.shape != input.shape[: len(bits.shape)]:
raise AssertionError(
"Batch and channel must be equal between bits and input. "
f"Got {bits.shape}, {input.shape[:len(bits.shape)]}."
)
_input = input.view(-1, *input.shape[len(bits.shape) :])
_bits = bits.flatten()
for i in range(input.shape[0]):
res.append(_posterize_one(_input[i], _bits[i]))
return torch.stack(res, dim=0).reshape(*input.shape)
[docs]@perform_keep_shape_image
def sharpness(input: Tensor, factor: Union[float, Tensor]) -> Tensor:
r"""Apply sharpness to the input tensor.
.. image:: _static/img/sharpness.png
Implemented Sharpness function from PIL using torch ops. This implementation refers to:
https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py#L326
Args:
input: image tensor with shape :math:`(*, C, H, W)` to sharpen.
factor: factor of sharpness strength. Must be above 0.
If float or one element tensor, input will be sharpened by the same factor across the whole batch.
If 1-d tensor, input will be sharpened element-wisely, len(factor) == len(input).
Returns:
Sharpened image or images with shape :math:`(*, C, H, W)`.
Example:
>>> x = torch.rand(1, 1, 5, 5)
>>> sharpness(x, 0.5).shape
torch.Size([1, 1, 5, 5])
"""
if not isinstance(factor, Tensor):
factor = torch.as_tensor(factor, device=input.device, dtype=input.dtype)
if len(factor.size()) != 0 and factor.shape != torch.Size([input.size(0)]):
raise AssertionError(
"Input batch size shall match with factor size if factor is not a 0-dim tensor. "
f"Got {input.size(0)} and {factor.shape}"
)
kernel = (
torch.as_tensor([[1, 1, 1], [1, 5, 1], [1, 1, 1]], dtype=input.dtype, device=input.device)
.view(1, 1, 3, 3)
.repeat(input.size(1), 1, 1, 1)
/ 13
)
# This shall be equivalent to depthwise conv2d:
# Ref: https://discuss.pytorch.org/t/depthwise-and-separable-convolutions-in-pytorch/7315/2
degenerate = torch.nn.functional.conv2d(input, kernel, bias=None, stride=1, groups=input.size(1))
degenerate = torch.clamp(degenerate, 0.0, 1.0)
# For the borders of the resulting image, fill in the values of the original image.
mask = torch.ones_like(degenerate)
padded_mask = torch.nn.functional.pad(mask, [1, 1, 1, 1])
padded_degenerate = torch.nn.functional.pad(degenerate, [1, 1, 1, 1])
result = torch.where(padded_mask == 1, padded_degenerate, input)
if len(factor.size()) == 0:
return _blend_one(result, input, factor)
return torch.stack([_blend_one(result[i], input[i], factor[i]) for i in range(len(factor))])
def _blend_one(input1: Tensor, input2: Tensor, factor: Tensor) -> Tensor:
r"""Blend two images into one.
Args:
input1: image tensor with shapes like :math:`(H, W)` or :math:`(D, H, W)`.
input2: image tensor with shapes like :math:`(H, W)` or :math:`(D, H, W)`.
factor: factor 0-dim tensor.
Returns:
: image tensor with the batch in the zero position.
"""
if not isinstance(input1, Tensor):
raise AssertionError(f"`input1` must be a tensor. Got {input1}.")
if not isinstance(input2, Tensor):
raise AssertionError(f"`input1` must be a tensor. Got {input2}.")
if isinstance(factor, Tensor) and len(factor.size()) != 0:
raise AssertionError(f"Factor shall be a float or single element tensor. Got {factor}.")
if factor == 0.0:
return input1
if factor == 1.0:
return input2
diff = (input2 - input1) * factor
res = input1 + diff
if factor > 0.0 and factor < 1.0:
return res
return torch.clamp(res, 0, 1)
def _build_lut(histo, step):
# Compute the cumulative sum, shifting by step // 2
# and then normalization by step.
step_trunc = torch.div(step, 2, rounding_mode='trunc')
lut = torch.div(torch.cumsum(histo, 0) + step_trunc, step, rounding_mode='trunc')
# Shift lut, prepending with 0.
lut = torch.cat([torch.zeros(1, device=lut.device, dtype=lut.dtype), lut[:-1]])
# Clip the counts to be in range. This is done
# in the C code for image.point.
return torch.clamp(lut, 0, 255)
# Code taken from: https://github.com/pytorch/vision/pull/796
def _scale_channel(im: Tensor) -> Tensor:
r"""Scale the data in the channel to implement equalize.
Args:
input: image tensor with shapes like :math:`(H, W)` or :math:`(D, H, W)`.
Returns:
image tensor with the batch in the zero position.
"""
min_ = im.min()
max_ = im.max()
if min_.item() < 0.0 and not torch.isclose(min_, torch.as_tensor(0.0, dtype=min_.dtype)):
raise ValueError(f"Values in the input tensor must greater or equal to 0.0. Found {min_.item()}.")
if max_.item() > 1.0 and not torch.isclose(max_, torch.as_tensor(1.0, dtype=max_.dtype)):
raise ValueError(f"Values in the input tensor must lower or equal to 1.0. Found {max_.item()}.")
ndims = len(im.shape)
if ndims not in (2, 3):
raise TypeError(f"Input tensor must have 2 or 3 dimensions. Found {ndims}.")
im = im * 255.0
# Compute the histogram of the image channel.
histo = _torch_histc_cast(im, bins=256, min=0, max=255)
# For the purposes of computing the step, filter out the nonzeros.
nonzero_histo = torch.reshape(histo[histo != 0], [-1])
step = torch.div(torch.sum(nonzero_histo) - nonzero_histo[-1], 255, rounding_mode='trunc')
# If step is zero, return the original image. Otherwise, build
# lut from the full histogram and step and then index from it.
if step == 0:
result = im
else:
# can't index using 2d index. Have to flatten and then reshape
result = torch.gather(_build_lut(histo, step), 0, im.flatten().long())
result = result.reshape_as(im)
return result / 255.0
[docs]@perform_keep_shape_image
def equalize(input: Tensor) -> Tensor:
r"""Apply equalize on the input tensor.
.. image:: _static/img/equalize.png
Implements Equalize function from PIL using PyTorch ops based on uint8 format:
https://github.com/tensorflow/tpu/blob/5f71c12a020403f863434e96982a840578fdd127/models/official/efficientnet/autoaugment.py#L355
Args:
input: image tensor to equalize with shape :math:`(*, C, H, W)`.
Returns:
Equalized image tensor with shape :math:`(*, C, H, W)`.
Example:
>>> x = torch.rand(1, 2, 3, 3)
>>> equalize(x).shape
torch.Size([1, 2, 3, 3])
"""
res = []
for image in input:
# Assumes RGB for now. Scales each channel independently
# and then stacks the result.
scaled_image = torch.stack([_scale_channel(image[i, :, :]) for i in range(len(image))])
res.append(scaled_image)
return torch.stack(res)
[docs]@perform_keep_shape_video
def equalize3d(input: Tensor) -> Tensor:
r"""Equalize the values for a 3D volumetric tensor.
Implements Equalize function for a sequence of images using PyTorch ops based on uint8 format:
https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py#L352
Args:
input: image tensor with shape :math:`(*, C, D, H, W)` to equalize.
Returns:
Equalized volume with shape :math:`(B, C, D, H, W)`.
"""
res = []
for volume in input:
# Assumes RGB for now. Scales each channel independently
# and then stacks the result.
scaled_input = torch.stack([_scale_channel(volume[i, :, :, :]) for i in range(len(volume))])
res.append(scaled_input)
return torch.stack(res)
[docs]def invert(image: Tensor, max_val: Tensor = Tensor([1.0])) -> Tensor:
r"""Invert the values of an input image tensor by its maximum value.
.. image:: _static/img/invert.png
Args:
image: The input tensor to invert with an arbitatry shape.
max_val: The expected maximum value in the input tensor. The shape has to
according to the input tensor shape, or at least has to work with broadcasting.
Example:
>>> img = torch.rand(1, 2, 4, 4)
>>> invert(img).shape
torch.Size([1, 2, 4, 4])
>>> img = 255. * torch.rand(1, 2, 3, 4, 4)
>>> invert(img, torch.as_tensor(255.)).shape
torch.Size([1, 2, 3, 4, 4])
>>> img = torch.rand(1, 3, 4, 4)
>>> invert(img, torch.as_tensor([[[[1.]]]])).shape
torch.Size([1, 3, 4, 4])
"""
if not isinstance(image, Tensor):
raise AssertionError(f"Input is not a Tensor. Got: {type(input)}")
if not isinstance(max_val, Tensor):
raise AssertionError(f"max_val is not a Tensor. Got: {type(max_val)}")
return max_val.to(image) - image
[docs]class AdjustSaturation(Module):
r"""Adjust color saturation of an image.
The input image is expected to be an RGB image in the range of [0, 1].
Args:
saturation_factor: How much to adjust the saturation. 0 will give a black
and white image, 1 will give the original image while 2 will enhance the saturation by a factor of 2.
saturation_mode: The mode to adjust saturation.
Shape:
- Input: Image/Tensor to be adjusted in the shape of :math:`(*, 3, H, W)`.
- Output: Adjusted image in the shape of :math:`(*, 3, H, W)`.
Example:
>>> x = torch.ones(1, 3, 3, 3)
>>> AdjustSaturation(2.)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]],
<BLANKLINE>
[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]],
<BLANKLINE>
[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 3, 3, 3)
>>> y = torch.ones(2)
>>> out = AdjustSaturation(y)(x)
>>> torch.nn.functional.mse_loss(x, out)
tensor(0.)
"""
def __init__(self, saturation_factor: Union[float, Tensor]) -> None:
super().__init__()
self.saturation_factor: Union[float, Tensor] = saturation_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_saturation(input, self.saturation_factor)
class AdjustSaturationWithGraySubtraction(Module):
r"""Adjust color saturation of an image.
This implementation aligns PIL. Hence, the output is close to TorchVision.
The input image is expected to be in the range of [0, 1].
The input image is expected to be an RGB or gray image in the range of [0, 1].
Args:
saturation_factor: How much to adjust the saturation. 0 will give a black
and white image, 1 will give the original image while 2 will enhance the saturation by a factor of 2.
saturation_mode: The mode to adjust saturation.
Shape:
- Input: Image/Tensor to be adjusted in the shape of :math:`(*, 3, H, W)`.
- Output: Adjusted image in the shape of :math:`(*, 3, H, W)`.
Example:
>>> x = torch.ones(1, 3, 3, 3)
>>> AdjustSaturationWithGraySubtraction(2.)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]],
<BLANKLINE>
[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]],
<BLANKLINE>
[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 3, 3, 3)
>>> y = torch.ones(2)
>>> out = AdjustSaturationWithGraySubtraction(y)(x)
>>> torch.nn.functional.mse_loss(x, out)
tensor(0.)
"""
def __init__(self, saturation_factor: Union[float, Tensor]) -> None:
super().__init__()
self.saturation_factor: Union[float, Tensor] = saturation_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_saturation_with_gray_subtraction(input, self.saturation_factor)
[docs]class AdjustHue(Module):
r"""Adjust hue of an image.
This implementation aligns PIL. Hence, the output is close to TorchVision.
The input image is expected to be in the range of [0, 1].
The input image is expected to be an RGB image in the range of [0, 1].
Args:
hue_factor: How much to shift the hue channel. Should be in [-PI, PI]. PI
and -PI give complete reversal of hue channel in HSV space in positive and negative
direction respectively. 0 means no shift. Therefore, both -PI and PI will give an
image with complementary colors while 0 gives the original image.
Shape:
- Input: Image/Tensor to be adjusted in the shape of :math:`(*, 3, H, W)`.
- Output: Adjusted image in the shape of :math:`(*, 3, H, W)`.
Example:
>>> x = torch.ones(1, 3, 3, 3)
>>> AdjustHue(3.141516)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]],
<BLANKLINE>
[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]],
<BLANKLINE>
[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 3, 3, 3)
>>> y = torch.ones(2) * 3.141516
>>> AdjustHue(y)(x).shape
torch.Size([2, 3, 3, 3])
"""
def __init__(self, hue_factor: Union[float, Tensor]) -> None:
super().__init__()
self.hue_factor: Union[float, Tensor] = hue_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_hue(input, self.hue_factor)
[docs]class AdjustGamma(Module):
r"""Perform gamma correction on an image.
The input image is expected to be in the range of [0, 1].
Args:
gamma: Non negative real number, same as γ\gammaγ in the equation.
gamma larger than 1 make the shadows darker, while gamma smaller than 1 make
dark regions lighter.
gain: The constant multiplier.
Shape:
- Input: Image to be adjusted in the shape of :math:`(*, N)`.
- Output: Adjusted image in the shape of :math:`(*, N)`.
Example:
>>> x = torch.ones(1, 1, 3, 3)
>>> AdjustGamma(1.0, 2.0)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y1 = torch.ones(2) * 1.0
>>> y2 = torch.ones(2) * 2.0
>>> AdjustGamma(y1, y2)(x).shape
torch.Size([2, 5, 3, 3])
"""
def __init__(self, gamma: Union[float, Tensor], gain: Union[float, Tensor] = 1.0) -> None:
super().__init__()
self.gamma: Union[float, Tensor] = gamma
self.gain: Union[float, Tensor] = gain
def forward(self, input: Tensor) -> Tensor:
return adjust_gamma(input, self.gamma, self.gain)
[docs]class AdjustContrast(Module):
r"""Adjust Contrast of an image.
This implementation aligns OpenCV, not PIL. Hence, the output differs from TorchVision.
The input image is expected to be in the range of [0, 1].
Args:
contrast_factor: Contrast adjust factor per element
in the batch. 0 generates a completely black image, 1 does not modify
the input image while any other non-negative number modify the
brightness by this factor.
Shape:
- Input: Image/Input to be adjusted in the shape of :math:`(*, N)`.
- Output: Adjusted image in the shape of :math:`(*, N)`.
Example:
>>> x = torch.ones(1, 1, 3, 3)
>>> AdjustContrast(0.5)(x)
tensor([[[[0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.ones(2)
>>> AdjustContrast(y)(x).shape
torch.Size([2, 5, 3, 3])
"""
def __init__(self, contrast_factor: Union[float, Tensor]) -> None:
super().__init__()
self.contrast_factor: Union[float, Tensor] = contrast_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_contrast(input, self.contrast_factor)
class AdjustContrastWithMeanSubtraction(Module):
r"""Adjust Contrast of an image.
This implementation aligns PIL. Hence, the output is close to TorchVision.
The input image is expected to be in the range of [0, 1].
Args:
contrast_factor: Contrast adjust factor per element
in the batch by subtracting its mean grayscaled version.
0 generates a completely black image, 1 does not modify
the input image while any other non-negative number modify the
brightness by this factor.
Shape:
- Input: Image/Input to be adjusted in the shape of :math:`(*, N)`.
- Output: Adjusted image in the shape of :math:`(*, N)`.
Example:
>>> x = torch.ones(1, 1, 3, 3)
>>> AdjustContrastWithMeanSubtraction(0.5)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.ones(2)
>>> AdjustContrastWithMeanSubtraction(y)(x).shape
torch.Size([2, 5, 3, 3])
"""
def __init__(self, contrast_factor: Union[float, Tensor]) -> None:
super().__init__()
self.contrast_factor: Union[float, Tensor] = contrast_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_contrast_with_mean_subtraction(input, self.contrast_factor)
[docs]class AdjustBrightness(Module):
r"""Adjust Brightness of an image.
This implementation aligns OpenCV, not PIL. Hence, the output differs from TorchVision.
The input image is expected to be in the range of [0, 1].
Args:
brightness_factor: Brightness adjust factor per element
in the batch. 0 does not modify the input image while any other number modify the
brightness.
Shape:
- Input: Image/Input to be adjusted in the shape of :math:`(*, N)`.
- Output: Adjusted image in the shape of :math:`(*, N)`.
Example:
>>> x = torch.ones(1, 1, 3, 3)
>>> AdjustBrightness(1.)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.ones(2)
>>> AdjustBrightness(y)(x).shape
torch.Size([2, 5, 3, 3])
"""
def __init__(self, brightness_factor: Union[float, Tensor]) -> None:
super().__init__()
self.brightness_factor: Union[float, Tensor] = brightness_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_brightness(input, self.brightness_factor)
[docs]class AdjustSigmoid(Module):
r"""Adjust the contrast of an image tensor or performs sigmoid correction on the input image tensor.
The input image is expected to be in the range of [0, 1].
Reference:
[1]: Gustav J. Braun, "Image Lightness Rescaling Using Sigmoidal Contrast Enhancement Functions",
http://markfairchild.org/PDFs/PAP07.pdf
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
cutoff: The cutoff of sigmoid function.
gain: The multiplier of sigmoid function.
inv: If is set to True the function will return the negative sigmoid correction.
Example:
>>> x = torch.ones(1, 1, 2, 2)
>>> AdjustSigmoid(gain=0)(x)
tensor([[[[0.5000, 0.5000],
[0.5000, 0.5000]]]])
"""
def __init__(self, cutoff: float = 0.5, gain: float = 10, inv: bool = False) -> None:
super().__init__()
self.cutoff: float = cutoff
self.gain: float = gain
self.inv: bool = inv
def forward(self, image: Tensor) -> Tensor:
return adjust_sigmoid(image, cutoff=self.cutoff, gain=self.gain, inv=self.inv)
[docs]class AdjustLog(Module):
"""Adjust log correction on the input image tensor.
The input image is expected to be in the range of [0, 1].
Reference:
[1]: http://www.ece.ucsb.edu/Faculty/Manjunath/courses/ece178W03/EnhancePart1.pdf
Args:
image: Image to be adjusted in the shape of :math:`(*, H, W)`.
gain: The multiplier of logarithmic function.
inv: If is set to True the function will return the inverse logarithmic correction.
clip_output: Whether to clip the output image with range of [0, 1].
Example:
>>> x = torch.zeros(1, 1, 2, 2)
>>> AdjustLog(inv=True)(x)
tensor([[[[0., 0.],
[0., 0.]]]])
"""
def __init__(self, gain: float = 1, inv: bool = False, clip_output: bool = True) -> None:
super().__init__()
self.gain: float = gain
self.inv: bool = inv
self.clip_output: bool = clip_output
def forward(self, image: Tensor) -> Tensor:
return adjust_log(image, gain=self.gain, inv=self.inv, clip_output=self.clip_output)
class AdjustBrightnessAccumulative(Module):
r"""Adjust Brightness of an image accumulatively.
This implementation aligns PIL. Hence, the output is close to TorchVision.
The input image is expected to be in the range of [0, 1].
Args:
brightness_factor: Brightness adjust factor per element
in the batch. 0 does not modify the input image while any other number modify the
brightness.
Shape:
- Input: Image/Input to be adjusted in the shape of :math:`(*, N)`.
- Output: Adjusted image in the shape of :math:`(*, N)`.
Example:
>>> x = torch.ones(1, 1, 3, 3)
>>> AdjustBrightnessAccumulative(1.)(x)
tensor([[[[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]]])
>>> x = torch.ones(2, 5, 3, 3)
>>> y = torch.ones(2)
>>> AdjustBrightnessAccumulative(y)(x).shape
torch.Size([2, 5, 3, 3])
"""
def __init__(self, brightness_factor: Union[float, Tensor]) -> None:
super().__init__()
self.brightness_factor: Union[float, Tensor] = brightness_factor
def forward(self, input: Tensor) -> Tensor:
return adjust_brightness_accumulative(input, self.brightness_factor)
[docs]class Invert(Module):
r"""Invert the values of an input tensor by its maximum value.
Args:
input: The input tensor to invert with an arbitatry shape.
max_val: The expected maximum value in the input tensor. The shape has to
according to the input tensor shape, or at least has to work with broadcasting. Default: 1.0.
Example:
>>> img = torch.rand(1, 2, 4, 4)
>>> Invert()(img).shape
torch.Size([1, 2, 4, 4])
>>> img = 255. * torch.rand(1, 2, 3, 4, 4)
>>> Invert(torch.as_tensor(255.))(img).shape
torch.Size([1, 2, 3, 4, 4])
>>> img = torch.rand(1, 3, 4, 4)
>>> Invert(torch.as_tensor([[[[1.]]]]))(img).shape
torch.Size([1, 3, 4, 4])
"""
def __init__(self, max_val: Tensor = torch.tensor(1.0)) -> None:
super().__init__()
if not isinstance(max_val, Parameter):
self.register_buffer("max_val", max_val)
else:
self.max_val = max_val
def forward(self, input: Tensor) -> Tensor:
return invert(input, self.max_val)