Move backend utils to proper module to avoid keras being loaded first

Author: ankandrew

fast_plate_ocr/backend_utils.py

@@ -0,0 +1,27 @@
+"""
+Utils for Keras supported backends.
+"""
+
+import os
+
+from fast_plate_ocr.custom_types import Framework
+
+
+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

fast_plate_ocr/utils.py

@@ -3,7 +3,6 @@
 """
 
 import logging
-import os
 import pathlib
 import random
 import time
@@ -25,7 +24,6 @@
     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
 
 
 def one_hot_plate(plate: str, alphabet: str = MODEL_ALPHABET) -> list[list[int]]:
@@ -45,26 +43,6 @@ def target_transform(
     return encoded_plate
 
 
-def set_tensorflow_backend() -> None:
-    """Set Keras backend to tensorflow."""
-    set_keras_backend("tensorflow")
-
-
-def set_jax_backend() -> None:
-    """Set Keras backend to jax."""
-    set_keras_backend("jax")
-
-
-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 read_plate_image(
     image_path: str, img_height: int = DEFAULT_IMG_HEIGHT, img_width: int = DEFAULT_IMG_WIDTH
 ) -> npt.NDArray:

test/fast_lp_ocr/test_custom.py

@@ -5,7 +5,7 @@
 # ruff: noqa: E402
 # pylint: disable=wrong-import-position,wrong-import-order,ungrouped-imports
 # fmt: off
-from fast_plate_ocr.utils import set_pytorch_backend
+from fast_plate_ocr.backend_utils import set_pytorch_backend
 
 set_pytorch_backend()
 # fmt: on