Add Rename Classes UQL Operation

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

@@ -1,4 +1,4 @@
-from typing import Any, List, Literal, Union
+from typing import Any, Dict, List, Literal, Union
 
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Annotated
@@ -450,6 +450,40 @@ class Divide(OperationDefinition):
     other: Union[int, float]
 
 
+class DetectionsRename(OperationDefinition):
+    model_config = ConfigDict(
+        json_schema_extra={
+            "description": "Renames classes in detections based on provided mapping",
+            "compound": False,
+            "input_kind": [
+                OBJECT_DETECTION_PREDICTION_KIND,
+                INSTANCE_SEGMENTATION_PREDICTION_KIND,
+                KEYPOINT_DETECTION_PREDICTION_KIND,
+            ],
+            "output_kind": [
+                OBJECT_DETECTION_PREDICTION_KIND,
+                INSTANCE_SEGMENTATION_PREDICTION_KIND,
+                KEYPOINT_DETECTION_PREDICTION_KIND,
+            ],
+        },
+    )
+    type: Literal["DetectionsRename"]
+    class_map: Dict[str, str] = Field(
+        description="Dictionary with classes replacement mapping"
+    )
+    strict: bool = Field(
+        description="Flag to decide if all class must be declared in `class_map`. When set `True` "
+        "all detections classes must be declared, otherwise error is raised.",
+        default=True,
+    )
+    new_classes_id_offset: int = Field(
+        description="When `strict` is `False`, this value determines the first "
+        "index given to re-mapped classes. This value let user create new class ids which"
+        "will not overlap with original identifiers.",
+        default=1024,
+    )
+
+
 AllOperationsType = Annotated[
     Union[
         StringToLowerCase,
@@ -469,6 +503,7 @@ class Divide(OperationDefinition):
         DetectionsFilter,
         DetectionsOffset,
         DetectionsShift,
+        DetectionsRename,
         RandomNumber,
         StringMatches,
         ExtractImageProperty,

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

@@ -28,6 +28,7 @@
     extract_detections_property,
     filter_detections,
     offset_detections,
+    rename_detections,
     select_detections,
     shift_detections,
     sort_detections,
@@ -187,6 +188,7 @@ def build_detections_filter_operation(
     "DetectionsSelection": select_detections,
     "SortDetections": sort_detections,
     "ClassificationPropertyExtract": extract_classification_property,
+    "DetectionsRename": rename_detections,
 }
 
 REGISTERED_COMPOUND_OPERATIONS_BUILDERS = {

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

@@ -15,6 +15,7 @@
 )
 from inference.core.workflows.core_steps.common.query_language.errors import (
     InvalidInputTypeError,
+    OperationError,
 )
 from inference.core.workflows.core_steps.common.query_language.operations.utils import (
     safe_stringify,
@@ -198,3 +199,82 @@ def sort_detections(
     if not ascending:
         sorted_indices = sorted_indices[::-1]
     return value[sorted_indices]
+
+
+def rename_detections(
+    detections: Any,
+    class_map: Dict[str, str],
+    strict: bool,
+    new_classes_id_offset: int,
+    **kwargs,
+) -> sv.Detections:
+    if not isinstance(detections, sv.Detections):
+        value_as_str = safe_stringify(value=detections)
+        raise InvalidInputTypeError(
+            public_message=f"Executing rename_detections(...), expected sv.Detections object as value, "
+            f"got {value_as_str} of type {type(detections)}",
+            context="step_execution | roboflow_query_language_evaluation",
+        )
+    detections_copy = deepcopy(detections)
+    original_class_names = detections_copy.data.get("class_name", []).tolist()
+    original_class_ids = detections_copy.class_id.tolist()
+    new_class_names = []
+    new_class_ids = []
+    if strict:
+        _ensure_all_classes_covered_in_new_mapping(
+            original_class_names=original_class_names,
+            class_map=class_map,
+        )
+        new_class_mapping = {
+            class_name: class_id
+            for class_id, class_name in enumerate(sorted(set(class_map.values())))
+        }
+    else:
+        new_class_mapping = _build_non_strict_class_to_id_mapping(
+            original_class_names=original_class_names,
+            original_class_ids=original_class_ids,
+            class_map=class_map,
+            new_classes_id_offset=new_classes_id_offset,
+        )
+    for class_name in original_class_names:
+        new_class_name = class_map.get(class_name, class_name)
+        new_class_id = new_class_mapping[new_class_name]
+        new_class_names.append(new_class_name)
+        new_class_ids.append(new_class_id)
+    detections_copy.data["class_name"] = np.array(new_class_names, dtype=object)
+    detections_copy.class_id = np.array(new_class_ids, dtype=int)
+    return detections_copy
+
+
+def _ensure_all_classes_covered_in_new_mapping(
+    original_class_names: List[str],
+    class_map: Dict[str, str],
+) -> None:
+    for original_class in original_class_names:
+        if original_class not in class_map:
+            raise OperationError(
+                public_message=f"Class '{original_class}' not found in class_map.",
+                context="step_execution | roboflow_query_language_evaluation",
+            )
+
+
+def _build_non_strict_class_to_id_mapping(
+    original_class_names: List[str],
+    original_class_ids: List[int],
+    class_map: Dict[str, str],
+    new_classes_id_offset: int,
+) -> Dict[str, int]:
+    original_mapping = {
+        class_name: class_id
+        for class_name, class_id in zip(original_class_names, original_class_ids)
+    }
+    new_target_classes = {
+        new_class_name
+        for new_class_name in class_map.values()
+        if new_class_name not in original_mapping
+    }
+    new_class_id = new_classes_id_offset
+    for new_target_class in sorted(new_target_classes):
+        original_mapping[new_target_class] = new_class_id
+        new_class_id += 1
+    return original_mapping

tests/workflows/integration_tests/execution/assets/image_credits.txt

@@ -1,3 +1,4 @@
 crowd.jpg: https://pixabay.com/users/wal_172619-12138562
 license_plate.jpg: https://www.pexels.com/photo/kia-niros-driving-on-the-road-11320632/
 dogs.jpg: https://www.pexels.com/photo/brown-and-white-dogs-sitting-on-field-3568134/
+multi-fruit.jpg: https://www.freepik.com/free-photo/front-close-view-organic-nutrition-source-fresh-bananas-bundle-red-apples-orange-with-stem-dark-background_17119128.htm

tests/workflows/integration_tests/execution/conftest.py

@@ -28,6 +28,11 @@ def red_image() -> np.ndarray:
     return cv2.imread(os.path.join(ASSETS_DIR, "red_image.png"))
 
 
+@pytest.fixture(scope="function")
+def fruit_image() -> np.ndarray:
+    return cv2.imread(os.path.join(ASSETS_DIR, "multi-fruit.jpg"))
+
+
 @pytest.fixture(scope="function")
 def left_scissors_right_paper() -> np.ndarray:
     return cv2.imread(

tests/workflows/integration_tests/execution/test_workflow_with_detection_renaming.py

@@ -0,0 +1,213 @@
+from typing import Dict
+
+import numpy as np
+import pytest
+
+from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS
+from inference.core.managers.base import ModelManager
+from inference.core.workflows.core_steps.common.entities import StepExecutionMode
+from inference.core.workflows.core_steps.common.query_language.errors import (
+    OperationError,
+)
+from inference.core.workflows.execution_engine.core import ExecutionEngine
+from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import (
+    add_to_workflows_gallery,
+)
+
+
+def build_class_remapping_workflow_definition(
+    class_map: Dict[str, str],
+    strict: bool,
+) -> dict:
+    return {
+        "version": "1.0",
+        "inputs": [
+            {"type": "WorkflowImage", "name": "image"},
+            {"type": "WorkflowParameter", "name": "confidence", "default_value": 0.4},
+        ],
+        "steps": [
+            {
+                "type": "ObjectDetectionModel",
+                "name": "model",
+                "image": "$inputs.image",
+                "model_id": "yolov8n-640",
+                "confidence": "$inputs.confidence",
+            },
+            {
+                "type": "DetectionsTransformation",
+                "name": "class_rename",
+                "predictions": "$steps.model.predictions",
+                "operations": [
+                    {
+                        "type": "DetectionsRename",
+                        "strict": strict,
+                        "class_map": class_map,
+                    }
+                ],
+            },
+        ],
+        "outputs": [
+            {
+                "type": "JsonField",
+                "name": "original_predictions",
+                "selector": "$steps.model.predictions",
+            },
+            {
+                "type": "JsonField",
+                "name": "renamed_predictions",
+                "selector": "$steps.class_rename.predictions",
+            },
+        ],
+    }
+
+
+@add_to_workflows_gallery(
+    category="Workflows with data transformations",
+    use_case_title="Workflow with detections class remapping",
+    use_case_description="""
+This workflow presents how to use Detections Transformation block that is going to 
+change the name of the following classes: `apple`, `banana` into `fruit`.
+
+In this example, we use non-strict mapping, causing new class `fruit` to be added to
+pool of classes - you can see that if `banana` or `apple` is detected, the
+class name changes to `fruit` and class id is 1024.
+
+You can test the execution submitting image like 
+[this](https://www.pexels.com/photo/four-trays-of-varieties-of-fruits-1300975/).
+    """,
+    workflow_definition=build_class_remapping_workflow_definition(
+        class_map={"apple": "fruit", "banana": "fruit"},
+        strict=False,
+    ),
+    workflow_name_in_app="detections-class-remapping",
+)
+def test_class_rename_workflow_with_non_strict_mapping(
+    model_manager: ModelManager,
+    fruit_image: np.ndarray,
+) -> None:
+    workflow_definition = build_class_remapping_workflow_definition(
+        class_map={"apple": "fruit", "banana": "fruit"},
+        strict=False,
+    )
+
+    workflow_init_parameters = {
+        "workflows_core.model_manager": model_manager,
+        "workflows_core.api_key": None,
+        "workflows_core.step_execution_mode": StepExecutionMode.LOCAL,
+    }
+    execution_engine = ExecutionEngine.init(
+        workflow_definition=workflow_definition,
+        init_parameters=workflow_init_parameters,
+        max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS,
+    )
+
+    # when
+    result = execution_engine.run(
+        runtime_parameters={
+            "image": fruit_image,
+            "model_id": "yolov8n-640",
+        },
+    )
+
+    # then
+    assert isinstance(result, list), "Expected result to be list"
+    assert len(result) == 1, "Single image provided - single output expected"
+
+    assert result[0]["renamed_predictions"]["class_name"].tolist() == [
+        "fruit",
+        "fruit",
+        "fruit",
+        "orange",
+        "fruit",
+    ], "Expected renamed set of classes to be the same as when test was created"
+    assert result[0]["renamed_predictions"].class_id.tolist() == [
+        1024,
+        1024,
+        1024,
+        49,
+        1024,
+    ], "Expected renamed set of class ids to be the same as when test was created"
+    assert len(result[0]["renamed_predictions"]) == len(
+        result[0]["original_predictions"]
+    ), "Expected length of predictions no to change"
+
+
+def test_class_rename_workflow_with_strict_mapping_when_all_classes_are_remapped(
+    model_manager: ModelManager,
+    fruit_image: np.ndarray,
+) -> None:
+    workflow_definition = build_class_remapping_workflow_definition(
+        class_map={"apple": "fruit", "banana": "fruit", "orange": "my-orange"},
+        strict=True,
+    )
+
+    workflow_init_parameters = {
+        "workflows_core.model_manager": model_manager,
+        "workflows_core.api_key": None,
+        "workflows_core.step_execution_mode": StepExecutionMode.LOCAL,
+    }
+    execution_engine = ExecutionEngine.init(
+        workflow_definition=workflow_definition,
+        init_parameters=workflow_init_parameters,
+        max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS,
+    )
+
+    # when
+    result = execution_engine.run(
+        runtime_parameters={
+            "image": fruit_image,
+            "model_id": "yolov8n-640",
+        },
+    )
+
+    # then
+    assert isinstance(result, list), "Expected result to be list"
+    assert len(result) == 1, "Single image provided - single output expected"
+
+    assert result[0]["renamed_predictions"]["class_name"].tolist() == [
+        "fruit",
+        "fruit",
+        "fruit",
+        "my-orange",
+        "fruit",
+    ], "Expected renamed set of classes to be the same as when test was created"
+    assert result[0]["renamed_predictions"].class_id.tolist() == [
+        0,
+        0,
+        0,
+        1,
+        0,
+    ], "Expected renamed set of class ids to be the same as when test was created"
+    assert len(result[0]["renamed_predictions"]) == len(
+        result[0]["original_predictions"]
+    ), "Expected length of predictions no to change"
+
+
+def test_class_rename_workflow_with_strict_mapping_when_not_all_classes_are_remapped(
+    model_manager: ModelManager,
+    fruit_image: np.ndarray,
+) -> None:
+    workflow_definition = build_class_remapping_workflow_definition(
+        class_map={"apple": "fruit", "banana": "fruit"},
+        strict=True,
+    )
+
+    workflow_init_parameters = {
+        "workflows_core.model_manager": model_manager,
+        "workflows_core.api_key": None,
+        "workflows_core.step_execution_mode": StepExecutionMode.LOCAL,
+    }
+    execution_engine = ExecutionEngine.init(
+        workflow_definition=workflow_definition,
+        init_parameters=workflow_init_parameters,
+        max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS,
+    )
+
+    # when
+    with pytest.raises(OperationError):
+        _ = execution_engine.run(
+            runtime_parameters={
+                "image": fruit_image,
+                "model_id": "yolov8n-640",
+            },
+        )