Bug in normalize_min_max

[vgilabert94]

Describe the bug

From #2860 (comment)

If the input tensor is an image without batch (C,H,W), where the normalisation would be applied incorrectly (without giving an error).

A possible solution is to convert the input to batch (B,C,H,W) inside the function and also modify the view by reshaping to ensure the continuity of the tensors in the torch. Another solution could be to adapt the normalize_min_max internally to accept non-batched inputs.

Reproduction steps

import torch
from kornia.enhance import normalize_min_max

rng = torch.manual_seed(1)
torch.set_printoptions(precision=10)
input_tensor = torch.rand(1,3,3)
res = normalize_min_max(input_tensor)
input_tensor = input_tensor.reshape(1,1,3,3)
expected = normalize_min_max(input_tensor)
print(input_tensor)
tensor([[[[0.4387779236, 0.6386804581, 0.5246658921],
          [0.6826140881, 0.3051494956, 0.4635456204],
          [0.4549863338, 0.5724719763, 0.4980025887]]]])

print(res)
tensor([[[0.9999978542, 0.0000000000, 0.2587345839],
         [0.9154204130, 0.0000000000, 0.9999986887],
         [0.0000000000, 0.9999971390, 0.4825139046]]])

print(expected)
tensor([[[[0.9451995492, 0.3244698048, 0.4850744009],
          [0.9154204130, 0.0000000000, 0.9999986887],
          [0.4773764014, 0.9409694076, 0.7010670900]]]])

Expected behavior

Excepted should be normalize having into account only the H and W channels of images (like in expected result).

Environment

PyTorch version: 2.1.2+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr  5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2803
DeviceID=CPU0
Family=198
L2CacheSize=5120
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2803
Name=11th Gen Intel(R) Core(TM) i7-1165G7 @ 2.80GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] mypy==1.8.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.3
[pip3] onnx==1.15.0
[pip3] torch==2.1.2
[pip3] torchaudio==2.1.2
[pip3] torchvision==0.16.2
[conda] Could not collect

Additional context

No response

[vgilabert94]

Something like this is working as expected. I can create a PR if you think it's right.

def normalize_min_max(x: torch.Tensor, min_val: float = 0.0, max_val: float = 1.0, eps: float = 1e-6) -> torch.Tensor:
    if not isinstance(x, torch.Tensor):
        raise TypeError(f"data should be a tensor. Got: {type(x)}.")

    if not isinstance(min_val, float):
        raise TypeError(f"'min_val' should be a float. Got: {type(min_val)}.")

    if not isinstance(max_val, float):
        raise TypeError(f"'b' should be a float. Got: {type(max_val)}.")

    orig_shape = x.shape
    if len(orig_shape) not in [3, 4]:
        raise ValueError(f"Input shape must be at least a 3d tensor. Got: {x.shape}.")

    if len(orig_shape) == 3:
        x = x.unsqueeze(0)

    (B, C) = x.shape[:2]
    x = x.reshape(B, C, -1)
    x_min: torch.Tensor = x.min(-1)[0].reshape(B, C, 1)
    x_max: torch.Tensor = x.max(-1)[0].reshape(B, C, 1)
    x_out: torch.Tensor = (max_val - min_val) * (x - x_min) / (x_max - x_min + eps) + min_val

    return x_out.reshape(orig_shape)

[edgarriba]

let's convert to batch and unbatch in the end if needed. Be aware that we have a perform_keep_shape_image decorator for that:

kornia/kornia/utils/image.py

Line 214 in d1a1cc0

def perform_keep_shape_image(f: Callable[..., Tensor]) -> Callable[..., Tensor]:

I remember there was a blocker with torchscript which we should revisit if with torch compile is working or not. /cc @johnnv1 @shijianjian