Repository: Lightning-AI/torchmetrics

Issue:
Unpredictable class order when `panoptic_quality(..., return_per_class=True)`
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

PanopticQuality can return per class score when `return_per_class=True` argument is given.
But the order of the output classes is unpredictable.

### To Reproduce

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
```py
def pq_demo(cats):
    a = [cats[0], 0]
    b = [cats[1], 0]
    c = [cats[2], 0]

    targs = torch.tensor([[a, a, b, b, c, c]])
    preds = torch.tensor([[a, a, b, b, b, c]])
    score = panoptic_quality(preds, targs, things=[], stuffs=cats, return_per_class=True)
    print(score)

pq_demo([0, 2, 1])
pq_demo([0, 3, 2])
pq_demo([0, 10, 2])

# OUTPUT:
# tensor([[1.0000, 0.0000, 0.6667]], dtype=torch.float64)
# tensor([[1.0000, 0.0000, 0.6667]], dtype=torch.float64)
# tensor([[1.0000, 0.6667, 0.0000]], dtype=torch.float64)
```

Three examples should return identical scores but actually the order has been changed.

(Note that it is not sure if this can be reproduced in a different environment as described below)

Internally, the order of the classes are determined by the value of `cat_id_to_continuous_id` the values are defined here:

https://github.com/Lightning-AI/torchmetrics/blob/4ecfde73216455ac61c406f6474bdc645fa3d13e/src/torchmetrics/functional/detection/_panoptic_quality_common.py#L139-L157

There, the `things` and `stuffs` are enumerated by `enumerate(things)`, but these values are objects of `Set` class that is "unordered collection". So the order of the resulting lists and the `cat_id_to_continuous_id` also has unpredictable order. 

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->


### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

The output classes are ordered by `things` and `stuffs`, and numerically sorted within each.
Ex. with `things=[4, 1], stuffs=[3, 2]`, the output classes are ordered by `[1, 4, 2, 3]`.

I guessed the expected behavior from the implementation because I could not find a description of the order of output classes.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): build the master branch
- Python & PyTorch Version (e.g., 1.0): '1.5.0dev'
- Any other relevant information such as OS (e.g., Linux): Linux (Ubuntu 22.04)

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/src/torchmetrics/detection/panoptic_qualities.py b/src/torchmetrics/detection/panoptic_qualities.py
--- a/src/torchmetrics/detection/panoptic_qualities.py
+++ b/src/torchmetrics/detection/panoptic_qualities.py
@@ -66,9 +66,11 @@ class PanopticQuality(Metric):
           single scalar tensor is returned with average panoptic quality over all classes. If ``return_sq_and_rq=True``
           and ``return_per_class=False`` a tensor of length 3 is returned with panoptic, segmentation and recognition
           quality (in that order). If If ``return_sq_and_rq=False`` and ``return_per_class=True`` a tensor of length
-          equal to the number of classes are returned, with panoptic quality for each class. Finally, if both arguments
-          are ``True`` a tensor of shape ``(3, C)`` is returned with individual panoptic, segmentation and recognition
-          quality for each class.
+          equal to the number of classes are returned, with panoptic quality for each class. The order of classes is
+          ``things`` first and then ``stuffs``, and numerically sorted within each.
+          (ex. with ``things=[4, 1], stuffs=[3, 2]``, the output classes are ordered by ``[1, 4, 2, 3]``)
+          Finally, if both arguments are ``True`` a tensor of shape ``(3, C)`` is returned with individual panoptic,
+          segmentation and recognition quality for each class.
 
     Args:
         things:
diff --git a/src/torchmetrics/functional/detection/_panoptic_quality_common.py b/src/torchmetrics/functional/detection/_panoptic_quality_common.py
--- a/src/torchmetrics/functional/detection/_panoptic_quality_common.py
+++ b/src/torchmetrics/functional/detection/_panoptic_quality_common.py
@@ -148,9 +148,9 @@ def _get_category_id_to_continuous_id(things: Set[int], stuffs: Set[int]) -> Dic
 
     """
     # things metrics are stored with a continuous id in [0, len(things)[,
-    thing_id_to_continuous_id = {thing_id: idx for idx, thing_id in enumerate(things)}
+    thing_id_to_continuous_id = {thing_id: idx for idx, thing_id in enumerate(sorted(things))}
     # stuff metrics are stored with a continuous id in [len(things), len(things) + len(stuffs)[
-    stuff_id_to_continuous_id = {stuff_id: idx + len(things) for idx, stuff_id in enumerate(stuffs)}
+    stuff_id_to_continuous_id = {stuff_id: idx + len(things) for idx, stuff_id in enumerate(sorted(stuffs))}
     cat_id_to_continuous_id = {}
     cat_id_to_continuous_id.update(thing_id_to_continuous_id)
     cat_id_to_continuous_id.update(stuff_id_to_continuous_id)


Test patch:
diff --git a/tests/unittests/detection/test_panoptic_quality.py b/tests/unittests/detection/test_panoptic_quality.py
--- a/tests/unittests/detection/test_panoptic_quality.py
+++ b/tests/unittests/detection/test_panoptic_quality.py
@@ -61,10 +61,26 @@
     .reshape((1, 1, 5, 2))
     .repeat(2, 1, 1, 1),
 )
+
 _ARGS_0 = {"things": {0, 1}, "stuffs": {6, 7}}
 _ARGS_1 = {"things": {2}, "stuffs": {3}, "allow_unknown_preds_category": True}
 _ARGS_2 = {"things": {0, 1}, "stuffs": {10, 11}}
 
+
+def _get_class_order_test_input_args(class_type, class1, class2, class3) -> (np.ndarray, dict):
+    a = [class1, 0]
+    b = [class2, 0]
+    c = [class3, 0]
+    _input = _Input(
+        # Shape of input tensors is (num_batches, batch_size, num_points, 2).
+        preds=torch.tensor([a, a, b, b, b, c]).reshape((1, 1, 6, 2)).repeat(2, 1, 1, 1),
+        target=torch.tensor([a, a, b, b, c, c]).reshape((1, 1, 6, 2)).repeat(2, 1, 1, 1),
+    )
+    _args = {"things": [], "stuffs": [], "return_per_class": True}
+    _args[class_type] = [class1, class2, class3]
+    return _input, _args
+
+
 # TODO: Improve _reference_fn by calling https://github.com/cocodataset/panopticapi/blob/master/panopticapi/evaluation.py
 # directly and compare at runtime on multiple examples.
 
@@ -84,6 +100,11 @@ def _reference_fn_1_2(preds, target) -> np.ndarray:
     return np.array([(2 / 3 + 1 + 2 / 3) / 3])
 
 
+def _reference_fn_class_order(preds, target) -> np.ndarray:
+    """Baseline result for the result of _get_class_order_test_input_args."""
+    return np.array([1, 0, 2 / 3])
+
+
 @pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_12, reason="PanopticQuality metric only supports PyTorch >= 1.12")
 class TestPanopticQuality(MetricTester):
     """Test class for `PanopticQuality` metric."""
@@ -95,6 +116,12 @@ class TestPanopticQuality(MetricTester):
             (_INPUTS_0, _ARGS_0, _reference_fn_0_0),
             (_INPUTS_0, _ARGS_1, _reference_fn_0_1),
             (_INPUTS_1, _ARGS_2, _reference_fn_1_2),
+            (*_get_class_order_test_input_args("stuffs", 0, 2, 1), _reference_fn_class_order),
+            (*_get_class_order_test_input_args("stuffs", 0, 3, 2), _reference_fn_class_order),
+            (*_get_class_order_test_input_args("stuffs", 0, 10, 2), _reference_fn_class_order),
+            (*_get_class_order_test_input_args("things", 0, 2, 1), _reference_fn_class_order),
+            (*_get_class_order_test_input_args("things", 0, 3, 2), _reference_fn_class_order),
+            (*_get_class_order_test_input_args("things", 0, 10, 2), _reference_fn_class_order),
         ],
     )
     def test_panoptic_quality_class(self, ddp, inputs, args, reference_metric):