Move train-related scripts to code package

ankandrew · Apr 6, 2024 · 9a8ddee · 9a8ddee
1 parent a1a31dd
commit 9a8ddee
Show file tree

Hide file tree

Showing 9 changed files with 31 additions and 8 deletions.
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

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
diff --git a/train.py → fast_plate_ocr/cli/train.py b/train.py → fast_plate_ocr/cli/train.py
diff --git a/valid.py → fast_plate_ocr/cli/valid.py b/valid.py → fast_plate_ocr/cli/valid.py
diff --git a/visualize_augmentation.py → fast_plate_ocr/cli/visualize_augmentation.py b/visualize_augmentation.py → fast_plate_ocr/cli/visualize_augmentation.py
diff --git a/visualize_predictions.py → fast_plate_ocr/cli/visualize_predictions.py b/visualize_predictions.py → fast_plate_ocr/cli/visualize_predictions.py
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -71,6 +71,9 @@ pandas-stubs = "^2.2.0.240218"
 pylint = "^3.1.0"
 types-pyyaml = "^6.0.12.20240311"
 
+[tool.poetry.scripts]
+plate_ocr = "fast_plate_ocr.cli.cli:main_cli"
+
 [tool.ruff]
 line-length = 100
 target-version = "py310"