Add MGP-STR image-to-text support

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"
+  }
+}

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"
+  }
+}

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)

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,

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,
+}

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,

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

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