Repository: kornia/kornia

Issue:
apply_colormap off-by-one error
### Describe the bug

Trying to use the `kornia.color.apply_colormap()` function gives the error:
```
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in <module>:20                                                                                   │
│                                                                                                  │
│   17                                                                                             │
│   18 index = torch.bucketize(input_tensor, keys)  # shape equals <input_tensor>: (B x H x W)     │
│   19                                                                                             │
│ ❱ 20 output = colormap.colors[:, index]  # (3 x B x H x W) or (3 x H x W)                        │
│   21                                                                                             │
│   22 if len(output.shape) == 4:                                                                  │
│   23 │   output = output.permute(1, 0, -2, -1)  # (B x 3 x H x W)                                │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
IndexError: index 64 is out of bounds for dimension 0 with size 64
```

### Reproduction steps

```bash
import kornia

input_tensor = torch.stack([torch.eye(4)* 255, torch.eye(4) * 150]).view(2, -1, 4, 4)
kornia.color.apply_colormap(input_tensor, kornia.color.AUTUMN())
```


### Expected behavior

Get a colormap, as shown in the documentation.

### Environment

```shell
- PyTorch Version (e.g., 1.0): 2.0.0+cu118
- OS (e.g., Linux): Linux
- Python version: 3.10
```


### Additional context

The solution is to change the line:
```
keys = torch.arange(0, len(colormap), dtype=input_tensor.dtype, device=input_tensor.device)  # (num_colors)
```

To:
```
keys = torch.arange(0, len(colormap) - 1, dtype=input_tensor.dtype, device=input_tensor.device)  # (num_colors)
```


Discussion:

Resolving patch:
diff --git a/kornia/color/colormap.py b/kornia/color/colormap.py
--- a/kornia/color/colormap.py
+++ b/kornia/color/colormap.py
@@ -159,7 +159,7 @@ def apply_colormap(input_tensor: Tensor, colormap: ColorMap) -> Tensor:
     elif len(input_tensor.shape) == 3 and input_tensor.shape[0] == 1:  # if (1 X H x W)
         input_tensor = input_tensor[0, ...]  # (H x W)
 
-    keys = torch.arange(0, len(colormap), dtype=input_tensor.dtype, device=input_tensor.device)  # (num_colors)
+    keys = torch.arange(0, len(colormap) - 1, dtype=input_tensor.dtype, device=input_tensor.device)  # (num_colors)
 
     index = torch.bucketize(input_tensor, keys)  # shape equals <input_tensor>: (B x H x W) or (H x W)
 


Test patch:
diff --git a/test/color/test_colormap.py b/test/color/test_colormap.py
--- a/test/color/test_colormap.py
+++ b/test/color/test_colormap.py
@@ -44,6 +44,23 @@ def test_smoke(self, device, dtype):
 
         self.assert_close(actual, expected_tensor)
 
+    def test_eye(self, device, dtype):
+        input_tensor = torch.stack(
+            [torch.eye(2, dtype=dtype, device=device) * 255, torch.eye(2, dtype=dtype, device=device) * 150]
+        ).view(2, -1, 2, 2)
+
+        expected_tensor = tensor(
+            [
+                [[[1.0, 1.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]],
+                [[[1.0, 1.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+        actual = apply_colormap(input_tensor, AUTUMN(device=device, dtype=dtype))
+        self.assert_close(actual, expected_tensor)
+
     def test_exception(self, device, dtype):
         cm = AUTUMN(device=device, dtype=dtype)
         with pytest.raises(TypeError):