Fix classification confidence shape handling (#2262)

* Fix classification confidence shape handling

* Cover short classification confidence vectors

* Run make style

Author: hansent

inference/core/models/inference_models_adapters.py

@@ -36,6 +36,7 @@
     RFDETR_ONNX_MAX_RESOLUTION,
     VALID_INFERENCE_MODELS_BACKENDS,
 )
+from inference.core.exceptions import PostProcessingError
 from inference.core.models.base import Model
 from inference.core.roboflow_api import get_extra_weights_provider_headers
 from inference.core.utils.image_utils import load_image_bgr, load_image_rgb
@@ -767,12 +768,17 @@ def prepare_multi_label_classification_response(
     """
     results = []
     for prediction, image_size in zip(post_processed_predictions, image_sizes):
+        class_confidences = _reshape_classification_confidences(
+            confidence=prediction.confidence.cpu(),
+            expected_num_images=1,
+            class_names=class_names,
+        )[0].tolist()
         image_predictions_dict = {
             class_names[class_id]: {
                 "confidence": confidence,
                 "class_id": class_id,
             }
-            for class_id, confidence in enumerate(prediction.confidence.cpu().tolist())
+            for class_id, confidence in enumerate(class_confidences)
         }
         predicted_classes = [
             class_names[class_id] for class_id in prediction.class_ids.tolist()
@@ -795,9 +801,12 @@ def prepare_classification_response(
     confidence_threshold: float,
 ) -> List[ClassificationInferenceResponse]:
     responses = []
-    for classes_confidence, image_size in zip(
-        post_processed_predictions.confidence.cpu().tolist(), image_sizes
-    ):
+    batch_confidences = _reshape_classification_confidences(
+        confidence=post_processed_predictions.confidence.cpu(),
+        expected_num_images=len(image_sizes),
+        class_names=class_names,
+    )
+    for classes_confidence, image_size in zip(batch_confidences.tolist(), image_sizes):
         individual_classes_predictions = []
         for i, cls_name in enumerate(class_names):
             class_score = float(classes_confidence[i])
@@ -831,6 +840,26 @@ def prepare_classification_response(
     return responses
 
 
+def _reshape_classification_confidences(
+    confidence: torch.Tensor,
+    expected_num_images: int,
+    class_names: List[str],
+) -> torch.Tensor:
+    expected_num_classes = len(class_names)
+    expected_num_scores = expected_num_images * expected_num_classes
+    actual_num_scores = confidence.numel()
+    if actual_num_scores != expected_num_scores:
+        raise PostProcessingError(
+            "Classification model output has shape "
+            f"{tuple(confidence.shape)} containing {actual_num_scores} confidence "
+            f"score(s), but response metadata expects {expected_num_images} image(s) "
+            f"x {expected_num_classes} class name(s) = {expected_num_scores} score(s). "
+            "This usually means the model package class names metadata does not match "
+            "the classifier head."
+        )
+    return confidence.reshape(expected_num_images, expected_num_classes)
+
+
 def draw_predictions(inference_request, inference_response, class_names: List[str]):
     """Draw prediction visuals on an image.
 

tests/inference/unit_tests/core/models/test_inference_models_adapters.py

@@ -4,9 +4,14 @@
 import torch
 
 from inference.core.models.inference_models_adapters import (
+    prepare_classification_response,
     prepare_multi_label_classification_response,
 )
-from inference_models import MultiLabelClassificationPrediction
+from inference.core.exceptions import PostProcessingError
+from inference_models import (
+    ClassificationPrediction,
+    MultiLabelClassificationPrediction,
+)
 
 
 def test_prepare_multi_label_response_uses_class_ids_for_predicted_classes() -> None:
@@ -46,3 +51,38 @@ def test_prepare_multi_label_response_uses_class_ids_for_predicted_classes() ->
     assert r.predictions["d"].confidence == pytest.approx(0.9)
     # Only the model's filtered class_ids show up in predicted_classes.
     assert r.predicted_classes == ["c"]
+
+
+def test_prepare_classification_response_flattens_singleton_output_dimensions() -> None:
+    class_names = ["cat", "dog"]
+    prediction = ClassificationPrediction(
+        class_id=torch.tensor([[1]], dtype=torch.long),
+        confidence=torch.tensor([[[0.1, 0.9]]]),
+    )
+
+    results = prepare_classification_response(
+        post_processed_predictions=prediction,
+        image_sizes=[(10, 20)],
+        class_names=class_names,
+        confidence_threshold=0.0,
+    )
+
+    assert len(results) == 1
+    assert results[0].top == "dog"
+    assert results[0].confidence == pytest.approx(0.9)
+    assert [p.class_name for p in results[0].predictions] == ["dog", "cat"]
+
+
+def test_prepare_classification_response_fails_on_class_count_mismatch() -> None:
+    prediction = ClassificationPrediction(
+        class_id=torch.tensor([0], dtype=torch.long),
+        confidence=torch.tensor([[0.7]]),
+    )
+
+    with pytest.raises(PostProcessingError, match="class names metadata"):
+        prepare_classification_response(
+            post_processed_predictions=prediction,
+            image_sizes=[(10, 20)],
+            class_names=["cat", "dog"],
+            confidence_threshold=0.0,
+        )