Merge pull request #1359 from roboflow/fix/add-default-fallback-to-class-replacement

Fix/add default fallback to class replacement

Author: PawelPeczek-Roboflow

inference/core/workflows/core_steps/fusion/detections_classes_replacement/v1.py

@@ -1,4 +1,5 @@
-from typing import Dict, List, Literal, Optional, Tuple, Type
+import sys
+from typing import Dict, List, Literal, Optional, Tuple, Type, Union
 from uuid import uuid4
 
 import numpy as np
@@ -17,8 +18,10 @@
 from inference.core.workflows.execution_engine.entities.types import (
     CLASSIFICATION_PREDICTION_KIND,
     INSTANCE_SEGMENTATION_PREDICTION_KIND,
+    INTEGER_KIND,
     KEYPOINT_DETECTION_PREDICTION_KIND,
     OBJECT_DETECTION_PREDICTION_KIND,
+    STRING_KIND,
     Selector,
 )
 from inference.core.workflows.prototypes.block import (
@@ -75,6 +78,19 @@ class BlockManifest(WorkflowBlockManifest):
         description="The output of classification model for crops taken based on RoIs pointed as the other parameter",
         examples=["$steps.my_classification_model.predictions"],
     )
+    fallback_class_name: Union[Optional[str], Selector(kind=[STRING_KIND])] = Field(
+        default=None,
+        title="Fallback class name",
+        description="The class name to be used as a fallback if no class is predicted for a bounding box",
+        examples=["unknown"],
+    )
+    fallback_class_id: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field(
+        default=None,
+        title="Fallback class id",
+        description="The class id to be used as a fallback if no class is predicted for a bounding box;"
+        f"if not specified or negative, the class id will be set to {sys.maxsize}",
+        examples=[77],
+    )
 
     @classmethod
     def accepts_empty_values(cls) -> bool:
@@ -119,6 +135,8 @@ def run(
         self,
         object_detection_predictions: Optional[sv.Detections],
         classification_predictions: Optional[Batch[Optional[dict]]],
+        fallback_class_name: Optional[str],
+        fallback_class_id: Optional[int],
     ) -> BlockResult:
         if object_detection_predictions is None:
             return {"predictions": None}
@@ -131,7 +149,9 @@ def run(
             return {"predictions": sv.Detections.empty()}
         detection_id_by_class: Dict[str, Optional[Tuple[str, int]]] = {
             prediction[PARENT_ID_KEY]: extract_leading_class_from_prediction(
-                prediction=prediction
+                prediction=prediction,
+                fallback_class_name=fallback_class_name,
+                fallback_class_id=fallback_class_id,
             )
             for prediction in classification_predictions
             if prediction is not None
@@ -180,8 +200,20 @@ def run(
 
 def extract_leading_class_from_prediction(
     prediction: dict,
+    fallback_class_name: Optional[str] = None,
+    fallback_class_id: Optional[int] = None,
 ) -> Optional[Tuple[str, int, float]]:
     if "top" in prediction:
+        if not prediction.get("predictions") and not fallback_class_name:
+            return None
+        elif not prediction.get("predictions") and fallback_class_name:
+            try:
+                fallback_class_id = int(fallback_class_id)
+            except ValueError:
+                fallback_class_id = None
+            if fallback_class_id is None or fallback_class_id < 0:
+                fallback_class_id = sys.maxsize
+            return fallback_class_name, fallback_class_id, 0
         class_name = prediction["top"]
         matching_class_ids = [
             (p["class_id"], p["confidence"])

tests/workflows/unit_tests/core_steps/fusion/test_detections_classes_replacement.py

@@ -24,6 +24,8 @@ def test_classes_replacement_when_object_detection_object_is_none() -> None:
     result = step.run(
         object_detection_predictions=None,
         classification_predictions=None,
+        fallback_class_name=None,
+        fallback_class_id=None,
     )
 
     # then
@@ -43,6 +45,8 @@ def test_classes_replacement_when_there_are_no_predictions_is_none() -> None:
     result = step.run(
         object_detection_predictions=detections,
         classification_predictions=None,
+        fallback_class_name=None,
+        fallback_class_id=None,
     )
 
     # then
@@ -100,6 +104,8 @@ def test_classes_replacement_when_replacement_to_happen_without_filtering_for_mu
     result = step.run(
         object_detection_predictions=detections,
         classification_predictions=classification_predictions,
+        fallback_class_name=None,
+        fallback_class_id=None,
     )
 
     # then
@@ -183,6 +189,8 @@ def test_classes_replacement_when_replacement_to_happen_without_filtering_for_mu
     result = step.run(
         object_detection_predictions=detections,
         classification_predictions=classification_predictions,
+        fallback_class_name=None,
+        fallback_class_id=None,
     )
 
     # then
@@ -245,6 +253,8 @@ def test_classes_replacement_when_replacement_to_happen_and_one_result_to_be_fil
     result = step.run(
         object_detection_predictions=detections,
         classification_predictions=classification_predictions,
+        fallback_class_name=None,
+        fallback_class_id=None,
     )
 
     # then
@@ -271,7 +281,7 @@ def test_classes_replacement_when_replacement_to_happen_and_one_result_to_be_fil
     ], "Expected to generate new detection id"
 
 
-def test_classes_replacement_when_empty_classification_predictions():
+def test_classes_replacement_when_empty_classification_predictions_no_fallback_class():
     # given
     step = DetectionsClassesReplacementBlockV1()
     detections = sv.Detections(
@@ -305,6 +315,8 @@ def test_classes_replacement_when_empty_classification_predictions():
     result = step.run(
         object_detection_predictions=detections,
         classification_predictions=classification_predictions,
+        fallback_class_name=None,
+        fallback_class_id=None,
     )
 
     # then
@@ -313,6 +325,72 @@ def test_classes_replacement_when_empty_classification_predictions():
     ), "Expected sv.Detections.empty(), as empty classification was passed"
 
 
+def test_classes_replacement_when_empty_classification_predictions_fallback_class_provided():
+    # given
+    step = DetectionsClassesReplacementBlockV1()
+    detections = sv.Detections(
+        xyxy=np.array(
+            [
+                [10, 20, 30, 40],
+                [11, 21, 31, 41],
+            ]
+        ),
+        class_id=np.array([7, 7]),
+        confidence=np.array([0.36, 0.91]),
+        data={
+            "class_name": np.array(["animal", "animal"]),
+            "detection_id": np.array(["zero", "one"]),
+        },
+    )
+    first_cls_prediction = ClassificationInferenceResponse(
+        image=InferenceResponseImage(width=128, height=256),
+        predictions=[
+            ClassificationPrediction(
+                **{"class": "cat", "class_id": 0, "confidence": 0.6}
+            ),
+            ClassificationPrediction(
+                **{"class": "dog", "class_id": 1, "confidence": 0.4}
+            ),
+        ],
+        top="cat",
+        confidence=0.6,
+        parent_id="some",
+    ).dict(by_alias=True, exclude_none=True)
+    first_cls_prediction["parent_id"] = "zero"
+    second_cls_prediction = ClassificationInferenceResponse(
+        image=InferenceResponseImage(width=128, height=256),
+        predictions=[],
+        top="cat",
+        confidence=0.6,
+        parent_id="some",
+    ).dict(by_alias=True, exclude_none=True)
+    second_cls_prediction["parent_id"] = "one"
+    classification_predictions = Batch(
+        content=[
+            first_cls_prediction,
+            second_cls_prediction,
+        ],
+        indices=[(0, 0), (0, 1)],
+    )
+
+    # when
+    result = step.run(
+        object_detection_predictions=detections,
+        classification_predictions=classification_predictions,
+        fallback_class_name="unknown",
+        fallback_class_id=123,
+    )
+
+    # then
+    assert (
+        len(result["predictions"]) == 2
+    ), "Expected sv.Detections.empty(), as empty classification was passed"
+    detections = result["predictions"]
+    assert detections.confidence[1] == 0, "Fallback class confidence expected to be set to 0"
+    assert detections.class_id[1] == 123, "class id expected to be set to value passed with fallback_class_id parameter"
+    assert detections.data["class_name"][1] == "unknown", "class name expected to be set to value passed with fallback_class_name parameter"
+
+
 def test_extract_leading_class_from_prediction_when_prediction_is_multi_label() -> None:
     # given
     prediction = ClassificationInferenceResponse(