Return single top class confidence for mulit-label predictions

inference/core/workflows/core_steps/common/query_language/entities/enums.py

@@ -21,6 +21,7 @@ class SequenceAggregationMode(Enum):
 class ClassificationProperty(Enum):
     TOP_CLASS = "top_class"
     TOP_CLASS_CONFIDENCE = "top_class_confidence"
+    TOP_CLASS_CONFIDENCE_SINGLE = "top_class_confidence_single"
     ALL_CLASSES = "all_classes"
     ALL_CONFIDENCES = "all_confidences"
 

inference/core/workflows/core_steps/common/query_language/operations/classification_results/base.py

@@ -27,6 +27,19 @@ def extract_top_class_confidence(prediction: dict) -> Union[float, List[float]]:
     ]
 
 
+def extract_top_class_confidence_single(prediction: dict) -> Union[float, List[float]]:
+    if "confidence" in prediction:
+        return prediction["confidence"]
+    predicted_classes = prediction.get("predicted_classes", [])
+    predicted_confidences = [
+        prediction["predictions"][class_name]["confidence"]
+        for class_name in predicted_classes
+    ]
+    if not predicted_confidences:
+        return 0.0
+    return max(predicted_confidences)
+
+
 def extract_all_class_names(prediction: dict) -> List[str]:
     predictions = prediction["predictions"]
     if isinstance(predictions, list):
@@ -52,6 +65,7 @@ def extract_all_classes_confidence(prediction: dict) -> List[float]:
 CLASSIFICATION_PROPERTY_EXTRACTORS = {
     ClassificationProperty.TOP_CLASS: extract_top_class,
     ClassificationProperty.TOP_CLASS_CONFIDENCE: extract_top_class_confidence,
+    ClassificationProperty.TOP_CLASS_CONFIDENCE_SINGLE: extract_top_class_confidence_single,
     ClassificationProperty.ALL_CLASSES: extract_all_class_names,
     ClassificationProperty.ALL_CONFIDENCES: extract_all_classes_confidence,
 }

tests/workflows/unit_tests/core_steps/common/query_language/operations/test_classification_results_operations.py

@@ -197,6 +197,58 @@ def test_classification_result_extraction_of_top_class_confidence_for_multi_labe
     assert result == [0.4]
 
 
+def test_classification_result_extraction_of_top_class_confidence_single_for_multi_label_classification_result_when_class_detected() -> (
+    None
+):
+    # given
+    operations = [
+        {
+            "type": "ClassificationPropertyExtract",
+            "property_name": "top_class_confidence_single",
+        }
+    ]
+    data = MultiLabelClassificationInferenceResponse(
+        image=InferenceResponseImage(width=128, height=256),
+        predictions={
+            "cat": MultiLabelClassificationPrediction(class_id=0, confidence=0.6),
+            "dog": MultiLabelClassificationPrediction(class_id=1, confidence=0.4),
+        },
+        predicted_classes=["cat", "dog"],
+    ).dict(by_alias=True, exclude_none=True)
+
+    # when
+    result = execute_operations(value=data, operations=operations)
+
+    # then
+    assert result == 0.6
+
+
+def test_classification_result_extraction_of_top_class_confidence_single_for_multi_label_classification_result_when_no_classes_detected() -> (
+    None
+):
+    # given
+    operations = [
+        {
+            "type": "ClassificationPropertyExtract",
+            "property_name": "top_class_confidence_single",
+        }
+    ]
+    data = MultiLabelClassificationInferenceResponse(
+        image=InferenceResponseImage(width=128, height=256),
+        predictions={
+            "cat": MultiLabelClassificationPrediction(class_id=0, confidence=0.6),
+            "dog": MultiLabelClassificationPrediction(class_id=1, confidence=0.4),
+        },
+        predicted_classes=[],
+    ).dict(by_alias=True, exclude_none=True)
+
+    # when
+    result = execute_operations(value=data, operations=operations)
+
+    # then
+    assert result == 0.0
+
+
 def test_classification_result_extraction_of_all_classes_for_multi_class_classification_result() -> (
     None
 ):