diff --git a/.flake8 b/.flake8 deleted file mode 100644 index 1663db5..0000000 --- a/.flake8 +++ /dev/null @@ -1,3 +0,0 @@ -[flake8] -max-line-length = 100 -extend-ignore = E203,E501 diff --git a/.gitattributes b/.gitattributes index d3574bc..cd1dbe1 100644 --- a/.gitattributes +++ b/.gitattributes @@ -2,3 +2,4 @@ examples/example3.wav filter=lfs diff=lfs merge=lfs -text examples/example4.wav filter=lfs diff=lfs merge=lfs -text examples/example1.wav filter=lfs diff=lfs merge=lfs -text examples/example2.wav filter=lfs diff=lfs merge=lfs -text +batbot/classifier/models/onnx/batbot.mobilenet.9dc57ea3.onnx filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/python-publish.yaml b/.github/workflows/python-publish.yaml index 2cbc3f0..1afdf74 100644 --- a/.github/workflows/python-publish.yaml +++ b/.github/workflows/python-publish.yaml @@ -1,4 +1,4 @@ -name: Wheel +name: Python distributions on: pull_request: @@ -9,43 +9,10 @@ on: - main tags: - v* - schedule: - - cron: '0 16 * * *' # Every day at 16:00 UTC (~09:00 PT) jobs: - build_wheels: - name: Build on ${{ matrix.os }} - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, macos-latest] - python-version: [3.12] - - steps: - - name: Checkout code - uses: nschloe/action-cached-lfs-checkout@v1.2.3 - with: - exclude: "examples/example[2-4].wav" - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v6 - with: - python-version: ${{ matrix.python-version }} - - - name: Build wheel - run: | - pip install --upgrade pip - pip install build - python -m build --wheel --outdir dist/ . - - - uses: actions/upload-artifact@v6 - with: - name: artifact-wheel-${{ matrix.os }}-${{ matrix.python-version }} - path: ./dist/*.whl - - build_sdist: - name: Build source distribution + build: + name: Build wheel and source distribution runs-on: ubuntu-latest steps: - name: Checkout code @@ -53,79 +20,87 @@ jobs: with: exclude: "examples/example[2-4].wav" - - name: Set up Python 3.12 + - name: Set up Python uses: actions/setup-python@v6 with: - python-version: '3.12' + python-version: '3.11' + cache: pip - - name: Build sdist + - name: Build and validate distributions run: | - pip install --upgrade pip - pip install build - python -m build --sdist --outdir dist/ . + python -m pip install --upgrade pip + python -m pip install build twine + python -m build --outdir dist/ . + python -m twine check dist/* - - uses: actions/upload-artifact@v6 + - name: Upload distributions + uses: actions/upload-artifact@v6 with: - name: artifact-sdist - path: ./dist/*.tar.gz + name: python-distributions + path: dist/ + if-no-files-found: error - test_wheel: - needs: [build_wheels, build_sdist] + test-wheel: + name: Test the installed wheel and bundled model + needs: build runs-on: ubuntu-latest - - # test wheel - if: github.event_name == 'push' steps: - - name: Set up Python 3.12 + - name: Set up Python uses: actions/setup-python@v6 with: - python-version: '3.12' + python-version: '3.11' + cache: pip - - uses: actions/download-artifact@v7 + - name: Download distributions + uses: actions/download-artifact@v7 with: - path: artifact + name: python-distributions + path: dist - name: Install wheel run: | - pip install --upgrade pip - pip install wheel - find . - mkdir dist - cp artifact/*-ubuntu-*/*.whl dist/ - cp artifact/*/*.tar.gz dist/ - pip install dist/*.whl + python -m pip install --upgrade pip + python -m pip install dist/*.whl - - name: Test module + - name: Smoke test wheel outside the checkout + working-directory: ${{ runner.temp }} run: | - python -c "import batbot;" - - # - name: Test CLI - # run: | - # batbot example - - upload_pypi: - needs: [test_wheel] - runs-on: ubuntu-latest - # upload to PyPI on every tag starting with 'v' + python -I -c "import batbot; assert batbot.__version__" + batbot --help + python -I - <<'PY' + import tempfile + from pathlib import Path + + import cv2 + import numpy as np + + from batbot import classifier + + with tempfile.TemporaryDirectory() as directory: + image = Path(directory) / 'smoke.png' + cv2.imwrite(str(image), np.zeros((300, 700, 3), dtype=np.uint8)) + result = classifier.Classifier().classify(image)[0] + assert result['label'] in classifier.CLASSES + assert len(result['scores']) == len(classifier.CLASSES) + PY + + publish: + name: Publish to PyPI + needs: test-wheel if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v') + runs-on: ubuntu-latest environment: name: pypi url: https://pypi.org/p/batbot permissions: id-token: write + contents: read steps: - - uses: actions/download-artifact@v7 + - name: Download distributions + uses: actions/download-artifact@v7 with: - path: artifact + name: python-distributions + path: dist - - name: Install wheel - run: | - find . - mkdir dist - cp artifact/*-ubuntu-*/*.whl dist/ - cp artifact/*/*.tar.gz dist/ - - - name: Publish package distributions to PyPI + - name: Publish distributions with trusted publishing uses: pypa/gh-action-pypi-publish@release/v1 - with: - password: ${{ secrets.BATBOT_PYPI_TOKEN }} diff --git a/.github/workflows/testing.yaml b/.github/workflows/testing.yaml index 7644c76..1aa855b 100644 --- a/.github/workflows/testing.yaml +++ b/.github/workflows/testing.yaml @@ -1,9 +1,10 @@ -# This workflow will install Python dependencies, run tests and lint with multiple versions of Python -# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions - name: Testing -on: push +on: + push: + pull_request: + branches: + - main jobs: test: @@ -11,14 +12,13 @@ jobs: strategy: fail-fast: false matrix: - python-version: ['3.9', '3.10', '3.11', '3.12'] + python-version: ['3.11', '3.12', '3.13', '3.14'] env: OS: ubuntu-latest PYTHON: ${{ matrix.python-version }} steps: - # Checkout and env setup - name: Checkout code uses: nschloe/action-cached-lfs-checkout@v1.2.3 with: @@ -28,25 +28,21 @@ jobs: uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} + cache: pip - - name: Install dependencies + - name: Install project and test tools run: | python -m pip install --upgrade pip - pip install -r requirements/runtime.txt - pip install -r requirements/optional.txt - pip install -e . + python -m pip install -e ".[test]" pre-commit - name: Check with pre-commit - run: | - SKIP=hadolint pre-commit - SKIP=hadolint pre-commit run --all-files + run: SKIP=hadolint pre-commit run --all-files - name: Run tests and coverage - run: | - set -ex - pytest --cov=batbot --cov-append --cov-report=xml --random-order-seed=1 + run: pytest --cov-report=xml --random-order-seed=1 - name: Upload coverage to Codecov + if: matrix.python-version == '3.11' continue-on-error: true uses: codecov/codecov-action@v5 with: diff --git a/.gitignore b/.gitignore index bc5bf9f..edfa416 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ example*.json .vscode/* assets/*.key +assets/*.png diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ea03ab9..26b38ba 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,46 +2,46 @@ # See https://pre-commit.com/hooks.html for more hooks repos: - repo: https://github.com/hadolint/hadolint - rev: v2.13.1-beta + rev: v2.15.1 hooks: - id: hadolint name: Hadolint for Dockerfiles - repo: https://github.com/asottile/pyupgrade - rev: v3.20.0 + rev: v3.21.2 hooks: - id: pyupgrade name: pyupgrade description: Run PyUpgrade on Python code. + args: [--py311-plus] - repo: https://github.com/pycqa/isort - rev: 6.0.1 + rev: 9.0.0b5 hooks: - id: isort - args: [--settings-path setup.cfg] + args: [--settings-path, pyproject.toml] name: isort description: Run import sorting (isort) on Python code. - repo: https://github.com/psf/black - rev: 25.1.0 + rev: 26.5.1 hooks: - id: black name: Black for Python code formatting language_version: python3 - args: ["--skip-string-normalization", "--target-version", "py310", "--line-length", "100"] - repo: https://github.com/pycqa/flake8 - rev: 7.2.0 + rev: 7.3.0 hooks: - id: flake8 name: Flake8 for Python code linting - args: ["--config", ".flake8"] + additional_dependencies: [flake8-pyproject==1.2.4] - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: double-quote-string-fixer name: Format single quotes - - id: requirements-txt-fixer - name: Format requirements.txt - id: check-yaml name: Format YAML files args: ['--unsafe'] + - id: check-toml + name: Format TOML files - id: trailing-whitespace name: Fix Whitespace - id: mixed-line-ending diff --git a/.readthedocs.yaml b/.readthedocs.yaml index c156481..8d19bb2 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -7,9 +7,9 @@ version: 2 # Set the version of Python and other tools you might need build: - os: ubuntu-22.04 + os: ubuntu-24.04 tools: - python: "3.10" + python: "3.11" # Build documentation in the docs/ directory with Sphinx sphinx: @@ -18,6 +18,7 @@ sphinx: # Optionally declare the Python requirements required to build your docs python: install: - - requirements: requirements/documentation.txt - method: pip path: . + extra_requirements: + - docs diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 86f0f4e..951e1bb 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,7 +1,13 @@ -========= Changelog ========= +Version 0.2.0 +------------- + +* Add MobileNet ONNX species classification for spectrogram images and WAV files. +* Add ``classify``, ``classify-wav``, and ``classify-bulk`` CLI commands. +* Add mirrored model retrieval, bulk species summaries, and a performance plotting example. + TODO ---- diff --git a/Dockerfile b/Dockerfile index 17f9ffb..d2be55c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -26,11 +26,8 @@ WORKDIR /code COPY ./ /code -RUN python3 -m venv /venv - # hadolint ignore=DL3003,DL3013 -RUN /venv/bin/pip install --no-cache-dir -r requirements/runtime.txt \ - && /venv/bin/pip install --no-cache-dir -r requirements/optional.txt \ +RUN python3 -m venv /venv \ && /venv/bin/pip install --no-cache-dir -e . # && if [ "$(uname -m)" != "aarch64" ] \ # ; then \ diff --git a/ISSUES.rst b/ISSUES.rst index 078a561..f5ceff6 100644 --- a/ISSUES.rst +++ b/ISSUES.rst @@ -1,4 +1,3 @@ -============ Known Issues ============ diff --git a/MANIFEST.in b/MANIFEST.in index e38272b..23195f6 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -6,3 +6,7 @@ include LICENSE # Include examples files for testing include examples/example1.wav +include examples/plot_classifier_performance.py + +# Include the classifier model in source distributions +recursive-include batbot/classifier/models/onnx *.onnx diff --git a/README.rst b/README.rst index d12579b..987e492 100644 --- a/README.rst +++ b/README.rst @@ -104,6 +104,8 @@ Here are the steps for extracting the compressed spectrogram: How to Install -------------- +BatBot requires Python 3.11 or newer. + .. code-block:: bash pip install batbot @@ -126,6 +128,61 @@ To then add GPU acceleration, you need to replace `onnxruntime` with `onnxruntim How to Run ---------- +Species Classification +~~~~~~~~~~~~~~~~~~~~~~ + +BatBot includes a 35-class MobileNet ONNX model for classifying spectrograms. +The model is used from the package when available and can also be downloaded +from its Kitware Data mirror with ``pooch``: + +.. code-block:: bash + + batbot fetch + batbot fetch --pull + +Classify an existing spectrogram, classify a WAV through BatBot's built-in +spectrogram step, or recursively process a large directory: + +.. code-block:: bash + + batbot classify recording.jpg + batbot classify-wav recording.wav --output recording.json + batbot classify-bulk ./recordings --input-type wav --num-workers 4 --output results.json + +Bulk JSON includes every per-file prediction plus ``label_counts``, +``species_counts`` (which excludes ``NOISE``), the noise count, failures, and +mean confidence. Generated WAV spectrograms are temporary by default; pass +``--spectrogram-dir ./spectrograms`` to retain them. +``--num-workers`` runs multiple spectrogram inference jobs concurrently while +sharing one validated ONNX Runtime session; results retain input order. + +The corresponding Python API follows the Scoutbot WIC ``pre`` / ``predict`` / +``post`` pattern. A convenience call is usually sufficient: + +.. code-block:: python + + from batbot import classifier + + result = classifier.classify('recording.jpg')[0] + wav_result = classifier.classify_wav('recording.wav') + folder = classifier.classify_bulk( + ['./recordings'], input_type='wav', num_workers=4 + ) + +To evaluate a labeled dataset arranged as ``LABEL/*.wav``, install +``scikit-learn`` and run the included performance example: + +.. code-block:: bash + + pip install "batbot[performance]" + python examples/plot_classifier_performance.py ./validation \ + --cache predictions.json --output performance.png + +The plot includes count, true-normalized, and prediction-normalized confusion +matrices, top-k accuracy, Matthews correlation, and ``NOISE`` precision-recall +and ROC diagnostics when that label is present. Species codes are reordered by +genus before plotting so the shaded error regions remain contiguous. + You can run the Gradio demo with: .. code-block:: bash @@ -212,11 +269,11 @@ PyPI To upload the latest BatBot version to the Python Package Index (PyPI), follow the steps below: -#. Edit ``batbot/__init__.py:65`` and set ``VERSION`` to the desired version +#. Edit ``batbot/_version.py`` and set ``__version__`` to the desired version .. code-block:: python - VERSION = 'X.Y.Z' + __version__ = 'X.Y.Z' #. Push any changes and version update to the ``main`` branch on GitHub and wait for CI tests to pass @@ -246,7 +303,7 @@ You can run the automated tests in the ``tests/`` folder by running: .. code-block:: bash - pip install -r requirements/optional.txt + pip install -e ".[test]" pytest You may also get a coverage percentage by running: @@ -265,16 +322,15 @@ There is Sphinx documentation in the ``docs/`` folder, which can be built by run .. code-block:: bash cd docs/ - pip install -r requirements/optional.txt + pip install -e "..[docs]" sphinx-build -M html . build/ Logging ------- -The script uses Python's built-in logging functionality called ``logging``. All print functions are replaced with ``log.info()``, which sends the output to two places: - -#. the terminal window, and -#. the file `batbot.log` +BatBot uses Python's standard ``logging`` package and installs a ``NullHandler`` +for library use. Applications can configure the ``batbot`` logger themselves or +call ``batbot.utils.init_logging()`` for Rich console and rotating-file output. Code Formatting --------------- @@ -286,10 +342,14 @@ Reference `pre-commit's installation instructions `_. Furthermore, try to conform to ``PEP8``. You should set up your preferred editor to use ``flake8`` as its Python linter, but pre-commit will ensure compliance before a git commit is completed. This will use the ``flake8`` configuration within ``setup.cfg``, which ignores several errors and stylistic considerations. See the ``setup.cfg`` file for a full and accurate listing of stylistic codes to ignore. +The code base is formatted by `Black `_ +and linted by Flake8. Black, isort, Flake8, pytest, coverage, packaging, and +dependency settings are centralized in ``pyproject.toml``. The +``flake8-pyproject`` plugin lets the unchanged Flake8 command consume that +configuration. .. |Tests| image:: https://github.com/Kitware/batbot/actions/workflows/testing.yaml/badge.svg?branch=main diff --git a/batbot/__init__.py b/batbot/__init__.py index fd6e2a4..8b36f5b 100644 --- a/batbot/__init__.py +++ b/batbot/__init__.py @@ -1,321 +1,40 @@ -""" -The above components must be run in the correct order, but BatBot also offers a processing pipeline. - -The machine learning (ML) model can be pre-downloaded and fetched by a single call to -:func:`batbot.fetch` and the unified pipeline can be run by the function :func:`batbot.pipeline`. -Below is example code for how these components interact. - -Furthermore, there is an application demo file (``app.py``) that shows how the entire pipeline can -be run on WAV files. - -.. code-block:: python +"""Public package interface for BatBot. - # Get WAV filepath - filepath = '/path/to/file.wav' - - # Run tiling - output_paths, metadata_path, metadata = spectrogram.compute(filepath) +Importing :mod:`batbot` intentionally avoids importing the scientific and ONNX +stacks. The larger submodules are loaded only when their APIs are used. """ -import concurrent.futures -from multiprocessing import Manager -from os.path import basename, exists, join, splitext -from pathlib import Path - -import pooch -from tqdm import tqdm - -from batbot import utils - -log = utils.init_logging() -QUIET = not utils.VERBOSE - - -from batbot import spectrogram # NOQA - -VERSION = '0.1.5' -version = VERSION -__version__ = VERSION - -PWD = Path(__file__).absolute().parent.parent - - -def fetch(pull=False, config=None): - """ - Fetch the Classifier ONNX model file from a CDN if it does not exist locally. - - This function will throw an AssertionError if the download fails or the - file otherwise does not exist locally on disk. - - Args: - pull (bool, optional): If :obj:`True`, force using the downloaded version - stored in the local system's cache. Defaults to :obj:`False`. - config (str or None, optional): the configuration to use. Defaults to :obj:`None`. - - Returns: - None - - Raises: - AssertionError: If the model cannot be fetched. - """ - raise NotImplementedError - - -def pipeline( - filepath, - out_file_stem=None, - output_folder=None, - fast_mode=False, - force_overwrite=False, - quiet=False, - plot_uncompressed_amplitude=False, - include_original_sr=False, - time_buffer_ms=1.0, - debug=False, -): - """ - Run the ML pipeline on a given WAV filepath and return the classification results - - The final output is a list of time windows where a bat exists. - Each dictionary has a structure with the following keys: - - :: - - { - 'l': class_label (str) - 'c': confidence (float) - 'x': x_top_left (float) - 'y': y_top_left (float) - 'w': width (float) - 'h': height (float) - } - - Args: - filepath (str): WAV filepath (relative or absolute) - config (str or None, optional): the configuration to use. Defaults to :obj:`None`. - classifier_thresh (float or None, optional): the confidence threshold for the classifier's - predictions. Defaults to the default configuration setting. - clean (bool, optional): a flag to clean up any on-disk spectrograms that were generated. - Defaults to :obj:`True`. - - Returns: - tuple ( float, list ( dict ) ): classifier score, list of time windows - """ - - # Generate spectrogram - output_paths, compressed_paths, metadata_path, metadata = spectrogram.compute( - filepath, - out_file_stem=out_file_stem, - output_folder=output_folder, - fast_mode=fast_mode, - force_overwrite=force_overwrite, - quiet=quiet, - plot_uncompressed_amplitude=plot_uncompressed_amplitude, - include_original_sr=include_original_sr, - time_buffer_ms=time_buffer_ms, - debug=debug, - ) - - return output_paths, compressed_paths, metadata_path - - -def pipeline_multi_wrapper( - filepaths, - out_file_stems=None, - fast_mode=False, - force_overwrite=False, - worker_position=None, - quiet=False, - tqdm_lock=None, -): - """Fault-tolerant wrapper for multiple inputs. - - Args: - filepaths (_type_): _description_ - out_file_stems (_type_, optional): _description_. Defaults to None. - fast_mode (bool, optional): _description_. Defaults to False. - force_overwrite (bool, optional): _description_. Defaults to False. - - Returns: - _type_: _description_ - """ - - if out_file_stems is not None: - assert len(filepaths) == len( - out_file_stems - ), 'Input filepaths and out_file_stems have different length.' - else: - out_file_stems = [None] * len(filepaths) - - outputs = {'output_paths': [], 'compressed_paths': [], 'metadata_paths': [], 'failed_files': []} - # print(filepaths, out_file_stems) - if tqdm_lock is not None: - tqdm.set_lock(tqdm_lock) - for in_file, out_stem in tqdm( - zip(filepaths, out_file_stems), - desc='Processing, worker {}'.format(worker_position), - position=worker_position, - total=len(filepaths), - leave=True, - ): - try: - output_paths, compressed_paths, metadata_path = pipeline( - in_file, - out_file_stem=out_stem, - fast_mode=fast_mode, - force_overwrite=force_overwrite, - quiet=quiet, - ) - outputs['output_paths'].extend(output_paths) - outputs['compressed_paths'].extend(compressed_paths) - outputs['metadata_paths'].append(metadata_path) - except Exception as e: - outputs['failed_files'].append((str(in_file), e)) - - return tuple(outputs.values()) - - -def parallel_pipeline( - in_file_chunks, - out_stem_chunks=None, - fast_mode=False, - force_overwrite=False, - num_workers=0, - threaded=False, - quiet=False, - desc=None, -): - - if out_stem_chunks is None: - out_stem_chunks = [None] * len(in_file_chunks) - - if len(in_file_chunks) == 0: - return None - else: - assert len(in_file_chunks) == len( - out_stem_chunks - ), 'in_file_chunks and out_stem_chunks must have the same length.' - - if threaded: - executor_cls = concurrent.futures.ThreadPoolExecutor - else: - executor_cls = concurrent.futures.ProcessPoolExecutor - - num_workers = min(len(in_file_chunks), num_workers) - - outputs = {'output_paths': [], 'compressed_paths': [], 'metadata_paths': [], 'failed_files': []} - - lock_manager = Manager() - tqdm_lock = lock_manager.Lock() - - with tqdm(total=len(in_file_chunks), disable=quiet, desc=desc) as progress: - with executor_cls(max_workers=num_workers) as executor: - - futures = [ - executor.submit( - pipeline_multi_wrapper, - filepaths=file_chunk, - out_file_stems=out_stem_chunk, - fast_mode=fast_mode, - force_overwrite=force_overwrite, - worker_position=index % num_workers, - quiet=quiet, - tqdm_lock=tqdm_lock, - ) - for index, (file_chunk, out_stem_chunk) in enumerate( - zip(in_file_chunks, out_stem_chunks) - ) - ] - - for future in concurrent.futures.as_completed(futures): - output_paths, compressed_paths, metadata_path, failed_files = future.result() - outputs['output_paths'].extend(output_paths) - outputs['compressed_paths'].extend(compressed_paths) - outputs['metadata_paths'].extend(metadata_path) - outputs['failed_files'].extend(failed_files) - progress.update(1) - - return tuple(outputs.values()) - - -def batch( - filepaths, - config=None, - # classifier_thresh=classifier.CONFIGS[None]['thresh'], - clean=True, -): - """ - Run the ML pipeline on a given batch of WAV filepaths and return the detections - in a corresponding list. The output is a list of outputs matching the output of - :func:`batbot.pipeline`, except the processing is done in batch and is much faster. - - The final output is a list of lists of dictionaries, each representing a - single detection. Each dictionary has a structure with the following keys: - - :: - - { - 'l': class_label (str) - 'c': confidence (float) - 'x': x_top_left (float) - 'y': y_top_left (float) - 'w': width (float) - 'h': height (float) - } - - Args: - filepaths (list): list of str WAV filepath (relative or absolute) - config (str or None, optional): the configuration to use. Defaults to :obj:`None`. - classifier_thresh (float or None, optional): the confidence threshold for the Classifier's - predictions. Defaults to the default configuration setting. - clean (bool, optional): a flag to clean up any on-disk spectrograms that were generated. - Defaults to :obj:`True`. - - Returns: - tuple ( list ( float ), list ( list ( dict ) ) : corresponding list of classifier scores, corresponding list of lists of predictions - """ - # Run tiling - batch = {} - for filepath in filepaths: - _, _, _, metadata = spectrogram.compute(filepath) - batch[filepath] = metadata - - raise NotImplementedError - - -def example(): - """ - Run the pipeline on an example WAV with the default configuration - """ - TEST_WAV = 'example1.wav' - TEST_WAV_HASH = '391efce5433d1057caddb4ce07b9712c523d6a815e4ee9e64b62973569982925' # NOQA - - wav_filepath = join(PWD, 'examples', 'example1.wav') - - if not exists(wav_filepath): - wav_filepath = pooch.retrieve( - url=f'https://raw.githubusercontent.com/Kitware/batbot/main/examples/{TEST_WAV}', - known_hash=TEST_WAV_HASH, - progressbar=True, - ) - assert exists(wav_filepath) - - log.debug(f'Running pipeline on WAV: {wav_filepath}') - - import time - - output_stem = join('output', splitext(basename(wav_filepath))[0]) - start_time = time.time() - results = pipeline( - wav_filepath, - out_file_stem=output_stem, - fast_mode=False, - force_overwrite=True, - plot_uncompressed_amplitude=True, - include_original_sr=True, - time_buffer_ms=5.0, - ) - stop_time = time.time() - print('Example pipeline completed in {} seconds.'.format(stop_time - start_time)) - - log.debug(results) +from importlib import import_module +from types import ModuleType +from typing import Any + +from batbot._config import QUIET, log +from batbot._version import VERSION, __version__ +from batbot.api import batch, example, fetch, parallel_pipeline, pipeline, pipeline_multi_wrapper + +version = __version__ + +__all__ = [ + 'QUIET', + 'VERSION', + '__version__', + 'batch', + 'classifier', + 'example', + 'fetch', + 'log', + 'parallel_pipeline', + 'pipeline', + 'pipeline_multi_wrapper', + 'spectrogram', + 'version', +] + + +def __getattr__(name: str) -> Any: + """Lazily expose the two computational subpackages.""" + if name in {'classifier', 'spectrogram'}: + module: ModuleType = import_module(f'.{name}', __name__) + globals()[name] = module + return module + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') diff --git a/batbot/_config.py b/batbot/_config.py new file mode 100644 index 0000000..86658e4 --- /dev/null +++ b/batbot/_config.py @@ -0,0 +1,10 @@ +"""Process-wide BatBot settings that are safe to import from any submodule.""" + +import logging +import os + +VERBOSE = os.getenv('BATBOT_VERBOSE', os.getenv('VERBOSE')) is not None +QUIET = not VERBOSE + +log = logging.getLogger('batbot') +log.addHandler(logging.NullHandler()) diff --git a/batbot/_version.py b/batbot/_version.py new file mode 100644 index 0000000..c158e3e --- /dev/null +++ b/batbot/_version.py @@ -0,0 +1,4 @@ +"""BatBot version information and the package metadata version source.""" + +__version__ = '0.2.0' +VERSION = __version__ diff --git a/batbot/api.py b/batbot/api.py new file mode 100644 index 0000000..8a106cb --- /dev/null +++ b/batbot/api.py @@ -0,0 +1,225 @@ +"""High-level BatBot processing APIs. + +Heavy dependencies are imported inside the functions that use them so that a +plain ``import batbot`` remains inexpensive and does not initialize logging or +create files on disk. +""" + +from __future__ import annotations + +import concurrent.futures +import time +from collections.abc import Sequence +from multiprocessing import Manager +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from batbot._config import log + +if TYPE_CHECKING: + from batbot.classifier.types import ClassificationResult + +PACKAGE_ROOT = Path(__file__).resolve().parent +PROJECT_ROOT = PACKAGE_ROOT.parent + + +def fetch(pull: bool = False, config: str | None = None) -> str: + """Return the local ONNX classifier model path.""" + from batbot.classifier import fetch as fetch_classifier + + return fetch_classifier(pull=pull, config=config) + + +def pipeline( + filepath: str | Path, + out_file_stem: str | None = None, + output_folder: str | None = None, + fast_mode: bool = False, + force_overwrite: bool = False, + quiet: bool = False, + plot_uncompressed_amplitude: bool = False, + include_original_sr: bool = False, + time_buffer_ms: float = 1.0, + debug: bool = False, +) -> tuple[list[str], list[str], str | None]: + """Generate spectrograms and metadata for one WAV file.""" + from batbot.spectrogram import compute + + output_paths, compressed_paths, metadata_path, _ = compute( + str(filepath), + out_file_stem=out_file_stem, + output_folder=output_folder, + fast_mode=fast_mode, + force_overwrite=force_overwrite, + quiet=quiet, + plot_uncompressed_amplitude=plot_uncompressed_amplitude, + include_original_sr=include_original_sr, + time_buffer_ms=time_buffer_ms, + debug=debug, + ) + return output_paths, compressed_paths, metadata_path + + +def pipeline_multi_wrapper( + filepaths: Sequence[str], + out_file_stems: Sequence[str | None] | None = None, + fast_mode: bool = False, + force_overwrite: bool = False, + worker_position: int | None = None, + quiet: bool = False, + tqdm_lock: Any = None, +) -> tuple[list[str], list[str], list[str | None], list[tuple[str, Exception]]]: + """Run :func:`pipeline` for a chunk while retaining per-file failures.""" + from tqdm import tqdm + + if out_file_stems is not None and len(filepaths) != len(out_file_stems): + raise ValueError('Input filepaths and out_file_stems have different length') + if out_file_stems is None: + out_file_stems = [None] * len(filepaths) + + output_paths: list[str] = [] + compressed_paths: list[str] = [] + metadata_paths: list[str | None] = [] + failed_files: list[tuple[str, Exception]] = [] + + if tqdm_lock is not None: + tqdm.set_lock(tqdm_lock) + for in_file, out_stem in tqdm( + zip(filepaths, out_file_stems), + desc=f'Processing, worker {worker_position}', + position=worker_position, + total=len(filepaths), + leave=True, + ): + try: + outputs, compressed, metadata = pipeline( + in_file, + out_file_stem=out_stem, + fast_mode=fast_mode, + force_overwrite=force_overwrite, + quiet=quiet, + ) + output_paths.extend(outputs) + compressed_paths.extend(compressed) + metadata_paths.append(metadata) + except Exception as error: # pragma: no cover - worker failures depend on input data + failed_files.append((str(in_file), error)) + + return output_paths, compressed_paths, metadata_paths, failed_files + + +def parallel_pipeline( + in_file_chunks: Sequence[Sequence[str]], + out_stem_chunks: Sequence[Sequence[str | None]] | None = None, + fast_mode: bool = False, + force_overwrite: bool = False, + num_workers: int = 0, + threaded: bool = False, + quiet: bool = False, + desc: str | None = None, +) -> tuple[list[str], list[str], list[str | None], list[tuple[str, Exception]]] | None: + """Run spectrogram processing chunks concurrently.""" + from tqdm import tqdm + + if not in_file_chunks: + return None + if out_stem_chunks is None: + out_stem_chunks = [[None] * len(chunk) for chunk in in_file_chunks] + if len(in_file_chunks) != len(out_stem_chunks): + raise ValueError('in_file_chunks and out_stem_chunks must have the same length') + + executor_cls = ( + concurrent.futures.ThreadPoolExecutor + if threaded + else concurrent.futures.ProcessPoolExecutor + ) + num_workers = min(len(in_file_chunks), num_workers) + if num_workers <= 0: + raise ValueError('num_workers must be positive') + + output_paths: list[str] = [] + compressed_paths: list[str] = [] + metadata_paths: list[str | None] = [] + failed_files: list[tuple[str, Exception]] = [] + + with Manager() as lock_manager: + tqdm_lock = lock_manager.Lock() + with tqdm(total=len(in_file_chunks), disable=quiet, desc=desc) as progress: + with executor_cls(max_workers=num_workers) as executor: + futures = [ + executor.submit( + pipeline_multi_wrapper, + filepaths=file_chunk, + out_file_stems=out_stem_chunk, + fast_mode=fast_mode, + force_overwrite=force_overwrite, + worker_position=index % num_workers, + quiet=quiet, + tqdm_lock=tqdm_lock, + ) + for index, (file_chunk, out_stem_chunk) in enumerate( + zip(in_file_chunks, out_stem_chunks) + ) + ] + for future in concurrent.futures.as_completed(futures): + outputs, compressed, metadata, failures = future.result() + output_paths.extend(outputs) + compressed_paths.extend(compressed) + metadata_paths.extend(metadata) + failed_files.extend(failures) + progress.update(1) + + return output_paths, compressed_paths, metadata_paths, failed_files + + +def batch( + filepaths: Sequence[str | Path], + config: str | None = None, + clean: bool = True, + num_workers: int = 1, +) -> list[ClassificationResult]: + """Classify multiple WAV files using one reusable ONNX session. + + ``clean`` remains for API compatibility; temporary spectrograms are always + removed by the classifier. + """ + del clean + from batbot.classifier import Classifier + + classifier = Classifier(config=config, num_workers=num_workers) + return [classifier.classify_wav(filepath) for filepath in filepaths] + + +def example() -> None: + """Run the spectrogram pipeline on the packaged example WAV.""" + import pooch + + wav_filepath = PROJECT_ROOT / 'examples' / 'example1.wav' + if not wav_filepath.exists(): + wav_filepath = Path( + pooch.retrieve( + url=( + 'https://media.githubusercontent.com/media/Kitware/batbot/' + 'main/examples/example1.wav' + ), + known_hash=( + 'sha256:391efce5433d1057caddb4ce07b9712c523d6a815e4ee9e64b62973569982925' + ), + progressbar=True, + ) + ) + + log.debug('Running pipeline on WAV: %s', wav_filepath) + output_stem = Path('output') / wav_filepath.stem + start_time = time.time() + results = pipeline( + wav_filepath, + out_file_stem=str(output_stem), + fast_mode=False, + force_overwrite=True, + plot_uncompressed_amplitude=True, + include_original_sr=True, + time_buffer_ms=5.0, + ) + print(f'Example pipeline completed in {time.time() - start_time} seconds.') + log.debug(results) diff --git a/batbot/batbot_cli.py b/batbot/batbot_cli.py index ee49094..dd81e94 100755 --- a/batbot/batbot_cli.py +++ b/batbot/batbot_cli.py @@ -2,6 +2,7 @@ """ CLI for BatBot """ + import json import pprint import warnings @@ -24,13 +25,13 @@ from tqdm import tqdm import batbot -from batbot import log +from batbot import classifier +from batbot._config import log def pipeline_filepath_validator(ctx, param, value): if not exists(value): - log.error(f'Input filepath does not exist: {value}') - ctx.exit() + raise click.BadParameter(f'Input filepath does not exist: {value}') return value @@ -39,13 +40,18 @@ def pipeline_filepath_validator(ctx, param, value): '--config', help='Which ML model to use for inference', default=None, - type=click.Choice(['usgs']), + type=click.Choice(['mobilenet']), ) -def fetch(config): +@click.option( + '--pull', + is_flag=True, + help='Download the mirrored model even when a bundled copy is available.', +) +def fetch(config, pull): """ Fetch the required machine learning ONNX model for the classifier """ - batbot.fetch(config=config) + print(batbot.fetch(config=config, pull=pull)) @click.command('pipeline') @@ -183,7 +189,7 @@ def preprocess( in_filepaths = sorted(list(set(in_filepaths))) if len(in_filepaths) == 0: - print('Found no files given filepaths input {}'.format(filepaths)) + print(f'Found no files given filepaths input {filepaths}') return # set up output paths for each input path @@ -209,7 +215,7 @@ def preprocess( if not force_overwrite: idx_remove = np.full((len(in_filepaths),), False) for ii, out_file_stem in enumerate(out_filepath_stems): - test_file = '{}.*'.format(out_file_stem) + test_file = f'{out_file_stem}.*' test_glob = glob(test_file) if len(test_glob) > 0: idx_remove[ii] = True @@ -229,7 +235,7 @@ def preprocess( # Find all "extra" files that would be deleted in cleanup mode all_files = set(glob(join(root_outpath, '**/*'), recursive=True)) for out_stem in out_filepath_stems_all: - out_files = glob('{}.*'.format(out_stem)) + out_files = glob(f'{out_stem}.*') all_files -= set(out_files) dir_files = [] # remove directories @@ -239,18 +245,18 @@ def preprocess( all_files -= set(dir_files) extra_files = all_files - print('Located {} total unprocessed files'.format(len(in_filepaths))) + print(f'Located {len(in_filepaths)} total unprocessed files') print('\tFast processing mode {}'.format('OFF' if process_metadata else 'ON')) if process_metadata: print('\t\tFull bat call metadata will be produced') print('\tForce output overwrite {}'.format('ON' if force_overwrite else 'OFF')) if not force_overwrite: - print('\t\tSkipped {} files with already preprocessed outputs'.format(n_skipped)) - print('\tNum parallel workers: {}'.format(num_workers)) + print(f'\t\tSkipped {n_skipped} files with already preprocessed outputs') + print(f'\tNum parallel workers: {num_workers}') if no_file_structure: print('\tFlattening output file structure') - print('\tCurrent working dir: {}'.format(getcwd())) - print('\tOutput root dir: {}'.format(output_dir)) + print(f'\tCurrent working dir: {getcwd()}') + print(f'\tOutput root dir: {output_dir}') print( '\tFirst input file -> output files: {} -> {}.*'.format( in_filepaths[0], out_filepath_stems[0] @@ -268,7 +274,7 @@ def preprocess( print('\nDry run mode active - skipping all processing') data = {} data['input file, output file stem'] = [ - (str(x), '{}.*'.format(y)) for x, y in zip(in_filepaths, out_filepath_stems) + (str(x), f'{y}.*') for x, y in zip(in_filepaths, out_filepath_stems) ] data['files to be deleted in cleanup'] = list(extra_files) if output_json is None: @@ -276,7 +282,7 @@ def preprocess( else: with open(output_json, 'w') as outfile: json.dump(data, outfile, indent=4) - print('Outputs written to {}'.format(output_json)) + print(f'Outputs written to {output_json}') print('Complete.') return @@ -294,13 +300,18 @@ def preprocess( print('Aborting cleanup mode.') return for file in extra_files: - print('Deleting file: {}'.format(file)) + print(f'Deleting file: {file}') remove(file) print('Complete.') return # Begin execution loop. - data = {'output_path': [], 'compressed_path': [], 'metadata_path': [], 'failed_files': []} + data = { + 'output_path': [], + 'compressed_path': [], + 'metadata_path': [], + 'failed_files': [], + } if num_workers is None or num_workers == 0: # Serial execution. @@ -321,7 +332,7 @@ def preprocess( data['compressed_path'].extend(compressed_paths) data['metadata_path'].append(metadata_path) except Exception as e: - warnings.warn('WARNING: Pipeline failed for file {}'.format(file)) + warnings.warn(f'WARNING: Pipeline failed for file {file}') data['failed_files'].append((str(file), e)) else: # Parallel execution. @@ -351,7 +362,7 @@ def preprocess( num_workers=num_workers, threaded=False, quiet=True, - desc='Preprocessing chunks of files with {} workers'.format(num_workers), + desc=f'Preprocessing chunks of files with {num_workers} workers', ) data['output_path'].extend(output_paths) data['compressed_path'].extend(compressed_paths) @@ -370,7 +381,7 @@ def preprocess( else: with open(output_json, 'w') as outfile: json.dump(data, outfile, indent=4) - print('Outputs written to {}'.format(output_json)) + print(f'Outputs written to {output_json}') print('\nComplete.') return data @@ -386,7 +397,7 @@ def preprocess( '--config', help='Which ML model to use for inference', default=None, - type=click.Choice(['usgs']), + type=click.Choice(['mobilenet']), ) @click.option( '--output', @@ -394,6 +405,13 @@ def preprocess( default=None, type=str, ) +@click.option( + '--num-workers', + default=1, + show_default=True, + type=click.IntRange(min=1), + help='Number of concurrent ONNX inference workers.', +) # @click.option( # '--classifier_thresh', # help='Classifier confidence threshold', @@ -404,48 +422,182 @@ def batch( filepaths, config, output, + num_workers, # classifier_thresh, ): - """ - Run the BatBot pipeline in batch on a list of input WAV filepaths. - An example output of the JSON can be seen below. - - .. code-block:: javascript - - { - '/path/to/file1.wav': { - 'classifier': 0.5, - }, - '/path/to/file2.wav': { - 'classifier': 0.8, - }, - ... - } - """ + """Classify a list of WAV files (legacy alias for ``classify-wav``).""" if config is not None: config = config.strip().lower() # classifier_thresh /= 100.0 log.debug(f'Running batch on {len(filepaths)} files...') - score_list = batbot.batch( + results = batbot.batch( filepaths, config=config, - # classifier_thresh=classifier_thresh, + num_workers=num_workers, ) - data = {} - for filepath, score in zip(filepaths, score_list): - data[filepath] = { - 'classifier': score, - } + data = { + 'results': results, + 'summary': classifier.summarize(results), + } log.debug('Outputting results...') if output: with open(output, 'w') as outfile: json.dump(data, outfile, indent=4) else: - print(data) + print(json.dumps(data, indent=2)) + + +def _write_classification_output(data, output): + encoded = json.dumps(data, indent=2) + if output: + with open(output, 'w') as outfile: + outfile.write(encoded) + outfile.write('\n') + click.echo(f'Outputs written to {output}') + else: + click.echo(encoded) + + +@click.command('classify') +@click.argument('spectrograms', nargs=-1, required=True, type=click.Path(exists=True)) +@click.option( + '--output', + '-o', + default=None, + type=click.Path(dir_okay=False), + help='Path to output JSON (defaults to stdout).', +) +@click.option( + '--batch-size', + default=classifier.BATCH_SIZE, + show_default=True, + type=click.IntRange(min=1), +) +@click.option( + '--num-workers', + default=1, + show_default=True, + type=click.IntRange(min=1), + help='Number of concurrent ONNX inference workers.', +) +@click.option('--top-k', default=5, show_default=True, type=click.IntRange(1, 35)) +def classify(spectrograms, output, batch_size, num_workers, top_k): + """Classify one or more spectrogram image files.""" + paths = classifier.discover_inputs(spectrograms, input_type='spectrogram') + results = classifier.classify( + paths, + batch_size=batch_size, + top_k=top_k, + num_workers=num_workers, + ) + data = {'results': results, 'summary': classifier.summarize(results)} + _write_classification_output(data, output) + + +@click.command('classify-wav') +@click.argument('wav_files', nargs=-1, required=True, type=click.Path(exists=True)) +@click.option( + '--output', + '-o', + default=None, + type=click.Path(dir_okay=False), + help='Path to output JSON (defaults to stdout).', +) +@click.option( + '--spectrogram-dir', + default=None, + type=click.Path(file_okay=False), + help='Keep generated spectrograms in this directory.', +) +@click.option( + '--batch-size', + default=classifier.BATCH_SIZE, + show_default=True, + type=click.IntRange(min=1), +) +@click.option( + '--num-workers', + default=1, + show_default=True, + type=click.IntRange(min=1), + help='Number of concurrent ONNX inference workers.', +) +@click.option('--top-k', default=5, show_default=True, type=click.IntRange(1, 35)) +def classify_wav(wav_files, output, spectrogram_dir, batch_size, num_workers, top_k): + """Create spectrograms from WAV files and classify each recording.""" + data = classifier.classify_bulk( + wav_files, + input_type='wav', + batch_size=batch_size, + top_k=top_k, + spectrogram_output=spectrogram_dir, + num_workers=num_workers, + ) + _write_classification_output(data, output) + + +@click.command('classify-bulk') +@click.argument('inputs', nargs=-1, required=True, type=click.Path(exists=True)) +@click.option( + '--input-type', + default='auto', + show_default=True, + type=click.Choice(['auto', 'spectrogram', 'wav']), + help='Limit inputs to spectrograms, WAV files, or detect both.', +) +@click.option('--recursive/--no-recursive', default=True, show_default=True) +@click.option( + '--output', + '-o', + default=None, + type=click.Path(dir_okay=False), + help='Path to output JSON (defaults to stdout).', +) +@click.option( + '--spectrogram-dir', + default=None, + type=click.Path(file_okay=False), + help='Keep spectrograms generated for WAV inputs in this directory.', +) +@click.option( + '--batch-size', + default=classifier.BATCH_SIZE, + show_default=True, + type=click.IntRange(min=1), +) +@click.option( + '--num-workers', + default=1, + show_default=True, + type=click.IntRange(min=1), + help='Number of concurrent ONNX inference workers.', +) +@click.option('--top-k', default=5, show_default=True, type=click.IntRange(1, 35)) +def classify_bulk( + inputs, + input_type, + recursive, + output, + spectrogram_dir, + batch_size, + num_workers, + top_k, +): + """Recursively classify a large folder and report species counts.""" + data = classifier.classify_bulk( + inputs, + input_type=input_type, + recursive=recursive, + batch_size=batch_size, + top_k=top_k, + spectrogram_output=spectrogram_dir, + num_workers=num_workers, + ) + _write_classification_output(data, output) @click.command('example') @@ -468,6 +620,9 @@ def cli(): cli.add_command(pipeline) cli.add_command(preprocess) cli.add_command(batch) +cli.add_command(classify) +cli.add_command(classify_wav) +cli.add_command(classify_bulk) cli.add_command(example) diff --git a/batbot/classifier/__init__.py b/batbot/classifier/__init__.py new file mode 100644 index 0000000..c0bb431 --- /dev/null +++ b/batbot/classifier/__init__.py @@ -0,0 +1,63 @@ +"""Classify BatBot spectrograms with the MobileNet ONNX model. + +The low-level :func:`pre`, :func:`predict`, and :func:`post` functions retain +the Scoutbot-style pipeline. :class:`Classifier` is the preferred API for +repeated work because it owns and reuses its ONNX Runtime session. +""" + +from batbot.classifier.bulk import classify_bulk, classify_wav, discover_inputs, summarize +from batbot.classifier.config import ( + CLASSES, + CONFIGS, + DEFAULT_CONFIG, + MODEL_HASH, + MODEL_NAME, + MODEL_URL, + SPECTROGRAM_EXTENSIONS, + WAV_EXTENSIONS, + resolve_config, +) +from batbot.classifier.dataloader import BATCH_SIZE, INPUT_SIZE, ImageFilePathList +from batbot.classifier.inference import Classifier, classify, post, pre, predict +from batbot.classifier.model import fetch +from batbot.classifier.types import ( + BulkClassification, + ClassificationFailure, + ClassificationItem, + ClassificationResult, + ClassificationSummary, + ClassifierConfig, + TopPrediction, +) + +__all__ = [ + 'BATCH_SIZE', + 'CLASSES', + 'CONFIGS', + 'DEFAULT_CONFIG', + 'INPUT_SIZE', + 'MODEL_HASH', + 'MODEL_NAME', + 'MODEL_URL', + 'SPECTROGRAM_EXTENSIONS', + 'WAV_EXTENSIONS', + 'BulkClassification', + 'ClassificationFailure', + 'ClassificationItem', + 'ClassificationResult', + 'ClassificationSummary', + 'Classifier', + 'ClassifierConfig', + 'ImageFilePathList', + 'TopPrediction', + 'classify', + 'classify_bulk', + 'classify_wav', + 'discover_inputs', + 'fetch', + 'post', + 'pre', + 'predict', + 'resolve_config', + 'summarize', +] diff --git a/batbot/classifier/bulk.py b/batbot/classifier/bulk.py new file mode 100644 index 0000000..8f848f5 --- /dev/null +++ b/batbot/classifier/bulk.py @@ -0,0 +1,226 @@ +"""WAV and fault-tolerant bulk classification orchestration.""" + +from __future__ import annotations + +import tempfile +from collections import Counter +from collections.abc import Iterable, MutableMapping, Sequence +from hashlib import sha256 +from pathlib import Path +from typing import Any, TypeGuard + +import numpy as np +import tqdm + +from batbot._config import QUIET +from batbot.classifier.config import ( + DEFAULT_CONFIG, + SPECTROGRAM_EXTENSIONS, + WAV_EXTENSIONS, +) +from batbot.classifier.dataloader import BATCH_SIZE +from batbot.classifier.inference import Classifier, _as_filepaths, _format_result +from batbot.classifier.types import ( + BulkClassification, + ClassificationFailure, + ClassificationItem, + ClassificationResult, + ClassificationSummary, + ClassifierConfig, + InputType, + PathInput, +) + + +def _is_success(result: ClassificationItem) -> TypeGuard[ClassificationResult]: + return 'label' in result + + +def _is_failure(result: ClassificationItem) -> TypeGuard[ClassificationFailure]: + return 'error' in result + + +def _aggregate_results( + filepath: PathInput, + results: Sequence[ClassificationResult], + spectrogram_paths: Sequence[PathInput], + top_k: int = 5, +) -> ClassificationResult: + if not results: + raise ValueError(f'No spectrograms were created for {filepath}') + + weights = np.asarray([result['window_count'] for result in results], dtype=np.float64) + classes = list(results[0]['scores']) + values = np.asarray( + [[result['scores'][class_name] for class_name in classes] for result in results] + ) + scores = { + class_name: float(score) + for class_name, score in zip(classes, np.average(values, axis=0, weights=weights)) + } + output = _format_result(filepath, scores, int(weights.sum()), top_k=top_k) + output['spectrogram_paths'] = [str(path) for path in spectrogram_paths] + return output + + +def classify_wav( + filepath: PathInput, + batch_size: int = BATCH_SIZE, + config: str | ClassifierConfig | None = DEFAULT_CONFIG, + providers: Sequence[str] | None = None, + top_k: int = 5, + output_folder: PathInput | None = None, + out_file_stem: PathInput | None = None, + keep_spectrograms: bool = False, + sessions: MutableMapping[str, Any] | None = None, + num_workers: int = 1, + *, + _runner: Classifier | None = None, +) -> ClassificationResult: + """Generate spectrograms for a WAV file and return one prediction.""" + from batbot.spectrogram import compute + + runner = _runner or Classifier( + config=config, + batch_size=batch_size, + providers=providers, + top_k=top_k, + sessions=sessions, + num_workers=num_workers, + ) + filepath_string = str(filepath) + if output_folder is not None or out_file_stem is not None or keep_spectrograms: + selected_output = './output' if output_folder is None else str(output_folder) + output_paths, _, _, _ = compute( + filepath_string, + output_folder=selected_output, + out_file_stem=None if out_file_stem is None else str(out_file_stem), + fast_mode=False, + force_overwrite=True, + quiet=True, + ) + results = runner.classify(output_paths, top_k=top_k) + return _aggregate_results(filepath_string, results, output_paths, top_k=top_k) + + with tempfile.TemporaryDirectory(prefix='batbot-classifier-') as temp_dir: + output_paths, _, _, _ = compute( + filepath_string, + output_folder=temp_dir, + fast_mode=False, + force_overwrite=True, + quiet=True, + ) + results = runner.classify(output_paths, top_k=top_k) + output = _aggregate_results(filepath_string, results, output_paths, top_k=top_k) + output['spectrogram_paths'] = [] + return output + + +def discover_inputs( + inputs: PathInput | Iterable[PathInput], + input_type: InputType = 'auto', + recursive: bool = True, +) -> list[str]: + """Resolve files and directories into deterministic classifier inputs.""" + if input_type not in {'auto', 'spectrogram', 'wav'}: + raise ValueError('input_type must be auto, spectrogram, or wav') + + extensions = SPECTROGRAM_EXTENSIONS | WAV_EXTENSIONS + if input_type == 'spectrogram': + extensions = SPECTROGRAM_EXTENSIONS + elif input_type == 'wav': + extensions = WAV_EXTENSIONS + + discovered = [] + for value in _as_filepaths(inputs): + path = Path(value) + if path.is_file(): + if path.suffix.lower() in extensions: + discovered.append(str(path)) + elif path.is_dir(): + iterator = path.rglob('*') if recursive else path.glob('*') + discovered.extend( + str(candidate) + for candidate in iterator + if candidate.is_file() and candidate.suffix.lower() in extensions + ) + else: + raise FileNotFoundError(f'Input does not exist: {value}') + return sorted(set(discovered)) + + +def summarize(results: Sequence[ClassificationItem]) -> ClassificationSummary: + """Generate species counts and basic confidence statistics.""" + successful = [result for result in results if _is_success(result)] + failures = [result for result in results if _is_failure(result)] + counts = Counter(result['label'] for result in successful) + species_counts = {label: count for label, count in counts.items() if label != 'NOISE'} + confidences = [result['confidence'] for result in successful] + return { + 'total': len(results), + 'classified': len(successful), + 'failed': len(failures), + 'label_counts': dict(sorted(counts.items())), + 'species_counts': dict(sorted(species_counts.items())), + 'noise_count': counts.get('NOISE', 0), + 'mean_confidence': float(np.mean(confidences)) if confidences else None, + } + + +def classify_bulk( + inputs: PathInput | Iterable[PathInput], + input_type: InputType = 'auto', + recursive: bool = True, + batch_size: int = BATCH_SIZE, + config: str | ClassifierConfig | None = DEFAULT_CONFIG, + providers: Sequence[str] | None = None, + top_k: int = 5, + spectrogram_output: PathInput | None = None, + num_workers: int = 1, + *, + _runner: Classifier | None = None, +) -> BulkClassification: + """Classify a directory tree while reusing a single ONNX session.""" + paths = discover_inputs(inputs, input_type=input_type, recursive=recursive) + runner = _runner or Classifier( + config=config, + batch_size=batch_size, + providers=providers, + top_k=top_k, + num_workers=num_workers, + ) + results_by_path: dict[str, ClassificationItem] = {} + + image_paths = [path for path in paths if Path(path).suffix.lower() in SPECTROGRAM_EXTENSIONS] + if image_paths: + try: + for result in runner.classify(image_paths, top_k=top_k): + results_by_path[result['path']] = result + except (OSError, ValueError): + for path in image_paths: + try: + results_by_path[path] = runner.classify([path], top_k=top_k)[0] + except Exception as error: # pragma: no cover - decoder errors vary by platform + results_by_path[path] = {'path': path, 'error': str(error)} + + wav_paths = [path for path in paths if Path(path).suffix.lower() in WAV_EXTENSIONS] + for path in tqdm.tqdm(wav_paths, disable=QUIET, desc='Classifying WAV files'): + try: + out_file_stem = None + if spectrogram_output is not None: + digest = sha256(str(Path(path).resolve()).encode('utf8')).hexdigest()[:10] + output_name = f'{Path(path).stem}.{digest}' + out_file_stem = str(Path(spectrogram_output) / output_name) + results_by_path[path] = classify_wav( + path, + top_k=top_k, + output_folder=spectrogram_output, + out_file_stem=out_file_stem, + keep_spectrograms=spectrogram_output is not None, + _runner=runner, + ) + except Exception as error: # pragma: no cover - spectrogram errors vary by input + results_by_path[path] = {'path': path, 'error': str(error)} + + results = [results_by_path[path] for path in paths] + return {'results': results, 'summary': summarize(results)} diff --git a/batbot/classifier/config.py b/batbot/classifier/config.py new file mode 100644 index 0000000..0939cc9 --- /dev/null +++ b/batbot/classifier/config.py @@ -0,0 +1,93 @@ +"""Classifier model declarations and configuration resolution.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from types import MappingProxyType + +from batbot.classifier.types import ClassifierConfig + +MODEL_NAME = 'batbot.mobilenet.9dc57ea3.onnx' +MODEL_URL = 'https://data.kitware.com/api/v1/file/6a8377e32688ba21262c3907/download' +MODEL_HASH = '351aa656ce5df717472d3c1cda7cef883b321d3bbb382740e16b3f2a1d74f621' + +CLASSES = ( + 'ANPA', + 'CORA', + 'COTO', + 'EPFU', + 'EUFL', + 'EUMA', + 'EUPE', + 'IDPH', + 'LABL', + 'LABO', + 'LACI', + 'LAIN', + 'LANO', + 'LASE', + 'LAXA', + 'MYAU', + 'MYCA', + 'MYCI', + 'MYEV', + 'MYGR', + 'MYLE', + 'MYLU', + 'MYSE', + 'MYSO', + 'MYTH', + 'MYVE', + 'MYVO', + 'MYYU', + 'NOISE', + 'NYFE', + 'NYHU', + 'NYMA', + 'PAHE', + 'PESU', + 'TABR', +) + +MOBILENET = ClassifierConfig( + key='mobilenet', + filename=MODEL_NAME, + url=MODEL_URL, + sha256=MODEL_HASH, + classes=CLASSES, +) + +_CONFIGS: dict[str | None, ClassifierConfig] = { + 'mobilenet': MOBILENET, + None: MOBILENET, +} +CONFIGS: Mapping[str | None, ClassifierConfig] = MappingProxyType(_CONFIGS) + +DEFAULT_CONFIG = ( + os.getenv( + 'BATBOT_CLASSIFIER_CONFIG', + os.getenv('CLASSIFIER_CONFIG', 'mobilenet'), + ) + .strip() + .lower() +) +if DEFAULT_CONFIG not in CONFIGS: + raise ValueError(f'Unknown classifier configuration: {DEFAULT_CONFIG}') + +SPECTROGRAM_EXTENSIONS = frozenset({'.jpg', '.jpeg', '.png', '.tif', '.tiff'}) +WAV_EXTENSIONS = frozenset({'.wav'}) + + +def resolve_config(config: str | ClassifierConfig | None = None) -> ClassifierConfig: + """Resolve a configuration name or return an existing configuration.""" + if isinstance(config, ClassifierConfig): + return config + key = DEFAULT_CONFIG if config is None else str(config).strip().lower() + try: + return CONFIGS[key] + except KeyError as error: + choices = ', '.join(sorted(key for key in CONFIGS if key is not None)) + raise ValueError( + f'Unknown classifier configuration {key!r}; choose from {choices}' + ) from error diff --git a/batbot/classifier/dataloader.py b/batbot/classifier/dataloader.py new file mode 100644 index 0000000..b385977 --- /dev/null +++ b/batbot/classifier/dataloader.py @@ -0,0 +1,133 @@ +"""Image loading and windowing for the BatBot species classifier.""" + +from __future__ import annotations + +import os +from collections.abc import Callable, Iterable +from typing import Any + +import cv2 +import numpy as np +from numpy.typing import NDArray + +from batbot.classifier.types import PathInput + +BATCH_SIZE = int( + os.getenv('BATBOT_CLASSIFIER_BATCH_SIZE', os.getenv('CLASSIFIER_BATCH_SIZE', '10')) +) +INPUT_SIZE = 224 +WINDOW_STRIDE = 100 +HORIZONTAL_SCALE = 0.5 + + +def _load_image(filepath: PathInput) -> NDArray[np.uint8]: + """Load a spectrogram in the BGR byte layout used to train the model.""" + image = cv2.imread(str(filepath), cv2.IMREAD_COLOR) + if image is None: + raise OSError(f'Unable to load spectrogram: {filepath}') + return np.asarray(image, dtype=np.uint8) + + +def _prepare_image( + image: NDArray[np.uint8] | None, + input_size: int = INPUT_SIZE, + window_stride: int = WINDOW_STRIDE, + horizontal_scale: float = HORIZONTAL_SCALE, +) -> NDArray[np.uint8]: + """Resize a spectrogram and split it into overlapping square windows. + + The training-time evaluation script resized a 300-pixel-high image by + ``224 / 300`` vertically and half that amount horizontally. Computing the + vertical scale from the actual image height retains that transform for the + original data while also accepting the spectrogram height emitted by + BatBot today. + """ + if image is None or image.ndim != 3 or image.shape[2] != 3: + raise ValueError('Expected a three-channel spectrogram image') + if input_size <= 0: + raise ValueError('input_size must be positive') + if window_stride <= 0: + raise ValueError('window_stride must be positive') + + height, width, _ = image.shape + ratio_y = input_size / float(height) + target_width = max(1, int(round(width * ratio_y * horizontal_scale))) + resized = cv2.resize( + image, + (target_width, input_size), + interpolation=cv2.INTER_LANCZOS4, + ) + + # The reference inference script pads narrow inputs to one pixel wider than + # a square so that range(0, width - height, stride) yields one window. + if resized.shape[1] <= input_size: + canvas = np.zeros((input_size, input_size + 1, 3), dtype=resized.dtype) + canvas[:, : resized.shape[1], :] = resized + resized = canvas + + starts = range(0, resized.shape[1] - input_size, window_stride) + windows = [resized[:, start : start + input_size, :] for start in starts] + if not windows: # Defensive fallback for custom transform arguments. + windows = [resized[:, :input_size, :]] + + return np.ascontiguousarray(np.stack(windows), dtype=np.uint8) + + +def _init_transforms( + input_size: int = INPUT_SIZE, + window_stride: int = WINDOW_STRIDE, + horizontal_scale: float = HORIZONTAL_SCALE, +) -> Callable[[NDArray[np.uint8]], NDArray[np.uint8]]: + """Return the deterministic preprocessing transform used for inference.""" + + def transform(image: NDArray[np.uint8]) -> NDArray[np.uint8]: + return _prepare_image( + image, + input_size=input_size, + window_stride=window_stride, + horizontal_scale=horizontal_scale, + ) + + return transform + + +class ImageFilePathList: + """Small, dependency-free equivalent of Scoutbot's image path dataset.""" + + def __init__( + self, + filepaths: Iterable[PathInput], + targets: Iterable[Any] | None = None, + transform: Callable[[NDArray[np.uint8]], NDArray[np.uint8]] | None = None, + target_transform: Callable[[Any], Any] | None = None, + ) -> None: + self.filepaths = [str(filepath) for filepath in filepaths] + self.target_values = list(targets) if targets is not None else None + if self.target_values is not None and len(self.filepaths) != len(self.target_values): + raise ValueError('filepaths and targets must have the same length') + + self.loader = _load_image + self.transform = transform + self.target_transform = target_transform + + if self.target_values is None: + self.classes, self.class_to_idx = None, None + else: + self.classes = sorted(set(self.target_values)) + self.class_to_idx = {class_name: index for index, class_name in enumerate(self.classes)} + + def __getitem__(self, index: int) -> tuple[NDArray[np.uint8], ...]: + sample = self.loader(self.filepaths[index]) + if self.transform is not None: + sample = self.transform(sample) + + if self.target_values is None: + return (sample,) + + target = self.target_values[index] + if self.target_transform is not None: + target = self.target_transform(target) + return sample, target + + def __len__(self) -> int: + return len(self.filepaths) diff --git a/batbot/classifier/inference.py b/batbot/classifier/inference.py new file mode 100644 index 0000000..d251e8f --- /dev/null +++ b/batbot/classifier/inference.py @@ -0,0 +1,355 @@ +"""Spectrogram preprocessing and ONNX inference.""" + +from __future__ import annotations + +import json +import os +import warnings +from collections import deque +from collections.abc import Iterable, Iterator, MutableMapping, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any + +import numpy as np +import tqdm + +from batbot._config import QUIET, log +from batbot.classifier.config import DEFAULT_CONFIG, resolve_config +from batbot.classifier.dataloader import BATCH_SIZE, INPUT_SIZE, ImageFilePathList, _init_transforms +from batbot.classifier.model import fetch +from batbot.classifier.types import ( + BulkClassification, + ClassificationResult, + ClassifierConfig, + InputType, + PathInput, +) + + +def _as_filepaths(inputs: PathInput | Iterable[PathInput]) -> list[str]: + if isinstance(inputs, (str, os.PathLike)): + return [str(inputs)] + return [str(filepath) for filepath in inputs] + + +def pre( + inputs: PathInput | Iterable[PathInput], + batch_size: int = BATCH_SIZE, + config: str | ClassifierConfig | None = DEFAULT_CONFIG, +) -> Iterator[tuple[np.ndarray[Any, Any], str]]: + """Load spectrograms and yield their model-ready sliding windows.""" + selected = resolve_config(config) + filepaths = _as_filepaths(inputs) + if batch_size <= 0: + raise ValueError('batch_size must be positive') + + log.debug( + 'Preprocessing %d classifier inputs with inference batches of %d', + len(filepaths), + batch_size, + ) + dataset = ImageFilePathList(filepaths, transform=_init_transforms()) + for index in range(len(dataset)): + (data,) = dataset[index] + yield data, selected.key + + +def _create_session(onnx_model: str, providers: Sequence[str] | None = None, reduced=False) -> Any: + # Official non-Windows builds may otherwise create a persistent telemetry + # device identifier as soon as ONNX Runtime initializes. Respect an + # explicit user setting while making private inference the default. + reduced = os.getenv('BATBOT_REDUCED', reduced) in [True, '1', 'Yes', 'yes', 'YES'] + + os.environ.setdefault('ORT_DISABLE_TELEMETRY', '1') + try: + import onnxruntime as ort + except ImportError as error: # pragma: no cover - a declared runtime dependency + raise ImportError( + 'ONNX inference requires onnxruntime; install batbot with its runtime dependencies' + ) from error + ort.disable_telemetry_events() + + selected_providers = providers + if selected_providers is None: + available = ort.get_available_providers() + preferred = ['CUDAExecutionProvider', 'CPUExecutionProvider'] + selected_providers = [provider for provider in preferred if provider in available] + + with warnings.catch_warnings(): + warnings.filterwarnings('ignore', category=UserWarning) + if reduced: + print('Reducing ONNX inference threads to 2') + opts = ort.SessionOptions() + opts.intra_op_num_threads = 2 + opts.inter_op_num_threads = 2 + else: + opts = None + return ort.InferenceSession(onnx_model, providers=selected_providers, sess_options=opts) + + +def _validate_session(session: Any, config: ClassifierConfig) -> None: + inputs = session.get_inputs() + outputs = session.get_outputs() + if len(inputs) != 1 or not outputs: + raise ValueError('Classifier model must have one input and at least one output') + + input_shape = inputs[0].shape + if list(input_shape[1:]) != [INPUT_SIZE, INPUT_SIZE, 3]: + raise ValueError(f'Unexpected classifier input shape: {input_shape}') + + metadata = session.get_modelmeta().custom_metadata_map + if 'labels' in metadata: + mapping = json.loads(metadata['labels']) + labels = [mapping['forward'][str(index)] for index in range(len(mapping['forward']))] + if labels != list(config.classes): + raise ValueError('Classifier labels do not match the selected configuration') + + +def _predict_windows( + windows: np.ndarray[Any, Any], + session: Any, + batch_size: int, +) -> np.ndarray[Any, Any]: + """Run one spectrogram's windows through a reusable ONNX session.""" + input_name = session.get_inputs()[0].name + outputs = [] + for start in range(0, len(windows), batch_size): + output = session.run(None, {input_name: windows[start : start + batch_size]}) + outputs.append(output[0]) + if not outputs: + raise ValueError('Classifier preprocessing produced no image windows') + return np.vstack(outputs).mean(axis=0, keepdims=True) + + +def predict( + gen: Iterable[tuple[np.ndarray[Any, Any], str | ClassifierConfig]], + batch_size: int = BATCH_SIZE, + providers: Sequence[str] | None = None, + pull: bool = False, + sessions: MutableMapping[str, Any] | None = None, + total: int | None = None, + num_workers: int = 1, +) -> Iterator[tuple[np.ndarray[Any, Any], str]]: + """Run ordered ONNX inference, optionally across multiple worker threads.""" + if batch_size <= 0: + raise ValueError('batch_size must be positive') + if num_workers <= 0: + raise ValueError('num_workers must be positive') + + active_sessions = {} if sessions is None else sessions + items = tqdm.tqdm( + gen, + disable=QUIET, + desc='Classifying spectrograms', + total=total, + ) + + def session_for(config_value: str | ClassifierConfig) -> tuple[Any, str]: + config = resolve_config(config_value) + session = active_sessions.get(config.key) + if session is None: + session = _create_session(fetch(pull=pull, config=config), providers=providers) + _validate_session(session, config) + active_sessions[config.key] = session + return session, config.key + + if num_workers == 1: + for windows, config_value in items: + session, config_key = session_for(config_value) + yield _predict_windows(windows, session, batch_size), config_key + return + + pending: deque[tuple[Future[np.ndarray[Any, Any]], str]] = deque() + with ThreadPoolExecutor( + max_workers=num_workers, + thread_name_prefix='batbot-onnx', + ) as executor: + for windows, config_value in items: + session, config_key = session_for(config_value) + pending.append( + (executor.submit(_predict_windows, windows, session, batch_size), config_key) + ) + if len(pending) >= num_workers: + future, completed_config = pending.popleft() + yield future.result(), completed_config + + while pending: + future, completed_config = pending.popleft() + yield future.result(), completed_config + + +def post( + gen: Iterable[tuple[np.ndarray[Any, Any], str | ClassifierConfig]], +) -> list[dict[str, float]]: + """Associate raw model scores with labels for each spectrogram.""" + outputs = [] + for predictions, config_value in gen: + config = resolve_config(config_value) + for prediction in predictions: + if len(prediction) != len(config.classes): + raise ValueError( + f'Model returned {len(prediction)} scores for {len(config.classes)} labels' + ) + outputs.append( + {class_name: float(score) for class_name, score in zip(config.classes, prediction)} + ) + return outputs + + +def _format_result( + filepath: PathInput, + scores: dict[str, float], + window_count: int, + top_k: int = 5, +) -> ClassificationResult: + ranked = sorted(scores.items(), key=lambda item: item[1], reverse=True) + selected_top_k = max(1, min(int(top_k), len(ranked))) + return { + 'path': str(filepath), + 'label': ranked[0][0], + 'confidence': ranked[0][1], + 'window_count': int(window_count), + 'top': [ + {'label': label, 'confidence': confidence} + for label, confidence in ranked[:selected_top_k] + ], + 'scores': scores, + } + + +class Classifier: + """Reusable classifier that owns a lazily-created ONNX Runtime session.""" + + def __init__( + self, + config: str | ClassifierConfig | None = DEFAULT_CONFIG, + batch_size: int = BATCH_SIZE, + providers: Sequence[str] | None = None, + top_k: int = 5, + pull: bool = False, + sessions: MutableMapping[str, Any] | None = None, + num_workers: int = 1, + ) -> None: + if batch_size <= 0: + raise ValueError('batch_size must be positive') + if top_k <= 0: + raise ValueError('top_k must be positive') + if num_workers <= 0: + raise ValueError('num_workers must be positive') + self.config = resolve_config(config) + self.batch_size = batch_size + self.providers = tuple(providers) if providers is not None else None + self.top_k = top_k + self.pull = pull + self._sessions = {} if sessions is None else sessions + self.num_workers = num_workers + + @property + def session(self) -> Any: + """Return this classifier's validated, reusable ONNX session.""" + session = self._sessions.get(self.config.key) + if session is None: + session = _create_session( + fetch(pull=self.pull, config=self.config), + providers=self.providers, + ) + _validate_session(session, self.config) + self._sessions[self.config.key] = session + return session + + def classify( + self, + inputs: PathInput | Iterable[PathInput], + top_k: int | None = None, + ) -> list[ClassificationResult]: + """Classify one or more spectrogram images.""" + if top_k is not None and top_k <= 0: + raise ValueError('top_k must be positive') + filepaths = _as_filepaths(inputs) + window_counts: list[int] = [] + + def track_windows() -> Iterator[tuple[np.ndarray[Any, Any], str]]: + for windows, selected_config in pre( + filepaths, + batch_size=self.batch_size, + config=self.config, + ): + window_counts.append(len(windows)) + yield windows, selected_config + + scores = post( + predict( + track_windows(), + batch_size=self.batch_size, + providers=self.providers, + pull=self.pull, + sessions={self.config.key: self.session}, + total=len(filepaths), + num_workers=self.num_workers, + ) + ) + result_top_k = self.top_k if top_k is None else top_k + return [ + _format_result(filepath, output, count, top_k=result_top_k) + for filepath, output, count in zip(filepaths, scores, window_counts) + ] + + def classify_wav( + self, + filepath: PathInput, + top_k: int | None = None, + output_folder: PathInput | None = None, + out_file_stem: PathInput | None = None, + keep_spectrograms: bool = False, + ) -> ClassificationResult: + """Generate spectrograms and classify a WAV using this session.""" + from batbot.classifier.bulk import classify_wav + + return classify_wav( + filepath, + top_k=self.top_k if top_k is None else top_k, + output_folder=output_folder, + out_file_stem=out_file_stem, + keep_spectrograms=keep_spectrograms, + _runner=self, + ) + + def classify_bulk( + self, + inputs: PathInput | Iterable[PathInput], + input_type: InputType = 'auto', + recursive: bool = True, + top_k: int | None = None, + spectrogram_output: PathInput | None = None, + ) -> BulkClassification: + """Classify a directory tree using this session.""" + from batbot.classifier.bulk import classify_bulk + + return classify_bulk( + inputs, + input_type=input_type, + recursive=recursive, + top_k=self.top_k if top_k is None else top_k, + spectrogram_output=spectrogram_output, + _runner=self, + ) + + +def classify( + inputs: PathInput | Iterable[PathInput], + batch_size: int = BATCH_SIZE, + config: str | ClassifierConfig | None = DEFAULT_CONFIG, + providers: Sequence[str] | None = None, + top_k: int = 5, + sessions: MutableMapping[str, Any] | None = None, + num_workers: int = 1, +) -> list[ClassificationResult]: + """Classify spectrograms with a one-shot or externally shared session.""" + return Classifier( + config=config, + batch_size=batch_size, + providers=providers, + top_k=top_k, + sessions=sessions, + num_workers=num_workers, + ).classify(inputs) diff --git a/batbot/classifier/model.py b/batbot/classifier/model.py new file mode 100644 index 0000000..645d38c --- /dev/null +++ b/batbot/classifier/model.py @@ -0,0 +1,90 @@ +"""Integrity-checked access to bundled and mirrored ONNX models.""" + +from __future__ import annotations + +import os +import shutil +import tempfile +from functools import lru_cache +from importlib.resources import as_file, files +from pathlib import Path + +import pooch + +from batbot._config import QUIET, log +from batbot._version import __version__ +from batbot.classifier.config import DEFAULT_CONFIG, resolve_config +from batbot.classifier.types import ClassifierConfig + + +def _cache_directory() -> Path: + """Return BatBot's versioned model cache directory.""" + return Path(pooch.os_cache('batbot')) / 'models' / __version__ + + +def _has_expected_hash(path: Path, config: ClassifierConfig) -> bool: + return pooch.file_hash(str(path), alg='sha256') == config.sha256 + + +def _materialize_resource(source: Path, config: ClassifierConfig) -> Path: + """Copy a resource extracted from a non-filesystem loader into the cache.""" + target = _cache_directory() / config.filename + target.parent.mkdir(parents=True, exist_ok=True) + if target.is_file() and _has_expected_hash(target, config): + return target + + with tempfile.NamedTemporaryFile(dir=target.parent, delete=False) as temporary: + temporary_path = Path(temporary.name) + try: + shutil.copyfile(source, temporary_path) + if not _has_expected_hash(temporary_path, config): + raise ValueError(f'Bundled classifier model failed its checksum: {config.filename}') + os.replace(temporary_path, target) + finally: + temporary_path.unlink(missing_ok=True) + return target + + +def _bundled_model(config: ClassifierConfig) -> Path | None: + resource = files(config.resource_package).joinpath(*config.resource_parts, config.filename) + if not resource.is_file(): + return None + + with as_file(resource) as candidate: + candidate = Path(candidate) + if not _has_expected_hash(candidate, config): + log.warning('Bundled classifier model failed checksum; using the mirror') + return None + if isinstance(resource, Path): + return candidate + return _materialize_resource(candidate, config) + + +def _download_model(config: ClassifierConfig) -> Path: + downloaded = pooch.retrieve( + url=config.url, + known_hash=f'sha256:{config.sha256}', + path=_cache_directory(), + fname=config.filename, + progressbar=not QUIET, + ) + return Path(downloaded) + + +@lru_cache(maxsize=None) +def _fetch_cached(pull: bool, config: ClassifierConfig) -> str: + model = None if pull else _bundled_model(config) + if model is None: + model = _download_model(config) + if not model.is_file(): # pragma: no cover - Pooch raises before this in normal failures + raise OSError('Classifier model could not be fetched') + log.debug('Classifier model: %s', model) + return str(model) + + +def fetch( + pull: bool = False, + config: str | ClassifierConfig | None = DEFAULT_CONFIG, +) -> str: + """Return a verified local model, downloading the mirror when necessary.""" + return _fetch_cached(pull, resolve_config(config)) diff --git a/batbot/classifier/models/onnx/batbot.mobilenet.9dc57ea3.onnx b/batbot/classifier/models/onnx/batbot.mobilenet.9dc57ea3.onnx new file mode 100644 index 0000000..d2f8d8c --- /dev/null +++ b/batbot/classifier/models/onnx/batbot.mobilenet.9dc57ea3.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:351aa656ce5df717472d3c1cda7cef883b321d3bbb382740e16b3f2a1d74f621 +size 12135629 diff --git a/batbot/classifier/types.py b/batbot/classifier/types.py new file mode 100644 index 0000000..58bea1c --- /dev/null +++ b/batbot/classifier/types.py @@ -0,0 +1,71 @@ +"""Public type definitions for classifier configuration and results.""" + +from __future__ import annotations + +from dataclasses import dataclass +from os import PathLike +from typing import Literal, NotRequired, TypeAlias, TypedDict + +PathInput: TypeAlias = str | PathLike[str] +InputType: TypeAlias = Literal['auto', 'spectrogram', 'wav'] + + +@dataclass(frozen=True, slots=True) +class ClassifierConfig: + """Immutable description of an ONNX classifier model.""" + + key: str + filename: str + url: str + sha256: str + classes: tuple[str, ...] + resource_package: str = 'batbot.classifier' + resource_parts: tuple[str, ...] = ('models', 'onnx') + + +class TopPrediction(TypedDict): + """One ranked classifier prediction.""" + + label: str + confidence: float + + +class ClassificationResult(TypedDict): + """Successful classification serialized by the Python API and CLI.""" + + path: str + label: str + confidence: float + window_count: int + top: list[TopPrediction] + scores: dict[str, float] + spectrogram_paths: NotRequired[list[str]] + + +class ClassificationFailure(TypedDict): + """Input that could not be classified during fault-tolerant bulk work.""" + + path: str + error: str + + +ClassificationItem: TypeAlias = ClassificationResult | ClassificationFailure + + +class ClassificationSummary(TypedDict): + """Aggregate species and confidence statistics.""" + + total: int + classified: int + failed: int + label_counts: dict[str, int] + species_counts: dict[str, int] + noise_count: int + mean_confidence: float | None + + +class BulkClassification(TypedDict): + """JSON-compatible bulk classification response.""" + + results: list[ClassificationItem] + summary: ClassificationSummary diff --git a/batbot/py.typed b/batbot/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/batbot/spectrogram/__init__.py b/batbot/spectrogram/__init__.py index 793d9a3..9ab7985 100644 --- a/batbot/spectrogram/__init__.py +++ b/batbot/spectrogram/__init__.py @@ -26,7 +26,7 @@ from shapely.geometry.polygon import Polygon from skimage import draw, measure -from batbot import log +from batbot._config import log # lp = LineProfiler() @@ -1448,7 +1448,7 @@ def compute_wrapper( - list of spectrogram filepaths, split by 50k horizontal pixels """ if not force_overwrite: - test_file = '{}.*'.format(out_file_stem) + test_file = f'{out_file_stem}.*' test_glob = glob(test_file) if len(test_glob) > 0: if not quiet: @@ -1944,6 +1944,8 @@ def compute_wrapper( masked_paths = [] waveplot_compressed_paths = [] waveplot_plots = [] + + datas = [] if not fast_mode: datas = [ (output_paths, 'jpg', stft_db), diff --git a/batbot/utils.py b/batbot/utils.py index 6cfcf78..952ccc4 100644 --- a/batbot/utils.py +++ b/batbot/utils.py @@ -3,15 +3,15 @@ """ import logging -import os from logging.handlers import TimedRotatingFileHandler +from batbot._config import VERBOSE + DAYS = 21 -VERBOSE = os.getenv('VERBOSE', None) is not None DEFAULT_LOG_LEVEL = logging.DEBUG if VERBOSE else logging.INFO -def init_logging(): +def init_logging() -> logging.Logger: """ Setup Python's built in logging functionality with on-disk logging, and prettier logging with Rich """ diff --git a/docs/batbot.rst b/docs/batbot.rst index 8696a9a..7112612 100644 --- a/docs/batbot.rst +++ b/docs/batbot.rst @@ -18,10 +18,23 @@ Spectrogram :undoc-members: :show-inheritance: +Classifier +---------- + +.. automodule:: batbot.classifier + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: batbot.classifier.dataloader + :members: + :undoc-members: + :show-inheritance: + Pipeline -------- -.. automodule:: batbot.__init__ +.. automodule:: batbot.api :members: :undoc-members: :show-inheritance: diff --git a/docs/cli.rst b/docs/cli.rst index 81b7601..e264a3c 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -1,5 +1,5 @@ BatBot CLI -============ +========== BatBot is the machine learning interface for the Kitware BatAI project. This page specifies the Command Line Interface (CLI) to interact with all of the algorithms and machine learning @@ -11,5 +11,3 @@ models that have been pretrained for inference in a production environment. .. click:: batbot.batbot_cli:cli :prog: batbot :nested: full - -.. include:: environment.rst diff --git a/docs/conf.py b/docs/conf.py index f514f9b..c0db1d7 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -27,6 +27,8 @@ 'sphinx_click', ] +suppress_warnings = ['autosectionlabel.*'] + intersphinx_mapping = { 'rtd': ('https://docs.readthedocs.io/en/stable/', None), 'python': ('https://docs.python.org/3/', None), diff --git a/docs/environment.rst b/docs/environment.rst index 559361d..d01b658 100644 --- a/docs/environment.rst +++ b/docs/environment.rst @@ -1,10 +1,15 @@ Environment Variables --------------------- -The BatBot API and CLI have two environment variables (envars) that allow you to configure global settings +The BatBot API and CLI have environment variables (envars) that allow you to configure global settings and configurations. - - ``VERBOSE`` (default: not set) + - ``BATBOT_VERBOSE`` or the legacy ``VERBOSE`` (default: not set) A verbosity flag that can be set to turn on debug logging. Defaults to "not set", which translates to no debug logging. Setting this value to anything will turn on debug logging (e.g., ``VERBOSE=1``). + - ``BATBOT_CLASSIFIER_CONFIG`` or the legacy ``CLASSIFIER_CONFIG`` (default: ``mobilenet``) + Selects the classifier model configuration. + - ``BATBOT_CLASSIFIER_BATCH_SIZE`` or the legacy ``CLASSIFIER_BATCH_SIZE`` (default: ``10``) + Limits the number of 224-by-224 spectrogram windows sent to ONNX Runtime + in one inference call. diff --git a/docs/index.rst b/docs/index.rst index 606784e..17a2175 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -19,6 +19,7 @@ Contents Home batbot cli + publishing Indices and tables ------------------ diff --git a/docs/publishing.rst b/docs/publishing.rst new file mode 100644 index 0000000..d005466 --- /dev/null +++ b/docs/publishing.rst @@ -0,0 +1,28 @@ +PyPI Trusted Publishing +======================= + +BatBot's ``python-publish.yaml`` workflow publishes tagged releases with +OpenID Connect (OIDC). Trusted publishing exchanges GitHub's short-lived OIDC +identity for a temporary PyPI credential, so the repository does not need a +long-lived PyPI API token. + +Publishing a release +-------------------- + +#. Update ``batbot/_version.py`` and merge the change into ``main``. +#. Wait for the test and distribution workflows to pass. +#. Create and push a matching semantic-version tag, for example: + + .. code-block:: bash + + git tag v0.3.0 + git push origin v0.3.0 + +#. Approve the ``pypi`` GitHub environment deployment if protection rules + require it. +#. Confirm that the ``Publish to PyPI`` job completed and verify the files on + `PyPI's BatBot page `_. + +The publish job needs ``permissions: id-token: write`` and the +``pypa/gh-action-pypi-publish`` step must omit ``password``. Those settings are +already present in this repository's workflow. diff --git a/examples/plot_classifier_performance.py b/examples/plot_classifier_performance.py new file mode 100755 index 0000000..811d388 --- /dev/null +++ b/examples/plot_classifier_performance.py @@ -0,0 +1,336 @@ +#!/usr/bin/env python +"""Evaluate BatBot on labeled WAV or JPG data and plot classifier performance. + +Example: + python examples/plot_classifier_performance.py ./validation --output performance.png + +The immediate parent directory of each input is its ground-truth label, for +example ``validation/EPFU/recording.wav`` or ``validation/EPFU/call.jpg``. +WAV, JPG, and JPEG inputs can be mixed in one dataset. Labels must use one of +the species codes embedded in the BatBot ONNX model. +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.patches as patches +import matplotlib.pyplot as plt +import numpy as np +from tqdm import tqdm + +from batbot import classifier + +CUSTOM_LABEL = 'NOISE' +SUPPORTED_EXTENSIONS = frozenset({'.wav', '.jpg', '.jpeg'}) +GENUS_ORDER_SWAPS = [ + (5, 4), + (12, 11), + (11, 10), + (10, 9), + (9, 8), + (30, 29), +] + + +def apply_genus_order(display, confidences, targets, predicted, custom_index): + """Keep labels and model outputs aligned while grouping species by genus.""" + display = list(display) + confidences = confidences.copy() + targets = targets.copy() + predicted = predicted.copy() + + for first, second in GENUS_ORDER_SWAPS: + assert custom_index not in [first, second] + display[first], display[second] = display[second], display[first] + confidences[:, [first, second]] = confidences[:, [second, first]] + + first_temporary = 100 + first + second_temporary = 100 + second + + targets[targets == first] = first_temporary + targets[targets == second] = second_temporary + targets[targets == first_temporary] = second + targets[targets == second_temporary] = first + + predicted[predicted == first] = first_temporary + predicted[predicted == second] = second_temporary + predicted[predicted == first_temporary] = second + predicted[predicted == second_temporary] = first + + return display, confidences, targets, predicted + + +def shade_regions(display, axis, plot): + """Shade errors between two-letter taxonomic groups.""" + aliases = {'EUMA': 'ETMA', 'LANO': 'L0N0', 'NYHU': 'NXHU'} + grouped = [aliases.get(label, label) for label in display] + regions = sorted({label[:2] for label in grouped}) + errors = 0 + for region in regions: + indices = [index for index, value in enumerate(grouped) if value[:2] == region] + minimum, maximum = min(indices), max(indices) + if minimum > 0: + errors += plot.confusion_matrix[minimum : maximum + 1, :minimum].sum() + axis.add_patch( + patches.Rectangle( + (-0.48, minimum - 0.52), + minimum, + len(indices), + edgecolor='none', + facecolor=(1.0, 0.0, 0.0, 0.2), + ) + ) + if maximum < len(display) - 1: + errors += plot.confusion_matrix[minimum : maximum + 1, maximum + 1 :].sum() + axis.add_patch( + patches.Rectangle( + (maximum + 1 - 0.48, minimum - 0.52), + len(display) - maximum - 1, + len(indices), + edgecolor='none', + facecolor=(1.0, 0.0, 0.0, 0.2), + ) + ) + return errors / max(1, plot.confusion_matrix.sum()) + + +def discover_inputs(data_path): + """Return supported WAV and JPG inputs in deterministic order.""" + return sorted( + path + for path in data_path.rglob('*') + if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS + ) + + +def run_predictions(paths, cache_path=None, batch_size=classifier.BATCH_SIZE, num_workers=1): + if cache_path is not None and cache_path.exists(): + with cache_path.open() as cache_file: + cached = json.load(cache_file) + predictions = cached['results'] + cached_paths = [result['path'] for result in predictions] + if cached_paths != [str(path) for path in paths]: + raise ValueError('Prediction cache does not match the discovered input files') + return predictions + + runner = classifier.Classifier(batch_size=batch_size, num_workers=num_workers) + predictions = [] + for path in tqdm(paths, desc='Classifying inputs'): + suffix = path.suffix.lower() + if suffix == '.wav': + predictions.append(runner.classify_wav(path)) + elif suffix in SUPPORTED_EXTENSIONS: + predictions.append(runner.classify(path)[0]) + else: + raise ValueError(f'Unsupported classifier input: {path}') + + if cache_path is not None: + cache_path.parent.mkdir(parents=True, exist_ok=True) + with cache_path.open('w') as cache_file: + json.dump({'results': predictions}, cache_file, indent=2) + return predictions + + +def plot_confusion(axis, targets, predicted, labels, display, normalize, title): + from sklearn import metrics + + plot = metrics.ConfusionMatrixDisplay.from_predictions( + targets, + predicted, + labels=labels, + display_labels=display, + normalize=normalize, + xticks_rotation='vertical', + ax=axis, + values_format='d' if normalize is None else '0.02f', + text_kw={'fontsize': 5.0}, + ) + for text in plot.text_.ravel(): + if text.get_text() in {'0', '0.00'}: + text.set_text('') + shade_regions(display, axis, plot) + axis.set_title(title, y=1.04) + return plot + + +def plot_performance(paths, predictions, output_path): + from sklearn import metrics + + classes = classifier.resolve_config().classes + backward = {label: index for index, label in enumerate(classes)} + unknown = sorted({path.parent.name for path in paths} - set(classes)) + if unknown: + raise ValueError('Unknown ground-truth labels: {}'.format(', '.join(unknown))) + + targets = np.asarray([backward[path.parent.name] for path in paths]) + confidences = np.asarray( + [[prediction['scores'][label] for label in classes] for prediction in predictions] + ) + predicted = np.argmax(confidences, axis=1) + custom_index = backward.get(CUSTOM_LABEL) + display, confidences, targets, predicted = apply_genus_order( + classes, + confidences, + targets, + predicted, + custom_index, + ) + labels = list(range(len(classes))) + + accuracy = metrics.accuracy_score(targets, predicted) + top_scores = {} + for top_k in [2, 3, 5]: + top_scores[top_k] = metrics.top_k_accuracy_score( + targets, + confidences, + k=top_k, + labels=labels, + ) + mcc = metrics.matthews_corrcoef(targets, predicted) + stats = ( + 'Top-1 = {:0.2f}% | Top-2 = {:0.2f}% | Top-3 = {:0.2f}% | ' + 'Top-5 = {:0.2f}% | MCC = {:0.4f}' + ).format( + 100 * accuracy, + 100 * top_scores[2], + 100 * top_scores[3], + 100 * top_scores[5], + mcc, + ) + + dataset_labels = {path.parent.name for path in paths} + has_noise_examples = CUSTOM_LABEL in dataset_labels and len(dataset_labels) > 1 + if has_noise_examples: + figure, axes = plt.subplots(2, 3, figsize=(45, 28)) + confusion_axes = axes[0] + else: + figure, confusion_axes = plt.subplots(1, 3, figsize=(45, 15)) + + absolute_plot = plot_confusion( + confusion_axes[0], + targets, + predicted, + labels, + display, + None, + f'Confusion Matrix (counts)\n{stats}', + ) + column_totals = absolute_plot.confusion_matrix.sum(axis=0) + row_totals = absolute_plot.confusion_matrix.sum(axis=1) + absolute_plot.ax_.set_xticklabels( + [f'({value}) {label}' for value, label in zip(column_totals, display)] + ) + absolute_plot.ax_.set_yticklabels( + [f'({value}) {label}' for value, label in zip(row_totals, display)] + ) + plot_confusion( + confusion_axes[1], + targets, + predicted, + labels, + display, + 'true', + f'Confusion Matrix (true-normalized)\n{stats}', + ) + plot_confusion( + confusion_axes[2], + targets, + predicted, + labels, + display, + 'pred', + f'Confusion Matrix (prediction-normalized)\n{stats}', + ) + + if has_noise_examples: + noise_index = display.index(CUSTOM_LABEL) + noise_targets = targets == noise_index + noise_scores = confidences[:, noise_index] + precision, recall, pr_thresholds = metrics.precision_recall_curve( + noise_targets, noise_scores + ) + average_precision = metrics.average_precision_score(noise_targets, noise_scores) + pr_operating_index = np.argmin( + np.linalg.norm(np.vstack((1.0 - precision, 1.0 - recall)), axis=0) + ) + pr_threshold = ( + pr_thresholds[min(pr_operating_index, len(pr_thresholds) - 1)] + if len(pr_thresholds) + else 0.5 + ) + metrics.PrecisionRecallDisplay(precision=precision, recall=recall).plot( + ax=axes[1, 0], + name='MobileNet (AP={:0.3f}, threshold={:0.3f})'.format( + average_precision, pr_threshold + ), + ) + axes[1, 0].set_title('NOISE precision-recall curve') + + false_positive_rate, true_positive_rate, roc_thresholds = metrics.roc_curve( + noise_targets, noise_scores + ) + auc = metrics.roc_auc_score(noise_targets, noise_scores) + roc_operating_index = np.argmin( + np.linalg.norm(np.vstack((false_positive_rate, 1.0 - true_positive_rate)), axis=0) + ) + roc_threshold = roc_thresholds[roc_operating_index] + metrics.RocCurveDisplay( + fpr=false_positive_rate, + tpr=true_positive_rate, + roc_auc=auc, + ).plot( + ax=axes[1, 1], + name=f'MobileNet (AUC={auc:0.3f}, threshold={roc_threshold:0.3f})', + ) + axes[1, 1].set_title('NOISE ROC curve') + + binary_predictions = noise_scores >= roc_threshold + binary_confusion = metrics.confusion_matrix(noise_targets, binary_predictions) + metrics.ConfusionMatrixDisplay( + binary_confusion, + display_labels=['species', 'NOISE'], + ).plot(ax=axes[1, 2], values_format='d') + axes[1, 2].set_title('NOISE operating-point confusion matrix') + + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(str(output_path), dpi=150, bbox_inches='tight') + plt.close(figure) + return stats + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + 'data', + type=Path, + help='Root directory containing labeled WAV, JPG, or JPEG files', + ) + parser.add_argument('--output', type=Path, default=Path('classifier-performance.png')) + parser.add_argument('--cache', type=Path, default=None, help='Optional prediction JSON cache') + parser.add_argument('--batch-size', type=int, default=classifier.BATCH_SIZE) + parser.add_argument( + '--num-workers', + type=int, + default=1, + help='Number of concurrent ONNX inference workers', + ) + args = parser.parse_args() + + paths = discover_inputs(args.data) + if not paths: + parser.error(f'No WAV or JPG files found beneath {args.data}') + + predictions = run_predictions( + paths, + cache_path=args.cache, + batch_size=args.batch_size, + num_workers=args.num_workers, + ) + stats = plot_performance(paths, predictions, args.output) + print(stats) + print(f'Saved performance plot: {args.output}') + + +if __name__ == '__main__': + main() diff --git a/pyproject.toml b/pyproject.toml index 9787c3b..59fcbf3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,143 @@ [build-system] -requires = ["setuptools", "wheel"] +requires = ["setuptools>=77.0.0"] build-backend = "setuptools.build_meta" + +[project] +name = "batbot" +dynamic = ["version"] +description = "Machine Learning app for the Kitware BatAI Project" +readme = {file = "README.rst", content-type = "text/x-rst"} +requires-python = ">=3.11" +license = "Apache-2.0" +license-files = ["LICENSE"] +authors = [ + {name = "Kitware", email = "vision@kitware.com"}, +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "click", + "librosa", + "matplotlib", + "numpy", + "onnxruntime", + "opencv-python-headless", + "pooch", + "pyastar2d>=1.1.1", + "rich", + "scikit-image", + "scipy", + "shapely", + "tqdm", +] + +[project.optional-dependencies] +test = [ + "pytest==8.4.2", + "pytest-cov", + "pytest-random-order", + "xdoctest", +] +docs = [ + "Sphinx>=5,<6", + "sphinx-click", + "sphinx-rtd-theme", + "standard-imghdr", +] +performance = [ + "scikit-learn", +] + +[project.scripts] +batbot = "batbot.batbot_cli:cli" + +[project.urls] +Documentation = "https://batbot.readthedocs.io" +Issues = "https://github.com/Kitware/batbot/issues" +Repository = "https://github.com/Kitware/batbot" + +[tool.setuptools] +include-package-data = true + +[tool.setuptools.dynamic] +version = {attr = "batbot._version.__version__"} + +[tool.setuptools.packages.find] +include = ["batbot*"] + +[tool.setuptools.package-data] +batbot = ["py.typed"] +"batbot.classifier" = ["models/onnx/*.onnx"] + +[tool.pytest.ini_options] +minversion = "8.4.2" +addopts = """ + -v + -p no:doctest + --xdoctest + --xdoctest-style=google + --random-order + --random-order-bucket=global + --cov=batbot + --cov-report=html + -m "not separate" + --durations-min=1.0 + --color=yes + --code-highlight=yes + --show-capture=log + -ra +""" +testpaths = ["batbot", "tests"] +filterwarnings = ["default"] + +[tool.coverage.run] +branch = true +source = ["batbot"] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "# NOCC", + "raise NotImplementedError", + "if __name__ == .__main__.:", +] +precision = 1 +ignore_errors = true +omit = [ + "*/__pycache__/*", + "tests/*", +] + +[tool.coverage.html] +directory = "coverage/html" + +[tool.coverage.xml] +output = "coverage/coverage.xml" + +[tool.black] +line-length = 100 +target-version = ["py311"] +skip-string-normalization = true + +[tool.isort] +profile = "black" +line_length = 100 +multi_line_output = 3 +include_trailing_comma = true +force_grid_wrap = 0 +use_parentheses = true +ensure_newline_before_comments = true + +[tool.flake8] +max-line-length = 100 +extend-ignore = ["E203", "E501"] +exclude = [".git"] diff --git a/requirements/documentation.txt b/requirements/documentation.txt deleted file mode 100644 index 4eee2cc..0000000 --- a/requirements/documentation.txt +++ /dev/null @@ -1,4 +0,0 @@ -Sphinx>=5,<6 -sphinx_click -sphinx_rtd_theme -standard-imghdr diff --git a/requirements/optional.txt b/requirements/optional.txt deleted file mode 100644 index 9891401..0000000 --- a/requirements/optional.txt +++ /dev/null @@ -1,25 +0,0 @@ -black -codecov -coverage -flake8 -# gradio -ipython -isort -# onnx -# onnxruntime -pre-commit -pytest -pytest-benchmark[histogram] -pytest-cov -pytest-pep8 -pytest-profiling -pytest-random-order -pytest-sugar -pytest-xdist -pyupgrade -pyyaml -rstcheck -rstcheck[sphinx] -# torch -# torchvision -xdoctest diff --git a/requirements/runtime.txt b/requirements/runtime.txt deleted file mode 100644 index a02e985..0000000 --- a/requirements/runtime.txt +++ /dev/null @@ -1,12 +0,0 @@ -click -librosa -matplotlib -numpy -opencv-python-headless -pooch -pyastar2d>=1.1.1 -rich -scikit-image -scipy -shapely -tqdm diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index 8854b72..0000000 --- a/setup.cfg +++ /dev/null @@ -1,66 +0,0 @@ -[metadata] -name = batbot -description = Machine Learning app for the Kitware BatAI Project -version = attr: batbot.VERSION -long_description = file: README.rst -long_description_content_type = text/x-rst -url = https://github.com/Kitware/batbot -author = Kitware -author_email = vision@kitware.com -license = MIT -license_file = LICENSE -project_urls = - Documentation = https://batbot.readthedocs.io - Source = https://github.com/Kitware/batbot - -[options] -packages = find: -platforms = any -include_package_data = true -install_requires = - click - librosa - matplotlib - numpy - opencv-python-headless - pooch - pyastar2d>=1.1.1 - rich - scikit-image - scipy - shapely - tqdm -python_requires = >=3.7 - -[options.entry_points] -console_scripts = - batbot = batbot.batbot_cli:cli - -[tool:pytest] -minversion = 5.4 -addopts = -v -p no:doctest --xdoctest --xdoctest-style=google --random-order --random-order-bucket=global --cov=./ --cov-report html -m "not separate" --durations-min=1.0 --color=yes --code-highlight=yes --show-capture=log -ra -testpaths = - batbot - tests -filterwarnings = - default - -[options.extras_require] -test = - pytest >= 6.2.2 - pycodestyle - pytest-cov -all = - %(test)s - -[flake8] -exclude = .git - -[tool:isort] -profile = black -line_length = 100 -multi_line_output = 3 -include_trailing_comma = true -force_grid_wrap = 0 -use_parentheses = true -ensure_newline_before_comments = true diff --git a/setup.py b/setup.py deleted file mode 100755 index 793f7bf..0000000 --- a/setup.py +++ /dev/null @@ -1,5 +0,0 @@ -#!/usr/bin/env python -import setuptools - -if __name__ == '__main__': - setuptools.setup() diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..51bedf3 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,213 @@ +from pathlib import Path + +import pooch +import pytest +import tqdm + +import batbot.api as api +from batbot import classifier, spectrogram + + +def _result(path): + return { + 'path': str(path), + 'label': 'EPFU', + 'confidence': 0.75, + 'window_count': 1, + 'top': [{'label': 'EPFU', 'confidence': 0.75}], + 'scores': {'EPFU': 0.75}, + } + + +def test_fetch_delegates_to_classifier(monkeypatch): + calls = [] + monkeypatch.setattr( + classifier, + 'fetch', + lambda **kwargs: calls.append(kwargs) or '/models/model.onnx', + ) + + assert api.fetch(pull=True, config='mobilenet') == '/models/model.onnx' + assert calls == [{'pull': True, 'config': 'mobilenet'}] + + +def test_pipeline_forwards_spectrogram_options(monkeypatch): + calls = [] + + def fake_compute(filepath, **kwargs): + calls.append((filepath, kwargs)) + return ['one.png'], ['one-compressed.png'], 'one.json', {'ignored': True} + + monkeypatch.setattr(spectrogram, 'compute', fake_compute) + + output = api.pipeline( + Path('recording.wav'), + out_file_stem='custom', + output_folder='output', + fast_mode=True, + force_overwrite=True, + quiet=True, + plot_uncompressed_amplitude=True, + include_original_sr=True, + time_buffer_ms=2.5, + debug=True, + ) + + assert output == (['one.png'], ['one-compressed.png'], 'one.json') + assert calls == [ + ( + 'recording.wav', + { + 'out_file_stem': 'custom', + 'output_folder': 'output', + 'fast_mode': True, + 'force_overwrite': True, + 'quiet': True, + 'plot_uncompressed_amplitude': True, + 'include_original_sr': True, + 'time_buffer_ms': 2.5, + 'debug': True, + }, + ) + ] + + +def test_pipeline_multi_wrapper_collects_outputs_and_failures(monkeypatch): + def fake_pipeline(filepath, **kwargs): + if filepath == 'bad.wav': + raise ValueError('bad recording') + return [f'{filepath}.png'], [f'{filepath}.compressed.png'], f'{filepath}.json' + + monkeypatch.setattr(api, 'pipeline', fake_pipeline) + + output, compressed, metadata, failures = api.pipeline_multi_wrapper( + ['good.wav', 'bad.wav'], + out_file_stems=['good', 'bad'], + quiet=True, + ) + + assert output == ['good.wav.png'] + assert compressed == ['good.wav.compressed.png'] + assert metadata == ['good.wav.json'] + assert failures[0][0] == 'bad.wav' + assert str(failures[0][1]) == 'bad recording' + + +def test_pipeline_multi_wrapper_builds_default_stems_and_sets_lock(monkeypatch): + lock = object() + captured = [] + monkeypatch.setattr( + api, + 'pipeline', + lambda filepath, **kwargs: captured.append((filepath, kwargs)) or ([], [], None), + ) + monkeypatch.setattr(tqdm.tqdm, 'set_lock', lambda value: captured.append(('lock', value))) + + output = api.pipeline_multi_wrapper(['one.wav'], tqdm_lock=lock, quiet=True) + + assert output == ([], [], [None], []) + assert captured[0] == ('lock', lock) + assert captured[1][1]['out_file_stem'] is None + + +def test_pipeline_multi_wrapper_validates_stem_count(): + with pytest.raises(ValueError, match='different length'): + api.pipeline_multi_wrapper(['one.wav'], out_file_stems=[]) + + +def test_parallel_pipeline_validates_work(monkeypatch): + assert api.parallel_pipeline([]) is None + with pytest.raises(ValueError, match='same length'): + api.parallel_pipeline([['one.wav']], out_stem_chunks=[]) + with pytest.raises(ValueError, match='num_workers'): + api.parallel_pipeline([['one.wav']], num_workers=0) + + +def test_parallel_pipeline_combines_threaded_results(monkeypatch): + class LockManager: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def Lock(self): + return object() + + def fake_wrapper(filepaths, **kwargs): + path = filepaths[0] + return [f'{path}.png'], [f'{path}.compressed.png'], [f'{path}.json'], [] + + monkeypatch.setattr(api, 'Manager', LockManager) + monkeypatch.setattr(api, 'pipeline_multi_wrapper', fake_wrapper) + + output = api.parallel_pipeline( + [['one.wav'], ['two.wav']], + num_workers=2, + threaded=True, + quiet=True, + ) + + assert output is not None + paths, compressed, metadata, failures = output + assert sorted(paths) == ['one.wav.png', 'two.wav.png'] + assert sorted(compressed) == ['one.wav.compressed.png', 'two.wav.compressed.png'] + assert sorted(metadata) == ['one.wav.json', 'two.wav.json'] + assert failures == [] + + +def test_batch_reuses_one_classifier(monkeypatch): + instances = [] + + class FakeClassifier: + def __init__(self, config=None, num_workers=1): + self.config = config + self.num_workers = num_workers + self.paths = [] + instances.append(self) + + def classify_wav(self, filepath): + self.paths.append(filepath) + return _result(filepath) + + monkeypatch.setattr(classifier, 'Classifier', FakeClassifier) + + results = api.batch( + ['one.wav', Path('two.wav')], + config='mobilenet', + clean=False, + num_workers=3, + ) + + assert len(instances) == 1 + assert instances[0].config == 'mobilenet' + assert instances[0].num_workers == 3 + assert instances[0].paths == ['one.wav', Path('two.wav')] + assert [result['path'] for result in results] == ['one.wav', 'two.wav'] + + +def test_example_downloads_missing_wav_and_runs_pipeline(monkeypatch, tmp_path, capsys): + downloaded = tmp_path / 'downloaded.wav' + downloaded.touch() + pipeline_calls = [] + retrieve_calls = [] + times = iter([10.0, 12.5]) + monkeypatch.setattr(api, 'PROJECT_ROOT', tmp_path / 'missing-project') + monkeypatch.setattr( + pooch, + 'retrieve', + lambda **kwargs: retrieve_calls.append(kwargs) or str(downloaded), + ) + monkeypatch.setattr( + api, + 'pipeline', + lambda filepath, **kwargs: pipeline_calls.append((filepath, kwargs)) or ([], [], None), + ) + monkeypatch.setattr(api.time, 'time', lambda: next(times)) + + api.example() + + assert retrieve_calls[0]['known_hash'].startswith('sha256:') + assert pipeline_calls[0][0] == downloaded + assert pipeline_calls[0][1]['out_file_stem'] == 'output/downloaded' + assert '2.5 seconds' in capsys.readouterr().out diff --git a/tests/test_batbot.py b/tests/test_batbot.py index 73bfea8..da2c769 100644 --- a/tests/test_batbot.py +++ b/tests/test_batbot.py @@ -1,5 +1,27 @@ +import subprocess +import sys + import batbot +def test_import_is_lightweight(): + result = subprocess.run( + [ + sys.executable, + '-c', + ( + 'import sys; import batbot; ' + 'assert "batbot.classifier" not in sys.modules; ' + 'assert "batbot.spectrogram" not in sys.modules' + ), + ], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + def test_example(): batbot.example() diff --git a/tests/test_classifier.py b/tests/test_classifier.py new file mode 100644 index 0000000..db3eafa --- /dev/null +++ b/tests/test_classifier.py @@ -0,0 +1,217 @@ +from dataclasses import FrozenInstanceError +from pathlib import Path +from threading import Barrier, Lock + +import cv2 +import numpy as np +import pooch +import pytest + +from batbot import classifier +from batbot.classifier import dataloader, inference, model + + +class _ModelValue: + def __init__(self, name='input', shape=None): + self.name = name + self.shape = shape + + +class _ModelMetadata: + custom_metadata_map = {} + + +class FakeSession: + def __init__(self): + self.batch_sizes = [] + + def get_inputs(self): + return [_ModelValue(shape=['batch_size', 224, 224, 3])] + + def get_outputs(self): + return [_ModelValue(name='output', shape=['batch_size', len(classifier.CLASSES)])] + + def get_modelmeta(self): + return _ModelMetadata() + + def run(self, output_names, inputs): + batch = inputs['input'] + self.batch_sizes.append(len(batch)) + scores = np.zeros((len(batch), len(classifier.CLASSES)), dtype=np.float32) + scores[:, classifier.CLASSES.index('EPFU')] = 0.75 + scores[:, classifier.CLASSES.index('NOISE')] = 0.25 + return [scores] + + +def test_prepare_image_matches_model_shape(): + image = np.zeros((300, 1200, 3), dtype=np.uint8) + windows = dataloader._prepare_image(image) + + assert windows.shape == (3, 224, 224, 3) + assert windows.dtype == np.uint8 + assert windows.flags['C_CONTIGUOUS'] + + +def test_prepare_image_pads_narrow_spectrogram(): + image = np.full((237, 100, 3), 255, dtype=np.uint8) + windows = dataloader._prepare_image(image) + + assert windows.shape == (1, 224, 224, 3) + assert np.all(windows[:, :, -1, :] == 0) + + +def test_predict_batches_and_averages(monkeypatch): + session = FakeSession() + monkeypatch.setattr(inference, 'fetch', lambda **kwargs: 'model.onnx') + monkeypatch.setattr(inference, '_create_session', lambda *args, **kwargs: session) + windows = np.zeros((5, 224, 224, 3), dtype=np.uint8) + + output = classifier.post(classifier.predict(iter([(windows, 'mobilenet')]), batch_size=2)) + + assert session.batch_sizes == [2, 2, 1] + assert output[0]['EPFU'] == 0.75 + assert output[0]['NOISE'] == 0.25 + + +def test_predict_runs_multiple_spectrograms_concurrently_and_preserves_order(): + class ConcurrentSession(FakeSession): + def __init__(self): + super().__init__() + self.barrier = Barrier(2) + self.lock = Lock() + self.active = 0 + self.maximum_active = 0 + self.calls = 0 + + def run(self, output_names, inputs): + batch = inputs['input'] + with self.lock: + self.active += 1 + self.maximum_active = max(self.maximum_active, self.active) + self.calls += 1 + call_number = self.calls + try: + if call_number <= 2: + self.barrier.wait(timeout=5) + scores = np.zeros((len(batch), len(classifier.CLASSES)), dtype=np.float32) + scores[:, 0] = batch[0, 0, 0, 0] + return [scores] + finally: + with self.lock: + self.active -= 1 + + session = ConcurrentSession() + inputs = [(np.full((1, 224, 224, 3), value, dtype=np.uint8), 'mobilenet') for value in range(4)] + + outputs = list( + classifier.predict( + inputs, + sessions={'mobilenet': session}, + num_workers=2, + total=len(inputs), + ) + ) + + assert session.maximum_active == 2 + assert [prediction[0, 0] for prediction, _ in outputs] == [0.0, 1.0, 2.0, 3.0] + assert [config for _, config in outputs] == ['mobilenet'] * len(inputs) + + +def test_classifier_reuses_its_session(monkeypatch, tmp_path): + path = tmp_path / 'spectrogram.jpg' + cv2.imwrite(str(path), np.zeros((300, 700, 3), dtype=np.uint8)) + session = FakeSession() + created = [] + + monkeypatch.setattr(inference, 'fetch', lambda **kwargs: 'model.onnx') + + def create_session(*args, **kwargs): + created.append(True) + return session + + monkeypatch.setattr(inference, '_create_session', create_session) + runner = classifier.Classifier(top_k=2) + + first = runner.classify(path)[0] + second = runner.classify(path)[0] + + assert len(created) == 1 + assert first == second + assert first['top'] == [ + {'label': 'EPFU', 'confidence': 0.75}, + {'label': 'NOISE', 'confidence': 0.25}, + ] + + +def test_config_is_immutable(): + config = classifier.resolve_config('mobilenet') + + with pytest.raises(FrozenInstanceError): + config.key = 'changed' + + +def test_fetch_uses_verified_bundled_model(): + model_path = Path(classifier.fetch()) + + assert model_path.name == classifier.MODEL_NAME + assert model_path.exists() + assert pooch.file_hash(str(model_path), alg='sha256') == classifier.MODEL_HASH + + +def test_fetch_uses_versioned_mirror_cache(monkeypatch, tmp_path): + downloaded = tmp_path / classifier.MODEL_NAME + downloaded.write_bytes(b'model') + calls = [] + model._fetch_cached.cache_clear() + monkeypatch.setattr(model, '_bundled_model', lambda config: None) + monkeypatch.setattr(model, '_cache_directory', lambda: tmp_path / 'models' / '0.2.0') + + def retrieve(**kwargs): + calls.append(kwargs) + return str(downloaded) + + monkeypatch.setattr(model.pooch, 'retrieve', retrieve) + + assert classifier.fetch() == str(downloaded) + assert calls[0]['known_hash'] == f'sha256:{classifier.MODEL_HASH}' + assert calls[0]['path'] == tmp_path / 'models' / '0.2.0' + model._fetch_cached.cache_clear() + + +def test_real_model_inference(tmp_path): + path = tmp_path / 'spectrogram.png' + cv2.imwrite(str(path), np.zeros((300, 700, 3), dtype=np.uint8)) + + results = classifier.Classifier(top_k=3, num_workers=2).classify([path, path]) + result = results[0] + + assert results[0] == results[1] + assert result['label'] in classifier.CLASSES + assert len(result['scores']) == len(classifier.CLASSES) + assert len(result['top']) == 3 + assert result['window_count'] > 0 + + +def test_discover_inputs_and_summary(tmp_path): + species = tmp_path / 'EPFU' + species.mkdir() + wav = species / 'one.WAV' + image = species / 'one.jpg' + ignored = species / 'notes.txt' + for path in [wav, image, ignored]: + path.touch() + + discovered = classifier.discover_inputs([tmp_path]) + summary = classifier.summarize( + [ + {'label': 'EPFU', 'confidence': 0.8}, + {'label': 'NOISE', 'confidence': 0.6}, + {'path': 'bad.wav', 'error': 'bad input'}, + ] + ) + + assert discovered == sorted([str(image), str(wav)]) + assert summary['species_counts'] == {'EPFU': 1} + assert summary['noise_count'] == 1 + assert summary['classified'] == 2 + assert summary['failed'] == 1 diff --git a/tests/test_classifier_cli.py b/tests/test_classifier_cli.py new file mode 100644 index 0000000..a5cd72a --- /dev/null +++ b/tests/test_classifier_cli.py @@ -0,0 +1,82 @@ +import json + +import cv2 +import numpy as np +from click.testing import CliRunner + +from batbot import classifier +from batbot.batbot_cli import classify, classify_bulk, classify_wav + + +def _result(path): + return { + 'path': str(path), + 'label': 'EPFU', + 'confidence': 0.75, + 'window_count': 1, + 'top': [{'label': 'EPFU', 'confidence': 0.75}], + 'scores': {'EPFU': 0.75}, + } + + +def _bulk(path): + result = _result(path) + return {'results': [result], 'summary': classifier.summarize([result])} + + +def test_classify_cli_writes_json(monkeypatch, tmp_path): + image = tmp_path / 'spectrogram.png' + output = tmp_path / 'result.json' + cv2.imwrite(str(image), np.zeros((300, 700, 3), dtype=np.uint8)) + calls = [] + monkeypatch.setattr( + classifier, + 'classify', + lambda paths, **kwargs: calls.append(kwargs) or [_result(paths[0])], + ) + + invocation = CliRunner().invoke( + classify, + [str(image), '--output', str(output), '--num-workers', '3'], + ) + + assert invocation.exit_code == 0, invocation.output + assert json.loads(output.read_text())['results'][0]['label'] == 'EPFU' + assert calls[0]['num_workers'] == 3 + + +def test_classify_wav_cli(monkeypatch, tmp_path): + wav = tmp_path / 'recording.wav' + wav.touch() + calls = [] + monkeypatch.setattr( + classifier, + 'classify_bulk', + lambda paths, **kwargs: calls.append(kwargs) or _bulk(paths[0]), + ) + + invocation = CliRunner().invoke(classify_wav, [str(wav), '--num-workers', '4']) + + assert invocation.exit_code == 0, invocation.output + assert json.loads(invocation.output)['summary']['species_counts'] == {'EPFU': 1} + assert calls[0]['num_workers'] == 4 + + +def test_classify_bulk_cli(monkeypatch, tmp_path): + wav = tmp_path / 'recording.wav' + wav.touch() + calls = [] + monkeypatch.setattr( + classifier, + 'classify_bulk', + lambda paths, **kwargs: calls.append(kwargs) or _bulk(paths[0]), + ) + + invocation = CliRunner().invoke( + classify_bulk, + [str(tmp_path), '--input-type', 'wav', '--num-workers', '5'], + ) + + assert invocation.exit_code == 0, invocation.output + assert json.loads(invocation.output)['results'][0]['path'] == str(tmp_path) + assert calls[0]['num_workers'] == 5 diff --git a/tests/test_classifier_performance.py b/tests/test_classifier_performance.py new file mode 100644 index 0000000..591a227 --- /dev/null +++ b/tests/test_classifier_performance.py @@ -0,0 +1,79 @@ +import json + +import pytest + +from examples import plot_classifier_performance as performance + + +def test_discover_inputs_accepts_wav_jpg_and_jpeg_case_insensitively(tmp_path): + label = tmp_path / 'EPFU' + label.mkdir() + expected = [label / 'call.JPG', label / 'call.jpeg', label / 'recording.wav'] + ignored = [label / 'notes.txt', label / 'spectrogram.png'] + for path in [*expected, *ignored]: + path.touch() + + assert performance.discover_inputs(tmp_path) == sorted(expected) + + +def test_run_predictions_dispatches_mixed_inputs_through_one_classifier(monkeypatch, tmp_path): + wav = tmp_path / 'EPFU' / 'recording.wav' + jpg = tmp_path / 'NOISE' / 'call.JPG' + jpeg = tmp_path / 'EPFU' / 'call.jpeg' + cache = tmp_path / 'predictions.json' + calls = [] + + class FakeClassifier: + def __init__(self, **kwargs): + calls.append(('init', kwargs)) + + def classify_wav(self, path): + calls.append(('wav', path)) + return {'path': str(path), 'source': 'wav'} + + def classify(self, path): + calls.append(('image', path)) + return [{'path': str(path), 'source': 'image'}] + + monkeypatch.setattr(performance.classifier, 'Classifier', FakeClassifier) + + predictions = performance.run_predictions( + [wav, jpg, jpeg], + cache_path=cache, + batch_size=6, + num_workers=3, + ) + + assert calls == [ + ('init', {'batch_size': 6, 'num_workers': 3}), + ('wav', wav), + ('image', jpg), + ('image', jpeg), + ] + assert [prediction['source'] for prediction in predictions] == ['wav', 'image', 'image'] + assert json.loads(cache.read_text()) == {'results': predictions} + + +def test_run_predictions_reuses_matching_cache_without_loading_model(monkeypatch, tmp_path): + path = tmp_path / 'EPFU' / 'call.jpg' + cache = tmp_path / 'predictions.json' + expected = [{'path': str(path), 'label': 'EPFU'}] + cache.write_text(json.dumps({'results': expected})) + monkeypatch.setattr( + performance.classifier, + 'Classifier', + lambda **kwargs: pytest.fail('Classifier should not be created for a valid cache'), + ) + + assert performance.run_predictions([path], cache_path=cache) == expected + + +def test_run_predictions_rejects_unsupported_explicit_input(monkeypatch, tmp_path): + class FakeClassifier: + def __init__(self, **kwargs): + pass + + monkeypatch.setattr(performance.classifier, 'Classifier', FakeClassifier) + + with pytest.raises(ValueError, match='Unsupported classifier input'): + performance.run_predictions([tmp_path / 'call.png']) diff --git a/tests/test_classifier_units.py b/tests/test_classifier_units.py new file mode 100644 index 0000000..b16148a --- /dev/null +++ b/tests/test_classifier_units.py @@ -0,0 +1,535 @@ +import hashlib +import importlib +import json +from contextlib import contextmanager +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from batbot import classifier +from batbot.classifier import bulk +from batbot.classifier import config as config_module +from batbot.classifier import dataloader, inference, model + +MOBILENET = classifier.resolve_config('mobilenet') + + +def _result(path, label='EPFU', confidence=0.75, window_count=1): + scores = {'EPFU': confidence, 'NOISE': 1.0 - confidence} + return { + 'path': str(path), + 'label': label, + 'confidence': confidence, + 'window_count': window_count, + 'top': [{'label': label, 'confidence': confidence}], + 'scores': scores, + } + + +class _Session: + def __init__(self, inputs=None, outputs=None, metadata=None): + self._inputs = ( + [SimpleNamespace(name='input', shape=['batch', 224, 224, 3])] + if inputs is None + else inputs + ) + self._outputs = [SimpleNamespace(name='output')] if outputs is None else outputs + self._metadata = {} if metadata is None else metadata + + def get_inputs(self): + return self._inputs + + def get_outputs(self): + return self._outputs + + def get_modelmeta(self): + return SimpleNamespace(custom_metadata_map=self._metadata) + + def run(self, output_names, inputs): + batch = inputs['input'] + return [np.zeros((len(batch), len(classifier.CLASSES)), dtype=np.float32)] + + +class _Runner: + def __init__(self): + self.calls = [] + + def classify(self, paths, top_k=5): + paths = [str(path) for path in paths] + self.calls.append((paths, top_k)) + return [_result(path) for path in paths] + + +@pytest.mark.parametrize( + 'image', + [None, np.zeros((20, 20), dtype=np.uint8), np.zeros((20, 20, 4), dtype=np.uint8)], +) +def test_prepare_image_rejects_non_color_images(image): + with pytest.raises(ValueError, match='three-channel'): + dataloader._prepare_image(image) + + +@pytest.mark.parametrize( + ('argument', 'value', 'message'), + [('input_size', 0, 'input_size'), ('window_stride', 0, 'window_stride')], +) +def test_prepare_image_rejects_nonpositive_geometry(argument, value, message): + kwargs = {argument: value} + + with pytest.raises(ValueError, match=message): + dataloader._prepare_image(np.zeros((20, 20, 3), dtype=np.uint8), **kwargs) + + +def test_load_image_reports_missing_file(tmp_path): + with pytest.raises(OSError, match='Unable to load spectrogram'): + dataloader._load_image(tmp_path / 'missing.png') + + +def test_dataset_applies_sample_and_target_transforms(): + dataset = dataloader.ImageFilePathList( + ['one.png', 'two.png'], + targets=['NOISE', 'EPFU'], + transform=lambda image: image + 1, + target_transform=str.lower, + ) + dataset.loader = lambda path: np.zeros((2, 2, 3), dtype=np.uint8) + + sample, target = dataset[0] + + assert len(dataset) == 2 + assert np.all(sample == 1) + assert target == 'noise' + assert dataset.classes == ['EPFU', 'NOISE'] + assert dataset.class_to_idx == {'EPFU': 0, 'NOISE': 1} + + +def test_dataset_without_targets_and_mismatched_targets(): + dataset = dataloader.ImageFilePathList(['one.png']) + dataset.loader = lambda path: np.zeros((2, 2, 3), dtype=np.uint8) + + assert len(dataset[0]) == 1 + assert dataset.classes is None + assert dataset.class_to_idx is None + with pytest.raises(ValueError, match='same length'): + dataloader.ImageFilePathList(['one.png'], targets=[]) + + +def test_resolve_config_normalizes_names_and_rejects_unknown_values(): + selected = classifier.resolve_config(' MOBILENET ') + + assert classifier.resolve_config(selected) is selected + assert classifier.resolve_config(None) is selected + with pytest.raises(ValueError, match='choose from mobilenet'): + classifier.resolve_config('unknown') + + +def test_invalid_default_config_fails_during_configuration_load(monkeypatch): + try: + with monkeypatch.context() as environment: + environment.setenv('BATBOT_CLASSIFIER_CONFIG', 'unknown') + with pytest.raises(ValueError, match='Unknown classifier configuration'): + importlib.reload(config_module) + finally: + importlib.reload(config_module) + + +@pytest.mark.parametrize( + ('session', 'message'), + [ + ( + _Session(inputs=[SimpleNamespace(shape=[]), SimpleNamespace(shape=[])]), + 'one input', + ), + (_Session(outputs=[]), 'one input'), + (_Session(inputs=[SimpleNamespace(shape=['batch', 3, 224, 224])]), 'input shape'), + ], +) +def test_validate_session_rejects_incompatible_graphs(session, message): + with pytest.raises(ValueError, match=message): + inference._validate_session(session, MOBILENET) + + +def test_validate_session_checks_embedded_labels(): + labels = { + 'labels': json.dumps( + {'forward': {str(index): label for index, label in enumerate(classifier.CLASSES)}} + ) + } + inference._validate_session(_Session(metadata=labels), MOBILENET) + + labels['labels'] = json.dumps({'forward': {'0': 'WRONG'}}) + with pytest.raises(ValueError, match='labels do not match'): + inference._validate_session(_Session(metadata=labels), MOBILENET) + + +def test_predict_rejects_invalid_batches_and_empty_windows(): + with pytest.raises(ValueError, match='batch_size'): + list(inference.predict([], batch_size=0)) + with pytest.raises(ValueError, match='num_workers'): + list(inference.predict([], num_workers=0)) + + empty = np.empty((0, 224, 224, 3), dtype=np.uint8) + with pytest.raises(ValueError, match='no image windows'): + list( + inference.predict( + [(empty, 'mobilenet')], + sessions={'mobilenet': _Session()}, + ) + ) + + +def test_pre_rejects_invalid_batch_size(): + with pytest.raises(ValueError, match='batch_size'): + list(inference.pre([], batch_size=0)) + + +def test_create_session_preserves_explicit_providers(monkeypatch): + import onnxruntime as ort + + calls = [] + session = object() + monkeypatch.delenv('ORT_DISABLE_TELEMETRY', raising=False) + monkeypatch.setattr(ort, 'disable_telemetry_events', lambda: calls.append('telemetry')) + monkeypatch.setattr( + ort, + 'InferenceSession', + lambda path, providers: calls.append((path, providers)) or session, + ) + + output = inference._create_session('model.onnx', providers=['CPUExecutionProvider']) + + assert output is session + assert calls == ['telemetry', ('model.onnx', ['CPUExecutionProvider'])] + + +def test_one_shot_classify_delegates_to_classifier(monkeypatch): + captured = [] + + class FakeClassifier: + def __init__(self, **kwargs): + captured.append(kwargs) + + def classify(self, inputs): + captured.append(inputs) + return [] + + monkeypatch.setattr(inference, 'Classifier', FakeClassifier) + + assert ( + inference.classify( + 'spectrogram.png', + batch_size=4, + config='mobilenet', + providers=['CPUExecutionProvider'], + top_k=2, + sessions={'mobilenet': object()}, + num_workers=6, + ) + == [] + ) + assert captured[0]['batch_size'] == 4 + assert captured[0]['providers'] == ['CPUExecutionProvider'] + assert captured[0]['top_k'] == 2 + assert captured[0]['num_workers'] == 6 + assert captured[1] == 'spectrogram.png' + + +def test_post_rejects_wrong_score_count(): + predictions = np.zeros((1, len(classifier.CLASSES) - 1), dtype=np.float32) + + with pytest.raises(ValueError, match='scores'): + inference.post([(predictions, 'mobilenet')]) + + +def test_classifier_validates_options_and_accepts_an_existing_session(): + with pytest.raises(ValueError, match='batch_size'): + classifier.Classifier(batch_size=0) + with pytest.raises(ValueError, match='top_k'): + classifier.Classifier(top_k=0) + with pytest.raises(ValueError, match='num_workers'): + classifier.Classifier(num_workers=0) + + session = _Session() + runner = classifier.Classifier(sessions={'mobilenet': session}) + + assert runner.session is session + with pytest.raises(ValueError, match='top_k'): + runner.classify([], top_k=0) + + +def test_classifier_convenience_methods_reuse_runner(monkeypatch): + captured = [] + + def fake_wav(filepath, **kwargs): + captured.append(('wav', filepath, kwargs)) + return _result(filepath) + + def fake_bulk(inputs, **kwargs): + captured.append(('bulk', inputs, kwargs)) + return {'results': [], 'summary': bulk.summarize([])} + + monkeypatch.setattr(bulk, 'classify_wav', fake_wav) + monkeypatch.setattr(bulk, 'classify_bulk', fake_bulk) + runner = classifier.Classifier(sessions={'mobilenet': _Session()}, top_k=3) + + runner.classify_wav('recording.wav', keep_spectrograms=True) + runner.classify_bulk('recordings', input_type='wav', recursive=False) + + assert captured[0][2]['top_k'] == 3 + assert captured[0][2]['keep_spectrograms'] is True + assert captured[0][2]['_runner'] is runner + assert captured[1][2]['input_type'] == 'wav' + assert captured[1][2]['recursive'] is False + assert captured[1][2]['_runner'] is runner + + +def test_model_resource_materialization_is_verified_and_reused(monkeypatch, tmp_path): + payload = b'verified ONNX model' + source = tmp_path / 'source.onnx' + source.write_bytes(payload) + config = replace( + MOBILENET, + filename='test.onnx', + sha256=hashlib.sha256(payload).hexdigest(), + ) + cache = tmp_path / 'cache' + monkeypatch.setattr(model, '_cache_directory', lambda: cache) + + materialized = model._materialize_resource(source, config) + source.write_bytes(b'changed after the copy') + + assert materialized.read_bytes() == payload + assert model._materialize_resource(source, config) == materialized + + +def test_model_resource_materialization_rejects_bad_checksum(monkeypatch, tmp_path): + source = tmp_path / 'source.onnx' + source.write_bytes(b'corrupt') + config = replace(MOBILENET, filename='test.onnx', sha256='0' * 64) + cache = tmp_path / 'cache' + monkeypatch.setattr(model, '_cache_directory', lambda: cache) + + with pytest.raises(ValueError, match='checksum'): + model._materialize_resource(source, config) + + assert not (cache / config.filename).exists() + assert list(cache.iterdir()) == [] + + +def test_bundled_model_missing_or_corrupt_uses_mirror(monkeypatch): + missing = replace(MOBILENET, resource_parts=('missing',)) + + assert model._bundled_model(missing) is None + + monkeypatch.setattr(model, '_has_expected_hash', lambda path, config: False) + assert model._bundled_model(MOBILENET) is None + + +def test_bundled_model_materializes_nonfilesystem_resource(monkeypatch, tmp_path): + payload = b'embedded model' + source = tmp_path / 'source.onnx' + source.write_bytes(payload) + config = replace( + MOBILENET, + filename='embedded.onnx', + sha256=hashlib.sha256(payload).hexdigest(), + ) + + class Resource: + def joinpath(self, *parts): + return self + + def is_file(self): + return True + + @contextmanager + def extracted(resource): + yield source + + monkeypatch.setattr(model, 'files', lambda package: Resource()) + monkeypatch.setattr(model, 'as_file', extracted) + monkeypatch.setattr(model, '_cache_directory', lambda: tmp_path / 'cache') + + materialized = model._bundled_model(config) + + assert materialized == tmp_path / 'cache' / config.filename + assert materialized.read_bytes() == payload + + +def test_model_cache_directory_is_versioned(monkeypatch, tmp_path): + monkeypatch.setattr(model.pooch, 'os_cache', lambda name: str(tmp_path / name)) + + assert model._cache_directory() == tmp_path / 'batbot' / 'models' / model.__version__ + + +def test_aggregate_results_weights_windows_and_rejects_empty_inputs(): + first = _result('first.png', confidence=0.2, window_count=1) + second = _result('second.png', confidence=0.8, window_count=3) + + combined = bulk._aggregate_results( + 'recording.wav', + [first, second], + ['first.png', 'second.png'], + top_k=2, + ) + + assert combined['path'] == 'recording.wav' + assert combined['scores']['EPFU'] == pytest.approx(0.65) + assert combined['scores']['NOISE'] == pytest.approx(0.35) + assert combined['window_count'] == 4 + assert combined['spectrogram_paths'] == ['first.png', 'second.png'] + with pytest.raises(ValueError, match='No spectrograms'): + bulk._aggregate_results('empty.wav', [], []) + + +def test_classify_wav_cleans_temporary_spectrograms(monkeypatch): + from batbot import spectrogram + + output_directories = [] + + def fake_compute(filepath, **kwargs): + output_directory = Path(kwargs['output_folder']) + output_directories.append(output_directory) + return [str(output_directory / 'one.png')], [], None, None + + monkeypatch.setattr(spectrogram, 'compute', fake_compute) + result = bulk.classify_wav('recording.wav', _runner=_Runner()) + + assert result['spectrogram_paths'] == [] + assert not output_directories[0].exists() + + +def test_classify_wav_retains_requested_spectrograms(monkeypatch, tmp_path): + from batbot import spectrogram + + captured = {} + + def fake_compute(filepath, **kwargs): + captured.update(kwargs) + return [str(tmp_path / 'one.png')], [], None, None + + monkeypatch.setattr(spectrogram, 'compute', fake_compute) + result = bulk.classify_wav( + 'recording.wav', + output_folder=tmp_path, + out_file_stem=tmp_path / 'custom', + _runner=_Runner(), + ) + + assert result['spectrogram_paths'] == [str(tmp_path / 'one.png')] + assert captured['output_folder'] == str(tmp_path) + assert captured['out_file_stem'] == str(tmp_path / 'custom') + + +def test_discover_inputs_filters_type_recursion_and_invalid_paths(tmp_path): + image = tmp_path / 'top.png' + wav = tmp_path / 'top.WAV' + nested = tmp_path / 'nested' + nested.mkdir() + nested_image = nested / 'nested.jpg' + for path in [image, wav, nested_image, tmp_path / 'ignored.txt']: + path.touch() + + assert bulk.discover_inputs(tmp_path, input_type='spectrogram', recursive=False) == [str(image)] + assert bulk.discover_inputs(tmp_path, input_type='wav') == [str(wav)] + assert bulk.discover_inputs(image) == [str(image)] + with pytest.raises(ValueError, match='input_type'): + bulk.discover_inputs(tmp_path, input_type='video') + with pytest.raises(FileNotFoundError, match='does not exist'): + bulk.discover_inputs(tmp_path / 'missing') + assert bulk.discover_inputs(tmp_path / 'ignored.txt') == [] + + +def test_summarize_empty_results(): + assert bulk.summarize([]) == { + 'total': 0, + 'classified': 0, + 'failed': 0, + 'label_counts': {}, + 'species_counts': {}, + 'noise_count': 0, + 'mean_confidence': None, + } + + +def test_classify_bulk_recovers_individual_images_and_names_wav_outputs(monkeypatch, tmp_path): + good = tmp_path / 'good.jpg' + bad = tmp_path / 'bad.png' + wav = tmp_path / 'call.wav' + for path in [good, bad, wav]: + path.touch() + spectrogram_output = tmp_path / 'spectrograms' + captured_wav = {} + + class FaultTolerantRunner(_Runner): + def classify(self, paths, top_k=5): + paths = [str(path) for path in paths] + if len(paths) > 1: + raise OSError('batch decoder failed') + if Path(paths[0]).name == 'bad.png': + raise ValueError('invalid image') + return [_result(paths[0])] + + def fake_classify_wav(path, **kwargs): + captured_wav.update(kwargs) + return _result(path) + + monkeypatch.setattr(bulk, 'classify_wav', fake_classify_wav) + output = bulk.classify_bulk( + tmp_path, + spectrogram_output=spectrogram_output, + _runner=FaultTolerantRunner(), + ) + + assert output['summary']['total'] == 3 + assert output['summary']['classified'] == 2 + assert output['summary']['failed'] == 1 + assert next(item for item in output['results'] if item['path'] == str(bad))['error'] == ( + 'invalid image' + ) + assert captured_wav['output_folder'] == spectrogram_output + assert captured_wav['keep_spectrograms'] is True + assert Path(captured_wav['out_file_stem']).parent == spectrogram_output + assert Path(captured_wav['out_file_stem']).name.startswith('call.') + + +def test_classify_bulk_batches_images_and_uses_temporary_wav_outputs(monkeypatch, tmp_path): + image = tmp_path / 'call.png' + wav = tmp_path / 'call.wav' + image.touch() + wav.touch() + runner = _Runner() + captured = [] + + def fake_classify_wav(path, **kwargs): + captured.append((path, kwargs)) + return _result(path) + + monkeypatch.setattr(bulk, 'classify_wav', fake_classify_wav) + + output = bulk.classify_bulk(tmp_path, _runner=runner) + + assert output['summary']['classified'] == 2 + assert runner.calls == [([str(image)], 5)] + assert captured[0][0] == str(wav) + assert captured[0][1]['output_folder'] is None + assert captured[0][1]['out_file_stem'] is None + assert captured[0][1]['keep_spectrograms'] is False + + +def test_bulk_constructs_classifier_with_requested_worker_count(monkeypatch): + captured = [] + + class Runner: + def __init__(self, **kwargs): + captured.append(kwargs) + + monkeypatch.setattr(bulk, 'Classifier', Runner) + + output = bulk.classify_bulk([], num_workers=4) + + assert output['results'] == [] + assert captured[0]['num_workers'] == 4 diff --git a/tests/test_cli_additional.py b/tests/test_cli_additional.py new file mode 100644 index 0000000..da978f3 --- /dev/null +++ b/tests/test_cli_additional.py @@ -0,0 +1,312 @@ +import json +import warnings +from pathlib import Path + +from click.testing import CliRunner + +import batbot +from batbot.batbot_cli import batch, cli, example, fetch, pipeline, preprocess + + +def _result(path): + return { + 'path': str(path), + 'label': 'EPFU', + 'confidence': 0.75, + 'window_count': 1, + 'top': [{'label': 'EPFU', 'confidence': 0.75}], + 'scores': {'EPFU': 0.75}, + } + + +def _wav_files(tmp_path): + input_directory = tmp_path / 'input' + input_directory.mkdir() + paths = [input_directory / 'one.wav', input_directory / 'two.wav'] + for path in paths: + path.touch() + return input_directory, paths + + +def test_fetch_cli_forwards_configuration_and_pull(monkeypatch): + calls = [] + monkeypatch.setattr( + batbot, + 'fetch', + lambda **kwargs: calls.append(kwargs) or '/models/batbot.onnx', + ) + + invocation = CliRunner().invoke(fetch, ['--config', 'mobilenet', '--pull']) + + assert invocation.exit_code == 0, invocation.output + assert invocation.output.strip() == '/models/batbot.onnx' + assert calls == [{'config': 'mobilenet', 'pull': True}] + + +def test_pipeline_cli_validates_input_and_forwards_output(monkeypatch, tmp_path): + wav = tmp_path / 'recording.wav' + wav.touch() + output = tmp_path / 'output' + calls = [] + monkeypatch.setattr(batbot, 'pipeline', lambda *args, **kwargs: calls.append((args, kwargs))) + + invocation = CliRunner().invoke(pipeline, [str(wav), '--output', str(output)]) + missing = CliRunner().invoke(pipeline, [str(tmp_path / 'missing.wav')]) + + assert invocation.exit_code == 0, invocation.output + assert calls == [((str(wav),), {'output_folder': str(output)})] + assert missing.exit_code == 2 + assert 'Input filepath does not exist' in missing.output + + +def test_batch_cli_writes_stdout_and_json(monkeypatch, tmp_path): + wav = tmp_path / 'recording.wav' + wav.touch() + output = tmp_path / 'batch.json' + calls = [] + + def fake_batch(paths, **kwargs): + calls.append((paths, kwargs)) + return [_result(paths[0])] + + monkeypatch.setattr(batbot, 'batch', fake_batch) + + printed = CliRunner().invoke( + batch, + [str(wav), '--config', 'mobilenet', '--num-workers', '3'], + ) + written = CliRunner().invoke(batch, [str(wav), '--output', str(output)]) + + assert printed.exit_code == 0, printed.output + assert json.loads(printed.output)['summary']['species_counts'] == {'EPFU': 1} + assert written.exit_code == 0, written.output + assert json.loads(output.read_text())['results'][0]['path'] == str(wav) + assert calls[0][1] == {'config': 'mobilenet', 'num_workers': 3} + assert calls[1][1] == {'config': None, 'num_workers': 1} + + +def test_example_and_root_cli(monkeypatch): + calls = [] + monkeypatch.setattr(batbot, 'example', lambda: calls.append(True)) + + example_result = CliRunner().invoke(example) + help_result = CliRunner().invoke(cli, ['--help']) + + assert example_result.exit_code == 0, example_result.output + assert calls == [True] + assert help_result.exit_code == 0 + assert 'classify-bulk' in help_result.output + + +def test_preprocess_reports_when_no_inputs_exist(tmp_path): + invocation = CliRunner().invoke(preprocess, [str(tmp_path / '*.wav')]) + + assert invocation.exit_code == 0, invocation.output + assert 'Found no files' in invocation.output + + +def test_preprocess_dry_run_writes_plan_without_processing(monkeypatch, tmp_path): + input_directory, paths = _wav_files(tmp_path) + output_directory = tmp_path / 'output' + output_directory.mkdir() + stale = output_directory / 'stale.txt' + stale.touch() + report = tmp_path / 'dry-run.json' + monkeypatch.setattr( + batbot, + 'pipeline', + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError('pipeline was called')), + ) + + invocation = CliRunner().invoke( + preprocess, + [ + str(input_directory), + '--output-dir', + str(output_directory), + '--force-overwrite', + '--dry-run', + '--output-json', + str(report), + ], + ) + + assert invocation.exit_code == 0, invocation.output + data = json.loads(report.read_text()) + assert [pair[0] for pair in data['input file, output file stem']] == [ + str(path) for path in paths + ] + assert data['files to be deleted in cleanup'] == [str(stale)] + assert 'Dry run mode active' in invocation.output + + +def test_preprocess_dry_run_prints_plan_and_flattens_structure(tmp_path): + _, paths = _wav_files(tmp_path) + output_directory = tmp_path / 'output' + + invocation = CliRunner().invoke( + preprocess, + [ + *(str(path) for path in paths), + '--output-dir', + str(output_directory), + '--force-overwrite', + '--dry-run', + '--no-file-structure', + ], + ) + + assert invocation.exit_code == 0, invocation.output + assert 'Flattening output file structure' in invocation.output + assert 'files to be deleted in cleanup' in invocation.output + + +def test_preprocess_skips_files_with_existing_outputs(tmp_path): + input_directory, _ = _wav_files(tmp_path) + output_directory = tmp_path / 'output' + output_directory.mkdir() + (output_directory / 'one.jpg').touch() + (output_directory / 'two.jpg').touch() + + invocation = CliRunner().invoke( + preprocess, + [str(input_directory), '--output-dir', str(output_directory)], + ) + + assert invocation.exit_code == 0, invocation.output + assert 'Found no unprocessed files' in invocation.output + assert 'use --force-overwrite' in invocation.output + + +def test_preprocess_serial_writes_pipeline_results(monkeypatch, tmp_path): + input_directory, paths = _wav_files(tmp_path) + output_directory = tmp_path / 'output' + report = tmp_path / 'results.json' + calls = [] + + def fake_pipeline(filepath, **kwargs): + calls.append((filepath, kwargs)) + stem = str(kwargs['out_file_stem']) + return [f'{stem}.png'], [f'{stem}.compressed.jpg'], f'{stem}.json' + + monkeypatch.setattr(batbot, 'pipeline', fake_pipeline) + + invocation = CliRunner().invoke( + preprocess, + [ + str(input_directory), + '--output-dir', + str(output_directory), + '--force-overwrite', + '--process-metadata', + '--output-json', + str(report), + ], + ) + + assert invocation.exit_code == 0, invocation.output + data = json.loads(report.read_text()) + assert len(calls) == len(paths) + assert all(call[1]['fast_mode'] is False for call in calls) + assert len(data['output_path']) == len(paths) + assert data['failed_files'] == [] + + +def test_preprocess_serial_retains_individual_failures(monkeypatch, tmp_path): + input_directory, _ = _wav_files(tmp_path) + + def fake_pipeline(filepath, **kwargs): + if Path(filepath).name == 'two.wav': + raise ValueError('corrupt WAV') + return ['one.png'], ['one.compressed.jpg'], 'one.json' + + monkeypatch.setattr(batbot, 'pipeline', fake_pipeline) + + with warnings.catch_warnings(record=True) as caught: + invocation = CliRunner().invoke( + preprocess, + [str(input_directory), '--force-overwrite'], + ) + + assert invocation.exit_code == 0, invocation.output + assert 'corrupt WAV' in invocation.output + assert any('Pipeline failed' in str(warning.message) for warning in caught) + + +def test_preprocess_parallel_delegates_chunked_work(monkeypatch, tmp_path): + input_directory, _ = _wav_files(tmp_path) + captured = [] + + def fake_parallel(**kwargs): + captured.append(kwargs) + return ['one.png'], ['one.compressed.jpg'], ['one.json'], [] + + monkeypatch.setattr(batbot, 'parallel_pipeline', fake_parallel) + + invocation = CliRunner().invoke( + preprocess, + [str(input_directory), '--force-overwrite', '--num-workers', '2'], + ) + + assert invocation.exit_code == 0, invocation.output + assert len(captured) == 1 + assert captured[0]['num_workers'] == 2 + assert len(captured[0]['in_file_chunks']) == 2 + assert captured[0]['fast_mode'] is True + assert 'one.compressed.jpg' in invocation.output + + +def test_preprocess_cleanup_can_abort_or_delete(monkeypatch, tmp_path): + input_directory, _ = _wav_files(tmp_path) + output_directory = tmp_path / 'output' + output_directory.mkdir() + stale = output_directory / 'stale.txt' + stale.touch() + monkeypatch.setattr( + batbot, + 'pipeline', + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError('pipeline was called')), + ) + arguments = [ + str(input_directory), + '--output-dir', + str(output_directory), + '--force-overwrite', + '--cleanup', + ] + + aborted = CliRunner().invoke(preprocess, arguments, input='n\n') + + assert aborted.exit_code == 0, aborted.output + assert 'Aborting cleanup mode' in aborted.output + assert stale.exists() + + deleted = CliRunner().invoke(preprocess, arguments, input='yes\n') + + assert deleted.exit_code == 0, deleted.output + assert f'Deleting file: {stale}' in deleted.output + assert not stale.exists() + + +def test_preprocess_cleanup_handles_no_extra_files(monkeypatch, tmp_path): + input_directory, _ = _wav_files(tmp_path) + output_directory = tmp_path / 'output' + monkeypatch.setattr( + batbot, + 'pipeline', + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError('pipeline was called')), + ) + + invocation = CliRunner().invoke( + preprocess, + [ + str(input_directory), + '--output-dir', + str(output_directory), + '--force-overwrite', + '--cleanup', + ], + ) + + assert invocation.exit_code == 0, invocation.output + assert 'No files to delete' in invocation.output diff --git a/tests/test_preprocess.py b/tests/test_preprocess.py index de5b4ca..710964c 100644 --- a/tests/test_preprocess.py +++ b/tests/test_preprocess.py @@ -21,19 +21,19 @@ def test_preprocess(): num_examples = 2 output_str = str(data.output).split('\n') for ii in range(num_examples): - expected_file = './output/example{}.01of01.compressed.jpg'.format(ii + 1) + expected_file = f'./output/example{ii + 1}.01of01.compressed.jpg' assert any( [expected_file in x for x in output_str] - ), 'Did not find file listed among outputs: {}'.format(expected_file) + ), f'Did not find file listed among outputs: {expected_file}' assert os.path.exists(expected_file), 'Did not find file on filesystem: {}'.format( expected_file ) num_min_call_segments = [65, 18, 149, 47] for ii in range(num_examples): - expected_file = './output/example{}.metadata.json'.format(ii + 1) + expected_file = f'./output/example{ii + 1}.metadata.json' assert any( [expected_file in x for x in output_str] - ), 'Did not find file listed among outputs: {}'.format(expected_file) + ), f'Did not find file listed among outputs: {expected_file}' assert os.path.exists(expected_file), 'Did not find file on filesystem: {}'.format( expected_file ) diff --git a/tests/test_preprocess_parallel.py b/tests/test_preprocess_parallel.py index 8e5fe83..0ee3995 100644 --- a/tests/test_preprocess_parallel.py +++ b/tests/test_preprocess_parallel.py @@ -16,18 +16,18 @@ def test_preprocess_parallel(): num_examples = 2 output_str = str(data.output).split('\n') for ii in range(num_examples): - expected_file = './output/example{}.01of01.compressed.jpg'.format(ii + 1) + expected_file = f'./output/example{ii + 1}.01of01.compressed.jpg' assert any( [expected_file in x for x in output_str] - ), 'Did not find file listed among outputs: {}'.format(expected_file) + ), f'Did not find file listed among outputs: {expected_file}' assert os.path.exists(expected_file), 'Did not find file in filesystem: {}'.format( expected_file ) for ii in range(num_examples): - expected_file = './output/example{}.metadata.json'.format(ii + 1) + expected_file = f'./output/example{ii + 1}.metadata.json' assert any( [expected_file in x for x in output_str] - ), 'Did not find file listed among outputs: {}'.format(expected_file) + ), f'Did not find file listed among outputs: {expected_file}' assert os.path.exists(expected_file), 'Did not find file in filesystem: {}'.format( expected_file ) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..369c096 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,30 @@ +import logging +from logging.handlers import TimedRotatingFileHandler + +import rich + +from batbot import utils + + +def test_init_logging_configures_file_and_rich_handlers(monkeypatch, tmp_path): + configured = {} + root = logging.getLogger() + original_level = root.level + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(rich, 'reconfigure', lambda **kwargs: configured.update(theme=kwargs)) + monkeypatch.setattr(logging, 'basicConfig', lambda **kwargs: configured.update(logging=kwargs)) + + try: + logger = utils.init_logging() + finally: + root.setLevel(original_level) + + handlers = configured['logging']['handlers'] + try: + assert logger.name == 'batbot' + assert configured['logging']['level'] == utils.DEFAULT_LOG_LEVEL + assert any(isinstance(handler, TimedRotatingFileHandler) for handler in handlers) + assert configured['theme']['theme'].styles['logging.level.error'].bold is True + finally: + for handler in handlers: + handler.close()