diff --git a/src/winml/modelkit/eval/keypoint_detection_evaluator.py b/src/winml/modelkit/eval/keypoint_detection_evaluator.py index a5a701f33..86f151418 100644 --- a/src/winml/modelkit/eval/keypoint_detection_evaluator.py +++ b/src/winml/modelkit/eval/keypoint_detection_evaluator.py @@ -86,12 +86,7 @@ def prepare_pipeline(self) -> Any: The processor size is forced to the exported ONNX input shape so the preprocessed crops match the static model input. """ - from transformers import AutoImageProcessor - - processor = AutoImageProcessor.from_pretrained( - self.config.model_id, - trust_remote_code=self.config.trust_remote_code, - ) + processor = self._load_image_processor() io_config = getattr(self.model, "io_config", None) or {} input_shapes = io_config.get("input_shapes", []) @@ -101,6 +96,52 @@ def prepare_pipeline(self) -> Any: return processor + def _load_image_processor(self) -> Any: + """Resolve image processor with a narrow vitpose-family metadata fallback.""" + from transformers import AutoImageProcessor, VitPoseImageProcessor + + model_id = self.config.model_id + if model_id is None: + raise ValueError("model_id is required for keypoint-detection evaluator") + + try: + return AutoImageProcessor.from_pretrained( + model_id, + trust_remote_code=self.config.trust_remote_code, + ) + except ValueError as e: + if not self._should_fallback_to_vitpose_image_processor(e): + raise + + logger.info( + "AutoImageProcessor resolution failed for vitpose metadata; " + "retrying with VitPoseImageProcessor" + ) + return VitPoseImageProcessor.from_pretrained( + model_id, + trust_remote_code=self.config.trust_remote_code, + ) + + def _should_fallback_to_vitpose_image_processor(self, error: ValueError) -> bool: + """Fallback only for known unrecognized-processor failures in vitpose family.""" + return "Unrecognized image processor" in str(error) and self._is_vitpose_family() + + def _is_vitpose_family(self) -> bool: + """Check model metadata for vitpose family without model-id branching.""" + hf_config = getattr(self.model, "config", None) + + model_type = getattr(hf_config, "model_type", None) + if isinstance(model_type, str) and model_type.lower() == "vitpose": + return True + + architectures = getattr(hf_config, "architectures", None) + if isinstance(architectures, list | tuple): + for architecture in architectures: + if isinstance(architecture, str) and "vitpose" in architecture.lower(): + return True + + return False + def compute(self) -> dict[str, Any]: """Run keypoint evaluation over all samples and return COCO AP/AR.""" from tqdm import tqdm diff --git a/tests/unit/eval/test_keypoint_detection_evaluator.py b/tests/unit/eval/test_keypoint_detection_evaluator.py index 543bcdb53..27cf6ff14 100644 --- a/tests/unit/eval/test_keypoint_detection_evaluator.py +++ b/tests/unit/eval/test_keypoint_detection_evaluator.py @@ -12,6 +12,7 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -66,6 +67,121 @@ def test_processor_forwards_trust_remote_code(self): trust_remote_code=True, ) + def test_vitpose_alias_mismatch_uses_gated_fallback(self): + ev = _make_evaluator() + ev.config = WinMLEvaluationConfig( + model_id="nielsr/vitpose-base-simple", + task="keypoint-detection", + trust_remote_code=True, + ) + ev.model = MagicMock( + io_config={}, + config=SimpleNamespace( + model_type="vitpose", + architectures=["ViTPoseForPoseEstimation"], + ), + ) + fallback_processor = MagicMock() + + with ( + patch( + "transformers.AutoImageProcessor.from_pretrained", + side_effect=ValueError( + "Unrecognized image processor in nielsr/vitpose-base-simple" + ), + ) as load_auto, + patch( + "transformers.VitPoseImageProcessor.from_pretrained", + return_value=fallback_processor, + ) as load_vitpose, + ): + assert ev.prepare_pipeline() is fallback_processor + + load_auto.assert_called_once_with( + "nielsr/vitpose-base-simple", + trust_remote_code=True, + ) + load_vitpose.assert_called_once_with( + "nielsr/vitpose-base-simple", + trust_remote_code=True, + ) + + def test_non_vitpose_success_keeps_strict_auto_path(self): + ev = _make_evaluator() + ev.config = WinMLEvaluationConfig( + model_id="microsoft/resnet-50", + task="keypoint-detection", + trust_remote_code=False, + ) + ev.model = MagicMock( + io_config={}, + config=SimpleNamespace( + model_type="resnet", architectures=["ResNetForImageClassification"] + ), + ) + processor = MagicMock() + + with ( + patch( + "transformers.AutoImageProcessor.from_pretrained", + return_value=processor, + ) as load_auto, + patch("transformers.VitPoseImageProcessor.from_pretrained") as load_vitpose, + ): + assert ev.prepare_pipeline() is processor + + load_auto.assert_called_once_with( + "microsoft/resnet-50", + trust_remote_code=False, + ) + load_vitpose.assert_not_called() + + def test_unrelated_value_error_is_not_swallowed(self): + ev = _make_evaluator() + ev.config = WinMLEvaluationConfig( + model_id="nielsr/vitpose-base-simple", + task="keypoint-detection", + trust_remote_code=False, + ) + ev.model = MagicMock( + io_config={}, + config=SimpleNamespace( + model_type="vitpose", + architectures=["ViTPoseForPoseEstimation"], + ), + ) + + with ( + patch( + "transformers.AutoImageProcessor.from_pretrained", + side_effect=ValueError("network timeout while loading processor"), + ) as load_auto, + patch("transformers.VitPoseImageProcessor.from_pretrained") as load_vitpose, + pytest.raises(ValueError, match="network timeout"), + ): + ev.prepare_pipeline() + + load_auto.assert_called_once_with( + "nielsr/vitpose-base-simple", + trust_remote_code=False, + ) + load_vitpose.assert_not_called() + + def test_missing_model_id_raises_before_processor_load(self): + ev = _make_evaluator() + ev.config = WinMLEvaluationConfig(task="keypoint-detection") + ev.model = MagicMock(io_config={}) + + with ( + patch("transformers.AutoImageProcessor.from_pretrained") as load_auto, + patch("transformers.VitPoseImageProcessor.from_pretrained") as load_vitpose, + pytest.raises(ValueError, match="model_id is required"), + ): + ev.prepare_pipeline() + + load_auto.assert_not_called() + load_vitpose.assert_not_called() + class TestPredictionFlattening: def test_flatten_interleaves_xy_and_score(self):