Source code for rising.transforms.functional.intensity

import torch

from typing import Union, Sequence, Optional

from rising.utils import check_scalar

__all__ = ["norm_range", "norm_min_max", "norm_zero_mean_unit_std", "norm_mean_std",
           "add_noise", "add_value", "gamma_correction", "scale_by_value", "clamp"]


def clamp(data: torch.Tensor, min: float, max: float,
          out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Clamp tensor to minimal and maximal value

    Args:
        data: tensor to clamp
        min: lower limit
        max: upper limit
        out: output tensor

    Returns:
        Tensor: clamped tensor
    """
    return torch.clamp(data, min=float(min), max=float(max), out=out)


def norm_range(data: torch.Tensor, min: float, max: float,
               per_channel: bool = True,
               out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Scale range of tensor

    Args:
        data: input data. Per channel option supports [C,H,W] and [C,H,W,D].
        min: minimal value
        max: maximal value
        per_channel: range is normalized per channel
        out: if provided, result is saved in here

    Returns:
        torch.Tensor: normalized data
    """
    if out is None:
        out = torch.zeros_like(data)

    out = norm_min_max(data, per_channel=per_channel, out=out)
    _range = max - min
    out = (out * _range) + min
    return out


def norm_min_max(data: torch.Tensor, per_channel: bool = True,
                 out: Optional[torch.Tensor] = None,
                 eps: Optional[float] = 1e-8) -> torch.Tensor:
    """
    Scale range to [0,1]

    Args:
        data: input data. Per channel option supports [C,H,W] and [C,H,W,D].
        per_channel: range is normalized per channel
        out:  if provided, result is saved in here
        eps: small constant for numerical stability.
            If None, no factor constant will be added

    Returns:
        torch.Tensor: scaled data
    """
    def _norm(_data: torch.Tensor, _out: torch.Tensor):
        _min = _data.min()
        _range = _data.max() - _min
        if eps is not None:
            _range = _range + eps
        _out = (_data - _min) / _range
        return _out

    if out is None:
        out = torch.zeros_like(data)

    if per_channel:
        for _c in range(data.shape[0]):
            out[_c] = _norm(data[_c], out[_c])
    else:
        out = _norm(data, out)
    return out


def norm_zero_mean_unit_std(data: torch.Tensor, per_channel: bool = True,
                            out: Optional[torch.Tensor] = None,
                            eps: Optional[float] = 1e-8) -> torch.Tensor:
    """
    Normalize mean to zero and std to one

    Args:
        data: input data. Per channel option supports [C,H,W] and [C,H,W,D].
        per_channel: range is normalized per channel
        out: if provided, result is saved in here
        eps: small constant for numerical stability.
            If None, no factor constant will be added

    Returns:
        torch.Tensor: normalized data
    """
    def _norm(_data: torch.Tensor, _out: torch.Tensor):
        denom = _data.std()
        if eps is not None:
            denom = denom + eps
        _out = (_data - _data.mean()) / denom
        return _out

    if out is None:
        out = torch.zeros_like(data)

    if per_channel:
        for _c in range(data.shape[0]):
            out[_c] = _norm(data[_c], out[_c])
    else:
        out = _norm(data, out)
    return out


def norm_mean_std(data: torch.Tensor, mean: Union[float, Sequence],
                  std: Union[float, Sequence],
                  per_channel: bool = True,
                  out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Normalize mean and std with provided values

    Args:
        data:input data. Per channel option supports [C,H,W] and [C,H,W,D].
        mean: used for mean normalization
        std: used for std normalization
        per_channel: range is normalized per channel
        out: if provided, result is saved into out

    Returns:
        torch.Tensor: normalized data
    """
    if out is None:
        out = torch.zeros_like(data)

    if per_channel:
        if check_scalar(mean):
            mean = [mean] * data.shape[0]
        if check_scalar(std):
            std = [std] * data.shape[0]
        for _c in range(data.shape[0]):
            out[_c] = (data[_c] - mean[_c]) / std[_c]
    else:
        out = (data - mean) / std

    return out


def add_noise(data: torch.Tensor, noise_type: str,
              out: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
    """
    Add noise to input

    Args:
        data: input data
        noise_type: supports all inplace functions of a pytorch tensor
        out: if provided, result is saved in here
        kwargs: keyword arguments passed to generating function

    Returns:
        torch.Tensor: data with added noise

    See Also:
        :func:`torch.Tensor.normal_`, :func:`torch.Tensor.exponential_`
    """
    if not noise_type.endswith('_'):
        noise_type = noise_type + '_'
    noise_tensor = torch.empty_like(data, requires_grad=False)
    getattr(noise_tensor, noise_type)(**kwargs)
    return torch.add(data, noise_tensor, out=out)


def gamma_correction(data: torch.Tensor, gamma: float) -> torch.Tensor:
    """
    Apply gamma correction to data
    (currently this functions is intended as an interface in case
    additional functionality should be added to transform)

    Args:
        data: input data
        gamma: gamma for correction

    Returns:
        torch.Tensor: gamma corrected data
    """
    if torch.is_tensor(gamma):
        gamma = gamma.to(data)
    return data.pow(gamma)


def add_value(data: torch.Tensor, value: float, out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Increase brightness additively by value
    (currently this functions is intended as an interface in case
    additional functionality should be added to transform)

    Args:
        data: input data
        value: additive value
        out: if provided, result is saved in here

    Returns:
        torch.Tensor: augmented data
    """
    return torch.add(data, value, out=out)


def scale_by_value(data: torch.Tensor, value: float,
                   out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Increase brightness scaled by value
    (currently this functions is intended as an interface in case
    additional functionality should be added to transform)

    Args:
        data: input data
        value: scaling value
        out: if provided, result is saved in here

    Returns:
        torch.Tensor: augmented data
    """
    return torch.mul(data, value, out=out)