matrice-analytics 0.1.60__py3-none-any.whl

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

@@ -0,0 +1,323 @@
+"""
+ONNX inference module.
+"""
+
+from __future__ import annotations
+
+import logging
+import pathlib
+from collections.abc import Sequence
+from typing import Literal, Union, Optional
+
+import numpy as np
+import numpy.typing as npt
+
+try:
+    import onnxruntime as ort
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "ONNX Runtime is not installed. Run: "
+        "pip install 'fast-plate-ocr[onnx]' (or [onnx-gpu], etc.)"
+    ) from e
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from fast_plate_ocr.core.process import (
+    postprocess_output,
+    preprocess_image,
+    read_and_resize_plate_image,
+    resize_image,
+)
+from fast_plate_ocr.core.types import BatchArray, BatchOrImgLike, ImgLike, PathLike
+from fast_plate_ocr.core.utils import measure_time
+from fast_plate_ocr.inference import hub
+from fast_plate_ocr.inference.config import PlateOCRConfig
+from fast_plate_ocr.inference.hub import OcrModel
+
+
+def _frame_from(item: ImgLike, cfg: PlateOCRConfig) -> BatchArray:
+    """
+    Converts a single image-like input into a normalized (H, W, C) NumPy array ready for model
+    inference. It handles both file paths and in-memory images. If input is a file path, the image
+    is read and resized using the configuration provided. If it's a NumPy array, it is validated and
+    resized accordingly.
+    """
+    # If it's a path, read and resize
+    if isinstance(item, (str, pathlib.PurePath)):
+        return read_and_resize_plate_image(
+            item,
+            img_height=cfg.img_height,
+            img_width=cfg.img_width,
+            image_color_mode=cfg.image_color_mode,
+            keep_aspect_ratio=cfg.keep_aspect_ratio,
+            interpolation_method=cfg.interpolation,
+            padding_color=cfg.padding_color,
+        )
+
+    # Otherwise it must be a numpy array
+    if not isinstance(item, np.ndarray):
+        raise TypeError(f"Unsupported element type: {type(item)}")
+
+    # If it has (N, H, W, C) shape we assume it's ready for inference
+    if item.ndim == 4:
+        return item
+
+    # If it's a single frame resize accordingly
+    return resize_image(
+        item,
+        cfg.img_height,
+        cfg.img_width,
+        image_color_mode=cfg.image_color_mode,
+        keep_aspect_ratio=cfg.keep_aspect_ratio,
+        interpolation_method=cfg.interpolation,
+        padding_color=cfg.padding_color,
+    )
+
+
+def _load_image_from_source(source: BatchOrImgLike, cfg: PlateOCRConfig) -> BatchArray:
+    """
+    Converts an image input or batch of inputs into a 4-D NumPy array (N, H, W, C).
+
+    This utility supports a wide range of input formats, including single images or batches, file
+    paths or NumPy arrays. It ensures the result is always a model-ready batch.
+
+    Supported input formats:
+    - Single path (`str` or `PathLike`) -> image is read and resized
+    - List or tuple of paths -> each image is read and resized
+    - Single 2D or 3D NumPy array -> resized and wrapped in a batch
+    - List or tuple of NumPy arrays -> each image is resized and batched
+    - Single 4D NumPy array with shape (N, H, W, C) -> returned as is
+
+    Args:
+        source: A single image or batch of images in path or NumPy array format.
+        cfg: The configuration object that defines image preprocessing parameters.
+
+    Returns:
+        A 4D NumPy array of shape (N, H, W, C), dtype uint8, ready for model inference.
+    """
+    if isinstance(source, np.ndarray) and source.ndim == 4:
+        return source
+
+    items: Sequence[ImgLike] = (
+        source
+        if isinstance(source, Sequence)
+        and not isinstance(source, (str, pathlib.PurePath, np.ndarray))
+        else [source]
+    )
+
+    frames: list[BatchArray] = [
+        frame
+        for item in items
+        for frame in (
+            _frame_from(item, cfg)  # type: ignore[attr-defined]
+            if isinstance(item, np.ndarray) and item.ndim == 4
+            else [_frame_from(item, cfg)]
+        )
+    ]
+
+    return np.stack(frames, axis=0, dtype=np.uint8)
+
+
+class LicensePlateRecognizer:
+    """
+    ONNX inference class for performing license plates OCR.
+    """
+
+    def __init__(
+        self,
+        hub_ocr_model: Optional[OcrModel] = None,
+        device: Literal["cuda", "cpu", "auto"] = "auto",
+        providers: Optional[Sequence[Union[str, tuple[str, dict]]]] = None,
+        sess_options: Optional[ort.SessionOptions] = None,
+        onnx_model_path: Optional[PathLike] = None,
+        plate_config_path: Optional[PathLike] = None,
+        force_download: bool = False,
+    ) -> None:
+        """
+        Initializes the `LicensePlateRecognizer` with the specified OCR model and inference device.
+
+        The current OCR models available from the HUB are:
+
+        - `cct-s-v1-global-model`: OCR model trained with **global** plates data. Based on Compact
+            Convolutional Transformer (CCT) architecture. This is the **S** variant.
+        - `cct-xs-v1-global-model`: OCR model trained with **global** plates data. Based on Compact
+            Convolutional Transformer (CCT) architecture. This is the **XS** variant.
+        - `argentinian-plates-cnn-model`: OCR for **Argentinian** license plates. Uses fully conv
+            architecture.
+        - `argentinian-plates-cnn-synth-model`: OCR for **Argentinian** license plates trained with
+            synthetic and real data. Uses fully conv architecture.
+        - `european-plates-mobile-vit-v2-model`: OCR for **European** license plates. Uses
+            MobileVIT-2 for the backbone.
+        - `global-plates-mobile-vit-v2-model`: OCR for **global** license plates (+65 countries).
+            Uses MobileVIT-2 for the backbone.
+
+        Args:
+            hub_ocr_model: Name of the OCR model to use from the HUB.
+            device: Device type for inference. Should be one of ('cpu', 'cuda', 'auto'). If
+                'auto' mode, the device will be deduced from
+                `onnxruntime.get_available_providers()`.
+            providers: Optional sequence of providers in order of decreasing precedence. If not
+                specified, all available providers are used based on the device argument.
+            sess_options: Advanced session options for ONNX Runtime.
+            onnx_model_path: Path to ONNX model file to use (In case you want to use a custom one).
+            plate_config_path: Path to config file to use (In case you want to use a custom one).
+            force_download: Force and download the model, even if it already exists.
+        Returns:
+            None.
+        """
+        self.logger = logging.getLogger(__name__)
+
+        if providers is not None:
+            self.providers = providers
+            self.logger.info("Using custom providers: %s", providers)
+        else:
+            if device == "cuda":
+                self.providers = ["CUDAExecutionProvider"]
+            elif device == "cpu":
+                self.providers = ["CPUExecutionProvider"]
+            elif device == "auto":
+                self.providers = ort.get_available_providers()
+            else:
+                raise ValueError(
+                    f"Device should be one of ('cpu', 'cuda', 'auto'). Got '{device}'."
+                )
+
+            self.logger.info("Using device '%s' with providers: %s", device, self.providers)
+
+        if onnx_model_path and plate_config_path:
+            onnx_model_path = pathlib.Path(onnx_model_path)
+            plate_config_path = pathlib.Path(plate_config_path)
+            if not onnx_model_path.exists() or not plate_config_path.exists():
+                raise FileNotFoundError("Missing model/config file!")
+            self.model_name = onnx_model_path.stem
+        elif hub_ocr_model:
+            self.model_name = hub_ocr_model
+            onnx_model_path, plate_config_path = hub.download_model(
+                model_name=hub_ocr_model, force_download=force_download
+            )
+        else:
+            raise ValueError(
+                "Either provide a model from the HUB or a custom model_path and config_path"
+            )
+
+        self.config = PlateOCRConfig.from_yaml(plate_config_path)
+        self.model = ort.InferenceSession(
+            onnx_model_path, providers=self.providers, sess_options=sess_options
+        )
+
+    def benchmark(
+        self,
+        n_iter: int = 2_500,
+        batch_size: int = 1,
+        include_processing: bool = False,
+        warmup: int = 250,
+    ) -> None:
+        """
+        Run an inference benchmark and pretty print the results.
+
+        It reports the following metrics:
+
+        * **Average latency per batch** (milliseconds)
+        * **Throughput** in *plates / second* (PPS), i.e., how many plates the pipeline can process
+          per second at the chosen ``batch_size``.
+
+        Args:
+            n_iter: The number of iterations to run the benchmark. This determines how many times
+                the inference will be executed to compute the average performance metrics.
+            batch_size : Batch size to use for the benchmark.
+            include_processing: Indicates whether the benchmark should include preprocessing and
+                postprocessing times in the measurement.
+            warmup: Number of warmup iterations to run before the benchmark.
+        """
+        x = np.random.randint(
+            0,
+            256,
+            size=(
+                batch_size,
+                self.config.img_height,
+                self.config.img_width,
+                self.config.num_channels,
+            ),
+            dtype=np.uint8,
+        )
+
+        # Warm-up
+        for _ in range(warmup):
+            if include_processing:
+                self.run(x)
+            else:
+                self.model.run(None, {"input": x})
+
+        # Timed loop
+        cum_time = 0.0
+        for _ in range(n_iter):
+            with measure_time() as time_taken:
+                if include_processing:
+                    self.run(x)
+                else:
+                    self.model.run(None, {"input": x})
+            cum_time += time_taken()
+
+        avg_time_ms = cum_time / n_iter if n_iter else 0.0
+        pps = (1_000 / avg_time_ms) * batch_size if n_iter else 0.0
+
+        console = Console()
+        model_info = Panel(
+            Text(f"Model: {self.model_name}\nProviders: {self.providers}", style="bold green"),
+            title="Model Information",
+            border_style="bright_blue",
+            expand=False,
+        )
+        console.print(model_info)
+        table = Table(title=f"Benchmark for '{self.model_name}'", border_style="bright_blue")
+        table.add_column("Metric", justify="center", style="cyan", no_wrap=True)
+        table.add_column("Value", justify="center", style="magenta")
+
+        table.add_row("Batch size", str(batch_size))
+        table.add_row("Warm-up iters", str(warmup))
+        table.add_row("Timed iterations", str(n_iter))
+        table.add_row("Average Time / batch (ms)", f"{avg_time_ms:.4f}")
+        table.add_row("Plates per Second (PPS)", f"{pps:.4f}")
+        console.print(table)
+
+    def run(
+        self,
+        source: Union[str, list[str], npt.NDArray, list[npt.NDArray]],
+        return_confidence: bool = False,
+    ) -> Union[tuple[list[str], npt.NDArray], list[str]]:
+        """
+        Performs OCR to recognize license plate characters from an image or a list of images.
+
+        Args:
+            source: One or more image inputs, which can be:
+
+                - A file path (`str` or `PathLike`) to an image.
+                - A list of file paths.
+                - A NumPy array of a single image, with shape (H, W), (H, W, 1) or (H, W, 3).
+                - A list of NumPy arrays, each representing an image.
+                - A 4D NumPy array of shape (N, H, W, C), ready for inference.
+
+                Images will be automatically resized and converted as needed based on the model's
+                configuration (including color mode and aspect ratio settings).
+
+            return_confidence: Whether to return confidence scores along with plate predictions.
+
+        Returns:
+            A list of recognized license plates (one per image). If `return_confidence` is True,
+            also returns a NumPy array of shape `(N, plate_slots)` containing the confidence scores
+            for each predicted character.
+        """
+        x = _load_image_from_source(source, self.config)
+        # Preprocess
+        x = preprocess_image(x)
+        # Run model
+        y: list[npt.NDArray] = self.model.run(None, {"input": x})
+        # Postprocess model output
+        return postprocess_output(
+            y[0],
+            self.config.max_plate_slots,
+            self.config.alphabet,
+            return_confidence=return_confidence,
+        )

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/train/data/augmentation.py

@@ -0,0 +1,101 @@
+"""
+Augmentations used for training the OCR model.
+"""
+
+import albumentations as A
+import cv2
+
+from fast_plate_ocr.core.types import ImageColorMode
+
+BORDER_COLOR_BLACK: tuple[int, int, int] = (0, 0, 0)
+
+
+def default_train_augmentation(img_color_mode: ImageColorMode) -> A.Compose:
+    """
+    Default training augmentation pipeline.
+    """
+    if img_color_mode == "grayscale":
+        return A.Compose(
+            [
+                A.Affine(
+                    translate_percent=(-0.02, 0.02),
+                    scale=(0.75, 1.10),
+                    rotate=(-15, 15),
+                    border_mode=cv2.BORDER_CONSTANT,
+                    fill=BORDER_COLOR_BLACK,
+                    shear=(0.0, 0.0),
+                    p=0.75,
+                ),
+                A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=1.0),
+                A.GaussianBlur(sigma_limit=(0.2, 0.5), p=0.25),
+                A.OneOf(
+                    [
+                        A.CoarseDropout(
+                            num_holes_range=(1, 14),
+                            hole_height_range=(1, 5),
+                            hole_width_range=(1, 5),
+                            p=0.2,
+                        ),
+                        A.PixelDropout(dropout_prob=0.02, p=0.2),
+                        A.GridDropout(ratio=0.3, fill="random", p=0.2),
+                    ],
+                    p=0.7,
+                ),
+            ]
+        )
+    if img_color_mode == "rgb":
+        return A.Compose(
+            [
+                A.Affine(
+                    translate_percent=(-0.02, 0.02),
+                    scale=(0.75, 1.10),
+                    rotate=(-15, 15),
+                    border_mode=cv2.BORDER_CONSTANT,
+                    fill=BORDER_COLOR_BLACK,
+                    shear=(0.0, 0.0),
+                    p=0.75,
+                ),
+                A.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=0.10, p=0.5),
+                A.OneOf(
+                    [
+                        A.HueSaturationValue(
+                            hue_shift_limit=5, sat_shift_limit=10, val_shift_limit=10, p=0.7
+                        ),
+                        A.RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10, p=0.3),
+                    ],
+                    p=0.3,
+                ),
+                A.RandomGamma(gamma_limit=(95, 105), p=0.20),
+                A.ToGray(p=0.05),
+                A.OneOf(
+                    [
+                        A.GaussianBlur(sigma_limit=(0.2, 0.5), p=0.5),
+                        A.MotionBlur(blur_limit=(3, 3), p=0.5),
+                    ],
+                    p=0.2,
+                ),
+                A.OneOf(
+                    [
+                        A.GaussNoise(std_range=(0.01, 0.03), p=0.2),
+                        A.MultiplicativeNoise(multiplier=(0.98, 1.02), p=0.1),
+                        A.ISONoise(intensity=(0.005, 0.02), p=0.1),
+                        A.ImageCompression(quality_range=(55, 90), p=0.1),
+                    ],
+                    p=0.3,
+                ),
+                A.OneOf(
+                    [
+                        A.CoarseDropout(
+                            num_holes_range=(1, 14),
+                            hole_height_range=(1, 5),
+                            hole_width_range=(1, 5),
+                            p=0.2,
+                        ),
+                        A.PixelDropout(dropout_prob=0.02, p=0.3),
+                        A.GridDropout(ratio=0.3, fill="random", p=0.3),
+                    ],
+                    p=0.5,
+                ),
+            ]
+        )
+    raise ValueError(f"Unsupported img_color_mode: {img_color_mode!r}. Expected 'grayscale'/'rgb'.")

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/train/data/dataset.py

@@ -0,0 +1,97 @@
+"""
+Dataset module.
+"""
+
+import math
+import os
+from typing import Union
+import albumentations as A
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+from keras.src.trainers.data_adapters.py_dataset_adapter import PyDataset
+
+from fast_plate_ocr.core.process import read_and_resize_plate_image
+from fast_plate_ocr.train.model.config import PlateOCRConfig
+from fast_plate_ocr.train.utilities import utils
+
+
+class PlateRecognitionPyDataset(PyDataset):
+    """
+    Custom PyDataset for OCR license plate recognition.
+    """
+
+    def __init__(
+        self,
+        annotations_file: Union[str, os.PathLike],
+        plate_config: PlateOCRConfig,
+        batch_size: int,
+        transform: Optional[A.Compose] = None,
+        shuffle: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        # Load annotations
+        annotations = pd.read_csv(annotations_file, dtype={"plate_text": str})
+        annotations["image_path"] = (
+            os.path.dirname(os.path.realpath(annotations_file)) + os.sep + annotations["image_path"]
+        )
+        # Check that plate lengths do not exceed max_plate_slots.
+        assert (annotations["plate_text"].str.len() <= plate_config.max_plate_slots).all(), (
+            "Plates are longer than max_plate_slots specified param. Change the parameter."
+        )
+        # Convert the dataframe to a NumPy array
+        self.annotations = annotations.to_numpy()
+
+        self.plate_config = plate_config
+        self.transform = transform
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+
+        # Shuffle once at initialization if `shuffle=True`
+        self._shuffle_data()
+
+    def __len__(self) -> int:
+        return math.ceil(len(self.annotations) / self.batch_size)
+
+    def __getitem__(self, idx: int) -> tuple[npt.NDArray, npt.NDArray]:
+        # Determine the idx-es of current batch
+        low = idx * self.batch_size
+        high = min(low + self.batch_size, len(self.annotations))
+        batch = self.annotations[low:high]
+
+        batch_x = []
+        batch_y = []
+        for image_path, plate_text in batch:
+            # Read and process image
+            x = read_and_resize_plate_image(
+                image_path=image_path,
+                img_height=self.plate_config.img_height,
+                img_width=self.plate_config.img_width,
+                image_color_mode=self.plate_config.image_color_mode,
+                keep_aspect_ratio=self.plate_config.keep_aspect_ratio,
+                interpolation_method=self.plate_config.interpolation,
+                padding_color=self.plate_config.padding_color,
+            )
+            # Transform target
+            y = utils.target_transform(
+                plate_text=plate_text,
+                max_plate_slots=self.plate_config.max_plate_slots,
+                alphabet=self.plate_config.alphabet,
+                pad_char=self.plate_config.pad_char,
+            )
+            # Apply augmentation if provided
+            if self.transform:
+                x = self.transform(image=x)["image"]
+            batch_x.append(x)
+            batch_y.append(y)
+
+        return np.array(batch_x), np.array(batch_y)
+
+    def _shuffle_data(self) -> None:
+        if self.shuffle:
+            np.random.shuffle(self.annotations)
+
+    def on_epoch_begin(self) -> None:
+        # Optionally shuffle the dataset at the start of each epoch
+        self._shuffle_data()

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/train/model/config.py

@@ -0,0 +1,114 @@
+"""
+License Plate OCR config. This config file defines how license plate images and text should be
+preprocessed for OCR model training and inference.
+"""
+
+from pathlib import Path
+from typing import Annotated, TypeAlias, Union
+
+import annotated_types
+import yaml
+from pydantic import (
+    BaseModel,
+    PositiveInt,
+    StringConstraints,
+    computed_field,
+    model_validator,
+)
+
+from fast_plate_ocr.core.types import ImageColorMode, ImageInterpolation, PathLike
+
+UInt8: TypeAlias = Annotated[int, annotated_types.Ge(0), annotated_types.Le(255)]
+"""
+An integer in the range [0, 255], used for color channel values.
+"""
+
+
+class PlateOCRConfig(BaseModel, extra="forbid", frozen=True):
+    """
+    Model License Plate OCR config.
+    """
+
+    max_plate_slots: PositiveInt
+    """
+    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: Annotated[str, StringConstraints(min_length=1, max_length=1)]
+    """
+    Padding character for plates which length is smaller than MAX_PLATE_SLOTS.
+    """
+    img_height: PositiveInt
+    """
+    Image height which is fed to the model.
+    """
+    img_width: PositiveInt
+    """
+    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[UInt8, UInt8, UInt8], UInt8] = (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.
+    """
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def vocabulary_size(self) -> int:
+        return len(self.alphabet)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def pad_idx(self) -> int:
+        return self.alphabet.index(self.pad_char)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def num_channels(self) -> int:
+        return 3 if self.image_color_mode == "rgb" else 1
+
+    @model_validator(mode="after")
+    def check_alphabet_and_pad(self) -> "PlateOCRConfig":
+        # `pad_char` must be in alphabet
+        if self.pad_char not in self.alphabet:
+            raise ValueError("Pad character must be present in model alphabet.")
+        # all chars in alphabet must be unique
+        if len(set(self.alphabet)) != len(self.alphabet):
+            raise ValueError("Alphabet must not contain duplicate characters.")
+        return self
+
+
+def load_plate_config_from_yaml(yaml_path: PathLike) -> PlateOCRConfig:
+    """
+    Reads and parses a YAML file containing the plate configuration.
+
+    Args:
+        yaml_path: Path to the YAML file containing the plate config.
+
+    Returns:
+        PlateOCRConfig: Parsed and validated plate configuration.
+
+    Raises:
+        FileNotFoundError: If the YAML file does not exist.
+    """
+    if not Path(yaml_path).is_file():
+        raise FileNotFoundError(f"Plate config '{yaml_path}' doesn't exist.")
+    with open(yaml_path, encoding="utf-8") as f_in:
+        yaml_content = yaml.safe_load(f_in)
+    config = PlateOCRConfig(**yaml_content)
+    return config