diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..9985f20 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,28 @@ +# This workflow will install Python dependencies, run tests and lint with a single version of Python +# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions + +name: Test + +on: [push, pull_request] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.8", "3.9", "3.10", "3.11"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Display Python version + run: python -c "import sys; print(sys.version)" + - name: Install dependencies + run: | + pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu + pip install pytest pytest-cov coveralls pytest-sugar + - name: Run Tests + run: | + pytest --doctest-modules --cov=yaltai --verbose diff --git a/README.md b/README.md index 2c98e89..e813d4f 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Convert (and split optionally) your data ```bash # Keeps .1 data in the validation set and convert all alto into YOLOv5 format # Keeps the segmonto information up to the regions -yaltai alto-to-yolo PATH/TO/ALTOorPAGE/*.xml my-dataset --shuffle .1 --segmonto region +yaltai convert alto-to-yolo PATH/TO/ALTOorPAGE/*.xml my-dataset --shuffle .1 --segmonto region ``` And then [train YOLO](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data) diff --git a/setup.py b/setup.py index 9af34a7..69c32c6 100644 --- a/setup.py +++ b/setup.py @@ -105,7 +105,7 @@ def run(self): # py_modules=['mypackage'], entry_points={ - 'console_scripts': ['yaltai=yaltai.yaltai:cli', 'kraken-yaltai=yaltai.kraken_yaltai:cli'], + 'console_scripts': ['yaltai=yaltai.cli.yaltai:yaltai_cli'], }, install_requires=REQUIRED, extras_require=EXTRAS, diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/nano-yolo-ladas.pt b/tests/nano-yolo-ladas.pt new file mode 100644 index 0000000..f3906d5 --- /dev/null +++ b/tests/nano-yolo-ladas.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5386ac4dc37ba05ec1f42f0b8224bd1e13b0ec4f273624c49c0045b601502b84 +size 6246382 diff --git a/tests/test_convert.py b/tests/test_convert.py new file mode 100644 index 0000000..086aa8a --- /dev/null +++ b/tests/test_convert.py @@ -0,0 +1,40 @@ +import tempfile +import os +from click.testing import CliRunner + +from yaltai.cli.yaltai import yaltai_cli + + +def test_yaltai_single_alto_to_xml(): + """Ensures that we can convert to YOLO format""" + runner = CliRunner() + + with tempfile.TemporaryDirectory() as tempdir: + # Run the Click command + result = runner.invoke( + yaltai_cli, + [ + "convert", + "alto-to-yolo", + os.path.join(os.path.abspath(os.path.dirname(__file__)), "test_files", "alto_dataset", "output.xml"), + tempdir, + ]) + + # Ensure the command ran successfully + assert result.exit_code == 0 + assert "Found 1 to convert." in result.output + assert "- 00001 NumberingZone" in result.output, "Correct number of zone types are found" + assert "- 00001 RunningTitleZone" in result.output, "Correct number of zone types are found" + assert "- 00001 MainZone-P-Continued" in result.output, "Correct number of zone types are found" + assert "- 00004 MainZone-P" in result.output, "Correct number of zone types are found" + with open(os.path.join(tempdir, "labels", "output.txt")) as f: + data = [line.split() for line in f.read().split("\n")] + assert data == [ + ['0', '0.345606', '0.117500', '0.037501', '0.020206'], + ['1', '0.612827', '0.118333', '0.212882', '0.020736'], + ['2', '0.616390', '0.210833', '0.578172', '0.142851'], + ['3', '0.616390', '0.365833', '0.577029', '0.162800'], + ['3', '0.616390', '0.547500', '0.574365', '0.199581'], + ['3', '0.616390', '0.727500', '0.576050', '0.161861'], + ['3', '0.611639', '0.819167', '0.583077', '0.022773'] + ] diff --git a/tests/test_files/.gitignore b/tests/test_files/.gitignore new file mode 100644 index 0000000..c7eabc8 --- /dev/null +++ b/tests/test_files/.gitignore @@ -0,0 +1 @@ +page1.xml \ No newline at end of file diff --git a/tests/test_files/alto_dataset/output.jpg b/tests/test_files/alto_dataset/output.jpg new file mode 100644 index 0000000..ecf2ad2 Binary files /dev/null and b/tests/test_files/alto_dataset/output.jpg differ diff --git a/tests/test_files/alto_dataset/output.xml b/tests/test_files/alto_dataset/output.xml new file mode 100644 index 0000000..6c38f73 --- /dev/null +++ b/tests/test_files/alto_dataset/output.xml @@ -0,0 +1,12430 @@ + + + + pixel + + output.jpg + + + contentGeneration + Baseline and region segmentation + model: blla.mlmodel; text_direction: horizontal-lr + + kraken + 5.2.9 + + + + contentGeneration + Text line recognition + text_direction: horizontal-tb; models: catmus-print-fondue-large.mlmodel; pad: 16; bidi_reordering: True + + kraken + 5.2.9 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/tests/test_files/annot1.txt b/tests/test_files/annot1.txt new file mode 100644 index 0000000..50e12d4 --- /dev/null +++ b/tests/test_files/annot1.txt @@ -0,0 +1,7 @@ +14 0.49139869729173813 0.3715823573386494 0.6673020226259856 0.07884874158983304 +14 0.49408639012684263 0.23613506105158236 0.6722831676379842 0.15367804634936458 +35 0.49086047308878983 0.1285771243458759 0.26186150154268084 0.03464490406179915 +31 0.8086732944806309 0.13971592324943932 0.03529310935893041 0.03310740094692251 +14 0.49913952691121016 0.594151507600299 0.6766986630099417 0.31005232992773485 +14 0.507857387727117 0.8290804884126588 0.6902228316763799 0.1168950909543982 +1 0.9231984916009598 0.4728781460254174 0.1535995886184436 0.18793919760777475 \ No newline at end of file diff --git a/tests/test_files/annot2.txt b/tests/test_files/annot2.txt new file mode 100644 index 0000000..60ac401 --- /dev/null +++ b/tests/test_files/annot2.txt @@ -0,0 +1,6 @@ +23 0.4914115952466691 0.45760844079718643 0.6449045732805185 0.4102297772567409 +23 0.4933705437522506 0.7764806565064478 0.6465250270075622 0.2320398593200469 +31 0.7885343896290963 0.11840562719812427 0.033172488296723084 0.026182883939038688 +35 0.4834569679510263 0.1191535756154748 0.4334101548433561 0.024438452520515828 +24 0.48678790061217136 0.20451113716295427 0.6372740367302845 0.10247831184056272 +28 0.47840835433921497 0.14436342321219228 0.24528628015844436 0.025828839390386868 \ No newline at end of file diff --git a/tests/test_files/label_map.txt b/tests/test_files/label_map.txt new file mode 100644 index 0000000..53ead56 --- /dev/null +++ b/tests/test_files/label_map.txt @@ -0,0 +1,4 @@ +Class0 +Class1 + +Stuff \ No newline at end of file diff --git a/tests/test_files/page1.jpg b/tests/test_files/page1.jpg new file mode 100644 index 0000000..21c37c4 Binary files /dev/null and b/tests/test_files/page1.jpg differ diff --git a/tests/test_kraken.py b/tests/test_kraken.py new file mode 100644 index 0000000..a7afa7f --- /dev/null +++ b/tests/test_kraken.py @@ -0,0 +1,108 @@ +import os +import logging +import pytest +from click.testing import CliRunner +from ultralytics.utils import LOGGER +from yaltai.cli.yaltai import yaltai_cli +from kraken.lib.xml import XMLPage + + +@pytest.fixture(scope='function') +def custom_logger(): + class CustomLoggingHandler(logging.Handler): + def __init__(self): + super().__init__() + # A list to store log records (level and message) + self.records = [] + + def emit(self, record): + # Append a tuple of the log level and the message to the records list + self.records.append((record.levelname, record.getMessage())) + + def clear(self): + self.records = [] + + # Create a logger + logger = LOGGER + logger.setLevel(logging.DEBUG) # Set the logger to handle all log levels + + # Create and add the custom handler to the logger + custom_handler = CustomLoggingHandler() + logger.addHandler(custom_handler) + + # Clear the records list before each test + custom_handler.records.clear() + + # Provide both the logger and handler for access during tests + yield custom_handler + + custom_handler.records.clear() + + # Optional: Remove the handler after the test to prevent interference + logger.removeHandler(custom_handler) + + +def test_yaltai_single_alto_to_xml(custom_logger): + """Ensures that we can convert to YOLO format""" + runner = CliRunner() + + # Trigger a warning. + result = runner.invoke( + yaltai_cli, + [ + "kraken", + "--alto", + "-i", + os.path.join(os.path.abspath(os.path.dirname(__file__)), "test_files", "page1.jpg"), + os.path.join(os.path.abspath(os.path.dirname(__file__)), "test_files", "page1.xml"), + "segment", + "-y", + os.path.join(os.path.abspath(os.path.dirname(__file__)), "nano-yolo-ladas.pt") + ] + ) + print(result.output) + assert result.exit_code == 0 + assert "page1.jpg: 640x352 2 GraphicZones, 1 MainZone-P-Continued, 1 MainZone-Sp, 2 QuireMarksZones" in "\n".join([ + record[1] + for record in custom_logger.records + ]) + + page = XMLPage(os.path.join(os.path.abspath(os.path.dirname(__file__)), "test_files", "page1.xml")) + assert { + region_type: [region.boundary for region in regions] + for region_type, regions in page.regions.items() + } == { + 'GraphicZone': [[(614.0, 12.0), + (2614.0, 12.0), + (2614.0, 819.0), + (614.0, 819.0), + (614.0, 12.0)], + [(720.0, 0.0), + (2634.0, 0.0), + (2634.0, 343.0), + (720.0, 343.0), + (720.0, 0.0)]], + 'MainZone-P-Continued': [[(122.0, 1536.0), + (2218.0, 1536.0), + (2218.0, 2888.0), + (122.0, 2888.0), + (122.0, 1536.0)]], + 'MainZone-Sp': [[(95.0, 2974.0), + (2201.0, 2974.0), + (2201.0, 4854.0), + (95.0, 4854.0), + (95.0, 2974.0)]], + 'QuireMarksZone': [[(1617.0, 4877.0), + (1825.0, 4877.0), + (1825.0, 4980.0), + (1617.0, 4980.0), + (1617.0, 4877.0)], + [(1531.0, 4841.0), + (1814.0, 4841.0), + (1814.0, 4983.0), + (1531.0, 4983.0), + (1531.0, 4841.0)]] + } + + assert len([line.baseline for line in page.lines.values()]) + # ToDo: Add a test to check for line being part of regions diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..773a760 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,60 @@ +import os +import numpy as np +from yaltai.utils import read_labelmap, parse_box_labels, XYXY + + +def test_read_labelmap(): + """Asserts that reading a label map works""" + labels = read_labelmap(os.path.join( + os.path.dirname(__file__), + "test_files", + "label_map.txt" + )) + assert labels == ["Class0", "Class1", "Stuff"] + + +def test_read_files(): + """Asserts that parsing COCO/YOLO formats work""" + annots, arrays = parse_box_labels([ + os.path.join( + os.path.dirname(__file__), + "test_files", + "annot1.txt" + ), + os.path.join( + os.path.dirname(__file__), + "test_files", + "annot2.txt" + ) + ]) + assert annots == { + 'boxes': [ + XYXY(x0=15, y0=33, x1=82, y1=41), + XYXY(x0=15, y0=15, x1=83, y1=31), + XYXY(x0=35, y0=11, x1=62, y1=14), + XYXY(x0=79, y0=12, x1=82, y1=15), + XYXY(x0=16, y0=43, x1=83, y1=74), + XYXY(x0=16, y0=77, x1=85, y1=88), + XYXY(x0=84, y0=37, x1=99, y1=56), + XYXY(x0=16, y0=25, x1=81, y1=66), + XYXY(x0=17, y0=66, x1=81, y1=89), + XYXY(x0=77, y0=10, x1=80, y1=13), + XYXY(x0=26, y0=10, x1=70, y1=13), + XYXY(x0=16, y0=15, x1=80, y1=25), + XYXY(x0=35, y0=13, x1=60, y1=15) + ], + 'labels': [14, 14, 35, 31, 14, 14, 1, 23, 23, 31, 35, 24, 28] + } + assert (arrays[0] == np.array([[15, 33, 82, 41, 14, 0, 0], + [15, 15, 83, 31, 14, 0, 0], + [35, 11, 62, 14, 35, 0, 0], + [79, 12, 82, 15, 31, 0, 0], + [16, 43, 83, 74, 14, 0, 0], + [16, 77, 85, 88, 14, 0, 0], + [84, 37, 99, 56, 1, 0, 0]])).all() + assert (arrays[1] == np.array([[16, 25, 81, 66, 23, 0, 0], + [17, 66, 81, 89, 23, 0, 0], + [77, 10, 80, 13, 31, 0, 0], + [26, 10, 70, 13, 35, 0, 0], + [16, 15, 80, 25, 24, 0, 0], + [35, 13, 60, 15, 28, 0, 0]])).all() diff --git a/yaltai/cli/__init__.py b/yaltai/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/yaltai/kraken_yaltai.py b/yaltai/cli/krakn.py similarity index 61% rename from yaltai/kraken_yaltai.py rename to yaltai/cli/krakn.py index bc844a0..4595673 100644 --- a/yaltai/kraken_yaltai.py +++ b/yaltai/cli/krakn.py @@ -1,25 +1,23 @@ import click import os import dataclasses -from typing import Dict, IO, Any, cast +from typing import cast from kraken.kraken import ( # Constants SEGMENTATION_DEFAULT_MODEL, # CLI Stuff - message, logger, log, - # Logics + message, logger, # Logics get_input_parser, partial ) from PIL import Image -from kraken.lib.progress import KrakenProgressBar from kraken.containers import Segmentation from ultralytics import YOLO def segmenter(model, text_direction, mask, device, yolo_model, ignore_lines, deskew, max_angle, input, output) -> None: import json - import yaltai.kraken_adapter - import yaltai.yolo_adapter + import yaltai.models.krakn + import yaltai.models.yolo ctx = click.get_current_context() @@ -44,11 +42,11 @@ def segmenter(model, text_direction, mask, device, yolo_model, ignore_lines, des message(f'Segmenting {ctx.meta["orig_file"]}\t', nl=False) try: - regions = yaltai.yolo_adapter.segment( + regions = yaltai.models.yolo.segment( yolo_model, input=input, apply_deskew=deskew, max_angle=max_angle ) - res: Segmentation = yaltai.kraken_adapter.segment( + res: Segmentation = yaltai.models.krakn.segment( im, text_direction, mask=mask, model=model, device=device, regions=regions, ignore_lignes=ignore_lines, raise_on_error=ctx.meta['raise_failed'], autocast=ctx.meta["autocast"] @@ -79,58 +77,11 @@ def segmenter(model, text_direction, mask, device, yolo_model, ignore_lines, des json.dump(dataclasses.asdict(res), fp) message('\u2713', fg='green') -# -# @click.group(chain=True) -# @click.version_option() -# @click.option('-i', '--input', -# type=(click.Path(exists=True), # type: ignore -# click.Path(writable=True)), -# multiple=True, -# help='Input-output file pairs. Each input file (first argument) is mapped to one ' -# 'output file (second argument), e.g. `-i input.png output.txt`') -# @click.option('-f', '--format-type', type=click.Choice(['image', 'pdf']), default='image', -# help='Sets the default input type. In image mode inputs are image ' -# 'files, pdf ' -# 'expects PDF files with numbered suffixes added to output file ' -# 'names as needed.') -# @click.option('-I', '--batch-input', multiple=True, help='Glob expression to add multiple files at once.') -# @click.option('-r', '--raise-on-error/--no-raise-on-error', default=False, show_default=True, -# help='Raises the exception that caused processing to fail in the case of an error') -# @click.option('-v', '--verbose', default=0, count=True, show_default=True) -# @click.option('-d', '--device', default='cpu', show_default=True, -# help='Select device to use (cpu, cuda:0, cuda:1, ...)') -# @click.option('-o', '--suffix', default='', show_default=True, -# help='Suffix for output files from batch and PDF inputs.') -# @click.option('-p', '--pdf-format', default='{src}_{idx:06d}', -# show_default=True, -# help='Format for output of PDF files. valid fields ' -# 'are `src` (source file), `idx` (page number), and `uuid` (v4 uuid). ' -# '`-o` suffixes are appended to this format string.') -# def cli(device, input, batch_input, raise_on_error, format_type, verbose, suffix, pdf_format): -# """ YALTAi is built as a group of command but only takes one command at the time: segment """ -# ctx = click.get_current_context() -# if device != 'cpu': -# import torch -# try: -# torch.ones(1, device=device) -# except AssertionError as e: -# if raise_on_error: -# raise -# logger.error(f'Device {device} not available: {e.args[0]}.') -# ctx.exit(1) -# ctx.meta['device'] = device -# ctx.meta['input_format_type'] = format_type if format_type != 'pdf' else 'image' -# ctx.meta['raise_failed'] = raise_on_error -# ctx.meta['output_mode'] = "alto" # Unlike Kraken, forces ALTO -# ctx.meta['verbose'] = verbose -# ctx.meta['steps'] = [] -# log.set_logger(logger, level=30 - min(10 * verbose, 20)) - - -from kraken.kraken import cli - - -@cli.command('segment') + +from kraken.kraken import cli as kcli + + +@kcli.command('segment') @click.pass_context @click.option('-i', '--model', default=None, @@ -166,7 +117,7 @@ def yaltai_segment(ctx, model, text_direction, mask, yolo, ignore_lines, deskew, if not model: model = SEGMENTATION_DEFAULT_MODEL if not yolo: - raise Exception("No YOLOv5 model given") + raise Exception("No YOLOv8 model given") ctx.meta['steps'].append({'category': 'processing', 'description': 'Baseline and region segmentation', 'settings': {'model': os.path.basename(model), @@ -189,7 +140,3 @@ def yaltai_segment(ctx, model, text_direction, mask, yolo, ignore_lines, deskew, yolo.to(ctx.meta["device"]) return partial(segmenter, model, text_direction, mask, ctx.meta['device'], yolo, ignore_lines, deskew, max_angle) - - -if __name__ == "__main__": - cli() diff --git a/yaltai/yaltai.py b/yaltai/cli/yaltai.py similarity index 91% rename from yaltai/yaltai.py rename to yaltai/cli/yaltai.py index 61b2eb4..738bdb3 100644 --- a/yaltai/yaltai.py +++ b/yaltai/cli/yaltai.py @@ -17,26 +17,26 @@ from PIL import Image import click import yaml -from kraken.kraken import message -#from kraken.lib.xml import parse_xml import tabulate -from yaltai.converter import AltoToYoloZone, parse_box_labels, read_labelmap, YoloV5Zone -from mean_average_precision import MetricBuilder - -def _read_manifest(manifest: Union[str, Path]) -> List[str]: - with open(manifest, 'r') as f: - out = [x.strip() for x in f.read().splitlines() if x.strip()] - return out +from kraken.kraken import message +from kraken.lib.xml import XMLPage +from yaltai.utils import AltoToYoloZone, parse_box_labels, read_labelmap, YoloV5Zone, read_manifest +from mean_average_precision import MetricBuilder @click.group() -def cli(): +def yaltai_cli(): """ `yaltai` commands provides conversion options """ -@cli.command("alto-to-yolo") +@yaltai_cli.group("convert") +def convert(): + """ Converts formats to various other formats """ + + +@convert.command("alto-to-yolo") @click.argument("input", type=click.Path(exists=True, dir_okay=False, file_okay=True), nargs=-1) @click.argument("output", type=click.Path(dir_okay=True, file_okay=False)) @click.option("--single-class", type=str, default=None, @@ -58,7 +58,7 @@ def cli(): help="Copy images when converting ALTO to YOLOv5") @click.option("--line-as-region", type=str, multiple=True, help="Line that should be added for zone detection") -def convert( +def alto_to_yolo( input: Optional[List[click.Path]], output: click.Path, single_class: Optional[str], @@ -79,16 +79,16 @@ def convert( if manifest: message("Using single manifest", fg="blue") - input_paths = _read_manifest(manifest) + input_paths = read_manifest(manifest) elif train and val: message("Using train and validation manifests", fg="blue") - train = _read_manifest(train) - val = _read_manifest(val) + train = read_manifest(train) + val = read_manifest(val) val_idx = len(train) input_paths = train + val else: - input_paths = input message(f"Using list of inputs.", fg="blue") + input_paths = input message(f"Found {len(input_paths)} to convert.", fg="blue") @@ -133,17 +133,18 @@ def map_zones(zone_type: str) -> str: # Count Zones for idx, file in tqdm(enumerate(input_paths)): - parsed = parse_xml(file) - image_path: Path = parsed["image"] - regions = parsed["regions"] + parsed = XMLPage(file) + image_path: Path = parsed.imagename + # We record each region identifier and map the region if required + regions = parsed.regions for region in regions: if map_zones(region) not in Zones: Zones.append(map_zones(region)) processed_lines: List[Dict] = [] - if line_as_region: - for line in parsed["lines"]: + if line_as_region: # ToDo: Adapt to new system + for line in parsed.lines: if line.get("tags", {}).get("type") in line_as_region: line_type = line["tags"]["type"] if line_type not in Zones: @@ -158,10 +159,10 @@ def map_zones(zone_type: str) -> str: local_file: List[AltoToYoloZone] = [] for region, examples in regions.items(): region_id = Zones.index(map_zones(region)) - for box in examples: + for region_obj in examples: local_file.append( AltoToYoloZone( - BOX=box, + BOX=region_obj.boundary, PAGE_WIDTH=width, PAGE_HEIGHT=height, tag=region_id @@ -242,7 +243,7 @@ def map_zones(zone_type: str) -> str: message(f"\t- {cnt:05} {zone}", fg='blue') -@cli.command("scores") +@yaltai_cli.command("scores") @click.argument("gt-directory", type=click.Path(exists=True, dir_okay=True, file_okay=False)) @click.argument("pred-directory", type=click.Path(dir_okay=True, file_okay=False, exists=True)) @click.option("-t", "--threshold", type=float, help="IoU Threshold", default=.5, show_default=True) @@ -310,7 +311,7 @@ def reclass_classes(array_list: List[np.array]) -> None: }, save_json) -@cli.command("yolo-to-alto") +@convert.command("yolo-to-alto") @click.argument("input", type=click.Path(exists=True, dir_okay=False, file_okay=True), nargs=-1) @click.option("-l", "--labelmap", type=click.Path(exists=True, file_okay=True, dir_okay=False), help="Format for the score table", default=None, show_default=True) @@ -327,19 +328,19 @@ def yolo_to_alto(input, labelmap): for idx, zone in enumerate(labelmap) ]) - with open(os.path.join(os.path.dirname(__file__), "template.xml")) as f: + with open(os.path.join(os.path.dirname(__file__), "../template.xml")) as f: TEMPLATE = f.read() for file in input: xml_name = file[:-4] + ".xml" img_file_name = os.path.basename(file[:-4]) + ".jpg" zones = [] - if os.path.exists(os.path.join(os.path.dirname(file), "..", "images", img_file_name)): - img_name = os.path.join(os.path.dirname(file), "..", "images", img_file_name) + if os.path.exists(os.path.join(os.path.dirname(file), "../..", "images", img_file_name)): + img_name = os.path.join(os.path.dirname(file), "../..", "images", img_file_name) img_for_xml_name = f"../images/{img_file_name}" - elif os.path.exists(os.path.join(os.path.dirname(file), "..", "images", img_file_name)): - img_name = os.path.join(os.path.dirname(file), "..", "images", img_file_name) - img_for_xml_name = os.path.join("..", img_file_name) + elif os.path.exists(os.path.join(os.path.dirname(file), "../..", "images", img_file_name)): + img_name = os.path.join(os.path.dirname(file), "../..", "images", img_file_name) + img_for_xml_name = os.path.join("../..", img_file_name) else: message(f"Can't find the image for {img_file_name}") sys.exit(0) @@ -379,10 +380,10 @@ def yolo_to_alto(input, labelmap): ) -import yaltai.kraken_yaltai as kyaltai +from yaltai.cli.krakn import kcli -cli.add_command(kyaltai.cli, "kraken") +yaltai_cli.add_command(kcli, "kraken") if __name__ == "__main__": - cli() + yaltai_cli() diff --git a/yaltai/models/__init__.py b/yaltai/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/yaltai/kraken_adapter.py b/yaltai/models/krakn.py similarity index 97% rename from yaltai/kraken_adapter.py rename to yaltai/models/krakn.py index 16577b8..7ccea6f 100644 --- a/yaltai/kraken_adapter.py +++ b/yaltai/models/krakn.py @@ -161,6 +161,11 @@ def segment(im: PIL.Image.Image, lines.extend(_lines) + # Rounding ! + for reg_class in regions: + for reg_obj in regions[reg_class]: + reg_obj.boundary = list(map(lambda x: list(map(round, x)), reg_obj.boundary)) + if len(rets['cls_map']['baselines']) > 1: script_detection = True else: diff --git a/yaltai/yolo_adapter.py b/yaltai/models/yolo.py similarity index 100% rename from yaltai/yolo_adapter.py rename to yaltai/models/yolo.py diff --git a/yaltai/preprocessing.py b/yaltai/preprocessing.py index 74156b1..567fcdf 100644 --- a/yaltai/preprocessing.py +++ b/yaltai/preprocessing.py @@ -19,14 +19,14 @@ def rotatebox(bbox: List[List[int]], image: Image.Image, angle: float): rotated_bbox = [] for i, coord in enumerate(bbox): - rot_matrix = cv2.getRotationMatrix2D((image_center_x, image_center_y), angle, 1.0) - cosinus, sinus = abs(rot_matrix[0, 0]), abs(rot_matrix[0, 1]) - new_width = int((height * sinus) + (width * cosinus)) - new_height = int((height * cosinus) + (width * sinus)) - rot_matrix[0, 2] += (new_width / 2) - image_center_x - rot_matrix[1, 2] += (new_height / 2) - image_center_y - v = [coord[0], coord[1], 1] # ? - adjusted_coord = np.dot(rot_matrix, v) - rotated_bbox.append((int(adjusted_coord[0]), int(adjusted_coord[1]))) + rot_matrix = cv2.getRotationMatrix2D((image_center_x, image_center_y), angle, 1.0) + cosinus, sinus = abs(rot_matrix[0, 0]), abs(rot_matrix[0, 1]) + new_width = int((height * sinus) + (width * cosinus)) + new_height = int((height * cosinus) + (width * sinus)) + rot_matrix[0, 2] += (new_width / 2) - image_center_x + rot_matrix[1, 2] += (new_height / 2) - image_center_y + v = [coord[0], coord[1], 1] # ? + adjusted_coord = np.dot(rot_matrix, v) + rotated_bbox.append((int(adjusted_coord[0]), int(adjusted_coord[1]))) return rotated_bbox diff --git a/yaltai/converter.py b/yaltai/utils.py similarity index 80% rename from yaltai/converter.py rename to yaltai/utils.py index 355e9dc..2ca3347 100644 --- a/yaltai/converter.py +++ b/yaltai/utils.py @@ -1,6 +1,7 @@ from dataclasses import dataclass from collections import namedtuple from typing import List, Tuple, Optional, Dict, Union +from pathlib import Path import numpy as np XYXY = namedtuple("XYXY", ["x0", "y0", "x1", "y1"]) @@ -89,6 +90,13 @@ def parse_box_labels( files: List[str], gt: bool = True ) -> Tuple[Dict[str, Union[List[Union[XYXY, int]]]], List[np.array]]: + """Parse a list of YOLO/COCO BB annotation files + + This function is only used to compute metrics + + :param files: List of file path in YOLO / COCO formats + :param gt: If data are ground Truth of if they are predicted (I don't remember why) + """ parsed = {"boxes": [], "labels": []} arrays = [] for file in sorted(files): @@ -109,7 +117,16 @@ def parse_box_labels( def read_labelmap(path: str) -> List[str]: - lines = [] + """ Reads a labelmap YAML file and parses the classes + """ with open(path) as f: - lines = f.read().split() + lines = [line.strip() for line in f.read().split() if line.strip()] return lines + + +def read_manifest(manifest: Union[str, Path]) -> List[str]: + """ Reads a manifest (for training purposes ?) + """ + with open(manifest, 'r') as f: + out = [x.strip() for x in f.read().splitlines() if x.strip()] + return out