diff --git a/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py b/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py index 0c1bab390..04c6921f7 100644 --- a/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py +++ b/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py @@ -403,7 +403,8 @@ def _pick_detections_by_parent_class( dependent_detections_to_keep.add(detection_idx) continue detections_to_keep_list = sorted(list(dependent_detections_to_keep)) - return dependent_detections[detections_to_keep_list] + filtered_dependent_detections = dependent_detections[detections_to_keep_list] + return sv.Detections.merge([parent_detections, filtered_dependent_detections]) def _is_point_within_box(point: np.ndarray, box: np.ndarray) -> bool: diff --git a/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py b/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py index a78a151a5..723ac1819 100644 --- a/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py +++ b/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py @@ -173,7 +173,9 @@ def test_picking_detections_by_parent_class_when_no_child_detections_matching() result = execute_operations(value=detections, operations=operations) # then - assert len(result) == 0 + assert len(result) == 2 + assert np.allclose(result.xyxy, np.array([[0, 0, 10, 10], [20, 20, 30, 30]])) + assert np.allclose(result.confidence, [0.3, 0.4]) def test_picking_detections_by_parent_class_when_there_are_child_detections_matching() -> ( @@ -198,9 +200,11 @@ def test_picking_detections_by_parent_class_when_there_are_child_detections_matc result = execute_operations(value=detections, operations=operations) # then - assert len(result) == 2 - assert np.allclose(result.xyxy, np.array([[20, 20, 30, 30], [40, 40, 50, 50]])) - assert np.allclose(result.confidence, [0.4, 0.5]) + assert len(result) == 3 + assert np.allclose( + result.xyxy, np.array([[0, 0, 50, 50], [20, 20, 30, 30], [40, 40, 50, 50]]) + ) + assert np.allclose(result.confidence, [0.3, 0.4, 0.5]) def test_picking_detections_by_parent_class_when_there_are_child_detections_matching_different_parents() -> ( @@ -216,20 +220,29 @@ def test_picking_detections_by_parent_class_when_there_are_child_detections_matc [40, 40, 50, 50], [100, 100, 200, 200], [150, 100, 250, 200], + [400, 400, 600, 600], ] ), - class_id=np.array([0, 1, 1, 0, 2]), - confidence=np.array([0.3, 0.4, 0.5, 0.6, 0.7]), - data={"class_name": np.array(["a", "b", "b", "a", "c"])}, + class_id=np.array([0, 1, 1, 0, 2, 3]), + confidence=np.array([0.3, 0.4, 0.5, 0.6, 0.7, 0.9]), + data={"class_name": np.array(["a", "b", "b", "a", "c", "d"])}, ) # when result = execute_operations(value=detections, operations=operations) # then - assert len(result) == 3 + assert len(result) == 5 assert np.allclose( result.xyxy, - np.array([[20, 20, 30, 30], [40, 40, 50, 50], [150, 100, 250, 200]]), + np.array( + [ + [0, 0, 50, 50], + [100, 100, 200, 200], + [20, 20, 30, 30], + [40, 40, 50, 50], + [150, 100, 250, 200], + ] + ), ) - assert np.allclose(result.confidence, [0.4, 0.5, 0.7]) + assert np.allclose(result.confidence, [0.3, 0.6, 0.4, 0.5, 0.7])