Merge branch 'rework-OCR-2.0' into develop

# Conflicts: # .github/workflows/main.yaml # pyproject.toml
ankandrew · Apr 8, 2024 · 32ee300 · 32ee300
2 parents 00b0a6f + ff24b8e
commit 32ee300
Show file tree

Hide file tree

Showing 174 changed files with 598 additions and 1,107 deletions.
diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml
@@ -2,6 +2,7 @@ name: Test and Deploy
 on: push
 jobs:
   test:
+    name: Test
     strategy:
       fail-fast: false
       matrix:
@@ -54,6 +55,10 @@ jobs:
         uses: actions/setup-python@v3
         with:
           python-version: '3.10'
+          cache: 'poetry'
+
+      - name: Install dependencies
+        run: poetry install --all-extras
 
       - name: Build a binary wheel
         run: poetry build

diff --git a/.gitignore b/.gitignore
@@ -157,3 +157,6 @@ cython_debug/
 
 # CUDA DNN
 cudnn64_7.dll
+
+# Train folder
+train_val_set/
diff --git a/LICENSE b/LICENSE
diff --git a/Makefile b/Makefile
@@ -1,11 +1,5 @@
 # Directories
-SRC_PATHS := fast_plate_ocr/ \
-            test/ \
-            train.py \
-            valid.py \
-            visualize_augmentation.py \
-            visualize_predictions.py \
-            onnx_converter.py
+SRC_PATHS := fast_plate_ocr/ test/
 
 # Tasks
 .PHONY: help
@@ -58,7 +52,7 @@ lint: ruff pylint mypy
 .PHONY: test
 test:
 	@echo "=====> Running tests..."
-	@poetry run pytest
+	@poetry run pytest test/
 
 .PHONY: clean
 clean:

diff --git a/assets/benchmark/annotations.csv b/assets/benchmark/annotations.csv
diff --git a/assets/benchmark/imgs/morning_10620_100.png b/assets/benchmark/imgs/morning_10620_100.png
diff --git a/assets/benchmark/imgs/morning_10710_65.png b/assets/benchmark/imgs/morning_10710_65.png
diff --git a/assets/benchmark/imgs/morning_11880_6.png b/assets/benchmark/imgs/morning_11880_6.png
diff --git a/assets/benchmark/imgs/morning_12060_22.png b/assets/benchmark/imgs/morning_12060_22.png
diff --git a/assets/benchmark/imgs/morning_16380_99.png b/assets/benchmark/imgs/morning_16380_99.png
diff --git a/assets/benchmark/imgs/morning_16740_10.png b/assets/benchmark/imgs/morning_16740_10.png
diff --git a/assets/benchmark/imgs/morning_1710_41.png b/assets/benchmark/imgs/morning_1710_41.png
diff --git a/assets/benchmark/imgs/morning_17190_68.png b/assets/benchmark/imgs/morning_17190_68.png
diff --git a/assets/benchmark/imgs/morning_19080_70.png b/assets/benchmark/imgs/morning_19080_70.png
diff --git a/assets/benchmark/imgs/morning_2160_87.png b/assets/benchmark/imgs/morning_2160_87.png
diff --git a/assets/benchmark/imgs/morning_3870_4.png b/assets/benchmark/imgs/morning_3870_4.png
diff --git a/assets/benchmark/imgs/morning_42480_61.png b/assets/benchmark/imgs/morning_42480_61.png
diff --git a/assets/benchmark/imgs/morning_43830_92.png b/assets/benchmark/imgs/morning_43830_92.png
diff --git a/assets/benchmark/imgs/morning_44010_51.png b/assets/benchmark/imgs/morning_44010_51.png
diff --git a/assets/benchmark/imgs/morning_45360_29.png b/assets/benchmark/imgs/morning_45360_29.png
diff --git a/assets/benchmark/imgs/morning_46080_44.png b/assets/benchmark/imgs/morning_46080_44.png
diff --git a/assets/benchmark/imgs/morning_46440_78.png b/assets/benchmark/imgs/morning_46440_78.png
diff --git a/assets/benchmark/imgs/morning_49770_49.png b/assets/benchmark/imgs/morning_49770_49.png
diff --git a/assets/benchmark/imgs/morning_49860_79.png b/assets/benchmark/imgs/morning_49860_79.png
diff --git a/assets/benchmark/imgs/morning_51120_77.png b/assets/benchmark/imgs/morning_51120_77.png
diff --git a/assets/benchmark/imgs/morning_52380_99.png b/assets/benchmark/imgs/morning_52380_99.png
diff --git a/assets/benchmark/imgs/morning_52470_3.png b/assets/benchmark/imgs/morning_52470_3.png
diff --git a/assets/benchmark/imgs/morning_52560_6.png b/assets/benchmark/imgs/morning_52560_6.png
diff --git a/assets/benchmark/imgs/morning_56880_50.png b/assets/benchmark/imgs/morning_56880_50.png
diff --git a/assets/benchmark/imgs/morning_65430_13.png b/assets/benchmark/imgs/morning_65430_13.png
diff --git a/assets/benchmark/imgs/morning_65430_2.png b/assets/benchmark/imgs/morning_65430_2.png
diff --git a/assets/benchmark/imgs/morning_66060_95.png b/assets/benchmark/imgs/morning_66060_95.png
diff --git a/assets/benchmark/imgs/morning_66150_20.png b/assets/benchmark/imgs/morning_66150_20.png
diff --git a/assets/benchmark/imgs/morning_68130_65.png b/assets/benchmark/imgs/morning_68130_65.png
diff --git a/assets/benchmark/imgs/morning_68130_67.png b/assets/benchmark/imgs/morning_68130_67.png
diff --git a/assets/benchmark/imgs/morning_68490_18.png b/assets/benchmark/imgs/morning_68490_18.png
diff --git a/assets/benchmark/imgs/morning_68940_92.png b/assets/benchmark/imgs/morning_68940_92.png
diff --git a/assets/benchmark/imgs/morning_71190_2.png b/assets/benchmark/imgs/morning_71190_2.png
diff --git a/assets/benchmark/imgs/morning_71280_14.png b/assets/benchmark/imgs/morning_71280_14.png
diff --git a/assets/benchmark/imgs/morning_71550_56.png b/assets/benchmark/imgs/morning_71550_56.png
diff --git a/assets/benchmark/imgs/morning_72720_34.png b/assets/benchmark/imgs/morning_72720_34.png
diff --git a/assets/benchmark/imgs/morning_76500_21.png b/assets/benchmark/imgs/morning_76500_21.png
diff --git a/assets/benchmark/imgs/morning_86220_67.png b/assets/benchmark/imgs/morning_86220_67.png
diff --git a/assets/benchmark/imgs/morning_86220_78.png b/assets/benchmark/imgs/morning_86220_78.png
diff --git a/assets/benchmark/imgs/morning_86220_82.png b/assets/benchmark/imgs/morning_86220_82.png
diff --git a/assets/benchmark/imgs/morning_86850_23.png b/assets/benchmark/imgs/morning_86850_23.png
diff --git a/assets/benchmark/imgs/morning_87480_36.png b/assets/benchmark/imgs/morning_87480_36.png
diff --git a/assets/benchmark/imgs/night_111540_61.png b/assets/benchmark/imgs/night_111540_61.png
diff --git a/assets/benchmark/imgs/night_11550_51.png b/assets/benchmark/imgs/night_11550_51.png
diff --git a/assets/benchmark/imgs/night_118668_81.png b/assets/benchmark/imgs/night_118668_81.png
diff --git a/assets/benchmark/imgs/night_118998_24.png b/assets/benchmark/imgs/night_118998_24.png
diff --git a/assets/benchmark/imgs/night_121110_88.png b/assets/benchmark/imgs/night_121110_88.png
diff --git a/assets/benchmark/imgs/night_124344_83.png b/assets/benchmark/imgs/night_124344_83.png
diff --git a/assets/benchmark/imgs/night_1254_62.png b/assets/benchmark/imgs/night_1254_62.png
diff --git a/assets/benchmark/imgs/night_130284_70.png b/assets/benchmark/imgs/night_130284_70.png
diff --git a/assets/benchmark/imgs/night_132660_40.png b/assets/benchmark/imgs/night_132660_40.png
diff --git a/assets/benchmark/imgs/night_133980_54.png b/assets/benchmark/imgs/night_133980_54.png
diff --git a/assets/benchmark/imgs/night_136026_77.png b/assets/benchmark/imgs/night_136026_77.png
diff --git a/assets/benchmark/imgs/night_137544_54.png b/assets/benchmark/imgs/night_137544_54.png
diff --git a/assets/benchmark/imgs/night_137874_52.png b/assets/benchmark/imgs/night_137874_52.png
diff --git a/assets/benchmark/imgs/night_138600_51.png b/assets/benchmark/imgs/night_138600_51.png
diff --git a/assets/benchmark/imgs/night_138666_32.png b/assets/benchmark/imgs/night_138666_32.png
diff --git a/assets/benchmark/imgs/night_138996_93.png b/assets/benchmark/imgs/night_138996_93.png
diff --git a/assets/benchmark/imgs/night_140448_9.png b/assets/benchmark/imgs/night_140448_9.png
diff --git a/assets/benchmark/imgs/night_142230_66.png b/assets/benchmark/imgs/night_142230_66.png
diff --git a/assets/benchmark/imgs/night_142362_36.png b/assets/benchmark/imgs/night_142362_36.png
diff --git a/assets/benchmark/imgs/night_14322_38.png b/assets/benchmark/imgs/night_14322_38.png
diff --git a/assets/benchmark/imgs/night_143352_20.png b/assets/benchmark/imgs/night_143352_20.png
diff --git a/assets/benchmark/imgs/night_144738_20.png b/assets/benchmark/imgs/night_144738_20.png
diff --git a/assets/benchmark/imgs/night_147048_42.png b/assets/benchmark/imgs/night_147048_42.png
diff --git a/assets/benchmark/imgs/night_15510_92.png b/assets/benchmark/imgs/night_15510_92.png
diff --git a/assets/benchmark/imgs/night_1584_30.png b/assets/benchmark/imgs/night_1584_30.png
diff --git a/assets/benchmark/imgs/night_1584_73.png b/assets/benchmark/imgs/night_1584_73.png
diff --git a/assets/benchmark/imgs/night_15972_7.png b/assets/benchmark/imgs/night_15972_7.png
diff --git a/assets/benchmark/imgs/night_16038_81.png b/assets/benchmark/imgs/night_16038_81.png
diff --git a/assets/benchmark/imgs/night_16170_61.png b/assets/benchmark/imgs/night_16170_61.png
diff --git a/assets/benchmark/imgs/night_1650_2.png b/assets/benchmark/imgs/night_1650_2.png
diff --git a/assets/benchmark/imgs/night_17490_91.png b/assets/benchmark/imgs/night_17490_91.png
diff --git a/assets/benchmark/imgs/night_22440_82.png b/assets/benchmark/imgs/night_22440_82.png
diff --git a/assets/benchmark/imgs/night_22968_59.png b/assets/benchmark/imgs/night_22968_59.png
diff --git a/assets/benchmark/imgs/night_24354_47.png b/assets/benchmark/imgs/night_24354_47.png
diff --git a/assets/benchmark/imgs/night_24618_10.png b/assets/benchmark/imgs/night_24618_10.png
diff --git a/assets/benchmark/imgs/night_25014_44.png b/assets/benchmark/imgs/night_25014_44.png
diff --git a/assets/benchmark/imgs/night_26136_6.png b/assets/benchmark/imgs/night_26136_6.png
diff --git a/assets/benchmark/imgs/night_31746_64.png b/assets/benchmark/imgs/night_31746_64.png
diff --git a/assets/benchmark/imgs/night_31812_31.png b/assets/benchmark/imgs/night_31812_31.png
diff --git a/assets/benchmark/imgs/night_36102_93.png b/assets/benchmark/imgs/night_36102_93.png
diff --git a/assets/benchmark/imgs/night_37752_76.png b/assets/benchmark/imgs/night_37752_76.png
diff --git a/assets/benchmark/imgs/night_43890_51.png b/assets/benchmark/imgs/night_43890_51.png
diff --git a/assets/benchmark/imgs/night_4686_0.png b/assets/benchmark/imgs/night_4686_0.png
diff --git a/assets/benchmark/imgs/night_46926_40.png b/assets/benchmark/imgs/night_46926_40.png
diff --git a/assets/benchmark/imgs/night_47058_34.png b/assets/benchmark/imgs/night_47058_34.png
diff --git a/assets/benchmark/imgs/night_47124_39.png b/assets/benchmark/imgs/night_47124_39.png
diff --git a/assets/benchmark/imgs/night_47124_86.png b/assets/benchmark/imgs/night_47124_86.png
diff --git a/assets/benchmark/imgs/night_47190_45.png b/assets/benchmark/imgs/night_47190_45.png
diff --git a/assets/benchmark/imgs/night_47190_86.png b/assets/benchmark/imgs/night_47190_86.png
diff --git a/assets/benchmark/imgs/night_47256_61.png b/assets/benchmark/imgs/night_47256_61.png
diff --git a/assets/benchmark/imgs/night_48114_97.png b/assets/benchmark/imgs/night_48114_97.png
diff --git a/assets/benchmark/imgs/night_49038_89.png b/assets/benchmark/imgs/night_49038_89.png
diff --git a/assets/benchmark/imgs/night_50226_38.png b/assets/benchmark/imgs/night_50226_38.png
diff --git a/assets/benchmark/imgs/night_50226_98.png b/assets/benchmark/imgs/night_50226_98.png
diff --git a/assets/benchmark/imgs/night_50952_32.png b/assets/benchmark/imgs/night_50952_32.png
diff --git a/assets/benchmark/imgs/night_5148_38.png b/assets/benchmark/imgs/night_5148_38.png
diff --git a/assets/benchmark/imgs/night_5676_15.png b/assets/benchmark/imgs/night_5676_15.png
diff --git a/assets/benchmark/imgs/night_6072_4.png b/assets/benchmark/imgs/night_6072_4.png
diff --git a/assets/benchmark/imgs/night_6138_71.png b/assets/benchmark/imgs/night_6138_71.png
diff --git a/assets/benchmark/imgs/night_6732_36.png b/assets/benchmark/imgs/night_6732_36.png
diff --git a/assets/benchmark/imgs/night_75372_90.png b/assets/benchmark/imgs/night_75372_90.png
diff --git a/assets/benchmark/imgs/night_86460_33.png b/assets/benchmark/imgs/night_86460_33.png
diff --git a/assets/benchmark/imgs/night_86526_58.png b/assets/benchmark/imgs/night_86526_58.png
diff --git a/assets/benchmark/imgs/night_86658_36.png b/assets/benchmark/imgs/night_86658_36.png
diff --git a/assets/benchmark/imgs/night_86790_79.png b/assets/benchmark/imgs/night_86790_79.png
diff --git a/assets/benchmark/imgs/night_88176_18.png b/assets/benchmark/imgs/night_88176_18.png
diff --git a/assets/benchmark/imgs/night_94182_9.png b/assets/benchmark/imgs/night_94182_9.png
diff --git a/assets/benchmark/imgs/night_97416_12.png b/assets/benchmark/imgs/night_97416_12.png
diff --git a/assets/benchmark/imgs/night_97878_2.png b/assets/benchmark/imgs/night_97878_2.png
diff --git a/assets/benchmark/imgs/night_9900_56.png b/assets/benchmark/imgs/night_9900_56.png
diff --git a/assets/benchmark/imgs/random_1.png b/assets/benchmark/imgs/random_1.png
diff --git a/assets/benchmark/imgs/random_1095.png b/assets/benchmark/imgs/random_1095.png
diff --git a/assets/benchmark/imgs/random_1240.png b/assets/benchmark/imgs/random_1240.png
diff --git a/assets/benchmark/imgs/random_194.png b/assets/benchmark/imgs/random_194.png
diff --git a/assets/benchmark/imgs/random_2.png b/assets/benchmark/imgs/random_2.png
diff --git a/assets/benchmark/imgs/random_2116.png b/assets/benchmark/imgs/random_2116.png
diff --git a/assets/benchmark/imgs/random_2591.png b/assets/benchmark/imgs/random_2591.png
diff --git a/assets/benchmark/imgs/random_3.png b/assets/benchmark/imgs/random_3.png
diff --git a/assets/benchmark/imgs/random_3028.png b/assets/benchmark/imgs/random_3028.png
diff --git a/assets/benchmark/imgs/random_3396.png b/assets/benchmark/imgs/random_3396.png
diff --git a/assets/benchmark/imgs/random_3746.png b/assets/benchmark/imgs/random_3746.png
diff --git a/assets/benchmark/imgs/random_4238.png b/assets/benchmark/imgs/random_4238.png
diff --git a/assets/benchmark/imgs/random_427.png b/assets/benchmark/imgs/random_427.png
diff --git a/assets/benchmark/imgs/random_4690.png b/assets/benchmark/imgs/random_4690.png
diff --git a/assets/benchmark/imgs/random_525.png b/assets/benchmark/imgs/random_525.png
diff --git a/assets/benchmark/imgs/random_5513.png b/assets/benchmark/imgs/random_5513.png
diff --git a/assets/benchmark/imgs/random_6101.png b/assets/benchmark/imgs/random_6101.png
diff --git a/assets/benchmark/imgs/random_6224.png b/assets/benchmark/imgs/random_6224.png
diff --git a/assets/benchmark/imgs/random_6547.png b/assets/benchmark/imgs/random_6547.png
diff --git a/assets/benchmark/imgs/random_670.png b/assets/benchmark/imgs/random_670.png
diff --git a/assets/benchmark/imgs/random_6925.png b/assets/benchmark/imgs/random_6925.png
diff --git a/assets/benchmark/imgs/random_7031.png b/assets/benchmark/imgs/random_7031.png
diff --git a/assets/benchmark/imgs/random_7457.png b/assets/benchmark/imgs/random_7457.png
diff --git a/assets/benchmark/imgs/random_7752.png b/assets/benchmark/imgs/random_7752.png
diff --git a/assets/benchmark/imgs/random_8104.png b/assets/benchmark/imgs/random_8104.png
diff --git a/assets/benchmark/imgs/random_8416.png b/assets/benchmark/imgs/random_8416.png
diff --git a/assets/benchmark/imgs/random_8818.png b/assets/benchmark/imgs/random_8818.png
diff --git a/assets/benchmark/imgs/random_8901.png b/assets/benchmark/imgs/random_8901.png
diff --git a/assets/benchmark/imgs/random_8910.png b/assets/benchmark/imgs/random_8910.png
diff --git a/assets/benchmark/imgs/random_9007.png b/assets/benchmark/imgs/random_9007.png
diff --git a/assets/benchmark/imgs/random_9202.png b/assets/benchmark/imgs/random_9202.png
diff --git a/assets/benchmark/imgs/random_9324.png b/assets/benchmark/imgs/random_9324.png
diff --git a/assets/benchmark/imgs/random_9542.png b/assets/benchmark/imgs/random_9542.png
diff --git a/assets/benchmark/imgs/random_9657.png b/assets/benchmark/imgs/random_9657.png
diff --git a/config/arg_cnn_ocr_config.yaml → config/arg_plate_example.yaml b/config/arg_cnn_ocr_config.yaml → config/arg_plate_example.yaml
@@ -1,3 +1,7 @@
+# Config example for Argentinian License Plates
+# The old license plates contain 6 slots/characters (i.e. JUH697)
+# and new 'Mercosur' contain 7 slots/characters (i.e. AB123CD)
+
 # Max number of plate slots supported. This represents the number of model classification heads.
 max_plate_slots: 7
 # All the possible character set for the model output.

diff --git a/fast_plate_ocr/__init__.py b/fast_plate_ocr/__init__.py
@@ -0,0 +1,7 @@
+"""
+Fast Plate OCR package.
+"""
+
+from fast_plate_ocr.inference.onnx_inference import ONNXPlateRecognizer
+
+__all__ = ["ONNXPlateRecognizer"]
diff --git a/train_val_set/train/.gitkeep → fast_plate_ocr/cli/__init__.py b/train_val_set/train/.gitkeep → fast_plate_ocr/cli/__init__.py
diff --git a/fast_plate_ocr/cli/cli.py b/fast_plate_ocr/cli/cli.py
@@ -0,0 +1,26 @@
+"""
+Main CLI used when training a FastPlateOCR model.
+"""
+
+try:
+    import click
+
+    from fast_plate_ocr.cli.onnx_converter import export_onnx
+    from fast_plate_ocr.cli.train import train
+    from fast_plate_ocr.cli.valid import valid
+    from fast_plate_ocr.cli.visualize_augmentation import visualize_augmentation
+    from fast_plate_ocr.cli.visualize_predictions import visualize_predictions
+except ImportError as e:
+    raise ImportError("Make sure to 'pip install fast-plate-ocr[train]' to run this!") from e
+
+
+@click.group(context_settings={"max_content_width": 120})
+def main_cli():
+    """FastPlateOCR CLI."""
+
+
+main_cli.add_command(visualize_predictions)
+main_cli.add_command(visualize_augmentation)
+main_cli.add_command(valid)
+main_cli.add_command(train)
+main_cli.add_command(export_onnx)
diff --git a/onnx_converter.py → fast_plate_ocr/cli/onnx_converter.py b/onnx_converter.py → fast_plate_ocr/cli/onnx_converter.py
@@ -16,10 +16,12 @@
 from fast_plate_ocr.train.model.config import load_config_from_yaml
 from fast_plate_ocr.train.utilities.utils import load_keras_model
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+)
 
 
-@click.command(context_settings={"max_content_width": 140})
+@click.command(context_settings={"max_content_width": 120})
 @click.option(
     "-m",
     "--model",
@@ -47,12 +49,15 @@
     show_default=True,
     help="Opset version for ONNX.",
 )
-def export_model_as_onnx(
+def export_onnx(
     model_path: pathlib.Path,
     output_path: str,
     config_file: pathlib.Path,
     opset: int,
 ) -> None:
+    """
+    Export Keras models to ONNX format.
+    """
     config = load_config_from_yaml(config_file)
     model = load_keras_model(
         model_path,
@@ -80,4 +85,4 @@ def export_model_as_onnx(
 
 
 if __name__ == "__main__":
-    export_model_as_onnx()
+    export_onnx()
diff --git a/train.py → fast_plate_ocr/cli/train.py b/train.py → fast_plate_ocr/cli/train.py
@@ -27,7 +27,7 @@
 # pylint: disable=too-many-arguments,too-many-locals
 
 
-@click.command(context_settings={"max_content_width": 140})
+@click.command(context_settings={"max_content_width": 120})
 @click.option(
     "--dense/--no-dense",
     default=True,
@@ -36,8 +36,7 @@
 )
 @click.option(
     "--config-file",
-    default="./config/arg_plates.yaml",
-    show_default=True,
+    required=True,
     type=click.Path(exists=True, file_okay=True, path_type=pathlib.Path),
     help="Path pointing to the model license plate OCR config.",
 )
@@ -128,6 +127,9 @@ def train(
     early_stopping_patience: int,
     reduce_lr_patience: int,
 ) -> None:
+    """
+    Train the License Plate OCR model.
+    """
     train_augmentation = (
         A.load(augmentation_path, data_format="yaml") if augmentation_path else TRAIN_AUGMENTATION
     )

diff --git a/valid.py → fast_plate_ocr/cli/valid.py b/valid.py → fast_plate_ocr/cli/valid.py
@@ -14,7 +14,7 @@
 from fast_plate_ocr.train.utilities import utils
 
 
-@click.command(context_settings={"max_content_width": 140})
+@click.command(context_settings={"max_content_width": 120})
 @click.option(
     "-m",
     "--model",
@@ -25,16 +25,14 @@
 )
 @click.option(
     "--config-file",
-    default="./config/arg_plates.yaml",
-    show_default=True,
+    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",
-    default="assets/benchmark/annotations.csv",
-    show_default=True,
+    required=True,
     type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path),
     help="Annotations file used for validation.",
 )
@@ -52,7 +50,9 @@ def valid(
     annotations: pathlib.Path,
     batch_size: int,
 ) -> None:
-    """Validate a model for a given annotated data."""
+    """
+    Validate the trained OCR model on a labeled set.
+    """
     config = load_config_from_yaml(config_file)
     model = utils.load_keras_model(
         model_path, vocab_size=config.vocabulary_size, max_plate_slots=config.max_plate_slots