matrice-analytics 0.1.60__py3-none-any.whl

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/cli/train.py

@@ -0,0 +1,447 @@
+"""
+Script for training the License Plate OCR models.
+"""
+
+from __future__ import annotations
+
+import json
+import pathlib
+import shutil
+from datetime import datetime
+from typing import Literal, Optional
+
+import albumentations as A
+import click
+import keras
+from keras.src.callbacks import (
+    CSVLogger,
+    EarlyStopping,
+    ModelCheckpoint,
+    SwapEMAWeights,
+    TensorBoard,
+    TerminateOnNaN,
+)
+from keras.src.optimizers import AdamW
+
+import fast_plate_ocr.train.model.model_builders
+from fast_plate_ocr.cli.utils import print_params, print_train_details
+from fast_plate_ocr.train.data.augmentation import (
+    default_train_augmentation,
+)
+from fast_plate_ocr.train.data.dataset import PlateRecognitionPyDataset
+from fast_plate_ocr.train.model.config import load_plate_config_from_yaml
+from fast_plate_ocr.train.model.loss import cce_loss, focal_cce_loss
+from fast_plate_ocr.train.model.metric import (
+    cat_acc_metric,
+    plate_acc_metric,
+    plate_len_acc_metric,
+    top_3_k_metric,
+)
+from fast_plate_ocr.train.model.model_schema import load_model_config_from_yaml
+
+# ruff: noqa: PLR0913
+# pylint: disable=too-many-arguments,too-many-locals
+
+
+EVAL_METRICS: dict[str, Literal["max", "min", "auto"]] = {
+    "val_plate_acc": "max",
+    "val_cat_acc": "max",
+    "val_top_3_k_acc": "max",
+    "val_plate_len_acc": "max",
+    "val_loss": "min",
+}
+"""Eval metric to monitor."""
+
+
+@click.command(context_settings={"max_content_width": 120})
+@click.option(
+    "--model-config-file",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="Path to the YAML config that describes the model architecture.",
+)
+@click.option(
+    "--plate-config-file",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="Path to the plate YAML config.",
+)
+@click.option(
+    "--annotations",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="Path pointing to the train annotations CSV file.",
+)
+@click.option(
+    "--val-annotations",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="Path pointing to the train validation CSV file.",
+)
+@click.option(
+    "--validation-freq",
+    default=1,
+    show_default=True,
+    type=int,
+    help="Frequency (in epochs) at which to evaluate the validation data.",
+)
+@click.option(
+    "--augmentation-path",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="YAML file pointing to the augmentation pipeline saved with Albumentations.save(...)",
+)
+@click.option(
+    "--lr",
+    default=0.001,
+    show_default=True,
+    type=float,
+    help="Initial learning rate.",
+)
+@click.option(
+    "--final-lr-factor",
+    default=1e-2,
+    show_default=True,
+    type=float,
+    help="Final learning rate factor for the cosine decay scheduler. It's the fraction of"
+    " the initial learning rate that remains after decay.",
+)
+@click.option(
+    "--warmup-fraction",
+    default=0.05,
+    show_default=True,
+    type=float,
+    help="Fraction of total training steps to linearly warm up.",
+)
+@click.option(
+    "--weight-decay",
+    default=0.001,
+    show_default=True,
+    type=float,
+    help="Weight decay for the AdamW optimizer.",
+)
+@click.option(
+    "--clipnorm",
+    default=1.0,
+    show_default=True,
+    type=float,
+    help="Gradient clipping norm value for the AdamW optimizer.",
+)
+@click.option(
+    "--loss",
+    default="cce",
+    type=click.Choice(["cce", "focal_cce"], case_sensitive=False),
+    show_default=True,
+    help="Loss function to use during training.",
+)
+@click.option(
+    "--focal-alpha",
+    default=0.25,
+    show_default=True,
+    type=float,
+    help="Alpha parameter for focal loss. Applicable only when '--loss' is 'focal_cce'.",
+)
+@click.option(
+    "--focal-gamma",
+    default=2.0,
+    show_default=True,
+    type=float,
+    help="Gamma parameter for focal loss. Applicable only when '--loss' is 'focal_cce'.",
+)
+@click.option(
+    "--label-smoothing",
+    default=0.01,
+    show_default=True,
+    type=float,
+    help="Amount of label smoothing to apply.",
+)
+@click.option(
+    "--mixed-precision-policy",
+    default=None,
+    type=click.Choice(["mixed_float16", "mixed_bfloat16", "float32"]),
+    help=(
+        "Optional mixed precision policy for training. Choose one of: mixed_float16, "
+        "mixed_bfloat16, or float32. If not provided, Keras uses its default global policy."
+    ),
+)
+@click.option(
+    "--batch-size",
+    default=64,
+    show_default=True,
+    type=int,
+    help="Batch size for training.",
+)
+@click.option(
+    "--workers",
+    default=1,
+    show_default=True,
+    type=int,
+    help="Number of worker threads/processes for parallel data loading.",
+)
+@click.option(
+    "--use-multiprocessing/--no-use-multiprocessing",
+    default=False,
+    show_default=True,
+    help="Use multiprocessing for data loading.",
+)
+@click.option(
+    "--max-queue-size",
+    default=10,
+    show_default=True,
+    type=int,
+    help="Maximum queue size for dataset workers.",
+)
+@click.option(
+    "--output-dir",
+    default="./trained_models",
+    type=click.Path(dir_okay=True, file_okay=False, path_type=pathlib.Path),
+    help="Output directory where model will be saved.",
+)
+@click.option(
+    "--epochs",
+    default=300,
+    show_default=True,
+    type=int,
+    help="Number of training epochs.",
+)
+@click.option(
+    "--tensorboard",
+    "-t",
+    is_flag=True,
+    help="Whether to use TensorBoard visualization tool.",
+)
+@click.option(
+    "--tensorboard-dir",
+    "-l",
+    default="tensorboard_logs",
+    show_default=True,
+    type=click.Path(path_type=pathlib.Path),
+    help="The path of the directory where to save the TensorBoard log files.",
+)
+@click.option(
+    "--early-stopping-patience",
+    default=100,
+    show_default=True,
+    type=int,
+    help="Stop training when the early stopping metric doesn't improve for X epochs.",
+)
+@click.option(
+    "--early-stopping-metric",
+    default="val_plate_acc",
+    show_default=True,
+    type=click.Choice(list(EVAL_METRICS), case_sensitive=False),
+    help="Metric to monitor for early stopping.",
+)
+@click.option(
+    "--weights-path",
+    type=click.Path(exists=True, file_okay=True, path_type=pathlib.Path),
+    help="Path to the pretrained model weights file.",
+)
+@click.option(
+    "--use-ema/--no-use-ema",
+    default=True,
+    show_default=True,
+    help=(
+        "Whether to use exponential moving averages in the AdamW optimizer. "
+        "Defaults to True; use --no-use-ema to disable."
+    ),
+)
+@click.option(
+    "--wd-ignore",
+    default="bias,layer_norm",
+    show_default=True,
+    type=str,
+    help="Comma-separated list of variable substrings to exclude from weight decay.",
+)
+@click.option(
+    "--seed",
+    type=int,
+    help="Sets all random seeds (Python, NumPy, and backend framework, e.g. TF).",
+)
+@print_params(table_title="CLI Training Parameters", c1_title="Parameter", c2_title="Details")
+def train(
+    model_config_file: pathlib.Path,
+    plate_config_file: pathlib.Path,
+    annotations: pathlib.Path,
+    val_annotations: pathlib.Path,
+    validation_freq: int,
+    augmentation_path: Optional[pathlib.Path],
+    lr: float,
+    final_lr_factor: float,
+    warmup_fraction: float,
+    weight_decay: float,
+    clipnorm: float,
+    loss: str,
+    focal_alpha: float,
+    focal_gamma: float,
+    label_smoothing: float,
+    mixed_precision_policy: Optional[str],
+    batch_size: int,
+    workers: int,
+    use_multiprocessing: bool,
+    max_queue_size: int,
+    output_dir: pathlib.Path,
+    epochs: int,
+    tensorboard: bool,
+    tensorboard_dir: pathlib.Path,
+    early_stopping_patience: int,
+    early_stopping_metric: str,
+    weights_path: Optional[pathlib.Path],
+    use_ema: bool,
+    wd_ignore: str,
+    seed: Optional[int],
+) -> None:
+    """
+    Train the License Plate OCR model.
+    """
+    if seed is not None:
+        keras.utils.set_random_seed(seed)
+
+    if mixed_precision_policy is not None:
+        keras.mixed_precision.set_global_policy(mixed_precision_policy)
+
+    plate_config = load_plate_config_from_yaml(plate_config_file)
+    model_config = load_model_config_from_yaml(model_config_file)
+    train_augmentation = (
+        A.load(augmentation_path, data_format="yaml")
+        if augmentation_path
+        else default_train_augmentation(img_color_mode=plate_config.image_color_mode)
+    )
+    print_train_details(train_augmentation, plate_config.model_dump())
+
+    train_dataset = PlateRecognitionPyDataset(
+        annotations_file=annotations,
+        transform=train_augmentation,
+        plate_config=plate_config,
+        batch_size=batch_size,
+        shuffle=True,
+        workers=workers,
+        use_multiprocessing=use_multiprocessing,
+        max_queue_size=max_queue_size,
+    )
+
+    val_dataset = PlateRecognitionPyDataset(
+        annotations_file=val_annotations,
+        plate_config=plate_config,
+        batch_size=batch_size,
+        shuffle=False,
+        workers=workers,
+        use_multiprocessing=use_multiprocessing,
+        max_queue_size=max_queue_size,
+    )
+
+    # Train
+    model = fast_plate_ocr.train.model.model_builders.build_model(model_config, plate_config)
+
+    if weights_path:
+        model.load_weights(weights_path, skip_mismatch=True)
+
+    total_steps = epochs * len(train_dataset)
+    warmup_steps = int(warmup_fraction * total_steps)
+
+    cosine_decay = keras.optimizers.schedules.CosineDecay(
+        initial_learning_rate=0.0 if warmup_steps > 0 else lr,
+        decay_steps=total_steps,
+        alpha=final_lr_factor,
+        warmup_steps=warmup_steps,
+        warmup_target=lr if warmup_steps > 0 else None,
+    )
+
+    optimizer = AdamW(cosine_decay, weight_decay=weight_decay, clipnorm=clipnorm, use_ema=use_ema)
+    optimizer.exclude_from_weight_decay(
+        var_names=[name.strip() for name in wd_ignore.split(",") if name.strip()]
+    )
+
+    if loss == "cce":
+        loss_fn = cce_loss(
+            vocabulary_size=plate_config.vocabulary_size, label_smoothing=label_smoothing
+        )
+    elif loss == "focal_cce":
+        loss_fn = focal_cce_loss(
+            vocabulary_size=plate_config.vocabulary_size,
+            alpha=focal_alpha,
+            gamma=focal_gamma,
+            label_smoothing=label_smoothing,
+        )
+    else:
+        raise ValueError(f"Unsupported loss type: {loss}")
+
+    model.compile(
+        loss=loss_fn,
+        jit_compile=False,
+        optimizer=optimizer,
+        metrics=[
+            cat_acc_metric(
+                max_plate_slots=plate_config.max_plate_slots,
+                vocabulary_size=plate_config.vocabulary_size,
+            ),
+            plate_acc_metric(
+                max_plate_slots=plate_config.max_plate_slots,
+                vocabulary_size=plate_config.vocabulary_size,
+            ),
+            top_3_k_metric(vocabulary_size=plate_config.vocabulary_size),
+            plate_len_acc_metric(
+                max_plate_slots=plate_config.max_plate_slots,
+                vocabulary_size=plate_config.vocabulary_size,
+                pad_token_index=plate_config.pad_idx,
+            ),
+        ],
+    )
+
+    output_dir /= datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model_file_path = output_dir / "ckpt-epoch_{epoch:02d}-acc_{val_plate_acc:.3f}.keras"
+
+    # Save params and configs used for training
+    shutil.copy(model_config_file, output_dir / "model_config.yaml")
+    shutil.copy(plate_config_file, output_dir / "plate_config.yaml")
+    A.save(train_augmentation, output_dir / "train_augmentation.yaml", "yaml")
+    with open(output_dir / "hyper_params.json", "w", encoding="utf-8") as f_out:
+        json.dump(
+            {k: v for k, v in locals().items() if k in click.get_current_context().params},
+            f_out,
+            indent=4,
+            default=str,
+        )
+
+    callbacks = [
+        # Stop training when early_stopping_metric doesn't improve for X epochs
+        EarlyStopping(
+            monitor=early_stopping_metric,
+            patience=early_stopping_patience,
+            mode=EVAL_METRICS[early_stopping_metric],
+            restore_best_weights=False,
+            verbose=1,
+        ),
+        # To save model checkpoint with EMA weights, we need to place this before `ModelCheckpoint`
+        *([SwapEMAWeights(swap_on_epoch=True)] if use_ema else []),
+        # We don't use EarlyStopping restore_best_weights=True because it won't restore the best
+        # weights when it didn't manage to EarlyStop but finished all epochs
+        ModelCheckpoint(output_dir / "last.keras", save_weights_only=False, save_best_only=False),
+        ModelCheckpoint(
+            model_file_path,
+            monitor=early_stopping_metric,
+            mode=EVAL_METRICS[early_stopping_metric],
+            save_weights_only=False,
+            save_best_only=True,
+            verbose=1,
+        ),
+        TerminateOnNaN(),
+        CSVLogger(str(output_dir / "training_log.csv")),
+    ]
+
+    if tensorboard:
+        run_dir = tensorboard_dir / datetime.now().strftime("run_%Y-%m-%d_%H-%M-%S")
+        run_dir.mkdir(parents=True, exist_ok=True)
+        callbacks.append(TensorBoard(log_dir=run_dir))
+
+    model.fit(
+        train_dataset,
+        epochs=epochs,
+        validation_data=val_dataset,
+        callbacks=callbacks,
+        validation_freq=validation_freq,
+    )
+
+
+if __name__ == "__main__":
+    train()

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

@@ -0,0 +1,129 @@
+"""
+Utils used for the CLI scripts.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import inspect
+import pathlib
+import random
+from collections.abc import Callable, Sequence
+from functools import wraps
+from typing import Any, Optional
+
+import albumentations as A
+import numpy as np
+from rich import box
+from rich.console import Console
+from rich.pretty import Pretty
+from rich.table import Table
+
+
+def print_variables_as_table(
+    c1_title: str, c2_title: str, title: str = "Variables Table", **kwargs: Any
+) -> None:
+    """
+    Prints variables in a formatted table using the rich library.
+
+    Args:
+        c1_title (str): Title of the first column.
+        c2_title (str): Title of the second column.
+        title (str): Title of the table.
+        **kwargs (Any): Variable names and values to be printed.
+    """
+    console = Console()
+    console.print("\n")
+    table = Table(title=title, show_header=True, header_style="bold blue", box=box.ROUNDED)
+    table.add_column(c1_title, min_width=20, justify="left", style="bold")
+    table.add_column(c2_title, min_width=60, justify="left", style="bold")
+
+    for key, value in kwargs.items():
+        if isinstance(value, pathlib.Path):
+            value = str(value)  # noqa: PLW2901
+        table.add_row(f"[bold]{key}[/bold]", Pretty(value))
+
+    console.print(table)
+
+
+def print_params(
+    table_title: str = "Parameters Table", c1_title: str = "Variable", c2_title: str = "Value"
+) -> Callable:
+    """
+    A decorator that prints the parameters of a function in a formatted table
+    using the rich library.
+
+    Args:
+        c1_title (str, optional): Title of the first column. Defaults to "Variable".
+        c2_title (str, optional): Title of the second column. Defaults to "Value".
+        table_title (str, optional): Title of the table. Defaults to "Parameters Table".
+
+    Returns:
+        Callable: The wrapped function with parameter printing functionality.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            func_signature = inspect.signature(func)
+            bound_arguments = func_signature.bind(*args, **kwargs)
+            bound_arguments.apply_defaults()
+            params = dict(bound_arguments.arguments.items())
+            print_variables_as_table(c1_title, c2_title, table_title, **params)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def print_train_details(augmentation: A.Compose, config: dict[str, Any]) -> None:
+    console = Console()
+    console.print("\n")
+    console.print("[bold blue]Augmentation Pipeline:[/bold blue]")
+    console.print(Pretty(augmentation))
+    console.print("\n")
+    console.print("[bold blue]Configuration:[/bold blue]")
+    console.print(Pretty(config))
+    console.print("\n")
+
+
+def requires(*modules: str, pkg_name: Optional[Sequence[str]] = None) -> Callable:
+    """
+    Decorator that checks if given modules are importable. If not, raises ModuleNotFoundError with
+    a hint to install the package(s).
+
+    Args:
+        modules (str): Names of modules to check (via importlib.util.find_spec).
+        pkg_name (Optional[Sequence[str]]): Names of packages to suggest installing.
+
+    Returns:
+        Callable: The wrapped function that checks for module availability.
+    """
+
+    def decorator(fn: Callable) -> Callable:
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any):
+            missing = [m for m in modules if importlib.util.find_spec(m) is None]
+            if missing:
+                pkg_missing = " ".join(pkg_name or missing)
+                raise ModuleNotFoundError(
+                    f"Cannot run `{fn.__name__}` because {missing!r} "
+                    f"is not installed. Please install the required package(s): {pkg_missing}"
+                )
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def seed_everything(seed: int) -> None:
+    """
+    Seed random number generators for reproducibility.
+
+    Args:
+        seed (int): The seed value to set.
+    """
+    random.seed(seed)
+    np.random.seed(seed)

matrice_analytics/post_processing/ocr/fast_plate_ocr_py38/cli/valid.py

@@ -0,0 +1,93 @@
+"""
+Script for validating trained OCR models.
+"""
+
+from __future__ import annotations
+
+import pathlib
+
+import click
+
+from fast_plate_ocr.train.data.dataset import PlateRecognitionPyDataset
+from fast_plate_ocr.train.model.config import load_plate_config_from_yaml
+from fast_plate_ocr.train.utilities.utils import load_keras_model
+
+
+@click.command(context_settings={"max_content_width": 120})
+@click.option(
+    "-m",
+    "--model",
+    "model_path",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="Path to the saved .keras model.",
+)
+@click.option(
+    "--plate-config-file",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, path_type=pathlib.Path),
+    help="Path pointing to the model license plate OCR config.",
+)
+@click.option(
+    "-a",
+    "--annotations",
+    required=True,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
+    help="Annotations file used for validation.",
+)
+@click.option(
+    "-b",
+    "--batch-size",
+    default=1,
+    show_default=True,
+    type=int,
+    help="Batch size.",
+)
+@click.option(
+    "--workers",
+    default=1,
+    show_default=True,
+    type=int,
+    help="Number of worker threads/processes for parallel data loading via PyDataset.",
+)
+@click.option(
+    "--use-multiprocessing/--no-use-multiprocessing",
+    default=False,
+    show_default=True,
+    help="Whether to use multiprocessing for data loading.",
+)
+@click.option(
+    "--max-queue-size",
+    default=10,
+    show_default=True,
+    type=int,
+    help="Maximum number of batches to prefetch for the dataset.",
+)
+def valid(
+    model_path: pathlib.Path,
+    plate_config_file: pathlib.Path,
+    annotations: pathlib.Path,
+    batch_size: int,
+    workers: int,
+    use_multiprocessing: bool,
+    max_queue_size: int,
+) -> None:
+    """
+    Validate the trained OCR model on a labeled set.
+    """
+    plate_config = load_plate_config_from_yaml(plate_config_file)
+    model = load_keras_model(model_path, plate_config)
+    val_dataset = PlateRecognitionPyDataset(
+        annotations_file=annotations,
+        plate_config=plate_config,
+        batch_size=batch_size,
+        shuffle=False,
+        workers=workers,
+        use_multiprocessing=use_multiprocessing,
+        max_queue_size=max_queue_size,
+    )
+    model.evaluate(val_dataset)
+
+
+if __name__ == "__main__":
+    valid()