matrice-analytics 0.1.60__py3-none-any.whl

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

@@ -0,0 +1,95 @@
+"""
+Model builder functions for supported architectures.
+"""
+
+from collections.abc import Sequence
+
+import keras
+import numpy as np
+from keras import layers
+
+from fast_plate_ocr.train.model.config import PlateOCRConfig
+from fast_plate_ocr.train.model.layers import (
+    PatchExtractor,
+    PositionEmbedding,
+    TokenReducer,
+    TransformerBlock,
+    VocabularyProjection,
+)
+from fast_plate_ocr.train.model.model_schema import AnyModelConfig, CCTModelConfig, LayerConfig
+
+
+def _build_stem_from_config(specs: Sequence[LayerConfig]) -> keras.Sequential:
+    return keras.Sequential([spec.to_keras_layer() for spec in specs], name="conv_stem")
+
+
+def _build_cct_model(
+    cfg: CCTModelConfig,
+    input_shape: tuple[int, int, int],
+    max_plate_slots: int,
+    vocabulary_size: int,
+) -> keras.Model:
+    # 1. Input
+    inputs = layers.Input(shape=input_shape)
+
+    # 2. Rescale & conv stem
+    data_rescale = cfg.rescaling.to_keras_layer()
+    x = _build_stem_from_config(cfg.tokenizer.blocks)(data_rescale(inputs))
+
+    # 3. Patch extraction: (B, H, W, C) -> (B, num_patches, C*patch_size**2)
+    x = PatchExtractor(patch_size=cfg.tokenizer.patch_size)(x)
+
+    # 5. Optional patch MLP
+    if cfg.tokenizer.patch_mlp is not None:
+        x = cfg.tokenizer.patch_mlp.to_keras_layer()(x)
+
+    # 6. Positional embeddings
+    if cfg.tokenizer.positional_emb:
+        seq_len = keras.ops.shape(x)[1]
+        x = x + PositionEmbedding(sequence_length=seq_len, name="pos_emb")(x)
+
+    # 7. N x TransformerBlock's
+    dpr = list(
+        np.linspace(0.0, cfg.transformer_encoder.stochastic_depth, cfg.transformer_encoder.layers)
+    )
+    for i, rate in enumerate(dpr, 1):
+        x = TransformerBlock(
+            projection_dim=cfg.transformer_encoder.projection_dim,
+            num_heads=cfg.transformer_encoder.heads,
+            mlp_units=cfg.transformer_encoder.units,
+            attention_dropout=cfg.transformer_encoder.attention_dropout,
+            mlp_dropout=cfg.transformer_encoder.mlp_dropout,
+            drop_path_rate=rate,
+            norm_type=cfg.transformer_encoder.normalization,
+            activation=cfg.transformer_encoder.activation,
+            name=f"transformer_block_{i}",
+        )(x)
+
+    # 8. Reduce to a fixed number of tokens, then project to vocab
+    x = TokenReducer(
+        num_tokens=max_plate_slots,
+        projection_dim=cfg.transformer_encoder.projection_dim,
+        num_heads=cfg.transformer_encoder.token_reducer_heads,
+    )(x)
+
+    logits = VocabularyProjection(
+        vocabulary_size=vocabulary_size,
+        dropout_rate=cfg.transformer_encoder.head_mlp_dropout,
+        name="vocab_projection",
+    )(x)
+
+    return keras.Model(inputs, logits, name="CCT_OCR")
+
+
+def build_model(model_cfg: AnyModelConfig, plate_cfg: PlateOCRConfig) -> keras.Model:
+    """
+    Build a Keras OCR model based on the specified model and plate configuration.
+    """
+    if model_cfg.model == "cct":
+        return _build_cct_model(
+            cfg=model_cfg,
+            input_shape=(plate_cfg.img_height, plate_cfg.img_width, plate_cfg.num_channels),
+            max_plate_slots=plate_cfg.max_plate_slots,
+            vocabulary_size=plate_cfg.vocabulary_size,
+        )
+    raise ValueError(f"Unsupported model type: {model_cfg.model!r}")

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

@@ -0,0 +1,395 @@
+"""
+Schema definitions for validating supported model architectures and layer blocks.
+"""
+
+from pathlib import Path
+from typing import Annotated, Literal, Optional, TypeAlias, Union
+
+import keras
+import yaml
+from keras.src.layers import RMSNormalization
+from pydantic import BaseModel, Field, PositiveFloat, PositiveInt, model_validator
+
+from fast_plate_ocr.core.types import PathLike
+from fast_plate_ocr.train.model.layers import (
+    MLP,
+    CoordConv2D,
+    DyT,
+    MaxBlurPooling2D,
+    SqueezeExcite,
+)
+
+UnitFloat: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)]
+"""A float that must be in range of [0, 1]."""
+PaddingTypeStr: TypeAlias = Literal["valid", "same"]
+"""Padding modes supported by Keras convolution and pooling layers."""
+PositiveIntTuple: TypeAlias = Annotated[Union[PositiveInt, tuple[PositiveInt, PositiveInt]], Field(discriminator="type")]
+"""A single positive integer or a tuple of two positive integers, usually used for sizes/strides."""
+NormalizationStr: TypeAlias = Literal["layer_norm", "rms_norm", "dyt"]
+"""Available normalization layers."""
+
+ActivationStr: TypeAlias = Literal[
+    "celu",
+    "elu",
+    "exponential",
+    "gelu",
+    "glu",
+    "hard_shrink",
+    "hard_sigmoid",
+    "hard_silu",
+    "hard_tanh",
+    "leaky_relu",
+    "linear",
+    "log_sigmoid",
+    "log_softmax",
+    "mish",
+    "relu",
+    "relu6",
+    "selu",
+    "sigmoid",
+    "silu",
+    "soft_shrink",
+    "softmax",
+    "softplus",
+    "softsign",
+    "sparse_plus",
+    "sparsemax",
+    "squareplus",
+    "tanh",
+    "tanh_shrink",
+    "threshold",
+]
+"""Supported Keras activation functions."""
+
+WeightInitializationStr: TypeAlias = Literal[
+    "glorot_normal",
+    "glorot_uniform",
+    "he_normal",
+    "he_uniform",
+    "lecun_normal",
+    "lecun_uniform",
+    "ones",
+    "random_normal",
+    "random_uniform",
+    "truncated_normal",
+    "variance_scaling",
+    "zeros",
+]
+"""Keras weight initialization strategies."""
+
+
+class _Rescaling(BaseModel):
+    scale: float = 1.0 / 255
+    offset: float = 0.0
+
+    def to_keras_layer(self):
+        return keras.layers.Rescaling(self.scale, self.offset)
+
+
+class _Activation(BaseModel):
+    layer: Literal["Activation"]
+    activation: ActivationStr
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.Activation(self.activation)
+
+
+class _Conv2DBase(BaseModel):
+    filters: PositiveInt
+    kernel_size: PositiveIntTuple
+    strides: PositiveIntTuple = 1
+    padding: PaddingTypeStr = "same"
+    activation: ActivationStr = "relu"
+    use_bias: bool = True
+    kernel_initializer: WeightInitializationStr = "he_normal"
+    bias_initializer: WeightInitializationStr = "zeros"
+
+
+class _Conv2D(_Conv2DBase):
+    layer: Literal["Conv2D"]
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        params = self.model_dump(exclude={"layer"})
+        return keras.layers.Conv2D(**params)
+
+
+class _CoordConv2D(_Conv2DBase):
+    layer: Literal["CoordConv2D"]
+    with_r: bool = False
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        params = self.model_dump(exclude={"layer", "with_r"})
+        return CoordConv2D(with_r=self.with_r, **params)
+
+
+class _DepthwiseConv2D(BaseModel):
+    layer: Literal["DepthwiseConv2D"]
+    kernel_size: PositiveIntTuple
+    strides: PositiveIntTuple = 1
+    padding: PaddingTypeStr = "same"
+    depth_multiplier: PositiveInt = 1
+    activation: ActivationStr = "relu"
+    use_bias: bool = True
+    depthwise_initializer: WeightInitializationStr = "he_normal"
+    bias_initializer: WeightInitializationStr = "zeros"
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.DepthwiseConv2D(
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            padding=self.padding,
+            depth_multiplier=self.depth_multiplier,
+            activation=self.activation,
+            use_bias=self.use_bias,
+            depthwise_initializer=self.depthwise_initializer,
+            bias_initializer=self.bias_initializer,
+        )
+
+
+class _SeparableConv2D(BaseModel):
+    layer: Literal["SeparableConv2D"]
+    filters: PositiveInt
+    kernel_size: PositiveIntTuple
+    strides: PositiveIntTuple = 1
+    padding: PaddingTypeStr = "same"
+    depth_multiplier: PositiveInt = 1
+    activation: ActivationStr = "relu"
+    use_bias: bool = True
+    depthwise_initializer: WeightInitializationStr = "he_normal"
+    pointwise_initializer: WeightInitializationStr = "glorot_uniform"
+    bias_initializer: WeightInitializationStr = "zeros"
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.SeparableConv2D(
+            filters=self.filters,
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            padding=self.padding,
+            depth_multiplier=self.depth_multiplier,
+            activation=self.activation,
+            use_bias=self.use_bias,
+            depthwise_initializer=self.depthwise_initializer,
+            pointwise_initializer=self.pointwise_initializer,
+            bias_initializer=self.bias_initializer,
+        )
+
+
+class _MLP(BaseModel):
+    layer: Literal["MLP"]
+    hidden_units: list[PositiveInt]
+    dropout_rate: UnitFloat = 0.1
+    activation: ActivationStr = "gelu"
+    use_bias: bool = True
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return MLP(
+            hidden_units=self.hidden_units,
+            dropout_rate=self.dropout_rate,
+            activation=self.activation,
+            use_bias=self.use_bias,
+        )
+
+
+class _MaxBlurPooling2D(BaseModel):
+    layer: Literal["MaxBlurPooling2D"]
+    pool_size: PositiveInt = 2
+    filter_size: PositiveInt = 3
+    padding: PaddingTypeStr = "same"
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return MaxBlurPooling2D(
+            pool_size=self.pool_size, filter_size=self.filter_size, padding=self.padding
+        )
+
+
+class _MaxPooling2D(BaseModel):
+    layer: Literal["MaxPooling2D"]
+    pool_size: PositiveIntTuple = 2
+    strides: Optional[PositiveInt] = None
+    padding: PaddingTypeStr = "valid"
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.MaxPooling2D(
+            pool_size=self.pool_size,
+            strides=self.strides,
+            padding=self.padding,
+        )
+
+
+class _AveragePooling2D(BaseModel):
+    layer: Literal["AveragePooling2D"]
+    pool_size: PositiveIntTuple = 2
+    strides: Optional[PositiveInt] = None
+    padding: PaddingTypeStr = "valid"
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.AveragePooling2D(
+            pool_size=self.pool_size,
+            strides=self.strides,
+            padding=self.padding,
+        )
+
+
+class _ZeroPadding2D(BaseModel):
+    layer: Literal["ZeroPadding2D"]
+    padding: PositiveIntTuple = 1
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.ZeroPadding2D(padding=self.padding)
+
+
+class _SqueezeExcite(BaseModel):
+    layer: Literal["SqueezeExcite"]
+    ratio: PositiveFloat = 1.0
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return SqueezeExcite(ratio=self.ratio)
+
+
+class _BatchNormalization(BaseModel):
+    layer: Literal["BatchNormalization"]
+    momentum: PositiveFloat = 0.99
+    epsilon: PositiveFloat = 1e-3
+    center: bool = True
+    scale: bool = True
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.BatchNormalization(
+            momentum=self.momentum,
+            epsilon=self.epsilon,
+            center=self.center,
+            scale=self.scale,
+        )
+
+
+class _Dropout(BaseModel):
+    layer: Literal["Dropout"]
+    rate: PositiveFloat
+
+    def to_keras_layer(self):
+        return keras.layers.Dropout(rate=self.rate)
+
+
+class _SpatialDropout2D(BaseModel):
+    layer: Literal["SpatialDropout2D"]
+    rate: PositiveFloat
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.SpatialDropout2D(rate=self.rate)
+
+
+class _GaussianNoise(BaseModel):
+    layer: Literal["GaussianNoise"]
+    stddev: PositiveFloat
+    seed: Optional[int] = None
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.GaussianNoise(stddev=self.stddev, seed=self.seed)
+
+
+class _LayerNorm(BaseModel):
+    layer: Literal["LayerNorm"]
+    epsilon: PositiveFloat = 1e-3
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return keras.layers.LayerNormalization(epsilon=self.epsilon)
+
+
+class _RMSNorm(BaseModel):
+    layer: Literal["RMSNorm"]
+    epsilon: PositiveFloat = 1e-6
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return RMSNormalization(epsilon=self.epsilon)
+
+
+class _DyT(BaseModel):
+    layer: Literal["DyT"]
+    alpha_init_value: PositiveFloat = 0.5
+
+    def to_keras_layer(self) -> keras.layers.Layer:
+        return DyT(alpha_init_value=self.alpha_init_value)
+
+
+LayerConfig = Annotated[
+    _Activation
+    | _Conv2D
+    | _CoordConv2D
+    | _DepthwiseConv2D
+    | _SeparableConv2D
+    | _MLP
+    | _MaxBlurPooling2D
+    | _MaxPooling2D
+    | _AveragePooling2D
+    | _ZeroPadding2D
+    | _SqueezeExcite
+    | _BatchNormalization
+    | _Dropout
+    | _SpatialDropout2D
+    | _GaussianNoise
+    | _LayerNorm
+    | _RMSNorm
+    | _DyT,
+    Field(discriminator="layer"),
+]
+
+
+class _CCTTokenizerConfig(BaseModel):
+    blocks: list[LayerConfig]
+    patch_size: PositiveIntTuple = 1
+    patch_mlp: Optional[_MLP] = None
+    positional_emb: bool = True
+
+
+class _CCTTransformerEncoderConfig(BaseModel):
+    layers: PositiveInt
+    heads: PositiveInt
+    projection_dim: PositiveInt
+    units: list[PositiveInt]
+    activation: ActivationStr = "gelu"
+    stochastic_depth: UnitFloat = 0.1
+    attention_dropout: UnitFloat = 0.1
+    mlp_dropout: UnitFloat = 0.1
+    head_mlp_dropout: UnitFloat = 0.2
+    token_reducer_heads: PositiveInt = 2
+    normalization: NormalizationStr = "layer_norm"
+
+    @model_validator(mode="after")
+    def _consistency_checks(self):
+        if self.units[-1] != self.projection_dim:
+            raise ValueError(
+                "'units[-1]' must equal 'projection_dim' "
+                f"(got {self.units[-1]} vs {self.projection_dim})."
+            )
+        return self
+
+
+class CCTModelConfig(BaseModel):
+    model: Literal["cct"] = "cct"
+    rescaling: _Rescaling
+    tokenizer: _CCTTokenizerConfig
+    transformer_encoder: _CCTTransformerEncoderConfig
+
+
+AnyModelConfig = Annotated[CCTModelConfig, Field(discriminator="model")]
+"""Supported model-architecture. New model configs should be added here."""
+
+
+def load_model_config_from_yaml(yaml_path: PathLike) -> AnyModelConfig:
+    """
+    Loads, parses, and validates a YAML file defining a model architecture.
+
+    Args:
+        yaml_path: Path to the YAML file.
+
+    Returns:
+        AnyModelConfig: Parsed and validated model configuration.
+
+    Raises:
+        FileNotFoundError: If the YAML file does not exist.
+    """
+    if not Path(yaml_path).is_file():
+        raise FileNotFoundError(f"Model config '{yaml_path}' doesn't exist.")
+    with open(yaml_path, encoding="utf-8") as f_in:
+        data = yaml.safe_load(f_in)
+    return AnyModelConfig(**data)

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/train/utilities/backend_utils.py

@@ -0,0 +1,38 @@
+"""
+Utils for Keras supported backends.
+"""
+
+import os
+from typing import Literal, TypeAlias
+
+Framework: TypeAlias = Literal["jax", "tensorflow", "torch"]
+"""Supported backend frameworks for Keras."""
+
+
+def set_jax_backend() -> None:
+    """Set Keras backend to jax."""
+    set_keras_backend("jax")
+
+
+def set_tensorflow_backend() -> None:
+    """Set Keras backend to tensorflow."""
+    set_keras_backend("tensorflow")
+
+
+def set_pytorch_backend() -> None:
+    """Set Keras backend to pytorch."""
+    set_keras_backend("torch")
+
+
+def set_keras_backend(framework: Framework) -> None:
+    """Set the Keras backend to a given framework."""
+    os.environ["KERAS_BACKEND"] = framework
+
+
+def reload_keras_backend(framework: Framework) -> None:
+    """Reload the Keras backend with a given framework."""
+    # ruff: noqa: PLC0415
+    # pylint: disable=import-outside-toplevel
+    import keras
+
+    keras.config.set_backend(framework)