From ead874105c7f12284c904358b4da7b15f12a46a2 Mon Sep 17 00:00:00 2001 From: Shiyi Zheng Date: Thu, 30 Jul 2026 12:29:32 +0800 Subject: [PATCH] feat(modelkit): add MGP-STR image-to-text support --- .../cpu/cpu/image-to-text_fp16_config.json | 53 ++++++++ .../cpu/cpu/image-to-text_fp32_config.json | 34 +++++ src/winml/modelkit/inference/pipeline.py | 19 +++ src/winml/modelkit/models/hf/__init__.py | 3 + src/winml/modelkit/models/hf/mgp_str.py | 22 ++++ src/winml/modelkit/models/winml/__init__.py | 3 + .../modelkit/models/winml/image_to_text.py | 118 ++++++++++++++++++ .../unit/eval/test_image_to_text_evaluator.py | 70 +++++++++++ tests/unit/export/test_mgp_str_onnx_config.py | 57 +++++++++ tests/unit/inference/test_pipeline.py | 29 +++++ tests/unit/loader/test_resolve_task.py | 15 +++ tests/unit/models/auto/test_auto_model.py | 47 +++++++ 12 files changed, 470 insertions(+) create mode 100644 examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json create mode 100644 examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json create mode 100644 src/winml/modelkit/models/hf/mgp_str.py create mode 100644 src/winml/modelkit/models/winml/image_to_text.py create mode 100644 tests/unit/export/test_mgp_str_onnx_config.py diff --git a/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json new file mode 100644 index 000000000..42dac7ae7 --- /dev/null +++ b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json @@ -0,0 +1,53 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [1, 3, 32, 128], + "value_range": [0, 1] + } + ], + "output_tensors": [ + {"name": "char_logits"}, + {"name": "bpe_logits"}, + {"name": "wp_logits"} + ] + }, + "optim": {}, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "image-to-text", + "model_class": "MgpstrForSceneTextRecognition", + "model_type": "mgp-str" + } +} \ No newline at end of file diff --git a/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json new file mode 100644 index 000000000..d7d7e349d --- /dev/null +++ b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json @@ -0,0 +1,34 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [1, 3, 32, 128], + "value_range": [0, 1] + } + ], + "output_tensors": [ + {"name": "char_logits"}, + {"name": "bpe_logits"}, + {"name": "wp_logits"} + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "image-to-text", + "model_class": "MgpstrForSceneTextRecognition", + "model_type": "mgp-str" + } +} \ No newline at end of file diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index e14de7af7..20ca1a629 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -38,6 +38,10 @@ "sentence-similarity": "feature-extraction", } +_SPECIALIZED_PIPELINE_REGISTRY: dict[tuple[str, str], str] = { + ("mgp-str", "image-to-text"): "MgpstrImageToTextPipeline", +} + _PIPELINE_COMPONENT_FLAGS = ( ("tokenizer", "_load_tokenizer"), ("feature_extractor", "_load_feature_extractor"), @@ -540,6 +544,21 @@ def create_pipeline( Returns: A configured task callable ready for inference. """ + pipe: Any + model_type = getattr(getattr(model, "config", None), "model_type", None) + if not isinstance(model_type, str): + model_type = "" + specialized = _SPECIALIZED_PIPELINE_REGISTRY.get((model_type, task)) + if specialized is not None: + if model_id is None: + raise ValueError(f"{specialized} requires a model ID to load its processor.") + from ..models.winml.image_to_text import MgpstrImageToTextPipeline + + pipe = MgpstrImageToTextPipeline(model, model_id) # type: ignore[arg-type] + _adapt_image_processor_size(pipe, task, model) + logger.info("Created specialized pipeline: task=%s model=%s", task, model_id) + return pipe + from transformers import pipeline hf_task = _HF_PIPELINE_TASK_MAP.get(task, task) diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index b11e7b416..244c187eb 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -55,6 +55,8 @@ from .marian import MODEL_CLASS_MAPPING as _MARIAN_CLASS_MAPPING from .marian import MarianDecoderIOConfig as _MarianDecoderIOConfig # triggers registration from .marian import MarianEncoderIOConfig as _MarianEncoderIOConfig # triggers registration +from .mgp_str import MODEL_CLASS_MAPPING as _MGP_STR_CLASS_MAPPING +from .mgp_str import MgpstrImage2TextOnnxConfig as _MgpstrImage2TextOnnxConfig from .mu2 import MODEL_CLASS_MAPPING as _MU2_CLASS_MAPPING from .mu2 import MU2_CONFIG from .mu2 import Mu2DecoderIOConfig as _Mu2DecoderIOConfig # triggers registration @@ -124,6 +126,7 @@ _BLIP_CLASS_MAPPING, _CLIP_CLASS_MAPPING, _MARIAN_CLASS_MAPPING, + _MGP_STR_CLASS_MAPPING, _MU2_CLASS_MAPPING, _QWEN_CLASS_MAPPING, _QWEN_TO_CLASS_MAPPING, diff --git a/src/winml/modelkit/models/hf/mgp_str.py b/src/winml/modelkit/models/hf/mgp_str.py new file mode 100644 index 000000000..bc880baac --- /dev/null +++ b/src/winml/modelkit/models/hf/mgp_str.py @@ -0,0 +1,22 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Configure MGP-STR for image-to-text export.""" + +from __future__ import annotations + +from optimum.exporters.onnx.model_configs import MgpstrOnnxConfig +from transformers import MgpstrForSceneTextRecognition + +from ...export import register_onnx_overwrite + + +@register_onnx_overwrite("mgp-str", "image-to-text", library_name="transformers") +class MgpstrImage2TextOnnxConfig(MgpstrOnnxConfig): # type: ignore[misc] + """Register the vendor MGP-STR export config for image-to-text.""" + + +MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = { + ("mgp-str", "image-to-text"): MgpstrForSceneTextRecognition, +} diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index 251c40dad..6adb0b076 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -59,6 +59,7 @@ # - Specific OPSET requirements # - Input remapping (non-standard tensor names) # - Custom pre/post-processing + ("mgp-str", "image-to-text"): "WinMLModelForMgpstrSceneTextRecognition", } @@ -82,6 +83,7 @@ def _import_winml_class(class_name: str) -> type[WinMLPreTrainedModel]: WinMLModelForImageSegmentation, WinMLModelForSemanticSegmentation, ) + from .image_to_text import WinMLModelForMgpstrSceneTextRecognition from .object_detection import WinMLModelForObjectDetection from .question_answering import WinMLModelForQuestionAnswering from .sequence_classification import WinMLModelForSequenceClassification @@ -91,6 +93,7 @@ def _import_winml_class(class_name: str) -> type[WinMLPreTrainedModel]: "WinMLModelForDepthEstimation": WinMLModelForDepthEstimation, "WinMLModelForFeatureExtraction": WinMLModelForFeatureExtraction, "WinMLModelForImageClassification": WinMLModelForImageClassification, + "WinMLModelForMgpstrSceneTextRecognition": WinMLModelForMgpstrSceneTextRecognition, "WinMLModelForImageSegmentation": WinMLModelForImageSegmentation, "WinMLModelForObjectDetection": WinMLModelForObjectDetection, "WinMLModelForQuestionAnswering": WinMLModelForQuestionAnswering, diff --git a/src/winml/modelkit/models/winml/image_to_text.py b/src/winml/modelkit/models/winml/image_to_text.py new file mode 100644 index 000000000..088c16499 --- /dev/null +++ b/src/winml/modelkit/models/winml/image_to_text.py @@ -0,0 +1,118 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Specialized MGP-STR image-to-text inference support.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, cast + +import numpy as np +import torch +from PIL import Image + +from .base import WinMLPreTrainedModel + + +def _convert_images_to_rgb(images: Any) -> Any: + """Normalize supported image inputs to RGB before tensor conversion.""" + from transformers.image_transforms import to_pil_image + from transformers.image_utils import load_image + + if isinstance(images, list | tuple): + return [_convert_images_to_rgb(image) for image in images] + if isinstance(images, np.ndarray | torch.Tensor) and images.ndim == 4: + return [_convert_images_to_rgb(image) for image in images] + if isinstance(images, str): + try: + return load_image(images).convert("RGB") + except ValueError: + return images + if isinstance(images, Image.Image): + return images.convert("RGB") + if isinstance(images, np.ndarray): + if images.ndim == 2: + return Image.fromarray(images).convert("RGB") + return to_pil_image(images).convert("RGB") + if isinstance(images, torch.Tensor): + if images.ndim == 2: + images = images.unsqueeze(0) + return to_pil_image(images).convert("RGB") + return images + + +class WinMLModelForMgpstrSceneTextRecognition(WinMLPreTrainedModel): + """Expose MGP-STR's three ONNX heads in the Transformers output contract.""" + + main_input_name = "pixel_values" + + def forward(self, **kwargs: Any) -> Any: + """Run the ONNX graph and return its ordered three-head logits tuple.""" + from transformers.models.mgp_str.modeling_mgp_str import MgpstrModelOutput + + try: + pixel_values = kwargs.pop("pixel_values") + except KeyError as exc: + raise ValueError("MGP-STR inference requires 'pixel_values'.") from exc + + outputs = self._run_inference(self._format_inputs(pixel_values=pixel_values)) + return MgpstrModelOutput( + logits=cast( + "Any", + ( + outputs["char_logits"], + outputs["bpe_logits"], + outputs["wp_logits"], + ), + ) + ) + + +class MgpstrImageToTextPipeline: + """Preprocess images and decode MGP-STR's character, BPE, and WordPiece heads.""" + + def __init__(self, model: WinMLPreTrainedModel, model_id: str) -> None: + from transformers import AutoProcessor + + self.model = model + self.processor = AutoProcessor.from_pretrained(model_id) + self.image_processor = self.processor.image_processor + self.tokenizer = getattr(self.processor, "char_tokenizer", None) + self._preprocess_params: dict[str, Any] = {} + + def _sanitize_parameters(self, **_kwargs: Any) -> tuple[dict, dict, dict]: + """Expose the pipeline parameter-introspection contract.""" + return {}, {}, {} + + def __call__(self, images: Any, *, prompt: str | None = None, **_kwargs: Any) -> Any: + """Recognize text in one image or a batch of images.""" + if prompt is not None: + raise ValueError("MGP-STR scene text recognition does not accept a text prompt.") + + model_inputs = self.processor( + images=_convert_images_to_rgb(images), + return_tensors="pt", + ) + outputs = self.model(**model_inputs) + if isinstance(outputs, Mapping) and all( + name in outputs for name in ("char_logits", "bpe_logits", "wp_logits") + ): + logits = tuple( + outputs[name] for name in ("char_logits", "bpe_logits", "wp_logits") + ) + elif isinstance(outputs, Mapping): + logits = outputs["logits"] + else: + logits = outputs.logits + decoded = self.processor.batch_decode(logits) + + records = [] + for index, text in enumerate(decoded["generated_text"]): + record: dict[str, Any] = {"generated_text": text} + if "scores" in decoded: + score = decoded["scores"][index] + record["score"] = float(score.item() if hasattr(score, "item") else score) + records.append(record) + return records diff --git a/tests/unit/eval/test_image_to_text_evaluator.py b/tests/unit/eval/test_image_to_text_evaluator.py index a1f0b5d03..3d44bcda4 100644 --- a/tests/unit/eval/test_image_to_text_evaluator.py +++ b/tests/unit/eval/test_image_to_text_evaluator.py @@ -7,10 +7,58 @@ from __future__ import annotations +from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch +import numpy as np +import torch +from PIL import Image + from winml.modelkit.eval.image_to_text_evaluator import WinMLImageToTextEvaluator from winml.modelkit.inference.pipeline import _HF_PIPELINE_TASK_MAP +from winml.modelkit.models.winml.image_to_text import ( + MgpstrImageToTextPipeline, + WinMLModelForMgpstrSceneTextRecognition, +) + + +class _MgpstrEvaluationModel(WinMLModelForMgpstrSceneTextRecognition): + def __init__(self) -> None: + self.config = SimpleNamespace(model_type="mgp-str") + self._io_config = { + "input_names": ["pixel_values"], + "input_shapes": [[1, 3, 32, 128]], + } + + @property + def io_config(self) -> dict: + return self._io_config + + def _run_inference(self, inputs: dict[str, np.ndarray]) -> dict[str, torch.Tensor]: + pixel_values = inputs["pixel_values"] + assert pixel_values.shape == (1, 3, 32, 128) + return { + "char_logits": torch.zeros((1, 2, 3)), + "bpe_logits": torch.zeros((1, 2, 3)), + "wp_logits": torch.zeros((1, 2, 3)), + } + + def __call__(self, **kwargs: Any) -> dict[str, torch.Tensor]: + return self._run_inference(self._format_inputs(**kwargs)) + + +class _MgpstrEvaluationProcessor: + def __call__(self, *, images: Image.Image, return_tensors: str) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + assert images.mode == "RGB" + resized = images.resize((128, 32)) + pixels = np.asarray(resized, dtype=np.float32).copy() + return {"pixel_values": torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0)} + + def batch_decode(self, logits: tuple[torch.Tensor, ...]) -> dict[str, list[Any]]: + assert len(logits) == 3 + return {"generated_text": ["TEXT"], "scores": [torch.tensor(1.0)]} def make_evaluator(columns_mapping=None): @@ -193,3 +241,25 @@ def test_empty_dataset(self): assert result["n_samples"] == 0 assert result["cer"] is None assert result["cider"] is None + + def test_mgpstr_evaluation_normalizes_grayscale_and_decodes_wrapper_output(self): + """Exercise evaluator -> MGP pipeline -> wrapper without mocks.""" + pipe = object.__new__(MgpstrImageToTextPipeline) + pipe.model = _MgpstrEvaluationModel() + pipe.processor = _MgpstrEvaluationProcessor() + + evaluator = object.__new__(WinMLImageToTextEvaluator) + evaluator._image_col = "image" + evaluator._label_col = "text" + evaluator.pipe = pipe + evaluator.data = [ + {"image": Image.new("L", (16, 8)), "text": "TEXT"}, + {"image": np.zeros((8, 16), dtype=np.uint8), "text": "TEXT"}, + {"image": torch.zeros((1, 8, 16), dtype=torch.uint8), "text": "TEXT"}, + ] + + result = evaluator.compute() + + assert result["cer"] == 0.0 + assert result["n_samples"] == 3 + assert "skipped" not in result diff --git a/tests/unit/export/test_mgp_str_onnx_config.py b/tests/unit/export/test_mgp_str_onnx_config.py new file mode 100644 index 000000000..bbcf83a91 --- /dev/null +++ b/tests/unit/export/test_mgp_str_onnx_config.py @@ -0,0 +1,57 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for MGP-STR image-to-text export registration.""" + +from __future__ import annotations + +import pytest +from optimum.exporters.tasks import TasksManager + +from winml.modelkit.export import resolve_io_specs +from winml.modelkit.models.hf.mgp_str import MODEL_CLASS_MAPPING + + +@pytest.fixture(scope="module") +def mgp_str_config(): + from transformers import MgpstrConfig + + return MgpstrConfig( + image_size=[32, 128], + patch_size=4, + num_channels=3, + hidden_size=48, + num_hidden_layers=1, + num_attention_heads=2, + mlp_ratio=2, + max_token_length=27, + num_character_labels=38, + num_bpe_labels=50257, + num_wordpiece_labels=30522, + ) + + +def test_mgp_str_config_registered() -> None: + config_constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="mgp-str", + task="image-to-text", + library_name="transformers", + ) + assert config_constructor.func.__name__ == "MgpstrImage2TextOnnxConfig" + + +def test_mgp_str_io_specs(mgp_str_config) -> None: + specs = resolve_io_specs("mgp-str", "image-to-text", mgp_str_config) + assert specs["input_names"] == ["pixel_values"] + assert specs["output_names"] == ["char_logits", "bpe_logits", "wp_logits"] + + +def test_mgp_str_model_class_mapping() -> None: + from transformers import MgpstrForSceneTextRecognition + + assert ( + MODEL_CLASS_MAPPING[("mgp-str", "image-to-text")] + is MgpstrForSceneTextRecognition + ) diff --git a/tests/unit/inference/test_pipeline.py b/tests/unit/inference/test_pipeline.py index 7755ecdb6..a6e47097f 100644 --- a/tests/unit/inference/test_pipeline.py +++ b/tests/unit/inference/test_pipeline.py @@ -164,6 +164,35 @@ def test_unknown_task_not_in_map(self) -> None: assert "image-classification" not in _HF_PIPELINE_TASK_MAP +class TestMgpstrPipeline: + def test_three_heads_are_decoded_to_generated_text(self) -> None: + import torch + + processor = MagicMock() + processor.image_processor.size = {"height": 32, "width": 128} + processor.return_value = {"pixel_values": torch.zeros((1, 3, 32, 128))} + processor.batch_decode.return_value = { + "generated_text": ["hello"], + "scores": [torch.tensor(0.9)], + } + model = MagicMock() + model.config.model_type = "mgp-str" + model.io_config = {"input_shapes": [[1, 3, 32, 128]]} + model.return_value.logits = ( + torch.zeros((1, 27, 38)), + torch.zeros((1, 27, 50257)), + torch.zeros((1, 27, 30522)), + ) + + with patch("transformers.AutoProcessor.from_pretrained", return_value=processor): + pipe = create_pipeline("image-to-text", model, "local/mgp-str") + result = pipe("image") + + processor.batch_decode.assert_called_once_with(model.return_value.logits) + assert result[0]["generated_text"] == "hello" + assert abs(result[0]["score"] - 0.9) < 1e-6 + + class TestPipelineComponentKwargs: def test_omits_components_disabled_by_pipeline(self) -> None: class ProcessorOnlyPipeline: diff --git a/tests/unit/loader/test_resolve_task.py b/tests/unit/loader/test_resolve_task.py index fffd20fe1..3d6f5fa8e 100644 --- a/tests/unit/loader/test_resolve_task.py +++ b/tests/unit/loader/test_resolve_task.py @@ -118,6 +118,21 @@ def test_unimportable_architecture_falls_back_to_hf_task_default(): assert r.source == TaskSource.HF_TASK_DEFAULT +def test_mgp_str_stale_architecture_uses_image_to_text_pipeline_tag(monkeypatch): + cfg = _cfg("mgp-str", ["MGPSTRModel"]) + cfg._name_or_path = "local-mgp-str-checkpoint" + monkeypatch.setattr( + "winml.modelkit.utils.hub_utils.get_pipeline_tag", + lambda _model_id: "image-to-text", + ) + + result = resolve_task(cfg) + + assert result.task == "image-to-text" + assert result.model_class.__name__ == "MgpstrForSceneTextRecognition" + assert result.source == TaskSource.PIPELINE_TAG + + # --- ported from the deleted test_detect_task_from_config.py ----------------- # These pin the architecture-class -> TasksManager task inference that the # (now-removed) ``_detect_task_from_config`` covered, asserted against the diff --git a/tests/unit/models/auto/test_auto_model.py b/tests/unit/models/auto/test_auto_model.py index ecc7d1fb2..3ffaf0c62 100644 --- a/tests/unit/models/auto/test_auto_model.py +++ b/tests/unit/models/auto/test_auto_model.py @@ -170,6 +170,53 @@ def test_get_winml_class_unsupported_task_returns_generic(self): model_class = get_winml_class("unknown", "unsupported-task-type") assert model_class == WinMLModelForGenericTask + def test_mgp_str_image_to_text_uses_specialized_wrapper(self): + from winml.modelkit.models import get_winml_class + from winml.modelkit.models.winml.image_to_text import ( + WinMLModelForMgpstrSceneTextRecognition, + ) + + assert ( + get_winml_class("mgp-str", "image-to-text") + is WinMLModelForMgpstrSceneTextRecognition + ) + + def test_mgp_str_wrapper_preserves_three_head_order(self): + from unittest.mock import MagicMock + + import torch + + from winml.modelkit.models.winml.image_to_text import ( + WinMLModelForMgpstrSceneTextRecognition, + ) + + model = object.__new__(WinMLModelForMgpstrSceneTextRecognition) + model._format_inputs = MagicMock(side_effect=lambda **kwargs: kwargs) + expected = { + "char_logits": torch.tensor([1.0]), + "bpe_logits": torch.tensor([2.0]), + "wp_logits": torch.tensor([3.0]), + } + model._run_inference = MagicMock(return_value=expected) + + result = model.forward(pixel_values=torch.zeros((1, 3, 32, 128))) + + assert result.logits == ( + expected["char_logits"], + expected["bpe_logits"], + expected["wp_logits"], + ) + + def test_mgp_str_wrapper_requires_pixel_values(self): + from winml.modelkit.models.winml.image_to_text import ( + WinMLModelForMgpstrSceneTextRecognition, + ) + + model = object.__new__(WinMLModelForMgpstrSceneTextRecognition) + + with pytest.raises(ValueError, match="requires 'pixel_values'"): + model.forward() + @pytest.mark.parametrize( "task,model_type,expected_class_name", [