Add VietOCR (Vietnamese OCR)

src/qai_hub_models/models/vietocr/README.md

@@ -0,0 +1,84 @@
+# [VietOCR: Vietnamese text recognition with a vgg19_bn CNN backbone and transformer decoder](https://aihub.qualcomm.com/models/vietocr)
+
+VietOCR is a text recognition model for Vietnamese, covering the full precomposed tone character set. It pairs a vgg19_bn CNN backbone with a transformer sequence decoder. This contribution exports the recognition CNN backbone for on-device deployment.
+
+This is based on the implementation of VietOCR found [here](https://github.com/pbcquoc/vietocr).
+This repository contains scripts for optimized on-device export suitable to run on Qualcomm® devices. More details on model performance across various devices, can be found [here](https://aihub.qualcomm.com/models/vietocr).
+
+Qualcomm AI Hub Models uses [Qualcomm AI Hub Workbench](https://workbench.aihub.qualcomm.com) to compile, profile, and evaluate this model. [Sign up](https://myaccount.qualcomm.com/signup) to run these models on a hosted Qualcomm® device.
+
+## Setup
+### 1. Install the package
+Install the package via pip:
+```bash
+# NOTE: 3.10 <= PYTHON_VERSION < 3.14 is supported.
+pip install "qai-hub-models[vietocr]"
+```
+
+### 2. Configure Qualcomm® AI Hub Workbench
+Sign-in to [Qualcomm® AI Hub Workbench](https://workbench.aihub.qualcomm.com/) with your
+Qualcomm® ID. Once signed in navigate to `Account -> Settings -> API Token`.
+
+With this API token, you can configure your client to run models on the cloud
+hosted devices.
+```bash
+qai-hub configure --api_token API_TOKEN
+```
+Navigate to [docs](https://workbench.aihub.qualcomm.com/docs/) for more information.
+
+## Run CLI Demo
+Run the following simple CLI demo to verify the model is working end to end:
+
+```bash
+python -m qai_hub_models.models.vietocr.demo
+```
+More details on the CLI tool can be found with the `--help` option. See
+[demo.py](demo.py) for sample usage of the model including pre/post processing
+scripts. Please refer to our [general instructions on using
+models](../../../#getting-started) for more usage instructions.
+
+## Export for on-device deployment
+To run the model on Qualcomm® devices, you must export the model for use with an edge runtime such as
+TensorFlow Lite, ONNX Runtime, or Qualcomm AI Engine Direct. Use the following command to export the model:
+```bash
+python -m qai_hub_models.models.vietocr.export
+```
+Additional options are documented with the `--help` option.
+
+## Scope of this contribution
+
+Vietnamese text recognition is currently absent from the AI Hub catalog. VietOCR is a
+widely used, Apache-2.0, PyTorch-native Vietnamese recognizer, which makes it a good fit
+for the standard trace-and-compile export path. This contribution covers the recognition
+**CNN backbone** (vgg19_bn); the transformer decoder is a planned follow-up component.
+
+## Performance
+
+Measured on a Samsung Galaxy S25 Ultra (Snapdragon 8 Elite) via Qualcomm AI Hub,
+float precision, all layers on the Hexagon NPU (no CPU fallback):
+
+| Component | On-device latency | NPU layer coverage |
+|-----------|-------------------|--------------------|
+| vgg19_bn backbone | 4.48 ms | 26 / 26 (100%) |
+
+An end-to-end recognition accuracy number is not reported here because this contribution
+covers only the CNN backbone; accuracy will be reported alongside the decoder follow-up.
+
+## Engineering note
+
+The original backbone tail uses `permute(-1, 0, 1)` and `transpose(-1, -2).flatten(2)`.
+Negative permutation axes and the implied dynamic reshape do not export to a static
+on-device graph. The tail is rebuilt with equivalent static, positive-axis operations
+(`transpose(2, 3)`, `permute(2, 0, 1)`), preserving semantics while producing a fully
+static graph. See `model.py` and `test.py`.
+
+## License
+* The license for the original implementation of VietOCR can be found
+  [here](https://github.com/pbcquoc/vietocr/blob/master/LICENSE).
+
+## References
+* [Source Model Implementation](https://github.com/pbcquoc/vietocr)
+
+## Community
+* Join [our AI Hub Slack community](https://aihub.qualcomm.com/community/slack) to collaborate, post questions and learn more about on-device AI.
+* For questions or feedback please [reach out to us](mailto:ai-hub-support@qti.qualcomm.com).

src/qai_hub_models/models/vietocr/__init__.py

@@ -0,0 +1,10 @@
+# ---------------------------------------------------------------------
+# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+# ---------------------------------------------------------------------
+
+from .app import VietOCRApp as App
+from .model import MODEL_ID
+from .model import VietOCR as Model
+
+__all__ = ["MODEL_ID", "App", "Model"]

src/qai_hub_models/models/vietocr/app.py

@@ -0,0 +1,140 @@
+# ---------------------------------------------------------------------
+# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+# ---------------------------------------------------------------------
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+from qai_hub_models.models.vietocr.model import IMAGE_HEIGHT, IMAGE_WIDTH
+
+
+class VietOCRApp:
+    """
+    Light-weight "app code" for running the VietOCR recognition CNN backbone.
+
+    For a given text-line image, the app will:
+        * convert to RGB and resize to the backbone input shape
+        * normalize pixel values to [0, 1]
+        * run the backbone to produce a per-column feature sequence
+    """
+
+    def __init__(
+        self,
+        backbone: Callable[[torch.Tensor], torch.Tensor],
+        image_shape: tuple[int, int] = (IMAGE_HEIGHT, IMAGE_WIDTH),
+    ) -> None:
+        self.backbone = backbone
+        self.image_shape = image_shape
+
+    def preprocess_image(self, image: Image.Image) -> torch.Tensor:
+        """Resize and normalize a PIL image into a backbone input tensor."""
+        height, width = self.image_shape
+        resized = image.convert("RGB").resize((width, height))
+        arr = np.asarray(resized, dtype=np.float32) / 255.0
+        chw = arr.transpose(2, 0, 1)[np.newaxis]
+        return torch.from_numpy(np.ascontiguousarray(chw))
+
+    def predict(self, *args: Any, **kwargs: Any) -> np.ndarray:
+        return self.predict_features_from_image(*args, **kwargs)
+
+    def predict_features_from_image(
+        self, pixel_values_or_image: np.ndarray | torch.Tensor | Image.Image
+    ) -> np.ndarray:
+        """
+        Produce the per-column feature sequence for a text-line image.
+
+        Parameters
+        ----------
+        pixel_values_or_image
+            Input PIL image (before pre-processing), or a tensor / array
+            already shaped [batch, 3, H, W] with values in [0, 1].
+
+        Returns
+        -------
+        features : np.ndarray
+            Per-column feature sequence. Shape [W', batch, 256].
+        """
+        pixel_values = self._to_input_tensor(pixel_values_or_image)
+        return np.asarray(self.backbone(pixel_values))
+
+    def _to_input_tensor(
+        self, pixel_values_or_image: np.ndarray | torch.Tensor | Image.Image
+    ) -> torch.Tensor:
+        if isinstance(pixel_values_or_image, Image.Image):
+            return self.preprocess_image(pixel_values_or_image)
+        if isinstance(pixel_values_or_image, np.ndarray):
+            return torch.from_numpy(pixel_values_or_image)
+        return pixel_values_or_image
+
+    # ------------------------------------------------------------------
+    # End-to-end recognition.
+    #
+    # The AI Hub model exported here is the *CNN backbone* (the first of the
+    # two VietOCR recognizer components). To turn the backbone's per-column
+    # feature sequence into recognized text, those features are fed through the
+    # *second* component -- the Transformer seq2seq encoder/decoder plus the
+    # output vocabulary -- taken from the installed ``vietocr`` package. The
+    # combination demonstrates full end-to-end Vietnamese text recognition
+    # while keeping the exported on-device artifact limited to the backbone.
+    # ------------------------------------------------------------------
+    def recognize_text(
+        self,
+        image: Image.Image,
+        transformer: torch.nn.Module,
+        vocab: Any,
+        max_seq_length: int = 128,
+        sos_token: int = 1,
+        eos_token: int = 2,
+    ) -> str:
+        """
+        Recognize the Vietnamese text in a text-line image, end to end.
+
+        The backbone (this AI Hub model) produces the feature sequence; the
+        ``transformer`` (VietOCR's Transformer seq2seq head) greedily decodes
+        that sequence into token ids, which ``vocab`` maps back to characters.
+
+        Parameters
+        ----------
+        image
+            A text-line PIL image.
+        transformer
+            VietOCR's ``LanguageTransformer`` head (provides
+            ``forward_encoder`` / ``forward_decoder``).
+        vocab
+            VietOCR's ``Vocab`` (provides ``decode``).
+        max_seq_length, sos_token, eos_token
+            Greedy-decode controls matching VietOCR's defaults.
+
+        Returns
+        -------
+        text : str
+            The recognized Vietnamese text.
+        """
+        pixel_values = self._to_input_tensor(image)
+
+        with torch.no_grad():
+            # Component 1: exported CNN backbone -> per-column features.
+            src = self.backbone(pixel_values)
+            if isinstance(src, np.ndarray):
+                src = torch.from_numpy(src)
+
+            # Component 2: Transformer seq2seq decode of the features.
+            memory = transformer.forward_encoder(src)
+
+            translated = [sos_token]
+            for _ in range(max_seq_length):
+                tgt = torch.LongTensor([translated]).transpose(0, 1)
+                output, memory = transformer.forward_decoder(tgt, memory)
+                next_token = int(output[:, -1, :].argmax(dim=-1).item())
+                translated.append(next_token)
+                if next_token == eos_token:
+                    break
+
+        return vocab.decode(translated)

src/qai_hub_models/models/vietocr/code-gen.yaml

@@ -0,0 +1,2 @@
+supported_precisions:
+- float

src/qai_hub_models/models/vietocr/conftest.py

@@ -0,0 +1,38 @@
+# ---------------------------------------------------------------------
+# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+# ---------------------------------------------------------------------
+# THIS FILE WAS AUTO-GENERATED. DO NOT EDIT MANUALLY.
+
+import gc
+import warnings
+
+import pytest
+import torch.jit._trace
+
+from qai_hub_models.models.vietocr import Model
+from qai_hub_models.scorecard.utils.testing import make_cached_from_pretrained_fixture
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    # pytest is unable to figure out how to silence several PyTorch warning types from pyproject.toml settings,
+    # so we apply a manual warning filter here instead.
+    warnings.filterwarnings(action="ignore", category=torch.jit._trace.TracerWarning)
+    warnings.filterwarnings(action="ignore", category=UserWarning, module="torch.*")
+    warnings.filterwarnings(action="ignore", category=FutureWarning, module="torch.*")
+    warnings.filterwarnings(
+        action="ignore", category=DeprecationWarning, module="torch.*"
+    )
+
+
+# Instantiate the model only once for all tests.
+# Mock from_pretrained to always return the initialized model.
+# This speeds up tests and limits memory leaks.
+cached_from_pretrained = make_cached_from_pretrained_fixture(
+    Model, skip_clone_repo=True
+)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def ensure_gc() -> None:
+    gc.collect()

src/qai_hub_models/models/vietocr/demo.py

@@ -0,0 +1,97 @@
+# ---------------------------------------------------------------------
+# Copyright (c) 2025 Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+# ---------------------------------------------------------------------
+
+from __future__ import annotations
+
+from typing import Any
+
+from qai_hub_models.models.vietocr.app import VietOCRApp
+from qai_hub_models.models.vietocr.model import (
+    MODEL_ASSET_VERSION,
+    MODEL_ID,
+    VIETOCR_CONFIG_NAME,
+    VietOCR,
+)
+from qai_hub_models.utils.args import get_model_cli_parser, model_from_cli_args
+from qai_hub_models.utils.asset_loaders import CachedWebModelAsset, load_image
+
+DEFAULT_SAMPLE_IMAGE = CachedWebModelAsset.from_asset_store(
+    MODEL_ID, MODEL_ASSET_VERSION, "sample_text.jpg"
+)
+
+
+def _load_transformer_head_and_vocab() -> tuple[Any, Any, Any]:
+    """Load VietOCR's Transformer seq2seq head and output vocab.
+
+    The AI Hub model in this directory is only the CNN backbone (the first of
+    VietOCR's two recognizer components). The Transformer seq2seq head plus its
+    output vocabulary -- the second component -- are loaded here from the
+    installed ``vietocr`` package so the demo can decode the backbone features
+    into recognized text end to end.
+    """
+    import torch
+    from vietocr.model.transformerocr import VietOCR as _VietOCR
+    from vietocr.model.vocab import Vocab
+    from vietocr.tool.config import Cfg
+    from vietocr.tool.utils import download_weights
+
+    cfg = Cfg.load_config_from_name(VIETOCR_CONFIG_NAME)
+    cfg["device"] = "cpu"
+    vocab = Vocab(cfg["vocab"])
+    full_model = _VietOCR(
+        len(vocab),
+        cfg["backbone"],
+        cfg["cnn"],
+        cfg["transformer"],
+        cfg["seq_modeling"],
+    ).eval()
+    weights = download_weights(cfg["pretrain"])
+    full_model.load_state_dict(torch.load(weights, map_location="cpu"))
+    return full_model, full_model.transformer, vocab
+
+
+def main(is_test: bool = False) -> None:
+    import numpy as np
+
+    # Demo parameters
+    parser = get_model_cli_parser(VietOCR)
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=DEFAULT_SAMPLE_IMAGE,
+        help="image file path or URL",
+    )
+    args = parser.parse_args([] if is_test else None)
+
+    # Load the exported CNN backbone (AI Hub model) and wrap it in the app.
+    app = VietOCRApp(model_from_cli_args(VietOCR, args))
+
+    # Load image and run the backbone -> per-column feature sequence.
+    image = load_image(args.image)
+    features = app.predict_features_from_image(image)
+
+    # Load the second recognizer component (Transformer head + vocab) from the
+    # vietocr package and decode the backbone features into text end to end.
+    full_model, transformer, vocab = _load_transformer_head_and_vocab()
+    text = app.recognize_text(image, transformer, vocab)
+
+    # Sanity check: the exported backbone reproduces the package backbone's
+    # features on the same input, proving the on-device CNN export is correct.
+    pixel_values = app.preprocess_image(image)
+    with __import__("torch").no_grad():
+        reference = np.asarray(full_model.cnn(pixel_values))
+    backbone_features = app.predict_features_from_image(pixel_values)
+    features_match = np.allclose(backbone_features, reference, atol=1e-4)
+
+    if not is_test:
+        print(f"Backbone feature sequence shape: {features.shape}")
+        print(f"Exported backbone matches package backbone: {features_match}")
+        print(f"Recognized text: {text}")
+    else:
+        assert features_match, "Exported backbone features diverge from VietOCR's."
+
+
+if __name__ == "__main__":
+    main()