add integration tests for per-class pathways

Author: leeclemnet

inference_models/tests/integration_tests/models/test_deep_lab_v3_segmentation_predictions_onnx.py

@@ -707,3 +707,32 @@ def test_onnx_package_with_static_crop_center_crop_batch_torch_list(
     assert (
         13700 <= torch.sum(predictions[1].segmentation_map.cpu() == 3).item() <= 13900
     )
+
+
+@pytest.mark.slow
+@pytest.mark.onnx_extras
+def test_onnx_per_class_confidence_reduces_balloon_pixels(
+    balloons_deep_lab_v3_onnx_stretch_package: str,
+    balloons_image_numpy: np.ndarray,
+) -> None:
+    """Baseline (see `test_onnx_package_with_stretch_numpy` above) has class 3
+    (balloon) occupying ~16600 pixels. Setting a 0.99 per-class threshold on
+    class 3 reassigns below-threshold pixels to background, reducing the
+    class 3 pixel count."""
+    from inference_models.models.deep_lab_v3_plus.deep_lab_v3_plus_segmentation_onnx import (
+        DeepLabV3PlusForSemanticSegmentationOnnx,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = DeepLabV3PlusForSemanticSegmentationOnnx.from_pretrained(
+        model_name_or_path=balloons_deep_lab_v3_onnx_stretch_package,
+        onnx_execution_providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+    )
+    class_names = list(model.class_names)
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={class_names[3]: 0.99},
+    )
+    predictions = model(balloons_image_numpy, confidence="best")
+    balloon_pixels = torch.sum(predictions[0].segmentation_map.cpu() == 3).item()
+    assert balloon_pixels < 16600

inference_models/tests/integration_tests/models/test_deep_lab_v3_segmentation_predictions_torch.py

@@ -649,3 +649,29 @@ def test_torch_package_with_static_crop_center_crop_batch_torch_list(
     assert (
         13700 <= torch.sum(predictions[1].segmentation_map.cpu() == 3).item() <= 13900
     )
+
+
+@pytest.mark.slow
+@pytest.mark.torch_models
+def test_torch_per_class_confidence_reduces_balloon_pixels(
+    balloons_deep_lab_v3_torch_stretch_package: str,
+    balloons_image_numpy: np.ndarray,
+) -> None:
+    """Baseline (see `test_torch_package_with_stretch_numpy` above) has class 3
+    (balloon) occupying ~16600 pixels. Setting a 0.99 per-class threshold on
+    class 3 reassigns below-threshold pixels to background, reducing the
+    class 3 pixel count."""
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = DeepLabV3PlusForSemanticSegmentationTorch.from_pretrained(
+        model_name_or_path=balloons_deep_lab_v3_torch_stretch_package,
+        device=DEFAULT_DEVICE,
+    )
+    class_names = list(model.class_names)
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={class_names[3]: 0.99},
+    )
+    predictions = model(balloons_image_numpy, confidence="best")
+    balloon_pixels = torch.sum(predictions[0].segmentation_map.cpu() == 3).item()
+    assert balloon_pixels < 16600

inference_models/tests/integration_tests/models/test_dinov3_classifier_predictions_onnx.py

@@ -691,3 +691,33 @@ def test_multi_label_onnx_static_package_batch_torch_tensor_list(
         predictions[1].class_ids.cpu(),
         torch.tensor([8], dtype=torch.int64).cpu(),
     )
+
+
+@pytest.mark.slow
+@pytest.mark.onnx_extras
+def test_multi_label_onnx_per_class_confidence_blocks_specific_class(
+    dinov3_multi_label_onnx_static_package: str,
+    chess_set_image_numpy: np.ndarray,
+) -> None:
+    """Baseline (see `test_multi_label_onnx_static_package_numpy` above) returns
+    class_ids=[0..11] (12 classes). Setting a 0.99 per-class threshold on class 0
+    leaves class_ids=[1..11]."""
+    from inference_models.models.dinov3.dinov3_classification_onnx import (
+        DinoV3ForMultiLabelClassificationOnnx,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = DinoV3ForMultiLabelClassificationOnnx.from_pretrained(
+        model_name_or_path=dinov3_multi_label_onnx_static_package,
+        onnx_execution_providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+    )
+    class_names = list(model.class_names)
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={class_names[0]: 0.99},
+    )
+    predictions = model(chess_set_image_numpy, confidence="best")
+    assert torch.equal(
+        predictions[0].class_ids.cpu(),
+        torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], dtype=torch.int64),
+    )

inference_models/tests/integration_tests/models/test_dinov3_classifier_predictions_torch.py

@@ -678,3 +678,25 @@ def test_multi_label_torch_static_package_batch_torch_tensor_list(
         predictions[1].class_ids.cpu(),
         torch.tensor([8], dtype=torch.int64).cpu(),
     )
+
+
+@pytest.mark.slow
+@pytest.mark.torch_models
+def test_multi_label_torch_per_class_confidence_filters_classes(
+    dinov3_multi_label_torch_static_package: str,
+    chess_set_image_numpy: np.ndarray,
+) -> None:
+    from inference_models.models.dinov3.dinov3_classification_torch import (
+        DinoV3ForMultiLabelClassificationTorch,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = DinoV3ForMultiLabelClassificationTorch.from_pretrained(
+        model_name_or_path=dinov3_multi_label_torch_static_package,
+    )
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={name: 0.99 for name in model.class_names},
+    )
+    predictions = model(chess_set_image_numpy, confidence="best")
+    assert predictions[0].class_ids.numel() == 0

inference_models/tests/integration_tests/models/test_mediapipe_face_detector_predictions.py

@@ -175,3 +175,26 @@ def test_face_detector_predictions_for_torch_list(
         ),
         atol=3,
     )
+
+
+@pytest.mark.torch_models
+def test_per_class_confidence_filters_faces(
+    mediapipe_face_detector_package: str,
+    man_image_numpy: np.ndarray,
+) -> None:
+    from inference_models.models.mediapipe_face_detection.face_detection import (
+        MediaPipeFaceDetector,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = MediaPipeFaceDetector.from_pretrained(mediapipe_face_detector_package)
+    baseline = model(man_image_numpy, confidence=0.25)
+    baseline_count = baseline[1][0].xyxy.shape[0]
+    assert baseline_count > 0
+
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.25,
+        per_class_confidence={name: 0.99 for name in model.class_names},
+    )
+    filtered = model(man_image_numpy, confidence="best")
+    assert filtered[1][0].xyxy.shape[0] < baseline_count

inference_models/tests/integration_tests/models/test_resnet_classifier_predictions_onnx.py

@@ -798,3 +798,30 @@ def test_multi_class_onnx_static_bs_package_batch_torch_list(
         predictions.class_id.cpu(),
         torch.tensor([2, 2], dtype=torch.int64).cpu(),
     )
+
+
+@pytest.mark.slow
+@pytest.mark.onnx_extras
+def test_multi_label_onnx_per_class_confidence_blocks_specific_class(
+    flowers_multi_label_resnet_onnx_dynamic_bs_package: str,
+    flowers_image_numpy: np.ndarray,
+) -> None:
+    """Baseline (see `test_multi_label_onnx_dynamic_bs_package_numpy` above)
+    returns class_ids=[2] (class 2 at conf 0.794). Setting a 0.99 per-class
+    threshold on class 2 leaves no classes above threshold."""
+    from inference_models.models.resnet.resnet_classification_onnx import (
+        ResNetForMultiLabelClassificationOnnx,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = ResNetForMultiLabelClassificationOnnx.from_pretrained(
+        model_name_or_path=flowers_multi_label_resnet_onnx_dynamic_bs_package,
+        onnx_execution_providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+    )
+    class_names = list(model.class_names)
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={class_names[2]: 0.99},
+    )
+    predictions = model(flowers_image_numpy, confidence="best")
+    assert predictions[0].class_ids.numel() == 0

inference_models/tests/integration_tests/models/test_resnet_classifier_predictions_torch.py

@@ -363,3 +363,23 @@ def test_multi_class_torch_package_batch_torch_list(
         predictions.class_id.cpu(),
         torch.tensor([2, 2], dtype=torch.int64).cpu(),
     )
+
+
+@pytest.mark.slow
+@pytest.mark.torch_models
+def test_multi_label_torch_per_class_confidence_filters_classes(
+    flowers_multi_label_resnet_torch_package: str,
+    flowers_image_numpy: np.ndarray,
+) -> None:
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = ResNetForMultiLabelClassificationTorch.from_pretrained(
+        model_name_or_path=flowers_multi_label_resnet_torch_package,
+        device=DEFAULT_DEVICE,
+    )
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={name: 0.99 for name in model.class_names},
+    )
+    predictions = model(flowers_image_numpy, confidence="best")
+    assert predictions[0].class_ids.numel() == 0

inference_models/tests/integration_tests/models/test_resnet_classifier_predictions_trt.py

@@ -285,3 +285,26 @@ def test_multi_label_trt_package_torch_batch(
     # then
     assert abs(predictions[0].confidence[2].item() - 0.99951171875) < 1e-3
     assert abs(predictions[1].confidence[2].item() - 0.99951171875) < 1e-3
+
+
+@pytest.mark.slow
+@pytest.mark.trt_extras
+def test_multi_label_trt_per_class_confidence_filters_classes(
+    resnet_multi_label_cls_trt_package: str,
+    dog_image_numpy: np.ndarray,
+) -> None:
+    from inference_models.models.resnet.resnet_classification_trt import (
+        ResNetForMultiLabelClassificationTRT,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = ResNetForMultiLabelClassificationTRT.from_pretrained(
+        model_name_or_path=resnet_multi_label_cls_trt_package,
+        engine_host_code_allowed=True,
+    )
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.5,
+        per_class_confidence={name: 0.99 for name in model.class_names},
+    )
+    predictions = model(dog_image_numpy, confidence="best")
+    assert predictions[0].class_ids.numel() == 0

inference_models/tests/integration_tests/models/test_rfdetr_predictions_onnx.py

@@ -1225,3 +1225,31 @@ def test_onnx_package_with_nonsquare_letterbox_torch_list(
         assert torch.allclose(
             pred.xyxy.cpu(), _NONSQUARE_LETTERBOX_ONNX_EXPECTED_XYXY, atol=2
         )
+
+
+@pytest.mark.slow
+@pytest.mark.onnx_extras
+def test_onnx_per_class_confidence_blocks_specific_class(
+    coin_counting_rfdetr_nano_onnx_cs_stretch_package: str,
+    coins_counting_image_numpy: np.ndarray,
+) -> None:
+    """RFDETR OD: Group B direct-indexing + DETR no-object handling.
+    Other tests in this file (e.g. center-crop runs) return all class_id=1 for
+    coin counting. Setting a 0.99 per-class threshold on class 1 leaves no
+    detections."""
+    from inference_models.models.rfdetr.rfdetr_object_detection_onnx import (
+        RFDetrForObjectDetectionONNX,
+    )
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = RFDetrForObjectDetectionONNX.from_pretrained(
+        model_name_or_path=coin_counting_rfdetr_nano_onnx_cs_stretch_package,
+        onnx_execution_providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+    )
+    class_names = list(model.class_names)
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.3,
+        per_class_confidence={class_names[1]: 0.99},
+    )
+    predictions = model(coins_counting_image_numpy, confidence="best")
+    assert 1 not in predictions[0].class_id.cpu().tolist()

inference_models/tests/integration_tests/models/test_rfdetr_predictions_torch.py

@@ -1394,3 +1394,22 @@ def test_torch_package_with_nonsquare_letterbox_torch_list(
         assert torch.allclose(
             pred.xyxy.cpu(), _NONSQUARE_LETTERBOX_TORCH_EXPECTED_XYXY_TORCH, atol=2
         )
+
+
+@pytest.mark.slow
+@pytest.mark.torch_models
+def test_torch_per_class_confidence_filters_detections(
+    coin_counting_rfdetr_nano_torch_cs_stretch_package: str,
+    coins_counting_image_numpy: np.ndarray,
+) -> None:
+    from inference_models.weights_providers.entities import RecommendedParameters
+
+    model = RFDetrForObjectDetectionTorch.from_pretrained(
+        model_name_or_path=coin_counting_rfdetr_nano_torch_cs_stretch_package,
+    )
+    model.recommended_parameters = RecommendedParameters(
+        confidence=0.3,
+        per_class_confidence={name: 0.99 for name in model.class_names},
+    )
+    predictions = model(coins_counting_image_numpy, confidence="best")
+    assert predictions[0].class_id.numel() == 0