Repository: kornia/kornia

Issue:
VisualPrompter for SAM with multiple keypoint prompts leads to IndexError
### Describe the bug

Calling VisualPrompter.predict when multiple keypoint prompts (but no boxes or masks) are given, leads to the following error: `IndexError: The shape of the mask [2, 3] at index 0 does not match the shape of the indexed tensor [1, 3, 256] at index 0`
I believe this is due to this line: https://github.com/kornia/kornia/blob/master/kornia/contrib/visual_prompter.py#L169 where the prompts are unpacked. However, the Augmentation returns an already unpacked list (https://github.com/kornia/kornia/blob/master/kornia/augmentation/container/augment.py#L368). Therefore, the iteration in `VisualPrompter._transform_prompts` is over the first dim of the keypoint tensor, discarding all but the first keypoint prompt. Below is an MWE to reproduce the error. 
I am using Kornia version 0.7.0. I tested it on CPU and GPU, Python 3.10 and 3.11, pytorch 2.1.0 compiled from source as well as 2.0.0 from conda.

### Reproduction steps

```bash
import torch
from kornia.contrib.models.sam import SamConfig
from kornia.contrib.visual_prompter import VisualPrompter

image = torch.rand(3, 256, 256)

keypoints = torch.tensor([[[10, 10], [20, 20]], [[100, 100], [150, 150]]], dtype=torch.float32)
keypoints_labels = torch.tensor([[0, 1], [0, 1]])

config = SamConfig('vit_b', pretrained=True)
prompter = VisualPrompter(config)

prompter.set_image(image)
prompts = prompter.predict(keypoints=keypoints, keypoints_labels=keypoints_labels)
```
```


### Expected behavior

The prompter should not unpack the outputs, when only a single tensor is given, thus no IndexError gets produced.

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
- PyTorch Version (e.g., 1.0): 2.0.0 / 2.1.0
- OS (e.g., Linux): Ubuntu 22.04
- How you installed PyTorch (`conda`, `pip`, source): conda / source
- Build command you used (if compiling from source): `python setup.py install`
- Python version: 3.10 / 3.11
- CUDA/cuDNN version: 8.9.5
- GPU models and configuration: 3x Tesla V100S
- Any other relevant information:
```


### Additional context

_No response_
VisualPrompter for SAM with multiple keypoint prompts leads to IndexError
### Describe the bug

Calling VisualPrompter.predict when multiple keypoint prompts (but no boxes or masks) are given, leads to the following error: `IndexError: The shape of the mask [2, 3] at index 0 does not match the shape of the indexed tensor [1, 3, 256] at index 0`
I believe this is due to this line: https://github.com/kornia/kornia/blob/master/kornia/contrib/visual_prompter.py#L169 where the prompts are unpacked. However, the Augmentation returns an already unpacked list (https://github.com/kornia/kornia/blob/master/kornia/augmentation/container/augment.py#L368). Therefore, the iteration in `VisualPrompter._transform_prompts` is over the first dim of the keypoint tensor, discarding all but the first keypoint prompt. Below is an MWE to reproduce the error. 
I am using Kornia version 0.7.0. I tested it on CPU and GPU, Python 3.10 and 3.11, pytorch 2.1.0 compiled from source as well as 2.0.0 from conda.

### Reproduction steps

```bash
import torch
from kornia.contrib.models.sam import SamConfig
from kornia.contrib.visual_prompter import VisualPrompter

image = torch.rand(3, 256, 256)

keypoints = torch.tensor([[[10, 10], [20, 20]], [[100, 100], [150, 150]]], dtype=torch.float32)
keypoints_labels = torch.tensor([[0, 1], [0, 1]])

config = SamConfig('vit_b', pretrained=True)
prompter = VisualPrompter(config)

prompter.set_image(image)
prompts = prompter.predict(keypoints=keypoints, keypoints_labels=keypoints_labels)
```
```


### Expected behavior

The prompter should not unpack the outputs, when only a single tensor is given, thus no IndexError gets produced.

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
- PyTorch Version (e.g., 1.0): 2.0.0 / 2.1.0
- OS (e.g., Linux): Ubuntu 22.04
- How you installed PyTorch (`conda`, `pip`, source): conda / source
- Build command you used (if compiling from source): `python setup.py install`
- Python version: 3.10 / 3.11
- CUDA/cuDNN version: 8.9.5
- GPU models and configuration: 3x Tesla V100S
- Any other relevant information:
```


### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/kornia/contrib/models/structures.py b/kornia/contrib/models/structures.py
--- a/kornia/contrib/models/structures.py
+++ b/kornia/contrib/models/structures.py
@@ -14,7 +14,7 @@ class SegmentationResults:
 
     Args:
         logits: Results logits with shape :math:`(B, C, H, W)`, where :math:`C` refers to the number of predicted masks
-        scores: The scores from the logits. Shape :math:`(B,)`
+        scores: The scores from the logits. Shape :math:`(B, C)`
         mask_threshold: The threshold value to generate the `binary_masks` from the `logits`
     """
 
diff --git a/kornia/contrib/visual_prompter.py b/kornia/contrib/visual_prompter.py
--- a/kornia/contrib/visual_prompter.py
+++ b/kornia/contrib/visual_prompter.py
@@ -166,6 +166,10 @@ def _transform_prompts(
         self, *prompts: Tensor | Boxes | Keypoints, data_keys: list[str] = []
     ) -> dict[str, Tensor | Boxes | Keypoints]:
         transformed_prompts = self.transforms(*prompts, data_keys=data_keys, params=self._tfs_params)
+
+        # prevent unpacking tensor when creating the output dict (issue #2627)
+        if not isinstance(transformed_prompts, (list, tuple)):
+            transformed_prompts = [transformed_prompts]
         return {key: transformed_prompts[idx] for idx, key in enumerate(data_keys)}
 
     def preprocess_prompts(
@@ -197,7 +201,7 @@ def preprocess_prompts(
         if 'keypoints' in data and isinstance(data['keypoints'], Keypoints):
             kpts_tensor = data['keypoints'].to_tensor()
             if KORNIA_CHECK_IS_TENSOR(kpts_tensor) and KORNIA_CHECK_IS_TENSOR(keypoints_labels):
-                points = (kpts_tensor[None, ...], keypoints_labels)
+                points = (kpts_tensor, keypoints_labels)
         else:
             points = None
 


Test patch:
diff --git a/test/contrib/test_prompter.py b/test/contrib/test_prompter.py
--- a/test/contrib/test_prompter.py
+++ b/test/contrib/test_prompter.py
@@ -33,11 +33,11 @@ def test_cardinality(self, device, batch_size, N, multimask_output):
 
         prompter.set_image(inpt)
 
-        out = prompter.predict((keypoints, labels), multimask_output=multimask_output)
+        out = prompter.predict(keypoints, labels, multimask_output=multimask_output)
 
         C = 3 if multimask_output else 1
-        assert out.logits.shape == (1, C, 256, 256)
-        assert out.scores.shape == (1, C)
+        assert out.logits.shape == (batch_size, C, 256, 256)
+        assert out.scores.shape == (batch_size, C)
 
     def test_exception(self):
         prompter = VisualPrompter(SamConfig('vit_b'))