matrice-analytics 0.1.60__py3-none-any.whl

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/core/process.py

@@ -0,0 +1,246 @@
+"""
+Utility functions for processing model input/output.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Union
+import cv2
+import numpy as np
+
+from fast_plate_ocr.core.types import (
+    ImageColorMode,
+    ImageInterpolation,
+    PaddingColor,
+    PathLike,
+)
+
+INTERPOLATION_MAP: dict[ImageInterpolation, int] = {
+    "nearest": cv2.INTER_NEAREST,
+    "linear": cv2.INTER_LINEAR,
+    "cubic": cv2.INTER_CUBIC,
+    "area": cv2.INTER_AREA,
+    "lanczos4": cv2.INTER_LANCZOS4,
+}
+"""Mapping from interpolation method name to OpenCV constant."""
+
+
+def read_plate_image(
+    image_path: PathLike,
+    image_color_mode: ImageColorMode = "grayscale",
+) -> np.ndarray:
+    """
+    Reads an image from disk in the requested colour mode.
+
+    Args:
+        image_path: Path to the image file.
+        image_color_mode: ``"grayscale"`` for single-channel or ``"rgb"`` for three-channel
+            colour. Defaults to ``"grayscale"``.
+
+    Returns:
+        The image as a NumPy array.
+            Grayscale images have shape ``(H, W)``, RGB images have shape ``(H, W, 3)``.
+
+    Raises:
+        FileNotFoundError: If the image file does not exist.
+        ValueError: If the image cannot be decoded.
+    """
+    image_path = str(image_path)
+
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image not found: {image_path}")
+
+    if image_color_mode == "rgb":
+        raw = cv2.imread(image_path, cv2.IMREAD_COLOR)
+        if raw is None:
+            raise ValueError(f"Failed to decode image: {image_path}")
+        img = cv2.cvtColor(raw, cv2.COLOR_BGR2RGB)
+    else:
+        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+        if img is None:
+            raise ValueError(f"Failed to decode image: {image_path}")
+
+    return img
+
+
+def resize_image(
+    img: np.ndarray,
+    img_height: int,
+    img_width: int,
+    image_color_mode: ImageColorMode = "grayscale",
+    keep_aspect_ratio: bool = False,
+    interpolation_method: ImageInterpolation = "linear",
+    padding_color: PaddingColor = (114, 114, 114),
+) -> np.ndarray:
+    """
+    Resizes an in-memory image with optional aspect-ratio preservation and padding.
+
+    Args:
+        img: Input image.
+        img_height: Target image height.
+        img_width: Target image width.
+        image_color_mode: Output colour mode, ``"grayscale"`` or ``"rgb"``.
+        keep_aspect_ratio: If ``True``, maintain the original aspect ratio using letter-box
+            padding. Defaults to ``False``.
+        interpolation_method: Interpolation method used for resizing. Defaults to ``"linear"``.
+        padding_color: Padding colour (scalar for grayscale, tuple for RGB). Defaults to
+            ``(114, 114, 114)``.
+
+    Returns:
+        The resized image with shape ``(H, W, C)`` (a channel axis is added for grayscale).
+
+    Raises:
+        ValueError: If ``padding_color`` length is not 3 for RGB output.
+    """
+    # pylint: disable=too-many-locals
+
+    interpolation = INTERPOLATION_MAP[interpolation_method]
+
+    if not keep_aspect_ratio:
+        img = cv2.resize(img, (img_width, img_height), interpolation=interpolation)
+
+    else:
+        orig_h, orig_w = img.shape[:2]
+        # Scale ratio (new / old) - choose the limiting dimension
+        r = min(img_height / orig_h, img_width / orig_w)
+        # Compute the size of the resized (unpadded) image
+        new_unpad_w, new_unpad_h = round(orig_w * r), round(orig_h * r)
+        # Resize if necessary
+        if (orig_w, orig_h) != (new_unpad_w, new_unpad_h):
+            img = cv2.resize(img, (new_unpad_w, new_unpad_h), interpolation=interpolation)
+        # Padding on each side
+        dw, dh = (img_width - new_unpad_w) / 2, (img_height - new_unpad_h) / 2
+        top, bottom, left, right = (
+            round(dh - 0.1),
+            round(dh + 0.1),
+            round(dw - 0.1),
+            round(dw + 0.1),
+        )
+        border_color: PaddingColor
+        # Ensure padding colour matches channel count
+        if image_color_mode == "grayscale":
+            if isinstance(padding_color, tuple):
+                border_color = int(padding_color[0])
+            else:
+                border_color = int(padding_color)
+        elif image_color_mode == "rgb":
+            if isinstance(padding_color, tuple):
+                if len(padding_color) != 3:
+                    raise ValueError("padding_color must be length-3 for RGB images")
+                border_color = tuple(int(c) for c in padding_color)  # type: ignore[assignment]
+            else:
+                border_color = (int(padding_color),) * 3
+        img = cv2.copyMakeBorder(
+            img,
+            top,
+            bottom,
+            left,
+            right,
+            borderType=cv2.BORDER_CONSTANT,
+            value=border_color,  # type: ignore[arg-type]
+        )
+    # Add channel axis for gray so output is HxWxC
+    if image_color_mode == "grayscale" and img.ndim == 2:
+        img = np.expand_dims(img, axis=-1)
+
+    return img
+
+
+def read_and_resize_plate_image(
+    image_path: PathLike,
+    img_height: int,
+    img_width: int,
+    image_color_mode: ImageColorMode = "grayscale",
+    keep_aspect_ratio: bool = False,
+    interpolation_method: ImageInterpolation = "linear",
+    padding_color: PaddingColor = (114, 114, 114),
+) -> np.ndarray:
+    """
+    Reads an image from disk and resizes it for model input.
+
+    Args:
+        image_path: Path to the image.
+        img_height: Desired output height.
+        img_width: Desired output width.
+        image_color_mode: ``"grayscale"`` or ``"rgb"``. Defaults to ``"grayscale"``.
+        keep_aspect_ratio: Whether to preserve aspect ratio via letter-boxing. Defaults to
+            ``False``.
+        interpolation_method: Interpolation method to use. Defaults to ``"linear"``.
+        padding_color: Colour used for padding when aspect ratio is preserved. Defaults to
+            ``(114, 114, 114)``.
+
+    Returns:
+        The resized (and possibly padded) image with shape ``(H, W, C)``.
+    """
+    img = read_plate_image(image_path, image_color_mode=image_color_mode)
+    return resize_image(
+        img,
+        img_height,
+        img_width,
+        image_color_mode=image_color_mode,
+        keep_aspect_ratio=keep_aspect_ratio,
+        interpolation_method=interpolation_method,
+        padding_color=padding_color,
+    )
+
+
+def preprocess_image(images: np.ndarray) -> np.ndarray:
+    """
+    Converts image data to the format expected by the model.
+
+    The model itself handles pixel-value normalisation, so this function only ensures the
+    batch-dimension and dtype are correct.
+
+    Args:
+        images: Image or batch of images with shape ``(H, W, C)`` or ``(N, H, W, C)``.
+
+    Returns:
+        A NumPy array with shape ``(N, H, W, C)`` and dtype ``uint8``.
+
+    Raises:
+        ValueError: If the input does not have 3 or 4 dimensions.
+    """
+    # single sample (H, W, C)
+    if images.ndim == 3:
+        images = np.expand_dims(images, axis=0)
+
+    if images.ndim != 4:
+        raise ValueError("Expected input of shape (N, H, W, C).")
+
+    return images.astype(np.uint8)
+
+
+def postprocess_output(
+    model_output: np.ndarray,
+    max_plate_slots: int,
+    model_alphabet: str,
+    return_confidence: bool = False,
+) -> tuple[list[str], Union[np.ndarray, list[str]]]:
+    """
+    Decodes model predictions into licence-plate strings.
+
+    Args:
+        model_output: Raw output tensor from the model.
+        max_plate_slots: Maximum number of character positions.
+        model_alphabet: Alphabet used by the model.
+        return_confidence: If ``True``, also return per-character confidence scores.
+            Defaults to ``False``.
+
+    Returns:
+        If ``return_confidence`` is ``False``: a list of decoded plate strings.
+            If ``True``: a two-tuple ``(plates, probs)`` where
+
+            * ``plates`` is the list of decoded strings, and
+            * ``probs`` is an array of shape ``(N, max_plate_slots)`` with the corresponding
+              confidence scores.
+    """
+    predictions = model_output.reshape((-1, max_plate_slots, len(model_alphabet)))
+    prediction_indices = np.argmax(predictions, axis=-1)
+    alphabet_array = np.array(list(model_alphabet))
+    plate_chars = alphabet_array[prediction_indices]
+    plates: list[str] = np.apply_along_axis("".join, 1, plate_chars).tolist()
+    if return_confidence:
+        probs = np.max(predictions, axis=-1)
+        return plates, probs
+    return plates

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/core/types.py

@@ -0,0 +1,60 @@
+"""
+Common custom types used across the lib.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Sequence
+from typing import Literal, Tuple, Union
+
+import numpy as np
+from numpy import typing as npt
+
+ImageInterpolation = Literal["nearest", "linear", "cubic", "area", "lanczos4"]
+"""Interpolation method used for resizing the input image."""
+ImageColorMode = Literal["grayscale", "rgb"]
+"""
+Input image color mode. Use ``grayscale`` for single-channel input or ``rgb`` for 3-channel input.
+"""
+PaddingColor = Union[Tuple[int, int, int], int]
+"""Padding colour for letterboxing (only used when keeping image aspect ratio)."""
+PathLike = Union[str, os.PathLike]
+"""Path-like objects."""
+ImgLike = Union[PathLike, npt.NDArray[np.uint8]]
+"""Image-like objects, including paths to image files and NumPy arrays of images."""
+BatchOrImgLike = Union[ImgLike, Sequence[ImgLike]]
+"""
+Image-like objects, including paths to image files and NumPy arrays of images, or a batch of images.
+"""
+BatchArray = npt.NDArray[np.uint8]
+"""Numpy array of images, representing a batch of images."""
+TensorDataFormat = Literal["channels_last", "channels_first"]
+"""
+Data format of the input tensor. It can be either ``channels_last`` or ``channels_first``.
+``channels_last`` corresponds to inputs with shape ``(batch, height, width, channels)``, while
+``channels_first`` corresponds to inputs with shape ``(batch, channels, height, width)``.
+"""
+KerasDtypes = Literal[
+    "float16",
+    "float32",
+    "float64",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "bfloat16",
+    "bool",
+    "string",
+    "float8_e4m3fn",
+    "float8_e5m2",
+    "complex64",
+    "complex128",
+]
+"""
+Keras data types supported by the library.
+"""

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/core/utils.py

@@ -0,0 +1,87 @@
+"""
+Common utilities used across the package.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import IO, Any, Optional, Union
+
+
+@contextmanager
+def log_time_taken(process_name: str) -> Iterator[None]:
+    """
+    A concise context manager to time code snippets and log the result.
+
+    Usage:
+        ```python
+        with log_time_taken("process_name"):
+            # Code snippet to be timed
+        ```
+
+    Args:
+        process_name: Name of the process being timed.
+    """
+    time_start: float = time.perf_counter()
+    try:
+        yield
+    finally:
+        time_end: float = time.perf_counter()
+        time_elapsed: float = time_end - time_start
+        logger = logging.getLogger(__name__)
+        logger.info("Computation time of '%s' = %.3fms", process_name, 1_000 * time_elapsed)
+
+
+@contextmanager
+def measure_time() -> Iterator[Callable[[], float]]:
+    """
+    A context manager for measuring execution time (in milliseconds) within its code block.
+
+    Usage:
+        ```python
+        with measure_time() as timer:
+            # Code snippet to be timed
+        print(f"Code took: {timer()} ms")
+        ```
+
+    Returns:
+        A function that returns the elapsed time in milliseconds.
+    """
+    start_time = end_time = time.perf_counter()
+    yield lambda: (end_time - start_time) * 1_000
+    end_time = time.perf_counter()
+
+
+@contextmanager
+def safe_write(
+    file: Union[str, os.PathLike[str]],
+    mode: str = "wb",
+    encoding: Optional[str] = None,
+    **kwargs: Any,
+) -> Iterator[IO]:
+    """
+    Context manager for safe file writing.
+
+    Opens the specified file for writing and yields a file object.
+    If an exception occurs during writing, the file is removed before raising the exception.
+
+    Args:
+        file: Path to the file to write.
+        mode: File open mode (e.g. ``"wb"``, ``"w"``, etc.). Defaults to ``"wb"``.
+        encoding: Encoding to use (for text modes). Ignored in binary mode.
+        **kwargs: Additional arguments passed to ``open()``.
+
+    Returns:
+        A writable file object.
+    """
+    try:
+        with open(file, mode, encoding=encoding, **kwargs) as f:
+            yield f
+    except Exception as e:
+        Path(file).unlink(missing_ok=True)
+        raise e

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/inference/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+"""Inference utilities for Fast Plate OCR (Python 3.8 compatible)."""

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/inference/config.py

@@ -0,0 +1,82 @@
+"""
+Model config reading/parsing for doing inference.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Union
+import yaml
+
+from fast_plate_ocr.core.types import ImageColorMode, ImageInterpolation, PathLike
+
+# pylint: disable=duplicate-code
+
+
+@dataclass(frozen=True)
+class PlateOCRConfig:  # pylint: disable=too-many-instance-attributes
+    """
+    Plate OCR Config used for inference.
+
+    This dataclass is used to read and parse the config file used for training the OCR model.
+    We prefer to keep the inference package with minimal dependencies and avoid using Pydantic here.
+    """
+
+    max_plate_slots: int
+    """
+    Max number of plate slots supported. This represents the number of model classification heads.
+    """
+    alphabet: str
+    """
+    All the possible character set for the model output.
+    """
+    pad_char: str
+    """
+    Padding character for plates which length is smaller than MAX_PLATE_SLOTS.
+    """
+    img_height: int
+    """
+    Image height which is fed to the model.
+    """
+    img_width: int
+    """
+    Image width which is fed to the model.
+    """
+    keep_aspect_ratio: bool = False
+    """
+    Keep aspect ratio of the input image.
+    """
+    interpolation: ImageInterpolation = "linear"
+    """
+    Interpolation method used for resizing the input image.
+    """
+    image_color_mode: ImageColorMode = "grayscale"
+    """
+    Input image color mode. Use 'grayscale' for single-channel input or 'rgb' for 3-channel input.
+    """
+    padding_color: Union[tuple[int, int, int], int] = (114, 114, 114)
+    """
+    Padding color used when keep_aspect_ratio is True. For grayscale images, this should be a single
+    integer and for RGB images, this must be a tuple of three integers.
+    """
+
+    @property
+    def vocabulary_size(self) -> int:
+        return len(self.alphabet)
+
+    @property
+    def pad_idx(self) -> int:
+        return self.alphabet.index(self.pad_char)
+
+    @property
+    def num_channels(self) -> int:
+        return 3 if self.image_color_mode == "rgb" else 1
+
+    @classmethod
+    def from_yaml(cls, path: PathLike) -> "PlateOCRConfig":
+        """
+        Read and parse a yaml containing the Plate OCR config.
+        """
+        with open(path, encoding="utf-8") as f_in:
+            data = yaml.safe_load(f_in)
+        return cls(**data)

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/inference/hub.py

@@ -0,0 +1,141 @@
+"""
+Utilities function used for doing inference with the OCR models.
+"""
+
+from __future__ import annotations
+
+import logging
+import pathlib
+import shutil
+import urllib.request
+from http import HTTPStatus
+from typing import Literal, Tuple, Dict, Optional
+
+from tqdm.asyncio import tqdm
+
+from fast_plate_ocr.core.utils import safe_write
+
+BASE_URL: str = "https://github.com/ankandrew/cnn-ocr-lp/releases/download"
+OcrModel = Literal[
+    "cct-s-v1-global-model",
+    "cct-xs-v1-global-model",
+    "cct-s-relu-v1-global-model",
+    "cct-xs-relu-v1-global-model",
+    "argentinian-plates-cnn-model",
+    "argentinian-plates-cnn-synth-model",
+    "european-plates-mobile-vit-v2-model",
+    "global-plates-mobile-vit-v2-model",
+]
+"""Available OCR models for doing inference."""
+
+
+AVAILABLE_ONNX_MODELS: Dict[OcrModel, Tuple[str, str]] = {
+    "cct-s-v1-global-model": (
+        f"{BASE_URL}/arg-plates/cct_s_v1_global.onnx",
+        f"{BASE_URL}/arg-plates/cct_s_v1_global_plate_config.yaml",
+    ),
+    "cct-xs-v1-global-model": (
+        f"{BASE_URL}/arg-plates/cct_xs_v1_global.onnx",
+        f"{BASE_URL}/arg-plates/cct_xs_v1_global_plate_config.yaml",
+    ),
+    "argentinian-plates-cnn-model": (
+        f"{BASE_URL}/arg-plates/arg_cnn_ocr.onnx",
+        f"{BASE_URL}/arg-plates/arg_cnn_ocr_config.yaml",
+    ),
+    "argentinian-plates-cnn-synth-model": (
+        f"{BASE_URL}/arg-plates/arg_cnn_ocr_synth.onnx",
+        f"{BASE_URL}/arg-plates/arg_cnn_ocr_config.yaml",
+    ),
+    "european-plates-mobile-vit-v2-model": (
+        f"{BASE_URL}/arg-plates/european_mobile_vit_v2_ocr.onnx",
+        f"{BASE_URL}/arg-plates/european_mobile_vit_v2_ocr_config.yaml",
+    ),
+    "global-plates-mobile-vit-v2-model": (
+        f"{BASE_URL}/arg-plates/global_mobile_vit_v2_ocr.onnx",
+        f"{BASE_URL}/arg-plates/global_mobile_vit_v2_ocr_config.yaml",
+    ),
+    "cct-s-relu-v1-global-model": (
+        f"{BASE_URL}/arg-plates/cct_s_relu_v1_global.onnx",
+        f"{BASE_URL}/arg-plates/cct_s_relu_v1_global_plate_config.yaml",
+    ),
+    "cct-xs-relu-v1-global-model": (
+        f"{BASE_URL}/arg-plates/cct_xs_relu_v1_global.onnx",
+        f"{BASE_URL}/arg-plates/cct_xs_relu_v1_global_plate_config.yaml",
+    ),
+}
+"""Dictionary of available OCR models and their URLs."""
+
+MODEL_CACHE_DIR: pathlib.Path = pathlib.Path.home() / ".cache" / "fast-plate-ocr"
+"""Default location where models will be stored."""
+
+
+def _download_with_progress(url: str, filename: pathlib.Path) -> None:
+    """
+    Download utility function with progress bar.
+
+    :param url: URL of the model to download.
+    :param filename: Where to save the OCR model.
+    """
+    with urllib.request.urlopen(url) as response, safe_write(filename, mode="wb") as out_file:
+        if response.getcode() != HTTPStatus.OK:
+            raise ValueError(f"Failed to download file from {url}. Status code: {response.status}")
+
+        file_size = int(response.headers.get("Content-Length", 0))
+        desc = f"Downloading {filename.name}"
+
+        with tqdm.wrapattr(out_file, "write", total=file_size, desc=desc) as f_out:
+            shutil.copyfileobj(response, f_out)
+
+
+def download_model(
+    model_name: OcrModel,
+    save_directory: Optional[pathlib.Path] = None,
+    force_download: bool = False,
+) -> tuple[pathlib.Path, pathlib.Path]:
+    """
+    Download an OCR model and the config to a given directory.
+
+    Args:
+        model_name: Which model to download.
+        save_directory: Directory to save the OCR model. It should point to a folder.
+            If not supplied, this will point to '~/.cache/<model_name>'.
+        force_download: Force and download the model if it already exists in
+            `save_directory`.
+
+    Returns:
+        A tuple consisting of (model_downloaded_path, config_downloaded_path).
+    """
+    if model_name not in AVAILABLE_ONNX_MODELS:
+        available_models = ", ".join(AVAILABLE_ONNX_MODELS.keys())
+        raise ValueError(f"Unknown model {model_name}. Use one of [{available_models}]")
+
+    if save_directory is None:
+        save_directory = MODEL_CACHE_DIR / model_name
+    elif save_directory.is_file():
+        raise ValueError(f"Expected a directory, but got {save_directory}")
+
+    save_directory.mkdir(parents=True, exist_ok=True)
+
+    model_url, plate_config_url = AVAILABLE_ONNX_MODELS[model_name]
+    model_filename = save_directory / model_url.split("/")[-1]
+    plate_config_filename = save_directory / plate_config_url.split("/")[-1]
+
+    if not force_download and model_filename.is_file() and plate_config_filename.is_file():
+        logging.info(
+            "Skipping download of '%s' model, already exists at %s",
+            model_name,
+            save_directory,
+        )
+        return model_filename, plate_config_filename
+
+    # Download the model if not present or if we want to force the download
+    if force_download or not model_filename.is_file():
+        logging.info("Downloading model to %s", model_filename)
+        _download_with_progress(url=model_url, filename=model_filename)
+
+    # Same for the config
+    if force_download or not plate_config_filename.is_file():
+        logging.info("Downloading config to %s", plate_config_filename)
+        _download_with_progress(url=plate_config_url, filename=plate_config_filename)
+
+    return model_filename, plate_config_filename