Add and use more utilities

Author: ankandrew

fast_plate_ocr/utils.py

@@ -2,13 +2,29 @@
 Utility functions module
 """
 
+import logging
 import os
+import pathlib
+import random
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
 
 import cv2
+import keras
 import numpy as np
 import numpy.typing as npt
-
-from fast_plate_ocr.config import MAX_PLATE_SLOTS, MODEL_ALPHABET, PAD_CHAR
+from keras.src.activations import softmax
+
+from fast_plate_ocr.config import (
+    DEFAULT_IMG_HEIGHT,
+    DEFAULT_IMG_WIDTH,
+    MAX_PLATE_SLOTS,
+    MODEL_ALPHABET,
+    PAD_CHAR,
+    VOCABULARY_SIZE,
+)
+from fast_plate_ocr.custom import cat_acc_metric, cce_loss, plate_acc_metric, top_3_k_metric
 from fast_plate_ocr.custom_types import Framework
 
 
@@ -49,7 +65,9 @@ def set_keras_backend(framework: Framework) -> None:
     os.environ["KERAS_BACKEND"] = framework
 
 
-def read_plate_image(image_path: str, img_height: int, img_width: int) -> npt.NDArray:
+def read_plate_image(
+    image_path: str, img_height: int = DEFAULT_IMG_HEIGHT, img_width: int = DEFAULT_IMG_WIDTH
+) -> npt.NDArray:
     """
     Read and resize a license plate image.
 
@@ -62,3 +80,59 @@ def read_plate_image(image_path: str, img_height: int, img_width: int) -> npt.ND
     img = cv2.resize(img, (img_width, img_height), interpolation=cv2.INTER_LINEAR)
     img = np.expand_dims(img, -1)
     return img
+
+
+def load_keras_model(
+    model_path: pathlib.Path,
+    vocab_size: int = VOCABULARY_SIZE,
+    max_plate_slots: int = MAX_PLATE_SLOTS,
+) -> keras.Model:
+    """
+    Utility helper function to load the keras OCR model.
+    """
+    custom_objects = {
+        "cce": cce_loss(vocabulary_size=vocab_size),
+        "cat_acc": cat_acc_metric(max_plate_slots=max_plate_slots, vocabulary_size=vocab_size),
+        "plate_acc": plate_acc_metric(max_plate_slots=max_plate_slots, vocabulary_size=vocab_size),
+        "top_3_k": top_3_k_metric(vocabulary_size=vocab_size),
+        "softmax": softmax,
+    }
+    model = keras.models.load_model(model_path, custom_objects=custom_objects)
+    return model
+
+
+IMG_EXTENSIONS: set[str] = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"}
+"""Valid image extensions for the scope of this script."""
+
+
+def load_images_from_folder(
+    img_dir: pathlib.Path,
+    width: int = DEFAULT_IMG_WIDTH,
+    height: int = DEFAULT_IMG_HEIGHT,
+    shuffle: bool = False,
+    limit: int | None = None,
+) -> list[npt.NDArray]:
+    """
+    Return all images read from a directory. This uses the same read function used during training.
+    """
+    image_paths = sorted(
+        str(f.resolve()) for f in img_dir.iterdir() if f.is_file() and f.suffix in IMG_EXTENSIONS
+    )
+    if limit:
+        image_paths = image_paths[:limit]
+    if shuffle:
+        random.shuffle(image_paths)
+    images = [read_plate_image(i, img_height=height, img_width=width) for i in image_paths]
+    return images
+
+
+@contextmanager
+def log_time_taken(process_name: str) -> Iterator[None]:
+    """A concise context manager to time code snippets and log the result."""
+    time_start: float = time.perf_counter()
+    try:
+        yield
+    finally:
+        time_end: float = time.perf_counter()
+        time_elapsed: float = time_end - time_start
+        logging.info("Computation time of '%s' = %.3fms", process_name, 1000 * time_elapsed)

valid.py

@@ -5,14 +5,12 @@
 import pathlib
 
 import click
-import keras
-from keras.src.activations import softmax
 from torch.utils.data import DataLoader
 
+from fast_plate_ocr import utils
 from fast_plate_ocr.config import MAX_PLATE_SLOTS, MODEL_ALPHABET, PAD_CHAR, VOCABULARY_SIZE
 
 # Custom metris / losses
-from fast_plate_ocr.custom import cat_acc_metric, cce_loss, plate_acc_metric, top_3_k_metric
 from fast_plate_ocr.dataset import LicensePlateDataset
 
 
@@ -79,14 +77,7 @@ def valid(
     pad_char: str,
 ) -> None:
     """Validate a model for a given annotated data."""
-    custom_objects = {
-        "cce": cce_loss(vocabulary_size=vocab_size),
-        "cat_acc": cat_acc_metric(max_plate_slots=plate_slots, vocabulary_size=vocab_size),
-        "plate_acc": plate_acc_metric(max_plate_slots=plate_slots, vocabulary_size=vocab_size),
-        "top_3_k": top_3_k_metric(vocabulary_size=vocab_size),
-        "softmax": softmax,
-    }
-    model = keras.models.load_model(model_path, custom_objects=custom_objects)
+    model = utils.load_keras_model(model_path, vocab_size=vocab_size, max_plate_slots=plate_slots)
     val_torch_dataset = LicensePlateDataset(
         annotations_file=annotations,
         max_plate_slots=plate_slots,

visualize_augmentation.py

@@ -11,12 +11,9 @@
 import numpy as np
 import numpy.typing as npt
 
+from fast_plate_ocr import utils
 from fast_plate_ocr.augmentation import TRAIN_AUGMENTATION
 from fast_plate_ocr.config import DEFAULT_IMG_HEIGHT, DEFAULT_IMG_WIDTH
-from fast_plate_ocr.utils import read_plate_image
-
-IMG_EXTENSIONS: set[str] = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"}
-"""Valid image extensions for the scope of this script."""
 
 
 def _set_seed(seed: int | None) -> None:
@@ -29,18 +26,13 @@ def _set_seed(seed: int | None) -> None:
 def load_images(
     img_dir: pathlib.Path,
     num_images: int,
-    shuffle_img: bool,
+    shuffle: bool,
     height: int,
     width: int,
 ) -> tuple[list[npt.NDArray[np.uint8]], list[npt.NDArray[np.uint8]]]:
-    img_paths = sorted(f for f in img_dir.iterdir() if f.is_file() and f.suffix in IMG_EXTENSIONS)
-    img_paths = img_paths[:num_images]
-    if shuffle_img:
-        random.shuffle(img_paths)
-    images = [
-        read_plate_image(image_path=str(img), img_height=height, img_width=width)
-        for img in img_paths
-    ]
+    images = utils.load_images_from_folder(
+        img_dir, height=height, width=width, shuffle=shuffle, limit=num_images
+    )
     augmented_images = [TRAIN_AUGMENTATION(image=i)["image"] for i in images]
     return images, augmented_images
 
@@ -94,7 +86,7 @@ def display_images(
     help="Maximum number of images to visualize.",
 )
 @click.option(
-    "--shuffle_img",
+    "--shuffle",
     "-s",
     is_flag=True,
     default=False,
@@ -146,7 +138,7 @@ def display_images(
 def visualize_augmentation(
     img_dir: pathlib.Path,
     num_images: int,
-    shuffle_img: bool,
+    shuffle: bool,
     columns: int,
     rows: int,
     height: int,
@@ -155,7 +147,7 @@ def visualize_augmentation(
     show_original: bool,
 ) -> None:
     _set_seed(seed)
-    images, augmented_images = load_images(img_dir, num_images, shuffle_img, height, width)
+    images, augmented_images = load_images(img_dir, num_images, shuffle, height, width)
     display_images(images, augmented_images, columns, rows, show_original)