diff --git a/.agents/skills/mblt-vision-readme/SKILL.md b/.agents/skills/mblt-vision-readme/SKILL.md new file mode 100644 index 0000000..2970e51 --- /dev/null +++ b/.agents/skills/mblt-vision-readme/SKILL.md @@ -0,0 +1,41 @@ +--- +name: mblt-vision-readme +description: >- + Write and maintain mblt-vision-python README documentation, API examples, model references, + and migration notes. +--- + +# Mobilint Vision README Writing + +## Documentation Ownership + +- Keep the root README concise: package purpose, installation, a minimal example, and a link to + mblt_vision/README.md. +- Keep the detailed Vision API reference in mblt_vision/README.md. It owns Python construction, + framework selection, model discovery, model-family tables, output taxonomy, and migration notes. +- Keep development-tool instructions in `benchmark/README.md` and `compile/README.md`; their + command scripts live directly in those directories. +- Document Model Zoo compatibility as migration context only. Do not present Model Zoo CLI, + validation, dataset organization, or compilation commands as features of this package. + +## Accuracy Rules + +- Use the public mblt_vision namespace in every executable example. +- Use model_path for new local-artifact examples. Mention mxq_path and onnx_path only as + compatibility aliases. +- State that .mxq and .onnx paths select their framework automatically when framework is omitted. + Document the explicit-framework conflict error. +- Describe file_cfg.filename as the MXQ source artifact and same-stem ONNX derivation. Mention + onnx_filename only for a genuinely different published artifact. +- Keep post_cfg.dataset terminology precise: it identifies output taxonomy, not just a task. +- Use obb as the only oriented-bounding-box name in standalone documentation. + +## Style and Validation + +- Use ATX headings, one blank line between blocks, hyphen lists, concise paragraphs, and + language-tagged code fences. +- Prefer generated discovery examples such as list_tasks() and list_models() over manually + maintained exhaustive name lists. +- When changing models, package metadata, dependencies, public APIs, or runtime behavior, update + the relevant README, `AGENTS.md`, and the canonical `mblt-vision` skill if the workflow changes. +- For documentation-only updates, run git diff --check and verify relative links and headings. diff --git a/.agents/skills/mblt-vision/SKILL.md b/.agents/skills/mblt-vision/SKILL.md new file mode 100644 index 0000000..e183c8b --- /dev/null +++ b/.agents/skills/mblt-vision/SKILL.md @@ -0,0 +1,80 @@ +--- +name: mblt-vision +description: >- + Work on the standalone Mobilint Vision Python API, model registry, preprocessing, + postprocessing, results, runtime integration, and package compatibility contracts. +--- + +# Mobilint Vision Python + +## Start Here + +1. Read AGENTS.md. +2. Run git status --short before changing files. +3. Read pyproject.toml, the affected package exports, matching model YAML, and relevant tests. +4. For a compatibility migration, compare against + ../mblt-model-zoo/mblt_model_zoo/vision deliberately; do not make it a runtime dependency. + +## Public API and Model Registry + +- Use mblt_vision.MBLT_Engine and task subpackages as the public surface. +- Keep mblt_vision as the sole intended import namespace. Use obb as the sole + oriented-bounding-box task name. +- Update a task package, top-level lazy exports, and list_models() discovery together. +- Preserve constructor arguments including model_path, mxq_path, onnx_path, + model_type, and core-selection options unless intentionally changing the API. +- Keep .mxq/.onnx suffix routing and explicit-framework conflict errors intact. +- Every model YAML must define stable file_cfg, pre_cfg, and post_cfg mappings. + Use file_cfg.filename for MXQ and derive the same-stem ONNX artifact unless + onnx_filename is required. +- Every post_cfg declares dataset; resolve output taxonomy from the dataset/task pair. + +## Processing and Results + +- Reuse the shared letterbox geometry for both preprocessing and inverse coordinate restoration. +- Detection requires pre_cfg.LetterBox. Keep semantic metadata (img0_shape and + ratio_pad) through postprocessing so logits restore to the original geometry before + argmax. +- Preserve decoded-output layout provenance through NMS. For ambiguous tensors without + provenance, prioritize channels-first raw-output normalization. +- Normalize dense depth and semantic outputs before inverse letterboxing. Validate baked semantic + maps are finite, integral, and in-range before converting them to integer class IDs. +- Keep result shapes, ordering, coordinates, dtype, and empty-result behavior compatible with + the Model Zoo reference. + +## Runtime and Packaging + +- Route NPU runtime access through mblt-npu-python; do not copy backend classes into Vision. +- Use the shared `ONNXBackend` for ONNX inference. Keep ONNX Runtime optional and lazy-imported; + raise a specific installation error when it is requested but unavailable. +- Normalize legacy `aries` and `regulus` target values through mblt-npu-python. MXQ artifacts and + compilation metadata must resolve only from the selected board folder, never a core-mode path or + a fallback board folder. +- Include model and dataset YAML files as package data. Build a wheel and inspect it after + changing metadata or assets. +- Do not require native bindings, GStreamer, hardware, downloaded models, or caches for normal + imports and unit tests. + +## Tooling Layout and Documentation + +- Keep all executable benchmark scripts directly in `benchmark/`; reusable reporting helpers belong + in `mblt_vision.benchmark`. +- Keep all executable compile scripts and the compile guide directly in `compile/`. +- Use `~/.mblt_model_zoo` as the shared artifact and dataset cache root. Keep organizer defaults, + dataset registry YAMLs, compilation defaults, and documented commands aligned to it. +- Keep imports free of cache-directory creation, write probes, downloads, and temporary-directory + allocation; resolve a writable cache only when an artifact or compilation output needs it. +- Make fallback caches stable, private, and user-owned. Never use a new temporary directory per + process or trust a shared fallback cache without validating it. +- For every significant package change (public API, CLI, runtime/dependency, artifact layout, or + tooling structure), update `AGENTS.md`, this canonical skill, the Claude skill entry point when + its workflow changes, and the relevant README in the same change. + +## Validate Proportionately + +- Begin with the smallest relevant test file or -k selection. +- Add deterministic differential tests for Model Zoo compatibility, including invalid inputs, + empty detections, threshold boundaries, task discovery, and image geometry. +- Run pre-commit run --files when available. For docs, run + git diff --check. +- Report unavailable hardware, downloads, or optional dependencies rather than weakening tests. diff --git a/.claude/skills/mblt-vision-readme/SKILL.md b/.claude/skills/mblt-vision-readme/SKILL.md new file mode 100644 index 0000000..ea6a148 --- /dev/null +++ b/.claude/skills/mblt-vision-readme/SKILL.md @@ -0,0 +1,9 @@ +--- +name: mblt-vision-readme +description: Write and maintain documentation for the standalone Mobilint Vision Python package. +--- + +# Mobilint Vision README Writing + +Read and follow the canonical skill at +../../../.agents/skills/mblt-vision-readme/SKILL.md. diff --git a/.claude/skills/mblt-vision/SKILL.md b/.claude/skills/mblt-vision/SKILL.md new file mode 100644 index 0000000..6120c85 --- /dev/null +++ b/.claude/skills/mblt-vision/SKILL.md @@ -0,0 +1,9 @@ +--- +name: mblt-vision +description: Work effectively on the standalone Mobilint Vision Python API and model registry. +--- + +# Mobilint Vision Python + +Read and follow the canonical skill at +../../../.agents/skills/mblt-vision/SKILL.md. diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..45f0af2 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,80 @@ +name: Publish to PyPI + +on: + release: + types: [published] + +permissions: + contents: read + +jobs: + build: + name: Build distribution 🛠️ + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.12' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + python -m pip install build twine + + - name: Build package + run: python -m build + + - name: Validate distributions + run: python -m twine check dist/* + + - name: Verify wheel package contents + run: | + python -c "from pathlib import Path; from zipfile import ZipFile; wheel, = Path('dist').glob('*.whl'); archive = ZipFile(wheel); assert 'mblt_vision/py.typed' in archive.namelist(), 'Wheel is missing mblt_vision/py.typed'" + + - name: Upload built distributions + uses: actions/upload-artifact@v7 + with: + name: python-package-distributions + path: dist/ + + publish-to-testpypi: + name: Publish to TestPyPI 🧪 + needs: [build] + runs-on: ubuntu-latest + environment: testpypi + permissions: + id-token: write + steps: + - name: Download distribution artifacts + uses: actions/download-artifact@v7 + with: + name: python-package-distributions + path: dist/ + + - name: Publish to TestPyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + repository-url: https://test.pypi.org/legacy/ + + publish-to-pypi: + name: Publish to PyPI 🚀 + needs: [build, publish-to-testpypi] + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write + steps: + - name: Download distribution artifacts + uses: actions/download-artifact@v7 + with: + name: python-package-distributions + path: dist/ + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index e69de29..2c98cff 100644 --- a/.gitignore +++ b/.gitignore @@ -0,0 +1,47 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +.Python +.venv/ +venv/ +env/ +.env +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.coverage +coverage.xml +htmlcov/ +build/ +dist/ +*.egg-info/ + +# C/C++ build artifacts +build/ +cmake-build-*/ +CMakeFiles/ +CMakeCache.txt +cmake_install.cmake +compile_commands.json +*.o +*.obj +*.a +*.lib +*.so +*.dylib +*.dll +*.exe +*.out + +# IDE and OS files +.idea/ +.vscode/ +*.swp +*.swo +.DS_Store +Thumbs.db + +*.zip +*.tar.gz +*.egg \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..b7350df --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,8 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.6.9 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + diff --git a/AGENTS.md b/AGENTS.md index c9b161e..b90b62b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,17 +9,18 @@ paths: ## Mission `mblt-vision-python` is the Python distribution and public compatibility layer for Mobilint -Vision. It binds to the supported native API in `mblt-vision`; it is not a second implementation -of the vision runtime. Its end-state is a drop-in replacement for the Vision API currently shipped -by `mblt-model-zoo`. +Vision. The immediate plan is a pure-Python implementation that can replace the Vision API +currently shipped by `mblt-model-zoo`. C++ bindings are deferred until `mblt-vision` has a stable, +supported native API. -The ownership boundary is deliberate: +The current ownership boundary is deliberate: -- `mblt-vision` owns inference, model loading, preprocessing, postprocessing, and native resource - management. -- This package owns Python ergonomics, Python object conversion, package metadata, wheels, API - documentation, and compatibility shims. -- Do not put GStreamer integration, C++ runtime logic, or duplicated numerical kernels in Python. +- This package owns the public Python API, model loading, preprocessing, postprocessing, runtime + integration, dataset organization, benchmark evaluation, model compilation, package metadata, + wheels, documentation, and compatibility shims. +- Keep implementation code in Python; do not block Python API progress on C++ or GStreamer work. +- Design internal seams so a future optional `mblt-vision` backend can replace implementation + details without changing the documented Python API or result contracts. ## Before Editing @@ -29,8 +30,9 @@ The ownership boundary is deliberate: - For a Model Zoo replacement item, inspect the matching behavior in `../mblt-model-zoo/mblt_model_zoo/vision`, including its tests and model YAML configuration. Treat it as the compatibility reference until the new package explicitly supersedes it. -- Coordinate native API changes with `../mblt-vision`. Do not bind private, undocumented, or - build-tree-only native symbols. +- Do not make `../mblt-vision`, a compiled extension, or GStreamer a dependency of normal package + development, installation, import, or unit tests. Revisit integration only after its native API + is documented and versioned. ## Python API Contract @@ -46,32 +48,78 @@ The ownership boundary is deliberate: buffer for as long as native code can access it. - Never expose raw native pointers or require callers to manage native lifetime. Wrap resources in deterministic `close()`/context-manager behavior and safe finalization as appropriate. -- Maintain task vocabulary compatibility: canonical `obb`, with - `oriented_bounding_boxes` accepted only as an input compatibility alias. - -## Binding and Native Dependency Rules - -- The binding must use only the documented, versioned `mblt-vision` interface. Add a native - capability/version check when a Python feature needs a newer library. -- Keep ownership and the GIL explicit. Release the GIL only around blocking native work that does - not access Python objects; reacquire it before callbacks, exceptions, or Python buffer access. -- Do not catch broad native errors and return empty or plausible-looking results. Fail loudly with - an exception that identifies the invalid input, unsupported feature, or native-library problem. -- Keep native library discovery relocatable and diagnosable. Avoid hard-coded developer paths, - `LD_LIBRARY_PATH` requirements for normal wheels, or importing an optional native backend at - module import time if that prevents useful error reporting. +- Use `obb` as the sole oriented-bounding-box task name. + +## Vision Models and Processing Contracts + +- Use mblt_vision.MBLT_Engine for loading. Prefer model_path; retain mxq_path and + onnx_path as compatibility aliases. +- Update task-package exports and lazy top-level exports together. Confirm that + list_models() discovers every public model class. +- Keep each model YAML's file_cfg, pre_cfg, and post_cfg shape stable. + file_cfg.filename is the canonical MXQ Hub artifact; derive the same-stem ONNX filename + unless the Hub artifact requires an explicit onnx_filename. +- Require post_cfg.dataset in every model YAML and resolve output class counts using the + dataset/task pair. Do not assume one output taxonomy for every model in a task. +- Preserve automatic .mxq/.onnx framework detection and the fail-fast error when a local + suffix conflicts with an explicitly selected framework. +- Preserve anchorless decoded-output layout provenance through NMS. When a tensor is ambiguous + and provenance is unavailable, normalize it as raw channels-first before candidates-first. +- Use the shared letterbox helpers for forward geometry and inverse output restoration. Detection + postprocessors require pre_cfg.LetterBox; metadata-aware semantic preprocessing returns the + original image shape and ratio_pad so logits can be restored before argmax. +- Normalize dense outputs before inverse letterboxing: upsample quarter-resolution depth maps by + four, preserve baked-resize depth maps, convert Cityscapes NHWC logits to NCHW, and reject + non-finite, fractional, or out-of-range baked semantic IDs before casting. +- Keep hardware-specific runtime access behind mblt-npu-python. Optional ONNX Runtime imports + must remain lazy and report the appropriate package extra when unavailable. + +## Python-First Architecture + +- Keep the public layer independent from a particular backend. Define small internal interfaces + for model execution and artifact resolution, but do not add speculative abstractions before a + second backend exists. +- Put preprocessing, postprocessing, model configuration, and compatibility behavior in tested + Python modules. Reuse Model Zoo semantics deliberately; do not copy code wholesale without + understanding its public contract and license context. +- Use established Python runtime dependencies only when they materially support the package goals. + Keep optional frameworks lazy-imported and raise a specific installation error when a requested + backend is unavailable. +- Do not catch broad runtime errors and return empty or plausible-looking results. Fail loudly with + an exception that identifies the invalid input, unsupported feature, or unavailable dependency. +- If/when a native backend is introduced, it must be optional, use a documented and versioned + `mblt-vision` interface, and preserve the Python public API, exceptions, result values, layouts, + and lifecycle behavior. Add native capability/version checks at that time. + +## Benchmark and Compilation Tooling + +- Keep executable benchmark organizers, the unified benchmark runner, and result comparison scripts + directly under `benchmark/`. Put reusable benchmark reporting helpers in `mblt_vision.benchmark`. +- Keep executable compilation helpers and their guide directly under `compile/`. Do not recreate a + Vision-only subdirectory under either tooling root. +- Use `~/.mblt_model_zoo` as the shared artifact and dataset cache root. Organizer defaults, + dataset registry YAMLs, compilation defaults, and documented commands must agree on that root. +- Keep package imports free of cache-directory creation, write probes, downloads, and temporary + directory allocation. Resolve a writable cache lazily only when an artifact or compilation output needs it. +- If the preferred cache is unavailable, use a stable, private, user-owned fallback cache. Do not + create a new temporary cache per process or reuse an unsafe shared directory. +- Benchmark and compilation commands are development tools; do not package them as public CLI + entry points without an explicit product decision. The supported end-user command is + `mblt-vision`. +- Compile and artifact resolution must use normalized board identifiers (`aries-rb`, `regulus-ra`, + or `regulus-rb`) and must not fall back to a different board folder. ## PyPI and Wheel Packaging - `pyproject.toml` is the source of truth for Python metadata, supported Python versions, dependencies, and build backend. Keep package versioning synchronized with the exposed API and native compatibility requirements. -- Distribute wheels that include or correctly depend on the matching native library according to - the chosen packaging strategy. Do not publish an sdist/wheel whose import or basic diagnostics - are broken without a locally installed development tree. +- Publish pure-Python wheels and sdists that install and import without a local C++ build, a native + library, or GStreamer. Do not publish artifacts whose import or basic diagnostics require a + developer environment. - Build and test each intended platform/architecture wheel in a clean environment. Verify wheel - contents, package metadata, install-from-wheel, import, and a minimal inference-free native - smoke test. Do not upload from a developer environment as the only validation. + contents, package metadata, install-from-wheel, import, and a minimal API smoke test. Do not + upload from a developer environment as the only validation. - Keep optional dependencies genuinely optional and avoid importing them from package top level. Do not add model weights, caches, test assets, or compiled build artifacts to source control. @@ -83,10 +131,14 @@ The ownership boundary is deliberate: - Use deterministic differential tests against Model Zoo for shared behavior. Cover edge cases, not only successful end-to-end examples: invalid layouts/dtypes, empty detections, threshold boundaries, image geometry, model aliases, task aliases, and resource cleanup. -- Numerical semantics belong to the native layer, but binding tests must verify that Python - conversion does not change values, layout, coordinates, ordering, dtype, or ownership. -- Avoid making hardware, downloaded models, or GStreamer a requirement for ordinary unit tests. +- Verify that Python preprocessing and postprocessing preserve expected values, layouts, + coordinates, ordering, dtype, and ownership. If a future backend is used, require the same + parity from its conversion boundary. +- Avoid making hardware, downloaded models, compiled extensions, or GStreamer a requirement for + ordinary unit tests. Mark and document integration prerequisites; run the narrowest relevant suite first. +- Use a deterministic default seed of 0 for any public API that samples or otherwise uses + randomness. ## Code Quality and Documentation @@ -95,6 +147,15 @@ The ownership boundary is deliberate: - Keep imports ordered as standard library, third-party, then local. Catch specific exceptions. - Update the README and API examples whenever installation, native-library discovery, supported platforms, imports, or migration compatibility changes. +- Keep the root README focused on installation and navigation. Maintain the complete Vision API, + model-family, runtime, and taxonomy reference in mblt_vision/README.md. +- Write documentation with ATX headings, one blank line between blocks, hyphen lists, + language-tagged code fences, and concise paragraphs. Keep examples executable against the + public mblt_vision namespace and do not document Model Zoo CLI commands as standalone features. +- When a durable public fact changes, update this guide, the matching agent skill, CLAUDE.md, and + the relevant README in the same change. Treat a significant package change—public API, + dependency/runtime, artifact layout, CLI, or tooling structure—as a required guide-and-skill + synchronization point. - For documentation-only changes, run `git diff --check` and verify headings and links. Report skipped platform, hardware, or native-runtime checks clearly. diff --git a/CLAUDE.md b/CLAUDE.md index 9d3f9ec..ddefdfd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,5 +2,8 @@ @AGENTS.md -`AGENTS.md` is the canonical guide for this repository. Follow it for the thin native-binding, -PyPI-wheel, lifecycle, and mblt-model-zoo Vision compatibility requirements. +`AGENTS.md` is the canonical guide for this repository. Follow it for the Python-first +implementation, PyPI-wheel, lifecycle, and mblt-model-zoo Vision compatibility requirements. + +For focused model, preprocessing, postprocessing, and model-registry work, also read +.claude/skills/mblt-vision/SKILL.md. diff --git a/README.md b/README.md index e69de29..b21c357 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,108 @@ +# Mobilint Vision Python + +Run pre-trained Mobilint Vision models from Python. `mblt-vision-python` provides +model configuration, artifact loading, preprocessing, inference integration, and +typed postprocessing results for image classification, depth estimation, face and +object detection, OBB, instance and semantic segmentation, +and pose estimation. + +Version `0.0.0` is the initial standalone release. + +## Installation + +```bash +pip install mblt-vision-python +``` + +MXQ inference requires a supported Mobilint NPU environment. Model artifacts are +downloaded from the Mobilint Hugging Face organization when no local `model_path` +is supplied. For ONNX execution, install one of the optional extras: + +```bash +pip install "mblt-vision-python[onnxruntime]" +# Or, on supported systems: +pip install "mblt-vision-python[onnxruntime-gpu]" +``` + +## Quick start + +Each model includes its matching preprocess and postprocess behavior: + +```python +from mblt_vision import ResNet50 + +model = ResNet50() +x = model.preprocess("image.jpg") +result = model.postprocess(model(x)) +``` + +For configurable model selection and local MXQ or ONNX artifacts, use +`MBLT_Engine`: + +```python +from mblt_vision import MBLT_Engine + +model = MBLT_Engine(model_cls="resnet50", model_type="DEFAULT") +try: + result = model.postprocess(model(model.preprocess("image.jpg"))) +finally: + model.dispose() +``` + +Discover supported tasks and models with `list_tasks()` and `list_models()`. New +code should use the task subpackages (for example, +`mblt_vision.object_detection`) or `MBLT_Engine`. Top-level model imports such as +`from mblt_vision import ResNet50` remain supported for convenience. + +`obb` is the canonical oriented-bounding-box task name. + +## Model Zoo migration + +Vision is now maintained in this package. `mblt-model-zoo` retains +`mblt_model_zoo.vision` as a compatibility facade for existing applications; new +projects should import from `mblt_vision` directly. Its `mblt-model-zoo predict`, +`val`, and `compile` commands also delegate to this package. + +## Command line + +The standalone package provides the `mblt-vision` command with `predict`, `val`, +and `compile` subcommands: + +```bash +mblt-vision predict --source image.jpg --model resnet50 +``` + +`predict` is the single inference command for classification, depth estimation, +object and face detection, instance and semantic segmentation, OBB, and pose +estimation. The selected model determines its task and processing pipeline. +By default it downloads the model artifact and saves a plotted result under +`runs/vision/predict/`. Use `--output` to choose the result-image path, +`--topk` for classification labels, and `--conf-thres`/`--iou-thres` for +detection-style tasks. `--framework onnx` selects ONNX Runtime inference; +`--target-device` and `--core-mode` select the MXQ board/runtime mode. + +```bash +mblt-vision predict --source image.jpg --model yolo11m --conf-thres 0.4 --output result.jpg +mblt-vision predict --source image.jpg --model yolo11m-pose --target-device regulus-ra --core-mode single +``` + +The corresponding `mblt-model-zoo` commands use the same standalone handlers for +backward compatibility. + +## Documentation and tests + +See [the Vision API guide](mblt_vision/README.md) for supported model families, +model details, artifact selection, and output taxonomy behavior. See the +[compilation guide](compile/README.md) for calibration-data preparation +and MXQ compilation. The [test guide](tests/TEST.md) explains offline, Hugging +Face, and NPU test runs. + +## Support and issues + +For installation, model, or runtime support, visit the +[Mobilint forum](https://discuss.mobilint.com/). Report reproducible package issues in the +[mblt-vision-python issue tracker](https://github.com/mobilint/mblt-vision-python/issues). + +## License + +Distributed under the [BSD 3-Clause License](LICENSE). diff --git a/benchmark/README.md b/benchmark/README.md new file mode 100644 index 0000000..d3dc3dd --- /dev/null +++ b/benchmark/README.md @@ -0,0 +1,50 @@ +# Vision benchmark commands + +`mblt-vision-python` is a Vision-only package, so every command script lives +directly in this directory. Reusable benchmark support—argument parsing, +artifact writing, charts, and summaries—is packaged under +`mblt_vision.benchmark` for use by library code and command scripts alike. + +The old per-dataset benchmark scripts and duplicated organizer wrappers were +removed. Use the unified runner and the organizer that matches your dataset. + +## Organize a dataset + +Public datasets can use their default download sources. ImageNet, DOTA, and +Cityscapes require the appropriate source archives or credentials. + +```bash +python benchmark/organize_coco.py +python benchmark/organize_ade20k.py +python benchmark/organize_nyu_depth.py +``` + +For Cityscapes, provide the official archives: + +```bash +python benchmark/organize_cityscapes.py \ + --image-dir path/to/leftImg8bit_trainvaltest.zip \ + --annotation-dir path/to/gtFine_trainvaltest.zip +``` + +## Run a benchmark + +Use the unified runner for every Vision task. It chooses the evaluator from +`--task`, writes JSON/CSV/Markdown artifacts, and can create an accuracy chart. + +```bash +python benchmark/benchmark_vision_models.py \ + --models ResNet50 \ + --task image_classification \ + --target-device aries-rb \ + --data-path ~/.mblt_model_zoo/datasets/imagenet +``` + +Use `--framework onnx` for ONNX Runtime, `--core-mode all` to compare supported +MXQ core modes, `--target-device regulus-ra` or `regulus-rb` for the corresponding +Regulus board artifact, and `--fail-fast` to stop on the first failed target. Compare +completed runs with: + +```bash +python benchmark/compare_benchmark_results.py run-a run-b +``` diff --git a/benchmark/benchmark_vision_models.py b/benchmark/benchmark_vision_models.py new file mode 100644 index 0000000..489f3d2 --- /dev/null +++ b/benchmark/benchmark_vision_models.py @@ -0,0 +1,518 @@ +"""Run reproducible multi-model vision accuracy benchmarks. + +This entry point follows the artifact contract used by the Transformers benchmark: +one results directory contains machine-readable JSON and CSV output, an optional +summary, and an accuracy chart. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Literal, cast + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.benchmark.argparse_utils import parse_positive_int +from mblt_vision.benchmark.io_utils import safe_filename, write_csv, write_json +from mblt_vision.benchmark.summary_utils import ( + collect_host_pc_info, + markdown_table, + write_summary_markdown, +) +from mblt_vision._model_paths import resolve_framework +from mblt_vision._tasks import VISION_TASKS, normalize_vision_task +from mblt_vision.wrapper import core_modes_for_target_device +from mblt_npu import normalize_target_device + +CoreMode = Literal["single", "multi", "global4", "global8"] +CORE_MODES: tuple[CoreMode, ...] = cast( + tuple[CoreMode, ...], core_modes_for_target_device("aries-rb") +) +TASK_CHOICES = VISION_TASKS +SUPPORTED_TARGET_DEVICES = frozenset({"aries-rb", "regulus-ra", "regulus-rb"}) + + +def parse_unit_interval(value: str) -> float: + """Parse a floating-point value strictly between zero and one.""" + + try: + parsed = float(value) + except ValueError as exc: + raise argparse.ArgumentTypeError( + "expected a number in the open interval (0, 1)" + ) from exc + if not 0 < parsed < 1: + raise argparse.ArgumentTypeError( + f"expected a number in the open interval (0, 1), got {value}" + ) + return parsed + + +def normalize_core_mode(core_mode: str) -> CoreMode: + """Validate and narrow a benchmark NPU core mode.""" + + if core_mode not in CORE_MODES: + raise ValueError( + f"Invalid core mode {core_mode!r}; expected one of {list(CORE_MODES)}." + ) + return cast(CoreMode, core_mode) + + +def _parse_task(value: str) -> str: + """Normalize a benchmark task for argparse.""" + + try: + return normalize_vision_task(value) + except (TypeError, ValueError) as exc: + raise argparse.ArgumentTypeError(str(exc)) from exc + + +def _parse_target_device(value: str) -> str: + """Normalize and validate a benchmark target board.""" + + try: + target_device = normalize_target_device(value) + except TypeError as exc: + raise argparse.ArgumentTypeError(str(exc)) from exc + if target_device not in SUPPORTED_TARGET_DEVICES: + raise argparse.ArgumentTypeError( + f"unsupported target device {value!r}; expected one of " + f"{sorted(SUPPORTED_TARGET_DEVICES)}." + ) + return target_device + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments for the standardized vision benchmark. + + Args: + argv: Optional argument sequence. ``None`` reads process arguments. + + Returns: + Parsed benchmark options. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--models", nargs="+", required=True, help="Vision model classes to benchmark." + ) + parser.add_argument( + "--task", + type=_parse_task, + choices=TASK_CHOICES, + required=True, + help="Task shared by all requested models.", + ) + parser.add_argument( + "--model-type", + default="DEFAULT", + help="Model variant from the YAML configuration.", + ) + parser.add_argument( + "--model-path", + default="", + help="Optional local MXQ or ONNX model path for one target.", + ) + parser.add_argument( + "--mxq-path", default="", help="Compatibility alias for a local MXQ path." + ) + parser.add_argument( + "--onnx-path", default="", help="Optional local ONNX path for one target." + ) + parser.add_argument( + "--framework", choices=["mxq", "onnx"], help="Explicit inference framework." + ) + parser.add_argument( + "--target-device", + type=_parse_target_device, + default="aries-rb", + help="NPU board: aries-rb, regulus-ra, or regulus-rb (legacy aries/regulus accepted).", + ) + parser.add_argument( + "--core-mode", + default=None, + choices=[*CORE_MODES, "all"], + help=( + "NPU core mode, or `all` to run every mode supported by the selected " + "board. Defaults to global8 on Aries and single on Regulus." + ), + ) + parser.add_argument("--dev-no", type=int, default=0, help="NPU device number.") + parser.add_argument( + "--batch-size", + type=parse_positive_int, + default=1, + help="Validation batch size.", + ) + parser.add_argument( + "--data-path", required=True, help="Path to an organized validation dataset." + ) + parser.add_argument( + "--conf-thres", + type=parse_unit_interval, + default=None, + help="Optional confidence threshold override.", + ) + parser.add_argument( + "--iou-thres", + type=parse_unit_interval, + default=None, + help="Optional IoU threshold override.", + ) + parser.add_argument( + "--results-dir", + type=Path, + default=Path("benchmark/results"), + help="Output directory.", + ) + parser.add_argument( + "--no-plot", action="store_true", help="Do not write an accuracy chart." + ) + parser.add_argument( + "--collect-host-info", + action="store_true", + help="Collect host metadata with mblt-tracker.", + ) + parser.add_argument( + "--fail-fast", + action="store_true", + help="Stop after the first failed model run.", + ) + args = parser.parse_args(argv) + local_paths = [args.model_path, args.mxq_path, args.onnx_path] + if len(args.models) != 1 and any(local_paths): + parser.error( + "--model-path, --mxq-path, and --onnx-path require exactly one --models target." + ) + return args + + +def _core_modes( + core_mode: str | None, + framework: str | None = None, + target_device: str = "aries-rb", +) -> tuple[str, ...]: + """Expand the multi-run core-mode shorthand. + + Args: + core_mode: Requested core mode, or ``None`` for the board default. + framework: Resolved inference framework. + target_device: Selected NPU board. + + Returns: + One or more concrete core modes. + """ + if framework == "onnx": + return ("onnx",) + supported_modes = core_modes_for_target_device(target_device) + if core_mode is None: + return (supported_modes[-1],) + if core_mode == "all": + return supported_modes + normalized_core_mode = normalize_core_mode(core_mode) + if normalized_core_mode not in supported_modes: + raise ValueError( + f"Core mode {normalized_core_mode!r} is not supported by " + f"{normalize_target_device(target_device)}; expected one of " + f"{list(supported_modes)}." + ) + return (normalized_core_mode,) + + +def _evaluate( + model: Any, args: argparse.Namespace, run_dir: Path +) -> tuple[float, str, dict[str, float]]: + """Evaluate one model and normalize task metrics for benchmark artifacts. + + Args: + model: Initialized vision engine. + args: Parsed benchmark options. + run_dir: Per-run output directory. + + Returns: + Primary score, its name, and all normalized metrics. + + Raises: + ValueError: If the model task differs from the requested benchmark task. + """ + from mblt_vision.utils.evaluation import ( + eval_ade20k, + eval_cityscapes, + eval_coco_metrics, + eval_dota, + eval_imagenet_metrics, + eval_nyu_depth, + eval_widerface, + ) + + model_task = normalize_vision_task(model.post_cfg.get("task", "")) + if model_task != args.task: + raise ValueError( + f"Model task '{model_task}' does not match requested task '{args.task}'." + ) + if args.task == "image_classification": + result = eval_imagenet_metrics(model, args.data_path, args.batch_size) + return ( + float(result.primary_score), + "top1_accuracy", + { + "top1_accuracy": float(result.top1), + "top5_accuracy": float(result.top5), + }, + ) + if args.task in {"object_detection", "instance_segmentation", "pose_estimation"}: + result = eval_coco_metrics( + model, args.data_path, args.batch_size, args.conf_thres, args.iou_thres + ) + return ( + float(result.primary_score), + "map50_95", + {"map50_95": float(result.map5095), "map50": float(result.map50)}, + ) + if args.task == "depth_estimation": + result = eval_nyu_depth(model, args.data_path, args.batch_size) + return ( + float(result.primary_score), + "delta1", + { + "delta1": float(result.delta1), + "abs_rel": float(result.abs_rel), + "rmse": float(result.rmse), + }, + ) + if args.task == "semantic_segmentation": + dataset = str(model.post_cfg.get("dataset", "")).lower() + if dataset == "ade20k": + result = eval_ade20k(model, args.data_path, args.batch_size) + elif dataset == "cityscapes": + result = eval_cityscapes(model, args.data_path, args.batch_size) + else: + raise ValueError( + f"Unsupported semantic segmentation benchmark taxonomy {dataset!r}; expected 'ade20k' or 'cityscapes'." + ) + return ( + float(result.primary_score), + "miou", + { + "miou": float(result.miou), + "pixel_accuracy": float(result.pixel_accuracy), + }, + ) + if args.task == "obb": + result = eval_dota( + model, + args.data_path, + args.batch_size, + args.conf_thres, + args.iou_thres, + str(run_dir / "dota_task1"), + ) + return ( + float(result.primary_score), + "map50_95", + { + "map50_95": float(result.map5095), + "map50": float(result.map50), + }, + ) + if args.task == "face_detection": + result = eval_widerface( + model, args.data_path, args.batch_size, args.conf_thres, args.iou_thres + ) + return ( + float(result.primary_score), + "mean_ap", + { + "easy_ap": float(result.easy_ap), + "medium_ap": float(result.medium_ap), + "hard_ap": float(result.hard_ap), + "mean_ap": float(result.mean_ap), + }, + ) + raise ValueError(f"Unsupported vision benchmark task: {args.task}") + + +def _run_target( + model_name: str, core_mode: str, args: argparse.Namespace, results_dir: Path +) -> dict[str, Any]: + """Run and record one model/core-mode benchmark target. + + Args: + model_name: Vision model class name. + core_mode: Concrete NPU core mode or the neutral ONNX runtime label. + args: Parsed benchmark options. + results_dir: Root directory for benchmark artifacts. + + Returns: + A normalized benchmark result row. + """ + from mblt_vision import MBLT_Engine + + label = f"{model_name}@{core_mode}" + run_dir = results_dir / "runs" / safe_filename(label) + run_dir.mkdir(parents=True, exist_ok=True) + row: dict[str, Any] = { + "model": model_name, + "core_mode": core_mode, + "task": args.task, + "batch_size": args.batch_size, + "target_device": args.target_device, + "status": "error", + } + model = None + started = time.perf_counter() + try: + engine_kwargs: dict[str, Any] = { + "model_cls": model_name, + "model_type": args.model_type, + "model_path": args.model_path, + "mxq_path": args.mxq_path, + "onnx_path": args.onnx_path, + "framework": args.framework, + "dev_no": args.dev_no, + "target_device": args.target_device, + } + if core_mode != "onnx": + engine_kwargs["core_mode"] = core_mode + model = MBLT_Engine( + **engine_kwargs, + ) + score, score_name, metrics = _evaluate(model, args, run_dir) + row.update( + {"status": "ok", "score": score, "score_name": score_name, **metrics} + ) + except ( + ImportError, + OSError, + RuntimeError, + TypeError, + ValueError, + NotImplementedError, + ) as exc: + row["error"] = f"{type(exc).__name__}: {exc}" + finally: + if model is not None: + try: + model.dispose() + except Exception as exc: + cleanup_error = f"{type(exc).__name__}: {exc}" + if "error" in row: + row["cleanup_error"] = cleanup_error + else: + row["error"] = cleanup_error + row["status"] = "error" + row["elapsed_s"] = round(time.perf_counter() - started, 6) + return row + + +def _write_outputs( + rows: list[dict[str, Any]], args: argparse.Namespace, results_dir: Path +) -> None: + """Write the shared JSON, CSV, chart, and Markdown benchmark artifacts. + + Args: + rows: Normalized benchmark rows. + args: Parsed benchmark options. + results_dir: Destination directory. + """ + results_dir.mkdir(parents=True, exist_ok=True) + payload = { + "schema_version": 1, + "benchmark": "vision", + "task": args.task, + "results": rows, + } + write_json(results_dir / "results.json", payload) + write_csv(results_dir / "results.csv", rows) + + successful = [row for row in rows if row["status"] == "ok"] + plot_paths: list[Path] = [] + if successful and not args.no_plot: + from mblt_vision.benchmark.chart_utils import plot_simple_barh + + chart_path = results_dir / "accuracy.png" + plot_simple_barh( + labels=[f"{row['model']} ({row['core_mode']})" for row in successful], + values=[float(row["score"]) for row in successful], + x_label="Accuracy score", + title=f"{args.task} benchmark accuracy", + output_path=chart_path, + ) + plot_paths.append(chart_path) + + host_info_path = ( + collect_host_pc_info(results_dir) if args.collect_host_info else None + ) + table = markdown_table( + ["Model", "Core mode", "Metric", "Score", "Elapsed (s)", "Status"], + [ + [ + row["model"], + row["core_mode"], + row.get("score_name", "-"), + f"{float(row['score']):.5f}" if row.get("score") is not None else "-", + row["elapsed_s"], + row["status"], + ] + for row in rows + ], + ) + table_path = results_dir / "results.md" + table_path.write_text(table, encoding="utf-8") + write_summary_markdown( + results_dir / "summary.md", + title=f"Vision benchmark: {args.task}", + host_info_path=host_info_path, + table_markdown_path=table_path, + plot_paths=plot_paths, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """Run the standardized vision benchmark. + + Args: + argv: Optional command-line arguments. + + Returns: + Zero when every target succeeds, otherwise one. + """ + args = _parse_args(argv) + results_dir = args.results_dir.expanduser().resolve() + # Keep compatibility-path routing aligned with MBLT_Engine: an explicit + # ONNX path selects the ONNX backend unless an MXQ path takes precedence. + framework_model_path = args.model_path + if not framework_model_path and not args.mxq_path: + framework_model_path = args.onnx_path + framework = resolve_framework(args.framework, framework_model_path) + rows: list[dict[str, Any]] = [] + try: + core_modes = _core_modes(args.core_mode, framework, args.target_device) + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + for model_name in args.models: + for core_mode in core_modes: + print(f"Benchmarking {model_name} with core mode {core_mode}...") + row = _run_target(model_name, core_mode, args, results_dir) + rows.append(row) + if row["status"] == "ok": + print(f" {row['score_name']}: {row['score']:.5f}") + else: + print(f" failed: {row['error']}") + if args.fail_fast: + _write_outputs(rows, args, results_dir) + return 1 + _write_outputs(rows, args, results_dir) + print(f"Saved benchmark artifacts to: {results_dir}") + return 0 if all(row["status"] == "ok" for row in rows) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/compare_benchmark_results.py b/benchmark/compare_benchmark_results.py new file mode 100644 index 0000000..b67d96a --- /dev/null +++ b/benchmark/compare_benchmark_results.py @@ -0,0 +1,179 @@ +"""Compare standardized vision benchmark result directories.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import NamedTuple, Sequence + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.benchmark.summary_utils import read_csv_rows + + +class BenchmarkScore(NamedTuple): + """One successful standardized vision benchmark score.""" + + task: str + metric: str + value: float + + +def _resolve_results_csv(raw: str) -> Path: + """Resolve a standard results CSV file from a path argument. + + Args: + raw: A results directory or explicit CSV path. + + Returns: + Resolved CSV path. + + Raises: + ValueError: If the path is not a standard benchmark results file. + """ + path = Path(raw).expanduser().resolve() + if path.is_dir(): + path = path / "results.csv" + if not path.is_file() or path.name != "results.csv": + raise ValueError( + f"Expected a vision benchmark results.csv file or containing directory: {path}" + ) + return path + + +def _collect_scores(path: Path) -> dict[str, BenchmarkScore]: + """Read successful model/core-mode scores from one results CSV. + + Args: + path: Standardized benchmark ``results.csv`` path. + + Returns: + Scores keyed by model and core mode. + + Raises: + ValueError: If a successful result has invalid standardized fields. + """ + scores: dict[str, BenchmarkScore] = {} + for row in read_csv_rows(path): + if row.get("status") != "ok": + continue + model = (row.get("model") or "").strip() + core_mode = (row.get("core_mode") or "").strip() + task = (row.get("task") or "").strip() + metric = (row.get("score_name") or "").strip() + if not model or not core_mode or not task or not metric: + raise ValueError( + f"Missing model, core_mode, task, or score_name in {path}." + ) + try: + score = float(row["score"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + f"Invalid score for {model}@{core_mode} in {path}." + ) from exc + scores.setdefault( + f"{model}@{core_mode}", + BenchmarkScore(task=task, metric=metric, value=score), + ) + return scores + + +def _common_targets(scores_by_source: Sequence[dict[str, BenchmarkScore]]) -> list[str]: + """Return model/core-mode targets shared by all comparison sources. + + Args: + scores_by_source: Scores parsed from each result source. + + Returns: + Sorted shared target labels. + """ + if not scores_by_source: + return [] + return sorted(set.intersection(*(set(scores) for scores in scores_by_source))) + + +def main(argv: Sequence[str] | None = None) -> int: + """Generate a grouped score chart from standardized vision result files. + + Args: + argv: Optional command-line arguments. + + Returns: + Zero on success. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "inputs", + nargs="+", + help="At least two result directories or results.csv files.", + ) + parser.add_argument( + "--output-dir", type=Path, help="Directory for the grouped comparison chart." + ) + args = parser.parse_args(argv) + if len(args.inputs) < 2: + parser.error("Provide at least two benchmark result sources.") + try: + sources = [_resolve_results_csv(raw) for raw in args.inputs] + scores_by_source = [_collect_scores(source) for source in sources] + except ValueError as exc: + parser.error(str(exc)) + + task_names = { + score.task for scores in scores_by_source for score in scores.values() + } + if len(task_names) != 1: + raise SystemExit( + f"Inputs contain incompatible benchmark tasks: {', '.join(sorted(task_names))}." + ) + targets = _common_targets(scores_by_source) + if not targets: + raise SystemExit( + "No successful model/core-mode target is shared by all input sources." + ) + metric_names = { + score.metric for scores in scores_by_source for score in scores.values() + } + if len(metric_names) != 1: + raise SystemExit( + f"Inputs contain incompatible benchmark metrics: {', '.join(sorted(metric_names))}." + ) + metric_name = next(iter(metric_names)) + from mblt_vision.benchmark.chart_utils import ( + default_charts_dir, + plot_grouped_scalar_barh, + source_labels, + ) + + source_dirs = [source.parent for source in sources] + output_dir = ( + args.output_dir.expanduser().resolve() + if args.output_dir + else default_charts_dir( + Path(__file__).resolve().parent, source_dirs, use_stem=False + ) + ) + output_dir.mkdir(parents=True, exist_ok=True) + plot_grouped_scalar_barh( + models=targets, + group_labels=source_labels(source_dirs, use_stem=False), + grouped_values=[ + {target: scores[target].value for target in targets} + for scores in scores_by_source + ], + x_label=metric_name, + y_label="model@core_mode", + title=f"Vision benchmark comparison: {metric_name}", + output_path=output_dir / "score.png", + ) + print(f"Compared {len(targets)} shared targets using {metric_name}.") + print(f"Saved chart to: {output_dir / 'score.png'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/organize_ade20k.py b/benchmark/organize_ade20k.py new file mode 100644 index 0000000..67290e9 --- /dev/null +++ b/benchmark/organize_ade20k.py @@ -0,0 +1,28 @@ +"""Organize the ADE20K validation dataset for local use.""" + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.utils.datasets import organize_ade20k +from mblt_vision.utils.datasets.organizer import ADE20K_URL + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Organize ADE20K validation dataset") + parser.add_argument( + "--dataset-path", + default=ADE20K_URL, + help="Local path or download URL for ADE20K", + ) + parser.add_argument( + "--output-dir", + default=None, + help="Path to the organized validation dataset (defaults to cache)", + ) + args = parser.parse_args() + organize_ade20k(dataset_path=args.dataset_path, output_dir=args.output_dir) diff --git a/benchmark/organize_cityscapes.py b/benchmark/organize_cityscapes.py new file mode 100644 index 0000000..1966865 --- /dev/null +++ b/benchmark/organize_cityscapes.py @@ -0,0 +1,47 @@ +"""Organize the Cityscapes validation dataset for local use.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.utils.datasets import organize_cityscapes + + +def main() -> None: + """Parse organizer options and materialize Cityscapes validation data.""" + + parser = argparse.ArgumentParser( + description="Organize Cityscapes validation data from official ZIP archives" + ) + parser.add_argument( + "--image-dir", + required=True, + help="Path to leftImg8bit_trainvaltest.zip", + ) + parser.add_argument( + "--annotation-dir", + required=True, + help="Path to gtFine_trainvaltest.zip", + ) + parser.add_argument( + "--output-dir", + default=None, + help="Destination for the flat images/ and annotations/ directories (defaults to cache)", + ) + args = parser.parse_args() + organize_cityscapes( + image_dir=args.image_dir, + annotation_dir=args.annotation_dir, + output_dir=args.output_dir, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/organize_coco.py b/benchmark/organize_coco.py new file mode 100644 index 0000000..6ad9cf6 --- /dev/null +++ b/benchmark/organize_coco.py @@ -0,0 +1,50 @@ +""" +Script to organize the COCO dataset. + +This script takes local archives or downloadable sources for the COCO dataset +and organizes them into a structure suitable for the model zoo. +""" + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.datasets import get_dataset_config +from mblt_vision.utils.datasets import organize_coco + +COCO_DOWNLOAD_CONFIG = get_dataset_config("coco")["download"] +DEFAULT_COCO_IMAGE_SOURCE = COCO_DOWNLOAD_CONFIG["images"] +DEFAULT_COCO_ANNOTATION_SOURCE = COCO_DOWNLOAD_CONFIG["annotations"] + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Organize COCO dataset") + parser.add_argument( + "--image-dir", + type=str, + default=DEFAULT_COCO_IMAGE_SOURCE, + help="Local path or download URL for the image zip file", + ) + parser.add_argument( + "--ann-dir", + type=str, + default=DEFAULT_COCO_ANNOTATION_SOURCE, + help="Local path or download URL for the annotation zip file", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Path to the directory to save the organized dataset (defaults to cache)", + ) + args = parser.parse_args() + + organize_coco( + image_dir=args.image_dir, + annotation_dir=args.ann_dir, + output_dir=args.output_dir, + ) diff --git a/benchmark/organize_dotav1.py b/benchmark/organize_dotav1.py new file mode 100644 index 0000000..4a9b4eb --- /dev/null +++ b/benchmark/organize_dotav1.py @@ -0,0 +1,41 @@ +""" +Script to organize the DOTAv1 validation dataset. + +This script takes a local archive, extracted directory, or downloadable source +for DOTAv1 and organizes only the validation split into a structure suitable for +the model zoo. +""" + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.datasets import get_dataset_config +from mblt_vision.utils.datasets import organize_dotav1 + +DEFAULT_DOTAV1_SOURCE = get_dataset_config("dotav1")["download"]["url"] +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Organize DOTAv1 validation dataset") + parser.add_argument( + "--dataset-path", + type=str, + default=DEFAULT_DOTAV1_SOURCE, + help="Local path, archive URL, or Google Drive folder URL for DOTAv1", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Path to the directory to save the organized dataset (defaults to cache)", + ) + args = parser.parse_args() + + organize_dotav1( + dataset_path=args.dataset_path, + output_dir=args.output_dir, + ) diff --git a/benchmark/organize_imagenet.py b/benchmark/organize_imagenet.py new file mode 100644 index 0000000..cec0c93 --- /dev/null +++ b/benchmark/organize_imagenet.py @@ -0,0 +1,51 @@ +""" +Script to organize the ImageNet dataset. + +This script takes local archives or downloadable sources for the ImageNet +dataset and organizes them into a structure suitable for the model zoo. +""" + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.utils.datasets import organize_imagenet + +DEFAULT_IMAGENET_IMAGE_SOURCE = ( + "https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar" +) +DEFAULT_IMAGENET_XML_SOURCE = ( + "https://www.image-net.org/data/ILSVRC/2012/ILSVRC2012_bbox_val_v3.tgz" +) +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Organize ImageNet dataset") + parser.add_argument( + "--image-dir", + type=str, + default=DEFAULT_IMAGENET_IMAGE_SOURCE, + help="Local path or download URL for the image tar file", + ) + parser.add_argument( + "--xml-dir", + type=str, + default=DEFAULT_IMAGENET_XML_SOURCE, + help="Local path or download URL for the XML tgz file", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Path to the directory to save the organized dataset (defaults to cache)", + ) + args = parser.parse_args() + + organize_imagenet( + image_dir=args.image_dir, + xml_dir=args.xml_dir, + output_dir=args.output_dir, + ) diff --git a/benchmark/organize_nyu_depth.py b/benchmark/organize_nyu_depth.py new file mode 100644 index 0000000..e8187a1 --- /dev/null +++ b/benchmark/organize_nyu_depth.py @@ -0,0 +1,35 @@ +"""Organize the NYU Depth dataset for local use. + +The script downloads the published NYU Depth archive by default, or accepts a +local zip file or extracted dataset directory. +""" + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.utils.datasets import organize_nyu_depth +from mblt_vision.utils.datasets.organizer import NYU_DEPTH_URL + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Organize NYU Depth dataset") + parser.add_argument( + "--dataset-path", + type=str, + default=NYU_DEPTH_URL, + help="Local path or download URL for the NYU Depth zip file or extracted dataset", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Path to the directory to save the organized dataset (defaults to cache)", + ) + args = parser.parse_args() + + organize_nyu_depth(dataset_path=args.dataset_path, output_dir=args.output_dir) diff --git a/benchmark/organize_widerface.py b/benchmark/organize_widerface.py new file mode 100644 index 0000000..bebf78f --- /dev/null +++ b/benchmark/organize_widerface.py @@ -0,0 +1,48 @@ +""" +Script to organize the WiderFace dataset. + +This script takes local archives or downloadable sources for the WiderFace +dataset and organizes them into a structure suitable for the model zoo. +""" + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.utils.datasets import organize_widerface + +DEFAULT_WIDERFACE_IMAGE_SOURCE = "https://huggingface.co/datasets/CUHK-CSE/wider_face/resolve/main/data/WIDER_val.zip" +DEFAULT_WIDERFACE_ANNOTATION_SOURCE = "https://huggingface.co/datasets/CUHK-CSE/wider_face/resolve/main/data/wider_face_split.zip" + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Organize WiderFace dataset") + parser.add_argument( + "--image-dir", + type=str, + default=DEFAULT_WIDERFACE_IMAGE_SOURCE, + help="Local path or download URL for the image zip file", + ) + parser.add_argument( + "--annotation-dir", + type=str, + default=DEFAULT_WIDERFACE_ANNOTATION_SOURCE, + help="Local path or download URL for the annotation zip file", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Path to the directory to save the organized dataset (defaults to cache)", + ) + args = parser.parse_args() + + organize_widerface( + image_dir=args.image_dir, + annotation_dir=args.annotation_dir, + output_dir=args.output_dir, + ) diff --git a/compile/README.md b/compile/README.md new file mode 100644 index 0000000..c9f1350 --- /dev/null +++ b/compile/README.md @@ -0,0 +1,124 @@ +# Vision Model Compilation + +The installable vision compiler uses each model's packaged YAML configuration and +`MBLT_Engine.preprocess` implementation to compile its ONNX artifact for a selected board. + +## Installation + +Install the package with the compilation extra: + +```bash +pip install -e ".[qbcompiler]" +``` + +`qbcompiler` is currently distributed through Mobilint's external package channel and may not be +available from the public Python Package Index. Obtain access to the compiler package before using +this extra. + +qbcompiler is loaded only when `compile_vision_model()` or `mblt-vision compile` begins an +actual compilation request. Importing the base package, vision APIs, this compilation module, or +the main CLI does not import qbcompiler. If qbcompiler is unavailable, only the compilation request +fails with an installation message; non-compile APIs and CLI commands remain usable. + +## Python API + +```python +from mblt_vision.compile.vision import compile_vision_model + +output_path = compile_vision_model( + "alexnet", + target_device="aries-rb", + data_path="~/.mblt_model_zoo/datasets/imagenet", + save_path="./alexnet.mxq", +) +print(output_path) +``` + +## CLI + +```bash +mblt-vision compile \ + --model-cls alexnet \ + --target-device aries-rb \ + --data-path ~/.mblt_model_zoo/datasets/imagenet \ + --save-path ./alexnet.mxq +``` + +Start from an existing sampled image subset or a ready calibration tensor directory when those +stages have already been completed: + +```bash +mblt-vision compile --model-cls alexnet --target-device aries-rb --subset-path ./sampled-images +mblt-vision compile --model-cls alexnet --target-device aries-rb --calib-data-path ./preprocessed-npy +``` + +Use `--model-type` for a non-default YAML variant and `--model-path` or `--onnx-path` to prefer a +local ONNX file. If that path is omitted or does not exist, the configured Hugging Face repository +supplies the ONNX artifact. Downloaded ONNX files and compiled outputs use +`~/.mblt_model_zoo`; without `--save-path`, the compiler writes +`~/.mblt_model_zoo/.mxq`. + +`target_device` / `--target-device` is required: choose `aries-rb`, `regulus-ra`, or `regulus-rb` +to compile an artifact for that board. + +The ONNX preprocessing engine uses ONNX Runtime's CPU provider by default, so compilation does not +probe TensorRT, CUDA, or other accelerators. Callers that construct `MBLT_Engine` directly can opt +into another provider order with `onnx_providers`. + +## Calibration Data + +The compilation data pipeline has three levels: + +1. `data_path` / `--data-path`: Original organized dataset containing all images. The pipeline + prepares the dataset when needed, samples a task-specific subset, and preprocesses it. +2. `subset_path` / `--subset-path`: Already-sampled images. The pipeline skips original dataset + preparation and subset generation, then preprocesses every image in the supplied folder. +3. `calib_data_path` / `--calib-data-path`: Ready HWC, three-channel, contiguous `float32` `.npy` + tensors. The directory is validated and passed directly to qbcompiler without image processing. + +Supply at most one of these paths. When none is supplied, the task's registry-backed original +dataset path is used. `--calib-data-dir` remains an alias for `--calib-data-path` in the CLI and +standalone compatibility wrapper. + +Compilation maps model tasks to the packaged dataset registry: + +- Image classification uses ImageNet. +- Depth estimation uses NYU Depth V2. +- Object detection, instance segmentation, and pose estimation use COCO. +- Semantic segmentation uses the dataset declared by the model's `post_cfg.dataset`: ADE20K or + Cityscapes. +- Face detection uses WiderFace. +- OBB uses DOTAv1. + +`--data-path` is an organized dataset root. An existing ready layout is reused; otherwise the +registry-backed organizer downloads and prepares the dataset. When no path is supplied, the +registry default under `~/.mblt_model_zoo/datasets` is used. + +NYU Depth and ADE20K use their registry download URLs. Cityscapes requires the manually downloaded +official `leftImg8bit_trainvaltest.zip` and `gtFine_trainvaltest.zip` archives; place both inside the +organized dataset path or its parent before compiling a Cityscapes model. + +Selection is deterministic with seed `0`. ImageNet and WiderFace select one image from every +category subfolder by default; COCO, DOTAv1, NYU Depth, ADE20K, and Cityscapes select 100 images +total. Override these values with `--subset-size` and `--seed`. Selected images and preprocessed +NumPy arrays live only in temporary directories and are removed after compilation, including when +compilation fails. + +The compiler uses explicit `--percentile` and `--topk-ratio` values first. Missing values are read +independently from `/best_result.json`; if optional hosted values are unavailable, defaults of +`0.9999` and `0.01` are used with a warning. + +## Compatibility Scripts + +The standalone scripts remain available: + +```bash +python compile/vision_model_compile.py --model-cls alexnet --target-device aries-rb +python compile/make_imagenet_subset.py --output-dir ./imagenet-calibration +python compile/make_coco_subset.py --output-dir ./coco-calibration +python compile/make_dotav1_subset.py --output-dir ./dotav1-calibration +python compile/make_widerface_subset.py --output-dir ./widerface-calibration +python compile/make_nyu_depth_subset.py --output-dir ./nyu-depth-calibration +python compile/make_ade20k_subset.py --output-dir ./ade20k-calibration +python compile/make_cityscapes_subset.py --output-dir ./cityscapes-calibration +``` diff --git a/compile/make_ade20k_subset.py b/compile/make_ade20k_subset.py new file mode 100644 index 0000000..43df262 --- /dev/null +++ b/compile/make_ade20k_subset.py @@ -0,0 +1,54 @@ +"""Compatibility entry point for ADE20K calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/ADEChallengeData2016" + + +def make_ade20k_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic ADE20K calibration subset. + + Args: + data_dir: Organized ADE20K root. + output_dir: Flat subset destination. + subset_size: Total image count. + seed: Random selection seed. + """ + + copied = make_calibration_subset( + "semantic_segmentation", data_dir, output_dir, subset_size, seed + ) + print(f"Created ADE20K subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Create an ADE20K calibration subset") + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized ADE20K dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", type=int, default=100, help="Number of images to select" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_ade20k_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/make_cityscapes_subset.py b/compile/make_cityscapes_subset.py new file mode 100644 index 0000000..eeaae62 --- /dev/null +++ b/compile/make_cityscapes_subset.py @@ -0,0 +1,56 @@ +"""Compatibility entry point for Cityscapes calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/cityscapes" + + +def make_cityscapes_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic Cityscapes calibration subset. + + Args: + data_dir: Organized Cityscapes root. + output_dir: Flat subset destination. + subset_size: Total image count. + seed: Random selection seed. + """ + + copied = make_calibration_subset( + "semantic_segmentation", data_dir, output_dir, subset_size, seed + ) + print(f"Created Cityscapes subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a Cityscapes calibration subset" + ) + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized Cityscapes dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", type=int, default=100, help="Number of images to select" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_cityscapes_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/make_coco_subset.py b/compile/make_coco_subset.py new file mode 100644 index 0000000..c6c1e2d --- /dev/null +++ b/compile/make_coco_subset.py @@ -0,0 +1,54 @@ +"""Compatibility entry point for COCO calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/coco" + + +def make_coco_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic COCO calibration subset. + + Args: + data_dir: Organized COCO root. + output_dir: Flat subset destination. + subset_size: Total image count. + seed: Random selection seed. + """ + + copied = make_calibration_subset( + "object_detection", data_dir, output_dir, subset_size, seed + ) + print(f"Created COCO subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Create a COCO calibration subset") + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized COCO dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", type=int, default=100, help="Number of images to select" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_coco_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/make_dotav1_subset.py b/compile/make_dotav1_subset.py new file mode 100644 index 0000000..5ff5b77 --- /dev/null +++ b/compile/make_dotav1_subset.py @@ -0,0 +1,52 @@ +"""Compatibility entry point for DOTAv1 calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/dotav1" + + +def make_dotav1_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic DOTAv1 calibration subset. + + Args: + data_dir: Organized DOTAv1 root. + output_dir: Flat subset destination. + subset_size: Total image count. + seed: Random selection seed. + """ + + copied = make_calibration_subset("obb", data_dir, output_dir, subset_size, seed) + print(f"Created DOTAv1 subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Create a DOTAv1 calibration subset") + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized DOTAv1 dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", type=int, default=100, help="Number of images to select" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_dotav1_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/make_imagenet_subset.py b/compile/make_imagenet_subset.py new file mode 100644 index 0000000..e108e5b --- /dev/null +++ b/compile/make_imagenet_subset.py @@ -0,0 +1,59 @@ +"""Compatibility entry point for ImageNet calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/imagenet" + + +def make_imagenet_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic ImageNet calibration subset. + + Args: + data_dir: Organized ImageNet root. + output_dir: Flat subset destination. + subset_size: Images selected per class. + seed: Random selection seed. + """ + + copied = make_calibration_subset( + "image_classification", data_dir, output_dir, subset_size, seed + ) + print(f"Created ImageNet subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create an ImageNet calibration subset" + ) + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized ImageNet dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", + type=int, + default=1, + help="Number of images to select per category", + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_imagenet_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/make_nyu_depth_subset.py b/compile/make_nyu_depth_subset.py new file mode 100644 index 0000000..04e6d25 --- /dev/null +++ b/compile/make_nyu_depth_subset.py @@ -0,0 +1,56 @@ +"""Compatibility entry point for NYU Depth calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/nyu-depth" + + +def make_nyu_depth_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic NYU Depth calibration subset. + + Args: + data_dir: Organized NYU Depth root. + output_dir: Flat subset destination. + subset_size: Total image count. + seed: Random selection seed. + """ + + copied = make_calibration_subset( + "depth_estimation", data_dir, output_dir, subset_size, seed + ) + print(f"Created NYU Depth subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create an NYU Depth calibration subset" + ) + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized NYU Depth dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", type=int, default=100, help="Number of images to select" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_nyu_depth_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/make_widerface_subset.py b/compile/make_widerface_subset.py new file mode 100644 index 0000000..1344aee --- /dev/null +++ b/compile/make_widerface_subset.py @@ -0,0 +1,59 @@ +"""Compatibility entry point for WiderFace calibration subsets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import make_calibration_subset + +DEFAULT_DATA_DIR = "~/.mblt_model_zoo/datasets/widerface" + + +def make_widerface_subset( + data_dir: str, output_dir: str, subset_size: int, seed: int = 0 +) -> None: + """Create a deterministic WiderFace calibration subset. + + Args: + data_dir: Organized WiderFace root. + output_dir: Flat subset destination. + subset_size: Images selected per category. + seed: Random selection seed. + """ + + copied = make_calibration_subset( + "face_detection", data_dir, output_dir, subset_size, seed + ) + print(f"Created WiderFace subset with {len(copied)} images at {output_dir}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a WiderFace calibration subset" + ) + parser.add_argument( + "--data-dir", + default=DEFAULT_DATA_DIR, + help="Path to the organized WiderFace dataset", + ) + parser.add_argument( + "--output-dir", required=True, help="Path to save the selected images" + ) + parser.add_argument( + "--subset-size", + type=int, + default=1, + help="Number of images to select per category", + ) + parser.add_argument( + "--seed", type=int, default=0, help="Random seed used to select images" + ) + args = parser.parse_args() + make_widerface_subset(args.data_dir, args.output_dir, args.subset_size, args.seed) diff --git a/compile/vision_model_compile.py b/compile/vision_model_compile.py new file mode 100644 index 0000000..da48c8b --- /dev/null +++ b/compile/vision_model_compile.py @@ -0,0 +1,97 @@ +"""Compatibility entry point for packaged vision compilation.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# ruff: noqa: E402 +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from mblt_vision.compile.vision import compile_vision_model + + +def build_parser() -> argparse.ArgumentParser: + """Build the standalone compatibility parser. + + Returns: + Configured argument parser. + """ + + parser = argparse.ArgumentParser( + description="Compile a configured vision ONNX model to MXQ" + ) + parser.add_argument("--model-cls", required=True, help="Vision model name.") + parser.add_argument( + "--model-type", + default="DEFAULT", + help="Model variant from the YAML configuration.", + ) + parser.add_argument( + "--target-device", + default="aries-rb", + choices=["aries-rb", "regulus-ra", "regulus-rb"], + help="Board target for the compiled MXQ artifact.", + ) + parser.add_argument( + "--model-path", + "--onnx-path", + dest="model_path", + help="Preferred local ONNX file.", + ) + data_group = parser.add_mutually_exclusive_group() + data_group.add_argument("--data-path", help="Original organized dataset root.") + data_group.add_argument("--subset-path", help="Already-sampled image subset root.") + data_group.add_argument( + "--calib-data-path", + "--calib-data-dir", + dest="calib_data_path", + help="Ready directory of preprocessed .npy tensors.", + ) + parser.add_argument( + "--save-path", + help="Output MXQ path. Defaults to the ONNX stem under ~/.mblt_model_zoo.", + ) + parser.add_argument( + "--subset-size", + type=int, + help="Per-category ImageNet/WiderFace count or total count for other datasets.", + ) + parser.add_argument( + "--seed", type=int, default=0, help="Calibration selection seed." + ) + parser.add_argument( + "--percentile", type=float, help="Quantization percentile override." + ) + parser.add_argument( + "--topk-ratio", type=float, help="Quantization top-k ratio override." + ) + return parser + + +def main() -> None: + """Run standalone vision compilation.""" + + args = build_parser().parse_args() + output_path = compile_vision_model( + model_cls=args.model_cls, + model_type=args.model_type, + target_device=args.target_device, + model_path=args.model_path, + data_path=args.data_path, + subset_path=args.subset_path, + calib_data_path=args.calib_data_path, + save_path=args.save_path, + subset_size=args.subset_size, + seed=args.seed, + percentile=args.percentile, + topk_ratio=args.topk_ratio, + ) + print(f"Compiled MXQ model to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/mblt_vision/README.md b/mblt_vision/README.md new file mode 100644 index 0000000..1ee0e47 --- /dev/null +++ b/mblt_vision/README.md @@ -0,0 +1,362 @@ +# Mobilint Vision API + +`mblt-vision-python` provides Python access to pre-trained Mobilint Vision models: +image classification, depth estimation, face and object detection, oriented bounding +boxes (OBB), instance and semantic segmentation, and pose estimation. Each model +configuration includes the artifact, preprocessing, output taxonomy, and +postprocessing contract needed to produce task-specific results. + +## Loading models + +Use a task subpackage for a known model, or `MBLT_Engine` to choose a model and +artifact at runtime. These imports are both supported: + +```python +from mblt_vision import ResNet50 +from mblt_vision import YOLO11m +``` + +```python +from mblt_vision import MBLT_Engine + +model = MBLT_Engine(model_cls="resnet50", model_type="DEFAULT", model_path="", core_mode="global8") +``` + +```python +from mblt_vision import MBLT_Engine + +# Install the optional `onnxruntime` or `onnxruntime-gpu` extra first. +model = MBLT_Engine(model_cls="alexnet", framework="onnx") +``` + +When `framework="onnx"` is selected, the engine uses ONNX Runtime's `CPUExecutionProvider` by +default. Pass `onnx_providers` to select another provider order explicitly when needed. If +`model_path` or `file_cfg.model_path` ends +with `.mxq` or `.onnx`, the engine auto-detects the framework from that suffix when `framework` is +omitted. + +For Hub-backed models, `file_cfg.filename` is the canonical MXQ artifact name. The engine derives +the corresponding ONNX artifact by replacing its `.mxq` suffix with `.onnx`, so every model uses +the same ONNX loading path. Set `file_cfg.onnx_filename` only when a Hub repository uses a +non-matching ONNX filename. + +Every model configuration also declares `post_cfg.dataset`. This identifies the output taxonomy, +not merely the broad task: COCO detection uses 80 classes, Cityscapes semantic segmentation uses +19, and ADE20K semantic segmentation uses 150. YOLO postprocessing resolves shape-sensitive class +counts from the dataset and task together. + +```python +from mblt_vision.image_classification import ResNet50 +from mblt_vision.object_detection import YOLO11m +``` + +The task subpackages are the clearest import surface; `MBLT_Engine`, `list_tasks()`, +and `list_models()` are the preferred discovery and loading APIs for new code. +Top-level model imports remain available for compatibility. `obb` is the sole +oriented-bounding-box task name. + +## Pre-Trained Vision Models + +This section lists the publicly pre-trained models supported by the vision framework. + +### Image Classification + +| Model | Input Size
(H,W,C) | AccTop1
(NPU) | AccTop1
(GPU) | FLOPs (B) | params (M) | Source | Note | +| --- | --- | --- | --- | --- | --- | --- | --- | +| AlexNet | (224,224,3) | 56.084 | 56.556 | 1.43 | 61.10 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.alexnet.html) | | +| CAFormer_S18 | (224,224,3) | 82.592 | 83.626 | 9.10 | 26.34 | [Link](https://huggingface.co/timm/caformer_s18.sail_in1k) | sail_in1k | +| CAFormer_B36 | (224,224,3) | 83.938 | 85.482 | 49.41 | 98.75 | [Link](https://huggingface.co/timm/caformer_b36.sail_in1k) | sail_in1k | +| CoAtNet_0_RW_224 | (224,224,3) | 81.842 | 82.418 | 10.06 | 30.46 | [Link](https://huggingface.co/timm/coatnet_0_rw_224.sw_in1k) | sw_in1k | +| CoAtNet_1_RW_224 | (224,224,3) | 83.506 | 83.600 | 18.17 | 47.91 | [Link](https://huggingface.co/timm/coatnet_1_rw_224.sw_in1k) | sw_in1k | +| CoAtNet_2_RW_224 | (224,224,3) | 86.084 | 86.534 | 32.80 | 82.11 | [Link](https://huggingface.co/timm/coatnet_2_rw_224.sw_in12k_ft_in1k) | sw_in12k_ft_in1k | +| ConvFormer S36 | (224,224,3) | 83.360 | 84.016 | 16.66 | 40.01 | [Link](https://huggingface.co/timm/convformer_s36.sail_in1k) | sail_in1k | +| ConvFormer M36 | (224,224,3) | 84.014 | 84.448 | 27.58 | 57.05 | [Link](https://huggingface.co/timm/convformer_m36.sail_in1k) | sail_in1k | +| ConvFormer B36 | (224,224,3) | 84.244 | 84.830 | 47.79 | 99.88 | [Link](https://huggingface.co/timm/convformer_b36.sail_in1k) | sail_in1k | +| ConvNeXt_Tiny | (224,224,3) | 82.354 | 82.458 | 9.11 | 28.59 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.convnext_tiny.html) | | +| ConvNeXt_Small | (224,224,3) | 83.434 | 83.560 | 17.68 | 50.22 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.convnext_small.html) | | +| ConvNeXt_Base | (224,224,3) | 83.940 | 84.048 | 31.13 | 88.59 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.convnext_base.html) | | +| ConvNeXt_Large | (224,224,3) | 84.316 | 84.410 | 69.35 | 197.77 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.convnext_large.html) | | +| DeiT_Tiny_Patch16_224 | (224,224,3) | 71.944 | 72.030 | 2.66 | 5.72 | [Link](https://huggingface.co/timm/deit_tiny_patch16_224.fb_in1k) | fb_in1k | +| DeiT_Small_Patch16_224 | (224,224,3) | 79.722 | 79.790 | 9.50 | 22.05 | [Link](https://huggingface.co/timm/deit_small_patch16_224.fb_in1k) | fb_in1k | +| DeiT_Base_Patch16_224 | (224,224,3) | 81.932 | 81.980 | 35.73 | 86.57 | [Link](https://huggingface.co/timm/deit_base_patch16_224.fb_in1k) | fb_in1k | +| DeiT_Base_Patch16_384 | (384,384,3) | 83.046 | 83.100 | 115.00 | 86.86 | [Link](https://huggingface.co/timm/deit_base_patch16_384.fb_in1k) | fb_in1k | +| DeiT3_Small_Patch16_224 | (224,224,3) | 81.368 | 81.390 | 9.50 | 22.06 | [Link](https://huggingface.co/timm/deit3_small_patch16_224.fb_in1k) | fb_in1k | +| DeiT3_Small_Patch16_384 | (384,384,3) | 83.350 | 83.420 | 33.01 | 22.21 | [Link](https://huggingface.co/timm/deit3_small_patch16_384.fb_in1k) | fb_in1k | +| DeiT3_Medium_Patch16_224 | (224,224,3) | 83.020 | 83.044 | 16.39 | 38.85 | [Link](https://huggingface.co/timm/deit3_medium_patch16_224.fb_in1k) | fb_in1k | +| DeiT3_Base_Patch16_224 | (224,224,3) | 83.708 | 83.766 | 35.74 | 86.59 | [Link](https://huggingface.co/timm/deit3_base_patch16_224.fb_in1k) | fb_in1k | +| DeiT3_Base_Patch16_384 | (384,384,3) | 84.986 | 85.064 | 115.02 | 86.88 | [Link](https://huggingface.co/timm/deit3_base_patch16_384.fb_in1k) | fb_in1k | +| DeiT3_Large_Patch16_224 | (224,224,3) | 84.734 | 84.736 | 124.73 | 304.37 | [Link](https://huggingface.co/timm/deit3_large_patch16_224.fb_in1k) | fb_in1k | +| DeiT3_Large_Patch16_384 | (384,384,3) | 85.806 | 85.826 | 392.94 | 304.76 | [Link](https://huggingface.co/timm/deit3_large_patch16_384.fb_in1k) | fb_in1k | +| ConvFormer S18 | (224,224,3) | 81.862 | 82.866 | 8.59 | 26.77 | [Link](https://huggingface.co/timm/convformer_s18.sail_in1k) | sail_in1k | +| EfficientFormer_L7 | (224,224,3) | 82.526 | 83.352 | 20.67 | 82.14 | [Link](https://huggingface.co/timm/efficientformer_l7.snap_dist_in1k) | snap_dist_in1k | +| DenseNet121 | (224,224,3) | 74.320 | 74.422 | 6.37 | 8.04 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.densenet121.html) | | +| DenseNet161 | (224,224,3) | 77.200 | 77.142 | 16.85 | 28.86 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.densenet161.html) | | +| DenseNet169 | (224,224,3) | 75.554 | 75.568 | 7.62 | 14.28 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.densenet169.html) | | +| DenseNet201 | (224,224,3) | 76.742 | 76.882 | 9.82 | 20.21 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.densenet201.html) | | +| FlexiViT_Small | (240,240,3) | 82.150 | 82.536 | 11.05 | 22.06 | [Link](https://huggingface.co/timm/flexivit_small.1200ep_in1k) | 1200ep_in1k | +| FlexiViT_Base | (240,240,3) | 84.484 | 84.664 | 41.30 | 86.59 | [Link](https://huggingface.co/timm/flexivit_base.1200ep_in1k) | 1200ep_in1k | +| FlexiViT_Large | (240,240,3) | 85.540 | 85.658 | 143.89 | 304.36 | [Link](https://huggingface.co/timm/flexivit_large.1200ep_in1k) | 1200ep_in1k | +| Inception_V3 | (299,299,3) | 77.246 | 77.278 | 11.51 | 23.82 | [Link](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.inception_v3.html) | | +| LeViT_Conv_128 | (224,224,3) | 77.282 | 78.488 | 0.88 | 9.19 | [Link](https://huggingface.co/timm/levit_conv_128.fb_dist_in1k) | fb_dist_in1k | +| LeViT_Conv_128S | (224,224,3) | 75.204 | 76.574 | 0.65 | 7.76 | [Link](https://huggingface.co/timm/levit_conv_128s.fb_dist_in1k) | fb_dist_in1k | +| LeViT_Conv_192 | (224,224,3) | 79.138 | 79.876 | 1.38 | 10.92 | [Link](https://huggingface.co/timm/levit_conv_192.fb_dist_in1k) | fb_dist_in1k | +| LeViT_Conv_256 | (224,224,3) | 80.960 | 81.538 | 2.34 | 18.86 | [Link](https://huggingface.co/timm/levit_conv_256.fb_dist_in1k) | fb_dist_in1k | +| LeViT_Conv_384 | (224,224,3) | 82.170 | 82.582 | 4.83 | 39.08 | [Link](https://huggingface.co/timm/levit_conv_384.fb_dist_in1k) | fb_dist_in1k | +| MNASNet1_0 | (224,224,3) | 72.806 | 73.416 | 0.65 | 4.36 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.mnasnet1_0.html) | | +| MobileNet_V2 | (224,224,3) | 71.728 | 72.142 | 0.64 | 3.49 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.mobilenet_v2.html) | IMAGENET1K_V2 | +| RegNet_X_400MF | (224,224,3) | 72.530 | 72.908 | 0.84 | 5.48 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_400mf.html) | IMAGENET1K_V1 | +| RegNet_X_400MF | (224,224,3) | 74.182 | 74.860 | 0.84 | 5.48 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_400mf.html) | IMAGENET1K_V2 | +| RegNet_X_800MF | (224,224,3) | 74.972 | 75.210 | 1.62 | 7.24 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_800mf.html) | IMAGENET1K_V1 | +| RegNet_X_800MF | (224,224,3) | 77.056 | 77.496 | 1.62 | 7.24 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_800mf.html) | IMAGENET1K_V2 | +| RegNet_X_1_6GF | (224,224,3) | 76.900 | 77.084 | 3.24 | 9.17 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_1_6gf.html) | IMAGENET1K_V1 | +| RegNet_X_1_6GF | (224,224,3) | 79.254 | 79.676 | 3.24 | 9.17 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_1_6gf.html) | IMAGENET1K_V2 | +| RegNet_X_3_2GF | (224,224,3) | 78.142 | 78.342 | 6.40 | 15.27 | [Link](https://docs.pytorch.org//vision/2.0/models/generated/torchvision.models.regnet_x_3_2gf.html) | IMAGENET1K_V1 | +| RegNet_X_3_2GF | (224,224,3) | 80.888 | 81.194 | 6.40 | 15.27 | [Link](https://docs.pytorch.org//vision/2.0/models/generated/torchvision.models.regnet_x_3_2gf.html) | IMAGENET1K_V2 | +| RegNet_X_8GF | (224,224,3) | 79.338 | 79.372 | 16.05 | 39.53 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_8gf.html) | IMAGENET1K_V1 | +| RegNet_X_8GF | (224,224,3) | 81.386 | 81.692 | 16.05 | 39.53 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_8gf.html) | IMAGENET1K_V2 | +| RegNet_X_16GF | (224,224,3) | 79.932 | 80.092 | 31.99 | 54.22 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_16gf.html) | IMAGENET1K_V1 | +| RegNet_X_16GF | (224,224,3) | 82.434 | 82.712 | 31.99 | 54.22 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_16gf.html) | IMAGENET1K_V2 | +| RegNet_X_32GF | (224,224,3) | 80.550 | 80.592 | 63.63 | 107.73 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_32gf.html) | IMAGENET1K_V1 | +| RegNet_X_32GF | (224,224,3) | 82.856 | 83.022 | 63.63 | 107.73 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_x_32gf.html) | IMAGENET1K_V2 | +| RegNet_Y_400MF | (224,224,3) | 73.690 | 74.004 | 0.82 | 4.33 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.regnet_y_400mf.html) | IMAGENET1K_V1 | +| RegNet_Y_400MF | (224,224,3) | 75.312 | 75.802 | 0.82 | 4.33 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.regnet_y_400mf.html) | IMAGENET1K_V2 | +| RegNet_Y_800MF | (224,224,3) | 76.100 | 76.396 | 1.70 | 6.42 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.regnet_y_800mf.html) | IMAGENET1K_V1 | +| RegNet_Y_800MF | (224,224,3) | 78.448 | 78.890 | 1.70 | 6.42 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.regnet_y_800mf.html) | IMAGENET1K_V2 | +| RegNet_Y_1_6GF | (224,224,3) | 77.370 | 77.926 | 3.27 | 11.18 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_1_6gf.html) | IMAGENET1K_V1 | +| RegNet_Y_1_6GF | (224,224,3) | 80.490 | 80.882 | 3.27 | 11.18 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_1_6gf.html) | IMAGENET1K_V2 | +| RegNet_Y_3_2GF | (224,224,3) | 78.744 | 78.962 | 6.41 | 19.40 | [Link](https://docs.pytorch.org//vision/2.0/models/generated/torchvision.models.regnet_y_3_2gf.html) | IMAGENET1K_V1 | +| RegNet_Y_3_2GF | (224,224,3) | 81.454 | 82.018 | 6.41 | 19.40 | [Link](https://docs.pytorch.org//vision/2.0/models/generated/torchvision.models.regnet_y_3_2gf.html) | IMAGENET1K_V2 | +| RegNet_Y_8GF | (224,224,3) | 79.874 | 80.052 | 17.05 | 39.34 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_8gf.html) | IMAGENET1K_V1 | +| RegNet_Y_8GF | (224,224,3) | 82.546 | 82.824 | 17.05 | 39.34 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_8gf.html) | IMAGENET1K_V2 | +| RegNet_Y_16GF | (224,224,3) | 80.326 | 80.424 | 31.95 | 83.53 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_16gf.html) | IMAGENET1K_V1 | +| RegNet_Y_16GF | (224,224,3) | 82.470 | 82.862 | 31.95 | 83.53 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_16gf.html) | IMAGENET1K_V2 | +| RegNet_Y_32GF | (224,224,3) | 80.700 | 80.834 | 64.72 | 144.97 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_32gf.html) | IMAGENET1K_V1 | +| RegNet_Y_32GF | (224,224,3) | 82.890 | 83.362 | 64.72 | 144.97 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.regnet_y_32gf.html) | IMAGENET1K_V2 | +| ResNet18 | (224,224,3) | 69.558 | 69.778 | 3.64 | 11.68 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet18.html) | | +| ResNet34 | (224,224,3) | 73.166 | 73.304 | 7.35 | 21.79 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet34.html) | | +| ResNet50 | (224,224,3) | 75.980 | 76.116 | 8.23 | 25.53 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet50.html) | IMAGENET1K_V1 | +| ResNet50 | (224,224,3) | 80.574 | 80.852 | 8.23 | 25.53 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet50.html) | IMAGENET1K_V2 | +| ResNet101 | (224,224,3) | 77.076 | 77.350 | 15.69 | 44.50 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet101.html) | IMAGENET1K_V1 | +| ResNet101 | (224,224,3) | 81.534 | 81.914 | 15.69 | 44.50 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet101.html) | IMAGENET1K_V2 | +| ResNet152 | (224,224,3) | 78.044 | 78.306 | 23.14 | 60.12 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet152.html) | IMAGENET1K_V1 | +| ResNet152 | (224,224,3) | 81.952 | 82.272 | 23.14 | 60.12 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnet152.html) | IMAGENET1K_V2 | +| ResNeXt50_32X4D | (224,224,3) | 77.568 | 77.634 | 8.53 | 24.99 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnext50_32x4d.html) | IMAGENET1K_V1 | +| ResNeXt50_32X4D | (224,224,3) | 80.896 | 81.212 | 8.53 | 24.99 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnext50_32x4d.html) | IMAGENET1K_V2 | +| ResNeXt101_32X8D | (224,224,3) | 79.234 | 79.290 | 32.97 | 88.69 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnext101_32x8d.html) | IMAGENET1K_V1 | +| ResNeXt101_32X8D | (224,224,3) | 82.598 | 82.784 | 32.97 | 88.69 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnext101_32x8d.html) | IMAGENET1K_V2 | +| ResNeXt101_64X4D | (224,224,3) | 82.980 | 83.234 | 31.06 | 83.35 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.resnext101_64x4d.html) | | +| ShuffleNet_V2_X1_0 | (224,224,3) | 68.734 | 69.312 | 0.30 | 2.27 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.shufflenet_v2_x1_0.html) | | +| ShuffleNet_V2_X1_5 | (224,224,3) | 72.458 | 72.966 | 0.60 | 3.49 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.shufflenet_v2_x1_5.html) | | +| ShuffleNet_V2_X2_0 | (224,224,3) | 75.614 | 76.224 | 1.18 | 7.38 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.shufflenet_v2_x2_0.html) | | +| Swin_B | (224,224,3) | 83.254 | 83.562 | 31.58 | 88.61 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.swin_b.html) | | +| Swin_S | (224,224,3) | 82.848 | 83.180 | 18.02 | 50.24 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.swin_s.html) | | +| Swin_T | (224,224,3) | 81.076 | 81.444 | 9.32 | 28.60 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.swin_t.html) | | +| VGG11 | (224,224,3) | 68.706 | 68.974 | 15.26 | 132.86 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg11.html) | | +| VGG11_BN | (224,224,3) | 70.074 | 70.328 | 15.26 | 132.86 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg11_bn.html) | | +| VGG13 | (224,224,3) | 69.742 | 69.888 | 22.68 | 133.05 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg13.html) | | +| VGG13_BN | (224,224,3) | 71.370 | 71.564 | 22.68 | 133.05 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg13_bn.html) | | +| VGG16 | (224,224,3) | 71.526 | 71.616 | 31.01 | 138.36 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg16.html) | | +| VGG16_BN | (224,224,3) | 73.276 | 73.406 | 31.01 | 138.36 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg16_bn.html) | | +| VGG19 | (224,224,3) | 72.284 | 72.386 | 39.34 | 143.67 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg19.html) | | +| VGG19_BN | (224,224,3) | 74.022 | 74.170 | 39.34 | 143.67 | [Link](https://docs.pytorch.org//vision/stable/models/generated/torchvision.models.vgg19_bn.html) | | +| ViT_B_16 | (224,224,3) | 81.002 | 81.040 | 35.73 | 86.57 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.vit_b_16.html) | IMAGENET1K_V1 | +| ViT_B_16 | (384,384,3) | 85.118 | 85.276 | 115.01 | 86.86 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.vit_b_16.html) | IMAGENET1K_SWAG_E2E_V1 | +| ViT_B_16 | (224,224,3) | 81.516 | 81.926 | 35.73 | 86.57 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.vit_b_16.html) | IMAGENET1K_SWAG_LINEAR_V1 | +| ViT_B_32 | (224,224,3) | 75.680 | 75.908 | 8.90 | 88.22 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.vit_b_32.html) | | +| RepViT_M1 | (224,224,3) | 77.668 | 78.510 | 1.70 | 5.46 | [Link](https://huggingface.co/timm/repvit_m1.dist_in1k) | dist_in1k | +| RepViT_M1_0 | (224,224,3) | 79.210 | 80.198 | 2.39 | 7.29 | [Link](https://huggingface.co/timm/repvit_m1_0.dist_300e_in1k) | dist_300e_in1k | +| RepViT_M1_1 | (224,224,3) | 79.838 | 80.828 | 2.87 | 8.79 | [Link](https://huggingface.co/timm/repvit_m1_1.dist_300e_in1k) | dist_300e_in1k | +| RepViT_M1_5 | (224,224,3) | 81.532 | 82.374 | 4.88 | 14.62 | [Link](https://huggingface.co/timm/repvit_m1_5.dist_300e_in1k) | dist_300e_in1k | +| RepViT_M2 | (224,224,3) | 79.882 | 80.526 | 2.77 | 8.77 | [Link](https://huggingface.co/timm/repvit_m2.dist_in1k) | dist_in1k | +| RepViT_M2_3 | (224,224,3) | 82.960 | 83.466 | 9.59 | 23.66 | [Link](https://huggingface.co/timm/repvit_m2_3.dist_300e_in1k) | dist_300e_in1k | +| RepViT_M3 | (224,224,3) | 80.896 | 81.478 | 3.85 | 10.65 | [Link](https://huggingface.co/timm/repvit_m3.dist_in1k) | dist_in1k | +| VisFormer_Small | (224,224,3) | 81.586 | 82.106 | 10.05 | 40.24 | [Link](https://huggingface.co/timm/visformer_small.in1k) | in1k | +| VisFormer_Tiny | (224,224,3) | 77.564 | 78.262 | 2.68 | 10.33 | [Link](https://huggingface.co/timm/visformer_tiny.in1k) | in1k | +| ViT_Tiny_Patch16_224 | (224,224,3) | 74.568 | 75.456 | 2.66 | 5.72 | [Link](https://huggingface.co/timm/vit_tiny_patch16_224.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Tiny_Patch16_384 | (384,384,3) | 77.312 | 78.470 | 10.37 | 5.79 | [Link](https://huggingface.co/timm/vit_tiny_patch16_384.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Small_Patch16_224 | (224,224,3) | 81.404 | 81.412 | 9.50 | 22.05 | [Link](https://huggingface.co/timm/vit_small_patch16_224.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Small_Patch16_384 | (384,384,3) | 83.560 | 83.782 | 33.00 | 22.20 | [Link](https://huggingface.co/timm/vit_small_patch16_384.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Small_Patch32_224 | (224,224,3) | 75.792 | 75.922 | 2.32 | 22.88 | [Link](https://huggingface.co/timm/vit_small_patch32_224.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Small_Patch32_384 | (384,384,3) | 80.092 | 80.452 | 7.07 | 22.92 | [Link](https://huggingface.co/timm/vit_small_patch32_384.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Base_Patch8_224 | (224,224,3) | 86.302 | 86.268 | 163.48 | 86.58 | [Link](https://huggingface.co/timm/vit_base_patch8_224.augreg2_in21k_ft_in1k) | augreg2_in21k_ft_in1k | +| ViT_Base_Patch16_224 | (224,224,3) | 85.080 | 85.112 | 35.73 | 86.57 | [Link](https://huggingface.co/timm/vit_base_patch16_224.augreg2_in21k_ft_in1k) | augreg2_in21k_ft_in1k | +| ViT_Base_Patch16_384 | (384,384,3) | 85.860 | 86.018 | 115.00 | 86.86 | [Link](https://huggingface.co/timm/vit_base_patch16_384.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Base_Patch32_224 | (224,224,3) | 80.612 | 80.698 | 8.89 | 88.22 | [Link](https://huggingface.co/timm/vit_base_patch32_224.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Base_Patch32_384 | (384,384,3) | 83.138 | 83.398 | 26.45 | 88.30 | [Link](https://huggingface.co/timm/vit_base_patch32_384.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Large_Patch16_224 | (224,224,3) | 85.880 | 85.870 | 124.71 | 304.33 | [Link](https://huggingface.co/timm/vit_large_patch16_224.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Large_Patch16_384 | (384,384,3) | 86.980 | 87.086 | 392.88 | 304.72 | [Link](https://huggingface.co/timm/vit_large_patch16_384.augreg_in21k_ft_in1k) | augreg_in21k_ft_in1k | +| ViT_Large_Patch32_384 | (384,384,3) | 81.122 | 81.512 | 91.52 | 306.63 | [Link](https://huggingface.co/timm/vit_large_patch32_384.orig_in21k_ft_in1k) | orig_in21k_ft_in1k | +| Wide_ResNet50_2 | (224,224,3) | 78.402 | 78.488 | 22.87 | 68.85 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.wide_resnet50_2.html) | IMAGENET1K_V1 | +| Wide_ResNet50_2 | (224,224,3) | 81.268 | 81.630 | 22.87 | 68.85 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.wide_resnet50_2.html) | IMAGENET1K_V2 | +| Wide_ResNet101_2 | (224,224,3) | 78.500 | 78.830 | 45.61 | 126.82 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.wide_resnet101_2.html) | IMAGENET1K_V1 | +| Wide_ResNet101_2 | (224,224,3) | 82.358 | 82.510 | 45.61 | 126.82 | [Link](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.wide_resnet101_2.html) | IMAGENET1K_V2 | +| YOLOv5n-cls | (224,224,3) | 63.538 | 63.982 | 0.43 | 2.49 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5s-cls | (224,224,3) | 70.698 | 70.854 | 1.42 | 5.45 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5m-cls | (224,224,3) | 75.348 | 75.418 | 4.03 | 12.95 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5l-cls | (224,224,3) | 77.536 | 77.528 | 8.82 | 26.54 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5x-cls | (224,224,3) | 78.338 | 78.312 | 16.46 | 48.07 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv8s-cls | (224,224,3) | 72.630 | 73.774 | 1.67 | 6.36 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8m-cls | (224,224,3) | 76.032 | 76.824 | 5.37 | 17.04 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8l-cls | (224,224,3) | 77.708 | 78.276 | 12.53 | 37.47 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8x-cls | (224,224,3) | 78.440 | 78.936 | 19.38 | 57.40 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLO11s-cls | (224,224,3) | 74.532 | 75.244 | 1.63 | 6.72 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11m-cls | (224,224,3) | 76.484 | 77.388 | 5.17 | 11.62 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11l-cls | (224,224,3) | 77.658 | 78.284 | 6.51 | 14.10 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11x-cls | (224,224,3) | 78.666 | 79.426 | 14.20 | 29.61 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO26s-cls | (224,224,3) | 75.300 | 75.910 | 1.63 | 6.72 | [Link](https://docs.ultralytics.com/models/yolo26/) | | +| YOLO26m-cls | (224,224,3) | 77.608 | 78.078 | 5.17 | 11.62 | [Link](https://docs.ultralytics.com/models/yolo26/) | | +| YOLO26l-cls | (224,224,3) | 78.614 | 79.060 | 6.51 | 14.10 | [Link](https://docs.ultralytics.com/models/yolo26/) | | +| YOLO26x-cls | (224,224,3) | 79.478 | 79.866 | 14.20 | 29.61 | [Link](https://docs.ultralytics.com/models/yolo26/) | | + +
+Image Classification (ImageNet) + +- AccTop1 values are the primary model accuracies on the + [ImageNet](https://www.image-net.org/index.php) validation set. Validation also reports + AccTop5 as the secondary metric. + +
+ +### Object Detection + +| Model | Input Size
(H,W,C) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{box}}}$
(NPU) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{box}}}$
(GPU) | FLOPs (B) | params (M) | Source | Note | +| --- | --- | --- | --- | --- | --- | --- | --- | +| YOLOv3 | (640,640,3) | 46.354 | 46.839 | 162.27 | 61.92 | [Link](https://docs.ultralytics.com/models/yolov3/) | | +| YOLOv3u | (640,640,3) | 51.214 | 51.582 | 289.23 | 103.73 | [Link](https://docs.ultralytics.com/models/yolov3/) | | +| YOLOv3-spp | (640,640,3) | 47.106 | 47.616 | 163.23 | 62.97 | [Link](https://docs.ultralytics.com/models/yolov3/) | | +| YOLOv3-sppu | (640,640,3) | 51.710 | 51.754 | 290.20 | 104.78 | [Link](https://docs.ultralytics.com/models/yolov3/) | | +| YOLOv5nu | (640,640,3) | 33.549 | 34.286 | 8.79 | 2.65 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5n6 | (1280,1280,3) | 35.021 | 35.892 | 22.10 | 3.24 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5n6u | (1280,1280,3) | 41.637 | 42.111 | 35.63 | 4.33 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5s | (640,640,3) | 36.780 | 37.540 | 18.21 | 7.23 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5su | (640,640,3) | 42.414 | 42.877 | 25.93 | 9.14 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5s6 | (1280,1280,3) | 43.878 | 44.512 | 74.33 | 12.61 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5s6u | (1280,1280,3) | 48.363 | 48.636 | 105.38 | 15.29 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5m | (640,640,3) | 44.446 | 45.238 | 52.14 | 21.17 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5mu | (640,640,3) | 48.388 | 48.910 | 67.70 | 25.09 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5m6 | (1280,1280,3) | 50.637 | 51.078 | 212.83 | 35.70 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5m6u | (1280,1280,3) | 53.259 | 53.476 | 275.41 | 41.19 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5l | (640,640,3) | 48.128 | 48.914 | 114.20 | 46.53 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5lu | (640,640,3) | 51.824 | 52.172 | 140.50 | 53.19 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5l6 | (1280,1280,3) | 52.892 | 53.376 | 466.00 | 76.73 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5l6u | (1280,1280,3) | 55.344 | 55.466 | 571.74 | 86.05 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5x | (640,640,3) | 49.986 | 50.554 | 213.04 | 86.71 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5xu | (640,640,3) | 52.792 | 53.090 | 254.38 | 97.23 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5x6 | (1280,1280,3) | 54.197 | 54.706 | 869.05 | 140.73 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5x6u | (1280,1280,3) | 56.516 | 56.488 | 1035.24 | 155.48 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv7 | (640,640,3) | 50.442 | 50.942 | 110.55 | 36.91 | [Link](https://github.com/WongKinYiu/yolov7/) | | +| YOLOv7x | (640,640,3) | 52.389 | 52.706 | 197.67 | 71.31 | [Link](https://github.com/WongKinYiu/yolov7/) | | +| YOLOv7w6 | (1280,1280,3) | 53.886 | 54.128 | 374.79 | 70.39 | [Link](https://github.com/WongKinYiu/yolov7/) | | +| YOLOv7d6 | (1280,1280,3) | 55.510 | 55.799 | 730.75 | 133.76 | [Link](https://github.com/WongKinYiu/yolov7/) | | +| YOLOv7e6 | (1280,1280,3) | 55.379 | 55.567 | 538.42 | 97.20 | [Link](https://github.com/WongKinYiu/yolov7/) | | +| YOLOv7e6e | (1280,1280,3) | 55.932 | 56.298 | 878.43 | 151.69 | [Link](https://github.com/WongKinYiu/yolov7/) | | +| YOLOv8n | (640,640,3) | 36.650 | 37.328 | 9.78 | 3.15 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8s | (640,640,3) | 44.243 | 44.918 | 30.49 | 11.16 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8m | (640,640,3) | 49.915 | 50.239 | 82.26 | 25.89 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8l | (640,640,3) | 52.485 | 52.772 | 170.26 | 43.67 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8x | (640,640,3) | 53.497 | 53.802 | 264.17 | 68.20 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| GELANs | (640,640,3) | 45.679 | 46.486 | 29.01 | 7.11 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv9s | (640,640,3) | 45.637 | 46.777 | 29.01 | 7.11 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| GELANm | (640,640,3) | 50.564 | 50.928 | 80.81 | 19.98 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv9m | (640,640,3) | 50.902 | 51.191 | 80.81 | 19.98 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| GELANc | (640,640,3) | 52.086 | 52.273 | 107.83 | 25.29 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv9c | (640,640,3) | 52.598 | 52.917 | 107.83 | 25.29 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| GELANe | (640,640,3) | 54.568 | 54.927 | 199.86 | 57.35 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv9e | (640,640,3) | 55.466 | 55.528 | 199.86 | 57.35 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv10n | (640,640,3) | 37.432 | 38.382 | 8.06 | 2.30 | [Link](https://docs.ultralytics.com/models/yolov10/) | | +| YOLOv10s | (640,640,3) | 45.386 | 46.015 | 24.10 | 7.25 | [Link](https://docs.ultralytics.com/models/yolov10/) | | +| YOLOv10m | (640,640,3) | 50.075 | 50.838 | 63.51 | 15.36 | [Link](https://docs.ultralytics.com/models/yolov10/) | | +| YOLOv10b | (640,640,3) | 51.525 | 52.096 | 97.77 | 19.07 | [Link](https://docs.ultralytics.com/models/yolov10/) | | +| YOLOv10l | (640,640,3) | 52.246 | 52.816 | 127.32 | 24.37 | [Link](https://docs.ultralytics.com/models/yolov10/) | | +| YOLOv10x | (640,640,3) | 53.427 | 53.993 | 170.13 | 29.47 | [Link](https://docs.ultralytics.com/models/yolov10/) | | +| YOLO11n | (640,640,3) | 38.567 | 39.298 | 7.76 | 2.62 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11s | (640,640,3) | 46.179 | 46.617 | 23.80 | 9.44 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11m | (640,640,3) | 50.887 | 51.310 | 72.81 | 20.09 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11l | (640,640,3) | 52.770 | 53.165 | 93.21 | 25.34 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11x | (640,640,3) | 54.136 | 54.478 | 204.31 | 56.92 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO12n | (640,640,3) | 40.172 | 40.740 | 9.27 | 2.59 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12s | (640,640,3) | 47.124 | 47.687 | 26.73 | 9.26 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12m | (640,640,3) | 51.899 | 52.297 | 77.22 | 20.17 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12l | (640,640,3) | 53.200 | 53.508 | 105.07 | 26.40 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12x | (640,640,3) | 54.758 | 55.061 | 223.27 | 59.14 | [Link](https://docs.ultralytics.com/models/yolo12/) | | + +
+Object Detection (COCO) + +- $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{box}}}$ values are for single-model single-scale on the [COCO val2017](https://cocodataset.org/) dataset. + +
+ +### Instance Segmentation + +| Model | Input Size
(H,W,C) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{mask}}}$
(NPU) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{mask}}}$
(GPU) | FLOPs (B) | params (M) | Source | Note | +| --- | --- | --- | --- | --- | --- | --- | --- | +| YOLOv5n-seg | (640,640,3) | 22.671 | 23.334 | 8.23 | 1.99 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5s-seg | (640,640,3) | 31.164 | 31.592 | 28.47 | 7.61 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5m-seg | (640,640,3) | 36.764 | 37.148 | 74.57 | 21.97 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5l-seg | (640,640,3) | 39.685 | 39.942 | 153.53 | 47.89 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv5x-seg | (640,640,3) | 41.185 | 41.318 | 273.99 | 88.77 | [Link](https://docs.ultralytics.com/models/yolov5/) | | +| YOLOv8n-seg | (640,640,3) | 29.951 | 30.464 | 13.91 | 3.40 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8s-seg | (640,640,3) | 36.534 | 36.691 | 44.86 | 11.81 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8m-seg | (640,640,3) | 40.362 | 40.596 | 114.06 | 27.27 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8l-seg | (640,640,3) | 42.316 | 42.462 | 226.27 | 45.97 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8x-seg | (640,640,3) | 43.048 | 43.155 | 351.31 | 71.80 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| GELANc-seg | (640,640,3) | 42.150 | 42.204 | 150.93 | 27.42 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv9c-seg | (640,640,3) | 42.574 | 42.655 | 152.44 | 27.45 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLOv9e-seg | (640,640,3) | 44.410 | 44.394 | 256.38 | 59.74 | [Link](https://github.com/WongKinYiu/yolov9/) | | +| YOLO11n-seg | (640,640,3) | 31.509 | 32.129 | 11.88 | 2.87 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11s-seg | (640,640,3) | 37.379 | 37.726 | 38.18 | 10.10 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11m-seg | (640,640,3) | 41.552 | 41.683 | 128.82 | 22.40 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11l-seg | (640,640,3) | 42.948 | 42.974 | 149.22 | 27.65 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11x-seg | (640,640,3) | 43.895 | 43.910 | 329.44 | 62.09 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO12n-seg | (640,640,3) | 32.516 | 32.941 | 12.94 | 2.80 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12s-seg | (640,640,3) | 38.470 | 38.787 | 39.26 | 9.76 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12m-seg | (640,640,3) | 42.270 | 42.433 | 125.74 | 21.94 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12l-seg | (640,640,3) | 43.429 | 43.456 | 154.99 | 28.76 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO12x-seg | (640,640,3) | 44.187 | 44.418 | 334.55 | 64.51 | [Link](https://docs.ultralytics.com/models/yolo12/) | | +| YOLO26m-seg | (640,640,3) | 42.961 | 43.732 | 137.37 | 23.57 | [Link](https://docs.ultralytics.com/models/yolo26/) | | +| YOLO26l-seg | (640,640,3) | 44.246 | 45.242 | 157.06 | 27.97 | [Link](https://docs.ultralytics.com/models/yolo26/) | | +| YOLO26x-seg | (640,640,3) | 46.033 | 46.717 | 346.96 | 62.82 | [Link](https://docs.ultralytics.com/models/yolo26/) | | + +
+Instance Segmentation (COCO) + +- $\underset{50\text{–}95}{\text{mAP}_{\text{val}}^{\text{mask}}}$ values are for single-model single-scale on the [COCO val2017](https://cocodataset.org/) dataset. + +
+ +### Pose Estimation + +| Model | Input Size
(H,W,C) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{pose}}}$
(NPU) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{pose}}}$
(GPU) | FLOPs (B) | params (M) | Source | Note | +| --- | --- | --- | --- | --- | --- | --- | --- | +| YOLOv8s-pose | (640,640,3) | 57.007 | 59.366 | 32.10 | 11.62 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8m-pose | (640,640,3) | 62.153 | 64.451 | 84.37 | 26.45 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8l-pose | (640,640,3) | 65.173 | 66.822 | 173.70 | 44.47 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8x-pose | (640,640,3) | 67.063 | 68.357 | 269.63 | 69.46 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLOv8x-pose-p6 | (1280,1280,3) | 69.221 | 70.758 | 1092.51 | 99.14 | [Link](https://docs.ultralytics.com/models/yolov8/) | | +| YOLO11s-pose | (640,640,3) | 55.731 | 57.846 | 25.41 | 9.90 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11m-pose | (640,640,3) | 62.136 | 64.280 | 76.25 | 20.89 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11l-pose | (640,640,3) | 63.174 | 65.414 | 96.65 | 26.14 | [Link](https://docs.ultralytics.com/models/yolo11/) | | +| YOLO11x-pose | (640,640,3) | 67.218 | 68.600 | 212.26 | 58.75 | [Link](https://docs.ultralytics.com/models/yolo11/) | | + +
+Pose Estimation (COCO) + +- $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{pose}}}$ values are for single-model single-scale on the [COCO Keypoints val2017](https://cocodataset.org/) dataset. + +
+ +### OBB + +| Model | Input Size
(H,W,C) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{obb}}}$
(NPU) | $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{obb}}}$
(GPU) | FLOPs (B) | params (M) | Source | Note | +| --- | --- | --- | --- | --- | --- | --- | --- | + +
+OBB (DOTA v1.0) + +- $\underset{\texttt{50-95}}{\texttt{mAP}_{\texttt{val}}^{\texttt{obb}}}$ is the primary metric + for single-model single-scale validation on the + [DOTA v1.0](https://docs.ultralytics.com/datasets/obb/dota-v2#) dataset. Validation also reports + rotated mAP50 as the secondary metric. + +
diff --git a/mblt_vision/__init__.py b/mblt_vision/__init__.py index e69de29..1731c14 100644 --- a/mblt_vision/__init__.py +++ b/mblt_vision/__init__.py @@ -0,0 +1,86 @@ +"""MBLT vision task exports and discovery helpers. + +The vision package keeps task subpackages as the preferred import surface while +also supporting legacy top-level model imports such as +``from mblt_model_zoo.vision import ResNet50``. +""" + +from __future__ import annotations + +from . import depth_estimation as depth_estimation +from . import face_detection as face_detection +from . import image_classification as image_classification +from . import instance_segmentation as instance_segmentation +from . import obb as obb +from . import object_detection as object_detection +from . import pose_estimation as pose_estimation +from . import semantic_segmentation as semantic_segmentation +from ._api import list_models as list_models +from ._api import list_tasks as list_tasks +from .wrapper import MBLT_Engine as MBLT_Engine + +__version__ = "0.0.0" + +_TASK_MODULES = ( + face_detection, + depth_estimation, + image_classification, + instance_segmentation, + object_detection, + obb, + pose_estimation, + semantic_segmentation, +) + +_LEGACY_MODEL_EXPORTS: dict[str, object] = {} +for _task_module in _TASK_MODULES: + for _export_name in getattr(_task_module, "__all__", ()): + if _export_name in _LEGACY_MODEL_EXPORTS: + raise RuntimeError( + f"Duplicate vision export detected for '{_export_name}'." + ) + _LEGACY_MODEL_EXPORTS[_export_name] = _task_module + +_PUBLIC_EXPORTS = [ + "MBLT_Engine", + "list_models", + "list_tasks", + "face_detection", + "depth_estimation", + "image_classification", + "instance_segmentation", + "object_detection", + "obb", + "pose_estimation", + "semantic_segmentation", +] + sorted(_LEGACY_MODEL_EXPORTS) +# Keep legacy compatibility exports synchronized with their task packages. +__all__: list[str] = _PUBLIC_EXPORTS # pyright: ignore[reportUnsupportedDunderAll] + + +def __getattr__(name: str) -> object: + """Lazily resolve legacy top-level model exports. + + Args: + name: Attribute requested from the vision package. + + Returns: + The exported model wrapper class for the requested legacy name. + + Raises: + AttributeError: If the requested name is not exported by the package. + """ + + task_module = _LEGACY_MODEL_EXPORTS.get(name) + if task_module is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + value = getattr(task_module, name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Return package attributes including lazy legacy exports.""" + + return sorted(set(globals()) | set(__all__)) diff --git a/mblt_vision/_api.py b/mblt_vision/_api.py new file mode 100644 index 0000000..42e6105 --- /dev/null +++ b/mblt_vision/_api.py @@ -0,0 +1,53 @@ +"""Public helpers for discovering available vision tasks and models.""" + +from __future__ import annotations + +import importlib +import inspect +from typing import Iterable + +from ._tasks import VISION_TASKS, normalize_vision_task +from .wrapper import MBLT_Engine + + +def list_tasks() -> list[str]: + """Lists the available vision tasks.""" + + return list(VISION_TASKS) + + +def list_models(tasks: str | Iterable[str] | None = None) -> dict[str, list[str]]: + """Lists available models for the selected vision tasks. + + Args: + tasks: Task name or names to inspect. When omitted, all tasks are used. + + Returns: + A mapping of task name to exported model class names. + + Raises: + ValueError: If an unknown task name is provided. + """ + + if tasks is None: + task_list = list(VISION_TASKS) + elif isinstance(tasks, str): + task_list = [tasks] + else: + task_list = list(tasks) + + available_models: dict[str, list[str]] = {} + for task in task_list: + module_name = normalize_vision_task(task) + module = importlib.import_module( + f".{module_name}", package=__name__.replace("._api", "") + ) + available_models[task] = sorted( + name + for name, obj in inspect.getmembers(module, inspect.isclass) + if issubclass(obj, MBLT_Engine) + and obj is not MBLT_Engine + and not getattr(obj, "_yaml_missing", False) + ) + + return available_models diff --git a/mblt_vision/_compat.py b/mblt_vision/_compat.py new file mode 100644 index 0000000..99ee148 --- /dev/null +++ b/mblt_vision/_compat.py @@ -0,0 +1,227 @@ +"""Compatibility helpers for YAML-backed vision model exports. + +This module rebuilds the legacy task package exports that were removed during +the YAML migration. The generated classes keep the familiar import paths and +constructor shape while delegating model loading to ``MBLT_Engine``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Callable, Iterable, Sequence, TypeAlias, cast + +from ._model_paths import uses_shifted_compat_model_path_layout +from .wrapper import CoreMode, MBLT_Engine + +_MODEL_DIR = Path(__file__).parent / "models" + + +class CompatMBLTEngine(MBLT_Engine): + """Typed base for dynamically generated legacy compatibility wrappers.""" + + def __init__( + self, + local_path: str | None = None, + model_type: str = "DEFAULT", + infer_mode: CoreMode = "global8", + product: str = "aries", + dev_no: int = 0, + target_cores: Sequence[str] | None = None, + target_clusters: Sequence[int] | None = None, + mxq_path: str | None = None, + onnx_path: str | None = None, + framework: str | None = None, + model_path: str | None = None, + ) -> None: + """Type-only constructor matching the generated legacy wrappers.""" + del ( + local_path, + model_type, + infer_mode, + product, + dev_no, + target_cores, + target_clusters, + model_path, + mxq_path, + onnx_path, + framework, + ) + raise NotImplementedError + + +CompatMBLTEngineClass: TypeAlias = type[CompatMBLTEngine] + + +def _normalize_name(value: str) -> str: + """Returns an alphanumeric-only, case-insensitive identifier.""" + + return "".join(char for char in value.lower() if char.isalnum()) + + +def _resolve_yaml_name(class_name: str) -> str: + """Resolves the YAML config stem associated with a legacy class name. + + Args: + class_name: Exported legacy class name. + + Returns: + The YAML filename stem without the ``.yaml`` suffix. + + Raises: + ValueError: If no unique matching YAML file can be determined. + """ + + yaml_stems = [path.stem for path in _MODEL_DIR.glob("*.yaml")] + + exact_matches = [stem for stem in yaml_stems if stem.lower() == class_name.lower()] + if len(exact_matches) == 1: + return exact_matches[0] + + normalized_name = _normalize_name(class_name) + normalized_matches = [ + stem for stem in yaml_stems if _normalize_name(stem) == normalized_name + ] + if len(normalized_matches) == 1: + return normalized_matches[0] + + if not normalized_matches: + raise ValueError( + f"Could not find a YAML config for legacy class '{class_name}'." + ) + + raise ValueError( + f"Found multiple YAML configs for legacy class '{class_name}': {sorted(normalized_matches)}." + ) + + +def _build_init(yaml_name: str) -> Callable[..., None]: + """Builds a legacy-compatible ``__init__`` implementation.""" + + def __init__( + self, + local_path: str | None = None, + model_type: str = "DEFAULT", + infer_mode: CoreMode = "global8", + product: str = "aries", + dev_no: int = 0, + target_cores: Sequence[str] | None = None, + target_clusters: Sequence[int] | None = None, + mxq_path: str | None = None, + onnx_path: str | None = None, + framework: str | None = None, + model_path: str | None = None, + ) -> None: + """Initializes a YAML-backed compatibility wrapper. + + Args: + local_path: Deprecated legacy MXQ path alias. Prefer ``model_path`` + for generic MXQ or ONNX loading, or ``mxq_path`` for an + explicit MXQ-only override. + model_type: YAML config variant to load. + infer_mode: Execution mode forwarded to ``MBLT_Engine``. + product: Legacy product/board value forwarded as ``target_device``. + dev_no: Accelerator device number. + target_cores: Optional core selection for single-core mode. + target_clusters: Optional cluster selection for multi/global modes. + mxq_path: Optional explicit MXQ path alias. + onnx_path: Optional explicit ONNX path. + framework: Execution framework, either ``"mxq"`` or ``"onnx"``. When + omitted, ``model_path`` suffix is used first, then MXQ is the fallback. + model_path: Optional explicit local model path for MXQ or ONNX. + """ + + if uses_shifted_compat_model_path_layout( + model_path, mxq_path, onnx_path, framework + ): + model_path, mxq_path, onnx_path, framework = ( + mxq_path, + onnx_path, + framework, + cast(str | None, model_path), + ) + MBLT_Engine.__init__( + self, + model_cls=yaml_name, + model_type=model_type, + model_path=model_path or "", + mxq_path=mxq_path or local_path or "", + onnx_path=onnx_path or "", + dev_no=dev_no, + core_mode=infer_mode, + target_cores=list(target_cores) if target_cores is not None else None, + target_clusters=list(target_clusters) + if target_clusters is not None + else None, + target_device=product, + framework=framework, + ) + + return __init__ + + +def _build_missing_init(class_name: str, reason: str) -> Callable[..., None]: + """Builds an ``__init__`` that fails with a clear compatibility message.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Raises an informative error for removed YAML-backed models.""" + + del args, kwargs + raise ValueError( + f"Legacy vision model '{class_name}' is not available in the YAML model registry: {reason}" + ) + + return __init__ + + +def create_model_class(class_name: str, module_name: str) -> CompatMBLTEngineClass: + """Creates a legacy model class that delegates to ``MBLT_Engine``. + + Args: + class_name: Class name to expose from the task module. + module_name: Module path that should own the generated class. + + Returns: + A dynamically generated ``MBLT_Engine`` subclass. + """ + + class_doc = f"Compatibility wrapper for the legacy ``{class_name}`` vision model." + try: + yaml_name = _resolve_yaml_name(class_name) + except ValueError as exc: + return type( + class_name, + (CompatMBLTEngine,), + { + "__doc__": class_doc, + "__init__": _build_missing_init(class_name, str(exc)), + "__module__": module_name, + "_yaml_missing": True, + }, + ) + + return type( + class_name, + (CompatMBLTEngine,), + { + "__doc__": class_doc, + "__init__": _build_init(yaml_name), + "__module__": module_name, + "_yaml_name": yaml_name, + }, + ) + + +def export_model_classes( + namespace: dict[str, Any], class_names: Iterable[str], module_name: str +) -> None: + """Populates a module namespace with generated model classes. + + Args: + namespace: Target module globals. + class_names: Legacy class names to generate. + module_name: Import path for generated classes. + """ + + for class_name in class_names: + namespace[class_name] = create_model_class(class_name, module_name) diff --git a/mblt_vision/_model_paths.py b/mblt_vision/_model_paths.py new file mode 100644 index 0000000..a4eb2f6 --- /dev/null +++ b/mblt_vision/_model_paths.py @@ -0,0 +1,137 @@ +"""Framework and local-artifact path resolution for vision engines.""" + +from __future__ import annotations + +from pathlib import Path + +SUPPORTED_FRAMEWORKS = {"mxq", "onnx"} + + +def framework_from_model_path(model_path: str) -> str | None: + """Infer the runtime framework from a local model path suffix.""" + + suffix = Path(model_path).suffix.lower() + if suffix == ".mxq": + return "mxq" + if suffix == ".onnx": + return "onnx" + return None + + +def uses_shifted_engine_model_path_layout( + model_path: object, + mxq_path: object, + dev_no: object, + core_mode: object, + target_cores: object, + postprocess_kwargs: object, + framework: object, + onnx_providers: object, +) -> bool: + """Return whether engine arguments use the model-path-first layout. + + Public constructor layouts have used the third positional argument for + either ``mxq_path`` or ``model_path``. An ONNX suffix identifies the + model-path-first layout because it changes runtime routing. For MXQ, + remapping is needed only when later values have the types produced by a + one-slot positional shift; a path by itself behaves identically as the + ``mxq_path`` alias. + """ + + if not isinstance(mxq_path, str): + return False + inferred_framework = framework_from_model_path(mxq_path) + if inferred_framework not in SUPPORTED_FRAMEWORKS: + return False + if ( + isinstance(model_path, str) + and model_path + and model_path.lower() not in SUPPORTED_FRAMEWORKS + ): + return False + return ( + isinstance(dev_no, str) + or isinstance(core_mode, int) + or isinstance(target_cores, str) + or (postprocess_kwargs is not None and not isinstance(postprocess_kwargs, dict)) + or isinstance(framework, dict) + or isinstance(onnx_providers, str) + or (model_path is not None and not isinstance(model_path, str)) + or (isinstance(model_path, str) and model_path.lower() in SUPPORTED_FRAMEWORKS) + ) + + +def uses_shifted_compat_model_path_layout( + model_path: object, + mxq_path: object, + onnx_path: object, + framework: object, +) -> bool: + """Return whether generated-wrapper arguments use a shifted model-path tail.""" + + if not isinstance(mxq_path, str): + return False + inferred_framework = framework_from_model_path(mxq_path) + if inferred_framework not in SUPPORTED_FRAMEWORKS: + return False + if ( + isinstance(model_path, str) + and model_path + and model_path.lower() not in SUPPORTED_FRAMEWORKS + ): + return False + if inferred_framework == "onnx": + return True + return ( + (isinstance(model_path, str) and model_path.lower() in SUPPORTED_FRAMEWORKS) + or ( + isinstance(onnx_path, str) and framework_from_model_path(onnx_path) == "mxq" + ) + or ( + isinstance(framework, str) + and framework_from_model_path(framework) is not None + ) + ) + + +def resolve_framework(framework: str | None, model_path: str = "") -> str: + """Resolve the execution framework from explicit input and model path.""" + + normalized_framework = framework.lower() if framework is not None else None + if ( + normalized_framework is not None + and normalized_framework not in SUPPORTED_FRAMEWORKS + ): + raise ValueError( + f"Unsupported framework: {framework}. Must be one of {sorted(SUPPORTED_FRAMEWORKS)}." + ) + inferred_framework = framework_from_model_path(model_path) if model_path else None + if ( + normalized_framework + and inferred_framework + and normalized_framework != inferred_framework + ): + raise ValueError( + f"Framework `{normalized_framework}` conflicts with model path `{model_path}`. " + f"Use framework `{inferred_framework}` or remove the explicit framework." + ) + return inferred_framework or normalized_framework or "mxq" + + +def split_model_paths( + *, framework: str, model_path: str = "", mxq_path: str = "", onnx_path: str = "" +) -> tuple[str, str]: + """Resolve generic and framework-specific local model path arguments.""" + + resolved_mxq_path = mxq_path + resolved_onnx_path = onnx_path + if not model_path: + return resolved_mxq_path, resolved_onnx_path + inferred_framework = framework_from_model_path(model_path) + if inferred_framework == "mxq": + resolved_mxq_path = resolved_mxq_path or model_path + elif inferred_framework == "onnx" or framework == "onnx": + resolved_onnx_path = resolved_onnx_path or model_path + else: + resolved_mxq_path = resolved_mxq_path or model_path + return resolved_mxq_path, resolved_onnx_path diff --git a/mblt_vision/_tasks.py b/mblt_vision/_tasks.py new file mode 100644 index 0000000..df9a462 --- /dev/null +++ b/mblt_vision/_tasks.py @@ -0,0 +1,30 @@ +"""Canonical task names used by the standalone Vision processing API.""" + +from __future__ import annotations + +from collections.abc import Iterable + +VISION_TASKS: tuple[str, ...] = ( + "image_classification", + "depth_estimation", + "object_detection", + "instance_segmentation", + "semantic_segmentation", + "obb", + "pose_estimation", + "face_detection", +) + + +def normalize_vision_task(task: str, *, supported: Iterable[str] | None = None) -> str: + """Normalize a Vision task name and validate it against supported tasks.""" + + if not isinstance(task, str): + raise TypeError(f"Vision task must be a string, got {type(task).__name__}.") + normalized = task.lower() + supported_tasks = tuple(VISION_TASKS if supported is None else supported) + if normalized not in supported_tasks: + raise ValueError( + f"Unsupported Vision task {task!r}; expected one of {sorted(supported_tasks)}." + ) + return normalized diff --git a/mblt_vision/benchmark/__init__.py b/mblt_vision/benchmark/__init__.py new file mode 100644 index 0000000..b7039b2 --- /dev/null +++ b/mblt_vision/benchmark/__init__.py @@ -0,0 +1 @@ +"""Reusable benchmark reporting and command support for Mobilint Vision.""" diff --git a/mblt_vision/benchmark/argparse_utils.py b/mblt_vision/benchmark/argparse_utils.py new file mode 100644 index 0000000..bfa9bf6 --- /dev/null +++ b/mblt_vision/benchmark/argparse_utils.py @@ -0,0 +1,106 @@ +"""Shared argparse validators for benchmark scripts.""" + +from __future__ import annotations + +import argparse + + +def parse_positive_int(raw: str) -> int: + """Parses a positive integer for argparse. + + Args: + raw: Raw command-line value. + + Returns: + Parsed positive integer. + + Raises: + argparse.ArgumentTypeError: If the value is not a positive integer. + """ + try: + value = int(raw) + except (TypeError, ValueError) as e: + raise argparse.ArgumentTypeError("expected a positive integer") from e + if value <= 0: + raise argparse.ArgumentTypeError("expected a positive integer") + return value + + +def parse_positive_int_optional(raw: str | None) -> int | None: + """Parses an optional positive integer for argparse. + + Args: + raw: Raw command-line value or ``None``. + + Returns: + Parsed positive integer, or ``None`` for empty input. + + Raises: + argparse.ArgumentTypeError: If the value is not empty and not a positive integer. + """ + if raw is None or raw == "": + return None + return parse_positive_int(raw) + + +def parse_range_arg(raw: str) -> tuple[int, int, int]: + """Parses ``start:end:step`` or ``start,end,step`` positive integer ranges. + + Args: + raw: Raw command-line value. + + Returns: + Parsed ``(start, end, step)`` tuple. + + Raises: + argparse.ArgumentTypeError: If the value is malformed. + """ + sep = ":" if ":" in raw else ("," if "," in raw else None) + if sep is None: + raise argparse.ArgumentTypeError( + "expected format 'start:end:step' or 'start,end,step'" + ) + parts = [part.strip() for part in raw.split(sep)] + if len(parts) != 3: + raise argparse.ArgumentTypeError( + "expected exactly 3 integers: 'start:end:step' or 'start,end,step'" + ) + try: + start, end, step = (int(part) for part in parts) + except ValueError as e: + raise argparse.ArgumentTypeError("range values must be integers") from e + if start <= 0 or end <= 0 or step <= 0: + raise argparse.ArgumentTypeError("range values must be positive integers") + if start > end: + raise argparse.ArgumentTypeError("range start must be <= end") + return start, end, step + + +def parse_int_csv( + raw: str, *, unique_sorted: bool = True, allow_empty: bool = False +) -> list[int]: + """Parses comma-separated positive integers. + + Args: + raw: Raw command-line value. + unique_sorted: Whether to sort and de-duplicate parsed values. + allow_empty: Whether an empty input should return an empty list. + + Returns: + Parsed positive integers. + + Raises: + argparse.ArgumentTypeError: If the value is malformed. + """ + parts = [item.strip() for item in str(raw).split(",") if item.strip()] + if not parts: + if allow_empty: + return [] + raise argparse.ArgumentTypeError("expected at least one integer") + try: + values = [int(item) for item in parts] + except ValueError as e: + raise argparse.ArgumentTypeError("all values must be integers") from e + if any(value <= 0 for value in values): + raise argparse.ArgumentTypeError("all values must be positive integers") + return sorted(set(values)) if unique_sorted else values diff --git a/mblt_vision/benchmark/chart_utils.py b/mblt_vision/benchmark/chart_utils.py new file mode 100644 index 0000000..24a70ce --- /dev/null +++ b/mblt_vision/benchmark/chart_utils.py @@ -0,0 +1,146 @@ +import os +import re +from pathlib import Path +from typing import Optional, Sequence + +import matplotlib + +if "MPLBACKEND" not in os.environ: + matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + + +def sanitize_text(text: str) -> str: + cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", text).strip("-") + return cleaned or "unnamed" + + +def source_prefix(sources: list[Path], *, use_stem: bool) -> str: + parts: list[str] = [] + for source in sources: + raw = (source.stem if use_stem else source.name) or str(source) + parts.append(sanitize_text(raw)) + return "_".join(parts) + + +def default_charts_dir( + script_dir: Path, + sources: list[Path], + *, + use_stem: bool, +) -> Path: + return script_dir / "results" / "charts" / source_prefix(sources, use_stem=use_stem) + + +def source_labels(sources: list[Path], *, use_stem: bool) -> list[str]: + labels = [ + ((source.stem if use_stem else source.name) or str(source)) + for source in sources + ] + counts: dict[str, int] = {} + for label in labels: + counts[label] = counts.get(label, 0) + 1 + + seen: dict[str, int] = {} + out: list[str] = [] + for idx, label in enumerate(labels): + if counts[label] == 1: + out.append(label) + continue + seen[label] = seen.get(label, 0) + 1 + out.append(f"{label} [{seen[label]}/{counts[label]}]") + + if len(set(out)) != len(out): + out = [f"{label}#{idx + 1}" for idx, label in enumerate(labels)] + return out + + +def plot_grouped_scalar_barh( + *, + models: list[str], + group_labels: list[str], + grouped_values: list[dict[str, Optional[float]]], + x_label: str, + y_label: str, + title: str, + output_path: Path, + fig_width: float = 14.0, +) -> None: + if not models: + return + + y = np.arange(len(models), dtype=float) + group_height = 0.82 + bar_h = group_height / max(len(grouped_values), 1) + start = -group_height / 2 + bar_h / 2 + fig_h = max(5.0, 0.45 * len(models) + 2.0) + fig, ax = plt.subplots(figsize=(fig_width, fig_h)) + cmap = plt.get_cmap("tab10") + + for idx, (label, source_values) in enumerate(zip(group_labels, grouped_values)): + x_vals = [] + y_vals = [] + for i, model in enumerate(models): + value = source_values.get(model) + if value is None: + continue + x_vals.append(float(value)) + y_vals.append(y[i] + start + idx * bar_h) + if x_vals: + ax.barh( + y_vals, + x_vals, + height=bar_h * 0.95, + label=label, + color=cmap(idx % 10), + ) + + ax.set_yticks(y) + ax.set_yticklabels(models) + ax.invert_yaxis() + ax.set_xlabel(x_label) + ax.set_ylabel(y_label) + ax.set_title(title) + ax.grid(axis="x", linestyle="--", alpha=0.3) + handles, labels = ax.get_legend_handles_labels() + if handles: + ax.legend(loc="best") + plt.tight_layout() + fig.savefig(output_path, dpi=220) + plt.close(fig) + + +def plot_simple_barh( + *, + labels: Sequence[str], + values: Sequence[float], + x_label: str, + title: str, + output_path: Path, + fig_width: float = 12.0, +) -> None: + """Plots a single horizontal bar chart. + + Args: + labels: Y-axis labels. + values: Numeric values aligned with labels. + x_label: X-axis label. + title: Chart title. + output_path: Destination PNG path. + fig_width: Figure width. + """ + if not labels: + return + fig, ax = plt.subplots(figsize=(fig_width, max(4.0, 0.45 * len(labels) + 2.0))) + y = list(range(len(labels))) + ax.barh(y, [float(value) for value in values]) + ax.set_yticks(y) + ax.set_yticklabels(list(labels)) + ax.invert_yaxis() + ax.set_xlabel(x_label) + ax.set_title(title) + ax.grid(axis="x", linestyle="--", alpha=0.3) + plt.tight_layout() + fig.savefig(output_path, dpi=220) + plt.close(fig) diff --git a/mblt_vision/benchmark/io_utils.py b/mblt_vision/benchmark/io_utils.py new file mode 100644 index 0000000..8d59f83 --- /dev/null +++ b/mblt_vision/benchmark/io_utils.py @@ -0,0 +1,64 @@ +"""Shared file I/O helpers for benchmark scripts.""" + +from __future__ import annotations + +import csv +import json +from pathlib import Path +from typing import Any, Mapping, Sequence + +JsonObject = dict[str, Any] +CsvRow = dict[str, Any] + + +def safe_filename(text: str, *, replace_slash_only: bool = False) -> str: + """Returns a filename-safe representation of text. + + Args: + text: Source text to convert. + replace_slash_only: Whether to preserve the legacy behavior of replacing only forward slashes. + + Returns: + Sanitized text suitable for benchmark output filenames. + """ + if replace_slash_only: + return text.replace("/", "__") + + cleaned = ( + text.replace("/", "__").replace("\\", "__").replace(":", "_").replace(" ", "_") + ) + return cleaned or "unnamed" + + +def write_json(path: Path, payload: Any) -> None: + """Writes an indented UTF-8 JSON file after creating its parent directory. + + Args: + path: Destination path. + payload: JSON-serializable payload. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + + +def write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None: + """Writes dictionaries to CSV using the union of row keys as the field order. + + Args: + path: Destination path. + rows: Row mappings to write. Empty rows are ignored to preserve existing script behavior. + """ + if not rows: + return + path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = list(rows[0].keys()) + for row in rows[1:]: + for key in row: + if key not in fieldnames: + fieldnames.append(key) + with path.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for row in rows: + writer.writerow({key: row.get(key, "") for key in fieldnames}) diff --git a/mblt_vision/benchmark/summary_utils.py b/mblt_vision/benchmark/summary_utils.py new file mode 100644 index 0000000..443ba35 --- /dev/null +++ b/mblt_vision/benchmark/summary_utils.py @@ -0,0 +1,770 @@ +"""Markdown summary helpers for benchmark outputs.""" + +from __future__ import annotations + +import json +import re +import subprocess +from collections.abc import Mapping, Sequence +from datetime import datetime +from pathlib import Path +from typing import Any + +HOST_PC_INFO_FILENAME = "host_pc_info.json" + +_HOST_INFO_SECTION_ORDER = ("CPU", "Motherboard", "DRAM", "NPU") + +_HOST_INFO_SECTION_ALIASES = { + "cpu": "CPU", + "processor": "CPU", + "motherboard": "Motherboard", + "mainboard": "Motherboard", + "baseboard": "Motherboard", + "board": "Motherboard", + "dram": "DRAM", + "memory": "DRAM", + "ram": "DRAM", + "dimm": "DRAM", + "npu": "NPU", + "npus": "NPU", + "neural": "NPU", + "accelerator": "NPU", +} + +_NPU_ARRAY_KEY_RE = re.compile(r"(?:^|\.)npus\[(\d+)\]\.(.+)$") + + +_PLOT_TITLES_BY_NAME = { + "rtf.png": "Real-Time Factor", + "inverse_rtf.png": "Inverse Real-Time Factor", + "sec_per_j.png": "Seconds Per Joule", + "j_per_sec.png": "Joules Per Audio Second", + "wer.png": "Word Error Rate", + "cer.png": "Character Error Rate", + "p95_latency_s.png": "P95 Latency", + "throughput_samples_per_s.png": "Throughput", + "decode_tokens_per_s.png": "Decode Tokens Per Second", + "prefill_tps.png": "Prefill Tokens Per Second", + "measure_prefill_tps.png": "Prefill Tokens Per Second", + "llm_prefill_tps.png": "Prefill Tokens Per Second", + "measure_llm_prefill_tps.png": "Prefill Tokens Per Second", + "prefill_tps_per_w.png": "Prefill TPS/W", + "measure_prefill_tps_per_w.png": "Prefill TPS/W", + "llm_prefill_tps_per_w.png": "Prefill TPS/W", + "measure_llm_prefill_tps_per_w.png": "Prefill TPS/W", + "decode_tps.png": "Decode Tokens Per Second", + "measure_decode_tps.png": "Decode Tokens Per Second", + "llm_decode_tps.png": "Decode Tokens Per Second", + "measure_llm_decode_tps.png": "Decode Tokens Per Second", + "decode_tps_per_w.png": "Decode TPS/W", + "measure_decode_tps_per_w.png": "Decode TPS/W", + "llm_decode_tps_per_w.png": "Decode TPS/W", + "measure_llm_decode_tps_per_w.png": "Decode TPS/W", + "avg_power_w.png": "Power", + "measure_avg_power_w.png": "Power", + "avg_temperature_c.png": "Temperature", + "measure_avg_temperature_c.png": "Temperature", + "avg_utilization_pct.png": "Utilization", + "measure_avg_utilization_pct.png": "Utilization", + "avg_memory_used_mb.png": "Memory Used Megabytes", + "measure_avg_memory_used_mb.png": "Memory Used Megabytes", + "total_energy_j.png": "Total Energy", + "measure_total_energy_j.png": "Total Energy", + "vision_fps.png": "Vision FPS", + "vision_encode_ms.png": "Vision Encode ms", + "vision_img_per_j.png": "Vision Images Per Joule", +} + +_PLOT_NAME_ORDER = {name: idx for idx, name in enumerate(_PLOT_TITLES_BY_NAME)} + + +def collect_host_pc_info( + results_dir: Path | str, *, filename: str = HOST_PC_INFO_FILENAME +) -> Path: + """Run ``mblt-tracker collect`` and save its JSON output. + + The benchmark should not fail just because host information collection is unavailable. Failures are therefore + converted into a JSON payload that can still be rendered in the summary. + + Args: + results_dir: Benchmark output directory. + filename: JSON filename to write under ``results_dir``. + + Returns: + Path to the written JSON file. + """ + out_dir = Path(results_dir) + out_dir.mkdir(parents=True, exist_ok=True) + output_path = out_dir / filename + payload: Any + try: + proc = subprocess.run( + ["mblt-tracker", "collect"], + check=False, + capture_output=True, + encoding="utf-8", + timeout=30, + ) + if proc.returncode == 0: + try: + payload = json.loads(proc.stdout) + except json.JSONDecodeError: + payload = { + "status": "error", + "message": "mblt-tracker collect did not return valid JSON.", + "stdout": proc.stdout, + "stderr": proc.stderr, + "returncode": proc.returncode, + } + else: + payload = { + "status": "error", + "message": "mblt-tracker collect failed.", + "stdout": proc.stdout, + "stderr": proc.stderr, + "returncode": proc.returncode, + } + except FileNotFoundError: + payload = {"status": "error", "message": "mblt-tracker CLI was not found."} + except subprocess.TimeoutExpired as e: + payload = { + "status": "error", + "message": "mblt-tracker collect timed out.", + "stdout": e.stdout, + "stderr": e.stderr, + "timeout_s": e.timeout, + } + + payload = _with_collection_metadata(payload) + with output_path.open("w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + print(f"Saved Host PC Info: {output_path.name}") + return output_path + + +def write_summary_markdown( + path: Path | str, + *, + title: str, + host_info_path: Path | str | None, + table_markdown_path: Path | str | None, + plot_paths: Sequence[Path | str], + plot_tables: Mapping[str, str] | None = None, + host_info_paths: Mapping[str, Path | str] | None = None, +) -> None: + """Write a benchmark summary Markdown file with host info, plots, and table. + + Args: + path: Summary Markdown destination. + title: Document title. + host_info_path: JSON file written by :func:`collect_host_pc_info`. + table_markdown_path: Existing Markdown table to include. + plot_paths: PNG files to embed in the summary. + plot_tables: Optional Markdown tables keyed by plot PNG filename. When provided, matching tables are rendered + directly below each plot and the bottom combined table is omitted. + host_info_paths: Optional source-labeled host info JSON files. When provided, host info is rendered per source. + """ + summary_path = Path(path) + summary_path.parent.mkdir(parents=True, exist_ok=True) + host_info_lines = ( + _host_infos_markdown( + {label: Path(src_path) for label, src_path in host_info_paths.items()} + ) + if host_info_paths + else _host_info_markdown(Path(host_info_path) if host_info_path else None) + ) + lines = [f"# {title}\n\n"] + lines.extend(_device_energy_note_markdown()) + lines.extend( + _plots_markdown( + summary_path.parent, [Path(p) for p in plot_paths], plot_tables=plot_tables + ) + ) + if not plot_tables: + lines.extend( + _table_markdown(Path(table_markdown_path) if table_markdown_path else None) + ) + lines.extend(host_info_lines) + summary_path.write_text("".join(lines), encoding="utf-8") + + +def _device_energy_note_markdown() -> list[str]: + """Return a reusable note describing trace-integrated energy limitations.""" + + return [ + "## Device energy note\n\n", + "Energy and energy-efficiency metrics are computed from mblt-tracker power traces using " + "trapezoidal integration. At least two valid power samples are required, so measurements shorter " + "than the tracker sampling interval may leave energy fields empty.\n\n", + ] + + +def read_csv_rows(path: Path | str) -> list[dict[str, str]]: + """Read CSV rows if the file exists. + + Args: + path: CSV path to read. + + Returns: + CSV rows as dictionaries, or an empty list when the file does not exist. + """ + csv_path = Path(path) + if not csv_path.is_file(): + return [] + import csv + + with csv_path.open("r", encoding="utf-8", newline="") as f: + return list(csv.DictReader(f)) + + +def markdown_table(headers: Sequence[str], rows: Sequence[Sequence[Any]]) -> str: + """Build a compact Markdown table with right-aligned metric columns. + + Args: + headers: Table headers. + rows: Table row values. + + Returns: + Markdown table text, or an empty string for empty rows. + """ + if not rows: + return "" + lines = [ + "| " + " | ".join(_escape_markdown(header) for header in headers) + " |\n", + "| " + " | ".join(["---"] + ["---:" for _ in headers[1:]]) + " |\n", + ] + for row in rows: + lines.append( + "| " + " | ".join(_format_summary_cell(value) for value in row) + " |\n" + ) + return "".join(lines) + + +def scalar_plot_table( + rows: Sequence[Mapping[str, Any]], *, value_key: str, unit_header: str +) -> str: + """Build a model/value table for one scalar plot. + + Args: + rows: Rows containing a ``model`` key and the requested scalar key. + value_key: Key containing the scalar value. + unit_header: Header for the value column. + + Returns: + Markdown table text. + """ + return markdown_table( + ["Model", unit_header], [[row.get("model"), row.get(value_key)] for row in rows] + ) + + +def token_sweep_plot_table( + models: Sequence[str], + metrics_by_model: Mapping[str, Any], + *, + value_key: str, +) -> str: + """Build a model/token table for one token-sweep plot. + + Args: + models: Model names to include. + metrics_by_model: Mapping from model name to metric object with token dictionaries. + value_key: Attribute name containing a ``dict[int, float]`` token metric. + + Returns: + Markdown table text. + """ + token_set: set[int] = set() + for model in models: + token_set.update(getattr(metrics_by_model[model], value_key).keys()) + tokens = sorted(token_set) + if not tokens: + return "" + table_rows = [] + for model in models: + values = getattr(metrics_by_model[model], value_key) + table_rows.append([model, *(values.get(token) for token in tokens)]) + return markdown_table( + ["Model", *(f"{token} tokens" for token in tokens)], table_rows + ) + + +def write_token_combined_markdown( + path: Path | str, + tps_rows: Sequence[Mapping[str, Any]], + device_rows: Sequence[Mapping[str, Any]], +) -> None: + """Write a wide-form token sweep Markdown table shared by transformer benchmarks. + + Args: + path: Markdown output path. + tps_rows: Long-form TPS rows from ``BenchmarkResult.iter_rows``. + device_rows: Per-model device metric rows. + """ + if not tps_rows: + return + models = sorted({str(r["model"]) for r in tps_rows}) + prefill_tokens = sorted( + { + int(r["tokens"]) + for r in tps_rows + if str(r.get("phase")) == "prefill" and _is_int_like(r.get("tokens")) + } + ) + decode_tokens = sorted( + { + int(r["tokens"]) + for r in tps_rows + if str(r.get("phase")) == "decode" and _is_int_like(r.get("tokens")) + } + ) + tps_map: dict[tuple[str, str, int], float] = {} + time_map: dict[tuple[str, str, int], float] = {} + npu_pct_map: dict[tuple[str, str, int], float] = {} + for row in tps_rows: + model = str(row["model"]) + phase = str(row["phase"]) + token = int(row["tokens"]) + tps_val = row.get("tps") + time_ms_val = row.get("time_ms") + npu_pct_val = row.get("avg_npu_token_latency_pct") + if isinstance(tps_val, (int, float)): + tps_map[(model, phase, token)] = float(tps_val) + if isinstance(time_ms_val, (int, float)): + time_map[(model, phase, token)] = float(time_ms_val) + if isinstance(npu_pct_val, (int, float)): + npu_pct_map[(model, phase, token)] = float(npu_pct_val) + + device_map = { + str(r["model"]): r for r in device_rows if isinstance(r.get("model"), str) + } + device_cols = [ + "avg_power_w", + "p99_power_w", + "avg_utilization_pct", + "p99_utilization_pct", + "avg_temperature_c", + "p99_temperature_c", + "avg_memory_used_mb", + "p99_memory_used_mb", + "total_memory_mb", + "avg_memory_used_pct", + "p99_memory_used_pct", + "total_energy_j", + "prefill_tps_last", + "decode_tps_last", + "prefill_tps_per_w_last", + "decode_tps_per_w_last", + "prefill_j_per_tok_last", + "decode_j_per_tok_last", + ] + + headers = ["model"] + headers.extend([f"prefill_tps_{t}" for t in prefill_tokens]) + headers.extend([f"decode_tps_{t}" for t in decode_tokens]) + headers.extend([f"prefill_latency_ms_{t}" for t in prefill_tokens]) + headers.extend([f"decode_duration_ms_{t}" for t in decode_tokens]) + headers.extend([f"prefill_npu_latency_pct_{t}" for t in prefill_tokens]) + headers.extend([f"decode_npu_latency_pct_{t}" for t in decode_tokens]) + headers.extend(device_cols) + + rows: list[list[str]] = [] + for model in models: + values: list[str] = [model] + for token in prefill_tokens: + values.append( + _format_optional_float(tps_map.get((model, "prefill", token))) + ) + for token in decode_tokens: + values.append(_format_optional_float(tps_map.get((model, "decode", token)))) + for token in prefill_tokens: + values.append( + _format_optional_float(time_map.get((model, "prefill", token))) + ) + for token in decode_tokens: + values.append( + _format_optional_float(time_map.get((model, "decode", token))) + ) + for token in prefill_tokens: + values.append( + _format_optional_float(npu_pct_map.get((model, "prefill", token))) + ) + for token in decode_tokens: + values.append( + _format_optional_float(npu_pct_map.get((model, "decode", token))) + ) + + drow = device_map.get(model, {}) + for col in device_cols: + v = drow.get(col) if isinstance(drow, Mapping) else None + values.append(_format_optional_float(v)) + rows.append(values) + + Path(path).write_text(markdown_table(headers, rows), encoding="utf-8") + + +def existing_png_paths( + results_dir: Path | str, *, prefixes: Sequence[str] | None = None +) -> list[Path]: + """Return sorted PNG paths under a benchmark results directory. + + Args: + results_dir: Directory to scan. + prefixes: Optional filename prefixes to include. + + Returns: + Sorted PNG files matching the optional prefixes. + """ + out_dir = Path(results_dir) + paths = sorted( + out_dir.glob("*.png"), + key=lambda path: (_plot_sort_key(path), path.name), + ) + if prefixes is None: + return [path for path in paths if _plot_title(path) is not None] + prefix_tuple = tuple(prefixes) + return [path for path in paths if path.name.startswith(prefix_tuple)] + + +def _with_collection_metadata(payload: Any) -> dict[str, Any]: + collected_at = datetime.now().astimezone().isoformat(timespec="seconds") + if isinstance(payload, dict): + out = dict(payload) + out.setdefault("status", "ok") + out.setdefault("collected_at", collected_at) + return out + return {"status": "ok", "collected_at": collected_at, "data": payload} + + +def _host_info_markdown(path: Path | None) -> list[str]: + lines = ["## Host PC Info\n\n"] + if path is None or not path.is_file(): + lines.append("Host PC info is not available.\n\n") + return lines + try: + with path.open("r", encoding="utf-8") as f: + payload = json.load(f) + except (OSError, json.JSONDecodeError) as e: + lines.append(f"Failed to read `{path.name}`: {e}\n\n") + return lines + + sections = _host_info_sections(payload) + if not sections: + lines.append("Host PC info is empty.\n\n") + return lines + lines.append(f"Source: `{path.name}`\n\n") + for title, rows in sections: + if title == "NPU": + _append_npu_info_markdown(lines, rows) + else: + _append_host_info_table(lines, title, rows) + return lines + + +def _host_infos_markdown(paths_by_label: Mapping[str, Path]) -> list[str]: + """Render source-labeled host info JSON files as merged comparison tables.""" + + lines = ["## Host PC Info\n\n"] + if not paths_by_label: + lines.append("Host PC info is not available.\n\n") + return lines + + metadata_rows: list[tuple[str, dict[str, str]]] = [("Source", {}), ("Status", {})] + sections_by_label: dict[str, dict[str, list[tuple[str, str]]]] = {} + section_order: list[str] = [] + for label, path in paths_by_label.items(): + metadata_rows[0][1][label] = path.as_posix() + if not path.is_file(): + metadata_rows[1][1][label] = ( + "missing: host_pc_info.json was not found in the input folder" + ) + sections_by_label[label] = {} + continue + try: + with path.open("r", encoding="utf-8") as f: + payload = json.load(f) + except (OSError, json.JSONDecodeError) as e: + metadata_rows[1][1][label] = f"error: {e}" + sections_by_label[label] = {} + continue + + sections = _host_info_sections(payload) + if not sections: + metadata_rows[1][1][label] = "empty" + sections_by_label[label] = {} + continue + metadata_rows[1][1][label] = "ok" + label_sections: dict[str, list[tuple[str, str]]] = {} + for title, rows in sections: + label_sections[title] = rows + if title not in section_order: + section_order.append(title) + sections_by_label[label] = label_sections + + labels = list(paths_by_label) + _append_merged_host_info_table(lines, "Sources", labels, metadata_rows) + for section in section_order: + field_values: dict[str, dict[str, str]] = {} + for label in labels: + for field, value in sections_by_label.get(label, {}).get(section, []): + field_values.setdefault(field, {})[label] = value + rows = [(field, values) for field, values in field_values.items()] + _append_merged_host_info_table(lines, section, labels, rows) + return lines + + +def _append_merged_host_info_table( + lines: list[str], + title: str, + labels: Sequence[str], + rows: Sequence[tuple[str, Mapping[str, str]]], +) -> None: + """Append a source-merged host info table.""" + + lines.append(f"### {title}\n\n") + lines.append( + "| Field | " + " | ".join(_escape_markdown(label) for label in labels) + " |\n" + ) + lines.append("| --- | " + " | ".join("---" for _ in labels) + " |\n") + for field, values_by_label in rows: + values = [_escape_markdown(values_by_label.get(label, "")) for label in labels] + lines.append(f"| `{_escape_markdown(field)}` | " + " | ".join(values) + " |\n") + lines.append("\n") + + +def _plots_markdown( + base_dir: Path, + plot_paths: Sequence[Path], + *, + plot_tables: Mapping[str, str] | None = None, +) -> list[str]: + lines = ["## Plots\n\n"] + existing = [path for path in plot_paths if path.is_file()] + if not existing: + lines.append("No plot PNG files were generated.\n\n") + return lines + tables = plot_tables or {} + for path in existing: + rel = path.relative_to(base_dir) if path.is_relative_to(base_dir) else path + title = _plot_title(path) or path.stem.replace("_", " ").title() + lines.append(f"### {title}\n\n") + lines.append(f"![{title}]({rel.as_posix()})\n\n") + table = tables.get(path.name) + if table: + lines.append(table) + if not lines[-1].endswith("\n"): + lines.append("\n") + lines.append("\n") + return lines + + +def _plot_title(path: Path) -> str | None: + title = _PLOT_TITLES_BY_NAME.get(path.name) + if title is not None: + return title + + stem = path.stem + if stem.startswith("rtf_beams"): + return "Real-Time Factor" + if stem.startswith("wer_beams"): + return "Word Error Rate" + if stem.startswith("cer_beams"): + return "Character Error Rate" + return None + + +def _plot_sort_key(path: Path) -> int: + if path.name in _PLOT_NAME_ORDER: + return _PLOT_NAME_ORDER[path.name] + stem = path.stem + if stem.startswith("rtf_beams"): + return len(_PLOT_NAME_ORDER) + if stem.startswith("wer_beams"): + return len(_PLOT_NAME_ORDER) + 1 + if stem.startswith("cer_beams"): + return len(_PLOT_NAME_ORDER) + 2 + return len(_PLOT_NAME_ORDER) + 100 + + +def _host_info_sections(payload: Any) -> list[tuple[str, list[tuple[str, str]]]]: + rows = _flatten_json(payload) + if not rows: + return [] + + grouped: dict[str, list[tuple[str, str]]] = { + section: [] for section in _HOST_INFO_SECTION_ORDER + } + general_rows: list[tuple[str, str]] = [] + for key, value in rows: + section = _host_info_section_for_key(key) + if section is None: + general_rows.append((key, value)) + else: + grouped[section].append((key, value)) + + sections: list[tuple[str, list[tuple[str, str]]]] = [] + if general_rows: + sections.append(("General", general_rows)) + sections.extend( + (section, grouped[section]) + for section in _HOST_INFO_SECTION_ORDER + if grouped[section] + ) + return sections + + +def _append_host_info_table( + lines: list[str], + title: str, + rows: Sequence[tuple[str, str]], + *, + heading_level: int = 3, +) -> None: + """Append one host info section as a Markdown table.""" + lines.append(f"{'#' * heading_level} {title}\n\n") + _append_field_value_table(lines, rows) + + +def _append_npu_info_markdown( + lines: list[str], rows: Sequence[tuple[str, str]], *, heading_level: int = 3 +) -> None: + """Append NPU host info with ``npus`` array entries grouped by index.""" + common_rows: list[tuple[str, str]] = [] + indexed_rows: dict[int, list[tuple[str, str]]] = {} + for key, value in rows: + npu_key = _split_npu_array_key(key) + if npu_key is None: + common_rows.append((key, value)) + continue + index, field = npu_key + indexed_rows.setdefault(index, []).append((field, value)) + + if not indexed_rows: + _append_host_info_table(lines, "NPU", rows, heading_level=heading_level) + return + + lines.append(f"{'#' * heading_level} NPU\n\n") + if common_rows: + lines.append(f"{'#' * (heading_level + 1)} General\n\n") + _append_field_value_table(lines, common_rows) + for index in sorted(indexed_rows): + lines.append(f"{'#' * (heading_level + 1)} NPU {index}\n\n") + _append_field_value_table(lines, indexed_rows[index]) + + +def _append_field_value_table( + lines: list[str], rows: Sequence[tuple[str, str]] +) -> None: + """Append field/value rows as a Markdown table.""" + lines.append("| Field | Value |\n") + lines.append("| --- | --- |\n") + for key, value in rows: + lines.append(f"| `{_escape_markdown(key)}` | {_escape_markdown(value)} |\n") + lines.append("\n") + + +def _host_info_section_for_key(key: str) -> str | None: + if _split_npu_array_key(key) is not None: + return "NPU" + normalized = key.replace("_", ".").replace("-", ".").replace(" ", ".").lower() + parts = [ + part + for part in normalized.replace("[", ".").replace("]", ".").split(".") + if part + ] + for part in parts: + section = _HOST_INFO_SECTION_ALIASES.get(part) + if section is not None: + return section + return None + + +def _split_npu_array_key(key: str) -> tuple[int, str] | None: + """Return the NPU array index and field for flattened ``npus`` keys.""" + match = _NPU_ARRAY_KEY_RE.search(key) + if match is None: + return None + return int(match.group(1)), match.group(2) + + +def _host_info_section_title(value: str) -> str: + upper_names = {"cpu": "CPU", "dram": "DRAM", "gpu": "GPU", "npu": "NPU", "os": "OS"} + normalized = value.replace("_", " ").replace("-", " ").strip() + lowered = normalized.lower() + if lowered in upper_names: + return upper_names[lowered] + return " ".join( + upper_names.get(part.lower(), part.capitalize()) for part in normalized.split() + ) + + +def _table_markdown(path: Path | None) -> list[str]: + lines = ["## Results Table\n\n"] + if path is None or not path.is_file(): + lines.append("Results table is not available.\n") + return lines + lines.append(path.read_text(encoding="utf-8")) + if not lines[-1].endswith("\n"): + lines.append("\n") + return lines + + +def _flatten_json(value: Any, *, prefix: str = "") -> list[tuple[str, str]]: + if isinstance(value, Mapping): + rows: list[tuple[str, str]] = [] + for key, child in value.items(): + child_key = f"{prefix}.{key}" if prefix else str(key) + rows.extend(_flatten_json(child, prefix=child_key)) + return rows + if isinstance(value, list): + if all(not isinstance(item, (Mapping, list)) for item in value): + return [(prefix, ", ".join(_scalar_to_text(item) for item in value))] + rows = [] + for idx, child in enumerate(value): + rows.extend(_flatten_json(child, prefix=f"{prefix}[{idx}]")) + return rows + return [(prefix, _scalar_to_text(value))] if prefix else [] + + +def _scalar_to_text(value: Any) -> str: + if value is None: + return "" + if isinstance(value, (str, int, float, bool)): + return str(value) + return json.dumps(value, ensure_ascii=False) + + +def _format_summary_cell(value: Any) -> str: + """Format one benchmark summary table value.""" + if value is None or value == "": + return "" + if isinstance(value, (int, float)): + return f"{float(value):.6f}" + if isinstance(value, str): + try: + return f"{float(value):.6f}" + except ValueError: + return _escape_markdown(value) + return _escape_markdown(str(value)) + + +def _format_optional_float(value: Any) -> str: + """Format a numeric value for compact benchmark tables.""" + return f"{float(value):.6f}" if isinstance(value, (int, float)) else "" + + +def _is_int_like(value: Any) -> bool: + """Return whether a value can be losslessly parsed as an integer token count.""" + if isinstance(value, int): + return True + if isinstance(value, str): + try: + int(value) + except ValueError: + return False + return True + return False + + +def _escape_markdown(value: str) -> str: + return value.replace("|", "\\|").replace("\n", "
") diff --git a/mblt_vision/cli/__init__.py b/mblt_vision/cli/__init__.py new file mode 100644 index 0000000..eb8e256 --- /dev/null +++ b/mblt_vision/cli/__init__.py @@ -0,0 +1,5 @@ +"""Command-line interface for standalone Mobilint Vision workflows.""" + +from .main import build_parser, main + +__all__ = ["build_parser", "main"] diff --git a/mblt_vision/cli/__main__.py b/mblt_vision/cli/__main__.py new file mode 100644 index 0000000..13522c2 --- /dev/null +++ b/mblt_vision/cli/__main__.py @@ -0,0 +1,5 @@ +"""Run the standalone Vision CLI with ``python -m mblt_vision.cli``.""" + +from .main import main + +raise SystemExit(main()) diff --git a/mblt_vision/cli/_vision.py b/mblt_vision/cli/_vision.py new file mode 100644 index 0000000..096a2ef --- /dev/null +++ b/mblt_vision/cli/_vision.py @@ -0,0 +1,337 @@ +"""Shared helpers for vision CLI commands.""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path +from typing import Any + +import torch +from mblt_vision._tasks import normalize_vision_task + +DEFAULT_OUTPUT_DIR = Path("runs") / "vision" + + +def parse_unit_interval(value: str) -> float: + """Parse a floating-point value strictly between zero and one.""" + + try: + parsed = float(value) + except ValueError as exc: + raise argparse.ArgumentTypeError( + "expected a number in the open interval (0, 1)" + ) from exc + if not 0 < parsed < 1: + raise argparse.ArgumentTypeError( + f"expected a number in the open interval (0, 1), got {value}" + ) + return parsed + + +def parse_target_cores(value: str | None) -> list[str] | None: + """Parses a semicolon-separated target core list.""" + + if value is None: + return None + cores = [item.strip() for item in value.split(";") if item.strip()] + return cores or None + + +def parse_target_clusters(value: str | None) -> list[int] | None: + """Parses a semicolon-separated target cluster list.""" + + if value is None: + return None + try: + clusters = [int(item.strip()) for item in value.split(";") if item.strip()] + except ValueError as exc: + raise argparse.ArgumentTypeError( + "target clusters must be semicolon-separated integers" + ) from exc + return clusters or None + + +def add_common_vision_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments shared by all vision inference commands.""" + + parser.add_argument("--source", required=True, help="Path to the source image.") + parser.add_argument( + "--model", + required=True, + help="Vision model name, for example `resnet50` or `yolo11m`.", + ) + parser.add_argument( + "--output", + "--save-path", + dest="output", + help="Path to save the plotted result image.", + ) + parser.add_argument( + "--framework", + default=None, + choices=["mxq", "onnx"], + help="Inference framework to use. When omitted, `--model-path` suffix is used first, then `mxq`.", + ) + parser.add_argument( + "--model-path", + dest="model_path", + default="", + help="Optional generic local model path for MXQ or ONNX inference.", + ) + parser.add_argument( + "--mxq-path", + dest="mxq_path", + default="", + help="Optional local MXQ model path. Preserved as a compatibility alias.", + ) + parser.add_argument( + "--onnx-path", + dest="onnx_path", + default="", + help="Optional local ONNX model path.", + ) + parser.add_argument( + "--model-type", + default="DEFAULT", + help="Model variant from the YAML configuration.", + ) + parser.add_argument( + "--core-mode", + default=None, + choices=["single", "multi", "global4", "global8"], + help="NPU core execution mode. Defaults to global8 on Aries and single on Regulus.", + ) + parser.add_argument("--dev-no", type=int, default=0, help="NPU device number.") + parser.add_argument( + "--target-device", + default="aries-rb", + choices=["aries-rb", "regulus-ra", "regulus-rb"], + help="NPU board target. Determines the backend implementation.", + ) + parser.add_argument( + "--target-cores", + type=parse_target_cores, + help="Optional semicolon-separated core list for single-core mode, for example `0:0;0:1`.", + ) + parser.add_argument( + "--target-clusters", + type=parse_target_clusters, + help="Optional semicolon-separated cluster list for multi/global modes, for example `0;1`.", + ) + + +def add_threshold_args( + parser: argparse.ArgumentParser, + *, + conf_default: float | None = 0.25, + iou_default: float | None = None, +) -> None: + """Adds postprocess threshold arguments for dense vision tasks.""" + + parser.add_argument( + "--conf-thres", + type=parse_unit_interval, + default=conf_default, + help="Confidence threshold.", + ) + parser.add_argument( + "--iou-thres", + type=parse_unit_interval, + default=iou_default, + help="IoU threshold.", + ) + + +def parse_bool(value: str) -> bool: + """Parses a case-insensitive boolean CLI value. + + Args: + value: Boolean text to parse. + + Returns: + Parsed boolean value. + + Raises: + argparse.ArgumentTypeError: If the value is not a supported boolean spelling. + """ + + normalized = value.strip().lower() + if normalized in {"true", "1", "yes", "on"}: + return True + if normalized in {"false", "0", "no", "off"}: + return False + raise argparse.ArgumentTypeError("expected a boolean value: true or false") + + +def add_e2e_arg(parser: argparse.ArgumentParser) -> None: + """Adds an optional YOLO end-to-end postprocessing mode override. + + Leaving the option unset preserves the model configuration's default. + """ + + parser.add_argument( + "--e2e", + nargs="?", + const=True, + type=parse_bool, + default=None, + help="Enable or disable YOLO end-to-end postprocessing (true/false). Bare `--e2e` means true.", + ) + + +def add_vision_parser( + subparsers: argparse._SubParsersAction[argparse.ArgumentParser], + *, + command: str, + help_text: str, + handler: Any, + description: str | None = None, + epilog: str | None = None, +) -> argparse.ArgumentParser: + """Creates a vision command parser with common arguments.""" + + parser = subparsers.add_parser( + command, + help=help_text, + description=description, + epilog=epilog, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.set_defaults(_handler=handler) + add_common_vision_args(parser) + return parser + + +def build_default_output_path(command: str, source: str, model: str) -> str: + """Builds the default path used for plotted vision command output.""" + + source_path = Path(source) + suffix = source_path.suffix or ".jpg" + return str(DEFAULT_OUTPUT_DIR / command / f"{source_path.stem}_{model}{suffix}") + + +def resolve_output_path( + output: str | None, command: str, source: str, model: str +) -> str: + """Returns an absolute result image path and ensures its parent exists.""" + + save_path = Path( + output or build_default_output_path(command, source, model) + ).expanduser() + save_path.parent.mkdir(parents=True, exist_ok=True) + return str(save_path.resolve()) + + +def require_source_file(source: str) -> None: + """Exits with a clear message when the source image is unavailable.""" + + source_path = Path(source).expanduser() + if not source_path.is_file(): + raise SystemExit(f"Source image not found: {source}") + + +def create_vision_engine(args: argparse.Namespace) -> Any: + """Creates a vision engine from shared CLI model options. + + Args: + args: Parsed command options containing common vision model arguments. + + Returns: + Initialized vision inference engine. + + Raises: + SystemExit: If the vision runtime dependencies are unavailable. + """ + + try: + from mblt_vision import MBLT_Engine + from mblt_vision.wrapper import normalize_core_mode + except ImportError as exc: + print(f"Missing dependencies for vision CLI: {exc}", file=sys.stderr) + raise SystemExit(2) from exc + + postprocess_kwargs: dict[str, Any] = {} + if getattr(args, "e2e", None) is not None: + postprocess_kwargs["e2e"] = args.e2e + + return MBLT_Engine( + model_cls=args.model, + model_type=args.model_type, + framework=args.framework, + model_path=args.model_path, + mxq_path=args.mxq_path, + onnx_path=args.onnx_path, + dev_no=args.dev_no, + target_device=args.target_device, + core_mode=normalize_core_mode( + args.core_mode + or ("single" if args.target_device.startswith("regulus-") else "global8") + ), + target_cores=args.target_cores, + target_clusters=args.target_clusters, + postprocess_kwargs=postprocess_kwargs, + ) + + +def run_vision_inference( + args: argparse.Namespace, + *, + command: str, +) -> Any: + """Runs a complete vision inference pipeline for a CLI command.""" + + require_source_file(args.source) + model = create_vision_engine(args) + try: + actual_task = normalize_vision_task(model.post_cfg.get("task", "")) + plot_kwargs: dict[str, Any] = {} + if actual_task == "image_classification": + plot_kwargs["topk"] = args.topk + elif actual_task in { + "object_detection", + "face_detection", + "instance_segmentation", + "pose_estimation", + "obb", + }: + model.set_postprocess_thresholds( + conf_thres=args.conf_thres, iou_thres=args.iou_thres + ) + + postprocess_kwargs: dict[str, Any] = {} + if actual_task == "semantic_segmentation": + input_img, metadata = model.preprocess_with_metadata(args.source) + postprocess_kwargs["img0_shape"] = metadata["img0_shape"] + postprocess_kwargs["ratio_pad"] = metadata.get("ratio_pad") + else: + input_img = model.preprocess(args.source) + output = model(input_img) + if not getattr(getattr(model, "postprocessor", None), "e2e", True): + if args.output: + raise SystemExit( + "`--output` is unavailable with `--e2e false`; use `--raw-output` instead." + ) + + raw_output = model.postprocessor(output) + raw_output_path = getattr(args, "raw_output", None) + if raw_output_path: + output_path = Path(raw_output_path).expanduser() + output_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(raw_output, output_path) + print(f"Saved raw postprocess output to {output_path}") + else: + print( + "Generated raw export-style postprocess output. Use `--raw-output` to save it." + ) + return raw_output + + result = model.postprocess(output, **postprocess_kwargs) + + save_path = resolve_output_path(args.output, command, args.source, args.model) + result.plot(source_path=args.source, save_path=save_path, **plot_kwargs) + print(f"Saved result to {os.path.relpath(save_path)}") + return result + finally: + model.dispose() diff --git a/mblt_vision/cli/compile.py b/mblt_vision/cli/compile.py new file mode 100644 index 0000000..4726243 --- /dev/null +++ b/mblt_vision/cli/compile.py @@ -0,0 +1,117 @@ +"""Vision compilation CLI command.""" + +from __future__ import annotations + +import argparse +import sys + + +def _run_compile(args: argparse.Namespace) -> int: + """Compile a vision model from parsed CLI arguments. + + Args: + args: Parsed compilation arguments. + + Returns: + Successful process status. + """ + + try: + from mblt_vision.compile import compile_vision_model + + output_path = compile_vision_model( + model_cls=args.model_cls, + target_device=args.target_device, + model_type=args.model_type, + model_path=args.model_path, + data_path=args.data_path, + subset_path=args.subset_path, + calib_data_path=args.calib_data_path, + save_path=args.save_path, + subset_size=args.subset_size, + seed=args.seed, + percentile=args.percentile, + topk_ratio=args.topk_ratio, + ) + except ImportError as exc: + print(str(exc), file=sys.stderr) + raise SystemExit(2) from exc + print(f"Compiled MXQ model to {output_path}") + return 0 + + +def add_compile_parser( + subparsers: argparse._SubParsersAction[argparse.ArgumentParser], +) -> argparse.ArgumentParser: + """Register the vision compilation command. + + Args: + subparsers: Main CLI subparser collection. + + Returns: + Registered compile command parser. + """ + + parser = subparsers.add_parser( + "compile", help="Compile a configured vision ONNX model to MXQ." + ) + parser.set_defaults(_handler=_run_compile) + parser.add_argument( + "--model-cls", + required=True, + help="Vision model name, for example `alexnet` or `yolo11m`.", + ) + parser.add_argument( + "--target-device", + required=True, + choices=["aries-rb", "regulus-ra", "regulus-rb"], + help="Required NPU board target for the compiled MXQ artifact.", + ) + parser.add_argument( + "--model-type", + default="DEFAULT", + help="Model variant from the YAML configuration.", + ) + parser.add_argument( + "--model-path", + "--onnx-path", + dest="model_path", + help="Preferred local ONNX file. Missing files fall back to the configured hosted artifact.", + ) + data_group = parser.add_mutually_exclusive_group() + data_group.add_argument( + "--data-path", + help="Original organized dataset root; organize, sample, and preprocess it as needed.", + ) + data_group.add_argument( + "--subset-path", + help="Already-sampled image subset; skip original dataset preparation and sampling.", + ) + data_group.add_argument( + "--calib-data-path", + "--calib-data-dir", + dest="calib_data_path", + help="Ready directory of preprocessed .npy tensors; pass it directly to qbcompiler.", + ) + parser.add_argument( + "--save-path", + help="Output MXQ path. Defaults to the ONNX stem under ~/.mblt_model_zoo.", + ) + parser.add_argument( + "--subset-size", + type=int, + help="Per-category ImageNet/WiderFace count or total count for other datasets.", + ) + parser.add_argument( + "--seed", type=int, default=0, help="Deterministic calibration subset seed." + ) + parser.add_argument( + "--percentile", type=float, help="Quantization percentile override." + ) + parser.add_argument( + "--topk-ratio", type=float, help="Quantization top-k ratio override." + ) + return parser + + +__all__ = ["add_compile_parser"] diff --git a/mblt_vision/cli/main.py b/mblt_vision/cli/main.py new file mode 100644 index 0000000..814b7d0 --- /dev/null +++ b/mblt_vision/cli/main.py @@ -0,0 +1,35 @@ +"""Standalone command-line entry point for Mobilint Vision.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence + +from .compile import add_compile_parser +from .predict import add_predict_parser +from .val import add_val_parser + + +def build_parser() -> argparse.ArgumentParser: + """Build the standalone Vision command parser.""" + + parser = argparse.ArgumentParser( + prog="mblt-vision", + description="Run, validate, and compile Mobilint Vision models.", + ) + subparsers = parser.add_subparsers(help="mblt-vision commands") + add_predict_parser(subparsers) + add_val_parser(subparsers) + add_compile_parser(subparsers) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Run the standalone Vision command-line interface.""" + + parser = build_parser() + args = parser.parse_args(argv) + if hasattr(args, "_handler"): + return int(args._handler(args)) + parser.print_help() + return 1 diff --git a/mblt_vision/cli/predict.py b/mblt_vision/cli/predict.py new file mode 100644 index 0000000..739900d --- /dev/null +++ b/mblt_vision/cli/predict.py @@ -0,0 +1,64 @@ +"""Vision prediction CLI command.""" + +from __future__ import annotations + +import argparse + +from ._vision import ( + add_e2e_arg, + add_threshold_args, + add_vision_parser, + run_vision_inference, +) + + +def _cmd_predict(args: argparse.Namespace) -> int: + """Runs vision inference on a source image.""" + + run_vision_inference(args, command="predict") + return 0 + + +def add_predict_parser( + subparsers: argparse._SubParsersAction[argparse.ArgumentParser], +) -> None: + """Registers the unified vision prediction CLI command.""" + + parser = add_vision_parser( + subparsers, + command="predict", + help_text=( + "Run vision inference for classification, depth estimation, detection, instance or semantic " + "segmentation, OBB, pose, and face detection." + ), + description=( + "Run a configured Vision model on one image. The selected model determines the task, " + "preprocessing, postprocessing, and output visualization automatically." + ), + epilog="""Supported tasks: + image classification, depth estimation, object and face detection, instance + and semantic segmentation, oriented bounding boxes (OBB), and pose estimation. + +The command downloads the default MXQ artifact when no local model path is supplied, +then writes a plotted result under runs/vision/predict/ by default. Use --output to +choose the image destination. Use --framework onnx with --model-path or --onnx-path +for ONNX Runtime inference; MXQ is the default framework. + +Examples: + mblt-vision predict --source image.jpg --model resnet50 --topk 3 + mblt-vision predict --source image.jpg --model yolo11m --conf-thres 0.4 --output result.jpg + mblt-vision predict --source image.jpg --model yolo11m --framework onnx + mblt-vision predict --source image.jpg --model yolo11m-pose --target-device regulus-ra --core-mode single + +For export-style YOLO output, use --e2e false and optionally save it with --raw-output.""", + handler=_cmd_predict, + ) + parser.add_argument( + "--topk", type=int, default=5, help="Number of classification labels to show." + ) + parser.add_argument( + "--raw-output", + help="Path to save raw export-style output with `--e2e false`.", + ) + add_threshold_args(parser, conf_default=0.25, iou_default=None) + add_e2e_arg(parser) diff --git a/mblt_vision/cli/val.py b/mblt_vision/cli/val.py new file mode 100644 index 0000000..ccdd66d --- /dev/null +++ b/mblt_vision/cli/val.py @@ -0,0 +1,562 @@ +"""Vision validation CLI command.""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +from mblt_vision._tasks import normalize_vision_task +from mblt_vision.benchmark.argparse_utils import parse_positive_int +from mblt_vision.datasets import get_dataset_config, get_dataset_config_for_task +from mblt_vision.utils.datasets.readiness import dataset_ready +from mblt_vision.wrapper import get_mobilint_cache_dir + +from ._vision import ( + add_e2e_arg, + add_threshold_args, + create_vision_engine, + parse_target_clusters, + parse_target_cores, +) + +DEFAULT_IMAGENET_IMAGE_SOURCE = get_dataset_config("imagenet")["download"]["images"] +DEFAULT_IMAGENET_XML_SOURCE = get_dataset_config("imagenet")["download"]["annotations"] +DEFAULT_COCO_IMAGE_SOURCE = get_dataset_config("coco")["download"]["images"] +DEFAULT_COCO_ANNOTATION_SOURCE = get_dataset_config("coco")["download"]["annotations"] +DEFAULT_WIDERFACE_IMAGE_SOURCE = get_dataset_config("widerface")["download"]["images"] +DEFAULT_WIDERFACE_ANNOTATION_SOURCE = get_dataset_config("widerface")["download"][ + "annotations" +] +DEFAULT_DOTAV1_SOURCE = get_dataset_config("dotav1")["download"]["url"] +DEFAULT_NYU_DEPTH_SOURCE = get_dataset_config("nyu-depth")["download"]["url"] +DEFAULT_ADE20K_SOURCE = get_dataset_config("ade20k")["download"]["url"] +CITYSCAPES_DOWNLOAD_CONFIG = get_dataset_config("cityscapes")["download"] +CITYSCAPES_IMAGE_ARCHIVE = CITYSCAPES_DOWNLOAD_CONFIG["images_archive"] +CITYSCAPES_ANNOTATION_ARCHIVE = CITYSCAPES_DOWNLOAD_CONFIG["annotations_archive"] + + +def _candidate_search_roots(data_path: str) -> list[Path]: + """Returns directories to inspect for existing raw dataset sources.""" + + root = Path(data_path).expanduser() + candidates = [root, root.parent, Path.cwd()] + ordered: list[Path] = [] + seen: set[Path] = set() + for candidate in candidates: + resolved = candidate.resolve() + if resolved in seen: + continue + seen.add(resolved) + ordered.append(resolved) + return ordered + + +def _find_existing_source(data_path: str, candidate_names: list[str]) -> str | None: + """Finds a nearby raw archive or extracted dataset directory.""" + + for root in _candidate_search_roots(data_path): + for name in candidate_names: + candidate = root / name + if candidate.exists(): + return str(candidate) + return None + + +def _normalize_coco_annotation_source(annotation_dir: str | None) -> str | None: + """Normalizes a COCO annotation source for the organizer contract. + + The COCO organizer expects either the annotation archive or the extracted + parent directory that contains an ``annotations`` subdirectory. When source + discovery finds the extracted leaf ``annotations`` directory directly, + return its parent so downstream code does not resolve ``annotations`` + twice. + """ + + if annotation_dir is None: + return None + + candidate = Path(annotation_dir).expanduser() + if candidate.is_dir() and candidate.name == "annotations": + return str(candidate.parent) + return annotation_dir + + +def _resolve_imagenet_sources( + args: argparse.Namespace, data_path: str +) -> tuple[str, str]: + """Resolves local or remote sources for ImageNet organization.""" + + image_dir = args.image_dir + xml_dir = args.xml_dir + if not args.force_organize: + image_dir = image_dir or _find_existing_source( + data_path, ["ILSVRC2012_img_val.tar", "ILSVRC2012_img_val"] + ) + xml_dir = xml_dir or _find_existing_source( + data_path, ["ILSVRC2012_bbox_val_v3.tgz", "ILSVRC2012_bbox_val_v3"] + ) + return ( + image_dir or DEFAULT_IMAGENET_IMAGE_SOURCE, + xml_dir or DEFAULT_IMAGENET_XML_SOURCE, + ) + + +def _resolve_coco_sources(args: argparse.Namespace, data_path: str) -> tuple[str, str]: + """Resolves local or remote sources for COCO organization.""" + + image_dir = args.image_dir + annotation_dir = _normalize_coco_annotation_source(args.annotation_dir) + if not args.force_organize: + image_dir = image_dir or _find_existing_source( + data_path, ["val2017.zip", "val2017"] + ) + annotation_dir = annotation_dir or _normalize_coco_annotation_source( + _find_existing_source( + data_path, + [ + "annotations_trainval2017.zip", + "annotations_trainval2017", + "annotations", + ], + ) + ) + return ( + image_dir or DEFAULT_COCO_IMAGE_SOURCE, + annotation_dir or DEFAULT_COCO_ANNOTATION_SOURCE, + ) + + +def _resolve_widerface_sources( + args: argparse.Namespace, data_path: str +) -> tuple[str, str]: + """Resolves local or remote sources for WiderFace organization.""" + + image_dir = args.image_dir + annotation_dir = args.annotation_dir + if not args.force_organize: + image_dir = image_dir or _find_existing_source( + data_path, ["WIDER_val.zip", "WIDER_val"] + ) + annotation_dir = annotation_dir or _find_existing_source( + data_path, + ["wider_face_split.zip", "wider_face_split"], + ) + return ( + image_dir or DEFAULT_WIDERFACE_IMAGE_SOURCE, + annotation_dir or DEFAULT_WIDERFACE_ANNOTATION_SOURCE, + ) + + +def _resolve_dotav1_source(args: argparse.Namespace, data_path: str) -> str: + """Resolves a local or remote source for DOTAv1 organization.""" + + dataset_path = args.annotation_dir or args.image_dir + if not args.force_organize: + dataset_path = dataset_path or _find_existing_source( + data_path, ["DOTAv1.zip", "DOTAv1"] + ) + return dataset_path or DEFAULT_DOTAV1_SOURCE + + +def _resolve_nyu_depth_source(args: argparse.Namespace, data_path: str) -> str: + """Resolve a local archive or URL for NYU Depth organization.""" + + dataset_path = args.annotation_dir or args.image_dir + if not args.force_organize: + dataset_path = dataset_path or _find_existing_source( + data_path, ["nyu-depth.zip", "nyu-depth"] + ) + return dataset_path or DEFAULT_NYU_DEPTH_SOURCE + + +def _resolve_ade20k_source(args: argparse.Namespace, data_path: str) -> str: + """Resolve a local archive, extracted directory, or URL for ADE20K organization.""" + + dataset_path = args.annotation_dir or args.image_dir + if not args.force_organize: + dataset_path = dataset_path or _find_existing_source( + data_path, + ["ADEChallengeData2016.zip", "ADEChallengeData2016"], + ) + return dataset_path or DEFAULT_ADE20K_SOURCE + + +def _resolve_cityscapes_sources( + args: argparse.Namespace, data_path: str +) -> tuple[str, str]: + """Resolve the two manually downloaded official Cityscapes archives. + + Args: + args: Parsed validation CLI arguments. + data_path: Organized Cityscapes output path used as a discovery anchor. + + Returns: + Image and annotation ZIP paths. + + Raises: + SystemExit: If either required archive cannot be found. + """ + + image_dir = args.image_dir or _find_existing_source( + data_path, [CITYSCAPES_IMAGE_ARCHIVE] + ) + annotation_dir = args.annotation_dir or _find_existing_source( + data_path, [CITYSCAPES_ANNOTATION_ARCHIVE] + ) + if image_dir is None or annotation_dir is None: + raise SystemExit( + "Cityscapes organization requires the official image and annotation ZIP archives. " + "Register at https://www.cityscapes-dataset.com/, then download them with:\n" + " csDownload -d gtFine_trainvaltest.zip leftImg8bit_trainvaltest.zip\n" + "Pass the resulting files with --image-dir and --annotation-dir, or place them near the dataset path." + ) + return image_dir, annotation_dir + + +def _default_data_path_for_task(task: str, dataset: str | None = None) -> str: + """Returns the default organized dataset path for a vision task.""" + + try: + configured_path = Path(get_dataset_config_for_task(task, dataset)["path"]) + except ValueError as exc: + raise SystemExit(f"Unsupported vision task for validation: {task}") from exc + + configured_path = configured_path.expanduser() + default_cache_root = Path.home() / ".mblt_model_zoo" + try: + relative_path = configured_path.relative_to(default_cache_root) + except ValueError: + return str(configured_path) + return str(Path(get_mobilint_cache_dir()) / relative_path) + + +def _dataset_ready(task: str, data_path: str, dataset: str | None = None) -> bool: + """Checks whether the organized dataset appears ready for validation.""" + + return dataset_ready(data_path, task, dataset) + + +def _ensure_dataset( + args: argparse.Namespace, task: str, dataset: str | None = None +) -> str: + """Organizes the dataset automatically when the expected layout is missing.""" + + task = normalize_vision_task(task) + data_path = os.path.expanduser( + args.data_path or _default_data_path_for_task(task, dataset) + ) + if _dataset_ready(task, data_path, dataset) and not args.force_organize: + print(f"Using organized dataset at {data_path}") + return data_path + + try: + from mblt_vision.utils.datasets import ( + organize_ade20k, + organize_cityscapes, + organize_coco, + organize_dotav1, + organize_imagenet, + organize_nyu_depth, + organize_widerface, + ) + except ImportError as exc: + print( + f"Missing dependencies for vision dataset organization: {exc}", + file=sys.stderr, + ) + raise SystemExit(2) from exc + + print(f"Preparing validation dataset for task `{task}` at {data_path}...") + if task == "image_classification": + image_dir, xml_dir = _resolve_imagenet_sources(args, data_path) + organize_imagenet( + image_dir=image_dir, + xml_dir=xml_dir, + output_dir=data_path, + ) + elif task in {"object_detection", "instance_segmentation", "pose_estimation"}: + image_dir, annotation_dir = _resolve_coco_sources(args, data_path) + organize_coco( + image_dir=image_dir, + annotation_dir=annotation_dir, + output_dir=data_path, + ) + elif task == "face_detection": + image_dir, annotation_dir = _resolve_widerface_sources(args, data_path) + organize_widerface( + image_dir=image_dir, + annotation_dir=annotation_dir, + output_dir=data_path, + ) + elif task == "obb": + organize_dotav1( + dataset_path=_resolve_dotav1_source(args, data_path), + output_dir=data_path, + ) + elif task == "depth_estimation": + organize_nyu_depth( + dataset_path=_resolve_nyu_depth_source(args, data_path), + output_dir=data_path, + ) + elif task == "semantic_segmentation": + if dataset == "cityscapes": + image_dir, annotation_dir = _resolve_cityscapes_sources(args, data_path) + organize_cityscapes( + image_dir=image_dir, + annotation_dir=annotation_dir, + output_dir=data_path, + ) + else: + organize_ade20k( + dataset_path=_resolve_ade20k_source(args, data_path), + output_dir=data_path, + ) + else: + raise SystemExit(f"Unsupported vision task for validation: {task}") + + if not _dataset_ready(task, data_path, dataset): + raise SystemExit( + f"Organized validation dataset at {data_path} is incomplete or does not match " + f"the expected {dataset or task} dataset." + ) + return data_path + + +def _run_validation(args: argparse.Namespace) -> float: + """Runs model validation on the dataset associated with the model task.""" + + try: + from mblt_vision.utils.evaluation import ( + eval_ade20k, + eval_cityscapes, + eval_coco_metrics, + eval_dota, + eval_imagenet_metrics, + eval_nyu_depth, + eval_widerface, + ) + except ImportError as exc: + print(f"Missing dependencies for vision CLI: {exc}", file=sys.stderr) + raise SystemExit(2) from exc + + model = create_vision_engine(args) + try: + if not getattr(getattr(model, "postprocessor", None), "e2e", True): + raise SystemExit( + "Validation requires end-to-end YOLO postprocessing. Use `--e2e true` or omit the option." + ) + + task = normalize_vision_task(model.post_cfg.get("task", "")) + dataset = model.post_cfg.get("dataset") + taxonomy = str(dataset).lower() if isinstance(dataset, str) else None + if task == "semantic_segmentation" and taxonomy not in {"ade20k", "cityscapes"}: + raise SystemExit( + f"Unsupported semantic segmentation taxonomy for validation: {taxonomy!r}. " + "Expected `ade20k` or `cityscapes`." + ) + data_path = _ensure_dataset(args, task, taxonomy) + + if task == "image_classification": + imagenet_result = eval_imagenet_metrics( + model=model, data_path=data_path, batch_size=args.batch_size + ) + print( + "Validation score " + f"(Top-1 accuracy): {imagenet_result.top1:.5f}, " + f"(Top-5 accuracy): {imagenet_result.top5:.5f}" + ) + return imagenet_result.primary_score + + if task == "depth_estimation": + depth_result = eval_nyu_depth( + model=model, data_path=data_path, batch_size=args.batch_size + ) + print( + "Validation score " + f"(delta1): {depth_result.delta1:.5f}, " + f"(abs_rel): {depth_result.abs_rel:.5f}, " + f"(rmse): {depth_result.rmse:.5f}" + ) + return depth_result.primary_score + + if task == "semantic_segmentation": + if taxonomy == "cityscapes": + semantic_result = eval_cityscapes( + model=model, data_path=data_path, batch_size=args.batch_size + ) + elif taxonomy == "ade20k": + semantic_result = eval_ade20k( + model=model, data_path=data_path, batch_size=args.batch_size + ) + else: + raise AssertionError( + f"Unexpected validated semantic taxonomy: {taxonomy!r}" + ) + print( + "Validation score " + f"(mIoU): {semantic_result.miou:.5f}, " + f"(pixel accuracy): {semantic_result.pixel_accuracy:.5f}" + ) + return semantic_result.primary_score + + if task in {"object_detection", "instance_segmentation", "pose_estimation"}: + coco_result = eval_coco_metrics( + model=model, + data_path=data_path, + batch_size=args.batch_size, + conf_thres=args.conf_thres, + iou_thres=args.iou_thres, + ) + print( + f"Validation score (mAP50-95): {coco_result.map5095:.5f}, (mAP50): {coco_result.map50:.5f}" + ) + return coco_result.primary_score + + if task == "obb": + dota_result = eval_dota( + model=model, + data_path=data_path, + batch_size=args.batch_size, + conf_thres=args.conf_thres, + iou_thres=args.iou_thres, + ) + print( + "Validation score " + f"(rotated mAP50-95): {dota_result.map5095:.5f}, " + f"(rotated mAP50): {dota_result.map50:.5f}" + ) + return dota_result.primary_score + + if task == "face_detection": + widerface_result = eval_widerface( + model=model, + data_path=data_path, + batch_size=args.batch_size, + conf_thres=args.conf_thres, + iou_thres=args.iou_thres, + ) + print( + "Validation score " + f"(Easy AP): {widerface_result.easy_ap:.5f}, " + f"(Medium AP): {widerface_result.medium_ap:.5f}, " + f"(Hard AP): {widerface_result.hard_ap:.5f}, " + f"(Mean AP): {widerface_result.mean_ap:.5f}" + ) + return widerface_result.primary_score + + raise SystemExit(f"Unsupported vision task for validation: {task}") + finally: + model.dispose() + + +def _cmd_val(args: argparse.Namespace) -> int: + """Runs vision validation on the task-appropriate benchmark dataset.""" + + _run_validation(args) + return 0 + + +def add_val_parser( + subparsers: argparse._SubParsersAction[argparse.ArgumentParser], +) -> None: + """Registers the unified vision validation CLI command.""" + + parser = subparsers.add_parser( + "val", help="Validate a vision model on its benchmark dataset." + ) + parser.set_defaults(_handler=_cmd_val) + parser.add_argument( + "--model", + required=True, + help="Vision model name, for example `resnet50` or `yolo11m`.", + ) + parser.add_argument( + "--framework", + default=None, + choices=["mxq", "onnx"], + help="Inference framework to use. When omitted, `--model-path` suffix is used first, then `mxq`.", + ) + parser.add_argument( + "--model-path", + dest="model_path", + default="", + help="Optional generic local model path for MXQ or ONNX inference.", + ) + parser.add_argument( + "--mxq-path", + dest="mxq_path", + default="", + help="Optional local MXQ model path. Preserved as a compatibility alias.", + ) + parser.add_argument( + "--onnx-path", + dest="onnx_path", + default="", + help="Optional local ONNX model path.", + ) + parser.add_argument( + "--model-type", + default="DEFAULT", + help="Model variant from the YAML configuration.", + ) + parser.add_argument( + "--core-mode", + default=None, + choices=["single", "multi", "global4", "global8"], + help="NPU core execution mode. Defaults to global8 on Aries and single on Regulus.", + ) + parser.add_argument("--dev-no", type=int, default=0, help="NPU device number.") + parser.add_argument( + "--target-device", + default="aries-rb", + choices=["aries-rb", "regulus-ra", "regulus-rb"], + help="NPU board target. Determines the backend implementation.", + ) + parser.add_argument( + "--target-cores", + type=parse_target_cores, + help="Optional semicolon-separated core list for single-core mode, for example `0:0;0:1`.", + ) + parser.add_argument( + "--target-clusters", + type=parse_target_clusters, + help="Optional semicolon-separated cluster list for multi/global modes, for example `0;1`.", + ) + parser.add_argument( + "--batch-size", + type=parse_positive_int, + default=1, + help="Positive batch size for validation.", + ) + parser.add_argument( + "--data-path", + help="Path to an already organized validation dataset. If omitted, the default cache path is used.", + ) + parser.add_argument( + "--force-organize", + "--force", + "--reload", + action="store_true", + dest="force_organize", + help="Rebuild the organized dataset even when the target directory already looks ready.", + ) + parser.add_argument( + "--image-dir", + help=( + "Local archive path or download URL for dataset images. Cityscapes requires leftImg8bit_trainvaltest.zip." + ), + ) + parser.add_argument( + "--xml-dir", + help="Local archive path or download URL for ImageNet annotations used by automatic organization.", + ) + parser.add_argument( + "--annotation-dir", + help=( + "Local archive path or download URL for dataset annotations. Cityscapes requires gtFine_trainvaltest.zip." + ), + ) + add_threshold_args(parser, conf_default=None, iou_default=None) + add_e2e_arg(parser) diff --git a/mblt_vision/compile/__init__.py b/mblt_vision/compile/__init__.py new file mode 100644 index 0000000..9c09c9d --- /dev/null +++ b/mblt_vision/compile/__init__.py @@ -0,0 +1,21 @@ +"""Vision model compilation and calibration-data preparation.""" + +from .vision import ( + compile_vision_model, + copy_calibration_subset, + ensure_calibration_dataset, + make_calibration_subset, + prepare_calibration_arrays, + resolve_quantization_values, + select_calibration_images, +) + +__all__ = [ + "compile_vision_model", + "copy_calibration_subset", + "ensure_calibration_dataset", + "make_calibration_subset", + "prepare_calibration_arrays", + "resolve_quantization_values", + "select_calibration_images", +] diff --git a/mblt_vision/compile/vision.py b/mblt_vision/compile/vision.py new file mode 100644 index 0000000..0b9c019 --- /dev/null +++ b/mblt_vision/compile/vision.py @@ -0,0 +1,1120 @@ +"""Vision model compilation and calibration-data preparation.""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import shutil +import warnings +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from tempfile import TemporaryDirectory, mkdtemp +from typing import Any + +import numpy as np +import torch +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import HfHubHTTPError +from mblt_vision.utils.datasets.readiness import ( + _path_has_symlink_component, + dataset_ready, +) + +from mblt_vision._tasks import normalize_vision_task +from mblt_vision.datasets import get_dataset_config_for_task +from mblt_npu import normalize_target_device +from mblt_vision.wrapper import ( + MOBILINT_CACHE_DIR, + MBLT_Engine, + get_mobilint_cache_dir, + resolve_model_config, +) + +DEFAULT_PERCENTILE = 0.9999 +DEFAULT_TOPK_RATIO = 0.01 +DEFAULT_SEED = 0 +DEFAULT_MODEL_DIR = Path(MOBILINT_CACHE_DIR) +SUPPORTED_TARGET_DEVICES = frozenset({"aries-rb", "regulus-ra", "regulus-rb"}) +DEFAULT_SUBSET_SIZES = { + "image_classification": 1, + "depth_estimation": 100, + "object_detection": 100, + "instance_segmentation": 100, + "semantic_segmentation": 100, + "pose_estimation": 100, + "face_detection": 1, + "obb": 100, +} +IMAGE_SUFFIXES = {".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp"} + + +def _default_model_dir() -> Path: + """Return the writable default compilation output directory lazily.""" + + if DEFAULT_MODEL_DIR != Path(MOBILINT_CACHE_DIR): + return DEFAULT_MODEL_DIR + return Path(get_mobilint_cache_dir()) + + +def _normalize_task(task: str) -> str: + """Validate and normalize a supported task name. + + Args: + task: Task value from a model postprocess configuration. + + Returns: + Canonical task name used by the dataset registry. + + Raises: + ValueError: If the task is not supported by vision compilation. + """ + + return normalize_vision_task(task, supported=DEFAULT_SUBSET_SIZES) + + +def _validate_calibration_output_dir(output_dir: str | Path) -> Path: + """Return an output path only when it and its ancestors are not symlinks.""" + + destination = Path(output_dir).expanduser() + if _path_has_symlink_component(destination): + raise ValueError( + f"Calibration output_dir must not be or contain a symlink: {destination}." + ) + return destination + + +def _dataset_ready(task: str, data_path: Path, dataset: str | None = None) -> bool: + """Return whether an organized dataset contains calibration images. + + Args: + task: Canonical vision task name. + data_path: Organized dataset root. + dataset: Optional dense dataset taxonomy. + + Returns: + Whether the expected dataset identity, metadata, and full validation split are present. + """ + + return dataset_ready(data_path, task, dataset) + + +def _find_dataset_source(data_path: Path, filename: str) -> Path | None: + """Find a manually supplied dataset source at or beside its organized root. + + Args: + data_path: Organized dataset destination. + filename: Archive filename to locate. + + Returns: + Existing archive path, or ``None`` when it is unavailable. + """ + + for directory in (data_path, data_path.parent): + candidate = directory / filename + if candidate.is_file(): + return candidate + return None + + +def _organize_dataset(task: str, data_path: Path, dataset: str | None = None) -> None: + """Organize the registry-backed dataset required by a task. + + Args: + task: Canonical vision task name. + data_path: Destination dataset root. + dataset: Optional dataset taxonomy used to disambiguate semantic datasets. + """ + + from mblt_vision.utils.datasets import ( + organize_ade20k, + organize_cityscapes, + organize_coco, + organize_dotav1, + organize_imagenet, + organize_nyu_depth, + organize_widerface, + ) + + config = get_dataset_config_for_task(task, dataset) + download = config.get("download") + if not isinstance(download, dict): + raise ValueError( + f"Dataset `{config.get('name', task)}` does not define download metadata." + ) + + data_path.parent.mkdir(parents=True, exist_ok=True) + if task == "image_classification": + organize_imagenet( + image_dir=str(download["images"]), + xml_dir=str(download["annotations"]), + output_dir=str(data_path), + ) + elif task in {"object_detection", "instance_segmentation", "pose_estimation"}: + organize_coco( + image_dir=str(download["images"]), + annotation_dir=str(download["annotations"]), + output_dir=str(data_path), + ) + elif task == "face_detection": + organize_widerface( + image_dir=str(download["images"]), + annotation_dir=str(download["annotations"]), + output_dir=str(data_path), + ) + elif task == "obb": + organize_dotav1(dataset_path=str(download["url"]), output_dir=str(data_path)) + elif task == "depth_estimation": + organize_nyu_depth(dataset_path=str(download["url"]), output_dir=str(data_path)) + elif config["name"] == "ade20k": + organize_ade20k(dataset_path=str(download["url"]), output_dir=str(data_path)) + elif config["name"] == "cityscapes": + image_archive = _find_dataset_source(data_path, str(download["images_archive"])) + annotation_archive = _find_dataset_source( + data_path, str(download["annotations_archive"]) + ) + if image_archive is None or annotation_archive is None: + raise ValueError( + "Cityscapes compilation requires leftImg8bit_trainvaltest.zip and gtFine_trainvaltest.zip " + f"at {data_path} or {data_path.parent}." + ) + organize_cityscapes( + image_dir=str(image_archive), + annotation_dir=str(annotation_archive), + output_dir=str(data_path), + ) + else: + raise ValueError( + f"Unsupported calibration dataset `{config['name']}` for task `{task}`." + ) + + +def ensure_calibration_dataset( + task: str, + data_path: str | Path | None = None, + dataset: str | None = None, +) -> Path: + """Resolve and, when needed, organize a calibration dataset. + + Args: + task: Vision task name. + data_path: Optional organized dataset root. + dataset: Optional dataset taxonomy from the model postprocess configuration. + + Returns: + Expanded organized dataset root. + """ + + normalized_task = _normalize_task(task) + config = get_dataset_config_for_task(normalized_task, dataset) + dataset_name = str(config["name"]) + resolved_path = ( + Path(data_path).expanduser() + if data_path is not None + else _resolve_default_dataset_path(config["path"]) + ) + if not _dataset_ready(normalized_task, resolved_path, dataset_name): + _organize_dataset(normalized_task, resolved_path, dataset) + if not _dataset_ready(normalized_task, resolved_path, dataset_name): + raise ValueError( + f"Organized calibration dataset at {resolved_path} is incomplete or does not match " + f"the expected {dataset_name} dataset." + ) + return resolved_path + + +def _resolve_default_dataset_path(configured_path: str | Path) -> Path: + """Translate registry cache paths to the active writable artifact cache.""" + + path = Path(configured_path).expanduser() + default_cache_root = Path.home() / ".mblt_model_zoo" + try: + relative_path = path.relative_to(default_cache_root) + except ValueError: + return path + return Path(get_mobilint_cache_dir()) / relative_path + + +def _validate_subset_size(subset_size: int) -> None: + """Validate a requested calibration subset size. + + Args: + subset_size: Requested per-class or total sample count. + + Raises: + ValueError: If the size is not positive. + """ + + if subset_size <= 0: + raise ValueError("subset_size must be greater than zero.") + + +def _validate_calibration_image(image_path: Path, dataset_root: Path) -> Path: + """Resolve and validate a calibration image within its dataset root. + + Args: + image_path: Candidate calibration image path. + dataset_root: Resolved organized dataset root. + + Returns: + Resolved regular image path. + + Raises: + ValueError: If the candidate is a symlink, not a regular image, or escapes the dataset root. + """ + + if image_path.is_symlink(): + raise ValueError(f"Calibration image must not be a symlink: {image_path}.") + try: + source = image_path.resolve(strict=True) + except OSError as exc: + raise ValueError( + f"Unable to resolve calibration image {image_path}: {exc}." + ) from exc + if not source.is_file() or source.suffix.lower() not in IMAGE_SUFFIXES: + raise ValueError( + f"Calibration image must be a supported regular image file: {image_path}." + ) + if not source.is_relative_to(dataset_root): + raise ValueError( + f"Calibration image must remain within dataset root: {image_path}." + ) + return source + + +def _selectable_calibration_images( + paths: Sequence[Path], dataset_root: Path +) -> list[Path]: + """Validate supported image candidates before they enter subset sampling. + + Args: + paths: Candidate filesystem paths. + dataset_root: Resolved organized dataset root. + + Returns: + Validated candidate paths, retaining their original names for stable sampling. + """ + + images: list[Path] = [] + for path in paths: + if path.suffix.lower() in IMAGE_SUFFIXES and ( + path.is_symlink() or path.is_file() + ): + _validate_calibration_image(path, dataset_root) + images.append(path) + return images + + +def select_calibration_images( + task: str, + data_path: str | Path, + subset_size: int | None = None, + seed: int = DEFAULT_SEED, +) -> list[Path]: + """Select deterministic calibration images from an organized dataset. + + For ImageNet and WiderFace, ``subset_size`` is the number selected from + every category. For all other datasets it is the total number of selected + validation images. + + Args: + task: Vision task name. + data_path: Organized dataset root. + subset_size: Optional sample count, using task-specific defaults when omitted. + seed: Random selection seed. + + Returns: + Selected image paths in deterministic order. + + Raises: + ValueError: If the dataset layout or requested size is invalid. + """ + + import random + + normalized_task = _normalize_task(task) + root = Path(data_path).expanduser() + resolved_root = root.resolve() + requested_size = ( + DEFAULT_SUBSET_SIZES[normalized_task] if subset_size is None else subset_size + ) + _validate_subset_size(requested_size) + random_generator = random.Random(seed) + + if normalized_task in {"image_classification", "face_detection"}: + category_root = ( + root if normalized_task == "image_classification" else root / "images" + ) + dataset_name = ( + "ImageNet" if normalized_task == "image_classification" else "WiderFace" + ) + category_dirs = ( + sorted(path for path in category_root.iterdir() if path.is_dir()) + if category_root.is_dir() + else [] + ) + if not category_dirs: + raise ValueError( + f"No {dataset_name} category directories found in {category_root}." + ) + selected: list[Path] = [] + for category_dir in category_dirs: + images = _selectable_calibration_images( + sorted(category_dir.iterdir()), resolved_root + ) + if requested_size > len(images): + raise ValueError( + f"subset_size ({requested_size}) exceeds the {len(images)} available images in {category_dir.name}." + ) + selected.extend(random_generator.sample(images, requested_size)) + return selected + + image_dir = { + "depth_estimation": root / "images", + "object_detection": root / "val2017", + "instance_segmentation": root / "val2017", + "semantic_segmentation": root / "images", + "pose_estimation": root / "val2017", + "obb": root / "images", + }[normalized_task] + images = ( + _selectable_calibration_images(sorted(image_dir.rglob("*")), resolved_root) + if image_dir.is_dir() + else [] + ) + if not images: + raise ValueError(f"No calibration images found in {image_dir}.") + if requested_size > len(images): + raise ValueError( + f"subset_size ({requested_size}) exceeds the {len(images)} available images in {image_dir}." + ) + return random_generator.sample(images, requested_size) + + +def copy_calibration_subset( + images: Sequence[Path], data_path: str | Path, output_dir: str | Path +) -> list[Path]: + """Copy selected images into a flat directory using collision-safe names. + + Args: + images: Selected source image paths. + data_path: Dataset root used to derive stable relative names. + output_dir: Destination directory. + + Returns: + Copied image paths in input order. + """ + + root = Path(data_path).expanduser().resolve() + destination = _validate_calibration_output_dir(output_dir) + destination.mkdir(parents=True, exist_ok=True) + copied: list[Path] = [] + for image_path in images: + source = _validate_calibration_image(image_path, root) + relative_name = source.relative_to(root).as_posix() + digest = hashlib.sha256(relative_name.encode()).hexdigest()[:12] + target = destination / f"{digest}_{source.name}" + if target.exists(): + raise ValueError(f"Calibration subset filename collision for {source}.") + shutil.copy2(source, target) + copied.append(target) + return copied + + +def make_calibration_subset( + task: str, + data_path: str | Path, + output_dir: str | Path, + subset_size: int | None = None, + seed: int = DEFAULT_SEED, +) -> list[Path]: + """Select and copy a deterministic flat calibration subset. + + Args: + task: Vision task name. + data_path: Organized dataset root. + output_dir: Destination for copied images. + subset_size: Optional task-specific selection count. + seed: Random selection seed. + + Returns: + Copied image paths. + """ + + source_root = Path(data_path).expanduser().resolve() + destination = _validate_calibration_output_dir(output_dir).resolve() + if source_root.is_relative_to(destination) or destination.is_relative_to( + source_root + ): + raise ValueError("output_dir must not overlap data_path.") + images = select_calibration_images(task, data_path, subset_size, seed) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary_destination = Path( + mkdtemp(prefix=f".{destination.name}-", dir=destination.parent) + ) + try: + copied = copy_calibration_subset(images, data_path, temporary_destination) + backup_destination: Path | None = None + if destination.exists(): + backup_destination = Path( + mkdtemp(prefix=f".{destination.name}-backup-", dir=destination.parent) + ) + backup_destination.rmdir() + destination.replace(backup_destination) + try: + temporary_destination.replace(destination) + except BaseException: + if backup_destination is not None and not destination.exists(): + backup_destination.replace(destination) + raise + if backup_destination is not None: + shutil.rmtree(backup_destination) + return [destination / path.name for path in copied] + finally: + if temporary_destination.exists(): + shutil.rmtree(temporary_destination) + + +def _as_hwc_float32(value: Any, image_path: Path) -> np.ndarray: + """Validate and convert engine preprocessing output for qbcompiler. + + Args: + value: Engine preprocessing result. + image_path: Source image used for error context. + + Returns: + Contiguous HWC float32 array. + + Raises: + TypeError: If preprocessing did not return a tensor or array. + ValueError: If preprocessing did not return HWC three-channel data. + """ + + if isinstance(value, torch.Tensor): + array = value.detach().cpu().numpy() + elif isinstance(value, np.ndarray): + array = value + else: + raise TypeError( + f"Preprocessing {image_path} returned unsupported type {type(value).__name__}." + ) + if array.ndim != 3 or array.shape[-1] != 3: + raise ValueError( + f"Preprocessing {image_path} must produce an HWC three-channel array; got shape {array.shape}." + ) + return np.ascontiguousarray(array, dtype=np.float32) + + +def prepare_calibration_arrays( + engine: MBLT_Engine, images: Sequence[Path], output_dir: str | Path +) -> list[Path]: + """Preprocess calibration images and save one NumPy array per sample. + + Args: + engine: ONNX vision engine providing authoritative preprocessing. + images: Flat calibration image paths. + output_dir: Destination directory for ``.npy`` arrays. + + Returns: + Saved NumPy paths. + """ + + destination = _validate_calibration_output_dir(output_dir) + destination.mkdir(parents=True, exist_ok=True) + saved: list[Path] = [] + for index, image_path in enumerate(images): + array = _as_hwc_float32(engine.preprocess(str(image_path)), image_path) + array_path = destination / f"{index:06d}.npy" + np.save(array_path, array) + saved.append(array_path) + return saved + + +def get_subset_images(subset_path: str | Path) -> list[Path]: + """Load all images from an already-sampled subset. + + Args: + subset_path: Root containing sampled images, either flat or nested. + + Returns: + Image paths in deterministic order. + + Raises: + ValueError: If the subset contains no supported image files. + """ + + root = Path(subset_path).expanduser() + images = ( + sorted( + path + for path in root.rglob("*") + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + if root.is_dir() + else [] + ) + if not images: + raise ValueError(f"No calibration subset images found in {root}.") + return images + + +def _calibration_image_shape(pre_cfg: Mapping[str, Any]) -> tuple[int, int]: + """Derive the final HWC spatial shape produced by a model preprocessor.""" + + image_shape: tuple[int, int] | None = None + for operation, config in pre_cfg.items(): + if operation not in {"LetterBox", "Resize", "CenterCrop"} or not isinstance( + config, Mapping + ): + continue + size = config.get("img_size", config.get("size")) + if isinstance(size, int) and not isinstance(size, bool) and size > 0: + image_shape = (size, size) + elif ( + isinstance(size, Sequence) + and not isinstance(size, (str, bytes)) + and len(size) == 2 + and all( + isinstance(value, int) and not isinstance(value, bool) and value > 0 + for value in size + ) + ): + image_shape = (size[0], size[1]) + if image_shape is None: + raise ValueError( + "Unable to derive calibration image shape from pre_cfg; expected a " + "positive LetterBox.img_size, Resize.size, or CenterCrop.size." + ) + return image_shape + + +def validate_calibration_dataset( + calib_data_path: str | Path, + *, + image_shape: tuple[int, int] | None = None, +) -> Path: + """Validate a ready directory of preprocessed calibration arrays. + + Args: + calib_data_path: Directory containing HWC float32 ``.npy`` tensors. + + Returns: + Expanded calibration directory path. + + Raises: + ValueError: If the directory is empty or contains an invalid calibration tensor. + """ + + root = Path(calib_data_path).expanduser() + array_paths = sorted(root.glob("*.npy")) if root.is_dir() else [] + if not array_paths: + raise ValueError(f"No calibration .npy files found in {root}.") + for array_path in array_paths: + try: + array = np.load(array_path, mmap_mode="r", allow_pickle=False) + except (OSError, ValueError) as exc: + raise ValueError( + f"Unable to load calibration tensor {array_path}: {exc}." + ) from exc + if array.ndim != 3 or array.shape[-1] != 3: + raise ValueError( + f"Calibration tensor {array_path} must be HWC with three channels; got {array.shape}." + ) + if image_shape is not None and array.shape[:2] != image_shape: + raise ValueError( + f"Calibration tensor {array_path} must match the model pre_cfg " + f"image shape {image_shape}, got {array.shape[:2]}." + ) + if array.dtype != np.float32: + raise ValueError( + f"Calibration tensor {array_path} must use float32; got {array.dtype}." + ) + if not np.isfinite(array).all(): + raise ValueError( + f"Calibration tensor {array_path} must contain only finite values." + ) + if not array.flags.c_contiguous: + raise ValueError(f"Calibration tensor {array_path} must be C-contiguous.") + return root + + +def _validate_data_level_paths( + data_path: str | Path | None, + subset_path: str | Path | None, + calib_data_path: str | Path | None, +) -> None: + """Ensure at most one calibration pipeline entry level is supplied. + + Args: + data_path: Original organized dataset root. + subset_path: Already-sampled image subset root. + calib_data_path: Ready preprocessed NumPy dataset root. + + Raises: + ValueError: If more than one data level is supplied. + """ + + supplied = [ + name + for name, value in ( + ("data_path", data_path), + ("subset_path", subset_path), + ("calib_data_path", calib_data_path), + ) + if value is not None + ] + if len(supplied) > 1: + raise ValueError( + "Provide only one calibration pipeline input: `data_path`, `subset_path`, or `calib_data_path`; " + f"got {', '.join(supplied)}." + ) + + +def _load_qbcompiler() -> tuple[Callable[..., Any], type[Any]]: + """Load qbcompiler only when compilation is requested. + + Returns: + The ``mxq_compile`` function and ``CalibrationConfig`` class. + + Raises: + ImportError: If qbcompiler is not installed. + """ + + try: + module = importlib.import_module("qbcompiler") + except ImportError as exc: + raise ImportError( + "Vision compilation requires qbcompiler>=1.2.0. Install the compiler package supplied by Mobilint." + ) from exc + return module.mxq_compile, module.CalibrationConfig + + +def _validate_ratio(name: str, value: Any) -> float: + """Validate a quantization ratio. + + Args: + name: Field name used in an error message. + value: Candidate numeric ratio. + + Returns: + Validated float value. + + Raises: + ValueError: If the value is non-numeric or outside zero to one. + """ + + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError( + f"Quantization metadata `{name}` must be numeric, got {value!r}." + ) + resolved = float(value) + if not 0.0 <= resolved <= 1.0: + raise ValueError( + f"Quantization metadata `{name}` must be between 0 and 1, got {resolved}." + ) + return resolved + + +def _fetch_quantization_config( + repo_id: str, revision: str, target_device: str = "aries-rb" +) -> dict[str, Any] | None: + """Fetch optional hosted board quantization metadata. + + Args: + repo_id: Hugging Face model repository ID. + revision: Repository revision. + target_device: Board folder containing the metadata. + + Returns: + Hosted ``config`` mapping, or ``None`` when the optional file is unavailable. + + Raises: + ValueError: If the hosted JSON or its ``config`` field is malformed. + """ + + try: + metadata_path = hf_hub_download( + repo_id=repo_id, + filename="best_result.json", + subfolder=normalize_target_device(target_device), + revision=revision, + ) + except (HfHubHTTPError, OSError) as exc: + warnings.warn( + f"Unable to load optional quantization metadata for {repo_id}: {exc}. Using fallback values.", + stacklevel=2, + ) + return None + + try: + with Path(metadata_path).open(encoding="utf-8") as metadata_file: + metadata = json.load(metadata_file) + except json.JSONDecodeError as exc: + raise ValueError( + f"Malformed quantization metadata JSON at {metadata_path}: {exc.msg}." + ) from exc + if not isinstance(metadata, dict) or not isinstance(metadata.get("config"), dict): + raise ValueError( + f"Quantization metadata at {metadata_path} must contain a `config` object." + ) + return metadata["config"] + + +def resolve_quantization_values( + file_cfg: dict[str, Any], + percentile: float | None, + topk_ratio: float | None, +) -> tuple[float, float]: + """Resolve explicit, hosted, and fallback quantization values independently. + + Args: + file_cfg: Resolved model file configuration. + percentile: Explicit compiler percentile override. + topk_ratio: Explicit compiler top-k ratio override. + + Returns: + Resolved percentile and top-k ratio. + + Raises: + ValueError: If explicit or hosted values are malformed. + """ + + resolved_percentile = ( + _validate_ratio("percentile", percentile) if percentile is not None else None + ) + resolved_topk = ( + _validate_ratio("topk_ratio", topk_ratio) if topk_ratio is not None else None + ) + hosted_config: dict[str, Any] | None = None + if resolved_percentile is None or resolved_topk is None: + repo_id = file_cfg.get("repo_id") + revision = file_cfg.get("revision", "main") + if isinstance(repo_id, str) and repo_id: + target_device = file_cfg.get("target_device") + if isinstance(target_device, str): + hosted_config = _fetch_quantization_config( + repo_id, str(revision), normalize_target_device(target_device) + ) + else: + hosted_config = _fetch_quantization_config(repo_id, str(revision)) + else: + warnings.warn( + "Model configuration has no repository ID; using fallback quantization values.", + stacklevel=2, + ) + + if ( + resolved_percentile is None + and hosted_config is not None + and "percentile" in hosted_config + ): + hosted_percentile = _validate_ratio( + "config.percentile", hosted_config["percentile"] + ) + resolved_percentile = 1.0 - hosted_percentile + if resolved_topk is None and hosted_config is not None: + hosted_topk = hosted_config.get("topk_ratio", hosted_config.get("topk")) + if hosted_topk is not None: + resolved_topk = _validate_ratio("config.topk_ratio", hosted_topk) + + if resolved_percentile is None: + warnings.warn( + f"Quantization percentile is unavailable; using {DEFAULT_PERCENTILE}.", + stacklevel=2, + ) + resolved_percentile = DEFAULT_PERCENTILE + if resolved_topk is None: + warnings.warn( + f"Quantization top-k ratio is unavailable; using {DEFAULT_TOPK_RATIO}.", + stacklevel=2, + ) + resolved_topk = DEFAULT_TOPK_RATIO + return resolved_percentile, resolved_topk + + +def _configured_onnx_filename(file_cfg: dict[str, Any]) -> str | None: + """Derive the configured ONNX filename. + + Args: + file_cfg: Resolved model file configuration. + + Returns: + ONNX filename when it can be determined. + """ + + onnx_filename = file_cfg.get("onnx_filename") + if isinstance(onnx_filename, str) and onnx_filename: + return onnx_filename + filename = file_cfg.get("filename") + if isinstance(filename, str) and filename: + return f"{Path(filename).stem}.onnx" + return None + + +def _resolve_compile_onnx_path( + file_cfg: dict[str, Any], model_path: str | Path | None +) -> Path: + """Resolve an ONNX artifact for compilation without constructing an inference runtime. + + Args: + file_cfg: Resolved model file configuration. + model_path: Optional caller-provided ONNX path. + + Returns: + Existing local or downloaded ONNX artifact path. + + Raises: + FileNotFoundError: If no ONNX artifact can be resolved. + """ + + if model_path is not None: + local_path = Path(model_path).expanduser() + if local_path.is_file(): + if local_path.suffix.lower() != ".onnx": + raise ValueError( + "Compilation requires an ONNX model; " + f"local model path must end with `.onnx`, got {local_path}." + ) + return local_path.resolve() + configured_path = Path(str(file_cfg.get("onnx_path", ""))).expanduser() + if configured_path.is_file(): + return configured_path.resolve() + repo_id = file_cfg.get("repo_id") + revision = file_cfg.get("revision") + filename = _configured_onnx_filename(file_cfg) + if isinstance(repo_id, str) and isinstance(revision, str) and filename: + downloaded_path = Path( + hf_hub_download( + repo_id=repo_id, + filename=filename, + revision=revision, + local_dir=get_mobilint_cache_dir(), + ) + ) + if downloaded_path.is_file(): + return downloaded_path.resolve() + configured_name = filename or "the configured ONNX artifact" + raise FileNotFoundError(f"Unable to resolve {configured_name} for compilation.") + + +def _resolve_compile_output_path( + save_path: str | Path | None, resolved_onnx: Path +) -> Path: + """Resolve and validate an MXQ compilation output path.""" + + requested_output_path = ( + Path(save_path).expanduser() + if save_path is not None + else _default_model_dir() / f"{resolved_onnx.stem}.mxq" + ) + if requested_output_path.suffix.lower() != ".mxq": + raise ValueError( + "Compilation output path must end with `.mxq`, " + f"got {requested_output_path}." + ) + output_path = requested_output_path.resolve() + if output_path == resolved_onnx.resolve(): + raise ValueError( + "Compilation output path must not be the same as the input ONNX path: " + f"{output_path}." + ) + if output_path.suffix.lower() != ".mxq": + raise ValueError( + "Compilation output path must resolve to a path ending with `.mxq`, " + f"got {output_path}." + ) + return output_path + + +def compile_vision_model( + model_cls: str, + *, + target_device: str, + model_type: str = "DEFAULT", + model_path: str | Path | None = None, + onnx_path: str | Path | None = None, + data_path: str | Path | None = None, + subset_path: str | Path | None = None, + calib_data_path: str | Path | None = None, + save_path: str | Path | None = None, + subset_size: int | None = None, + seed: int = DEFAULT_SEED, + percentile: float | None = None, + topk_ratio: float | None = None, +) -> Path: + """Compile a configured vision ONNX model into a board-specific MXQ artifact. + + Args: + model_cls: Vision model name or YAML path. + target_device: Required target board: ``aries-rb``, ``regulus-ra``, or ``regulus-rb``. + model_type: Model variant from the YAML configuration. + model_path: Preferred local ONNX path compatibility option. + onnx_path: Local ONNX path alias. + data_path: Original organized dataset root. The pipeline organizes it when needed, samples + an image subset, and preprocesses that subset. + subset_path: Already-sampled image subset. The pipeline preprocesses every image directly + without organizing or sampling an original dataset. + calib_data_path: Ready directory of HWC float32 ``.npy`` tensors. The pipeline passes it + directly to qbcompiler without dataset preparation, sampling, or preprocessing. + save_path: Output MXQ path. Defaults to the ONNX stem under ``~/.mblt_model_zoo``. + subset_size: ImageNet/WiderFace per-category count or total count for other datasets. + seed: Deterministic calibration selection seed. + percentile: Explicit quantization percentile. + topk_ratio: Explicit quantization top-k ratio. + + Returns: + Output MXQ path. + + Raises: + ImportError: If qbcompiler is unavailable. + ValueError: If model metadata, calibration data, or quantization values are invalid. + """ + + _validate_data_level_paths(data_path, subset_path, calib_data_path) + target_device = normalize_target_device(target_device) + if target_device not in SUPPORTED_TARGET_DEVICES: + raise ValueError( + f"Unsupported target_device {target_device!r}; expected one of " + f"{sorted(SUPPORTED_TARGET_DEVICES)}." + ) + mxq_compile, calibration_config_class = _load_qbcompiler() + model_config = resolve_model_config(model_cls, model_type) + file_cfg = model_config.get("file_cfg") + pre_cfg = model_config.get("pre_cfg") + post_cfg = model_config.get("post_cfg") + if ( + not isinstance(file_cfg, dict) + or not isinstance(pre_cfg, dict) + or not isinstance(post_cfg, dict) + ): + raise ValueError( + "Resolved vision model configuration requires `file_cfg`, `pre_cfg`, and `post_cfg` objects." + ) + file_cfg["target_device"] = target_device + + task = _normalize_task(str(post_cfg.get("task", ""))) + dataset = post_cfg.get("dataset") + if not isinstance(dataset, str) or not dataset: + raise ValueError( + "Resolved vision model configuration requires a non-empty `post_cfg.dataset` value." + ) + selected_local_path = model_path or onnx_path + if calib_data_path is not None: + resolved_onnx = _resolve_compile_onnx_path(file_cfg, selected_local_path) + output_path = _resolve_compile_output_path(save_path, resolved_onnx) + output_path.parent.mkdir(parents=True, exist_ok=True) + resolved_percentile, resolved_topk = resolve_quantization_values( + file_cfg, percentile, topk_ratio + ) + calibration_config = calibration_config_class( + method=1, + output=0 if task == "image_classification" else 1, + mode=1, + max_percentile={ + "percentile": resolved_percentile, + "topk_ratio": resolved_topk, + }, + ) + mxq_compile( + model=str(resolved_onnx), + calib_data_path=str( + validate_calibration_dataset( + calib_data_path, + image_shape=_calibration_image_shape(pre_cfg), + ) + ), + save_path=str(output_path), + image_channels=3, + backend="onnx", + device="gpu", + target_device=target_device, + inference_scheme="all", + calibration_config=calibration_config, + ) + return output_path + engine_kwargs: dict[str, Any] = { + "model_cls": model_cls, + "model_type": model_type, + "framework": "onnx", + "onnx_providers": ["CPUExecutionProvider"], + } + if selected_local_path is not None: + expanded_local_path = Path(selected_local_path).expanduser() + if expanded_local_path.is_file(): + engine_kwargs["model_path"] = str(expanded_local_path) + + engine: MBLT_Engine | None = None + try: + engine = MBLT_Engine(**engine_kwargs) + resolved_onnx = Path(str(engine.file_cfg.get("onnx_path", ""))).expanduser() + if not resolved_onnx.is_file(): + configured_name = ( + _configured_onnx_filename(file_cfg) or "the configured ONNX artifact" + ) + raise FileNotFoundError( + f"Unable to resolve {configured_name} for model `{model_cls}`." + ) + + output_path = _resolve_compile_output_path(save_path, resolved_onnx) + output_path.parent.mkdir(parents=True, exist_ok=True) + resolved_percentile, resolved_topk = resolve_quantization_values( + file_cfg, percentile, topk_ratio + ) + calibration_config = calibration_config_class( + method=1, + output=0 if task == "image_classification" else 1, + mode=1, + max_percentile={ + "percentile": resolved_percentile, + "topk_ratio": resolved_topk, + }, + ) + + def _compile(calibration_path: Path) -> None: + """Run qbcompiler with a resolved calibration directory.""" + + mxq_compile( + model=str(resolved_onnx), + calib_data_path=str(calibration_path), + save_path=str(output_path), + image_channels=3, + backend="onnx", + device="gpu", + target_device=target_device, + inference_scheme="all", + calibration_config=calibration_config, + ) + + with TemporaryDirectory(prefix="mblt-vision-calibration-") as temporary_root: + temporary_path = Path(temporary_root) + if subset_path is not None: + subset_images = get_subset_images(subset_path) + else: + dataset_path = ensure_calibration_dataset(task, data_path, dataset) + subset_images = make_calibration_subset( + task, + dataset_path, + temporary_path / "images", + subset_size=subset_size, + seed=seed, + ) + array_dir = temporary_path / "arrays" + prepare_calibration_arrays(engine, subset_images, array_dir) + _compile(array_dir) + return output_path + finally: + if engine is not None: + engine.dispose() + + +__all__ = [ + "compile_vision_model", + "copy_calibration_subset", + "ensure_calibration_dataset", + "get_subset_images", + "make_calibration_subset", + "prepare_calibration_arrays", + "resolve_quantization_values", + "select_calibration_images", + "validate_calibration_dataset", +] diff --git a/mblt_vision/datasets/__init__.py b/mblt_vision/datasets/__init__.py new file mode 100644 index 0000000..0d4dae4 --- /dev/null +++ b/mblt_vision/datasets/__init__.py @@ -0,0 +1,15 @@ +"""YAML-backed definitions for vision validation datasets.""" + +from .registry import ( + get_dataset_category_ids, + get_dataset_class_names, + get_dataset_config, + get_dataset_config_for_task, +) + +__all__ = [ + "get_dataset_category_ids", + "get_dataset_class_names", + "get_dataset_config", + "get_dataset_config_for_task", +] diff --git a/mblt_vision/datasets/ade20k.yaml b/mblt_vision/datasets/ade20k.yaml new file mode 100644 index 0000000..2423bc6 --- /dev/null +++ b/mblt_vision/datasets/ade20k.yaml @@ -0,0 +1,164 @@ +# ADE20K validation dataset. Organized paths are relative to `path`. +name: ade20k +path: ~/.mblt_model_zoo/datasets/ADEChallengeData2016 +val: images +tasks: + - semantic_segmentation +names: + 0: wall + 1: building + 2: sky + 3: floor + 4: tree + 5: ceiling + 6: road + 7: bed + 8: windowpane + 9: grass + 10: cabinet + 11: sidewalk + 12: person + 13: earth + 14: door + 15: table + 16: mountain + 17: plant + 18: curtain + 19: chair + 20: car + 21: water + 22: painting + 23: sofa + 24: shelf + 25: house + 26: sea + 27: mirror + 28: rug + 29: field + 30: armchair + 31: seat + 32: fence + 33: desk + 34: rock + 35: wardrobe + 36: lamp + 37: bathtub + 38: railing + 39: cushion + 40: base + 41: box + 42: column + 43: signboard + 44: chest of drawers + 45: counter + 46: sand + 47: sink + 48: skyscraper + 49: fireplace + 50: refrigerator + 51: grandstand + 52: path + 53: stairs + 54: runway + 55: case + 56: pool table + 57: pillow + 58: screen door + 59: stairway + 60: river + 61: bridge + 62: bookcase + 63: blind + 64: coffee table + 65: toilet + 66: flower + 67: book + 68: hill + 69: bench + 70: countertop + 71: stove + 72: palm + 73: kitchen island + 74: computer + 75: swivel chair + 76: boat + 77: bar + 78: arcade machine + 79: hovel + 80: bus + 81: towel + 82: light + 83: truck + 84: tower + 85: chandelier + 86: awning + 87: streetlight + 88: booth + 89: television receiver + 90: airplane + 91: dirt track + 92: apparel + 93: pole + 94: land + 95: bannister + 96: escalator + 97: ottoman + 98: bottle + 99: buffet + 100: poster + 101: stage + 102: van + 103: ship + 104: fountain + 105: conveyor belt + 106: canopy + 107: washer + 108: plaything + 109: swimming pool + 110: stool + 111: barrel + 112: basket + 113: waterfall + 114: tent + 115: bag + 116: minibike + 117: cradle + 118: oven + 119: ball + 120: food + 121: step + 122: tank + 123: trade name + 124: microwave + 125: pot + 126: animal + 127: bicycle + 128: lake + 129: dishwasher + 130: screen + 131: blanket + 132: sculpture + 133: hood + 134: sconce + 135: vase + 136: traffic light + 137: tray + 138: ashcan + 139: fan + 140: pier + 141: crt screen + 142: plate + 143: monitor + 144: bulletin board + 145: shower + 146: radiator + 147: glass + 148: clock + 149: flag +download: + url: https://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip + sha256: 7ff1be44964418441f542a7cc1e1a650e7dc0fc275f5d23252bc9bbdbc977b29 +nc: 150 +source_ignore_label: 0 +ignore_label: 255 +label_offset: 1 diff --git a/mblt_vision/datasets/cityscapes.yaml b/mblt_vision/datasets/cityscapes.yaml new file mode 100644 index 0000000..e33b106 --- /dev/null +++ b/mblt_vision/datasets/cityscapes.yaml @@ -0,0 +1,54 @@ +# Cityscapes validation taxonomy from the official train/validation/test archives. +# Organized paths are relative to `path`. +name: cityscapes +path: ~/.mblt_model_zoo/datasets/cityscapes +val: images +tasks: + - semantic_segmentation +nc: 19 +names: + 0: road + 1: sidewalk + 2: building + 3: wall + 4: fence + 5: pole + 6: traffic light + 7: traffic sign + 8: vegetation + 9: terrain + 10: sky + 11: person + 12: rider + 13: car + 14: truck + 15: bus + 16: train + 17: motorcycle + 18: bicycle +category_ids: [7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33] +palette: + 0: [128, 64, 128] + 1: [244, 35, 232] + 2: [70, 70, 70] + 3: [102, 102, 156] + 4: [190, 153, 153] + 5: [153, 153, 153] + 6: [250, 170, 30] + 7: [220, 220, 0] + 8: [107, 142, 35] + 9: [152, 251, 152] + 10: [70, 130, 180] + 11: [220, 20, 60] + 12: [255, 0, 0] + 13: [0, 0, 142] + 14: [0, 0, 70] + 15: [0, 60, 100] + 16: [0, 80, 100] + 17: [0, 0, 230] + 18: [119, 11, 32] +download: + type: cityscapes + source: https://github.com/mcordts/cityscapesScripts + images_archive: leftImg8bit_trainvaltest.zip + annotations_archive: gtFine_trainvaltest.zip diff --git a/mblt_vision/datasets/coco.yaml b/mblt_vision/datasets/coco.yaml new file mode 100644 index 0000000..020e094 --- /dev/null +++ b/mblt_vision/datasets/coco.yaml @@ -0,0 +1,178 @@ +# COCO 2017 validation dataset. Paths are relative to `path` unless absolute. +name: coco +path: ~/.mblt_model_zoo/datasets/coco +tasks: + - object_detection + - instance_segmentation + - pose_estimation +val: val2017 +download: + # The official COCO object-store hostname has no valid TLS certificate. Use a + # pinned HTTPS mirror and verify each archive digest before extraction. + images: https://huggingface.co/datasets/pcuenq/coco-2017-mirror/resolve/5200d2cffae8121713bec767b4693bdafc0eeb0b/val2017.zip + images_sha256: 4f7e2ccb2866ec5041993c9cf2a952bbed69647b115d0f74da7ce8f4bef82f05 + annotations: https://huggingface.co/datasets/pcuenq/coco-2017-mirror/resolve/5200d2cffae8121713bec767b4693bdafc0eeb0b/annotations_trainval2017.zip + annotations_sha256: 113a836d90195ee1f884e704da6304dfaaecff1f023f49b6ca93c4aaae470268 + +names: + 0: person + 1: bicycle + 2: car + 3: motorcycle + 4: airplane + 5: bus + 6: train + 7: truck + 8: boat + 9: traffic light + 10: fire hydrant + 11: stop sign + 12: parking meter + 13: bench + 14: bird + 15: cat + 16: dog + 17: horse + 18: sheep + 19: cow + 20: elephant + 21: bear + 22: zebra + 23: giraffe + 24: backpack + 25: umbrella + 26: handbag + 27: tie + 28: suitcase + 29: frisbee + 30: skis + 31: snowboard + 32: sports ball + 33: kite + 34: baseball bat + 35: baseball glove + 36: skateboard + 37: surfboard + 38: tennis racket + 39: bottle + 40: wine glass + 41: cup + 42: fork + 43: knife + 44: spoon + 45: bowl + 46: banana + 47: apple + 48: sandwich + 49: orange + 50: broccoli + 51: carrot + 52: hot dog + 53: pizza + 54: donut + 55: cake + 56: chair + 57: couch + 58: potted plant + 59: bed + 60: dining table + 61: toilet + 62: tv + 63: laptop + 64: mouse + 65: remote + 66: keyboard + 67: cell phone + 68: microwave + 69: oven + 70: toaster + 71: sink + 72: refrigerator + 73: book + 74: clock + 75: vase + 76: scissors + 77: teddy bear + 78: hair drier + 79: toothbrush +category_ids: +- 1 +- 2 +- 3 +- 4 +- 5 +- 6 +- 7 +- 8 +- 9 +- 10 +- 11 +- 13 +- 14 +- 15 +- 16 +- 17 +- 18 +- 19 +- 20 +- 21 +- 22 +- 23 +- 24 +- 25 +- 27 +- 28 +- 31 +- 32 +- 33 +- 34 +- 35 +- 36 +- 37 +- 38 +- 39 +- 40 +- 41 +- 42 +- 43 +- 44 +- 46 +- 47 +- 48 +- 49 +- 50 +- 51 +- 52 +- 53 +- 54 +- 55 +- 56 +- 57 +- 58 +- 59 +- 60 +- 61 +- 62 +- 63 +- 64 +- 65 +- 67 +- 70 +- 72 +- 73 +- 74 +- 75 +- 76 +- 77 +- 78 +- 79 +- 80 +- 81 +- 82 +- 84 +- 85 +- 86 +- 87 +- 88 +- 89 +- 90 diff --git a/mblt_vision/datasets/dotav1.yaml b/mblt_vision/datasets/dotav1.yaml new file mode 100644 index 0000000..827c28a --- /dev/null +++ b/mblt_vision/datasets/dotav1.yaml @@ -0,0 +1,28 @@ +# DOTAv1 validation dataset. Paths are relative to `path` unless absolute. +name: dotav1 +path: ~/.mblt_model_zoo/datasets/dotav1 +tasks: + - obb +val: images +labels: labels/val_original +names: + 0: plane + 1: ship + 2: storage-tank + 3: baseball-diamond + 4: tennis-court + 5: basketball-court + 6: ground-track-field + 7: harbor + 8: bridge + 9: large-vehicle + 10: small-vehicle + 11: helicopter + 12: roundabout + 13: soccer-ball-field + 14: swimming-pool +download: + type: google_drive_folder + url: https://drive.google.com/drive/folders/1n5w45suVOyaqY84hltJhIZdtVFD9B224 + images_archive: images/part1.zip + labels_archive: labelTxt-v1.0/labelTxt.zip diff --git a/mblt_vision/datasets/imagenet.yaml b/mblt_vision/datasets/imagenet.yaml new file mode 100644 index 0000000..973c153 --- /dev/null +++ b/mblt_vision/datasets/imagenet.yaml @@ -0,0 +1,1019 @@ +# ImageNet-1k validation dataset. Paths are relative to `path` unless absolute. +name: imagenet +path: ~/.mblt_model_zoo/datasets/imagenet +tasks: + - image_classification +val: . +download: + images: https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar + annotations: https://www.image-net.org/data/ILSVRC/2012/ILSVRC2012_bbox_val_v3.tgz + +names: + 0: tench, Tinca tinca + 1: goldfish, Carassius auratus + 2: great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias + 3: tiger shark, Galeocerdo cuvieri + 4: hammerhead, hammerhead shark + 5: electric ray, crampfish, numbfish, torpedo + 6: stingray + 7: cock + 8: hen + 9: ostrich, Struthio camelus + 10: brambling, Fringilla montifringilla + 11: goldfinch, Carduelis carduelis + 12: house finch, linnet, Carpodacus mexicanus + 13: junco, snowbird + 14: indigo bunting, indigo finch, indigo bird, Passerina cyanea + 15: robin, American robin, Turdus migratorius + 16: bulbul + 17: jay + 18: magpie + 19: chickadee + 20: water ouzel, dipper + 21: kite + 22: bald eagle, American eagle, Haliaeetus leucocephalus + 23: vulture + 24: great grey owl, great gray owl, Strix nebulosa + 25: European fire salamander, Salamandra salamandra + 26: common newt, Triturus vulgaris + 27: eft + 28: spotted salamander, Ambystoma maculatum + 29: axolotl, mud puppy, Ambystoma mexicanum + 30: bullfrog, Rana catesbeiana + 31: tree frog, tree-frog + 32: tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui + 33: loggerhead, loggerhead turtle, Caretta caretta + 34: leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea + 35: mud turtle + 36: terrapin + 37: box turtle, box tortoise + 38: banded gecko + 39: common iguana, iguana, Iguana iguana + 40: American chameleon, anole, Anolis carolinensis + 41: whiptail, whiptail lizard + 42: agama + 43: frilled lizard, Chlamydosaurus kingi + 44: alligator lizard + 45: Gila monster, Heloderma suspectum + 46: green lizard, Lacerta viridis + 47: African chameleon, Chamaeleo chamaeleon + 48: Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis + 49: African crocodile, Nile crocodile, Crocodylus niloticus + 50: American alligator, Alligator mississipiensis + 51: triceratops + 52: thunder snake, worm snake, Carphophis amoenus + 53: ringneck snake, ring-necked snake, ring snake + 54: hognose snake, puff adder, sand viper + 55: green snake, grass snake + 56: king snake, kingsnake + 57: garter snake, grass snake + 58: water snake + 59: vine snake + 60: night snake, Hypsiglena torquata + 61: boa constrictor, Constrictor constrictor + 62: rock python, rock snake, Python sebae + 63: Indian cobra, Naja naja + 64: green mamba + 65: sea snake + 66: horned viper, cerastes, sand viper, horned asp, Cerastes cornutus + 67: diamondback, diamondback rattlesnake, Crotalus adamanteus + 68: sidewinder, horned rattlesnake, Crotalus cerastes + 69: trilobite + 70: harvestman, daddy longlegs, Phalangium opilio + 71: scorpion + 72: black and gold garden spider, Argiope aurantia + 73: barn spider, Araneus cavaticus + 74: garden spider, Aranea diademata + 75: black widow, Latrodectus mactans + 76: tarantula + 77: wolf spider, hunting spider + 78: tick + 79: centipede + 80: black grouse + 81: ptarmigan + 82: ruffed grouse, partridge, Bonasa umbellus + 83: prairie chicken, prairie grouse, prairie fowl + 84: peacock + 85: quail + 86: partridge + 87: African grey, African gray, Psittacus erithacus + 88: macaw + 89: sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita + 90: lorikeet + 91: coucal + 92: bee eater + 93: hornbill + 94: hummingbird + 95: jacamar + 96: toucan + 97: drake + 98: red-breasted merganser, Mergus serrator + 99: goose + 100: black swan, Cygnus atratus + 101: tusker + 102: echidna, spiny anteater, anteater + 103: platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus + anatinus + 104: wallaby, brush kangaroo + 105: koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus + 106: wombat + 107: jellyfish + 108: sea anemone, anemone + 109: brain coral + 110: flatworm, platyhelminth + 111: nematode, nematode worm, roundworm + 112: conch + 113: snail + 114: slug + 115: sea slug, nudibranch + 116: chiton, coat-of-mail shell, sea cradle, polyplacophore + 117: chambered nautilus, pearly nautilus, nautilus + 118: Dungeness crab, Cancer magister + 119: rock crab, Cancer irroratus + 120: fiddler crab + 121: king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica + 122: American lobster, Northern lobster, Maine lobster, Homarus americanus + 123: spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish + 124: crayfish, crawfish, crawdad, crawdaddy + 125: hermit crab + 126: isopod + 127: white stork, Ciconia ciconia + 128: black stork, Ciconia nigra + 129: spoonbill + 130: flamingo + 131: little blue heron, Egretta caerulea + 132: American egret, great white heron, Egretta albus + 133: bittern + 134: crane + 135: limpkin, Aramus pictus + 136: European gallinule, Porphyrio porphyrio + 137: American coot, marsh hen, mud hen, water hen, Fulica americana + 138: bustard + 139: ruddy turnstone, Arenaria interpres + 140: red-backed sandpiper, dunlin, Erolia alpina + 141: redshank, Tringa totanus + 142: dowitcher + 143: oystercatcher, oyster catcher + 144: pelican + 145: king penguin, Aptenodytes patagonica + 146: albatross, mollymawk + 147: grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus + 148: killer whale, killer, orca, grampus, sea wolf, Orcinus orca + 149: dugong, Dugong dugon + 150: sea lion + 151: Chihuahua + 152: Japanese spaniel + 153: Maltese dog, Maltese terrier, Maltese + 154: Pekinese, Pekingese, Peke + 155: Shih-Tzu + 156: Blenheim spaniel + 157: papillon + 158: toy terrier + 159: Rhodesian ridgeback + 160: Afghan hound, Afghan + 161: basset, basset hound + 162: beagle + 163: bloodhound, sleuthhound + 164: bluetick + 165: black-and-tan coonhound + 166: Walker hound, Walker foxhound + 167: English foxhound + 168: redbone + 169: borzoi, Russian wolfhound + 170: Irish wolfhound + 171: Italian greyhound + 172: whippet + 173: Ibizan hound, Ibizan Podenco + 174: Norwegian elkhound, elkhound + 175: otterhound, otter hound + 176: Saluki, gazelle hound + 177: Scottish deerhound, deerhound + 178: Weimaraner + 179: Staffordshire bullterrier, Staffordshire bull terrier + 180: American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, + pit bull terrier + 181: Bedlington terrier + 182: Border terrier + 183: Kerry blue terrier + 184: Irish terrier + 185: Norfolk terrier + 186: Norwich terrier + 187: Yorkshire terrier + 188: wire-haired fox terrier + 189: Lakeland terrier + 190: Sealyham terrier, Sealyham + 191: Airedale, Airedale terrier + 192: cairn, cairn terrier + 193: Australian terrier + 194: Dandie Dinmont, Dandie Dinmont terrier + 195: Boston bull, Boston terrier + 196: miniature schnauzer + 197: giant schnauzer + 198: standard schnauzer + 199: Scotch terrier, Scottish terrier, Scottie + 200: Tibetan terrier, chrysanthemum dog + 201: silky terrier, Sydney silky + 202: soft-coated wheaten terrier + 203: West Highland white terrier + 204: Lhasa, Lhasa apso + 205: flat-coated retriever + 206: curly-coated retriever + 207: golden retriever + 208: Labrador retriever + 209: Chesapeake Bay retriever + 210: German short-haired pointer + 211: vizsla, Hungarian pointer + 212: English setter + 213: Irish setter, red setter + 214: Gordon setter + 215: Brittany spaniel + 216: clumber, clumber spaniel + 217: English springer, English springer spaniel + 218: Welsh springer spaniel + 219: cocker spaniel, English cocker spaniel, cocker + 220: Sussex spaniel + 221: Irish water spaniel + 222: kuvasz + 223: schipperke + 224: groenendael + 225: malinois + 226: briard + 227: kelpie + 228: komondor + 229: Old English sheepdog, bobtail + 230: Shetland sheepdog, Shetland sheep dog, Shetland + 231: collie + 232: Border collie + 233: Bouvier des Flandres, Bouviers des Flandres + 234: Rottweiler + 235: German shepherd, German shepherd dog, German police dog, alsatian + 236: Doberman, Doberman pinscher + 237: miniature pinscher + 238: Greater Swiss Mountain dog + 239: Bernese mountain dog + 240: Appenzeller + 241: EntleBucher + 242: boxer + 243: bull mastiff + 244: Tibetan mastiff + 245: French bulldog + 246: Great Dane + 247: Saint Bernard, St Bernard + 248: Eskimo dog, husky + 249: malamute, malemute, Alaskan malamute + 250: Siberian husky + 251: dalmatian, coach dog, carriage dog + 252: affenpinscher, monkey pinscher, monkey dog + 253: basenji + 254: pug, pug-dog + 255: Leonberg + 256: Newfoundland, Newfoundland dog + 257: Great Pyrenees + 258: Samoyed, Samoyede + 259: Pomeranian + 260: chow, chow chow + 261: keeshond + 262: Brabancon griffon + 263: Pembroke, Pembroke Welsh corgi + 264: Cardigan, Cardigan Welsh corgi + 265: toy poodle + 266: miniature poodle + 267: standard poodle + 268: Mexican hairless + 269: timber wolf, grey wolf, gray wolf, Canis lupus + 270: white wolf, Arctic wolf, Canis lupus tundrarum + 271: red wolf, maned wolf, Canis rufus, Canis niger + 272: coyote, prairie wolf, brush wolf, Canis latrans + 273: dingo, warrigal, warragal, Canis dingo + 274: dhole, Cuon alpinus + 275: African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus + 276: hyena, hyaena + 277: red fox, Vulpes vulpes + 278: kit fox, Vulpes macrotis + 279: Arctic fox, white fox, Alopex lagopus + 280: grey fox, gray fox, Urocyon cinereoargenteus + 281: tabby, tabby cat + 282: tiger cat + 283: Persian cat + 284: Siamese cat, Siamese + 285: Egyptian cat + 286: cougar, puma, catamount, mountain lion, painter, panther, Felis concolor + 287: lynx, catamount + 288: leopard, Panthera pardus + 289: snow leopard, ounce, Panthera uncia + 290: jaguar, panther, Panthera onca, Felis onca + 291: lion, king of beasts, Panthera leo + 292: tiger, Panthera tigris + 293: cheetah, chetah, Acinonyx jubatus + 294: brown bear, bruin, Ursus arctos + 295: American black bear, black bear, Ursus americanus, Euarctos americanus + 296: ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus + 297: sloth bear, Melursus ursinus, Ursus ursinus + 298: mongoose + 299: meerkat, mierkat + 300: tiger beetle + 301: ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle + 302: ground beetle, carabid beetle + 303: long-horned beetle, longicorn, longicorn beetle + 304: leaf beetle, chrysomelid + 305: dung beetle + 306: rhinoceros beetle + 307: weevil + 308: fly + 309: bee + 310: ant, emmet, pismire + 311: grasshopper, hopper + 312: cricket + 313: walking stick, walkingstick, stick insect + 314: cockroach, roach + 315: mantis, mantid + 316: cicada, cicala + 317: leafhopper + 318: lacewing, lacewing fly + 319: dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, + snake doctor, mosquito hawk, skeeter hawk + 320: damselfly + 321: admiral + 322: ringlet, ringlet butterfly + 323: monarch, monarch butterfly, milkweed butterfly, Danaus plexippus + 324: cabbage butterfly + 325: sulphur butterfly, sulfur butterfly + 326: lycaenid, lycaenid butterfly + 327: starfish, sea star + 328: sea urchin + 329: sea cucumber, holothurian + 330: wood rabbit, cottontail, cottontail rabbit + 331: hare + 332: Angora, Angora rabbit + 333: hamster + 334: porcupine, hedgehog + 335: fox squirrel, eastern fox squirrel, Sciurus niger + 336: marmot + 337: beaver + 338: guinea pig, Cavia cobaya + 339: sorrel + 340: zebra + 341: hog, pig, grunter, squealer, Sus scrofa + 342: wild boar, boar, Sus scrofa + 343: warthog + 344: hippopotamus, hippo, river horse, Hippopotamus amphibius + 345: ox + 346: water buffalo, water ox, Asiatic buffalo, Bubalus bubalis + 347: bison + 348: ram, tup + 349: bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, + Ovis canadensis + 350: ibex, Capra ibex + 351: hartebeest + 352: impala, Aepyceros melampus + 353: gazelle + 354: Arabian camel, dromedary, Camelus dromedarius + 355: llama + 356: weasel + 357: mink + 358: polecat, fitch, foulmart, foumart, Mustela putorius + 359: black-footed ferret, ferret, Mustela nigripes + 360: otter + 361: skunk, polecat, wood pussy + 362: badger + 363: armadillo + 364: three-toed sloth, ai, Bradypus tridactylus + 365: orangutan, orang, orangutang, Pongo pygmaeus + 366: gorilla, Gorilla gorilla + 367: chimpanzee, chimp, Pan troglodytes + 368: gibbon, Hylobates lar + 369: siamang, Hylobates syndactylus, Symphalangus syndactylus + 370: guenon, guenon monkey + 371: patas, hussar monkey, Erythrocebus patas + 372: baboon + 373: macaque + 374: langur + 375: colobus, colobus monkey + 376: proboscis monkey, Nasalis larvatus + 377: marmoset + 378: capuchin, ringtail, Cebus capucinus + 379: howler monkey, howler + 380: titi, titi monkey + 381: spider monkey, Ateles geoffroyi + 382: squirrel monkey, Saimiri sciureus + 383: Madagascar cat, ring-tailed lemur, Lemur catta + 384: indri, indris, Indri indri, Indri brevicaudatus + 385: Indian elephant, Elephas maximus + 386: African elephant, Loxodonta africana + 387: lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens + 388: giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca + 389: barracouta, snoek + 390: eel + 391: coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch + 392: rock beauty, Holocanthus tricolor + 393: anemone fish + 394: sturgeon + 395: gar, garfish, garpike, billfish, Lepisosteus osseus + 396: lionfish + 397: puffer, pufferfish, blowfish, globefish + 398: abacus + 399: abaya + 400: academic gown, academic robe, judge's robe + 401: accordion, piano accordion, squeeze box + 402: acoustic guitar + 403: aircraft carrier, carrier, flattop, attack aircraft carrier + 404: airliner + 405: airship, dirigible + 406: altar + 407: ambulance + 408: amphibian, amphibious vehicle + 409: analog clock + 410: apiary, bee house + 411: apron + 412: ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, + trash barrel, trash bin + 413: assault rifle, assault gun + 414: backpack, back pack, knapsack, packsack, rucksack, haversack + 415: bakery, bakeshop, bakehouse + 416: balance beam, beam + 417: balloon + 418: ballpoint, ballpoint pen, ballpen, Biro + 419: Band Aid + 420: banjo + 421: bannister, banister, balustrade, balusters, handrail + 422: barbell + 423: barber chair + 424: barbershop + 425: barn + 426: barometer + 427: barrel, cask + 428: barrow, garden cart, lawn cart, wheelbarrow + 429: baseball + 430: basketball + 431: bassinet + 432: bassoon + 433: bathing cap, swimming cap + 434: bath towel + 435: bathtub, bathing tub, bath, tub + 436: beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, + waggon + 437: beacon, lighthouse, beacon light, pharos + 438: beaker + 439: bearskin, busby, shako + 440: beer bottle + 441: beer glass + 442: bell cote, bell cot + 443: bib + 444: bicycle-built-for-two, tandem bicycle, tandem + 445: bikini, two-piece + 446: binder, ring-binder + 447: binoculars, field glasses, opera glasses + 448: birdhouse + 449: boathouse + 450: bobsled, bobsleigh, bob + 451: bolo tie, bolo, bola tie, bola + 452: bonnet, poke bonnet + 453: bookcase + 454: bookshop, bookstore, bookstall + 455: bottlecap + 456: bow + 457: bow tie, bow-tie, bowtie + 458: brass, memorial tablet, plaque + 459: brassiere, bra, bandeau + 460: breakwater, groin, groyne, mole, bulwark, seawall, jetty + 461: breastplate, aegis, egis + 462: broom + 463: bucket, pail + 464: buckle + 465: bulletproof vest + 466: bullet train, bullet + 467: butcher shop, meat market + 468: cab, hack, taxi, taxicab + 469: caldron, cauldron + 470: candle, taper, wax light + 471: cannon + 472: canoe + 473: can opener, tin opener + 474: cardigan + 475: car mirror + 476: carousel, carrousel, merry-go-round, roundabout, whirligig + 477: carpenter's kit, tool kit + 478: carton + 479: car wheel + 480: cash machine, cash dispenser, automated teller machine, automatic teller machine, + automated teller, automatic teller, ATM + 481: cassette + 482: cassette player + 483: castle + 484: catamaran + 485: CD player + 486: cello, violoncello + 487: cellular telephone, cellular phone, cellphone, cell, mobile phone + 488: chain + 489: chainlink fence + 490: chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour + 491: chain saw, chainsaw + 492: chest + 493: chiffonier, commode + 494: chime, bell, gong + 495: china cabinet, china closet + 496: Christmas stocking + 497: church, church building + 498: cinema, movie theater, movie theatre, movie house, picture palace + 499: cleaver, meat cleaver, chopper + 500: cliff dwelling + 501: cloak + 502: clog, geta, patten, sabot + 503: cocktail shaker + 504: coffee mug + 505: coffeepot + 506: coil, spiral, volute, whorl, helix + 507: combination lock + 508: computer keyboard, keypad + 509: confectionery, confectionary, candy store + 510: container ship, containership, container vessel + 511: convertible + 512: corkscrew, bottle screw + 513: cornet, horn, trumpet, trump + 514: cowboy boot + 515: cowboy hat, ten-gallon hat + 516: cradle + 517: crane + 518: crash helmet + 519: crate + 520: crib, cot + 521: Crock Pot + 522: croquet ball + 523: crutch + 524: cuirass + 525: dam, dike, dyke + 526: desk + 527: desktop computer + 528: dial telephone, dial phone + 529: diaper, nappy, napkin + 530: digital clock + 531: digital watch + 532: dining table, board + 533: dishrag, dishcloth + 534: dishwasher, dish washer, dishwashing machine + 535: disk brake, disc brake + 536: dock, dockage, docking facility + 537: dogsled, dog sled, dog sleigh + 538: dome + 539: doormat, welcome mat + 540: drilling platform, offshore rig + 541: drum, membranophone, tympan + 542: drumstick + 543: dumbbell + 544: Dutch oven + 545: electric fan, blower + 546: electric guitar + 547: electric locomotive + 548: entertainment center + 549: envelope + 550: espresso maker + 551: face powder + 552: feather boa, boa + 553: file, file cabinet, filing cabinet + 554: fireboat + 555: fire engine, fire truck + 556: fire screen, fireguard + 557: flagpole, flagstaff + 558: flute, transverse flute + 559: folding chair + 560: football helmet + 561: forklift + 562: fountain + 563: fountain pen + 564: four-poster + 565: freight car + 566: French horn, horn + 567: frying pan, frypan, skillet + 568: fur coat + 569: garbage truck, dustcart + 570: gasmask, respirator, gas helmet + 571: gas pump, gasoline pump, petrol pump, island dispenser + 572: goblet + 573: go-kart + 574: golf ball + 575: golfcart, golf cart + 576: gondola + 577: gong, tam-tam + 578: gown + 579: grand piano, grand + 580: greenhouse, nursery, glasshouse + 581: grille, radiator grille + 582: grocery store, grocery, food market, market + 583: guillotine + 584: hair slide + 585: hair spray + 586: half track + 587: hammer + 588: hamper + 589: hand blower, blow dryer, blow drier, hair dryer, hair drier + 590: hand-held computer, hand-held microcomputer + 591: handkerchief, hankie, hanky, hankey + 592: hard disc, hard disk, fixed disk + 593: harmonica, mouth organ, harp, mouth harp + 594: harp + 595: harvester, reaper + 596: hatchet + 597: holster + 598: home theater, home theatre + 599: honeycomb + 600: hook, claw + 601: hoopskirt, crinoline + 602: horizontal bar, high bar + 603: horse cart, horse-cart + 604: hourglass + 605: iPod + 606: iron, smoothing iron + 607: jack-o'-lantern + 608: jean, blue jean, denim + 609: jeep, landrover + 610: jersey, T-shirt, tee shirt + 611: jigsaw puzzle + 612: jinrikisha, ricksha, rickshaw + 613: joystick + 614: kimono + 615: knee pad + 616: knot + 617: lab coat, laboratory coat + 618: ladle + 619: lampshade, lamp shade + 620: laptop, laptop computer + 621: lawn mower, mower + 622: lens cap, lens cover + 623: letter opener, paper knife, paperknife + 624: library + 625: lifeboat + 626: lighter, light, igniter, ignitor + 627: limousine, limo + 628: liner, ocean liner + 629: lipstick, lip rouge + 630: Loafer + 631: lotion + 632: loudspeaker, speaker, speaker unit, loudspeaker system, speaker system + 633: loupe, jeweler's loupe + 634: lumbermill, sawmill + 635: magnetic compass + 636: mailbag, postbag + 637: mailbox, letter box + 638: maillot + 639: maillot, tank suit + 640: manhole cover + 641: maraca + 642: marimba, xylophone + 643: mask + 644: matchstick + 645: maypole + 646: maze, labyrinth + 647: measuring cup + 648: medicine chest, medicine cabinet + 649: megalith, megalithic structure + 650: microphone, mike + 651: microwave, microwave oven + 652: military uniform + 653: milk can + 654: minibus + 655: miniskirt, mini + 656: minivan + 657: missile + 658: mitten + 659: mixing bowl + 660: mobile home, manufactured home + 661: Model T + 662: modem + 663: monastery + 664: monitor + 665: moped + 666: mortar + 667: mortarboard + 668: mosque + 669: mosquito net + 670: motor scooter, scooter + 671: mountain bike, all-terrain bike, off-roader + 672: mountain tent + 673: mouse, computer mouse + 674: mousetrap + 675: moving van + 676: muzzle + 677: nail + 678: neck brace + 679: necklace + 680: nipple + 681: notebook, notebook computer + 682: obelisk + 683: oboe, hautboy, hautbois + 684: ocarina, sweet potato + 685: odometer, hodometer, mileometer, milometer + 686: oil filter + 687: organ, pipe organ + 688: oscilloscope, scope, cathode-ray oscilloscope, CRO + 689: overskirt + 690: oxcart + 691: oxygen mask + 692: packet + 693: paddle, boat paddle + 694: paddlewheel, paddle wheel + 695: padlock + 696: paintbrush + 697: pajama, pyjama, pj's, jammies + 698: palace + 699: panpipe, pandean pipe, syrinx + 700: paper towel + 701: parachute, chute + 702: parallel bars, bars + 703: park bench + 704: parking meter + 705: passenger car, coach, carriage + 706: patio, terrace + 707: pay-phone, pay-station + 708: pedestal, plinth, footstall + 709: pencil box, pencil case + 710: pencil sharpener + 711: perfume, essence + 712: Petri dish + 713: photocopier + 714: pick, plectrum, plectron + 715: pickelhaube + 716: picket fence, paling + 717: pickup, pickup truck + 718: pier + 719: piggy bank, penny bank + 720: pill bottle + 721: pillow + 722: ping-pong ball + 723: pinwheel + 724: pirate, pirate ship + 725: pitcher, ewer + 726: plane, carpenter's plane, woodworking plane + 727: planetarium + 728: plastic bag + 729: plate rack + 730: plow, plough + 731: plunger, plumber's helper + 732: Polaroid camera, Polaroid Land camera + 733: pole + 734: police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria + 735: poncho + 736: pool table, billiard table, snooker table + 737: pop bottle, soda bottle + 738: pot, flowerpot + 739: potter's wheel + 740: power drill + 741: prayer rug, prayer mat + 742: printer + 743: prison, prison house + 744: projectile, missile + 745: projector + 746: puck, hockey puck + 747: punching bag, punch bag, punching ball, punchball + 748: purse + 749: quill, quill pen + 750: quilt, comforter, comfort, puff + 751: racer, race car, racing car + 752: racket, racquet + 753: radiator + 754: radio, wireless + 755: radio telescope, radio reflector + 756: rain barrel + 757: recreational vehicle, RV, R.V. + 758: reel + 759: reflex camera + 760: refrigerator, icebox + 761: remote control, remote + 762: restaurant, eating house, eating place, eatery + 763: revolver, six-gun, six-shooter + 764: rifle + 765: rocking chair, rocker + 766: rotisserie + 767: rubber eraser, rubber, pencil eraser + 768: rugby ball + 769: rule, ruler + 770: running shoe + 771: safe + 772: safety pin + 773: saltshaker, salt shaker + 774: sandal + 775: sarong + 776: sax, saxophone + 777: scabbard + 778: scale, weighing machine + 779: school bus + 780: schooner + 781: scoreboard + 782: screen, CRT screen + 783: screw + 784: screwdriver + 785: seat belt, seatbelt + 786: sewing machine + 787: shield, buckler + 788: shoe shop, shoe-shop, shoe store + 789: shoji + 790: shopping basket + 791: shopping cart + 792: shovel + 793: shower cap + 794: shower curtain + 795: ski + 796: ski mask + 797: sleeping bag + 798: slide rule, slipstick + 799: sliding door + 800: slot, one-armed bandit + 801: snorkel + 802: snowmobile + 803: snowplow, snowplough + 804: soap dispenser + 805: soccer ball + 806: sock + 807: solar dish, solar collector, solar furnace + 808: sombrero + 809: soup bowl + 810: space bar + 811: space heater + 812: space shuttle + 813: spatula + 814: speedboat + 815: spider web, spider's web + 816: spindle + 817: sports car, sport car + 818: spotlight, spot + 819: stage + 820: steam locomotive + 821: steel arch bridge + 822: steel drum + 823: stethoscope + 824: stole + 825: stone wall + 826: stopwatch, stop watch + 827: stove + 828: strainer + 829: streetcar, tram, tramcar, trolley, trolley car + 830: stretcher + 831: studio couch, day bed + 832: stupa, tope + 833: submarine, pigboat, sub, U-boat + 834: suit, suit of clothes + 835: sundial + 836: sunglass + 837: sunglasses, dark glasses, shades + 838: sunscreen, sunblock, sun blocker + 839: suspension bridge + 840: swab, swob, mop + 841: sweatshirt + 842: swimming trunks, bathing trunks + 843: swing + 844: switch, electric switch, electrical switch + 845: syringe + 846: table lamp + 847: tank, army tank, armored combat vehicle, armoured combat vehicle + 848: tape player + 849: teapot + 850: teddy, teddy bear + 851: television, television system + 852: tennis ball + 853: thatch, thatched roof + 854: theater curtain, theatre curtain + 855: thimble + 856: thresher, thrasher, threshing machine + 857: throne + 858: tile roof + 859: toaster + 860: tobacco shop, tobacconist shop, tobacconist + 861: toilet seat + 862: torch + 863: totem pole + 864: tow truck, tow car, wrecker + 865: toyshop + 866: tractor + 867: trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi + 868: tray + 869: trench coat + 870: tricycle, trike, velocipede + 871: trimaran + 872: tripod + 873: triumphal arch + 874: trolleybus, trolley coach, trackless trolley + 875: trombone + 876: tub, vat + 877: turnstile + 878: typewriter keyboard + 879: umbrella + 880: unicycle, monocycle + 881: upright, upright piano + 882: vacuum, vacuum cleaner + 883: vase + 884: vault + 885: velvet + 886: vending machine + 887: vestment + 888: viaduct + 889: violin, fiddle + 890: volleyball + 891: waffle iron + 892: wall clock + 893: wallet, billfold, notecase, pocketbook + 894: wardrobe, closet, press + 895: warplane, military plane + 896: washbasin, handbasin, washbowl, lavabo, wash-hand basin + 897: washer, automatic washer, washing machine + 898: water bottle + 899: water jug + 900: water tower + 901: whiskey jug + 902: whistle + 903: wig + 904: window screen + 905: window shade + 906: Windsor tie + 907: wine bottle + 908: wing + 909: wok + 910: wooden spoon + 911: wool, woolen, woollen + 912: worm fence, snake fence, snake-rail fence, Virginia fence + 913: wreck + 914: yawl + 915: yurt + 916: web site, website, internet site, site + 917: comic book + 918: crossword puzzle, crossword + 919: street sign + 920: traffic light, traffic signal, stoplight + 921: book jacket, dust cover, dust jacket, dust wrapper + 922: menu + 923: plate + 924: guacamole + 925: consomme + 926: hot pot, hotpot + 927: trifle + 928: ice cream, icecream + 929: ice lolly, lolly, lollipop, popsicle + 930: French loaf + 931: bagel, beigel + 932: pretzel + 933: cheeseburger + 934: hotdog, hot dog, red hot + 935: mashed potato + 936: head cabbage + 937: broccoli + 938: cauliflower + 939: zucchini, courgette + 940: spaghetti squash + 941: acorn squash + 942: butternut squash + 943: cucumber, cuke + 944: artichoke, globe artichoke + 945: bell pepper + 946: cardoon + 947: mushroom + 948: Granny Smith + 949: strawberry + 950: orange + 951: lemon + 952: fig + 953: pineapple, ananas + 954: banana + 955: jackfruit, jak, jack + 956: custard apple + 957: pomegranate + 958: hay + 959: carbonara + 960: chocolate sauce, chocolate syrup + 961: dough + 962: meat loaf, meatloaf + 963: pizza, pizza pie + 964: potpie + 965: burrito + 966: red wine + 967: espresso + 968: cup + 969: eggnog + 970: alp + 971: bubble + 972: cliff, drop, drop-off + 973: coral reef + 974: geyser + 975: lakeside, lakeshore + 976: promontory, headland, head, foreland + 977: sandbar, sand bar + 978: seashore, coast, seacoast, sea-coast + 979: valley, vale + 980: volcano + 981: ballplayer, baseball player + 982: groom, bridegroom + 983: scuba diver + 984: rapeseed + 985: daisy + 986: yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium + parviflorum + 987: corn + 988: acorn + 989: hip, rose hip, rosehip + 990: buckeye, horse chestnut, conker + 991: coral fungus + 992: agaric + 993: gyromitra + 994: stinkhorn, carrion fungus + 995: earthstar + 996: hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa + 997: bolete + 998: ear, spike, capitulum + 999: toilet tissue, toilet paper, bathroom tissue diff --git a/mblt_vision/datasets/imagenet_synsets.txt b/mblt_vision/datasets/imagenet_synsets.txt new file mode 100644 index 0000000..88aa58f --- /dev/null +++ b/mblt_vision/datasets/imagenet_synsets.txt @@ -0,0 +1,1000 @@ +n01440764 +n01443537 +n01484850 +n01491361 +n01494475 +n01496331 +n01498041 +n01514668 +n01514859 +n01518878 +n01530575 +n01531178 +n01532829 +n01534433 +n01537544 +n01558993 +n01560419 +n01580077 +n01582220 +n01592084 +n01601694 +n01608432 +n01614925 +n01616318 +n01622779 +n01629819 +n01630670 +n01631663 +n01632458 +n01632777 +n01641577 +n01644373 +n01644900 +n01664065 +n01665541 +n01667114 +n01667778 +n01669191 +n01675722 +n01677366 +n01682714 +n01685808 +n01687978 +n01688243 +n01689811 +n01692333 +n01693334 +n01694178 +n01695060 +n01697457 +n01698640 +n01704323 +n01728572 +n01728920 +n01729322 +n01729977 +n01734418 +n01735189 +n01737021 +n01739381 +n01740131 +n01742172 +n01744401 +n01748264 +n01749939 +n01751748 +n01753488 +n01755581 +n01756291 +n01768244 +n01770081 +n01770393 +n01773157 +n01773549 +n01773797 +n01774384 +n01774750 +n01775062 +n01776313 +n01784675 +n01795545 +n01796340 +n01797886 +n01798484 +n01806143 +n01806567 +n01807496 +n01817953 +n01818515 +n01819313 +n01820546 +n01824575 +n01828970 +n01829413 +n01833805 +n01843065 +n01843383 +n01847000 +n01855032 +n01855672 +n01860187 +n01871265 +n01872401 +n01873310 +n01877812 +n01882714 +n01883070 +n01910747 +n01914609 +n01917289 +n01924916 +n01930112 +n01943899 +n01944390 +n01945685 +n01950731 +n01955084 +n01968897 +n01978287 +n01978455 +n01980166 +n01981276 +n01983481 +n01984695 +n01985128 +n01986214 +n01990800 +n02002556 +n02002724 +n02006656 +n02007558 +n02009229 +n02009912 +n02011460 +n02012849 +n02013706 +n02017213 +n02018207 +n02018795 +n02025239 +n02027492 +n02028035 +n02033041 +n02037110 +n02051845 +n02056570 +n02058221 +n02066245 +n02071294 +n02074367 +n02077923 +n02085620 +n02085782 +n02085936 +n02086079 +n02086240 +n02086646 +n02086910 +n02087046 +n02087394 +n02088094 +n02088238 +n02088364 +n02088466 +n02088632 +n02089078 +n02089867 +n02089973 +n02090379 +n02090622 +n02090721 +n02091032 +n02091134 +n02091244 +n02091467 +n02091635 +n02091831 +n02092002 +n02092339 +n02093256 +n02093428 +n02093647 +n02093754 +n02093859 +n02093991 +n02094114 +n02094258 +n02094433 +n02095314 +n02095570 +n02095889 +n02096051 +n02096177 +n02096294 +n02096437 +n02096585 +n02097047 +n02097130 +n02097209 +n02097298 +n02097474 +n02097658 +n02098105 +n02098286 +n02098413 +n02099267 +n02099429 +n02099601 +n02099712 +n02099849 +n02100236 +n02100583 +n02100735 +n02100877 +n02101006 +n02101388 +n02101556 +n02102040 +n02102177 +n02102318 +n02102480 +n02102973 +n02104029 +n02104365 +n02105056 +n02105162 +n02105251 +n02105412 +n02105505 +n02105641 +n02105855 +n02106030 +n02106166 +n02106382 +n02106550 +n02106662 +n02107142 +n02107312 +n02107574 +n02107683 +n02107908 +n02108000 +n02108089 +n02108422 +n02108551 +n02108915 +n02109047 +n02109525 +n02109961 +n02110063 +n02110185 +n02110341 +n02110627 +n02110806 +n02110958 +n02111129 +n02111277 +n02111500 +n02111889 +n02112018 +n02112137 +n02112350 +n02112706 +n02113023 +n02113186 +n02113624 +n02113712 +n02113799 +n02113978 +n02114367 +n02114548 +n02114712 +n02114855 +n02115641 +n02115913 +n02116738 +n02117135 +n02119022 +n02119789 +n02120079 +n02120505 +n02123045 +n02123159 +n02123394 +n02123597 +n02124075 +n02125311 +n02127052 +n02128385 +n02128757 +n02128925 +n02129165 +n02129604 +n02130308 +n02132136 +n02133161 +n02134084 +n02134418 +n02137549 +n02138441 +n02165105 +n02165456 +n02167151 +n02168699 +n02169497 +n02172182 +n02174001 +n02177972 +n02190166 +n02206856 +n02219486 +n02226429 +n02229544 +n02231487 +n02233338 +n02236044 +n02256656 +n02259212 +n02264363 +n02268443 +n02268853 +n02276258 +n02277742 +n02279972 +n02280649 +n02281406 +n02281787 +n02317335 +n02319095 +n02321529 +n02325366 +n02326432 +n02328150 +n02342885 +n02346627 +n02356798 +n02361337 +n02363005 +n02364673 +n02389026 +n02391049 +n02395406 +n02396427 +n02397096 +n02398521 +n02403003 +n02408429 +n02410509 +n02412080 +n02415577 +n02417914 +n02422106 +n02422699 +n02423022 +n02437312 +n02437616 +n02441942 +n02442845 +n02443114 +n02443484 +n02444819 +n02445715 +n02447366 +n02454379 +n02457408 +n02480495 +n02480855 +n02481823 +n02483362 +n02483708 +n02484975 +n02486261 +n02486410 +n02487347 +n02488291 +n02488702 +n02489166 +n02490219 +n02492035 +n02492660 +n02493509 +n02493793 +n02494079 +n02497673 +n02500267 +n02504013 +n02504458 +n02509815 +n02510455 +n02514041 +n02526121 +n02536864 +n02606052 +n02607072 +n02640242 +n02641379 +n02643566 +n02655020 +n02666196 +n02667093 +n02669723 +n02672831 +n02676566 +n02687172 +n02690373 +n02692877 +n02699494 +n02701002 +n02704792 +n02708093 +n02727426 +n02730930 +n02747177 +n02749479 +n02769748 +n02776631 +n02777292 +n02782093 +n02783161 +n02786058 +n02787622 +n02788148 +n02790996 +n02791124 +n02791270 +n02793495 +n02794156 +n02795169 +n02797295 +n02799071 +n02802426 +n02804414 +n02804610 +n02807133 +n02808304 +n02808440 +n02814533 +n02814860 +n02815834 +n02817516 +n02823428 +n02823750 +n02825657 +n02834397 +n02835271 +n02837789 +n02840245 +n02841315 +n02843684 +n02859443 +n02860847 +n02865351 +n02869837 +n02870880 +n02871525 +n02877765 +n02879718 +n02883205 +n02892201 +n02892767 +n02894605 +n02895154 +n02906734 +n02909870 +n02910353 +n02916936 +n02917067 +n02927161 +n02930766 +n02939185 +n02948072 +n02950826 +n02951358 +n02951585 +n02963159 +n02965783 +n02966193 +n02966687 +n02971356 +n02974003 +n02977058 +n02978881 +n02979186 +n02980441 +n02981792 +n02988304 +n02992211 +n02992529 +n02999410 +n03000134 +n03000247 +n03000684 +n03014705 +n03016953 +n03017168 +n03018349 +n03026506 +n03028079 +n03032252 +n03041632 +n03042490 +n03045698 +n03047690 +n03062245 +n03063599 +n03063689 +n03065424 +n03075370 +n03085013 +n03089624 +n03095699 +n03100240 +n03109150 +n03110669 +n03124043 +n03124170 +n03125729 +n03126707 +n03127747 +n03127925 +n03131574 +n03133878 +n03134739 +n03141823 +n03146219 +n03160309 +n03179701 +n03180011 +n03187595 +n03188531 +n03196217 +n03197337 +n03201208 +n03207743 +n03207941 +n03208938 +n03216828 +n03218198 +n03220513 +n03223299 +n03240683 +n03249569 +n03250847 +n03255030 +n03259280 +n03271574 +n03272010 +n03272562 +n03290653 +n03291819 +n03297495 +n03314780 +n03325584 +n03337140 +n03344393 +n03345487 +n03347037 +n03355925 +n03372029 +n03376595 +n03379051 +n03384352 +n03388043 +n03388183 +n03388549 +n03393912 +n03394916 +n03400231 +n03404251 +n03417042 +n03424325 +n03425413 +n03443371 +n03444034 +n03445777 +n03445924 +n03447447 +n03447721 +n03450230 +n03452741 +n03457902 +n03459775 +n03461385 +n03467068 +n03476684 +n03476991 +n03478589 +n03481172 +n03482405 +n03483316 +n03485407 +n03485794 +n03492542 +n03494278 +n03495258 +n03496892 +n03498962 +n03527444 +n03529860 +n03530642 +n03532672 +n03534580 +n03535780 +n03538406 +n03544143 +n03584254 +n03584829 +n03590841 +n03594734 +n03594945 +n03595614 +n03598930 +n03599486 +n03602883 +n03617480 +n03623198 +n03627232 +n03630383 +n03633091 +n03637318 +n03642806 +n03649909 +n03657121 +n03658185 +n03661043 +n03662601 +n03666591 +n03670208 +n03673027 +n03676483 +n03680355 +n03690938 +n03691459 +n03692522 +n03697007 +n03706229 +n03709823 +n03710193 +n03710637 +n03710721 +n03717622 +n03720891 +n03721384 +n03724870 +n03729826 +n03733131 +n03733281 +n03733805 +n03742115 +n03743016 +n03759954 +n03761084 +n03763968 +n03764736 +n03769881 +n03770439 +n03770679 +n03773504 +n03775071 +n03775546 +n03776460 +n03777568 +n03777754 +n03781244 +n03782006 +n03785016 +n03786901 +n03787032 +n03788195 +n03788365 +n03791053 +n03792782 +n03792972 +n03793489 +n03794056 +n03796401 +n03803284 +n03804744 +n03814639 +n03814906 +n03825788 +n03832673 +n03837869 +n03838899 +n03840681 +n03841143 +n03843555 +n03854065 +n03857828 +n03866082 +n03868242 +n03868863 +n03871628 +n03873416 +n03874293 +n03874599 +n03876231 +n03877472 +n03877845 +n03884397 +n03887697 +n03888257 +n03888605 +n03891251 +n03891332 +n03895866 +n03899768 +n03902125 +n03903868 +n03908618 +n03908714 +n03916031 +n03920288 +n03924679 +n03929660 +n03929855 +n03930313 +n03930630 +n03933933 +n03935335 +n03937543 +n03938244 +n03942813 +n03944341 +n03947888 +n03950228 +n03954731 +n03956157 +n03958227 +n03961711 +n03967562 +n03970156 +n03976467 +n03976657 +n03977966 +n03980874 +n03982430 +n03983396 +n03991062 +n03992509 +n03995372 +n03998194 +n04004767 +n04005630 +n04008634 +n04009552 +n04019541 +n04023962 +n04026417 +n04033901 +n04033995 +n04037443 +n04039381 +n04040759 +n04041544 +n04044716 +n04049303 +n04065272 +n04067472 +n04069434 +n04070727 +n04074963 +n04081281 +n04086273 +n04090263 +n04099969 +n04111531 +n04116512 +n04118538 +n04118776 +n04120489 +n04125021 +n04127249 +n04131690 +n04133789 +n04136333 +n04141076 +n04141327 +n04141975 +n04146614 +n04147183 +n04149813 +n04152593 +n04153751 +n04154565 +n04162706 +n04179913 +n04192698 +n04200800 +n04201297 +n04204238 +n04204347 +n04208210 +n04209133 +n04209239 +n04228054 +n04229816 +n04235860 +n04238763 +n04239074 +n04243546 +n04251144 +n04252077 +n04252225 +n04254120 +n04254680 +n04254777 +n04258138 +n04259630 +n04263257 +n04264628 +n04265275 +n04266014 +n04270147 +n04273569 +n04275548 +n04277352 +n04285008 +n04286575 +n04296562 +n04310018 +n04311004 +n04311174 +n04317175 +n04325704 +n04326547 +n04328186 +n04330267 +n04332243 +n04335435 +n04336792 +n04344873 +n04346328 +n04347754 +n04350905 +n04355338 +n04355933 +n04356056 +n04357314 +n04366367 +n04367480 +n04370456 +n04371430 +n04371774 +n04372370 +n04376876 +n04380533 +n04389033 +n04392985 +n04398044 +n04399382 +n04404412 +n04409515 +n04417672 +n04418357 +n04423845 +n04428191 +n04429376 +n04435653 +n04442312 +n04443257 +n04447861 +n04456115 +n04458633 +n04461696 +n04462240 +n04465501 +n04467665 +n04476259 +n04479046 +n04482393 +n04483307 +n04485082 +n04486054 +n04487081 +n04487394 +n04493381 +n04501370 +n04505470 +n04507155 +n04509417 +n04515003 +n04517823 +n04522168 +n04523525 +n04525038 +n04525305 +n04532106 +n04532670 +n04536866 +n04540053 +n04542943 +n04548280 +n04548362 +n04550184 +n04552348 +n04553703 +n04554684 +n04557648 +n04560804 +n04562935 +n04579145 +n04579432 +n04584207 +n04589890 +n04590129 +n04591157 +n04591713 +n04592741 +n04596742 +n04597913 +n04599235 +n04604644 +n04606251 +n04612504 +n04613696 +n06359193 +n06596364 +n06785654 +n06794110 +n06874185 +n07248320 +n07565083 +n07579787 +n07583066 +n07584110 +n07590611 +n07613480 +n07614500 +n07615774 +n07684084 +n07693725 +n07695742 +n07697313 +n07697537 +n07711569 +n07714571 +n07714990 +n07715103 +n07716358 +n07716906 +n07717410 +n07717556 +n07718472 +n07718747 +n07720875 +n07730033 +n07734744 +n07742313 +n07745940 +n07747607 +n07749582 +n07753113 +n07753275 +n07753592 +n07754684 +n07760859 +n07768694 +n07802026 +n07831146 +n07836838 +n07860988 +n07871810 +n07873807 +n07875152 +n07880968 +n07892512 +n07920052 +n07930864 +n07932039 +n09193705 +n09229709 +n09246464 +n09256479 +n09288635 +n09332890 +n09399592 +n09421951 +n09428293 +n09468604 +n09472597 +n09835506 +n10148035 +n10565667 +n11879895 +n11939491 +n12057211 +n12144580 +n12267677 +n12620546 +n12768682 +n12985857 +n12998815 +n13037406 +n13040303 +n13044778 +n13052670 +n13054560 +n13133613 +n15075141 diff --git a/mblt_vision/datasets/nyu-depth.yaml b/mblt_vision/datasets/nyu-depth.yaml new file mode 100644 index 0000000..43aded5 --- /dev/null +++ b/mblt_vision/datasets/nyu-depth.yaml @@ -0,0 +1,7 @@ +name: nyu-depth +path: ~/.mblt_model_zoo/datasets/nyu-depth +val: images +tasks: + - depth_estimation +download: + url: https://github.com/ultralytics/assets/releases/download/v0.0.0/nyu-depth.zip diff --git a/mblt_vision/datasets/registry.py b/mblt_vision/datasets/registry.py new file mode 100644 index 0000000..887f8b9 --- /dev/null +++ b/mblt_vision/datasets/registry.py @@ -0,0 +1,158 @@ +"""Load YAML-backed vision dataset definitions.""" + +from __future__ import annotations + +import copy +from functools import lru_cache +from pathlib import Path +from typing import Any + +import yaml + +from .._tasks import normalize_vision_task + +DATASET_CONFIG_DIR = Path(__file__).parent + + +@lru_cache(maxsize=None) +def _load_dataset_config(name: str) -> dict[str, Any]: + """Load and validate a dataset definition without mutating its cached value.""" + + config_path = DATASET_CONFIG_DIR / f"{name.lower()}.yaml" + if not config_path.is_file(): + raise FileNotFoundError(f"Vision dataset definition not found: {config_path}") + + with config_path.open(encoding="utf-8") as config_file: + config = yaml.safe_load(config_file) + if not isinstance(config, dict): + raise ValueError(f"Vision dataset definition must be a mapping: {config_path}") + if not isinstance(config.get("path"), str): + raise ValueError( + f"Vision dataset definition requires a string `path`: {config_path}" + ) + if not isinstance(config.get("name"), str): + raise ValueError( + f"Vision dataset definition requires a string `name`: {config_path}" + ) + if not isinstance(config.get("tasks"), list) or not all( + isinstance(task, str) for task in config["tasks"] + ): + raise ValueError( + f"Vision dataset definition requires a string-list `tasks`: {config_path}" + ) + config["path"] = str(Path(config["path"]).expanduser()) + return config + + +def get_dataset_config(name: str) -> dict[str, Any]: + """Load a named vision dataset definition. + + Args: + name: Dataset filename stem, such as ``dotav1``. + + Returns: + Parsed dataset configuration with its path expanded. + + Raises: + FileNotFoundError: If the dataset definition does not exist. + ValueError: If the definition is malformed. + """ + + return copy.deepcopy(_load_dataset_config(name)) + + +def get_dataset_class_names(name: str) -> tuple[str, ...]: + """Return ordered class names from a dataset definition. + + Args: + name: Dataset filename stem, such as ``coco``. + + Returns: + Class names ordered by contiguous zero-based index. + + Raises: + ValueError: If the dataset has no valid contiguous ``names`` mapping. + """ + + names = _load_dataset_config(name).get("names") + if not isinstance(names, dict) or not all( + isinstance(index, int) and isinstance(label, str) + for index, label in names.items() + ): + raise ValueError( + f"Vision dataset definition requires an integer-keyed `names` mapping: {name}" + ) + if set(names) != set(range(len(names))): + raise ValueError( + f"Vision dataset class IDs must be contiguous and zero-based: {name}" + ) + return tuple(names[index] for index in range(len(names))) + + +@lru_cache(maxsize=None) +def get_dataset_category_ids(name: str) -> tuple[int, ...]: + """Return immutable source category IDs from a dataset definition. + + Args: + name: Dataset filename stem, such as ``coco``. + + Returns: + Source category IDs ordered by contiguous model-output index. + + Raises: + ValueError: If the dataset has no integer ``category_ids`` list. + """ + + category_ids = _load_dataset_config(name).get("category_ids") + if not isinstance(category_ids, list) or not all( + isinstance(category_id, int) for category_id in category_ids + ): + raise ValueError( + f"Vision dataset definition requires an integer-list `category_ids`: {name}" + ) + return tuple(category_ids) + + +def get_dataset_config_for_task( + task: str, dataset: str | None = None +) -> dict[str, Any]: + """Return the validation dataset definition associated with a vision task. + + Args: + task: Vision task name from a model postprocess configuration. + dataset: Optional output-taxonomy name. When omitted, the first configured + task match is returned for backward compatibility. + + Returns: + Matching dataset configuration. + + Raises: + ValueError: If no configured dataset supports the task. + """ + + normalized_task = normalize_vision_task(task) + if dataset is not None: + try: + config = get_dataset_config(dataset) + except FileNotFoundError as exc: + raise ValueError( + f"No vision dataset definition exists for taxonomy `{dataset}`." + ) from exc + configured_tasks = { + normalize_vision_task(configured_task) + for configured_task in config["tasks"] + } + if normalized_task not in configured_tasks: + raise ValueError( + f"Vision dataset `{dataset}` does not support task `{task}`." + ) + return config + + for config_path in sorted(DATASET_CONFIG_DIR.glob("*.yaml")): + config = get_dataset_config(config_path.stem) + if normalized_task in { + normalize_vision_task(configured_task) + for configured_task in config["tasks"] + }: + return config + raise ValueError(f"No vision dataset definition supports task `{task}`.") diff --git a/mblt_vision/datasets/widerface.yaml b/mblt_vision/datasets/widerface.yaml new file mode 100644 index 0000000..2724a8b --- /dev/null +++ b/mblt_vision/datasets/widerface.yaml @@ -0,0 +1,9 @@ +# WiderFace validation dataset. Paths are relative to `path` unless absolute. +name: widerface +path: ~/.mblt_model_zoo/datasets/widerface +tasks: + - face_detection +val: images +download: + images: https://huggingface.co/datasets/CUHK-CSE/wider_face/resolve/main/data/WIDER_val.zip + annotations: https://huggingface.co/datasets/CUHK-CSE/wider_face/resolve/main/data/wider_face_split.zip diff --git a/mblt_vision/depth_estimation/__init__.py b/mblt_vision/depth_estimation/__init__.py new file mode 100644 index 0000000..6682234 --- /dev/null +++ b/mblt_vision/depth_estimation/__init__.py @@ -0,0 +1,19 @@ +"""Depth estimation model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO26lDepth", + "YOLO26mDepth", + "YOLO26nDepth", + "YOLO26sDepth", + "YOLO26xDepth", +] + +YOLO26lDepth = create_model_class("YOLO26lDepth", __name__) +YOLO26mDepth = create_model_class("YOLO26mDepth", __name__) +YOLO26nDepth = create_model_class("YOLO26nDepth", __name__) +YOLO26sDepth = create_model_class("YOLO26sDepth", __name__) +YOLO26xDepth = create_model_class("YOLO26xDepth", __name__) diff --git a/mblt_vision/face_detection/__init__.py b/mblt_vision/face_detection/__init__.py new file mode 100644 index 0000000..0e1ca62 --- /dev/null +++ b/mblt_vision/face_detection/__init__.py @@ -0,0 +1,43 @@ +"""Face detection exports for the vision package.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO11l_face", + "YOLO11m_face", + "YOLO11n_face", + "YOLO11s_face", + "YOLO12l_face", + "YOLO12m_face", + "YOLO12n_face", + "YOLO12s_face", + "YOLOv10l_face", + "YOLOv10m_face", + "YOLOv10n_face", + "YOLOv10s_face", + "YOLOv6m_face", + "YOLOv6n_face", + "YOLOv8l_face", + "YOLOv8m_face", + "YOLOv8n_face", +] + +YOLO11l_face = create_model_class("YOLO11l_face", __name__) +YOLO11m_face = create_model_class("YOLO11m_face", __name__) +YOLO11n_face = create_model_class("YOLO11n_face", __name__) +YOLO11s_face = create_model_class("YOLO11s_face", __name__) +YOLO12l_face = create_model_class("YOLO12l_face", __name__) +YOLO12m_face = create_model_class("YOLO12m_face", __name__) +YOLO12n_face = create_model_class("YOLO12n_face", __name__) +YOLO12s_face = create_model_class("YOLO12s_face", __name__) +YOLOv10l_face = create_model_class("YOLOv10l_face", __name__) +YOLOv10m_face = create_model_class("YOLOv10m_face", __name__) +YOLOv10n_face = create_model_class("YOLOv10n_face", __name__) +YOLOv10s_face = create_model_class("YOLOv10s_face", __name__) +YOLOv6m_face = create_model_class("YOLOv6m_face", __name__) +YOLOv6n_face = create_model_class("YOLOv6n_face", __name__) +YOLOv8l_face = create_model_class("YOLOv8l_face", __name__) +YOLOv8m_face = create_model_class("YOLOv8m_face", __name__) +YOLOv8n_face = create_model_class("YOLOv8n_face", __name__) diff --git a/mblt_vision/image_classification/__init__.py b/mblt_vision/image_classification/__init__.py new file mode 100644 index 0000000..3223c59 --- /dev/null +++ b/mblt_vision/image_classification/__init__.py @@ -0,0 +1,311 @@ +"""Image classification model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "AlexNet", + "CAFormer_B36", + "CAFormer_M36", + "CAFormer_S18", + "CAFormer_S36", + "CoAtNet_0_RW_224", + "CoAtNet_1_RW_224", + "CoAtNet_2_RW_224", + "ConvFormer_B36", + "ConvFormer_M36", + "ConvFormer_S18", + "ConvFormer_S36", + "ConvNeXt_Base", + "ConvNeXt_Large", + "ConvNeXt_Small", + "ConvNeXt_Tiny", + "DeiT_Base_Patch16_224", + "DeiT_Base_Patch16_384", + "DeiT_Small_Patch16_224", + "DeiT_Tiny_Patch16_224", + "DeiT3_Base_Patch16_224", + "DeiT3_Base_Patch16_384", + "DeiT3_Large_Patch16_224", + "DeiT3_Large_Patch16_384", + "DeiT3_Medium_Patch16_224", + "DeiT3_Small_Patch16_224", + "DeiT3_Small_Patch16_384", + "DenseNet121", + "DenseNet161", + "DenseNet169", + "DenseNet201", + "EfficientFormer_L1", + "EfficientFormer_L3", + "EfficientFormer_L7", + "EfficientNet_B0", + "EfficientNet_B1", + "EfficientNet_B2", + "EfficientNet_B3", + "EfficientNet_B4", + "EfficientNet_B5", + "EfficientNet_B6", + "EfficientNet_B7", + "EfficientNet_V2_L", + "EfficientNet_V2_M", + "EfficientNet_V2_S", + "FlexiViT_Base", + "FlexiViT_Large", + "FlexiViT_Small", + "GoogLeNet", + "Inception_V3", + "LeViT_Conv_128", + "LeViT_Conv_128S", + "LeViT_Conv_192", + "LeViT_Conv_256", + "LeViT_Conv_384", + "MNASNet0_5", + "MNASNet0_75", + "MNASNet1_0", + "MNASNet1_3", + "MobileNet_V2", + "MobileNet_V3_Large", + "MobileNet_V3_Small", + "RegNet_X_1_6GF", + "RegNet_X_3_2GF", + "RegNet_X_8GF", + "RegNet_X_16GF", + "RegNet_X_32GF", + "RegNet_X_400MF", + "RegNet_X_800MF", + "RegNet_Y_1_6GF", + "RegNet_Y_3_2GF", + "RegNet_Y_8GF", + "RegNet_Y_16GF", + "RegNet_Y_32GF", + "RegNet_Y_400MF", + "RegNet_Y_800MF", + "RepViT_M0_9", + "RepViT_M1", + "RepViT_M1_0", + "RepViT_M1_1", + "RepViT_M1_5", + "RepViT_M2", + "RepViT_M2_3", + "RepViT_M3", + "ResNet18", + "ResNet34", + "ResNet50", + "ResNet101", + "ResNet152", + "ResNeXt50_32x4d", + "ResNeXt101_32x8d", + "ResNeXt101_64x4d", + "ShuffleNet_V2_X0_5", + "ShuffleNet_V2_X1_0", + "ShuffleNet_V2_X1_5", + "ShuffleNet_V2_X2_0", + "SqueezeNet1_0", + "SqueezeNet1_1", + "Swin_B", + "Swin_S", + "Swin_T", + "VGG11", + "VGG11_BN", + "VGG13", + "VGG13_BN", + "VGG16", + "VGG16_BN", + "VGG19", + "VGG19_BN", + "VisFormer_Small", + "VisFormer_Tiny", + "ViT_B_16", + "ViT_B_32", + "ViT_Base_Patch8_224", + "ViT_Base_Patch16_224", + "ViT_Base_Patch16_384", + "ViT_Base_Patch32_224", + "ViT_Base_Patch32_384", + "ViT_L_16", + "ViT_L_32", + "ViT_Large_Patch16_224", + "ViT_Large_Patch16_384", + "ViT_Large_Patch32_384", + "ViT_Small_Patch16_224", + "ViT_Small_Patch16_384", + "ViT_Small_Patch32_224", + "ViT_Small_Patch32_384", + "ViT_Tiny_Patch16_224", + "ViT_Tiny_Patch16_384", + "Wide_ResNet50_2", + "Wide_ResNet101_2", + "YOLO11lCls", + "YOLO11mCls", + "YOLO11nCls", + "YOLO11sCls", + "YOLO11xCls", + "YOLO26lCls", + "YOLO26mCls", + "YOLO26nCls", + "YOLO26sCls", + "YOLO26xCls", + "YOLOv5lCls", + "YOLOv5mCls", + "YOLOv5nCls", + "YOLOv5sCls", + "YOLOv5xCls", + "YOLOv8lCls", + "YOLOv8mCls", + "YOLOv8nCls", + "YOLOv8sCls", + "YOLOv8xCls", +] + +AlexNet = create_model_class("AlexNet", __name__) +CAFormer_B36 = create_model_class("CAFormer_B36", __name__) +CAFormer_M36 = create_model_class("CAFormer_M36", __name__) +CAFormer_S18 = create_model_class("CAFormer_S18", __name__) +CAFormer_S36 = create_model_class("CAFormer_S36", __name__) +CoAtNet_0_RW_224 = create_model_class("CoAtNet_0_RW_224", __name__) +CoAtNet_1_RW_224 = create_model_class("CoAtNet_1_RW_224", __name__) +CoAtNet_2_RW_224 = create_model_class("CoAtNet_2_RW_224", __name__) +ConvFormer_B36 = create_model_class("ConvFormer_B36", __name__) +ConvFormer_M36 = create_model_class("ConvFormer_M36", __name__) +ConvFormer_S18 = create_model_class("ConvFormer_S18", __name__) +ConvFormer_S36 = create_model_class("ConvFormer_S36", __name__) +ConvNeXt_Base = create_model_class("ConvNeXt_Base", __name__) +ConvNeXt_Large = create_model_class("ConvNeXt_Large", __name__) +ConvNeXt_Small = create_model_class("ConvNeXt_Small", __name__) +ConvNeXt_Tiny = create_model_class("ConvNeXt_Tiny", __name__) +DeiT_Base_Patch16_224 = create_model_class("DeiT_Base_Patch16_224", __name__) +DeiT_Base_Patch16_384 = create_model_class("DeiT_Base_Patch16_384", __name__) +DeiT_Small_Patch16_224 = create_model_class("DeiT_Small_Patch16_224", __name__) +DeiT_Tiny_Patch16_224 = create_model_class("DeiT_Tiny_Patch16_224", __name__) +DeiT3_Base_Patch16_224 = create_model_class("DeiT3_Base_Patch16_224", __name__) +DeiT3_Base_Patch16_384 = create_model_class("DeiT3_Base_Patch16_384", __name__) +DeiT3_Large_Patch16_224 = create_model_class("DeiT3_Large_Patch16_224", __name__) +DeiT3_Large_Patch16_384 = create_model_class("DeiT3_Large_Patch16_384", __name__) +DeiT3_Medium_Patch16_224 = create_model_class("DeiT3_Medium_Patch16_224", __name__) +DeiT3_Small_Patch16_224 = create_model_class("DeiT3_Small_Patch16_224", __name__) +DeiT3_Small_Patch16_384 = create_model_class("DeiT3_Small_Patch16_384", __name__) +DenseNet121 = create_model_class("DenseNet121", __name__) +DenseNet161 = create_model_class("DenseNet161", __name__) +DenseNet169 = create_model_class("DenseNet169", __name__) +DenseNet201 = create_model_class("DenseNet201", __name__) +EfficientFormer_L1 = create_model_class("EfficientFormer_L1", __name__) +EfficientFormer_L3 = create_model_class("EfficientFormer_L3", __name__) +EfficientFormer_L7 = create_model_class("EfficientFormer_L7", __name__) +EfficientNet_B0 = create_model_class("EfficientNet_B0", __name__) +EfficientNet_B1 = create_model_class("EfficientNet_B1", __name__) +EfficientNet_B2 = create_model_class("EfficientNet_B2", __name__) +EfficientNet_B3 = create_model_class("EfficientNet_B3", __name__) +EfficientNet_B4 = create_model_class("EfficientNet_B4", __name__) +EfficientNet_B5 = create_model_class("EfficientNet_B5", __name__) +EfficientNet_B6 = create_model_class("EfficientNet_B6", __name__) +EfficientNet_B7 = create_model_class("EfficientNet_B7", __name__) +EfficientNet_V2_L = create_model_class("EfficientNet_V2_L", __name__) +EfficientNet_V2_M = create_model_class("EfficientNet_V2_M", __name__) +EfficientNet_V2_S = create_model_class("EfficientNet_V2_S", __name__) +FlexiViT_Base = create_model_class("FlexiViT_Base", __name__) +FlexiViT_Large = create_model_class("FlexiViT_Large", __name__) +FlexiViT_Small = create_model_class("FlexiViT_Small", __name__) +GoogLeNet = create_model_class("GoogLeNet", __name__) +Inception_V3 = create_model_class("Inception_V3", __name__) +LeViT_Conv_128 = create_model_class("LeViT_Conv_128", __name__) +LeViT_Conv_128S = create_model_class("LeViT_Conv_128S", __name__) +LeViT_Conv_192 = create_model_class("LeViT_Conv_192", __name__) +LeViT_Conv_256 = create_model_class("LeViT_Conv_256", __name__) +LeViT_Conv_384 = create_model_class("LeViT_Conv_384", __name__) +MNASNet0_5 = create_model_class("MNASNet0_5", __name__) +MNASNet0_75 = create_model_class("MNASNet0_75", __name__) +MNASNet1_0 = create_model_class("MNASNet1_0", __name__) +MNASNet1_3 = create_model_class("MNASNet1_3", __name__) +MobileNet_V2 = create_model_class("MobileNet_V2", __name__) +MobileNet_V3_Large = create_model_class("MobileNet_V3_Large", __name__) +MobileNet_V3_Small = create_model_class("MobileNet_V3_Small", __name__) +RegNet_X_1_6GF = create_model_class("RegNet_X_1_6GF", __name__) +RegNet_X_3_2GF = create_model_class("RegNet_X_3_2GF", __name__) +RegNet_X_8GF = create_model_class("RegNet_X_8GF", __name__) +RegNet_X_16GF = create_model_class("RegNet_X_16GF", __name__) +RegNet_X_32GF = create_model_class("RegNet_X_32GF", __name__) +RegNet_X_400MF = create_model_class("RegNet_X_400MF", __name__) +RegNet_X_800MF = create_model_class("RegNet_X_800MF", __name__) +RegNet_Y_1_6GF = create_model_class("RegNet_Y_1_6GF", __name__) +RegNet_Y_3_2GF = create_model_class("RegNet_Y_3_2GF", __name__) +RegNet_Y_8GF = create_model_class("RegNet_Y_8GF", __name__) +RegNet_Y_16GF = create_model_class("RegNet_Y_16GF", __name__) +RegNet_Y_32GF = create_model_class("RegNet_Y_32GF", __name__) +RegNet_Y_400MF = create_model_class("RegNet_Y_400MF", __name__) +RegNet_Y_800MF = create_model_class("RegNet_Y_800MF", __name__) +RepViT_M0_9 = create_model_class("RepViT_M0_9", __name__) +RepViT_M1 = create_model_class("RepViT_M1", __name__) +RepViT_M1_0 = create_model_class("RepViT_M1_0", __name__) +RepViT_M1_1 = create_model_class("RepViT_M1_1", __name__) +RepViT_M1_5 = create_model_class("RepViT_M1_5", __name__) +RepViT_M2 = create_model_class("RepViT_M2", __name__) +RepViT_M2_3 = create_model_class("RepViT_M2_3", __name__) +RepViT_M3 = create_model_class("RepViT_M3", __name__) +ResNet18 = create_model_class("ResNet18", __name__) +ResNet34 = create_model_class("ResNet34", __name__) +ResNet50 = create_model_class("ResNet50", __name__) +ResNet101 = create_model_class("ResNet101", __name__) +ResNet152 = create_model_class("ResNet152", __name__) +ResNeXt50_32x4d = create_model_class("ResNeXt50_32x4d", __name__) +ResNeXt101_32x8d = create_model_class("ResNeXt101_32x8d", __name__) +ResNeXt101_64x4d = create_model_class("ResNeXt101_64x4d", __name__) +ShuffleNet_V2_X0_5 = create_model_class("ShuffleNet_V2_X0_5", __name__) +ShuffleNet_V2_X1_0 = create_model_class("ShuffleNet_V2_X1_0", __name__) +ShuffleNet_V2_X1_5 = create_model_class("ShuffleNet_V2_X1_5", __name__) +ShuffleNet_V2_X2_0 = create_model_class("ShuffleNet_V2_X2_0", __name__) +SqueezeNet1_0 = create_model_class("SqueezeNet1_0", __name__) +SqueezeNet1_1 = create_model_class("SqueezeNet1_1", __name__) +Swin_B = create_model_class("Swin_B", __name__) +Swin_S = create_model_class("Swin_S", __name__) +Swin_T = create_model_class("Swin_T", __name__) +VGG11 = create_model_class("VGG11", __name__) +VGG11_BN = create_model_class("VGG11_BN", __name__) +VGG13 = create_model_class("VGG13", __name__) +VGG13_BN = create_model_class("VGG13_BN", __name__) +VGG16 = create_model_class("VGG16", __name__) +VGG16_BN = create_model_class("VGG16_BN", __name__) +VGG19 = create_model_class("VGG19", __name__) +VGG19_BN = create_model_class("VGG19_BN", __name__) +VisFormer_Small = create_model_class("VisFormer_Small", __name__) +VisFormer_Tiny = create_model_class("VisFormer_Tiny", __name__) +ViT_B_16 = create_model_class("ViT_B_16", __name__) +ViT_B_32 = create_model_class("ViT_B_32", __name__) +ViT_Base_Patch8_224 = create_model_class("ViT_Base_Patch8_224", __name__) +ViT_Base_Patch16_224 = create_model_class("ViT_Base_Patch16_224", __name__) +ViT_Base_Patch16_384 = create_model_class("ViT_Base_Patch16_384", __name__) +ViT_Base_Patch32_224 = create_model_class("ViT_Base_Patch32_224", __name__) +ViT_Base_Patch32_384 = create_model_class("ViT_Base_Patch32_384", __name__) +ViT_L_16 = create_model_class("ViT_L_16", __name__) +ViT_L_32 = create_model_class("ViT_L_32", __name__) +ViT_Large_Patch16_224 = create_model_class("ViT_Large_Patch16_224", __name__) +ViT_Large_Patch16_384 = create_model_class("ViT_Large_Patch16_384", __name__) +ViT_Large_Patch32_384 = create_model_class("ViT_Large_Patch32_384", __name__) +ViT_Small_Patch16_224 = create_model_class("ViT_Small_Patch16_224", __name__) +ViT_Small_Patch16_384 = create_model_class("ViT_Small_Patch16_384", __name__) +ViT_Small_Patch32_224 = create_model_class("ViT_Small_Patch32_224", __name__) +ViT_Small_Patch32_384 = create_model_class("ViT_Small_Patch32_384", __name__) +ViT_Tiny_Patch16_224 = create_model_class("ViT_Tiny_Patch16_224", __name__) +ViT_Tiny_Patch16_384 = create_model_class("ViT_Tiny_Patch16_384", __name__) +Wide_ResNet50_2 = create_model_class("Wide_ResNet50_2", __name__) +Wide_ResNet101_2 = create_model_class("Wide_ResNet101_2", __name__) +YOLO11lCls = create_model_class("YOLO11lCls", __name__) +YOLO11mCls = create_model_class("YOLO11mCls", __name__) +YOLO11nCls = create_model_class("YOLO11nCls", __name__) +YOLO11sCls = create_model_class("YOLO11sCls", __name__) +YOLO11xCls = create_model_class("YOLO11xCls", __name__) +YOLO26lCls = create_model_class("YOLO26lCls", __name__) +YOLO26mCls = create_model_class("YOLO26mCls", __name__) +YOLO26nCls = create_model_class("YOLO26nCls", __name__) +YOLO26sCls = create_model_class("YOLO26sCls", __name__) +YOLO26xCls = create_model_class("YOLO26xCls", __name__) +YOLOv5lCls = create_model_class("YOLOv5lCls", __name__) +YOLOv5mCls = create_model_class("YOLOv5mCls", __name__) +YOLOv5nCls = create_model_class("YOLOv5nCls", __name__) +YOLOv5sCls = create_model_class("YOLOv5sCls", __name__) +YOLOv5xCls = create_model_class("YOLOv5xCls", __name__) +YOLOv8lCls = create_model_class("YOLOv8lCls", __name__) +YOLOv8mCls = create_model_class("YOLOv8mCls", __name__) +YOLOv8nCls = create_model_class("YOLOv8nCls", __name__) +YOLOv8sCls = create_model_class("YOLOv8sCls", __name__) +YOLOv8xCls = create_model_class("YOLOv8xCls", __name__) diff --git a/mblt_vision/instance_segmentation/__init__.py b/mblt_vision/instance_segmentation/__init__.py new file mode 100644 index 0000000..421e9ff --- /dev/null +++ b/mblt_vision/instance_segmentation/__init__.py @@ -0,0 +1,65 @@ +"""Instance segmentation model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO11lSeg", + "YOLO11mSeg", + "YOLO11nSeg", + "YOLO11sSeg", + "YOLO11xSeg", + "YOLO12lSeg", + "YOLO12mSeg", + "YOLO12nSeg", + "YOLO12sSeg", + "YOLO12xSeg", + "YOLO26lSeg", + "YOLO26mSeg", + "YOLO26nSeg", + "YOLO26sSeg", + "YOLO26xSeg", + "YOLOv5lSeg", + "YOLOv5mSeg", + "YOLOv5nSeg", + "YOLOv5sSeg", + "YOLOv5xSeg", + "YOLOv8lSeg", + "YOLOv8mSeg", + "YOLOv8nSeg", + "YOLOv8sSeg", + "YOLOv8xSeg", + "GELANcSeg", + "YOLOv9cSeg", + "YOLOv9eSeg", +] + +YOLO11lSeg = create_model_class("YOLO11lSeg", __name__) +YOLO11mSeg = create_model_class("YOLO11mSeg", __name__) +YOLO11nSeg = create_model_class("YOLO11nSeg", __name__) +YOLO11sSeg = create_model_class("YOLO11sSeg", __name__) +YOLO11xSeg = create_model_class("YOLO11xSeg", __name__) +YOLO12lSeg = create_model_class("YOLO12lSeg", __name__) +YOLO12mSeg = create_model_class("YOLO12mSeg", __name__) +YOLO12nSeg = create_model_class("YOLO12nSeg", __name__) +YOLO12sSeg = create_model_class("YOLO12sSeg", __name__) +YOLO12xSeg = create_model_class("YOLO12xSeg", __name__) +YOLO26lSeg = create_model_class("YOLO26lSeg", __name__) +YOLO26mSeg = create_model_class("YOLO26mSeg", __name__) +YOLO26nSeg = create_model_class("YOLO26nSeg", __name__) +YOLO26sSeg = create_model_class("YOLO26sSeg", __name__) +YOLO26xSeg = create_model_class("YOLO26xSeg", __name__) +YOLOv5lSeg = create_model_class("YOLOv5lSeg", __name__) +YOLOv5mSeg = create_model_class("YOLOv5mSeg", __name__) +YOLOv5nSeg = create_model_class("YOLOv5nSeg", __name__) +YOLOv5sSeg = create_model_class("YOLOv5sSeg", __name__) +YOLOv5xSeg = create_model_class("YOLOv5xSeg", __name__) +YOLOv8lSeg = create_model_class("YOLOv8lSeg", __name__) +YOLOv8mSeg = create_model_class("YOLOv8mSeg", __name__) +YOLOv8nSeg = create_model_class("YOLOv8nSeg", __name__) +YOLOv8sSeg = create_model_class("YOLOv8sSeg", __name__) +YOLOv8xSeg = create_model_class("YOLOv8xSeg", __name__) +GELANcSeg = create_model_class("GELANcSeg", __name__) +YOLOv9cSeg = create_model_class("YOLOv9cSeg", __name__) +YOLOv9eSeg = create_model_class("YOLOv9eSeg", __name__) diff --git a/mblt_vision/models/AlexNet.yaml b/mblt_vision/models/AlexNet.yaml new file mode 100644 index 0000000..b9dad06 --- /dev/null +++ b/mblt_vision/models/AlexNet.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/AlexNet + filename: alexnet_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CAFormer_B36.yaml b/mblt_vision/models/CAFormer_B36.yaml new file mode 100644 index 0000000..691448e --- /dev/null +++ b/mblt_vision/models/CAFormer_B36.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CAFormer_B36 + filename: caformer_b36.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CAFormer_M36.yaml b/mblt_vision/models/CAFormer_M36.yaml new file mode 100644 index 0000000..0896e02 --- /dev/null +++ b/mblt_vision/models/CAFormer_M36.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CAFormer_M36 + filename: caformer_m36.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CAFormer_S18.yaml b/mblt_vision/models/CAFormer_S18.yaml new file mode 100644 index 0000000..0a381d6 --- /dev/null +++ b/mblt_vision/models/CAFormer_S18.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CAFormer_S18 + filename: caformer_s18.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CAFormer_S36.yaml b/mblt_vision/models/CAFormer_S36.yaml new file mode 100644 index 0000000..f0911b4 --- /dev/null +++ b/mblt_vision/models/CAFormer_S36.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CAFormer_S36 + filename: caformer_s36.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CoAtNet_0_RW_224.yaml b/mblt_vision/models/CoAtNet_0_RW_224.yaml new file mode 100644 index 0000000..d9b553f --- /dev/null +++ b/mblt_vision/models/CoAtNet_0_RW_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CoAtNet_0_RW_224 + filename: coatnet_0_rw_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CoAtNet_1_RW_224.yaml b/mblt_vision/models/CoAtNet_1_RW_224.yaml new file mode 100644 index 0000000..9d1f493 --- /dev/null +++ b/mblt_vision/models/CoAtNet_1_RW_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CoAtNet_1_RW_224 + filename: coatnet_1_rw_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/CoAtNet_2_RW_224.yaml b/mblt_vision/models/CoAtNet_2_RW_224.yaml new file mode 100644 index 0000000..d1985c9 --- /dev/null +++ b/mblt_vision/models/CoAtNet_2_RW_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/CoAtNet_2_RW_224 + filename: coatnet_2_rw_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvFormer_B36.yaml b/mblt_vision/models/ConvFormer_B36.yaml new file mode 100644 index 0000000..144b6d2 --- /dev/null +++ b/mblt_vision/models/ConvFormer_B36.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ConvFormer_B36 + filename: convformer_b36.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvFormer_M36.yaml b/mblt_vision/models/ConvFormer_M36.yaml new file mode 100644 index 0000000..ba679c9 --- /dev/null +++ b/mblt_vision/models/ConvFormer_M36.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ConvFormer_M36 + filename: convformer_m36.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvFormer_S18.yaml b/mblt_vision/models/ConvFormer_S18.yaml new file mode 100644 index 0000000..2fbb09f --- /dev/null +++ b/mblt_vision/models/ConvFormer_S18.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ConvFormer_S18 + filename: convformer_s18.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvFormer_S36.yaml b/mblt_vision/models/ConvFormer_S36.yaml new file mode 100644 index 0000000..d39c46e --- /dev/null +++ b/mblt_vision/models/ConvFormer_S36.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ConvFormer_S36 + filename: convformer_s36.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvNext_Base.yaml b/mblt_vision/models/ConvNext_Base.yaml new file mode 100644 index 0000000..0de192a --- /dev/null +++ b/mblt_vision/models/ConvNext_Base.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ConvNext_Base + filename: convnext_base_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvNext_Large.yaml b/mblt_vision/models/ConvNext_Large.yaml new file mode 100644 index 0000000..3047676 --- /dev/null +++ b/mblt_vision/models/ConvNext_Large.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ConvNext_Large + filename: convnext_large_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvNext_Small.yaml b/mblt_vision/models/ConvNext_Small.yaml new file mode 100644 index 0000000..c986f90 --- /dev/null +++ b/mblt_vision/models/ConvNext_Small.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ConvNext_Small + filename: convnext_small_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 230 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ConvNext_Tiny.yaml b/mblt_vision/models/ConvNext_Tiny.yaml new file mode 100644 index 0000000..241cfe3 --- /dev/null +++ b/mblt_vision/models/ConvNext_Tiny.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ConvNext_Tiny + filename: convnext_tiny_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Base_Patch16_224.yaml b/mblt_vision/models/DeiT3_Base_Patch16_224.yaml new file mode 100644 index 0000000..3e4c4da --- /dev/null +++ b/mblt_vision/models/DeiT3_Base_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Base_Patch16_224 + filename: deit3_base_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Base_Patch16_384.yaml b/mblt_vision/models/DeiT3_Base_Patch16_384.yaml new file mode 100644 index 0000000..5eab39b --- /dev/null +++ b/mblt_vision/models/DeiT3_Base_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Base_Patch16_384 + filename: deit3_base_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Large_Patch16_224.yaml b/mblt_vision/models/DeiT3_Large_Patch16_224.yaml new file mode 100644 index 0000000..1db68c3 --- /dev/null +++ b/mblt_vision/models/DeiT3_Large_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Large_Patch16_224 + filename: deit3_large_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Large_Patch16_384.yaml b/mblt_vision/models/DeiT3_Large_Patch16_384.yaml new file mode 100644 index 0000000..92efd1e --- /dev/null +++ b/mblt_vision/models/DeiT3_Large_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Large_Patch16_384 + filename: deit3_large_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Medium_Patch16_224.yaml b/mblt_vision/models/DeiT3_Medium_Patch16_224.yaml new file mode 100644 index 0000000..c4bb1da --- /dev/null +++ b/mblt_vision/models/DeiT3_Medium_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Medium_Patch16_224 + filename: deit3_medium_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Small_Patch16_224.yaml b/mblt_vision/models/DeiT3_Small_Patch16_224.yaml new file mode 100644 index 0000000..1aa66e9 --- /dev/null +++ b/mblt_vision/models/DeiT3_Small_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Small_Patch16_224 + filename: deit3_small_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT3_Small_Patch16_384.yaml b/mblt_vision/models/DeiT3_Small_Patch16_384.yaml new file mode 100644 index 0000000..b576828 --- /dev/null +++ b/mblt_vision/models/DeiT3_Small_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT3_Small_Patch16_384 + filename: deit3_small_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT_Base_Patch16_224.yaml b/mblt_vision/models/DeiT_Base_Patch16_224.yaml new file mode 100644 index 0000000..872adc0 --- /dev/null +++ b/mblt_vision/models/DeiT_Base_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT_Base_Patch16_224 + filename: deit_base_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT_Base_Patch16_384.yaml b/mblt_vision/models/DeiT_Base_Patch16_384.yaml new file mode 100644 index 0000000..5b23b09 --- /dev/null +++ b/mblt_vision/models/DeiT_Base_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT_Base_Patch16_384 + filename: deit_base_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT_Small_Patch16_224.yaml b/mblt_vision/models/DeiT_Small_Patch16_224.yaml new file mode 100644 index 0000000..ab5372e --- /dev/null +++ b/mblt_vision/models/DeiT_Small_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT_Small_Patch16_224 + filename: deit_small_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DeiT_Tiny_Patch16_224.yaml b/mblt_vision/models/DeiT_Tiny_Patch16_224.yaml new file mode 100644 index 0000000..80f6f9d --- /dev/null +++ b/mblt_vision/models/DeiT_Tiny_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/DeiT_Tiny_Patch16_224 + filename: deit_tiny_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DenseNet121.yaml b/mblt_vision/models/DenseNet121.yaml new file mode 100644 index 0000000..17ac464 --- /dev/null +++ b/mblt_vision/models/DenseNet121.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/DenseNet121 + filename: densenet121_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DenseNet161.yaml b/mblt_vision/models/DenseNet161.yaml new file mode 100644 index 0000000..6ef129b --- /dev/null +++ b/mblt_vision/models/DenseNet161.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/DenseNet161 + filename: densenet161_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DenseNet169.yaml b/mblt_vision/models/DenseNet169.yaml new file mode 100644 index 0000000..02ef675 --- /dev/null +++ b/mblt_vision/models/DenseNet169.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/DenseNet169 + filename: densenet169_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/DenseNet201.yaml b/mblt_vision/models/DenseNet201.yaml new file mode 100644 index 0000000..661c37b --- /dev/null +++ b/mblt_vision/models/DenseNet201.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/DenseNet201 + filename: densenet201_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientFormer_L1.yaml b/mblt_vision/models/EfficientFormer_L1.yaml new file mode 100644 index 0000000..453b1b7 --- /dev/null +++ b/mblt_vision/models/EfficientFormer_L1.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/EfficientFormer_L1 + filename: efficientformer_l1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientFormer_L3.yaml b/mblt_vision/models/EfficientFormer_L3.yaml new file mode 100644 index 0000000..7ecf9c3 --- /dev/null +++ b/mblt_vision/models/EfficientFormer_L3.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/EfficientFormer_L3 + filename: efficientformer_l3.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientFormer_L7.yaml b/mblt_vision/models/EfficientFormer_L7.yaml new file mode 100644 index 0000000..403333f --- /dev/null +++ b/mblt_vision/models/EfficientFormer_L7.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/EfficientFormer_L7 + filename: efficientformer_l7.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B0.yaml b/mblt_vision/models/EfficientNet_B0.yaml new file mode 100644 index 0000000..64b50b1 --- /dev/null +++ b/mblt_vision/models/EfficientNet_B0.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B0 + filename: efficientnet_b0_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B1.yaml b/mblt_vision/models/EfficientNet_B1.yaml new file mode 100644 index 0000000..53f63ac --- /dev/null +++ b/mblt_vision/models/EfficientNet_B1.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B1.tv1_in1k + filename: efficientnet_b1_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bicubic + CenterCrop: + size: + - 240 + - 240 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/EfficientNet_B1.tv2_in1k + filename: efficientnet_b1_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 255 + interpolation: bilinear diff --git a/mblt_vision/models/EfficientNet_B2.yaml b/mblt_vision/models/EfficientNet_B2.yaml new file mode 100644 index 0000000..595086a --- /dev/null +++ b/mblt_vision/models/EfficientNet_B2.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B2 + filename: efficientnet_b2_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 288 + interpolation: bicubic + CenterCrop: + size: + - 288 + - 288 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B3.yaml b/mblt_vision/models/EfficientNet_B3.yaml new file mode 100644 index 0000000..686906c --- /dev/null +++ b/mblt_vision/models/EfficientNet_B3.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B3 + filename: efficientnet_b3_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 320 + interpolation: bicubic + CenterCrop: + size: + - 300 + - 300 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B4.yaml b/mblt_vision/models/EfficientNet_B4.yaml new file mode 100644 index 0000000..21f2694 --- /dev/null +++ b/mblt_vision/models/EfficientNet_B4.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B4 + filename: efficientnet_b4_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 380 + - 380 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B5.yaml b/mblt_vision/models/EfficientNet_B5.yaml new file mode 100644 index 0000000..94855c7 --- /dev/null +++ b/mblt_vision/models/EfficientNet_B5.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B5 + filename: efficientnet_b5_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 456 + interpolation: bicubic + CenterCrop: + size: + - 456 + - 456 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B6.yaml b/mblt_vision/models/EfficientNet_B6.yaml new file mode 100644 index 0000000..61723c4 --- /dev/null +++ b/mblt_vision/models/EfficientNet_B6.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B6 + filename: efficientnet_b6_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 528 + interpolation: bicubic + CenterCrop: + size: + - 528 + - 528 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_B7.yaml b/mblt_vision/models/EfficientNet_B7.yaml new file mode 100644 index 0000000..00143ab --- /dev/null +++ b/mblt_vision/models/EfficientNet_B7.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_B7 + filename: efficientnet_b7_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 600 + interpolation: bicubic + CenterCrop: + size: + - 600 + - 600 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_V2_L.yaml b/mblt_vision/models/EfficientNet_V2_L.yaml new file mode 100644 index 0000000..03ec40f --- /dev/null +++ b/mblt_vision/models/EfficientNet_V2_L.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_V2_L + filename: efficientnet_v2_l_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 480 + interpolation: bicubic + CenterCrop: + size: + - 480 + - 480 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_V2_M.yaml b/mblt_vision/models/EfficientNet_V2_M.yaml new file mode 100644 index 0000000..f93ac0c --- /dev/null +++ b/mblt_vision/models/EfficientNet_V2_M.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_V2_M + filename: efficientnet_v2_m_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 480 + interpolation: bilinear + CenterCrop: + size: + - 480 + - 480 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/EfficientNet_V2_S.yaml b/mblt_vision/models/EfficientNet_V2_S.yaml new file mode 100644 index 0000000..1ced73c --- /dev/null +++ b/mblt_vision/models/EfficientNet_V2_S.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/EfficientNet_V2_S + filename: efficientnet_v2_s_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bilinear + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/FlexiVit_Base.yaml b/mblt_vision/models/FlexiVit_Base.yaml new file mode 100644 index 0000000..9999d9c --- /dev/null +++ b/mblt_vision/models/FlexiVit_Base.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/FlexiVit_Base + filename: flexivit_base.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 252 + interpolation: bicubic + CenterCrop: + size: + - 240 + - 240 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/FlexiVit_Large.yaml b/mblt_vision/models/FlexiVit_Large.yaml new file mode 100644 index 0000000..5d3e166 --- /dev/null +++ b/mblt_vision/models/FlexiVit_Large.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/FlexiVit_Large + filename: flexivit_large.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 252 + interpolation: bicubic + CenterCrop: + size: + - 240 + - 240 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/FlexiVit_Small.yaml b/mblt_vision/models/FlexiVit_Small.yaml new file mode 100644 index 0000000..81e87c9 --- /dev/null +++ b/mblt_vision/models/FlexiVit_Small.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/FlexiVit_Small + filename: flexivit_small.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 252 + interpolation: bicubic + CenterCrop: + size: + - 240 + - 240 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/GELANc-seg.yaml b/mblt_vision/models/GELANc-seg.yaml new file mode 100644 index 0000000..9ef40b4 --- /dev/null +++ b/mblt_vision/models/GELANc-seg.yaml @@ -0,0 +1,29 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/GELANc-seg + filename: gelanc-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nc: 80 + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/GELANc.yaml b/mblt_vision/models/GELANc.yaml new file mode 100644 index 0000000..02b40e8 --- /dev/null +++ b/mblt_vision/models/GELANc.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/GELANc + filename: gelanc.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nc: 80 + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/GELANe.yaml b/mblt_vision/models/GELANe.yaml new file mode 100644 index 0000000..e73e031 --- /dev/null +++ b/mblt_vision/models/GELANe.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/GELANe + filename: gelane.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nc: 80 + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/GELANm.yaml b/mblt_vision/models/GELANm.yaml new file mode 100644 index 0000000..71beb02 --- /dev/null +++ b/mblt_vision/models/GELANm.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/GELANm + filename: gelanm.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nc: 80 + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/GELANs.yaml b/mblt_vision/models/GELANs.yaml new file mode 100644 index 0000000..a3454e3 --- /dev/null +++ b/mblt_vision/models/GELANs.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/GELANs + filename: gelans.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nc: 80 + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/GoogLeNet.yaml b/mblt_vision/models/GoogLeNet.yaml new file mode 100644 index 0000000..e07dfd7 --- /dev/null +++ b/mblt_vision/models/GoogLeNet.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/GoogLeNet + filename: googlenet_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/Inception_V3.yaml b/mblt_vision/models/Inception_V3.yaml new file mode 100644 index 0000000..cec6a41 --- /dev/null +++ b/mblt_vision/models/Inception_V3.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/Inception_V3 + filename: inception_v3_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 342 + interpolation: bilinear + CenterCrop: + size: + - 299 + - 299 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/LeViT_Conv_128.yaml b/mblt_vision/models/LeViT_Conv_128.yaml new file mode 100644 index 0000000..9e20421 --- /dev/null +++ b/mblt_vision/models/LeViT_Conv_128.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/LeViT_Conv_128 + filename: levit_conv_128.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/LeViT_Conv_128S.yaml b/mblt_vision/models/LeViT_Conv_128S.yaml new file mode 100644 index 0000000..5a39599 --- /dev/null +++ b/mblt_vision/models/LeViT_Conv_128S.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/LeViT_Conv_128S + filename: levit_conv_128s.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/LeViT_Conv_192.yaml b/mblt_vision/models/LeViT_Conv_192.yaml new file mode 100644 index 0000000..86027f8 --- /dev/null +++ b/mblt_vision/models/LeViT_Conv_192.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/LeViT_Conv_192 + filename: levit_conv_192.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/LeViT_Conv_256.yaml b/mblt_vision/models/LeViT_Conv_256.yaml new file mode 100644 index 0000000..5d78350 --- /dev/null +++ b/mblt_vision/models/LeViT_Conv_256.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/LeViT_Conv_256 + filename: levit_conv_256.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/LeViT_Conv_384.yaml b/mblt_vision/models/LeViT_Conv_384.yaml new file mode 100644 index 0000000..9c5ebae --- /dev/null +++ b/mblt_vision/models/LeViT_Conv_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/LeViT_Conv_384 + filename: levit_conv_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/MNASNet0_5.yaml b/mblt_vision/models/MNASNet0_5.yaml new file mode 100644 index 0000000..12834ac --- /dev/null +++ b/mblt_vision/models/MNASNet0_5.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MNASNet0_5 + filename: mnasnet0_5_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/MNASNet0_75.yaml b/mblt_vision/models/MNASNet0_75.yaml new file mode 100644 index 0000000..e47c643 --- /dev/null +++ b/mblt_vision/models/MNASNet0_75.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MNASNet0_75 + filename: mnasnet0_75_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/MNASNet1_0.yaml b/mblt_vision/models/MNASNet1_0.yaml new file mode 100644 index 0000000..6c0f2ab --- /dev/null +++ b/mblt_vision/models/MNASNet1_0.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MNASNet1_0 + filename: mnasnet1_0_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/MNASNet1_3.yaml b/mblt_vision/models/MNASNet1_3.yaml new file mode 100644 index 0000000..d20ccc1 --- /dev/null +++ b/mblt_vision/models/MNASNet1_3.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MNASNet1_3 + filename: mnasnet1_3_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/MobileNet_V2.yaml b/mblt_vision/models/MobileNet_V2.yaml new file mode 100644 index 0000000..67349da --- /dev/null +++ b/mblt_vision/models/MobileNet_V2.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MobileNet_V2.tv1_in1k + filename: mobilenet_v2_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/MobileNet_V2.tv2_in1k + filename: mobilenet_v2_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/MobileNet_V3_Large.yaml b/mblt_vision/models/MobileNet_V3_Large.yaml new file mode 100644 index 0000000..a94f87f --- /dev/null +++ b/mblt_vision/models/MobileNet_V3_Large.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MobileNet_V3_Large.tv1_in1k + filename: mobilenet_v3_large_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/MobileNet_V3_Large.tv2_in1k + filename: mobilenet_v3_large_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/MobileNet_V3_Small.yaml b/mblt_vision/models/MobileNet_V3_Small.yaml new file mode 100644 index 0000000..c1e3fd6 --- /dev/null +++ b/mblt_vision/models/MobileNet_V3_Small.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/MobileNet_V3_Small + filename: mobilenet_v3_small_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RegNet_X_16GF.yaml b/mblt_vision/models/RegNet_X_16GF.yaml new file mode 100644 index 0000000..95dd2aa --- /dev/null +++ b/mblt_vision/models/RegNet_X_16GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_16GF.tv1_in1k + filename: regnet_x_16gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_16GF.tv2_in1k + filename: regnet_x_16gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_X_1_6GF.yaml b/mblt_vision/models/RegNet_X_1_6GF.yaml new file mode 100644 index 0000000..f5c3d3c --- /dev/null +++ b/mblt_vision/models/RegNet_X_1_6GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_1_6GF.tv1_in1k + filename: regnet_x_1_6gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_1_6GF.tv2_in1k + filename: regnet_x_1_6gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_X_32GF.yaml b/mblt_vision/models/RegNet_X_32GF.yaml new file mode 100644 index 0000000..e6250a2 --- /dev/null +++ b/mblt_vision/models/RegNet_X_32GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_32GF.tv1_in1k + filename: regnet_x_32gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_32GF.tv2_in1k + filename: regnet_x_32gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_X_3_2GF.yaml b/mblt_vision/models/RegNet_X_3_2GF.yaml new file mode 100644 index 0000000..61f01ee --- /dev/null +++ b/mblt_vision/models/RegNet_X_3_2GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_3_2GF.tv1_in1k + filename: regnet_x_3_2gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_3_2GF.tv2_in1k + filename: regnet_x_3_2gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_X_400MF.yaml b/mblt_vision/models/RegNet_X_400MF.yaml new file mode 100644 index 0000000..e6039d0 --- /dev/null +++ b/mblt_vision/models/RegNet_X_400MF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_400MF.tv1_in1k + filename: regnet_x_400mf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_400MF.tv2_in1k + filename: regnet_x_400mf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_X_800MF.yaml b/mblt_vision/models/RegNet_X_800MF.yaml new file mode 100644 index 0000000..7952fea --- /dev/null +++ b/mblt_vision/models/RegNet_X_800MF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_800MF.tv1_in1k + filename: regnet_x_800mf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_800MF.tv2_in1k + filename: regnet_x_800mf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_X_8GF.yaml b/mblt_vision/models/RegNet_X_8GF.yaml new file mode 100644 index 0000000..8033e6e --- /dev/null +++ b/mblt_vision/models/RegNet_X_8GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_X_8GF.tv1_in1k + filename: regnet_x_8gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_X_8GF.tv2_in1k + filename: regnet_x_8gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_Y_16GF.yaml b/mblt_vision/models/RegNet_Y_16GF.yaml new file mode 100644 index 0000000..67a39e5 --- /dev/null +++ b/mblt_vision/models/RegNet_Y_16GF.yaml @@ -0,0 +1,76 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_16GF.tv1_in1k + filename: regnet_y_16gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_16GF.tv2_in1k + filename: regnet_y_16gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear +IMAGENET1K_SWAG_E2E_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_16GF.swag_e2e_in1k + filename: regnet_y_16gf_IMAGENET1K_SWAG_E2E_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_SWAG_LINEAR_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_16GF.swag_linear_in1k + filename: regnet_y_16gf_IMAGENET1K_SWAG_LINEAR_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RegNet_Y_1_6GF.yaml b/mblt_vision/models/RegNet_Y_1_6GF.yaml new file mode 100644 index 0000000..216c58d --- /dev/null +++ b/mblt_vision/models/RegNet_Y_1_6GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_1_6GF.tv1_in1k + filename: regnet_y_1_6gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_1_6GF.tv2_in1k + filename: regnet_y_1_6gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_Y_32GF.yaml b/mblt_vision/models/RegNet_Y_32GF.yaml new file mode 100644 index 0000000..ddeab69 --- /dev/null +++ b/mblt_vision/models/RegNet_Y_32GF.yaml @@ -0,0 +1,76 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_32GF.tv1_in1k + filename: regnet_y_32gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_32GF.tv2_in1k + filename: regnet_y_32gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear +IMAGENET1K_SWAG_E2E_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_32GF.swag_e2e_in1k + filename: regnet_y_32gf_IMAGENET1K_SWAG_E2E_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_SWAG_LINEAR_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_32GF.swag_linear_in1k + filename: regnet_y_32gf_IMAGENET1K_SWAG_LINEAR_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RegNet_Y_3_2GF.yaml b/mblt_vision/models/RegNet_Y_3_2GF.yaml new file mode 100644 index 0000000..ed10c4f --- /dev/null +++ b/mblt_vision/models/RegNet_Y_3_2GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_3_2GF.tv1_in1k + filename: regnet_y_3_2gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_3_2GF.tv2_in1k + filename: regnet_y_3_2gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_Y_400MF.yaml b/mblt_vision/models/RegNet_Y_400MF.yaml new file mode 100644 index 0000000..477037f --- /dev/null +++ b/mblt_vision/models/RegNet_Y_400MF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_400MF.tv1_in1k + filename: regnet_y_400mf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_400MF.tv2_in1k + filename: regnet_y_400mf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_Y_800MF.yaml b/mblt_vision/models/RegNet_Y_800MF.yaml new file mode 100644 index 0000000..5df53df --- /dev/null +++ b/mblt_vision/models/RegNet_Y_800MF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_800MF.tv1_in1k + filename: regnet_y_800mf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_800MF.tv2_in1k + filename: regnet_y_800mf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RegNet_Y_8GF.yaml b/mblt_vision/models/RegNet_Y_8GF.yaml new file mode 100644 index 0000000..20553d9 --- /dev/null +++ b/mblt_vision/models/RegNet_Y_8GF.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/RegNet_Y_8GF.tv1_in1k + filename: regnet_y_8gf_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/RegNet_Y_8GF.tv2_in1k + filename: regnet_y_8gf_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/RepViT_M0_9.yaml b/mblt_vision/models/RepViT_M0_9.yaml new file mode 100644 index 0000000..e2d2bf4 --- /dev/null +++ b/mblt_vision/models/RepViT_M0_9.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M0_9 + filename: repvit_m0_9.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M1.yaml b/mblt_vision/models/RepViT_M1.yaml new file mode 100644 index 0000000..0ce92f7 --- /dev/null +++ b/mblt_vision/models/RepViT_M1.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M1 + filename: repvit_m1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M1_0.yaml b/mblt_vision/models/RepViT_M1_0.yaml new file mode 100644 index 0000000..80f92b2 --- /dev/null +++ b/mblt_vision/models/RepViT_M1_0.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M1_0 + filename: repvit_m1_0.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M1_1.yaml b/mblt_vision/models/RepViT_M1_1.yaml new file mode 100644 index 0000000..31b36e3 --- /dev/null +++ b/mblt_vision/models/RepViT_M1_1.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M1_1 + filename: repvit_m1_1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M1_5.yaml b/mblt_vision/models/RepViT_M1_5.yaml new file mode 100644 index 0000000..8ce2309 --- /dev/null +++ b/mblt_vision/models/RepViT_M1_5.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M1_5 + filename: repvit_m1_5.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M2.yaml b/mblt_vision/models/RepViT_M2.yaml new file mode 100644 index 0000000..1c3baef --- /dev/null +++ b/mblt_vision/models/RepViT_M2.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M2 + filename: repvit_m2.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M2_3.yaml b/mblt_vision/models/RepViT_M2_3.yaml new file mode 100644 index 0000000..7ce1790 --- /dev/null +++ b/mblt_vision/models/RepViT_M2_3.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M2_3 + filename: repvit_m2_3.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/RepViT_M3.yaml b/mblt_vision/models/RepViT_M3.yaml new file mode 100644 index 0000000..ed8aaea --- /dev/null +++ b/mblt_vision/models/RepViT_M3.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/RepViT_M3 + filename: repvit_m3.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 236 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ResNeXt101_32X8D.yaml b/mblt_vision/models/ResNeXt101_32X8D.yaml new file mode 100644 index 0000000..a996ce1 --- /dev/null +++ b/mblt_vision/models/ResNeXt101_32X8D.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNeXt101_32X8D.tv1_in1k + filename: resnext101_32x8d_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/ResNeXt101_32X8D.tv2_in1k + filename: resnext101_32x8d_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/ResNeXt101_64X4D.yaml b/mblt_vision/models/ResNeXt101_64X4D.yaml new file mode 100644 index 0000000..bfe20ac --- /dev/null +++ b/mblt_vision/models/ResNeXt101_64X4D.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNeXt101_64X4D + filename: resnext101_64x4d_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ResNeXt50_32X4D.yaml b/mblt_vision/models/ResNeXt50_32X4D.yaml new file mode 100644 index 0000000..c3c8265 --- /dev/null +++ b/mblt_vision/models/ResNeXt50_32X4D.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNeXt50_32X4D.tv1_in1k + filename: resnext50_32x4d_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/ResNeXt50_32X4D.tv2_in1k + filename: resnext50_32x4d_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/ResNet101.yaml b/mblt_vision/models/ResNet101.yaml new file mode 100644 index 0000000..2367602 --- /dev/null +++ b/mblt_vision/models/ResNet101.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNet101.tv1_in1k + filename: resnet101_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/ResNet101.tv2_in1k + filename: resnet101_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/ResNet152.yaml b/mblt_vision/models/ResNet152.yaml new file mode 100644 index 0000000..af6e238 --- /dev/null +++ b/mblt_vision/models/ResNet152.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNet152.tv1_in1k + filename: resnet152_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/ResNet152.tv2_in1k + filename: resnet152_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/ResNet18.yaml b/mblt_vision/models/ResNet18.yaml new file mode 100644 index 0000000..4d64cf1 --- /dev/null +++ b/mblt_vision/models/ResNet18.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNet18 + filename: resnet18_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ResNet34.yaml b/mblt_vision/models/ResNet34.yaml new file mode 100644 index 0000000..b24ca76 --- /dev/null +++ b/mblt_vision/models/ResNet34.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNet34 + filename: resnet34_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ResNet50.yaml b/mblt_vision/models/ResNet50.yaml new file mode 100644 index 0000000..1459e8f --- /dev/null +++ b/mblt_vision/models/ResNet50.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ResNet50.tv1_in1k + filename: resnet50_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/ResNet50.tv2_in1k + filename: resnet50_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/ShuffleNet_V2_X0_5.yaml b/mblt_vision/models/ShuffleNet_V2_X0_5.yaml new file mode 100644 index 0000000..047e802 --- /dev/null +++ b/mblt_vision/models/ShuffleNet_V2_X0_5.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ShuffleNet_V2_X0_5 + filename: shufflenet_v2_x0_5_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ShuffleNet_V2_X1_0.yaml b/mblt_vision/models/ShuffleNet_V2_X1_0.yaml new file mode 100644 index 0000000..98b62ee --- /dev/null +++ b/mblt_vision/models/ShuffleNet_V2_X1_0.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ShuffleNet_V2_X1_0 + filename: shufflenet_v2_x1_0_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ShuffleNet_V2_X1_5.yaml b/mblt_vision/models/ShuffleNet_V2_X1_5.yaml new file mode 100644 index 0000000..e177b17 --- /dev/null +++ b/mblt_vision/models/ShuffleNet_V2_X1_5.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ShuffleNet_V2_X1_5 + filename: shufflenet_v2_x1_5_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ShuffleNet_V2_X2_0.yaml b/mblt_vision/models/ShuffleNet_V2_X2_0.yaml new file mode 100644 index 0000000..3656877 --- /dev/null +++ b/mblt_vision/models/ShuffleNet_V2_X2_0.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ShuffleNet_V2_X2_0 + filename: shufflenet_v2_x2_0_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/SqueezeNet1_0.yaml b/mblt_vision/models/SqueezeNet1_0.yaml new file mode 100644 index 0000000..38d77ef --- /dev/null +++ b/mblt_vision/models/SqueezeNet1_0.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/SqueezeNet1_0 + filename: squeezenet1_0_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/SqueezeNet1_1.yaml b/mblt_vision/models/SqueezeNet1_1.yaml new file mode 100644 index 0000000..43ee001 --- /dev/null +++ b/mblt_vision/models/SqueezeNet1_1.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/SqueezeNet1_1 + filename: squeezenet1_1_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/Swin_B.yaml b/mblt_vision/models/Swin_B.yaml new file mode 100644 index 0000000..a591788 --- /dev/null +++ b/mblt_vision/models/Swin_B.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/Swin_B + filename: swin_b_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 238 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/Swin_S.yaml b/mblt_vision/models/Swin_S.yaml new file mode 100644 index 0000000..ffa04ac --- /dev/null +++ b/mblt_vision/models/Swin_S.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/Swin_S + filename: swin_s_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 246 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/Swin_T.yaml b/mblt_vision/models/Swin_T.yaml new file mode 100644 index 0000000..fbe120e --- /dev/null +++ b/mblt_vision/models/Swin_T.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/Swin_T + filename: swin_t_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 232 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG11.yaml b/mblt_vision/models/VGG11.yaml new file mode 100644 index 0000000..89bdad5 --- /dev/null +++ b/mblt_vision/models/VGG11.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG11 + filename: vgg11_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG11_BN.yaml b/mblt_vision/models/VGG11_BN.yaml new file mode 100644 index 0000000..983e185 --- /dev/null +++ b/mblt_vision/models/VGG11_BN.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG11_BN + filename: vgg11_bn_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG13.yaml b/mblt_vision/models/VGG13.yaml new file mode 100644 index 0000000..517e451 --- /dev/null +++ b/mblt_vision/models/VGG13.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG13 + filename: vgg13_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG13_BN.yaml b/mblt_vision/models/VGG13_BN.yaml new file mode 100644 index 0000000..1c0c89b --- /dev/null +++ b/mblt_vision/models/VGG13_BN.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG13_BN + filename: vgg13_bn_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG16.yaml b/mblt_vision/models/VGG16.yaml new file mode 100644 index 0000000..4284f44 --- /dev/null +++ b/mblt_vision/models/VGG16.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG16 + filename: vgg16_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG16_BN.yaml b/mblt_vision/models/VGG16_BN.yaml new file mode 100644 index 0000000..0decf5a --- /dev/null +++ b/mblt_vision/models/VGG16_BN.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG16_BN + filename: vgg16_bn_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG19.yaml b/mblt_vision/models/VGG19.yaml new file mode 100644 index 0000000..782a541 --- /dev/null +++ b/mblt_vision/models/VGG19.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG19 + filename: vgg19_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VGG19_BN.yaml b/mblt_vision/models/VGG19_BN.yaml new file mode 100644 index 0000000..716d500 --- /dev/null +++ b/mblt_vision/models/VGG19_BN.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/VGG19_BN + filename: vgg19_bn_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_B_16.yaml b/mblt_vision/models/ViT_B_16.yaml new file mode 100644 index 0000000..2111588 --- /dev/null +++ b/mblt_vision/models/ViT_B_16.yaml @@ -0,0 +1,67 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ViT_B_16.tv1_in1k + filename: vit_b_16_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_SWAG_E2E_V1: + file_cfg: + repo_id: mobilint/ViT_B_16.swag_e2e_in1k + filename: vit_b_16_IMAGENET1K_SWAG_E2E_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_SWAG_LINEAR_V1: + file_cfg: + repo_id: mobilint/ViT_B_16.swag_linear_in1k + filename: vit_b_16_IMAGENET1K_SWAG_LINEAR_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_B_32.yaml b/mblt_vision/models/ViT_B_32.yaml new file mode 100644 index 0000000..ef29cc0 --- /dev/null +++ b/mblt_vision/models/ViT_B_32.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ViT_B_32 + filename: vit_b_32_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Base_Patch16_224.yaml b/mblt_vision/models/ViT_Base_Patch16_224.yaml new file mode 100644 index 0000000..d4fb6c0 --- /dev/null +++ b/mblt_vision/models/ViT_Base_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Base_Patch16_224 + filename: vit_base_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Base_Patch16_384.yaml b/mblt_vision/models/ViT_Base_Patch16_384.yaml new file mode 100644 index 0000000..c4cc9c0 --- /dev/null +++ b/mblt_vision/models/ViT_Base_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Base_Patch16_384 + filename: vit_base_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Base_Patch32_224.yaml b/mblt_vision/models/ViT_Base_Patch32_224.yaml new file mode 100644 index 0000000..d4c7f1c --- /dev/null +++ b/mblt_vision/models/ViT_Base_Patch32_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Base_Patch32_224 + filename: vit_base_patch32_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Base_Patch32_384.yaml b/mblt_vision/models/ViT_Base_Patch32_384.yaml new file mode 100644 index 0000000..f6bcd9c --- /dev/null +++ b/mblt_vision/models/ViT_Base_Patch32_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Base_Patch32_384 + filename: vit_base_patch32_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Base_Patch8_224.yaml b/mblt_vision/models/ViT_Base_Patch8_224.yaml new file mode 100644 index 0000000..af6ea08 --- /dev/null +++ b/mblt_vision/models/ViT_Base_Patch8_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Base_Patch8_224 + filename: vit_base_patch8_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_L_16.yaml b/mblt_vision/models/ViT_L_16.yaml new file mode 100644 index 0000000..02c0bb4 --- /dev/null +++ b/mblt_vision/models/ViT_L_16.yaml @@ -0,0 +1,67 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ViT_L_16.tv1_in1k + filename: vit_l_16_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 242 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_SWAG_E2E_V1: + file_cfg: + repo_id: mobilint/ViT_L_16.swag_e2e_in1k + filename: vit_l_16_IMAGENET1K_SWAG_E2E_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 512 + interpolation: bicubic + CenterCrop: + size: + - 512 + - 512 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_SWAG_LINEAR_V1: + file_cfg: + repo_id: mobilint/ViT_L_16.swag_linear_in1k + filename: vit_l_16_IMAGENET1K_SWAG_LINEAR_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_L_32.yaml b/mblt_vision/models/ViT_L_32.yaml new file mode 100644 index 0000000..b91f32b --- /dev/null +++ b/mblt_vision/models/ViT_L_32.yaml @@ -0,0 +1,23 @@ +DEFAULT: IMAGENET1K_V1 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/ViT_L_32 + filename: vit_l_32_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Large_Patch16_224.yaml b/mblt_vision/models/ViT_Large_Patch16_224.yaml new file mode 100644 index 0000000..75a871b --- /dev/null +++ b/mblt_vision/models/ViT_Large_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Large_Patch16_224 + filename: vit_large_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Large_Patch16_384.yaml b/mblt_vision/models/ViT_Large_Patch16_384.yaml new file mode 100644 index 0000000..f871425 --- /dev/null +++ b/mblt_vision/models/ViT_Large_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Large_Patch16_384 + filename: vit_large_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Large_Patch32_384.yaml b/mblt_vision/models/ViT_Large_Patch32_384.yaml new file mode 100644 index 0000000..d4a7199 --- /dev/null +++ b/mblt_vision/models/ViT_Large_Patch32_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Large_Patch32_384 + filename: vit_large_patch32_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Small_Patch16_224.yaml b/mblt_vision/models/ViT_Small_Patch16_224.yaml new file mode 100644 index 0000000..a55c062 --- /dev/null +++ b/mblt_vision/models/ViT_Small_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Small_Patch16_224 + filename: vit_small_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Small_Patch16_384.yaml b/mblt_vision/models/ViT_Small_Patch16_384.yaml new file mode 100644 index 0000000..99dff13 --- /dev/null +++ b/mblt_vision/models/ViT_Small_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Small_Patch16_384 + filename: vit_small_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Small_Patch32_224.yaml b/mblt_vision/models/ViT_Small_Patch32_224.yaml new file mode 100644 index 0000000..960ac46 --- /dev/null +++ b/mblt_vision/models/ViT_Small_Patch32_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Small_Patch32_224 + filename: vit_small_patch32_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Small_Patch32_384.yaml b/mblt_vision/models/ViT_Small_Patch32_384.yaml new file mode 100644 index 0000000..f634c45 --- /dev/null +++ b/mblt_vision/models/ViT_Small_Patch32_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Small_Patch32_384 + filename: vit_small_patch32_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Tiny_Patch16_224.yaml b/mblt_vision/models/ViT_Tiny_Patch16_224.yaml new file mode 100644 index 0000000..c5cb9af --- /dev/null +++ b/mblt_vision/models/ViT_Tiny_Patch16_224.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Tiny_Patch16_224 + filename: vit_tiny_patch16_224.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/ViT_Tiny_Patch16_384.yaml b/mblt_vision/models/ViT_Tiny_Patch16_384.yaml new file mode 100644 index 0000000..5ddd8ba --- /dev/null +++ b/mblt_vision/models/ViT_Tiny_Patch16_384.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/ViT_Tiny_Patch16_384 + filename: vit_tiny_patch16_384.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 384 + interpolation: bicubic + CenterCrop: + size: + - 384 + - 384 + SetOrder: + shape: HWC + Normalize: + style: tf + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VisFormer_Small.yaml b/mblt_vision/models/VisFormer_Small.yaml new file mode 100644 index 0000000..5a9e6a3 --- /dev/null +++ b/mblt_vision/models/VisFormer_Small.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/VisFormer_Small + filename: visformer_small.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/VisFormer_Tiny.yaml b/mblt_vision/models/VisFormer_Tiny.yaml new file mode 100644 index 0000000..ab8609e --- /dev/null +++ b/mblt_vision/models/VisFormer_Tiny.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/VisFormer_Tiny + filename: visformer_tiny.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 249 + interpolation: bicubic + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/Wide_ResNet101_2.yaml b/mblt_vision/models/Wide_ResNet101_2.yaml new file mode 100644 index 0000000..decd589 --- /dev/null +++ b/mblt_vision/models/Wide_ResNet101_2.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/Wide_ResNet101_2.tv1_in1k + filename: wide_resnet101_2_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/Wide_ResNet101_2.tv2_in1k + filename: wide_resnet101_2_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/Wide_ResNet50_2.yaml b/mblt_vision/models/Wide_ResNet50_2.yaml new file mode 100644 index 0000000..d1fd424 --- /dev/null +++ b/mblt_vision/models/Wide_ResNet50_2.yaml @@ -0,0 +1,32 @@ +DEFAULT: IMAGENET1K_V2 +IMAGENET1K_V1: + file_cfg: + repo_id: mobilint/Wide_ResNet50_2.tv1_in1k + filename: wide_resnet50_2_IMAGENET1K_V1.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 256 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet +IMAGENET1K_V2: + update: IMAGENET1K_V1 + file_cfg: + repo_id: mobilint/Wide_ResNet50_2.tv2_in1k + filename: wide_resnet50_2_IMAGENET1K_V2.mxq + pre_cfg: + Resize: + size: 232 + interpolation: bilinear diff --git a/mblt_vision/models/YOLO11l-cls.yaml b/mblt_vision/models/YOLO11l-cls.yaml new file mode 100644 index 0000000..f0d9845 --- /dev/null +++ b/mblt_vision/models/YOLO11l-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11l-cls + filename: yolo11l-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO11l-face.yaml b/mblt_vision/models/YOLO11l-face.yaml new file mode 100644 index 0000000..8c6982a --- /dev/null +++ b/mblt_vision/models/YOLO11l-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11l-face + filename: yolo11l-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11l-obb.yaml b/mblt_vision/models/YOLO11l-obb.yaml new file mode 100644 index 0000000..bb8d4af --- /dev/null +++ b/mblt_vision/models/YOLO11l-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11l-obb + filename: yolo11l-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11l-pose.yaml b/mblt_vision/models/YOLO11l-pose.yaml new file mode 100644 index 0000000..6bda47d --- /dev/null +++ b/mblt_vision/models/YOLO11l-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11l-pose + filename: yolo11l-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11l-seg.yaml b/mblt_vision/models/YOLO11l-seg.yaml new file mode 100644 index 0000000..55a7343 --- /dev/null +++ b/mblt_vision/models/YOLO11l-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11l-seg + filename: yolo11l-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11l.yaml b/mblt_vision/models/YOLO11l.yaml new file mode 100644 index 0000000..84faf43 --- /dev/null +++ b/mblt_vision/models/YOLO11l.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11l + filename: yolo11l.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11m-cls.yaml b/mblt_vision/models/YOLO11m-cls.yaml new file mode 100644 index 0000000..4dd25e6 --- /dev/null +++ b/mblt_vision/models/YOLO11m-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11m-cls + filename: yolo11m-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO11m-face.yaml b/mblt_vision/models/YOLO11m-face.yaml new file mode 100644 index 0000000..93d007e --- /dev/null +++ b/mblt_vision/models/YOLO11m-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11m-face + filename: yolo11m-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11m-obb.yaml b/mblt_vision/models/YOLO11m-obb.yaml new file mode 100644 index 0000000..91fccc5 --- /dev/null +++ b/mblt_vision/models/YOLO11m-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11m-obb + filename: yolo11m-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11m-pose.yaml b/mblt_vision/models/YOLO11m-pose.yaml new file mode 100644 index 0000000..bf7bc8d --- /dev/null +++ b/mblt_vision/models/YOLO11m-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11m-pose + filename: yolo11m-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11m-seg.yaml b/mblt_vision/models/YOLO11m-seg.yaml new file mode 100644 index 0000000..46f3e45 --- /dev/null +++ b/mblt_vision/models/YOLO11m-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11m-seg + filename: yolo11m-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11m.yaml b/mblt_vision/models/YOLO11m.yaml new file mode 100644 index 0000000..4831817 --- /dev/null +++ b/mblt_vision/models/YOLO11m.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11m + filename: yolo11m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11n-cls.yaml b/mblt_vision/models/YOLO11n-cls.yaml new file mode 100644 index 0000000..e7658ac --- /dev/null +++ b/mblt_vision/models/YOLO11n-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11n-cls + filename: yolo11n-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO11n-face.yaml b/mblt_vision/models/YOLO11n-face.yaml new file mode 100644 index 0000000..0fd8395 --- /dev/null +++ b/mblt_vision/models/YOLO11n-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11n-face + filename: yolo11n-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11n-obb.yaml b/mblt_vision/models/YOLO11n-obb.yaml new file mode 100644 index 0000000..c2f00a9 --- /dev/null +++ b/mblt_vision/models/YOLO11n-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11n-obb + filename: yolo11n-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11n-pose.yaml b/mblt_vision/models/YOLO11n-pose.yaml new file mode 100644 index 0000000..fe7b8dc --- /dev/null +++ b/mblt_vision/models/YOLO11n-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11n-pose + filename: yolo11n-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11n-seg.yaml b/mblt_vision/models/YOLO11n-seg.yaml new file mode 100644 index 0000000..736e6bf --- /dev/null +++ b/mblt_vision/models/YOLO11n-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11n-seg + filename: yolo11n-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11n.yaml b/mblt_vision/models/YOLO11n.yaml new file mode 100644 index 0000000..809df97 --- /dev/null +++ b/mblt_vision/models/YOLO11n.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11n + filename: yolo11n.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11s-cls.yaml b/mblt_vision/models/YOLO11s-cls.yaml new file mode 100644 index 0000000..fd53c03 --- /dev/null +++ b/mblt_vision/models/YOLO11s-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11s-cls + filename: yolo11s-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO11s-face.yaml b/mblt_vision/models/YOLO11s-face.yaml new file mode 100644 index 0000000..ec966ef --- /dev/null +++ b/mblt_vision/models/YOLO11s-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11s-face + filename: yolo11s-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11s-obb.yaml b/mblt_vision/models/YOLO11s-obb.yaml new file mode 100644 index 0000000..0aae543 --- /dev/null +++ b/mblt_vision/models/YOLO11s-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11s-obb + filename: yolo11s-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11s-pose.yaml b/mblt_vision/models/YOLO11s-pose.yaml new file mode 100644 index 0000000..b0ed955 --- /dev/null +++ b/mblt_vision/models/YOLO11s-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11s-pose + filename: yolo11s-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11s-seg.yaml b/mblt_vision/models/YOLO11s-seg.yaml new file mode 100644 index 0000000..65a6554 --- /dev/null +++ b/mblt_vision/models/YOLO11s-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11s-seg + filename: yolo11s-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11s.yaml b/mblt_vision/models/YOLO11s.yaml new file mode 100644 index 0000000..4dc6a4f --- /dev/null +++ b/mblt_vision/models/YOLO11s.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11s + filename: yolo11s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11x-cls.yaml b/mblt_vision/models/YOLO11x-cls.yaml new file mode 100644 index 0000000..c014ad6 --- /dev/null +++ b/mblt_vision/models/YOLO11x-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11x-cls + filename: yolo11x-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO11x-obb.yaml b/mblt_vision/models/YOLO11x-obb.yaml new file mode 100644 index 0000000..0d51b51 --- /dev/null +++ b/mblt_vision/models/YOLO11x-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11x-obb + filename: yolo11x-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11x-pose.yaml b/mblt_vision/models/YOLO11x-pose.yaml new file mode 100644 index 0000000..46e3132 --- /dev/null +++ b/mblt_vision/models/YOLO11x-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11x-pose + filename: yolo11x-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11x-seg.yaml b/mblt_vision/models/YOLO11x-seg.yaml new file mode 100644 index 0000000..fa7d992 --- /dev/null +++ b/mblt_vision/models/YOLO11x-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11x-seg + filename: yolo11x-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO11x.yaml b/mblt_vision/models/YOLO11x.yaml new file mode 100644 index 0000000..c3008c9 --- /dev/null +++ b/mblt_vision/models/YOLO11x.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO11x + filename: yolo11x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12l-face.yaml b/mblt_vision/models/YOLO12l-face.yaml new file mode 100644 index 0000000..6a0ccbf --- /dev/null +++ b/mblt_vision/models/YOLO12l-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12l-face + filename: yolo12l-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12l-seg.yaml b/mblt_vision/models/YOLO12l-seg.yaml new file mode 100644 index 0000000..e67c1da --- /dev/null +++ b/mblt_vision/models/YOLO12l-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12l-seg + filename: yolo12l-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12l.yaml b/mblt_vision/models/YOLO12l.yaml new file mode 100644 index 0000000..63c3059 --- /dev/null +++ b/mblt_vision/models/YOLO12l.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12l + filename: yolo12l.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12m-face.yaml b/mblt_vision/models/YOLO12m-face.yaml new file mode 100644 index 0000000..5d8340d --- /dev/null +++ b/mblt_vision/models/YOLO12m-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12m-face + filename: yolo12m-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12m-seg.yaml b/mblt_vision/models/YOLO12m-seg.yaml new file mode 100644 index 0000000..2572df0 --- /dev/null +++ b/mblt_vision/models/YOLO12m-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12m-seg + filename: yolo12m-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12m.yaml b/mblt_vision/models/YOLO12m.yaml new file mode 100644 index 0000000..c6f4fa3 --- /dev/null +++ b/mblt_vision/models/YOLO12m.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12m + filename: yolo12m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12n-face.yaml b/mblt_vision/models/YOLO12n-face.yaml new file mode 100644 index 0000000..686e61c --- /dev/null +++ b/mblt_vision/models/YOLO12n-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12n-face + filename: yolo12n-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12n-seg.yaml b/mblt_vision/models/YOLO12n-seg.yaml new file mode 100644 index 0000000..0672bd8 --- /dev/null +++ b/mblt_vision/models/YOLO12n-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12n-seg + filename: yolo12n-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12n.yaml b/mblt_vision/models/YOLO12n.yaml new file mode 100644 index 0000000..86897b6 --- /dev/null +++ b/mblt_vision/models/YOLO12n.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12n + filename: yolo12n.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12s-face.yaml b/mblt_vision/models/YOLO12s-face.yaml new file mode 100644 index 0000000..7573356 --- /dev/null +++ b/mblt_vision/models/YOLO12s-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12s-face + filename: yolo12s-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12s-seg.yaml b/mblt_vision/models/YOLO12s-seg.yaml new file mode 100644 index 0000000..8c9fb4e --- /dev/null +++ b/mblt_vision/models/YOLO12s-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12s-seg + filename: yolo12s-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12s.yaml b/mblt_vision/models/YOLO12s.yaml new file mode 100644 index 0000000..5e1c2d0 --- /dev/null +++ b/mblt_vision/models/YOLO12s.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12s + filename: yolo12s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12x-seg.yaml b/mblt_vision/models/YOLO12x-seg.yaml new file mode 100644 index 0000000..1671116 --- /dev/null +++ b/mblt_vision/models/YOLO12x-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12x-seg + filename: yolo12x-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO12x.yaml b/mblt_vision/models/YOLO12x.yaml new file mode 100644 index 0000000..3484f8a --- /dev/null +++ b/mblt_vision/models/YOLO12x.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO12x + filename: yolo12x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26l-cls.yaml b/mblt_vision/models/YOLO26l-cls.yaml new file mode 100644 index 0000000..d492a4c --- /dev/null +++ b/mblt_vision/models/YOLO26l-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-cls + filename: yolo26l-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO26l-depth.yaml b/mblt_vision/models/YOLO26l-depth.yaml new file mode 100644 index 0000000..226a1d8 --- /dev/null +++ b/mblt_vision/models/YOLO26l-depth.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-depth + filename: yolo26l-depth.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [768, 768] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: depth_estimation + dataset: nyu-depth diff --git a/mblt_vision/models/YOLO26l-distill.yaml b/mblt_vision/models/YOLO26l-distill.yaml new file mode 100644 index 0000000..52b0869 --- /dev/null +++ b/mblt_vision/models/YOLO26l-distill.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-distill + filename: yolo26l-distill.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26l-obb.yaml b/mblt_vision/models/YOLO26l-obb.yaml new file mode 100644 index 0000000..df61c05 --- /dev/null +++ b/mblt_vision/models/YOLO26l-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-obb + filename: yolo26l-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + dflfree: true + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26l-pose.yaml b/mblt_vision/models/YOLO26l-pose.yaml new file mode 100644 index 0000000..17fc25d --- /dev/null +++ b/mblt_vision/models/YOLO26l-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-pose + filename: yolo26l-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26l-seg.yaml b/mblt_vision/models/YOLO26l-seg.yaml new file mode 100644 index 0000000..7b861fb --- /dev/null +++ b/mblt_vision/models/YOLO26l-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-seg + filename: yolo26l-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26l-sem-ade20k.yaml b/mblt_vision/models/YOLO26l-sem-ade20k.yaml new file mode 100644 index 0000000..7fa0a5b --- /dev/null +++ b/mblt_vision/models/YOLO26l-sem-ade20k.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-sem-ade20k + filename: yolo26l-sem-ade20k.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [640, 640] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: ade20k diff --git a/mblt_vision/models/YOLO26l-sem.yaml b/mblt_vision/models/YOLO26l-sem.yaml new file mode 100644 index 0000000..822123e --- /dev/null +++ b/mblt_vision/models/YOLO26l-sem.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l-sem + filename: yolo26l-sem.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [1024, 2048] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: cityscapes diff --git a/mblt_vision/models/YOLO26l.yaml b/mblt_vision/models/YOLO26l.yaml new file mode 100644 index 0000000..3df6898 --- /dev/null +++ b/mblt_vision/models/YOLO26l.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26l + filename: yolo26l.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26m-cls.yaml b/mblt_vision/models/YOLO26m-cls.yaml new file mode 100644 index 0000000..57abe6b --- /dev/null +++ b/mblt_vision/models/YOLO26m-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-cls + filename: yolo26m-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO26m-depth.yaml b/mblt_vision/models/YOLO26m-depth.yaml new file mode 100644 index 0000000..ace9c16 --- /dev/null +++ b/mblt_vision/models/YOLO26m-depth.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-depth + filename: yolo26m-depth.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [768, 768] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: depth_estimation + dataset: nyu-depth diff --git a/mblt_vision/models/YOLO26m-distill.yaml b/mblt_vision/models/YOLO26m-distill.yaml new file mode 100644 index 0000000..79e8550 --- /dev/null +++ b/mblt_vision/models/YOLO26m-distill.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-distill + filename: yolo26m-distill.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26m-obb.yaml b/mblt_vision/models/YOLO26m-obb.yaml new file mode 100644 index 0000000..f6353fd --- /dev/null +++ b/mblt_vision/models/YOLO26m-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-obb + filename: yolo26m-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + dflfree: true + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26m-pose.yaml b/mblt_vision/models/YOLO26m-pose.yaml new file mode 100644 index 0000000..eddd6fa --- /dev/null +++ b/mblt_vision/models/YOLO26m-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-pose + filename: yolo26m-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26m-seg.yaml b/mblt_vision/models/YOLO26m-seg.yaml new file mode 100644 index 0000000..32be031 --- /dev/null +++ b/mblt_vision/models/YOLO26m-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-seg + filename: yolo26m-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26m-sem-ade20k.yaml b/mblt_vision/models/YOLO26m-sem-ade20k.yaml new file mode 100644 index 0000000..1ca8020 --- /dev/null +++ b/mblt_vision/models/YOLO26m-sem-ade20k.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-sem-ade20k + filename: yolo26m-sem-ade20k.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [640, 640] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: ade20k diff --git a/mblt_vision/models/YOLO26m-sem.yaml b/mblt_vision/models/YOLO26m-sem.yaml new file mode 100644 index 0000000..00ac704 --- /dev/null +++ b/mblt_vision/models/YOLO26m-sem.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m-sem + filename: yolo26m-sem.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [1024, 2048] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: cityscapes diff --git a/mblt_vision/models/YOLO26m.yaml b/mblt_vision/models/YOLO26m.yaml new file mode 100644 index 0000000..817db03 --- /dev/null +++ b/mblt_vision/models/YOLO26m.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26m + filename: yolo26m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26n-cls.yaml b/mblt_vision/models/YOLO26n-cls.yaml new file mode 100644 index 0000000..df2663e --- /dev/null +++ b/mblt_vision/models/YOLO26n-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-cls + filename: yolo26n-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO26n-depth.yaml b/mblt_vision/models/YOLO26n-depth.yaml new file mode 100644 index 0000000..fc31ab5 --- /dev/null +++ b/mblt_vision/models/YOLO26n-depth.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-depth + filename: yolo26n-depth.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [768, 768] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: depth_estimation + dataset: nyu-depth diff --git a/mblt_vision/models/YOLO26n-distill.yaml b/mblt_vision/models/YOLO26n-distill.yaml new file mode 100644 index 0000000..b6e4d14 --- /dev/null +++ b/mblt_vision/models/YOLO26n-distill.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-distill + filename: yolo26n-distill.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26n-obb.yaml b/mblt_vision/models/YOLO26n-obb.yaml new file mode 100644 index 0000000..aa92ffe --- /dev/null +++ b/mblt_vision/models/YOLO26n-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-obb + filename: yolo26n-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + dflfree: true + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26n-pose.yaml b/mblt_vision/models/YOLO26n-pose.yaml new file mode 100644 index 0000000..110b715 --- /dev/null +++ b/mblt_vision/models/YOLO26n-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-pose + filename: yolo26n-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26n-seg.yaml b/mblt_vision/models/YOLO26n-seg.yaml new file mode 100644 index 0000000..d67e642 --- /dev/null +++ b/mblt_vision/models/YOLO26n-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-seg + filename: yolo26n-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26n-sem-ade20k.yaml b/mblt_vision/models/YOLO26n-sem-ade20k.yaml new file mode 100644 index 0000000..9223741 --- /dev/null +++ b/mblt_vision/models/YOLO26n-sem-ade20k.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-sem-ade20k + filename: yolo26n-sem-ade20k.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [640, 640] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: ade20k diff --git a/mblt_vision/models/YOLO26n-sem.yaml b/mblt_vision/models/YOLO26n-sem.yaml new file mode 100644 index 0000000..c97909a --- /dev/null +++ b/mblt_vision/models/YOLO26n-sem.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n-sem + filename: yolo26n-sem.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [1024, 2048] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: cityscapes diff --git a/mblt_vision/models/YOLO26n.yaml b/mblt_vision/models/YOLO26n.yaml new file mode 100644 index 0000000..c1a3154 --- /dev/null +++ b/mblt_vision/models/YOLO26n.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26n + filename: yolo26n.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26s-cls.yaml b/mblt_vision/models/YOLO26s-cls.yaml new file mode 100644 index 0000000..9c8074a --- /dev/null +++ b/mblt_vision/models/YOLO26s-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-cls + filename: yolo26s-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO26s-depth.yaml b/mblt_vision/models/YOLO26s-depth.yaml new file mode 100644 index 0000000..8e9a89a --- /dev/null +++ b/mblt_vision/models/YOLO26s-depth.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-depth + filename: yolo26s-depth.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [768, 768] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: depth_estimation + dataset: nyu-depth diff --git a/mblt_vision/models/YOLO26s-distill.yaml b/mblt_vision/models/YOLO26s-distill.yaml new file mode 100644 index 0000000..a29d729 --- /dev/null +++ b/mblt_vision/models/YOLO26s-distill.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-distill + filename: yolo26s-distill.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26s-obb.yaml b/mblt_vision/models/YOLO26s-obb.yaml new file mode 100644 index 0000000..8d4d302 --- /dev/null +++ b/mblt_vision/models/YOLO26s-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-obb + filename: yolo26s-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + dflfree: true + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26s-pose.yaml b/mblt_vision/models/YOLO26s-pose.yaml new file mode 100644 index 0000000..6f1007d --- /dev/null +++ b/mblt_vision/models/YOLO26s-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-pose + filename: yolo26s-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26s-seg.yaml b/mblt_vision/models/YOLO26s-seg.yaml new file mode 100644 index 0000000..ccd86f4 --- /dev/null +++ b/mblt_vision/models/YOLO26s-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-seg + filename: yolo26s-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26s-sem-ade20k.yaml b/mblt_vision/models/YOLO26s-sem-ade20k.yaml new file mode 100644 index 0000000..cff9c9f --- /dev/null +++ b/mblt_vision/models/YOLO26s-sem-ade20k.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-sem-ade20k + filename: yolo26s-sem-ade20k.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [640, 640] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: ade20k diff --git a/mblt_vision/models/YOLO26s-sem.yaml b/mblt_vision/models/YOLO26s-sem.yaml new file mode 100644 index 0000000..5040771 --- /dev/null +++ b/mblt_vision/models/YOLO26s-sem.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s-sem + filename: yolo26s-sem.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [1024, 2048] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: cityscapes diff --git a/mblt_vision/models/YOLO26s.yaml b/mblt_vision/models/YOLO26s.yaml new file mode 100644 index 0000000..2167b14 --- /dev/null +++ b/mblt_vision/models/YOLO26s.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26s + filename: yolo26s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26x-cls.yaml b/mblt_vision/models/YOLO26x-cls.yaml new file mode 100644 index 0000000..03963e2 --- /dev/null +++ b/mblt_vision/models/YOLO26x-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-cls + filename: yolo26x-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLO26x-depth.yaml b/mblt_vision/models/YOLO26x-depth.yaml new file mode 100644 index 0000000..60cc02b --- /dev/null +++ b/mblt_vision/models/YOLO26x-depth.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-depth + filename: yolo26x-depth.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [768, 768] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: depth_estimation + dataset: nyu-depth diff --git a/mblt_vision/models/YOLO26x-distill.yaml b/mblt_vision/models/YOLO26x-distill.yaml new file mode 100644 index 0000000..4a5ba55 --- /dev/null +++ b/mblt_vision/models/YOLO26x-distill.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-distill + filename: yolo26x-distill.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26x-obb.yaml b/mblt_vision/models/YOLO26x-obb.yaml new file mode 100644 index 0000000..9e178d0 --- /dev/null +++ b/mblt_vision/models/YOLO26x-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-obb + filename: yolo26x-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + dflfree: true + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26x-pose.yaml b/mblt_vision/models/YOLO26x-pose.yaml new file mode 100644 index 0000000..2f3a9cf --- /dev/null +++ b/mblt_vision/models/YOLO26x-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-pose + filename: yolo26x-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26x-seg.yaml b/mblt_vision/models/YOLO26x-seg.yaml new file mode 100644 index 0000000..0469a92 --- /dev/null +++ b/mblt_vision/models/YOLO26x-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-seg + filename: yolo26x-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLO26x-sem-ade20k.yaml b/mblt_vision/models/YOLO26x-sem-ade20k.yaml new file mode 100644 index 0000000..38f3546 --- /dev/null +++ b/mblt_vision/models/YOLO26x-sem-ade20k.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-sem-ade20k + filename: yolo26x-sem-ade20k.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [640, 640] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: ade20k diff --git a/mblt_vision/models/YOLO26x-sem.yaml b/mblt_vision/models/YOLO26x-sem.yaml new file mode 100644 index 0000000..3826871 --- /dev/null +++ b/mblt_vision/models/YOLO26x-sem.yaml @@ -0,0 +1,17 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x-sem + filename: yolo26x-sem.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: [1024, 2048] + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: semantic_segmentation + dataset: cityscapes diff --git a/mblt_vision/models/YOLO26x.yaml b/mblt_vision/models/YOLO26x.yaml new file mode 100644 index 0000000..777fc7f --- /dev/null +++ b/mblt_vision/models/YOLO26x.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLO26x + filename: yolo26x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + dflfree: true + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10b.yaml b/mblt_vision/models/YOLOv10b.yaml new file mode 100644 index 0000000..013c1c4 --- /dev/null +++ b/mblt_vision/models/YOLOv10b.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10b + filename: yolov10b.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10l-face.yaml b/mblt_vision/models/YOLOv10l-face.yaml new file mode 100644 index 0000000..e9d9fde --- /dev/null +++ b/mblt_vision/models/YOLOv10l-face.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10l-face + filename: yolov10l-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10l.yaml b/mblt_vision/models/YOLOv10l.yaml new file mode 100644 index 0000000..a1ce2c8 --- /dev/null +++ b/mblt_vision/models/YOLOv10l.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10l + filename: yolov10l.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10m-face.yaml b/mblt_vision/models/YOLOv10m-face.yaml new file mode 100644 index 0000000..61b9df3 --- /dev/null +++ b/mblt_vision/models/YOLOv10m-face.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10m-face + filename: yolov10m-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10m.yaml b/mblt_vision/models/YOLOv10m.yaml new file mode 100644 index 0000000..cef0e70 --- /dev/null +++ b/mblt_vision/models/YOLOv10m.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10m + filename: yolov10m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10n-face.yaml b/mblt_vision/models/YOLOv10n-face.yaml new file mode 100644 index 0000000..0b3c266 --- /dev/null +++ b/mblt_vision/models/YOLOv10n-face.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10n-face + filename: yolov10n-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10n.yaml b/mblt_vision/models/YOLOv10n.yaml new file mode 100644 index 0000000..bd5585d --- /dev/null +++ b/mblt_vision/models/YOLOv10n.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10n + filename: yolov10n.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10s-face.yaml b/mblt_vision/models/YOLOv10s-face.yaml new file mode 100644 index 0000000..a66ebcb --- /dev/null +++ b/mblt_vision/models/YOLOv10s-face.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10s-face + filename: yolov10s-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10s.yaml b/mblt_vision/models/YOLOv10s.yaml new file mode 100644 index 0000000..374cdae --- /dev/null +++ b/mblt_vision/models/YOLOv10s.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10s + filename: yolov10s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv10x.yaml b/mblt_vision/models/YOLOv10x.yaml new file mode 100644 index 0000000..ce5a4c3 --- /dev/null +++ b/mblt_vision/models/YOLOv10x.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv10x + filename: yolov10x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + nmsfree: true + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv3-spp.yaml b/mblt_vision/models/YOLOv3-spp.yaml new file mode 100644 index 0000000..48e8c42 --- /dev/null +++ b/mblt_vision/models/YOLOv3-spp.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv3-spp + filename: yolov3-spp.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv3-sppu.yaml b/mblt_vision/models/YOLOv3-sppu.yaml new file mode 100644 index 0000000..cd2c845 --- /dev/null +++ b/mblt_vision/models/YOLOv3-sppu.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv3-sppu + filename: yolov3-sppu.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv3-tiny.yaml b/mblt_vision/models/YOLOv3-tiny.yaml new file mode 100644 index 0000000..047451b --- /dev/null +++ b/mblt_vision/models/YOLOv3-tiny.yaml @@ -0,0 +1,38 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv3-tiny + filename: yolov3-tiny.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 14 + - 23 + - 27 + - 37 + - 58 + - - 81 + - 82 + - 135 + - 169 + - 344 + - 319 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv3-tinyu.yaml b/mblt_vision/models/YOLOv3-tinyu.yaml new file mode 100644 index 0000000..a263aca --- /dev/null +++ b/mblt_vision/models/YOLOv3-tinyu.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv3-tinyu + filename: yolov3-tinyu.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 2 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv3.yaml b/mblt_vision/models/YOLOv3.yaml new file mode 100644 index 0000000..29be313 --- /dev/null +++ b/mblt_vision/models/YOLOv3.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv3 + filename: yolov3.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv3u.yaml b/mblt_vision/models/YOLOv3u.yaml new file mode 100644 index 0000000..bc1cf32 --- /dev/null +++ b/mblt_vision/models/YOLOv3u.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv3u + filename: yolov3u.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5l-cls.yaml b/mblt_vision/models/YOLOv5l-cls.yaml new file mode 100644 index 0000000..2f0a17c --- /dev/null +++ b/mblt_vision/models/YOLOv5l-cls.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5l-cls + filename: yolov5l-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/YOLOv5l-seg.yaml b/mblt_vision/models/YOLOv5l-seg.yaml new file mode 100644 index 0000000..ba386b8 --- /dev/null +++ b/mblt_vision/models/YOLOv5l-seg.yaml @@ -0,0 +1,45 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5l-seg + filename: yolov5l-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + n_extra: 32 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5l.yaml b/mblt_vision/models/YOLOv5l.yaml new file mode 100644 index 0000000..181aae6 --- /dev/null +++ b/mblt_vision/models/YOLOv5l.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5l + filename: yolov5l.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5l6.yaml b/mblt_vision/models/YOLOv5l6.yaml new file mode 100644 index 0000000..ad68e5c --- /dev/null +++ b/mblt_vision/models/YOLOv5l6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5l6 + filename: yolov5l6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5l6u.yaml b/mblt_vision/models/YOLOv5l6u.yaml new file mode 100644 index 0000000..380b80a --- /dev/null +++ b/mblt_vision/models/YOLOv5l6u.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5l6u + filename: yolov5l6u.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 4 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5lu.yaml b/mblt_vision/models/YOLOv5lu.yaml new file mode 100644 index 0000000..b1dc1e4 --- /dev/null +++ b/mblt_vision/models/YOLOv5lu.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5lu + filename: yolov5lu.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5m-cls.yaml b/mblt_vision/models/YOLOv5m-cls.yaml new file mode 100644 index 0000000..8fdb619 --- /dev/null +++ b/mblt_vision/models/YOLOv5m-cls.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5m-cls + filename: yolov5m-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/YOLOv5m-seg.yaml b/mblt_vision/models/YOLOv5m-seg.yaml new file mode 100644 index 0000000..6696be4 --- /dev/null +++ b/mblt_vision/models/YOLOv5m-seg.yaml @@ -0,0 +1,45 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5m-seg + filename: yolov5m-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + n_extra: 32 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5m.yaml b/mblt_vision/models/YOLOv5m.yaml new file mode 100644 index 0000000..505d44f --- /dev/null +++ b/mblt_vision/models/YOLOv5m.yaml @@ -0,0 +1,45 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5m + filename: yolov5m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5m6.yaml b/mblt_vision/models/YOLOv5m6.yaml new file mode 100644 index 0000000..4a21683 --- /dev/null +++ b/mblt_vision/models/YOLOv5m6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5m6 + filename: yolov5m6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5m6u.yaml b/mblt_vision/models/YOLOv5m6u.yaml new file mode 100644 index 0000000..fb12e16 --- /dev/null +++ b/mblt_vision/models/YOLOv5m6u.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5m6u + filename: yolov5m6u.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 4 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5mu.yaml b/mblt_vision/models/YOLOv5mu.yaml new file mode 100644 index 0000000..6fa0880 --- /dev/null +++ b/mblt_vision/models/YOLOv5mu.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5mu + filename: yolov5mu.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5n-cls.yaml b/mblt_vision/models/YOLOv5n-cls.yaml new file mode 100644 index 0000000..1447532 --- /dev/null +++ b/mblt_vision/models/YOLOv5n-cls.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5n-cls + filename: yolov5n-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/YOLOv5n-seg.yaml b/mblt_vision/models/YOLOv5n-seg.yaml new file mode 100644 index 0000000..1437343 --- /dev/null +++ b/mblt_vision/models/YOLOv5n-seg.yaml @@ -0,0 +1,45 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5n-seg + filename: yolov5n-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + n_extra: 32 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5n.yaml b/mblt_vision/models/YOLOv5n.yaml new file mode 100644 index 0000000..f135ba1 --- /dev/null +++ b/mblt_vision/models/YOLOv5n.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5n + filename: yolov5n.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5n6.yaml b/mblt_vision/models/YOLOv5n6.yaml new file mode 100644 index 0000000..b0c86b8 --- /dev/null +++ b/mblt_vision/models/YOLOv5n6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5n6 + filename: yolov5n6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5n6u.yaml b/mblt_vision/models/YOLOv5n6u.yaml new file mode 100644 index 0000000..207f82a --- /dev/null +++ b/mblt_vision/models/YOLOv5n6u.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5n6u + filename: yolov5n6u.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 4 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5nu.yaml b/mblt_vision/models/YOLOv5nu.yaml new file mode 100644 index 0000000..29872aa --- /dev/null +++ b/mblt_vision/models/YOLOv5nu.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5nu + filename: yolov5nu.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5s-cls.yaml b/mblt_vision/models/YOLOv5s-cls.yaml new file mode 100644 index 0000000..fc928f8 --- /dev/null +++ b/mblt_vision/models/YOLOv5s-cls.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5s-cls + filename: yolov5s-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/YOLOv5s-seg.yaml b/mblt_vision/models/YOLOv5s-seg.yaml new file mode 100644 index 0000000..e551e73 --- /dev/null +++ b/mblt_vision/models/YOLOv5s-seg.yaml @@ -0,0 +1,45 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5s-seg + filename: yolov5s-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + n_extra: 32 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5s.yaml b/mblt_vision/models/YOLOv5s.yaml new file mode 100644 index 0000000..b633770 --- /dev/null +++ b/mblt_vision/models/YOLOv5s.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5s + filename: yolov5s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5s6.yaml b/mblt_vision/models/YOLOv5s6.yaml new file mode 100644 index 0000000..d31e09a --- /dev/null +++ b/mblt_vision/models/YOLOv5s6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5s6 + filename: yolov5s6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5s6u.yaml b/mblt_vision/models/YOLOv5s6u.yaml new file mode 100644 index 0000000..cb4dbd2 --- /dev/null +++ b/mblt_vision/models/YOLOv5s6u.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5s6u + filename: yolov5s6u.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 4 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5su.yaml b/mblt_vision/models/YOLOv5su.yaml new file mode 100644 index 0000000..a7fe72d --- /dev/null +++ b/mblt_vision/models/YOLOv5su.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5su + filename: yolov5su.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5x-cls.yaml b/mblt_vision/models/YOLOv5x-cls.yaml new file mode 100644 index 0000000..93f8073 --- /dev/null +++ b/mblt_vision/models/YOLOv5x-cls.yaml @@ -0,0 +1,22 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5x-cls + filename: yolov5x-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: torch + post_cfg: + task: image_classification + dataset: imagenet diff --git a/mblt_vision/models/YOLOv5x-seg.yaml b/mblt_vision/models/YOLOv5x-seg.yaml new file mode 100644 index 0000000..efd2e09 --- /dev/null +++ b/mblt_vision/models/YOLOv5x-seg.yaml @@ -0,0 +1,45 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5x-seg + filename: yolov5x-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + n_extra: 32 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5x.yaml b/mblt_vision/models/YOLOv5x.yaml new file mode 100644 index 0000000..d20e4f1 --- /dev/null +++ b/mblt_vision/models/YOLOv5x.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5x + filename: yolov5x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 10 + - 13 + - 16 + - 30 + - 33 + - 23 + - - 30 + - 61 + - 62 + - 45 + - 59 + - 119 + - - 116 + - 90 + - 156 + - 198 + - 373 + - 326 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5x6.yaml b/mblt_vision/models/YOLOv5x6.yaml new file mode 100644 index 0000000..c9928f2 --- /dev/null +++ b/mblt_vision/models/YOLOv5x6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5x6 + filename: yolov5x6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.65 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5x6u.yaml b/mblt_vision/models/YOLOv5x6u.yaml new file mode 100644 index 0000000..5ceae88 --- /dev/null +++ b/mblt_vision/models/YOLOv5x6u.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5x6u + filename: yolov5x6u.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 4 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv5xu.yaml b/mblt_vision/models/YOLOv5xu.yaml new file mode 100644 index 0000000..2d155bf --- /dev/null +++ b/mblt_vision/models/YOLOv5xu.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv5xu + filename: yolov5xu.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv6m-face.yaml b/mblt_vision/models/YOLOv6m-face.yaml new file mode 100644 index 0000000..5c56065 --- /dev/null +++ b/mblt_vision/models/YOLOv6m-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv6m-face + filename: yolov6m-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv6n-face.yaml b/mblt_vision/models/YOLOv6n-face.yaml new file mode 100644 index 0000000..f1e7541 --- /dev/null +++ b/mblt_vision/models/YOLOv6n-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv6n-face + filename: yolov6n-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv7.yaml b/mblt_vision/models/YOLOv7.yaml new file mode 100644 index 0000000..65f88c7 --- /dev/null +++ b/mblt_vision/models/YOLOv7.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv7 + filename: yolov7.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 12 + - 16 + - 19 + - 36 + - 40 + - 28 + - - 36 + - 75 + - 76 + - 55 + - 72 + - 146 + - - 142 + - 110 + - 192 + - 243 + - 459 + - 401 + conf_thres: 0.001 + iou_thres: 0.6 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv7d6.yaml b/mblt_vision/models/YOLOv7d6.yaml new file mode 100644 index 0000000..7c777d7 --- /dev/null +++ b/mblt_vision/models/YOLOv7d6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv7d6 + filename: yolov7d6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.6 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv7e6.yaml b/mblt_vision/models/YOLOv7e6.yaml new file mode 100644 index 0000000..9c2dd47 --- /dev/null +++ b/mblt_vision/models/YOLOv7e6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv7e6 + filename: yolov7e6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.6 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv7e6e.yaml b/mblt_vision/models/YOLOv7e6e.yaml new file mode 100644 index 0000000..cc13bcb --- /dev/null +++ b/mblt_vision/models/YOLOv7e6e.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv7e6e + filename: yolov7e6e.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.6 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv7w6.yaml b/mblt_vision/models/YOLOv7w6.yaml new file mode 100644 index 0000000..776dc21 --- /dev/null +++ b/mblt_vision/models/YOLOv7w6.yaml @@ -0,0 +1,50 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv7w6 + filename: yolov7w6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 19 + - 27 + - 44 + - 40 + - 38 + - 94 + - - 96 + - 68 + - 86 + - 152 + - 180 + - 137 + - - 140 + - 301 + - 303 + - 264 + - 238 + - 542 + - - 436 + - 615 + - 739 + - 380 + - 925 + - 792 + conf_thres: 0.001 + iou_thres: 0.6 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv7x.yaml b/mblt_vision/models/YOLOv7x.yaml new file mode 100644 index 0000000..ae4e3b2 --- /dev/null +++ b/mblt_vision/models/YOLOv7x.yaml @@ -0,0 +1,44 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv7x + filename: yolov7x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + anchors: + - - 12 + - 16 + - 19 + - 36 + - 40 + - 28 + - - 36 + - 75 + - 76 + - 55 + - 72 + - 146 + - - 142 + - 110 + - 192 + - 243 + - 459 + - 401 + conf_thres: 0.001 + iou_thres: 0.6 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8l-cls.yaml b/mblt_vision/models/YOLOv8l-cls.yaml new file mode 100644 index 0000000..08c800e --- /dev/null +++ b/mblt_vision/models/YOLOv8l-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8l-cls + filename: yolov8l-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLOv8l-face.yaml b/mblt_vision/models/YOLOv8l-face.yaml new file mode 100644 index 0000000..12f7093 --- /dev/null +++ b/mblt_vision/models/YOLOv8l-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8l-face + filename: yolov8l-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8l-obb.yaml b/mblt_vision/models/YOLOv8l-obb.yaml new file mode 100644 index 0000000..280cb4b --- /dev/null +++ b/mblt_vision/models/YOLOv8l-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8l-obb + filename: yolov8l-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8l-pose.yaml b/mblt_vision/models/YOLOv8l-pose.yaml new file mode 100644 index 0000000..ea366b9 --- /dev/null +++ b/mblt_vision/models/YOLOv8l-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8l-pose + filename: yolov8l-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8l-seg.yaml b/mblt_vision/models/YOLOv8l-seg.yaml new file mode 100644 index 0000000..c34ac0f --- /dev/null +++ b/mblt_vision/models/YOLOv8l-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8l-seg + filename: yolov8l-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8l.yaml b/mblt_vision/models/YOLOv8l.yaml new file mode 100644 index 0000000..62609bf --- /dev/null +++ b/mblt_vision/models/YOLOv8l.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8l + filename: yolov8l.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8m-cls.yaml b/mblt_vision/models/YOLOv8m-cls.yaml new file mode 100644 index 0000000..37c82f5 --- /dev/null +++ b/mblt_vision/models/YOLOv8m-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8m-cls + filename: yolov8m-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLOv8m-face.yaml b/mblt_vision/models/YOLOv8m-face.yaml new file mode 100644 index 0000000..135340c --- /dev/null +++ b/mblt_vision/models/YOLOv8m-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8m-face + filename: yolov8m-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8m-obb.yaml b/mblt_vision/models/YOLOv8m-obb.yaml new file mode 100644 index 0000000..fcb5307 --- /dev/null +++ b/mblt_vision/models/YOLOv8m-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8m-obb + filename: yolov8m-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8m-pose.yaml b/mblt_vision/models/YOLOv8m-pose.yaml new file mode 100644 index 0000000..098fc23 --- /dev/null +++ b/mblt_vision/models/YOLOv8m-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8m-pose + filename: yolov8m-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8m-seg.yaml b/mblt_vision/models/YOLOv8m-seg.yaml new file mode 100644 index 0000000..10f5b81 --- /dev/null +++ b/mblt_vision/models/YOLOv8m-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8m-seg + filename: yolov8m-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8m.yaml b/mblt_vision/models/YOLOv8m.yaml new file mode 100644 index 0000000..06fa61e --- /dev/null +++ b/mblt_vision/models/YOLOv8m.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8m + filename: yolov8m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8n-cls.yaml b/mblt_vision/models/YOLOv8n-cls.yaml new file mode 100644 index 0000000..6a8f357 --- /dev/null +++ b/mblt_vision/models/YOLOv8n-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8n-cls + filename: yolov8n-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLOv8n-face.yaml b/mblt_vision/models/YOLOv8n-face.yaml new file mode 100644 index 0000000..5c7ee1f --- /dev/null +++ b/mblt_vision/models/YOLOv8n-face.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8n-face + filename: yolov8n-face.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: face_detection + dataset: widerface + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8n-obb.yaml b/mblt_vision/models/YOLOv8n-obb.yaml new file mode 100644 index 0000000..2a827b8 --- /dev/null +++ b/mblt_vision/models/YOLOv8n-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8n-obb + filename: yolov8n-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8n-pose.yaml b/mblt_vision/models/YOLOv8n-pose.yaml new file mode 100644 index 0000000..34f5d6f --- /dev/null +++ b/mblt_vision/models/YOLOv8n-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8n-pose + filename: yolov8n-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8n-seg.yaml b/mblt_vision/models/YOLOv8n-seg.yaml new file mode 100644 index 0000000..94920f6 --- /dev/null +++ b/mblt_vision/models/YOLOv8n-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8n-seg + filename: yolov8n-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8n.yaml b/mblt_vision/models/YOLOv8n.yaml new file mode 100644 index 0000000..1ea9850 --- /dev/null +++ b/mblt_vision/models/YOLOv8n.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8n + filename: yolov8n.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8s-cls.yaml b/mblt_vision/models/YOLOv8s-cls.yaml new file mode 100644 index 0000000..a70abb9 --- /dev/null +++ b/mblt_vision/models/YOLOv8s-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8s-cls + filename: yolov8s-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLOv8s-obb.yaml b/mblt_vision/models/YOLOv8s-obb.yaml new file mode 100644 index 0000000..755b321 --- /dev/null +++ b/mblt_vision/models/YOLOv8s-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8s-obb + filename: yolov8s-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8s-pose.yaml b/mblt_vision/models/YOLOv8s-pose.yaml new file mode 100644 index 0000000..be4d4f5 --- /dev/null +++ b/mblt_vision/models/YOLOv8s-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8s-pose + filename: yolov8s-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8s-seg.yaml b/mblt_vision/models/YOLOv8s-seg.yaml new file mode 100644 index 0000000..26fc679 --- /dev/null +++ b/mblt_vision/models/YOLOv8s-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8s-seg + filename: yolov8s-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8s.yaml b/mblt_vision/models/YOLOv8s.yaml new file mode 100644 index 0000000..f705f91 --- /dev/null +++ b/mblt_vision/models/YOLOv8s.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8s + filename: yolov8s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8x-cls.yaml b/mblt_vision/models/YOLOv8x-cls.yaml new file mode 100644 index 0000000..bb7242e --- /dev/null +++ b/mblt_vision/models/YOLOv8x-cls.yaml @@ -0,0 +1,23 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8x-cls + filename: yolov8x-cls.mxq + revision: main + pre_cfg: + Reader: + style: pil + Resize: + size: 224 + interpolation: bilinear + CenterCrop: + size: + - 224 + - 224 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: image_classification + dataset: imagenet + softmax: true diff --git a/mblt_vision/models/YOLOv8x-obb.yaml b/mblt_vision/models/YOLOv8x-obb.yaml new file mode 100644 index 0000000..0eccaf9 --- /dev/null +++ b/mblt_vision/models/YOLOv8x-obb.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8x-obb + filename: yolov8x-obb.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1024 + - 1024 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: obb + dataset: dotav1 + nl: 3 + n_extra: 1 + reg_max: 16 + conf_thres: 0.01 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8x-pose-p6.yaml b/mblt_vision/models/YOLOv8x-pose-p6.yaml new file mode 100644 index 0000000..e42b723 --- /dev/null +++ b/mblt_vision/models/YOLOv8x-pose-p6.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8x-pose-p6 + filename: yolov8x-pose-p6.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 1280 + - 1280 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 4 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8x-pose.yaml b/mblt_vision/models/YOLOv8x-pose.yaml new file mode 100644 index 0000000..694d47e --- /dev/null +++ b/mblt_vision/models/YOLOv8x-pose.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8x-pose + filename: yolov8x-pose.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: pose_estimation + dataset: coco + nl: 3 + n_extra: 51 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8x-seg.yaml b/mblt_vision/models/YOLOv8x-seg.yaml new file mode 100644 index 0000000..ddbc1dd --- /dev/null +++ b/mblt_vision/models/YOLOv8x-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8x-seg + filename: yolov8x-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv8x.yaml b/mblt_vision/models/YOLOv8x.yaml new file mode 100644 index 0000000..808e7f9 --- /dev/null +++ b/mblt_vision/models/YOLOv8x.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv8x + filename: yolov8x.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9c-seg.yaml b/mblt_vision/models/YOLOv9c-seg.yaml new file mode 100644 index 0000000..4eb0730 --- /dev/null +++ b/mblt_vision/models/YOLOv9c-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9c-seg + filename: yolov9c-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9c.yaml b/mblt_vision/models/YOLOv9c.yaml new file mode 100644 index 0000000..3c73d9a --- /dev/null +++ b/mblt_vision/models/YOLOv9c.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9c + filename: yolov9c.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9e-seg.yaml b/mblt_vision/models/YOLOv9e-seg.yaml new file mode 100644 index 0000000..a645374 --- /dev/null +++ b/mblt_vision/models/YOLOv9e-seg.yaml @@ -0,0 +1,28 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9e-seg + filename: yolov9e-seg.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: instance_segmentation + dataset: coco + nl: 3 + n_extra: 32 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9e.yaml b/mblt_vision/models/YOLOv9e.yaml new file mode 100644 index 0000000..e92b7b0 --- /dev/null +++ b/mblt_vision/models/YOLOv9e.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9e + filename: yolov9e.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9m.yaml b/mblt_vision/models/YOLOv9m.yaml new file mode 100644 index 0000000..b83d2a9 --- /dev/null +++ b/mblt_vision/models/YOLOv9m.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9m + filename: yolov9m.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9s.yaml b/mblt_vision/models/YOLOv9s.yaml new file mode 100644 index 0000000..cd4c9af --- /dev/null +++ b/mblt_vision/models/YOLOv9s.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9s + filename: yolov9s.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/models/YOLOv9t.yaml b/mblt_vision/models/YOLOv9t.yaml new file mode 100644 index 0000000..cf0166b --- /dev/null +++ b/mblt_vision/models/YOLOv9t.yaml @@ -0,0 +1,27 @@ +DEFAULT: + file_cfg: + repo_id: mobilint/YOLOv9t + filename: yolov9t.mxq + revision: main + pre_cfg: + Reader: + style: numpy + LetterBox: + img_size: + - 640 + - 640 + SetOrder: + shape: HWC + Normalize: + style: cv + post_cfg: + task: object_detection + dataset: coco + nl: 3 + reg_max: 16 + conf_thres: 0.001 + iou_thres: 0.7 +TURBO: + update: DEFAULT + file_cfg: + revision: TURBO diff --git a/mblt_vision/obb/__init__.py b/mblt_vision/obb/__init__.py new file mode 100644 index 0000000..fc20fd7 --- /dev/null +++ b/mblt_vision/obb/__init__.py @@ -0,0 +1,39 @@ +"""Oriented-bounding-box Vision model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO11lOBB", + "YOLO11mOBB", + "YOLO11nOBB", + "YOLO11sOBB", + "YOLO11xOBB", + "YOLO26lOBB", + "YOLO26mOBB", + "YOLO26nOBB", + "YOLO26sOBB", + "YOLO26xOBB", + "YOLOv8lOBB", + "YOLOv8mOBB", + "YOLOv8nOBB", + "YOLOv8sOBB", + "YOLOv8xOBB", +] + +YOLO11lOBB = create_model_class("YOLO11lOBB", __name__) +YOLO11mOBB = create_model_class("YOLO11mOBB", __name__) +YOLO11nOBB = create_model_class("YOLO11nOBB", __name__) +YOLO11sOBB = create_model_class("YOLO11sOBB", __name__) +YOLO11xOBB = create_model_class("YOLO11xOBB", __name__) +YOLO26lOBB = create_model_class("YOLO26lOBB", __name__) +YOLO26mOBB = create_model_class("YOLO26mOBB", __name__) +YOLO26nOBB = create_model_class("YOLO26nOBB", __name__) +YOLO26sOBB = create_model_class("YOLO26sOBB", __name__) +YOLO26xOBB = create_model_class("YOLO26xOBB", __name__) +YOLOv8lOBB = create_model_class("YOLOv8lOBB", __name__) +YOLOv8mOBB = create_model_class("YOLOv8mOBB", __name__) +YOLOv8nOBB = create_model_class("YOLOv8nOBB", __name__) +YOLOv8sOBB = create_model_class("YOLOv8sOBB", __name__) +YOLOv8xOBB = create_model_class("YOLOv8xOBB", __name__) diff --git a/mblt_vision/object_detection/__init__.py b/mblt_vision/object_detection/__init__.py new file mode 100644 index 0000000..3e2cd2f --- /dev/null +++ b/mblt_vision/object_detection/__init__.py @@ -0,0 +1,153 @@ +"""Object detection model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO11l", + "YOLO11m", + "YOLO11n", + "YOLO11s", + "YOLO11x", + "YOLO12l", + "YOLO12m", + "YOLO12n", + "YOLO12s", + "YOLO12x", + "YOLO26l", + "YOLO26lDistill", + "YOLO26m", + "YOLO26mDistill", + "YOLO26n", + "YOLO26nDistill", + "YOLO26s", + "YOLO26sDistill", + "YOLO26x", + "YOLO26xDistill", + "YOLOv3", + "YOLOv3_spp", + "YOLOv3_sppu", + "YOLOv3_tiny", + "YOLOv3_tinyu", + "YOLOv3u", + "YOLOv5l", + "YOLOv5l6", + "YOLOv5l6u", + "YOLOv5lu", + "YOLOv5m", + "YOLOv5m6", + "YOLOv5m6u", + "YOLOv5mu", + "YOLOv5n", + "YOLOv5n6", + "YOLOv5n6u", + "YOLOv5nu", + "YOLOv5s", + "YOLOv5s6", + "YOLOv5s6u", + "YOLOv5su", + "YOLOv5x", + "YOLOv5x6", + "YOLOv5x6u", + "YOLOv5xu", + "YOLOv7", + "YOLOv7d6", + "YOLOv7e6", + "YOLOv7e6e", + "YOLOv7w6", + "YOLOv7x", + "YOLOv8l", + "YOLOv8m", + "YOLOv8n", + "YOLOv8s", + "YOLOv8x", + "GELANc", + "GELANe", + "GELANm", + "GELANs", + "YOLOv9c", + "YOLOv9e", + "YOLOv9m", + "YOLOv9s", + "YOLOv9t", + "YOLOv10b", + "YOLOv10l", + "YOLOv10m", + "YOLOv10n", + "YOLOv10s", + "YOLOv10x", +] + +YOLO11l = create_model_class("YOLO11l", __name__) +YOLO11m = create_model_class("YOLO11m", __name__) +YOLO11n = create_model_class("YOLO11n", __name__) +YOLO11s = create_model_class("YOLO11s", __name__) +YOLO11x = create_model_class("YOLO11x", __name__) +YOLO12l = create_model_class("YOLO12l", __name__) +YOLO12m = create_model_class("YOLO12m", __name__) +YOLO12n = create_model_class("YOLO12n", __name__) +YOLO12s = create_model_class("YOLO12s", __name__) +YOLO12x = create_model_class("YOLO12x", __name__) +YOLO26l = create_model_class("YOLO26l", __name__) +YOLO26lDistill = create_model_class("YOLO26lDistill", __name__) +YOLO26m = create_model_class("YOLO26m", __name__) +YOLO26mDistill = create_model_class("YOLO26mDistill", __name__) +YOLO26n = create_model_class("YOLO26n", __name__) +YOLO26nDistill = create_model_class("YOLO26nDistill", __name__) +YOLO26s = create_model_class("YOLO26s", __name__) +YOLO26sDistill = create_model_class("YOLO26sDistill", __name__) +YOLO26x = create_model_class("YOLO26x", __name__) +YOLO26xDistill = create_model_class("YOLO26xDistill", __name__) +YOLOv3 = create_model_class("YOLOv3", __name__) +YOLOv3_spp = create_model_class("YOLOv3_spp", __name__) +YOLOv3_sppu = create_model_class("YOLOv3_sppu", __name__) +YOLOv3_tiny = create_model_class("YOLOv3_tiny", __name__) +YOLOv3_tinyu = create_model_class("YOLOv3_tinyu", __name__) +YOLOv3u = create_model_class("YOLOv3u", __name__) +YOLOv5l = create_model_class("YOLOv5l", __name__) +YOLOv5l6 = create_model_class("YOLOv5l6", __name__) +YOLOv5l6u = create_model_class("YOLOv5l6u", __name__) +YOLOv5lu = create_model_class("YOLOv5lu", __name__) +YOLOv5m = create_model_class("YOLOv5m", __name__) +YOLOv5m6 = create_model_class("YOLOv5m6", __name__) +YOLOv5m6u = create_model_class("YOLOv5m6u", __name__) +YOLOv5mu = create_model_class("YOLOv5mu", __name__) +YOLOv5n = create_model_class("YOLOv5n", __name__) +YOLOv5n6 = create_model_class("YOLOv5n6", __name__) +YOLOv5n6u = create_model_class("YOLOv5n6u", __name__) +YOLOv5nu = create_model_class("YOLOv5nu", __name__) +YOLOv5s = create_model_class("YOLOv5s", __name__) +YOLOv5s6 = create_model_class("YOLOv5s6", __name__) +YOLOv5s6u = create_model_class("YOLOv5s6u", __name__) +YOLOv5su = create_model_class("YOLOv5su", __name__) +YOLOv5x = create_model_class("YOLOv5x", __name__) +YOLOv5x6 = create_model_class("YOLOv5x6", __name__) +YOLOv5x6u = create_model_class("YOLOv5x6u", __name__) +YOLOv5xu = create_model_class("YOLOv5xu", __name__) +YOLOv7 = create_model_class("YOLOv7", __name__) +YOLOv7d6 = create_model_class("YOLOv7d6", __name__) +YOLOv7e6 = create_model_class("YOLOv7e6", __name__) +YOLOv7e6e = create_model_class("YOLOv7e6e", __name__) +YOLOv7w6 = create_model_class("YOLOv7w6", __name__) +YOLOv7x = create_model_class("YOLOv7x", __name__) +YOLOv8l = create_model_class("YOLOv8l", __name__) +YOLOv8m = create_model_class("YOLOv8m", __name__) +YOLOv8n = create_model_class("YOLOv8n", __name__) +YOLOv8s = create_model_class("YOLOv8s", __name__) +YOLOv8x = create_model_class("YOLOv8x", __name__) +GELANc = create_model_class("GELANc", __name__) +GELANe = create_model_class("GELANe", __name__) +GELANm = create_model_class("GELANm", __name__) +GELANs = create_model_class("GELANs", __name__) +YOLOv9c = create_model_class("YOLOv9c", __name__) +YOLOv9e = create_model_class("YOLOv9e", __name__) +YOLOv9m = create_model_class("YOLOv9m", __name__) +YOLOv9s = create_model_class("YOLOv9s", __name__) +YOLOv9t = create_model_class("YOLOv9t", __name__) +YOLOv10b = create_model_class("YOLOv10b", __name__) +YOLOv10l = create_model_class("YOLOv10l", __name__) +YOLOv10m = create_model_class("YOLOv10m", __name__) +YOLOv10n = create_model_class("YOLOv10n", __name__) +YOLOv10s = create_model_class("YOLOv10s", __name__) +YOLOv10x = create_model_class("YOLOv10x", __name__) diff --git a/mblt_vision/pose_estimation/__init__.py b/mblt_vision/pose_estimation/__init__.py new file mode 100644 index 0000000..ab136d7 --- /dev/null +++ b/mblt_vision/pose_estimation/__init__.py @@ -0,0 +1,41 @@ +"""Pose estimation model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO11lPose", + "YOLO11mPose", + "YOLO11nPose", + "YOLO11sPose", + "YOLO11xPose", + "YOLO26lPose", + "YOLO26mPose", + "YOLO26nPose", + "YOLO26sPose", + "YOLO26xPose", + "YOLOv8lPose", + "YOLOv8mPose", + "YOLOv8nPose", + "YOLOv8sPose", + "YOLOv8xPose", + "YOLOv8xPoseP6", +] + +YOLO11lPose = create_model_class("YOLO11lPose", __name__) +YOLO11mPose = create_model_class("YOLO11mPose", __name__) +YOLO11nPose = create_model_class("YOLO11nPose", __name__) +YOLO11sPose = create_model_class("YOLO11sPose", __name__) +YOLO11xPose = create_model_class("YOLO11xPose", __name__) +YOLO26lPose = create_model_class("YOLO26lPose", __name__) +YOLO26mPose = create_model_class("YOLO26mPose", __name__) +YOLO26nPose = create_model_class("YOLO26nPose", __name__) +YOLO26sPose = create_model_class("YOLO26sPose", __name__) +YOLO26xPose = create_model_class("YOLO26xPose", __name__) +YOLOv8lPose = create_model_class("YOLOv8lPose", __name__) +YOLOv8mPose = create_model_class("YOLOv8mPose", __name__) +YOLOv8nPose = create_model_class("YOLOv8nPose", __name__) +YOLOv8sPose = create_model_class("YOLOv8sPose", __name__) +YOLOv8xPose = create_model_class("YOLOv8xPose", __name__) +YOLOv8xPoseP6 = create_model_class("YOLOv8xPoseP6", __name__) diff --git a/mblt_vision/py.typed b/mblt_vision/py.typed new file mode 100644 index 0000000..be1a4ff --- /dev/null +++ b/mblt_vision/py.typed @@ -0,0 +1 @@ +# PEP 561 marker: mblt_vision ships inline type information. diff --git a/mblt_vision/semantic_segmentation/__init__.py b/mblt_vision/semantic_segmentation/__init__.py new file mode 100644 index 0000000..6fd6ff5 --- /dev/null +++ b/mblt_vision/semantic_segmentation/__init__.py @@ -0,0 +1,29 @@ +"""Semantic segmentation model exports.""" + +from __future__ import annotations + +from .._compat import create_model_class + +__all__: list[str] = [ + "YOLO26lSem", + "YOLO26lSemADE20K", + "YOLO26mSem", + "YOLO26mSemADE20K", + "YOLO26nSem", + "YOLO26nSemADE20K", + "YOLO26sSem", + "YOLO26sSemADE20K", + "YOLO26xSem", + "YOLO26xSemADE20K", +] + +YOLO26lSem = create_model_class("YOLO26lSem", __name__) +YOLO26lSemADE20K = create_model_class("YOLO26lSemADE20K", __name__) +YOLO26mSem = create_model_class("YOLO26mSem", __name__) +YOLO26mSemADE20K = create_model_class("YOLO26mSemADE20K", __name__) +YOLO26nSem = create_model_class("YOLO26nSem", __name__) +YOLO26nSemADE20K = create_model_class("YOLO26nSemADE20K", __name__) +YOLO26sSem = create_model_class("YOLO26sSem", __name__) +YOLO26sSemADE20K = create_model_class("YOLO26sSemADE20K", __name__) +YOLO26xSem = create_model_class("YOLO26xSem", __name__) +YOLO26xSemADE20K = create_model_class("YOLO26xSemADE20K", __name__) diff --git a/mblt_vision/utils/__init__.py b/mblt_vision/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mblt_vision/utils/datasets/__init__.py b/mblt_vision/utils/datasets/__init__.py new file mode 100644 index 0000000..b0751a8 --- /dev/null +++ b/mblt_vision/utils/datasets/__init__.py @@ -0,0 +1,85 @@ +""" +Datasets utilities and loaders. +""" + +from __future__ import annotations + +from .ade20k import get_ade20k_palette +from .cityscapes import get_cityscapes_palette +from .coco import ( + get_coco_class_num, + get_coco_det_palette, + get_coco_inv, + get_coco_keypoint_palette, + get_coco_label, + get_coco_limb_palette, + get_coco_pose_skeleton, +) +from .dataloader import ( + CustomADE20K, + CustomCityscapes, + CustomCocodata, + CustomCOCODataset, + CustomDOTAv1, + CustomImageFolder, + CustomNYUDepth, + CustomWiderface, + CustomWiderFaceDataset, + get_ade20k_loader, + get_cityscapes_loader, + get_coco_loader, + get_dota_loader, + get_imagenet_loader, + get_nyu_depth_loader, + get_widerface_loader, +) +from .dotav1 import get_dotav1_class_num, get_dotav1_label, get_dotav1_palette +from .imagenet import get_imagenet_label +from .organizer import ( + organize_ade20k, + organize_cityscapes, + organize_coco, + organize_dotav1, + organize_imagenet, + organize_nyu_depth, + organize_widerface, +) + +__all__: list[str] = [ + "get_ade20k_palette", + "get_cityscapes_palette", + "get_coco_class_num", + "get_coco_det_palette", + "get_coco_inv", + "get_coco_keypoint_palette", + "get_coco_label", + "get_coco_limb_palette", + "get_coco_pose_skeleton", + "get_dotav1_class_num", + "get_dotav1_label", + "get_dotav1_palette", + "CustomADE20K", + "CustomCityscapes", + "CustomCOCODataset", + "CustomCocodata", + "CustomDOTAv1", + "CustomImageFolder", + "CustomNYUDepth", + "CustomWiderFaceDataset", + "CustomWiderface", + "get_ade20k_loader", + "get_cityscapes_loader", + "get_coco_loader", + "get_dota_loader", + "get_imagenet_loader", + "get_nyu_depth_loader", + "get_widerface_loader", + "get_imagenet_label", + "organize_coco", + "organize_ade20k", + "organize_cityscapes", + "organize_dotav1", + "organize_imagenet", + "organize_nyu_depth", + "organize_widerface", +] diff --git a/mblt_vision/utils/datasets/ade20k.py b/mblt_vision/utils/datasets/ade20k.py new file mode 100644 index 0000000..8075db9 --- /dev/null +++ b/mblt_vision/utils/datasets/ade20k.py @@ -0,0 +1,41 @@ +""" +ADE20K dataset constants and utilities. +""" + +from __future__ import annotations + +ADE20K_PALETTE_BGR = [ + (255, 42, 4), + (235, 219, 11), + (243, 243, 243), + (183, 223, 0), + (104, 31, 17), + (221, 111, 255), + (79, 68, 255), + (0, 237, 204), + (68, 243, 0), + (255, 0, 189), + (255, 180, 0), + (186, 0, 221), + (255, 255, 0), + (0, 192, 38), + (179, 255, 1), + (255, 36, 125), + (104, 0, 123), + (108, 27, 255), + (47, 109, 252), + (11, 255, 162), +] + + +def get_ade20k_palette(idx: int) -> tuple[int, int, int]: + """Get an ADE20K visualization color in OpenCV BGR order. + + Args: + idx: ADE20K class index. + + Returns: + BGR color tuple. Colors repeat when the index exceeds the palette length. + """ + + return ADE20K_PALETTE_BGR[idx % len(ADE20K_PALETTE_BGR)] diff --git a/mblt_vision/utils/datasets/cityscapes.py b/mblt_vision/utils/datasets/cityscapes.py new file mode 100644 index 0000000..161a666 --- /dev/null +++ b/mblt_vision/utils/datasets/cityscapes.py @@ -0,0 +1,66 @@ +"""Cityscapes dataset visualization utilities.""" + +from __future__ import annotations + +import numpy as np + +# OpenCV images use BGR channel order. These are the official Cityscapes train-ID +# colors converted from RGB so familiar classes retain their standard appearance. +CITYSCAPES_PALETTE_BGR: tuple[tuple[int, int, int], ...] = ( + (128, 64, 128), # road + (232, 35, 244), # sidewalk + (70, 70, 70), # building + (156, 102, 102), # wall + (153, 153, 190), # fence + (153, 153, 153), # pole + (30, 170, 250), # traffic light + (0, 220, 220), # traffic sign + (35, 142, 107), # vegetation + (152, 251, 152), # terrain + (180, 130, 70), # sky + (60, 20, 220), # person + (0, 0, 255), # rider + (142, 0, 0), # car + (70, 0, 0), # truck + (100, 60, 0), # bus + (100, 80, 0), # train + (230, 0, 0), # motorcycle + (32, 11, 119), # bicycle +) + +CITYSCAPES_SOURCE_IDS: tuple[int, ...] = ( + 7, + 8, + 11, + 12, + 13, + 17, + *range(19, 29), + 31, + 32, + 33, +) +CITYSCAPES_SOURCE_TO_TRAIN_ID = np.full(256, 255, dtype=np.uint8) +CITYSCAPES_SOURCE_TO_TRAIN_ID[list(CITYSCAPES_SOURCE_IDS)] = np.arange( + 19, dtype=np.uint8 +) + + +def get_cityscapes_palette(class_id: int) -> tuple[int, int, int]: + """Return the BGR visualization color for a Cityscapes train ID. + + Args: + class_id: Cityscapes train ID in the range 0 through 18. + + Returns: + The corresponding color in OpenCV BGR order. + + Raises: + ValueError: If ``class_id`` is outside the Cityscapes train-ID range. + """ + + if not 0 <= class_id < len(CITYSCAPES_PALETTE_BGR): + raise ValueError( + f"Cityscapes class ID must be in [0, {len(CITYSCAPES_PALETTE_BGR) - 1}]." + ) + return CITYSCAPES_PALETTE_BGR[class_id] diff --git a/mblt_vision/utils/datasets/coco.py b/mblt_vision/utils/datasets/coco.py new file mode 100644 index 0000000..b2f06e4 --- /dev/null +++ b/mblt_vision/utils/datasets/coco.py @@ -0,0 +1,243 @@ +""" +COCO dataset constants and utilities. +""" + +from __future__ import annotations + +from ...datasets import get_dataset_category_ids, get_dataset_class_names + +DET_PALETTE = [ + (220, 20, 60), + (119, 11, 32), + (0, 0, 142), + (0, 0, 230), + (106, 0, 228), + (0, 60, 100), + (0, 80, 100), + (0, 0, 70), + (0, 0, 192), + (250, 170, 30), + (100, 170, 30), + (220, 220, 0), + (175, 116, 175), + (250, 0, 30), + (165, 42, 42), + (255, 77, 255), + (0, 226, 252), + (182, 182, 255), + (0, 82, 0), + (120, 166, 157), + (110, 76, 0), + (174, 57, 255), + (199, 100, 0), + (72, 0, 118), + (255, 179, 240), + (0, 125, 92), + (209, 0, 151), + (188, 208, 182), + (0, 220, 176), + (255, 99, 164), + (92, 0, 73), + (133, 129, 255), + (78, 180, 255), + (0, 228, 0), + (174, 255, 243), + (45, 89, 255), + (134, 134, 103), + (145, 148, 174), + (255, 208, 186), + (197, 226, 255), + (171, 134, 1), + (109, 63, 54), + (207, 138, 255), + (151, 0, 95), + (9, 80, 61), + (84, 105, 51), + (74, 65, 105), + (166, 196, 102), + (208, 195, 210), + (255, 109, 65), + (0, 143, 149), + (179, 0, 194), + (209, 99, 106), + (5, 121, 0), + (227, 255, 205), + (147, 186, 208), + (153, 69, 1), + (3, 95, 161), + (163, 255, 0), + (119, 0, 170), + (0, 182, 199), + (0, 165, 120), + (183, 130, 88), + (95, 32, 0), + (130, 114, 135), + (110, 129, 133), + (166, 74, 118), + (219, 142, 185), + (79, 210, 114), + (178, 90, 62), + (65, 70, 15), + (127, 167, 115), + (59, 105, 106), + (142, 108, 45), + (196, 172, 0), + (95, 54, 80), + (128, 76, 255), + (201, 57, 1), + (246, 0, 122), + (191, 162, 208), +] + +POSE_PALETTE = [ + [255, 128, 0], + [255, 153, 51], + [255, 178, 102], + [230, 230, 0], + [255, 153, 255], + [153, 204, 255], + [255, 102, 255], + [255, 51, 255], + [102, 178, 255], + [51, 153, 255], + [255, 153, 153], + [255, 102, 102], + [255, 51, 51], + [153, 255, 153], + [102, 255, 102], + [51, 255, 51], + [0, 255, 0], + [0, 0, 255], + [255, 0, 0], + [255, 255, 255], +] + +POSE_SKELETON = [ + [16, 14], + [14, 12], + [17, 15], + [15, 13], + [12, 13], + [6, 12], + [7, 13], + [6, 7], + [6, 8], + [7, 9], + [8, 10], + [9, 11], + [2, 3], + [1, 2], + [1, 3], + [2, 4], + [3, 5], + [4, 6], + [5, 7], +] + +LIMB_PALETTE = [ + POSE_PALETTE[i] + for i in [9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16] +] +KEYPOINT_PALETTE = [ + POSE_PALETTE[i] for i in [16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9] +] + + +def get_coco_class_num() -> int: + """Get the number of COCO classes. + + Returns: + int: The number of COCO classes. + """ + return len(get_dataset_class_names("coco")) + + +def get_coco_label(idx: int) -> str: + """ + Get the COCO class label for a given index. + + Args: + idx (int): Zero-based class index. + + Returns: + str: Descriptive label for the class (e.g., "person"). + + Raises: + ValueError: If index is out of range. + """ + if not 0 <= idx < get_coco_class_num(): + raise ValueError( + f"COCO class index must be in [0, {get_coco_class_num() - 1}], got {idx}." + ) + + return get_dataset_class_names("coco")[idx] + + +def get_coco_inv(idx: int) -> int: + """Get the original COCO category ID for a given model output index. + + Args: + idx (int): Model output class index. + + Returns: + int: Original COCO category ID. + """ + return get_dataset_category_ids("coco")[idx] + + +def get_coco_det_palette(idx: int) -> tuple[int, int, int]: + """ + Get a distinct color for a COCO detection class. + + Args: + idx (int): Class index. + + Returns: + tuple[int, int, int]: (R, G, B) color tuple. + """ + return DET_PALETTE[idx] + + +def get_coco_pose_palette(idx: int) -> list[int]: + """Get the COCO pose palette by index. + + Args: + idx (int): The index of the COCO pose palette. + + Returns: + list[int]: The COCO pose palette as an [R, G, B] list. + """ + return POSE_PALETTE[idx] + + +def get_coco_pose_skeleton() -> list[list[int]]: + """Get the COCO pose skeleton. + + Returns: + list[list[int]]: The COCO pose skeleton. + """ + return POSE_SKELETON + + +def get_coco_limb_palette(idx: int) -> list[int]: + """Get the COCO limb palette by index. + + Args: + idx (int): The index of the COCO limb palette. + + Returns: + list[int]: The COCO limb palette as an [R, G, B] list. + """ + return LIMB_PALETTE[idx] + + +def get_coco_keypoint_palette(idx: int) -> list[int]: + """Get the COCO keypoint palette by index. + + Args: + idx (int): The index of the COCO keypoint palette. + + Returns: + list[int]: The COCO keypoint palette as an [R, G, B] list. + """ + return KEYPOINT_PALETTE[idx] diff --git a/mblt_vision/utils/datasets/dataloader.py b/mblt_vision/utils/datasets/dataloader.py new file mode 100644 index 0000000..d23d719 --- /dev/null +++ b/mblt_vision/utils/datasets/dataloader.py @@ -0,0 +1,1016 @@ +""" +Custom dataloaders for vision datasets. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any, Callable + +import cv2 +import numpy as np +import torch +from faster_coco_eval import COCO +from mblt_vision.utils.preprocess.letterbox import letterbox_semantic_mask +from PIL import Image + +from .cityscapes import CITYSCAPES_SOURCE_TO_TRAIN_ID +from .readiness import IMAGE_SUFFIXES + + +def _unique_paths_by_stem( + directory: str, suffixes: tuple[str, ...], description: str +) -> dict[str, str]: + """Return supported direct-child files keyed by unique case-preserving stems.""" + + names = [name for name in os.listdir(directory) if name.lower().endswith(suffixes)] + stems = [os.path.splitext(name)[0] for name in names] + if len(stems) != len(set(stems)): + raise ValueError(f"{description} contain duplicate filename stems.") + return { + stem: os.path.join(directory, name) + for stem, name in zip(stems, names, strict=True) + } + + +class CustomCOCODataset(torch.utils.data.Dataset[tuple[np.ndarray, int, int, int]]): + """Custom COCO dataset class for loading images and metadata. + + This class provides a simple interface for accessing COCO formatted data + without requiring external library dependencies like torchvision. + + Attributes: + root (str): Root directory path containing the images. + coco (COCO): COCO helper object from faster_coco_eval. + ids (list[int]): Sorted list of image IDs in the dataset. + """ + + def __init__( + self, root: str, annFile: str, min_keypoints: int | None = None + ) -> None: + """Initialize the custom COCO dataset. + + Args: + root (str): Path to the directory containing images. + annFile (str): Path to the COCO annotation JSON file. + min_keypoints: If set, keep only images with at least one + annotation whose ``num_keypoints`` is greater than this value. + """ + self.root = root + try: + raw_annotation = json.loads(Path(annFile).read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise ValueError( + f"Unable to read COCO annotation file {annFile}: {exc}." + ) from exc + if not isinstance(raw_annotation, dict): + raise ValueError( + f"COCO annotation file {annFile} must contain a JSON object." + ) + self.raw_annotation = raw_annotation + self.coco = COCO(annFile) + if min_keypoints is None: + self.ids = list(sorted(self.coco.imgs.keys())) + else: + self.ids = list( + sorted( + { + ann["image_id"] + for ann in self.coco.anns.values() + if ann.get("num_keypoints", 0) > min_keypoints + } + ) + ) + + def _load_image(self, image_id: int) -> np.ndarray: + """Load image by ID""" + file_name = self.coco.loadImgs(image_id)[0]["file_name"] + if not isinstance(file_name, str) or not file_name: + raise ValueError(f"COCO image ID {image_id} has an invalid file_name.") + relative_path = Path(file_name) + if relative_path.is_absolute() or ".." in relative_path.parts: + raise ValueError( + f"COCO image ID {image_id} has an unsafe file_name: {file_name!r}." + ) + image_root = Path(self.root).resolve() + image_path = (image_root / relative_path).resolve() + try: + image_path.relative_to(image_root) + except ValueError as exc: + raise ValueError( + f"COCO image ID {image_id} resolves outside the image root: {file_name!r}." + ) from exc + image = cv2.imread(str(image_path)) # Load image (BGR format) + + if image is None: + raise FileNotFoundError(f"Image not found: {image_path}") + + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB + + def __getitem__(self, index: int) -> tuple[np.ndarray, int, int, int]: + """Get the image and target by index""" + image_id = self.ids[index] + image = self._load_image(image_id) + height = self.coco.imgs[image_id]["height"] + width = self.coco.imgs[image_id]["width"] + if (height, width) != image.shape[:2]: + raise ValueError( + "COCO annotation geometry does not match decoded image for " + f"image ID {image_id}: annotation {(height, width)}, " + f"image {image.shape[:2]}." + ) + return image, index, height, width + + def __len__(self) -> int: + """Return the total number of images""" + return len(self.ids) + + +CustomCocodata = CustomCOCODataset + + +def get_coco_loader( + dataset: CustomCOCODataset, + batch_size: int, + preprocess_fn: Callable, +) -> torch.utils.data.DataLoader: + """Creates a DataLoader for the COCO dataset. + + Args: + dataset (CustomCOCODataset): The dataset instance to load from. + batch_size (int): Number of samples per batch. + preprocess_fn (Callable): Function used to preprocess images. + + Returns: + torch.utils.data.DataLoader: A configured DataLoader for the COCO dataset. + """ + + def loader( + batch: list[Any], + ) -> tuple[np.ndarray, np.ndarray, list[Any], tuple[int, ...]]: + """Collate function for COCO DataLoader.""" + batch = list(filter(lambda x: x is not None, batch)) + images, idx, height, width = zip(*batch) + + processed_images = [] + ratio_pads = [] + for img in images: + processed = preprocess_fn(img) + if ( + isinstance(processed, tuple) + and len(processed) == 2 + and isinstance(processed[1], dict) + ): + processed_img, metadata = processed + ratio_pads.append(metadata.get("ratio_pad")) + else: + processed_img = processed + ratio_pads.append(None) + processed_images.append(processed_img) + + height_arr = np.array(height) + width_arr = np.array(width) + + return ( + np.stack(processed_images, axis=0), + np.stack((height_arr, width_arr), axis=1), + ratio_pads, + idx, + ) + + return torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=0, + collate_fn=loader, + ) + + +class CustomNYUDepth(torch.utils.data.Dataset[tuple[np.ndarray, np.ndarray, str]]): + """NYU Depth V2 validation dataset with paired RGB images and ``.npy`` depth maps.""" + + IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp") + + def __init__(self, root: str) -> None: + """Validate the organizer's ``images/`` and ``depth/`` validation-only layout.""" + + self.root = root + image_root, depth_root = ( + os.path.join(root, "images"), + os.path.join(root, "depth"), + ) + if not os.path.isdir(image_root) or not os.path.isdir(depth_root): + raise FileNotFoundError( + f"NYU Depth requires images/ and depth/ directories under: {root}" + ) + images = _unique_paths_by_stem( + image_root, self.IMG_EXTENSIONS, "NYU Depth images" + ) + depths = _unique_paths_by_stem(depth_root, (".npy",), "NYU Depth depth maps") + missing_depths, missing_images = ( + sorted(set(images) - set(depths)), + sorted(set(depths) - set(images)), + ) + if missing_depths or missing_images: + details = [] + if missing_depths: + details.append( + f"images without depth maps: {', '.join(missing_depths[:5])}" + ) + if missing_images: + details.append( + f"depth maps without images: {', '.join(missing_images[:5])}" + ) + raise ValueError(f"NYU Depth image/depth mismatch ({'; '.join(details)}).") + if not images: + raise ValueError(f"NYU Depth contains no image/depth pairs: {root}") + self.samples = [(images[stem], depths[stem], stem) for stem in sorted(images)] + + def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray, str]: + """Load an RGB image and finite-safe depth target.""" + + image_path, depth_path, stem = self.samples[index] + image = cv2.imread(image_path) + if image is None: + raise FileNotFoundError(f"NYU Depth image not found: {image_path}") + raw_depth = np.load(depth_path, allow_pickle=False) + if not np.issubdtype(raw_depth.dtype, np.number) or np.issubdtype( + raw_depth.dtype, np.complexfloating + ): + raise ValueError( + "NYU Depth target must use a real numeric dtype, " + f"got {raw_depth.dtype}: {depth_path}" + ) + depth = np.asarray(raw_depth, dtype=np.float32) + if depth.ndim != 2: + raise ValueError( + f"NYU Depth target must be two-dimensional, got {depth.shape}: {depth_path}" + ) + if depth.shape != image.shape[:2]: + raise ValueError( + "NYU Depth image and target shapes must match for " + f"{stem}: image {image.shape[:2]}, depth {depth.shape}." + ) + if not bool(np.isfinite(depth).all()): + raise ValueError( + f"NYU Depth target must contain only finite values: {depth_path}" + ) + if bool((depth < 0).any()): + raise ValueError( + f"NYU Depth target must not contain negative values: {depth_path}" + ) + return ( + cv2.cvtColor(image, cv2.COLOR_BGR2RGB), + depth, + stem, + ) + + def __len__(self) -> int: + """Return the number of paired validation samples.""" + + return len(self.samples) + + +def get_nyu_depth_loader( + dataset: CustomNYUDepth, + batch_size: int, + preprocess_fn: Callable, + image_size: tuple[int, int] | None = None, +) -> torch.utils.data.DataLoader: + """Create a NYU Depth loader with optional stretch-to-size validation preprocessing. + + Args: + dataset: Paired NYU Depth validation dataset. + batch_size: Number of samples per batch. + preprocess_fn: Preprocessing applied after an optional validation resize. + image_size: Optional ``(height, width)`` used to stretch RGB inputs with bilinear + interpolation and depth targets with nearest-neighbor interpolation. This + matches the Ultralytics depth validation pipeline. + + Returns: + Configured NYU Depth validation loader. + """ + + def loader( + batch: list[Any], + ) -> tuple[ + np.ndarray, list[np.ndarray], list[tuple[int, int]], list[Any], tuple[str, ...] + ]: + images, targets, stems = zip(*batch) + processed_images, shapes, ratio_pads = [], [], [] + processed_targets = [] + for image, target in zip(images, targets): + if image_size is not None: + height, width = image_size + image = cv2.resize( + image, (width, height), interpolation=cv2.INTER_LINEAR + ) + target = cv2.resize( + target, (width, height), interpolation=cv2.INTER_NEAREST + ) + shapes.append(tuple(image.shape[:2])) + processed = preprocess_fn(image) + if ( + isinstance(processed, tuple) + and len(processed) == 2 + and isinstance(processed[1], dict) + ): + processed_image, metadata = processed + ratio_pads.append(metadata.get("ratio_pad")) + else: + processed_image = processed + ratio_pads.append(None) + processed_images.append(processed_image) + processed_targets.append(target) + return np.stack(processed_images), processed_targets, shapes, ratio_pads, stems + + return torch.utils.data.DataLoader( + dataset, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=loader + ) + + +class CustomADE20K(torch.utils.data.Dataset[tuple[np.ndarray, np.ndarray, str]]): + """ADE20K validation dataset with paired RGB images and semantic PNG masks.""" + + IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp") + + def __init__(self, root: str) -> None: + """Validate the organizer's flat ``images/`` and ``annotations/`` layout.""" + + self.root = root + image_root = os.path.join(root, "images") + annotation_root = os.path.join(root, "annotations") + if not os.path.isdir(image_root) or not os.path.isdir(annotation_root): + raise FileNotFoundError( + f"ADE20K requires images/ and annotations/ directories under: {root}" + ) + images = _unique_paths_by_stem(image_root, self.IMG_EXTENSIONS, "ADE20K images") + annotations = _unique_paths_by_stem( + annotation_root, (".png",), "ADE20K annotations" + ) + missing_annotations = sorted(set(images) - set(annotations)) + missing_images = sorted(set(annotations) - set(images)) + if missing_annotations or missing_images: + details = [] + if missing_annotations: + details.append( + f"images without annotations: {', '.join(missing_annotations[:5])}" + ) + if missing_images: + details.append( + f"annotations without images: {', '.join(missing_images[:5])}" + ) + raise ValueError( + f"ADE20K image/annotation mismatch ({'; '.join(details)})." + ) + if not images: + raise ValueError(f"ADE20K contains no image/annotation pairs: {root}") + self.samples = [ + (images[stem], annotations[stem], stem) for stem in sorted(images) + ] + + def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray, str]: + """Load one RGB image and map its source labels to model class IDs.""" + + image_path, annotation_path, stem = self.samples[index] + image = cv2.imread(image_path, cv2.IMREAD_COLOR) + if image is None: + raise FileNotFoundError(f"ADE20K image not found: {image_path}") + try: + with Image.open(annotation_path) as annotation_image: + annotation = np.asarray(annotation_image) + except OSError as exc: + raise FileNotFoundError( + f"ADE20K annotation not found: {annotation_path}" + ) from exc + if annotation.ndim != 2 or annotation.dtype != np.uint8: + raise ValueError( + "ADE20K annotations must be single-channel 8-bit PNG masks: " + f"{annotation_path}" + ) + if image.shape[:2] != annotation.shape: + raise ValueError( + f"ADE20K image and annotation shapes must match, got {image.shape[:2]} and {annotation.shape}: {stem}" + ) + if annotation.size and int(annotation.max()) > 150: + raise ValueError( + f"ADE20K annotation values must be in [0, 150]: {annotation_path}" + ) + target = np.full(annotation.shape, 255, dtype=np.uint8) + valid = annotation > 0 + target[valid] = annotation[valid] - 1 + if not valid.any(): + raise ValueError( + f"ADE20K annotation contains no evaluable class IDs: {annotation_path}" + ) + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), target, stem + + def __len__(self) -> int: + """Return the number of paired validation samples.""" + + return len(self.samples) + + +def get_ade20k_loader( + dataset: CustomADE20K, + batch_size: int, + preprocess_fn: Callable, + image_size: tuple[int, int], +) -> torch.utils.data.DataLoader: + """Create an ADE20K loader that applies matching letterbox geometry to masks. + + Args: + dataset: Paired ADE20K validation dataset. + batch_size: Number of samples per batch. + preprocess_fn: Image preprocessing function that returns letterbox metadata. + image_size: Configured model input size as ``(height, width)``. + + Returns: + Configured ADE20K validation loader. + """ + + def loader( + batch: list[Any], + ) -> tuple[ + np.ndarray, np.ndarray, list[tuple[int, int]], list[Any], tuple[str, ...] + ]: + images, targets, stems = zip(*batch) + processed_images, processed_targets, shapes, ratio_pads = [], [], [], [] + input_height, input_width = image_size + for image, target in zip(images, targets): + shapes.append(tuple(image.shape[:2])) + processed = preprocess_fn(image) + if not ( + isinstance(processed, tuple) + and len(processed) == 2 + and isinstance(processed[1], dict) + ): + raise ValueError( + "ADE20K preprocessing must return image data and letterbox metadata." + ) + processed_image, metadata = processed + ratio_pad = metadata.get("ratio_pad") + if ratio_pad is None: + raise ValueError( + "ADE20K preprocessing requires LetterBox ratio_pad metadata." + ) + processed_target, target_ratio_pad = letterbox_semantic_mask( + target, + [input_height, input_width], + ) + if target_ratio_pad != ratio_pad: + raise ValueError( + "ADE20K image and mask LetterBox geometry do not match." + ) + processed_images.append(processed_image) + processed_targets.append(processed_target) + ratio_pads.append(ratio_pad) + return ( + np.stack(processed_images), + np.stack(processed_targets), + shapes, + ratio_pads, + stems, + ) + + return torch.utils.data.DataLoader( + dataset, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=loader + ) + + +class CustomCityscapes(torch.utils.data.Dataset[tuple[np.ndarray, np.ndarray, str]]): + """Cityscapes validation dataset with paired RGB images and source-ID masks.""" + + IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp") + + def __init__(self, root: str) -> None: + """Validate the organizer's flat ``images/`` and ``annotations/`` layout.""" + + self.root = root + image_root = os.path.join(root, "images") + annotation_root = os.path.join(root, "annotations") + if not os.path.isdir(image_root) or not os.path.isdir(annotation_root): + raise FileNotFoundError( + f"Cityscapes requires images/ and annotations/ directories under: {root}" + ) + images = _unique_paths_by_stem( + image_root, self.IMG_EXTENSIONS, "Cityscapes images" + ) + annotations = _unique_paths_by_stem( + annotation_root, (".png",), "Cityscapes annotations" + ) + missing_annotations = sorted(set(images) - set(annotations)) + missing_images = sorted(set(annotations) - set(images)) + if missing_annotations or missing_images: + details = [] + if missing_annotations: + details.append( + f"images without annotations: {', '.join(missing_annotations[:5])}" + ) + if missing_images: + details.append( + f"annotations without images: {', '.join(missing_images[:5])}" + ) + raise ValueError( + f"Cityscapes image/annotation mismatch ({'; '.join(details)})." + ) + if not images: + raise ValueError(f"Cityscapes contains no image/annotation pairs: {root}") + self.samples = [ + (images[stem], annotations[stem], stem) for stem in sorted(images) + ] + + def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray, str]: + """Load an image and map Cityscapes source IDs to contiguous train IDs.""" + + image_path, annotation_path, stem = self.samples[index] + image = cv2.imread(image_path, cv2.IMREAD_COLOR) + if image is None: + raise FileNotFoundError(f"Cityscapes image not found: {image_path}") + with Image.open(annotation_path) as annotation_image: + annotation = np.asarray(annotation_image) + if annotation.ndim == 3: + if annotation.shape[2] not in {3, 4} or not np.array_equal( + annotation[..., 0], annotation[..., 1] + ): + raise ValueError( + f"Cityscapes RGB annotation channels must contain identical source IDs: {annotation_path}" + ) + if not np.array_equal(annotation[..., 0], annotation[..., 2]): + raise ValueError( + f"Cityscapes RGB annotation channels must contain identical source IDs: {annotation_path}" + ) + annotation = annotation[..., 0] + if annotation.ndim != 2: + raise ValueError( + f"Cityscapes annotation must be grayscale or RGB-grayscale: {annotation_path}" + ) + if image.shape[:2] != annotation.shape: + raise ValueError( + "Cityscapes image and annotation shapes must match, " + f"got {image.shape[:2]} and {annotation.shape}: {stem}" + ) + if annotation.size and ( + int(annotation.min()) < 0 or int(annotation.max()) > 255 + ): + raise ValueError( + f"Cityscapes annotation values must be in [0, 255]: {annotation_path}" + ) + source_ids = np.unique(annotation.astype(np.uint8)) + known_ids = np.array([*range(34), 255], dtype=np.uint8) + unknown_ids = source_ids[~np.isin(source_ids, known_ids)] + if unknown_ids.size: + raise ValueError( + "Cityscapes annotation contains unsupported source IDs " + f"{unknown_ids.tolist()}: {annotation_path}" + ) + target = CITYSCAPES_SOURCE_TO_TRAIN_ID[annotation.astype(np.uint8)] + if not (target != 255).any(): + raise ValueError( + f"Cityscapes annotation contains no evaluable class IDs: {annotation_path}" + ) + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), target, stem + + def __len__(self) -> int: + """Return the number of paired validation samples.""" + + return len(self.samples) + + +def get_cityscapes_loader( + dataset: CustomCityscapes, + batch_size: int, + preprocess_fn: Callable, + image_size: tuple[int, int], +) -> torch.utils.data.DataLoader: + """Create a Cityscapes loader with image-matching letterbox geometry.""" + + def loader( + batch: list[Any], + ) -> tuple[ + np.ndarray, np.ndarray, list[tuple[int, int]], list[Any], tuple[str, ...] + ]: + images, targets, stems = zip(*batch) + processed_images, processed_targets, shapes, ratio_pads = [], [], [], [] + input_height, input_width = image_size + for image, target in zip(images, targets): + shapes.append(tuple(image.shape[:2])) + processed = preprocess_fn(image) + if not ( + isinstance(processed, tuple) + and len(processed) == 2 + and isinstance(processed[1], dict) + ): + raise ValueError( + "Cityscapes preprocessing must return image data and letterbox metadata." + ) + processed_image, metadata = processed + ratio_pad = metadata.get("ratio_pad") + if ratio_pad is None: + raise ValueError( + "Cityscapes preprocessing requires LetterBox ratio_pad metadata." + ) + processed_target, target_ratio_pad = letterbox_semantic_mask( + target, + [input_height, input_width], + ) + if target_ratio_pad != ratio_pad: + raise ValueError( + "Cityscapes image and mask LetterBox geometry do not match." + ) + processed_images.append(processed_image) + processed_targets.append(processed_target) + ratio_pads.append(ratio_pad) + return ( + np.stack(processed_images), + np.stack(processed_targets), + shapes, + ratio_pads, + stems, + ) + + return torch.utils.data.DataLoader( + dataset, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=loader + ) + + +class CustomDOTAv1(torch.utils.data.Dataset[tuple[np.ndarray, str, int, int]]): + """Custom DOTAv1 validation dataset for OBB evaluation. + + Attributes: + root: DOTAv1 dataset root. + image_root: Directory containing validation images. + ids: Image IDs derived from file stems. + """ + + IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff") + + def __init__(self, root: str) -> None: + """Initializes the DOTAv1 validation dataset. + + Args: + root: DOTAv1 root containing flat ``images/`` or legacy + ``images/val`` validation images. + + Raises: + FileNotFoundError: If the validation image directory is missing. + ValueError: If neither supported layout contains validation images. + """ + self.root = root + self.image_root = os.path.join(root, "images") + if not os.path.isdir(self.image_root): + raise FileNotFoundError( + f"DOTAv1 image directory not found: {self.image_root}" + ) + self.image_paths = self._find_image_paths(self.image_root) + legacy_image_root = os.path.join(self.image_root, "val") + if not self.image_paths and os.path.isdir(legacy_image_root): + self.image_root = legacy_image_root + self.image_paths = self._find_image_paths(self.image_root) + if not self.image_paths: + raise ValueError( + f"DOTAv1 validation images not found directly under {os.path.join(root, 'images')} " + "or its legacy `val` subdirectory." + ) + self.ids = [ + os.path.splitext(os.path.basename(path))[0] for path in self.image_paths + ] + if len(self.ids) != len(set(self.ids)): + raise ValueError( + "DOTAv1 validation images contain duplicate filename stems." + ) + + def _find_image_paths(self, image_root: str) -> list[str]: + """Return supported image files directly under a DOTAv1 image directory.""" + + return [ + os.path.join(image_root, file_name) + for file_name in sorted(os.listdir(image_root)) + if file_name.lower().endswith(self.IMG_EXTENSIONS) + ] + + def _load_image(self, image_path: str) -> np.ndarray: + """Load an image as RGB.""" + image = cv2.imread(image_path) + if image is None: + raise FileNotFoundError(f"Image not found: {image_path}") + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + def __getitem__(self, index: int) -> tuple[np.ndarray, str, int, int]: + """Get the image and metadata by index.""" + image_path = self.image_paths[index] + image = self._load_image(image_path) + height, width = image.shape[:2] + return image, self.ids[index], height, width + + def __len__(self) -> int: + """Return the number of validation images.""" + return len(self.image_paths) + + +def get_dota_loader( + dataset: CustomDOTAv1, batch_size: int, preprocess_fn: Callable +) -> torch.utils.data.DataLoader: + """Creates a DataLoader for DOTAv1 validation. + + Args: + dataset: The DOTAv1 dataset instance. + batch_size: Number of samples per batch. + preprocess_fn: Function used to preprocess images. + + Returns: + Configured DataLoader for DOTAv1. + """ + + def loader( + batch: list[Any], + ) -> tuple[np.ndarray, np.ndarray, list[Any], tuple[str, ...]]: + """Collate function for DOTAv1 DataLoader.""" + batch = list(filter(lambda x: x is not None, batch)) + images, image_ids, height, width = zip(*batch) + + processed_images = [] + ratio_pads = [] + for img in images: + processed = preprocess_fn(img) + if ( + isinstance(processed, tuple) + and len(processed) == 2 + and isinstance(processed[1], dict) + ): + processed_img, metadata = processed + ratio_pads.append(metadata.get("ratio_pad")) + else: + processed_img = processed + ratio_pads.append(None) + processed_images.append(processed_img) + + return ( + np.stack(processed_images, axis=0), + np.stack((np.array(height), np.array(width)), axis=1), + ratio_pads, + image_ids, + ) + + return torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=0, + collate_fn=loader, + ) + + +class CustomImageFolder(torch.utils.data.Dataset[tuple[Image.Image, int]]): + """Custom ImageFolder dataset for loading images from class-based directory structures. + + Expects data to be organized in the format: root/class_name/image.jpg. + + Attributes: + root (str): Root directory path. + classes (list[str]): List of class names found in the root directory. + class_to_idx (dict): Mapping from class name to class index. + samples (list[tuple]): List of (image_path, class_index) tuples. + """ + + def __init__(self, root: str) -> None: + """Initializes the CustomImageFolder instance. + + Args: + root (str): Path to the root directory. + """ + self.root = root + self.classes, self.class_to_idx = self.find_classes(root) + self.samples: list[tuple[str, int]] = [] + self.make_dataset() + + def make_dataset(self) -> None: + """Scans the root directory to create a list of samples.""" + instances = [] + for target_class in sorted(self.class_to_idx.keys()): + class_index = self.class_to_idx[target_class] + target_dir = os.path.join(self.root, target_class) + if not os.path.isdir(target_dir): + continue + for fname in sorted(os.listdir(target_dir)): + path = os.path.join(target_dir, fname) + if not os.path.isfile(path): + continue + if os.path.splitext(fname)[1].lower() not in IMAGE_SUFFIXES: + continue + instances.append((path, class_index)) + + self.samples = instances + + def loader(self, path: str) -> Image.Image: + """Load image from path using PIL.""" + with open(path, "rb") as f: + img = Image.open(f) + return img.convert("RGB") + + def find_classes(self, directory: str) -> tuple[list[str], dict[str, int]]: + """Find classes in the specified directory.""" + classes = sorted([d.name for d in os.scandir(directory) if d.is_dir()]) + class_to_idx = {cls: i for i, cls in enumerate(classes)} + return classes, class_to_idx + + def __getitem__(self, index: int) -> tuple[Image.Image, int]: + """ + Get sample and target at the specified index. + Args: + index (int): Index of the sample to retrieve. + Returns: + tuple: (sample, target) where sample is the loaded image and target is the class index. + """ + path, target = self.samples[index] + sample = self.loader(path) + return sample, target + + def __len__(self) -> int: + """ + Return the total number of samples. + Returns: + int: Number of samples in the dataset. + """ + return len(self.samples) + + +def get_imagenet_loader( + dataset: CustomImageFolder, batch_size: int, preprocess_fn: Callable +) -> torch.utils.data.DataLoader: + """Creates a DataLoader for the ImageNet dataset. + + Args: + dataset (CustomImageFolder): The dataset instance to load from. + batch_size (int): Number of samples per batch. + preprocess_fn (Callable): Function used to preprocess images. + + Returns: + torch.utils.data.DataLoader: A configured DataLoader for the ImageNet dataset. + """ + + def loader(batch: list[Any]) -> tuple[np.ndarray, np.ndarray]: + """Collate function for ImageNet DataLoader.""" + batch = list(filter(lambda x: x is not None, batch)) # remove None + images, labels = zip(*batch) + processed_images = [] + for img in images: + img = preprocess_fn(img) + processed_images.append(img) + + return ( + np.stack(processed_images, axis=0), + np.array(labels), + ) # BHWC, labels + + return torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=0, + collate_fn=loader, + ) + + +class CustomWiderFaceDataset(torch.utils.data.Dataset[tuple[np.ndarray, str, str]]): + """Custom dataset class for the WiderFace dataset. + + Attributes: + root (str): Path to the root directory containing WiderFace images. + classes (list[str]): List of class/event names found in the root. + samples (list[tuple]): List of (image_path, class_name, file_name) tuples. + """ + + def __init__(self, root: str) -> None: + """Initialize the custom WiderFace dataset. + + Args: + root (str): Path to the directory containing WiderFace images. + """ + self.root = root + self.classes = self.find_classes(root) + self.samples: list[tuple[str, str, str]] = [] + self.make_dataset() + + def make_dataset(self) -> None: + """Scans the root directory to create a list of samples.""" + instances = [] + for target_class in self.classes: + target_dir = os.path.join(self.root, target_class) + if not os.path.isdir(target_dir): + continue + for fname in sorted(os.listdir(target_dir)): + path = os.path.join(target_dir, fname) + if ( + os.path.islink(path) + or not os.path.isfile(path) + or os.path.splitext(fname)[1].lower() not in IMAGE_SUFFIXES + ): + continue + instances.append((path, target_class, fname)) + + self.samples = instances + + def loader(self, image_path: str) -> np.ndarray: + """Load image by image path""" + image = cv2.imread(image_path) # Load image (BGR format) + if image is None: + raise FileNotFoundError(f"Image not found: {image_path}") + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB + + def find_classes(self, directory: str) -> list[str]: + """Find classes in the specified directory.""" + unsorted_classes = [d.name for d in os.scandir(directory) if d.is_dir()] + class_to_idx = {} + for cls_name in unsorted_classes: + cls_idx = int(cls_name.split("--")[0]) + class_to_idx[cls_name] = cls_idx + + sorted_classes = sorted( + class_to_idx.keys(), key=lambda x: class_to_idx[x] + ) # sort by dictionary value with ascending order + + return sorted_classes + + def __getitem__(self, index: int) -> tuple[np.ndarray, str, str]: + """ + Get the image and target by index. + Args: + index (int): Index of the sample to retrieve. + Returns: + tuple: (image, target_class, fname) where image is the loaded image in RGB format. + """ + image_path, target_class, fname = self.samples[index] + image = self.loader(image_path) + + return image, target_class, fname + + def __len__(self) -> int: + """ + Return the total number of images. + Returns: + int: Number of images in the dataset. + """ + return len(self.samples) + + +CustomWiderface = CustomWiderFaceDataset + + +def get_widerface_loader( + dataset: CustomWiderFaceDataset, batch_size: int, preprocess_fn: Callable +) -> torch.utils.data.DataLoader: + """Creates a DataLoader for the WiderFace dataset. + + Args: + dataset (CustomWiderFaceDataset): The dataset instance to load from. + batch_size (int): Number of samples per batch. + preprocess_fn (Callable): Function used to preprocess images. + + Returns: + torch.utils.data.DataLoader: A configured DataLoader for the WiderFace dataset. + """ + + def loader( + batch: list[Any], + ) -> tuple[ + np.ndarray, np.ndarray, list[Any | None], tuple[str, ...], tuple[str, ...] + ]: + """Collate function for WiderFace DataLoader.""" + batch = list(filter(lambda x: x is not None, batch)) + images, target_classes, fnames = zip(*batch) + processed_images = [] + heights = [] + widths = [] + ratio_pads = [] + for img in images: + height, width = img.shape[:2] + processed = preprocess_fn(img) + if isinstance(processed, tuple): + processed_img, metadata = processed + ratio_pads.append(metadata.get("ratio_pad")) + else: + processed_img = processed + ratio_pads.append(None) + processed_images.append(processed_img) + heights.append(height) + widths.append(width) + + return ( + np.stack(processed_images, axis=0), + np.stack((heights, widths), axis=1), + ratio_pads, + target_classes, + fnames, + ) + + return torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=0, + collate_fn=loader, + ) diff --git a/mblt_vision/utils/datasets/dotav1.py b/mblt_vision/utils/datasets/dotav1.py new file mode 100644 index 0000000..6e5464f --- /dev/null +++ b/mblt_vision/utils/datasets/dotav1.py @@ -0,0 +1,59 @@ +""" +DOTAv1 dataset constants and utilities. +""" + +from __future__ import annotations + +from ...datasets import get_dataset_class_names + +DOTAV1_PALETTE_BGR = [ + (220, 20, 60), + (119, 11, 32), + (0, 0, 142), + (0, 0, 230), + (106, 0, 228), + (0, 60, 100), + (0, 80, 100), + (0, 0, 70), + (0, 0, 192), + (250, 170, 30), + (100, 170, 30), + (220, 220, 0), + (175, 116, 175), + (250, 0, 30), + (165, 42, 42), +] + + +def get_dotav1_class_num() -> int: + """Returns the number of DOTAv1 classes. + + Returns: + Number of DOTAv1 classes. + """ + return len(get_dataset_class_names("dotav1")) + + +def get_dotav1_palette(index: int) -> tuple[int, int, int]: + """Returns the DOTAv1 visualization color for a class index. + + Args: + index: Class index. + + Returns: + OpenCV BGR color tuple. + """ + + return DOTAV1_PALETTE_BGR[index % len(DOTAV1_PALETTE_BGR)] + + +def get_dotav1_label(index: int) -> str: + """Returns the DOTAv1 class name for a class index. + + Args: + index: Class index. + + Returns: + DOTAv1 class name. + """ + return get_dataset_class_names("dotav1")[index] diff --git a/mblt_vision/utils/datasets/imagenet.py b/mblt_vision/utils/datasets/imagenet.py new file mode 100644 index 0000000..b5a7596 --- /dev/null +++ b/mblt_vision/utils/datasets/imagenet.py @@ -0,0 +1,29 @@ +""" +ImageNet dataset constants and utilities. +""" + +from __future__ import annotations + +from ...datasets import get_dataset_class_names + + +def get_imagenet_label(idx: int) -> str: + """Get the descriptive label for an ImageNet class index. + + Args: + idx (int): ImageNet class index (0-999). + + Returns: + str: Primary class label string. + """ + label_str = get_dataset_class_names("imagenet")[idx] + return label_str.split(",")[0] + + +def get_imagenet_class_num() -> int: + """Get the total number of classes in the ImageNet dataset. + + Returns: + int: Total number of ImageNet classes. + """ + return len(get_dataset_class_names("imagenet")) diff --git a/mblt_vision/utils/datasets/organizer.py b/mblt_vision/utils/datasets/organizer.py new file mode 100644 index 0000000..f8c3484 --- /dev/null +++ b/mblt_vision/utils/datasets/organizer.py @@ -0,0 +1,2145 @@ +""" +Utilities for organizing datasets. +""" + +from __future__ import annotations + +import concurrent.futures +import cv2 +import hashlib +import json +import math +import os +import re +import shutil +import stat +import tarfile +import xml.etree.ElementTree as ET +import zipfile +from collections.abc import Iterable +from pathlib import Path +from tempfile import TemporaryDirectory, mkdtemp +from time import sleep +from typing import Protocol, TypeGuard +from urllib.parse import urlparse + +import requests +import numpy as np +from gdown.download import download +from gdown.download_folder import download_folder +from PIL import Image +from tqdm import tqdm + +from ...datasets import get_dataset_config +from .cityscapes import CITYSCAPES_SOURCE_TO_TRAIN_ID +from .readiness import ( + ADE20K_METADATA_FILES, + ADE20K_VALIDATION_SAMPLE_COUNT, + CITYSCAPES_SAMPLE_ID_PATTERN, + CITYSCAPES_VALIDATION_SAMPLE_COUNT, + DOTAV1_VALIDATION_SAMPLE_COUNT, + IMAGE_SUFFIXES, + NYU_DEPTH_VALIDATION_SAMPLE_COUNT, + _canonicalize_quadrilateral, + _path_has_symlink_component, + _polygon_has_positive_image_overlap, + dataset_ready, +) + +DOWNLOAD_CHUNK_SIZE = 1 * 1024 * 1024 +DOWNLOAD_RETRY_LIMIT = 4 +DOWNLOAD_RETRY_BACKOFF_SECONDS = 2.0 +DOWNLOAD_TIMEOUT = (10, 30) +DOTAV1_DOWNLOAD_CONFIG = get_dataset_config("dotav1")["download"] +DOTAV1_GOOGLE_DRIVE_ARCHIVES = { + DOTAV1_DOWNLOAD_CONFIG["images_archive"], + DOTAV1_DOWNLOAD_CONFIG["labels_archive"], +} +DOTAV1_CLASS_TO_IDX = { + name: int(index) for index, name in get_dataset_config("dotav1")["names"].items() +} +COCO_DOWNLOAD_CONFIG = get_dataset_config("coco")["download"] +ADE20K_DOWNLOAD_CONFIG = get_dataset_config("ade20k")["download"] +NYU_DEPTH_URL = ( + "https://github.com/ultralytics/assets/releases/download/v0.0.0/nyu-depth.zip" +) +ADE20K_URL = ADE20K_DOWNLOAD_CONFIG["url"] +CITYSCAPES_IMAGE_SUFFIX = "_leftImg8bit.png" +CITYSCAPES_ANNOTATION_SUFFIX = "_gtFine_labelIds.png" +IMAGENET_SYNSET_PATTERN = re.compile(r"n\d{8}") +RETRYABLE_HTTP_STATUS_CODES = frozenset({408, 429}) +CONTENT_RANGE_PATTERN = re.compile(r"^bytes (\d+)-(\d+)/(\d+|\*)$") +UNSATISFIABLE_CONTENT_RANGE_PATTERN = re.compile(r"^bytes \*/(\d+)$") +PINNED_ARCHIVE_SHA256 = { + COCO_DOWNLOAD_CONFIG["images"]: COCO_DOWNLOAD_CONFIG["images_sha256"], + COCO_DOWNLOAD_CONFIG["annotations"]: COCO_DOWNLOAD_CONFIG["annotations_sha256"], + ADE20K_DOWNLOAD_CONFIG["url"]: ADE20K_DOWNLOAD_CONFIG["sha256"], +} + + +def _resolve_organizer_output_dir(output_dir: str | None, dataset_name: str) -> str: + """Return an explicit output directory or the lazily resolved artifact cache.""" + + if output_dir is not None: + return os.path.expanduser(output_dir) + from mblt_vision.wrapper import get_mobilint_cache_dir + + return os.path.join(get_mobilint_cache_dir(), "datasets", dataset_name) + + +def _replace_staged_directories( + replacements: Iterable[tuple[str, str]], + output_parent_dir: str, + backup_prefix: str, +) -> None: + """Atomically install staged directories while preserving failed rollback backups. + + Args: + replacements: Pairs of staged and destination directories. + output_parent_dir: Parent directory where the backup directory is created. + backup_prefix: Prefix identifying the temporary backup directory. + + Raises: + OSError: If installation or rollback fails. A failed rollback leaves its + backup directory in place and includes its path in the error. + """ + + replacement_list = list(replacements) + backup_dir = mkdtemp(dir=output_parent_dir, prefix=backup_prefix) + backups: dict[str, str] = {} + installed_dirs: list[str] = [] + try: + for _, destination_dir in replacement_list: + if os.path.lexists(destination_dir): + backup_path = os.path.join( + backup_dir, os.path.basename(destination_dir) + ) + os.replace(destination_dir, backup_path) + backups[destination_dir] = backup_path + for staged_dir, destination_dir in replacement_list: + os.makedirs(os.path.dirname(destination_dir), exist_ok=True) + os.replace(staged_dir, destination_dir) + installed_dirs.append(destination_dir) + except OSError: + try: + for directory in installed_dirs: + if os.path.isdir(directory) and not os.path.islink(directory): + shutil.rmtree(directory) + elif os.path.lexists(directory): + os.remove(directory) + for destination_dir, backup_path in backups.items(): + os.makedirs(os.path.dirname(destination_dir), exist_ok=True) + os.replace(backup_path, destination_dir) + except OSError as rollback_error: + raise OSError( + f"Dataset installation rollback failed; backups are preserved at {backup_dir}." + ) from rollback_error + shutil.rmtree(backup_dir) + raise + shutil.rmtree(backup_dir) + + +def _validate_staged_dataset( + staged_output_dir: str, + dataset: str, + tasks: Iterable[str], +) -> None: + """Validate a complete staged dataset before replacing its managed cache. + + Args: + staged_output_dir: Root of the staged organized dataset. + dataset: Validation dataset taxonomy. + tasks: Tasks whose required metadata and files must all be ready. + + Raises: + ValueError: If the staged dataset is incomplete or has mismatched identity. + """ + + if not all(dataset_ready(staged_output_dir, task, dataset) for task in tasks): + raise ValueError( + f"Staged {dataset} validation dataset is incomplete or has mismatched metadata; " + "the existing dataset cache was not replaced." + ) + _validate_staged_payloads(Path(staged_output_dir), dataset) + + +def _validate_staged_payloads(staged_root: Path, dataset: str) -> None: + """Decode staged data files before a structurally valid cache is replaced.""" + + image_roots = { + "imagenet": (staged_root,), + "widerface": (staged_root / "images",), + "dotav1": (staged_root / "images",), + "ade20k": (staged_root / "images",), + "cityscapes": (staged_root / "images",), + } + for image_root in image_roots.get(dataset, ()): + for image_path in image_root.rglob("*"): + if image_path.is_file() and image_path.suffix.lower() in { + ".bmp", + ".jpeg", + ".jpg", + ".png", + ".tif", + ".tiff", + ".webp", + }: + if cv2.imread(str(image_path), cv2.IMREAD_COLOR) is None: + raise ValueError( + f"Staged {dataset} image is unreadable: {image_path}." + ) + + if dataset == "coco": + _validate_staged_coco_image_geometry(staged_root) + elif dataset in {"ade20k", "cityscapes"}: + _validate_staged_semantic_masks(staged_root, dataset) + elif dataset == "dotav1": + _validate_staged_dotav1_labels(staged_root) + + +def _validate_staged_coco_image_geometry(staged_root: Path) -> None: + """Compare every staged COCO image with its JSON-declared geometry.""" + + image_root = staged_root / "val2017" + decoded_shapes: dict[str, tuple[int, int]] = {} + for annotation_path in sorted(staged_root.glob("*_val2017.json")): + try: + annotation = json.loads(annotation_path.read_text(encoding="utf-8")) + image_records = annotation["images"] + except (json.JSONDecodeError, KeyError, OSError, TypeError) as exc: + raise ValueError( + f"Staged COCO annotation is unreadable: {annotation_path}." + ) from exc + if not isinstance(image_records, list): + raise ValueError( + f"Staged COCO annotation has an invalid images table: {annotation_path}." + ) + for record in image_records: + if not isinstance(record, dict): + raise ValueError( + f"Staged COCO annotation has an invalid image record: {annotation_path}." + ) + file_name, height, width = ( + record.get("file_name"), + record.get("height"), + record.get("width"), + ) + if ( + not isinstance(file_name, str) + or not file_name + or Path(file_name).is_absolute() + or ".." in Path(file_name).parts + or not isinstance(height, int) + or isinstance(height, bool) + or not isinstance(width, int) + or isinstance(width, bool) + or height <= 0 + or width <= 0 + ): + raise ValueError( + f"Staged COCO image metadata is invalid: {annotation_path}." + ) + image_shape = decoded_shapes.get(file_name) + if image_shape is None: + image = cv2.imread(str(image_root / file_name), cv2.IMREAD_COLOR) + if image is None: + raise ValueError( + f"Staged COCO image is unreadable: {image_root / file_name}." + ) + image_shape = (int(image.shape[0]), int(image.shape[1])) + decoded_shapes[file_name] = image_shape + if image_shape != (height, width): + raise ValueError( + "Staged COCO image geometry does not match annotation metadata for " + f"{file_name}: image {image_shape}, annotation {(height, width)}." + ) + + +def _validate_staged_semantic_masks(staged_root: Path, dataset: str) -> None: + """Validate decoded semantic targets against their paired staged images.""" + + image_dir = staged_root / "images" + annotation_dir = staged_root / "annotations" + for annotation_path in sorted(annotation_dir.glob("*.png")): + image_path = next( + ( + candidate + for candidate in image_dir.glob(f"{annotation_path.stem}.*") + if candidate.suffix.lower() in {".jpg", ".jpeg", ".png"} + ), + None, + ) + if image_path is None: + raise ValueError( + f"Staged {dataset} target has no paired image: {annotation_path}." + ) + image = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + try: + with Image.open(annotation_path) as annotation_image: + annotation = np.asarray(annotation_image) + except OSError as exc: + raise ValueError( + f"Staged {dataset} annotation is unreadable: {annotation_path}." + ) from exc + if dataset == "cityscapes" and annotation.ndim == 3: + if ( + annotation.shape[2] not in {3, 4} + or not np.array_equal(annotation[..., 0], annotation[..., 1]) + or not np.array_equal(annotation[..., 0], annotation[..., 2]) + ): + raise ValueError( + f"Staged Cityscapes annotation must be grayscale or RGB-grayscale: {annotation_path}." + ) + annotation = annotation[..., 0] + if image is None or annotation.ndim != 2 or annotation.shape != image.shape[:2]: + raise ValueError( + f"Staged {dataset} image and annotation geometry is invalid: {annotation_path}." + ) + if dataset == "ade20k" and ( + annotation.dtype != np.uint8 + or (annotation.size and int(annotation.max()) > 150) + ): + raise ValueError( + f"Staged ADE20K annotation must be an 8-bit mask with values in [0, 150]: {annotation_path}." + ) + if dataset == "cityscapes": + valid_ids = (annotation <= 33) | (annotation == 255) + if not np.all(valid_ids): + invalid_ids = np.unique(annotation[~valid_ids]) + raise ValueError( + "Staged Cityscapes annotation contains unsupported source IDs " + f"{invalid_ids.tolist()}; expected IDs in [0, 33] or 255: {annotation_path}." + ) + has_evaluable_class = bool( + ( + CITYSCAPES_SOURCE_TO_TRAIN_ID[annotation.astype(np.uint8)] != 255 + ).any() + ) + else: + has_evaluable_class = bool((annotation > 0).any()) + if not has_evaluable_class: + raise ValueError( + f"Staged {dataset} annotation contains no evaluable class IDs: {annotation_path}." + ) + + +def _validate_staged_dotav1_labels(staged_root: Path) -> None: + """Validate both DOTAv1 label representations before cache replacement.""" + + label_dirs = { + "normalized": staged_root / "labels" / "val", + "original": staged_root / "labels" / "val_original", + } + image_paths = { + path.stem: path + for path in (staged_root / "images").iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + } + valid_indices = set(DOTAV1_CLASS_TO_IDX.values()) + label_stems = {kind: set() for kind in label_dirs} + positive_stems = {kind: set() for kind in label_dirs} + for kind, label_dir in label_dirs.items(): + for label_path in sorted(label_dir.glob("*.txt")): + label_stems[kind].add(label_path.stem) + image_path = image_paths.get(label_path.stem) + if image_path is None: + raise ValueError( + f"Staged DOTAv1 label has no matching image: {label_path}." + ) + image = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + if image is None: + raise ValueError(f"Unable to decode staged DOTAv1 image: {image_path}.") + height, width = image.shape[:2] + has_positive_target = False + seen_targets: set[tuple[int | str, tuple[float, ...]]] = set() + for line_number, line in enumerate( + label_path.read_text(encoding="utf-8").splitlines(), start=1 + ): + fields = line.split() + if not fields or ( + kind == "original" + and ( + fields[0].startswith("imagesource:") + or fields[0].startswith("gsd:") + ) + ): + continue + min_fields = 9 if kind == "normalized" else 10 + if len(fields) < min_fields: + raise ValueError( + f"Malformed staged DOTAv1 {kind} annotation at " + f"{label_path}:{line_number}: expected at least {min_fields} fields." + ) + try: + coordinates = [ + float(value) + for value in ( + fields[1:9] if kind == "normalized" else fields[:8] + ) + ] + except ValueError as exc: + raise ValueError( + f"Malformed staged DOTAv1 coordinates at {label_path}:{line_number}." + ) from exc + if not all(math.isfinite(value) for value in coordinates): + raise ValueError( + f"Staged DOTAv1 coordinates must be finite at {label_path}:{line_number}." + ) + points = np.asarray(coordinates, dtype=np.float64).reshape(4, 2) + signed_double_area = np.dot( + points[:, 0], np.roll(points[:, 1], -1) + ) - np.dot(points[:, 1], np.roll(points[:, 0], -1)) + if abs(signed_double_area) <= 0: + raise ValueError( + f"Staged DOTAv1 polygon must have positive area at {label_path}:{line_number}." + ) + _validate_dotav1_polygon_vertices(coordinates, label_path, line_number) + image_coordinates = coordinates.copy() + if kind == "normalized": + image_coordinates = [ + coordinate * (width if index % 2 == 0 else height) + for index, coordinate in enumerate(coordinates) + ] + if not _polygon_has_positive_image_overlap( + image_coordinates, (height, width) + ): + raise ValueError( + "Staged DOTAv1 polygon must overlap its source image at " + f"{label_path}:{line_number}." + ) + if kind == "normalized": + try: + class_index = int(fields[0]) + except ValueError as exc: + raise ValueError( + f"Malformed staged DOTAv1 class index at {label_path}:{line_number}." + ) from exc + difficulty = fields[9] if len(fields) >= 10 else "0" + if class_index not in valid_indices: + raise ValueError( + f"Unsupported staged DOTAv1 class index at {label_path}:{line_number}." + ) + target_class: int | str = class_index + else: + difficulty = fields[9] + if fields[8] not in DOTAV1_CLASS_TO_IDX: + raise ValueError( + f"Unsupported staged DOTAv1 class at {label_path}:{line_number}." + ) + target_class = fields[8] + if difficulty not in {"0", "1", "2"}: + raise ValueError( + f"Unsupported staged DOTAv1 difficulty flag at {label_path}:{line_number}." + ) + target_key = ( + target_class, + _canonicalize_quadrilateral(image_coordinates), + ) + if target_key in seen_targets: + raise ValueError( + "Duplicate staged DOTAv1 annotation target at " + f"{label_path}:{line_number}." + ) + seen_targets.add(target_key) + has_positive_target |= difficulty == "0" + if has_positive_target: + positive_stems[kind].add(label_path.stem) + authoritative_positive_stems = positive_stems["original"] | ( + positive_stems["normalized"] - label_stems["original"] + ) + if not authoritative_positive_stems: + raise ValueError( + "Staged DOTAv1 dataset must contain at least one non-difficult target." + ) + + +def _validate_dotav1_polygon_vertices( + coordinates: list[float], annotation_path: str | Path, line_number: int +) -> None: + """Require four distinct, consistently ordered DOTAv1 quadrilateral vertices.""" + + points = np.asarray(coordinates, dtype=np.float64).reshape(4, 2) + if len(np.unique(points, axis=0)) != 4: + raise ValueError( + "DOTAv1 polygon must contain four distinct vertices in " + f"{annotation_path} at line {line_number}." + ) + edges = np.roll(points, -1, axis=0) - points + next_edges = np.roll(edges, -1, axis=0) + turns = edges[:, 0] * next_edges[:, 1] - edges[:, 1] * next_edges[:, 0] + if not (np.all(turns > 0) or np.all(turns < 0)): + raise ValueError( + "DOTAv1 polygon vertices must be consistently ordered in " + f"{annotation_path} at line {line_number}." + ) + + +class _GoogleDriveDownloadEntry(Protocol): + """The public attributes needed from a gdown folder-listing entry.""" + + id: str + path: str + + +def _is_google_drive_download_entry( + value: object, +) -> TypeGuard[_GoogleDriveDownloadEntry]: + """Returns whether a folder-listing value has the Google Drive file attributes needed here.""" + + return isinstance(getattr(value, "id", None), str) and isinstance( + getattr(value, "path", None), str + ) + + +def _is_url(path_or_url: str) -> bool: + """Returns whether the given string looks like an HTTP(S) URL.""" + parsed = urlparse(path_or_url) + return parsed.scheme in {"http", "https"} and bool(parsed.netloc) + + +def _verify_archive_sha256( + archive_path: str, expected_sha256: str, source_url: str +) -> None: + """Verify a downloaded archive before it can be extracted.""" + + digest = hashlib.sha256() + with open(archive_path, "rb") as archive: + for chunk in iter(lambda: archive.read(DOWNLOAD_CHUNK_SIZE), b""): + digest.update(chunk) + actual_sha256 = digest.hexdigest() + if actual_sha256 != expected_sha256: + Path(archive_path).unlink(missing_ok=True) + raise ValueError( + f"Downloaded archive SHA-256 mismatch for {source_url}: " + f"expected {expected_sha256}, got {actual_sha256}." + ) + + +def _has_expected_resume_offset(content_range: str | None, existing_size: int) -> bool: + """Return whether a partial response begins exactly after local archive bytes.""" + + if content_range is None: + return False + match = CONTENT_RANGE_PATTERN.fullmatch(content_range) + if match is None: + return False + start, end, total = match.groups() + return ( + int(start) == existing_size + and int(end) >= int(start) + and (total == "*" or int(end) < int(total)) + ) + + +def _is_completed_range_response(content_range: str | None, existing_size: int) -> bool: + """Return whether a 416 confirms that the local archive is complete.""" + + if content_range is None: + return False + match = UNSATISFIABLE_CONTENT_RANGE_PATTERN.fullmatch(content_range) + return match is not None and int(match.group(1)) == existing_size + + +def _restart_partial_download(local_path: str, url: str) -> None: + """Discard an invalid partial archive before retrying from byte zero.""" + + Path(local_path).unlink(missing_ok=True) + print( + f"Server returned an invalid resume response for {os.path.basename(local_path)}; " + "restarting from byte zero." + ) + + +def _download_url(url: str, local_path: str, expected_sha256: str | None = None) -> str: + """Downloads a URL to a local file with progress and resume support. + + Args: + url: HTTP(S) URL to download. + local_path: Destination file path. + expected_sha256: Optional pinned SHA-256 digest to verify before return. + + Returns: + The local destination path. + + Raises: + RuntimeError: If all download attempts fail. + """ + os.makedirs(os.path.dirname(local_path), exist_ok=True) + + for attempt in range(1, DOWNLOAD_RETRY_LIMIT + 1): + existing_size = os.path.getsize(local_path) if os.path.exists(local_path) else 0 + headers: dict[str, str] = {} + mode = "wb" + if existing_size > 0: + headers["Range"] = f"bytes={existing_size}-" + mode = "ab" + + try: + with requests.get( + url, stream=True, timeout=DOWNLOAD_TIMEOUT, headers=headers + ) as response: + if response.status_code == 416 and existing_size > 0: + if _is_completed_range_response( + response.headers.get("Content-Range"), existing_size + ): + if expected_sha256 is not None: + _verify_archive_sha256(local_path, expected_sha256, url) + return local_path + _restart_partial_download(local_path, url) + continue + response.raise_for_status() + + if response.status_code == 200 and existing_size > 0: + existing_size = 0 + mode = "wb" + elif existing_size > 0 and ( + response.status_code != 206 + or not _has_expected_resume_offset( + response.headers.get("Content-Range"), existing_size + ) + ): + _restart_partial_download(local_path, url) + continue + + total_size = response.headers.get("Content-Length") + total_bytes = ( + existing_size + int(total_size) if total_size is not None else None + ) + + desc = f"Downloading {os.path.basename(local_path)}" + with tqdm( + total=total_bytes, + initial=existing_size, + unit="B", + unit_scale=True, + unit_divisor=1024, + desc=desc, + ) as pbar: + with open(local_path, mode) as file_obj: + for chunk in response.iter_content( + chunk_size=DOWNLOAD_CHUNK_SIZE + ): + if not chunk: + continue + file_obj.write(chunk) + pbar.update(len(chunk)) + if expected_sha256 is not None: + _verify_archive_sha256(local_path, expected_sha256, url) + return local_path + except ( + requests.ConnectionError, + requests.Timeout, + requests.exceptions.ChunkedEncodingError, + requests.HTTPError, + ) as exc: + if isinstance(exc, requests.HTTPError): + status_code = getattr(exc.response, "status_code", None) + if not isinstance(status_code, int) or ( + status_code not in RETRYABLE_HTTP_STATUS_CODES + and not 500 <= status_code < 600 + ): + raise + if attempt == DOWNLOAD_RETRY_LIMIT: + raise RuntimeError( + f"Failed to download {url} after {DOWNLOAD_RETRY_LIMIT} attempts." + ) from exc + resumed_size = ( + os.path.getsize(local_path) if os.path.exists(local_path) else 0 + ) + print( + f"Download attempt failed for {os.path.basename(local_path)}; " + f"retrying from {resumed_size} bytes (attempt {attempt + 1}/{DOWNLOAD_RETRY_LIMIT})..." + ) + sleep(DOWNLOAD_RETRY_BACKOFF_SECONDS * attempt) + + raise RuntimeError( + f"Failed to download {url} after {DOWNLOAD_RETRY_LIMIT} attempts." + ) + + +def _should_download_serially(path_or_urls: list[str]) -> bool: + """Returns whether URL inputs should be downloaded one by one. + + Dataset hosts such as ImageNet often throttle concurrent archive downloads + from the same origin. Serializing same-host downloads is slower in the best + case, but much more stable for the large validation archives used here. + """ + + hosts = [ + urlparse(path_or_url).netloc + for path_or_url in path_or_urls + if _is_url(path_or_url) + ] + return len(hosts) > 1 and len(set(hosts)) == 1 + + +def _download_if_url(path_or_url: str, download_dir: str) -> str: + """Downloads a remote dataset archive when needed. + + Args: + path_or_url: Local path or HTTP(S) URL pointing to a dataset archive. + download_dir: Directory to store downloaded archives. + + Returns: + A local filesystem path to the archive or directory. + + Raises: + ValueError: If the URL path does not contain a filename. + """ + if not _is_url(path_or_url): + return path_or_url + + parsed = urlparse(path_or_url) + if parsed.scheme != "https": + raise ValueError( + "Dataset archive URLs must use HTTPS. Download the archive locally " + f"and provide its path instead: {path_or_url}" + ) + filename = os.path.basename(parsed.path) + if not filename: + raise ValueError(f"Unable to determine a filename from URL: {path_or_url}") + + local_path = os.path.join(download_dir, filename) + print(f"Downloading dataset archive from {path_or_url} to {local_path}...") + _download_url( + path_or_url, + local_path, + expected_sha256=PINNED_ARCHIVE_SHA256.get(path_or_url), + ) + print("Download completed") + return local_path + + +def _resolve_source(path_or_url: str, download_dir: str) -> str: + """Resolves a local path for a dataset source.""" + + return _download_if_url(path_or_url, download_dir) + + +def _resolve_sources(path_or_urls: list[str], download_dir: str) -> list[str]: + """Resolves multiple dataset sources, downloading URL inputs in parallel.""" + + if _should_download_serially(path_or_urls): + return [ + _resolve_source(path_or_url, download_dir) for path_or_url in path_or_urls + ] + + local_paths: list[str | None] = [None] * len(path_or_urls) + futures: dict[concurrent.futures.Future[str], int] = {} + with concurrent.futures.ThreadPoolExecutor( + max_workers=min(4, len(path_or_urls)) + ) as executor: + for idx, path_or_url in enumerate(path_or_urls): + if _is_url(path_or_url): + futures[executor.submit(_resolve_source, path_or_url, download_dir)] = ( + idx + ) + else: + local_paths[idx] = path_or_url + + for future in concurrent.futures.as_completed(futures): + local_paths[futures[future]] = future.result() + + return [path for path in local_paths if path is not None] + + +def _get_object_name(obj: ET.Element, xml_file: str) -> str: + """Extracts a non-empty object name from an ImageNet annotation node. + + Args: + obj: XML ``object`` element from an annotation file. + xml_file: Source XML filename used for error context. + + Returns: + The validated object name. + + Raises: + ValueError: If the object name node is missing or empty. + """ + name_element = obj.find("name") + if name_element is None or name_element.text is None: + raise ValueError(f"XML file {xml_file} has an object without a valid name") + + object_name = name_element.text.strip() + if not object_name: + raise ValueError(f"XML file {xml_file} has an object with an empty name") + if IMAGENET_SYNSET_PATTERN.fullmatch(object_name) is None: + raise ValueError( + f"XML file {xml_file} has invalid ImageNet synset name {object_name!r}; " + "expected n########." + ) + + return object_name + + +def _imagenet_class_output_dir(staged_output_dir: str, object_name: str) -> Path: + """Return a containment-checked class directory below the staging root.""" + + staged_root = Path(staged_output_dir).resolve() + class_dir = (staged_root / object_name).resolve() + if class_dir.parent != staged_root: + raise ValueError( + f"ImageNet class directory escapes staging root: {object_name!r}." + ) + return class_dir + + +def construct_imagenet(image_dir: str, xml_dir: str, output_dir: str) -> None: + """Constructs the ImageNet dataset by organizing images into category folders. + + Args: + image_dir (str): Directory containing the ImageNet validation images. + xml_dir (str): Directory containing the ImageNet bounding box XML files. + output_dir (str): Directory where the organized dataset will be stored. + + Raises: + ValueError: If an XML file has no objects or contains multiple object names. + ValueError: If the number of XML files and images do not match. + """ + + xml_count = len(os.listdir(xml_dir + "/val")) + image_count = len(os.listdir(image_dir)) + if xml_count != image_count: + raise ValueError( + f"Number of XML and image files do not match: {xml_count} != {image_count}." + ) + + # validate the XML files + pbar = tqdm(os.listdir(xml_dir + "/val"), desc="Validating XML files") + for xml_file in pbar: + xml_path = os.path.join(xml_dir + "/val", xml_file) + xml_tree = ET.parse(xml_path) + root = xml_tree.getroot() + + if len(root.findall("object")) < 1: + raise ValueError( + f"XML file {xml_file} has no object, but expected at least 1" + ) + + # check whether the object names in the XML files are the same + object_names = [ + _get_object_name(obj, xml_file) for obj in root.findall("object") + ] + if len(set(object_names)) != 1: + raise ValueError( + f"Object names in XML file {xml_file} are not the same. " + f"It has {len(set(object_names))} different object names." + ) + + pbar.close() + + output_dir = os.path.abspath(output_dir) + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".imagenet-staging-" + ) as staging_dir: + staged_output_dir = os.path.join(staging_dir, "imagenet") + + # construct the ImageNet dataset + pbar = tqdm(os.listdir(xml_dir + "/val"), desc="Constructing ImageNet dataset") + for xml_file in pbar: + xml_path = os.path.join(xml_dir + "/val", xml_file) + xml_tree = ET.parse(xml_path) + root = xml_tree.getroot() + object_name = _get_object_name(root.findall("object")[0], xml_file) + image_path = os.path.join(image_dir, xml_file.replace(".xml", ".JPEG")) + if not os.path.isfile(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + + class_dir = _imagenet_class_output_dir(staged_output_dir, object_name) + os.makedirs(class_dir, exist_ok=True) + shutil.copy( + image_path, + class_dir / os.path.basename(image_path), + ) + pbar.close() + + # validate the staged ImageNet dataset before replacing the managed output root + pbar = tqdm(os.listdir(staged_output_dir), desc="Validating ImageNet dataset") + print(f"Number of categories: {len(os.listdir(staged_output_dir))}") + for object_name in pbar: + num_images = len(os.listdir(os.path.join(staged_output_dir, object_name))) + if num_images != 50: + raise ValueError( + f"Object {object_name} has {num_images} images, but expected 50" + ) + pbar.close() + _validate_staged_dataset( + staged_output_dir, "imagenet", ("image_classification",) + ) + _replace_staged_directories( + ((staged_output_dir, output_dir),), + output_parent_dir, + ".imagenet-backup-", + ) + print("Each category has 50 images") + print("ImageNet dataset constructed successfully") + + +def organize_imagenet( + image_dir: str, + xml_dir: str, + output_dir: str | None = None, +) -> None: + """Organizes the ImageNet dataset, unpacking archives if necessary. + + Args: + image_dir (str): Path or URL to the image directory or archive (.tar). + xml_dir (str): Path or URL to the XML directory or archive (.tgz). + output_dir: Directory to store the organized dataset. Defaults to the + resolved Mobilint cache directory. + """ + output_dir = _resolve_organizer_output_dir(output_dir, "imagenet") + with TemporaryDirectory() as temp_dir: + local_image_dir, local_xml_dir = _resolve_sources( + [image_dir, xml_dir], temp_dir + ) + + if local_image_dir.endswith(".tar") and local_xml_dir.endswith(".tgz"): + print("Unpacking image and XML files to temporary directory...") + _safe_unpack_archive( + local_image_dir, os.path.join(temp_dir, "ILSVRC2012_img_val") + ) + _safe_unpack_archive( + local_xml_dir, os.path.join(temp_dir, "ILSVRC2012_bbox_val_v3") + ) + print("Unpacking completed") + construct_imagenet( + os.path.join(temp_dir, "ILSVRC2012_img_val"), + os.path.join(temp_dir, "ILSVRC2012_bbox_val_v3"), + output_dir, + ) + return + + construct_imagenet(local_image_dir, local_xml_dir, output_dir) + + +def construct_coco(image_dir: str, annotation_dir: str, output_dir: str) -> None: + """Constructs the COCO dataset by copying images and annotations to a target directory. + + Args: + image_dir (str): Directory containing COCO images. + annotation_dir (str): Directory containing COCO annotations. + output_dir (str): Directory where the organized dataset will be stored. + """ + print( + f"Constructing COCO dataset from {image_dir} and {annotation_dir} to {output_dir}" + ) + output_dir = os.path.abspath(output_dir) + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".coco-staging-" + ) as staging_dir: + staged_output_dir = os.path.join(staging_dir, "coco") + shutil.copytree(image_dir, os.path.join(staged_output_dir, "val2017")) + for file in os.listdir(os.path.join(annotation_dir, "annotations")): + if file.endswith("_val2017.json"): + shutil.copy( + os.path.join(annotation_dir, "annotations", file), + os.path.join(staged_output_dir, file), + ) + _validate_staged_dataset( + staged_output_dir, + "coco", + ("object_detection", "instance_segmentation", "pose_estimation"), + ) + _replace_staged_directories( + ((staged_output_dir, output_dir),), + output_parent_dir, + ".coco-backup-", + ) + print("Constructing COCO dataset completed") + + +def organize_coco( + image_dir: str, + annotation_dir: str, + output_dir: str | None = None, +) -> None: + """Organizes the COCO dataset, unpacking archives if necessary. + + Args: + image_dir (str): Path or URL to the image zip file or directory. + annotation_dir (str): Path or URL to the annotation zip file or directory. + output_dir: Directory to store the organized dataset. Defaults to the + resolved Mobilint cache directory. + """ + output_dir = _resolve_organizer_output_dir(output_dir, "coco") + with TemporaryDirectory() as temp_dir: + local_image_dir, local_annotation_dir = _resolve_sources( + [image_dir, annotation_dir], temp_dir + ) + + if local_image_dir.endswith(".zip") and local_annotation_dir.endswith(".zip"): + print("Unpacking image and annotation files to temporary directory...") + _safe_unpack_archive(local_image_dir, temp_dir) + _safe_unpack_archive( + local_annotation_dir, os.path.join(temp_dir, "annotations_trainval2017") + ) + print("Unpacking completed") + construct_coco( + os.path.join(temp_dir, "val2017"), + os.path.join(temp_dir, "annotations_trainval2017"), + output_dir, + ) + return + + construct_coco(local_image_dir, local_annotation_dir, output_dir) + + +def construct_widerface(image_dir: str, annotation_dir: str, output_dir: str) -> None: + """Constructs the WiderFace dataset by copying images and annotations to a target directory. + + Args: + image_dir (str): Directory containing WiderFace images. + annotation_dir (str): Directory containing WiderFace annotations. + output_dir (str): Directory where the organized dataset will be stored. + """ + print( + f"Constructing WiderFace dataset from {image_dir} and {annotation_dir} to {output_dir}" + ) + output_dir = os.path.abspath(output_dir) + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".widerface-staging-" + ) as staging_dir: + staged_output_dir = os.path.join(staging_dir, "widerface") + shutil.copytree( + os.path.join(image_dir, "images"), + os.path.join(staged_output_dir, "images"), + ) + for file in os.listdir(annotation_dir): + if "_val" in file: + shutil.copy(os.path.join(annotation_dir, file), staged_output_dir) + _validate_staged_dataset(staged_output_dir, "widerface", ("face_detection",)) + _replace_staged_directories( + ((staged_output_dir, output_dir),), + output_parent_dir, + ".widerface-backup-", + ) + print("Constructing WiderFace dataset completed") + + +def organize_widerface( + image_dir: str, + annotation_dir: str, + output_dir: str | None = None, +) -> None: + """Organizes the WiderFace dataset, unpacking archives if necessary. + + Args: + image_dir (str): Path or URL to the image zip file or directory. + annotation_dir (str): Path or URL to the annotation zip file or directory. + output_dir: Directory to store the organized dataset. Defaults to the + resolved Mobilint cache directory. + """ + output_dir = _resolve_organizer_output_dir(output_dir, "widerface") + with TemporaryDirectory() as temp_dir: + local_image_dir, local_annotation_dir = _resolve_sources( + [image_dir, annotation_dir], temp_dir + ) + + if local_image_dir.endswith(".zip") and local_annotation_dir.endswith(".zip"): + print("Unpacking image and annotation files to temporary directory...") + _safe_unpack_archive(local_image_dir, temp_dir) + _safe_unpack_archive(local_annotation_dir, temp_dir) + print("Unpacking completed") + construct_widerface( + os.path.join(temp_dir, "WIDER_val"), + os.path.join(temp_dir, "wider_face_split"), + output_dir, + ) + return + + construct_widerface(local_image_dir, local_annotation_dir, output_dir) + + +def _resolve_nyu_depth_validation_dirs(dataset_dir: str) -> tuple[str, str, str]: + """Resolves NYU Depth validation image and depth directories. + + Args: + dataset_dir: Directory containing the NYU Depth root or its parent. + + Returns: + Paths to the selected dataset root, validation image directory, and + validation depth directory. + + Raises: + ValueError: If the expected NYU Depth layout is not present. + """ + + roots = (os.path.join(dataset_dir, "nyu-depth"), dataset_dir) + for root in roots: + candidates = ( + (os.path.join(root, "images", "val"), os.path.join(root, "depth", "val")), + (os.path.join(root, "val", "images"), os.path.join(root, "val", "depth")), + (os.path.join(root, "images"), os.path.join(root, "depth")), + ) + for image_dir, depth_dir in candidates: + if os.path.isdir(image_dir) and os.path.isdir(depth_dir): + return root, image_dir, depth_dir + raise ValueError( + f"NYU Depth dataset must contain matching images/ and depth/ directories: {dataset_dir}" + ) + + +def _validate_dense_source_file(source_path: str, dataset_root: Path) -> str: + """Resolve a non-symlink regular file contained by a dense dataset root. + + Args: + source_path: Candidate data or metadata file. + dataset_root: Resolved root of the extracted dataset. + + Returns: + Resolved source path safe to copy. + + Raises: + ValueError: If the source is a symlink, is not a regular file, cannot be + resolved, or escapes the dataset root. + """ + + source = Path(source_path) + if source.is_symlink(): + raise ValueError(f"Dense dataset source file must not be a symlink: {source}.") + try: + resolved_source = source.resolve(strict=True) + except OSError as exc: + raise ValueError( + f"Unable to resolve dense dataset source file {source}: {exc}." + ) from exc + if not resolved_source.is_file(): + raise ValueError(f"Dense dataset source must be a regular file: {source}.") + if not resolved_source.is_relative_to(dataset_root): + raise ValueError( + f"Dense dataset source must remain within dataset root: {source}." + ) + return str(resolved_source) + + +def _collect_unique_dense_sources( + source_paths: Iterable[str], + dataset_root: Path, + source_description: str, +) -> dict[str, str]: + """Return validated dense sources keyed by unique filename stem. + + Dense organizers flatten source files into a single output directory. Reject + repeated stems up front instead of silently retaining whichever recursive + traversal entry happened to be processed last. + """ + + sources: dict[str, str] = {} + for source_path in sorted(source_paths): + sample_id = Path(source_path).stem + validated_path = _validate_dense_source_file(source_path, dataset_root) + previous_path = sources.get(sample_id) + if previous_path is not None: + raise ValueError( + f"{source_description} contain duplicate filename stem {sample_id!r}: " + f"{previous_path} and {validated_path}." + ) + sources[sample_id] = validated_path + return sources + + +def _validate_dense_output_root( + output_dir: str, + dataset_name: str, + layout_names: Iterable[str], +) -> str: + """Reject symlinks in a dense managed root before organization. + + Args: + output_dir: Requested managed dataset root. + dataset_name: Human-readable dataset name for error reporting. + layout_names: Dataset-specific directories managed below the root. + + Returns: + Expanded absolute output path. + + Raises: + ValueError: If the managed root, an ancestor, or a managed layout + directory is a symlink. + """ + + requested_path = Path(output_dir).expanduser() + output_path = Path(os.path.abspath(requested_path)) + if _path_has_symlink_component(requested_path): + raise ValueError( + f"{dataset_name} output directory and its existing parents must not be symlinks: {output_path}. " + "Remove the symlink or choose a path beneath regular directories." + ) + for layout_name in layout_names: + layout_path = output_path / layout_name + if layout_path.is_symlink(): + raise ValueError( + f"{dataset_name} output layout directories must not be symlinks: {layout_path}. " + "Remove the symlink or choose a different output directory." + ) + return str(output_path) + + +def _collect_nyu_depth_validation_files( + image_dir: str, + depth_dir: str, + dataset_root: Path, +) -> tuple[dict[str, str], dict[str, str]]: + """Validates and returns matching NYU Depth validation image/depth pairs.""" + + images = _collect_unique_dense_sources( + _iter_files(image_dir, [".jpg", ".jpeg", ".png"]), + dataset_root, + "NYU Depth images", + ) + depths = _collect_unique_dense_sources( + _iter_files(depth_dir, [".npy"]), + dataset_root, + "NYU Depth depth maps", + ) + missing_depths = sorted(set(images) - set(depths)) + missing_images = sorted(set(depths) - set(images)) + if missing_depths or missing_images: + details = [] + if missing_depths: + details.append( + f"images without depth maps: {', '.join(missing_depths[:5])}" + ) + if missing_images: + details.append( + f"depth maps without images: {', '.join(missing_images[:5])}" + ) + raise ValueError( + f"NYU Depth validation image/depth mismatch ({'; '.join(details)})." + ) + if len(images) != NYU_DEPTH_VALIDATION_SAMPLE_COUNT: + raise ValueError( + "NYU Depth validation dataset must contain " + f"{NYU_DEPTH_VALIDATION_SAMPLE_COUNT} matching image/depth pairs, found {len(images)}." + ) + return images, depths + + +def construct_nyu_depth(dataset_dir: str, output_dir: str) -> None: + """Constructs the NYU Depth layout from an extracted dataset directory. + + Args: + dataset_dir: Directory containing the NYU Depth root or its parent. + output_dir: Directory where the organized dataset will be stored. + """ + + output_dir = _validate_dense_output_root( + output_dir, "NYU Depth", ("images", "depth") + ) + selected_root, image_dir, depth_dir = _resolve_nyu_depth_validation_dirs( + dataset_dir + ) + try: + dataset_root = Path(selected_root).resolve(strict=True) + except OSError as exc: + raise ValueError( + f"Unable to resolve NYU Depth dataset root {selected_root}: {exc}." + ) from exc + images, depths = _collect_nyu_depth_validation_files( + image_dir, depth_dir, dataset_root + ) + print( + f"Constructing NYU Depth validation dataset from {dataset_dir} to {output_dir}" + ) + + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".nyu-depth-staging-" + ) as staging_dir: + staged_image_dir = os.path.join(staging_dir, "images") + staged_depth_dir = os.path.join(staging_dir, "depth") + os.makedirs(staged_image_dir) + os.makedirs(staged_depth_dir) + for sample_id in sorted(images): + shutil.copy2( + images[sample_id], + os.path.join(staged_image_dir, os.path.basename(images[sample_id])), + ) + shutil.copy2( + depths[sample_id], + os.path.join(staged_depth_dir, os.path.basename(depths[sample_id])), + ) + + _validate_staged_nyu_depth(staging_dir) + + replacements = ( + (staged_image_dir, os.path.join(output_dir, "images")), + (staged_depth_dir, os.path.join(output_dir, "depth")), + ) + _replace_staged_directories( + replacements, output_parent_dir, ".nyu-depth-backup-" + ) + print( + f"Constructed NYU Depth validation dataset with {len(images)} image/depth pairs" + ) + + +def _validate_staged_nyu_depth(staging_dir: str) -> None: + """Decode staged NYU pairs before they can replace an existing cache.""" + + image_dir = Path(staging_dir) / "images" + depth_dir = Path(staging_dir) / "depth" + for image_path in sorted(image_dir.iterdir()): + if image_path.suffix.lower() not in {".jpg", ".jpeg", ".png"}: + continue + depth_path = depth_dir / f"{image_path.stem}.npy" + image = cv2.imread(str(image_path)) + if image is None: + raise ValueError(f"Staged NYU Depth image is unreadable: {image_path}.") + try: + raw_depth = np.load(depth_path, allow_pickle=False) + except (OSError, ValueError) as exc: + raise ValueError( + f"Unable to load staged NYU Depth target {depth_path}: {exc}." + ) from exc + if not np.issubdtype(raw_depth.dtype, np.number) or np.issubdtype( + raw_depth.dtype, np.complexfloating + ): + raise ValueError( + "Staged NYU Depth target must use a real numeric dtype, " + f"got {raw_depth.dtype}: {depth_path}." + ) + depth = np.asarray(raw_depth, dtype=np.float32) + if depth.ndim != 2 or depth.shape != image.shape[:2]: + raise ValueError( + "Staged NYU Depth image and target shapes must match: " + f"image {image.shape[:2]}, depth {depth.shape}: {image_path}." + ) + if not bool(np.isfinite(depth).all()): + raise ValueError( + f"Staged NYU Depth target must contain only finite values: {depth_path}." + ) + if bool((depth < 0).any()): + raise ValueError( + f"Staged NYU Depth target must not contain negative values: {depth_path}." + ) + if not bool(((depth > 0.001) & (depth < 100.0)).any()): + raise ValueError( + "Staged NYU Depth target must contain at least one valid metric depth " + f"in the (0.001, 100.0) range: {depth_path}." + ) + + +def organize_nyu_depth( + dataset_path: str = NYU_DEPTH_URL, + output_dir: str | None = None, +) -> None: + """Organizes NYU Depth, downloading and unpacking an archive when necessary. + + Args: + dataset_path: Path or URL to the NYU Depth zip file or extracted dataset directory. + output_dir: Directory to store the organized dataset. Defaults to the + resolved Mobilint cache directory. + """ + + output_dir = _resolve_organizer_output_dir(output_dir, "nyu-depth") + output_dir = _validate_dense_output_root( + output_dir, "NYU Depth", ("images", "depth") + ) + with TemporaryDirectory() as temp_dir: + local_dataset_path = _resolve_source(dataset_path, temp_dir) + if local_dataset_path.endswith(".zip"): + print("Unpacking NYU Depth files to temporary directory...") + _safe_unpack_archive(local_dataset_path, temp_dir) + print("Unpacking completed") + construct_nyu_depth(temp_dir, output_dir) + return + + construct_nyu_depth(local_dataset_path, output_dir) + + +def _resolve_ade20k_validation_dirs(dataset_dir: str) -> tuple[str, str, str]: + """Resolves the ADE20K root and validation image/mask directories.""" + + for root in (os.path.join(dataset_dir, "ADEChallengeData2016"), dataset_dir): + for image_dir, annotation_dir in ( + ( + os.path.join(root, "images", "validation"), + os.path.join(root, "annotations", "validation"), + ), + (os.path.join(root, "images"), os.path.join(root, "annotations")), + ): + if os.path.isdir(image_dir) and os.path.isdir(annotation_dir): + return root, image_dir, annotation_dir + raise ValueError( + f"ADE20K dataset must contain matching images/ and annotations/ directories: {dataset_dir}" + ) + + +def construct_ade20k(dataset_dir: str, output_dir: str) -> None: + """Constructs the flat ADE20K validation layout from an extracted dataset. + + Args: + dataset_dir: Directory containing the ADE20K root or its parent. + output_dir: Directory where the organized validation dataset will be stored. + + Raises: + ValueError: If the source does not contain 2,000 matched validation image/mask pairs. + """ + + output_dir = _validate_dense_output_root( + output_dir, "ADE20K", ("images", "annotations") + ) + dataset_root, image_dir, annotation_dir = _resolve_ade20k_validation_dirs( + dataset_dir + ) + try: + resolved_dataset_root = Path(dataset_root).resolve(strict=True) + except OSError as exc: + raise ValueError( + f"Unable to resolve ADE20K dataset root {dataset_root}: {exc}." + ) from exc + images = _collect_unique_dense_sources( + ( + os.path.join(image_dir, file_name) + for file_name in os.listdir(image_dir) + if file_name.startswith("ADE_val_") and file_name.lower().endswith(".jpg") + ), + resolved_dataset_root, + "ADE20K images", + ) + annotations = _collect_unique_dense_sources( + ( + os.path.join(annotation_dir, file_name) + for file_name in os.listdir(annotation_dir) + if file_name.startswith("ADE_val_") and file_name.lower().endswith(".png") + ), + resolved_dataset_root, + "ADE20K annotations", + ) + if set(images) != set(annotations): + raise ValueError( + "ADE20K validation images and annotations must have matching file stems." + ) + if len(images) != ADE20K_VALIDATION_SAMPLE_COUNT: + raise ValueError( + f"ADE20K validation dataset must contain {ADE20K_VALIDATION_SAMPLE_COUNT} pairs, found {len(images)}." + ) + metadata: dict[str, str] = {} + for file_name in ADE20K_METADATA_FILES: + metadata_path = os.path.join(dataset_root, file_name) + if not os.path.lexists(metadata_path): + raise ValueError( + f"ADE20K dataset is missing required metadata files: {file_name}." + ) + metadata[file_name] = _validate_dense_source_file( + metadata_path, resolved_dataset_root + ) + + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".ade20k-staging-" + ) as staging_dir: + staged_output_dir = os.path.join(staging_dir, "ade20k") + staged_image_dir = os.path.join(staged_output_dir, "images") + staged_annotation_dir = os.path.join(staged_output_dir, "annotations") + os.makedirs(staged_image_dir) + os.makedirs(staged_annotation_dir) + for sample_id in sorted(images): + shutil.copy2( + images[sample_id], + os.path.join(staged_image_dir, os.path.basename(images[sample_id])), + ) + shutil.copy2( + annotations[sample_id], + os.path.join( + staged_annotation_dir, os.path.basename(annotations[sample_id]) + ), + ) + for file_name in ADE20K_METADATA_FILES: + shutil.copy2( + metadata[file_name], + os.path.join(staged_output_dir, file_name), + ) + + _validate_staged_dataset( + staged_output_dir, "ade20k", ("semantic_segmentation",) + ) + _replace_staged_directories( + ((staged_output_dir, output_dir),), + output_parent_dir, + ".ade20k-backup-", + ) + print(f"Constructed ADE20K validation dataset with {len(images)} image/mask pairs") + + +def organize_ade20k( + dataset_path: str = ADE20K_URL, + output_dir: str | None = None, +) -> None: + """Organizes ADE20K validation data, downloading and unpacking when necessary.""" + + output_dir = _resolve_organizer_output_dir(output_dir, "ade20k") + output_dir = _validate_dense_output_root( + output_dir, "ADE20K", ("images", "annotations") + ) + with TemporaryDirectory() as temp_dir: + local_dataset_path = _resolve_source(dataset_path, temp_dir) + if local_dataset_path.endswith(".zip"): + _safe_unpack_archive(local_dataset_path, temp_dir) + construct_ade20k(temp_dir, output_dir) + return + construct_ade20k(local_dataset_path, output_dir) + + +def _validate_cityscapes_zip(archive_path: str, source_name: str) -> str: + """Validate one official Cityscapes ZIP source. + + Args: + archive_path: Path to the raw Cityscapes archive. + source_name: Human-readable source description for errors. + + Returns: + Expanded absolute archive path. + + Raises: + ValueError: If the path is missing, is not a file, is not a ZIP, or contains duplicate members. + """ + + resolved_path = os.path.abspath(os.path.expanduser(archive_path)) + if not os.path.isfile(resolved_path): + raise ValueError( + f"Cityscapes {source_name} archive does not exist or is not a file: {resolved_path}." + ) + if not zipfile.is_zipfile(resolved_path): + raise ValueError( + f"Cityscapes {source_name} source must be a valid ZIP archive: {resolved_path}." + ) + + with zipfile.ZipFile(resolved_path) as archive: + seen_members: set[str] = set() + duplicate_members: set[str] = set() + for member in archive.infolist(): + if member.filename in seen_members: + duplicate_members.add(member.filename) + seen_members.add(member.filename) + if duplicate_members: + raise ValueError( + f"Cityscapes {source_name} archive contains duplicate members: {', '.join(sorted(duplicate_members)[:3])}." + ) + return resolved_path + + +def _collect_cityscapes_validation_files( + split_dir: str, + suffix: str, + source_name: str, +) -> dict[str, str]: + """Collect official Cityscapes validation files keyed by shared sample ID. + + Args: + split_dir: Extracted ``leftImg8bit/val`` or ``gtFine/val`` directory. + suffix: Required official file suffix. + source_name: Human-readable source description for errors. + + Returns: + Mapping from ``__`` to source path. + + Raises: + ValueError: If a candidate filename is malformed, misplaced, or duplicates an ID. + """ + + files: dict[str, str] = {} + if not os.path.isdir(split_dir): + return files + + for current_root, _, file_names in os.walk(split_dir): + relative_root = os.path.relpath(current_root, split_dir) + for file_name in file_names: + if not file_name.endswith(suffix): + continue + if relative_root == "." or os.sep in relative_root: + raise ValueError( + f"Malformed Cityscapes {source_name} path: " + f"{os.path.relpath(os.path.join(current_root, file_name), split_dir)}." + ) + sample_id = file_name.removesuffix(suffix) + match = CITYSCAPES_SAMPLE_ID_PATTERN.fullmatch(sample_id) + if match is None or match.group("city") != relative_root: + raise ValueError( + f"Malformed Cityscapes {source_name} filename: {file_name}." + ) + if sample_id in files: + raise ValueError( + f"Duplicate Cityscapes {source_name} sample ID: {sample_id}." + ) + files[sample_id] = os.path.join(current_root, file_name) + return files + + +def organize_cityscapes( + image_dir: str, + annotation_dir: str, + output_dir: str | None = None, +) -> None: + """Install official Cityscapes validation archives as lossless flat PNG pairs. + + Only validation RGB images and ``gtFine_labelIds`` masks are selected. + Training, test, and auxiliary annotation files remain excluded. + + Args: + image_dir: Path to ``leftImg8bit_trainvaltest.zip``. + annotation_dir: Path to ``gtFine_trainvaltest.zip``. + output_dir: Directory where the organized validation dataset is stored. + Defaults to the resolved Mobilint cache directory. + + Raises: + ValueError: If either source is invalid or does not contain exactly 500 matching pairs. + OSError: If extraction, copying, or atomic installation fails. + """ + + output_dir = _resolve_organizer_output_dir(output_dir, "cityscapes") + output_dir = _validate_dense_output_root( + output_dir, "Cityscapes", ("images", "annotations") + ) + image_archive = _validate_cityscapes_zip(image_dir, "image") + annotation_archive = _validate_cityscapes_zip(annotation_dir, "annotation") + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".cityscapes-staging-" + ) as staging_dir: + extracted_image_dir = os.path.join(staging_dir, "raw-images") + extracted_annotation_dir = os.path.join(staging_dir, "raw-annotations") + _safe_unpack_archive(image_archive, extracted_image_dir) + _safe_unpack_archive(annotation_archive, extracted_annotation_dir) + images = _collect_cityscapes_validation_files( + os.path.join(extracted_image_dir, "leftImg8bit", "val"), + CITYSCAPES_IMAGE_SUFFIX, + "image", + ) + annotations = _collect_cityscapes_validation_files( + os.path.join(extracted_annotation_dir, "gtFine", "val"), + CITYSCAPES_ANNOTATION_SUFFIX, + "annotation", + ) + missing_annotations = sorted(images.keys() - annotations.keys()) + missing_images = sorted(annotations.keys() - images.keys()) + if missing_annotations or missing_images: + details = [] + if missing_annotations: + details.append( + f"missing annotations for {', '.join(missing_annotations[:3])}" + ) + if missing_images: + details.append(f"missing images for {', '.join(missing_images[:3])}") + raise ValueError( + f"Cityscapes validation image/annotation mismatch ({'; '.join(details)})." + ) + if len(images) != CITYSCAPES_VALIDATION_SAMPLE_COUNT: + raise ValueError( + "Cityscapes validation archives must contain " + f"{CITYSCAPES_VALIDATION_SAMPLE_COUNT} pairs, found {len(images)}." + ) + + staged_image_dir = os.path.join(staging_dir, "images") + staged_annotation_dir = os.path.join(staging_dir, "annotations") + os.makedirs(staged_image_dir) + os.makedirs(staged_annotation_dir) + for sample_id in sorted(images): + shutil.copy2( + images[sample_id], os.path.join(staged_image_dir, f"{sample_id}.png") + ) + shutil.copy2( + annotations[sample_id], + os.path.join(staged_annotation_dir, f"{sample_id}.png"), + ) + + if not dataset_ready(staging_dir, "semantic_segmentation", "cityscapes"): + raise ValueError( + "Staged Cityscapes validation data failed identity and completeness checks." + ) + _validate_staged_payloads(Path(staging_dir), "cityscapes") + + replacements = ( + (staged_image_dir, os.path.join(output_dir, "images")), + (staged_annotation_dir, os.path.join(output_dir, "annotations")), + ) + os.makedirs(output_dir, exist_ok=True) + _replace_staged_directories( + replacements, output_parent_dir, ".cityscapes-backup-" + ) + print( + f"Constructed Cityscapes validation dataset with {len(images)} image/mask pairs" + ) + + +def _resolve_dotav1_root(dataset_dir: str) -> str: + """Resolves a DOTAv1 dataset root from a directory path. + + Args: + dataset_dir: Directory containing the DOTAv1 dataset or its parent. + + Returns: + Path to the DOTAv1 dataset root. + """ + dotav1_dir = os.path.join(dataset_dir, "DOTAv1") + if os.path.isdir(dotav1_dir): + return dotav1_dir + return dataset_dir + + +def _is_google_drive_folder_url(path_or_url: str) -> bool: + """Returns whether a URL points to a Google Drive folder.""" + + parsed = urlparse(path_or_url) + return parsed.hostname == "drive.google.com" and bool( + re.fullmatch(r"/drive(?:/u/[^/]+)?/folders/[^/]+/?", parsed.path) + ) + + +def _download_dotav1_google_drive_archives( + folder_url: str, download_dir: str +) -> tuple[str, str]: + """Downloads the DOTAv1 image and v1.0-label archives from a Google Drive folder. + + Args: + folder_url: Public Google Drive folder URL containing the DOTAv1 archives. + download_dir: Directory where the selected archives will be stored. + + Returns: + Paths to the image archive and original v1.0-label archive. + + Raises: + ValueError: If the required archives are absent from the Drive folder. + RuntimeError: If gdown fails to download a required archive. + """ + + print(f"Retrieving DOTAv1 archive list from {folder_url}...") + folder_entries = download_folder( + url=folder_url, output=download_dir, quiet=True, skip_download=True + ) + if folder_entries is None: + raise RuntimeError( + f"Failed to retrieve the DOTAv1 Google Drive folder listing: {folder_url}" + ) + files = [ + entry for entry in folder_entries if _is_google_drive_download_entry(entry) + ] + archives: dict[str, _GoogleDriveDownloadEntry] = {} + for archive_path in DOTAV1_GOOGLE_DRIVE_ARCHIVES: + matches = [ + drive_file + for drive_file in files + if drive_file.path == archive_path + or drive_file.path.endswith(f"/{archive_path}") + ] + if len(matches) == 1: + archives[archive_path] = matches[0] + continue + + available = ", ".join(sorted(drive_file.path for drive_file in files)) or "none" + if not matches: + raise ValueError( + f"DOTAv1 Drive folder is missing {archive_path}. Available files: {available}." + ) + ambiguous = ", ".join(sorted(drive_file.path for drive_file in matches)) + raise ValueError( + f"DOTAv1 Drive folder has ambiguous matches for {archive_path}: {ambiguous}." + ) + + local_archives: dict[str, str] = {} + for archive_path in sorted(DOTAV1_GOOGLE_DRIVE_ARCHIVES): + drive_file = archives[archive_path] + local_path = os.path.join(download_dir, os.path.basename(archive_path)) + print(f"Downloading DOTAv1 {archive_path}...") + downloaded_path = download( + id=drive_file.id, output=local_path, quiet=False, resume=True + ) + if not isinstance(downloaded_path, str): + raise RuntimeError( + f"Failed to download DOTAv1 archive {archive_path} from {folder_url}." + ) + local_archives[archive_path] = downloaded_path + + return ( + local_archives[DOTAV1_DOWNLOAD_CONFIG["images_archive"]], + local_archives[DOTAV1_DOWNLOAD_CONFIG["labels_archive"]], + ) + + +def _iter_files(root: str, extensions: Iterable[str]) -> Iterable[str]: + """Yields files below a directory with one of the requested suffixes.""" + + suffixes = tuple(extension.lower() for extension in extensions) + for current_root, _, file_names in os.walk(root): + for file_name in file_names: + if file_name.lower().endswith(suffixes): + yield os.path.join(current_root, file_name) + + +def _safe_archive_member_path(member_name: str, destination: str) -> str: + """Return an archive member destination after enforcing staging-directory containment. + + Args: + member_name: Path stored in an archive member. + destination: Archive extraction directory. + + Returns: + Absolute destination path for the member. + + Raises: + ValueError: If the member path is absolute or escapes the extraction directory. + """ + + root = os.path.abspath(destination) + target = os.path.abspath(os.path.join(root, member_name)) + if os.path.commonpath((root, target)) != root: + raise ValueError(f"Unsafe archive member path: {member_name!r}.") + return target + + +def _safe_unpack_archive(archive_path: str, destination: str) -> None: + """Extract an archive while rejecting links, special files, and escaping paths. + + Args: + archive_path: ZIP or tar-family dataset archive. + destination: Empty staging directory where archive members are written. + + Raises: + ValueError: If the archive format or any member is unsafe or unsupported. + OSError: If a validated archive cannot be read or written. + """ + + if zipfile.is_zipfile(archive_path): + with zipfile.ZipFile(archive_path) as archive: + members = archive.infolist() + targets: set[str] = set() + for member in members: + target = _safe_archive_member_path(member.filename, destination) + if target in targets: + raise ValueError( + f"Duplicate archive member path: {member.filename!r}." + ) + targets.add(target) + file_type = stat.S_IFMT(member.external_attr >> 16) + if file_type and not ( + stat.S_ISREG(file_type) or stat.S_ISDIR(file_type) + ): + raise ValueError( + f"Unsafe archive member type: {member.filename!r}." + ) + for member in members: + target = _safe_archive_member_path(member.filename, destination) + if member.is_dir(): + os.makedirs(target, exist_ok=True) + continue + os.makedirs(os.path.dirname(target), exist_ok=True) + with archive.open(member) as source, open(target, "wb") as output_file: + shutil.copyfileobj(source, output_file) + return + + if tarfile.is_tarfile(archive_path): + with tarfile.open(archive_path) as archive: + members = archive.getmembers() + targets: set[str] = set() + for member in members: + target = _safe_archive_member_path(member.name, destination) + if target in targets: + raise ValueError(f"Duplicate archive member path: {member.name!r}.") + targets.add(target) + if not (member.isfile() or member.isdir()): + raise ValueError(f"Unsafe archive member type: {member.name!r}.") + for member in members: + target = _safe_archive_member_path(member.name, destination) + if member.isdir(): + os.makedirs(target, exist_ok=True) + continue + source = archive.extractfile(member) + if source is None: + raise ValueError(f"Unable to read archive member: {member.name!r}.") + os.makedirs(os.path.dirname(target), exist_ok=True) + with source, open(target, "wb") as output_file: + shutil.copyfileobj(source, output_file) + return + + raise ValueError(f"Unsupported archive format: {archive_path}.") + + +def _write_dotav1_yolo_labels( + image_path: str, original_label_path: str, output_path: str +) -> None: + """Converts one official DOTAv1 label file into normalized OBB label format.""" + + with Image.open(image_path) as image: + width, height = image.size + converted_lines: list[str] = [] + seen_targets: set[tuple[str, tuple[float, ...]]] = set() + with open(original_label_path, encoding="utf-8") as label_file: + for line_number, line in enumerate(label_file, start=1): + fields = line.split() + if fields and ( + fields[0].startswith("imagesource:") or fields[0].startswith("gsd:") + ): + continue + if len(fields) < 10: + raise ValueError( + "Malformed DOTAv1 annotation in " + f"{original_label_path} at line {line_number}: expected at least " + f"10 fields, got {len(fields)}." + ) + class_name = fields[8] + if class_name not in DOTAV1_CLASS_TO_IDX: + raise ValueError( + f"Unsupported DOTAv1 class in {original_label_path}: {class_name}" + ) + coordinates = [float(value) for value in fields[:8]] + if not all(math.isfinite(coordinate) for coordinate in coordinates): + raise ValueError( + f"DOTAv1 coordinates must be finite in {original_label_path} " + f"at line {line_number}." + ) + if not _polygon_has_positive_image_overlap(coordinates, (height, width)): + raise ValueError( + "DOTAv1 polygon must overlap its source image in " + f"{original_label_path} at line {line_number}." + ) + _validate_dotav1_polygon_vertices( + coordinates, original_label_path, line_number + ) + if fields[9] not in {"0", "1", "2"}: + raise ValueError( + f"Unsupported DOTAv1 difficulty flag {fields[9]!r} in " + f"{original_label_path} at line {line_number}." + ) + target_key = (class_name, _canonicalize_quadrilateral(coordinates)) + if target_key in seen_targets: + raise ValueError( + "Duplicate DOTAv1 annotation target in " + f"{original_label_path} at line {line_number}." + ) + seen_targets.add(target_key) + normalized = [ + coordinate / (width if index % 2 == 0 else height) + for index, coordinate in enumerate(coordinates) + ] + converted_lines.append( + f"{DOTAV1_CLASS_TO_IDX[class_name]} " + + " ".join(f"{coordinate:.8g}" for coordinate in normalized) + # The trailing flag is normalized-label metadata, not a YOLO OBB + # coordinate. It preserves official difficult regions for evaluation. + + f" {int(fields[9] in {'1', '2'})}" + ) + with open(output_path, "w", encoding="utf-8") as output_file: + output_file.write("\n".join(converted_lines)) + if converted_lines: + output_file.write("\n") + + +def construct_dotav1_from_archives( + image_archive: str, label_archive: str, output_dir: str +) -> None: + """Constructs the DOTAv1 validation layout from the Google Drive archives. + + Args: + image_archive: Path to the DOTAv1 validation-image archive. + label_archive: Path to the original DOTAv1 v1.0 label archive. + output_dir: Directory where the organized validation dataset will be stored. + + Raises: + ValueError: If the archives have no validation files or their image and label stems differ. + OSError: If staging or replacing the organized dataset files fails. + """ + + with TemporaryDirectory() as extract_dir: + image_dir = os.path.join(extract_dir, "images") + label_dir = os.path.join(extract_dir, "labels") + _safe_unpack_archive(image_archive, image_dir) + _safe_unpack_archive(label_archive, label_dir) + + label_paths = list(_iter_files(label_dir, [".txt"])) + labels = { + os.path.splitext(os.path.basename(path))[0]: path for path in label_paths + } + if len(labels) != len(label_paths): + raise ValueError("DOTAv1 archive contains duplicate label stems.") + if not labels: + raise ValueError(f"No DOTAv1 label files found in {label_archive}.") + + image_paths = list( + _iter_files(image_dir, [".bmp", ".jpg", ".jpeg", ".png", ".tif", ".tiff"]) + ) + images = { + os.path.splitext(os.path.basename(path))[0]: path for path in image_paths + } + if len(images) != len(image_paths): + raise ValueError("DOTAv1 archive contains duplicate image stems.") + image_ids = set(images) + label_ids = set(labels) + missing_labels = sorted(image_ids - label_ids) + missing_images = sorted(label_ids - image_ids) + if missing_labels or missing_images: + details = [] + if missing_labels: + details.append( + f"images without labels: {', '.join(missing_labels[:5])}" + ) + if missing_images: + details.append( + f"labels without images: {', '.join(missing_images[:5])}" + ) + raise ValueError(f"DOTAv1 archive stem mismatch ({'; '.join(details)}).") + matching_ids = sorted(image_ids) + if len(matching_ids) != DOTAV1_VALIDATION_SAMPLE_COUNT: + raise ValueError( + "DOTAv1 validation dataset must contain " + f"{DOTAV1_VALIDATION_SAMPLE_COUNT} matching image/label pairs, found {len(matching_ids)}." + ) + + output_dir = os.path.abspath(output_dir) + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".dotav1-staging-" + ) as staging_dir: + staged_output_dir = os.path.join(staging_dir, "dotav1") + staged_image_dir = os.path.join(staged_output_dir, "images") + staged_label_dir = os.path.join(staged_output_dir, "labels", "val") + staged_original_label_dir = os.path.join( + staged_output_dir, "labels", "val_original" + ) + os.makedirs(staged_image_dir) + os.makedirs(staged_label_dir) + os.makedirs(staged_original_label_dir) + + for image_id in matching_ids: + image_path = images[image_id] + shutil.copy2( + image_path, + os.path.join(staged_image_dir, os.path.basename(image_path)), + ) + + for image_id in matching_ids: + shutil.copy2( + labels[image_id], + os.path.join(staged_original_label_dir, f"{image_id}.txt"), + ) + _write_dotav1_yolo_labels( + images[image_id], + labels[image_id], + os.path.join(staged_label_dir, f"{image_id}.txt"), + ) + + _validate_staged_dataset(staged_output_dir, "dotav1", ("obb",)) + _replace_staged_directories( + ((staged_output_dir, output_dir),), + output_parent_dir, + ".dotav1-backup-", + ) + + print(f"Constructed DOTAv1 validation dataset with {len(matching_ids)} images") + + +def _copy_dotav1_layout_to_staging(dataset_root: str, staged_output_dir: str) -> None: + """Copy a flat or legacy DOTAv1 validation layout into canonical staging.""" + + image_root = os.path.join(dataset_root, "images") + supported_image_suffixes = (".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff") + flat_image_files = ( + [ + file_name + for file_name in os.listdir(image_root) + if os.path.isfile(os.path.join(image_root, file_name)) + and file_name.lower().endswith(supported_image_suffixes) + ] + if os.path.isdir(image_root) + else [] + ) + source_image_dir = ( + image_root if flat_image_files else os.path.join(image_root, "val") + ) + if not os.path.isdir(source_image_dir): + raise ValueError(f"No DOTAv1 validation images found in {dataset_root}") + + staged_image_dir = os.path.join(staged_output_dir, "images") + os.makedirs(staged_image_dir) + for file_name in os.listdir(source_image_dir): + source_path = os.path.join(source_image_dir, file_name) + if os.path.isfile(source_path) and file_name.lower().endswith( + supported_image_suffixes + ): + shutil.copy2(source_path, os.path.join(staged_image_dir, file_name)) + + for label_directory in ("val", "val_original"): + source_label_dir = os.path.join(dataset_root, "labels", label_directory) + if not os.path.isdir(source_label_dir): + continue + staged_label_dir = os.path.join(staged_output_dir, "labels", label_directory) + os.makedirs(staged_label_dir, exist_ok=True) + for file_name in os.listdir(source_label_dir): + source_path = os.path.join(source_label_dir, file_name) + if os.path.isfile(source_path) and file_name.lower().endswith(".txt"): + shutil.copy2(source_path, os.path.join(staged_label_dir, file_name)) + + +def construct_dotav1(dataset_dir: str, output_dir: str) -> None: + """Constructs a validation-only DOTAv1 dataset. + + Args: + dataset_dir: Directory containing a DOTAv1 dataset or its parent. + output_dir: Directory where the organized validation dataset will be stored. + + Raises: + ValueError: If the staged validation dataset is incomplete or mismatched. + OSError: If staging or replacing the organized dataset files fails. + """ + dataset_root = _resolve_dotav1_root(dataset_dir) + print(f"Constructing DOTAv1 validation dataset from {dataset_root} to {output_dir}") + output_dir = os.path.abspath(output_dir) + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + with TemporaryDirectory( + dir=output_parent_dir, prefix=".dotav1-staging-" + ) as staging_dir: + staged_output_dir = os.path.join(staging_dir, "dotav1") + os.makedirs(staged_output_dir) + _copy_dotav1_layout_to_staging(dataset_root, staged_output_dir) + _validate_staged_dataset(staged_output_dir, "dotav1", ("obb",)) + _replace_staged_directories( + ((staged_output_dir, output_dir),), + output_parent_dir, + ".dotav1-backup-", + ) + + print("Constructing DOTAv1 validation dataset completed") + + +def organize_dotav1( + dataset_path: str, + output_dir: str | None = None, +) -> None: + """Organizes a validation-only DOTAv1 dataset. + + Args: + dataset_path: Path or URL to the DOTAv1 zip file or extracted dataset directory. + output_dir: Directory to store the organized dataset. Defaults to the + resolved Mobilint cache directory. + """ + output_dir = _resolve_organizer_output_dir(output_dir, "dotav1") + with TemporaryDirectory() as temp_dir: + if _is_google_drive_folder_url(dataset_path): + image_archive, label_archive = _download_dotav1_google_drive_archives( + dataset_path, temp_dir + ) + construct_dotav1_from_archives(image_archive, label_archive, output_dir) + return + + local_dataset_path = _resolve_source(dataset_path, temp_dir) + + if local_dataset_path.endswith(".zip"): + print("Unpacking DOTAv1 files to temporary directory...") + _safe_unpack_archive(local_dataset_path, temp_dir) + print("Unpacking completed") + construct_dotav1(temp_dir, output_dir) + return + + construct_dotav1(local_dataset_path, output_dir) diff --git a/mblt_vision/utils/datasets/readiness.py b/mblt_vision/utils/datasets/readiness.py new file mode 100644 index 0000000..3aeaa5d --- /dev/null +++ b/mblt_vision/utils/datasets/readiness.py @@ -0,0 +1,1063 @@ +"""Identity and completeness checks for organized vision validation datasets.""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import re +from importlib.resources import files +from pathlib import Path +from typing import Any + +import numpy as np +from faster_coco_eval import mask as coco_mask +from PIL import Image +from scipy.io import loadmat +from scipy.io.matlab import MatReadError + +from ..._tasks import normalize_vision_task +from ...datasets import get_dataset_category_ids +from .cityscapes import CITYSCAPES_SOURCE_TO_TRAIN_ID + +IMAGE_SUFFIXES = {".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp"} +IMAGENET_CLASS_COUNT = 1000 +IMAGENET_IMAGES_PER_CLASS = 50 +COCO_VALIDATION_SAMPLE_COUNT = 5000 +DOTAV1_VALIDATION_SAMPLE_COUNT = 458 +WIDERFACE_EVENT_COUNT = 61 +WIDERFACE_VALIDATION_SAMPLE_COUNT = 3226 +NYU_DEPTH_VALIDATION_SAMPLE_COUNT = 654 +ADE20K_VALIDATION_SAMPLE_COUNT = 2000 +CITYSCAPES_VALIDATION_SAMPLE_COUNT = 500 +ADE20K_METADATA_FILES = ("objectInfo150.txt", "sceneCategories.txt") +IMAGENET_CLASS_PATTERN = re.compile(r"n\d{8}") +IMAGENET_IMAGE_PATTERN = re.compile(r"ILSVRC2012_val_\d{8}") +COCO_IMAGE_PATTERN = re.compile(r"\d{12}") +WIDERFACE_EVENT_PATTERN = re.compile(r"\d+--\S.*") +CITYSCAPES_SAMPLE_ID_PATTERN = re.compile( + r"^(?P[A-Za-z][A-Za-z0-9-]*)_\d{6}_\d{6}$" +) +CITYSCAPES_VALIDATION_CITY_COUNTS = {"frankfurt": 267, "lindau": 59, "munster": 174} +IMAGENET_SYNSET_ORDER = tuple( + files("mblt_vision.datasets") + .joinpath("imagenet_synsets.txt") + .read_text(encoding="utf-8") + .splitlines() +) +IMAGENET_SYNSETS = frozenset(IMAGENET_SYNSET_ORDER) +COCO_ANNOTATION_COUNTS = { + "instances_val2017.json": 36781, + "person_keypoints_val2017.json": 11004, +} +COCO_CATEGORY_COUNTS = { + "instances_val2017.json": 80, + "person_keypoints_val2017.json": 1, +} +COCO_CATEGORY_IDS = frozenset(get_dataset_category_ids("coco")) +COCO_PERSON_KEYPOINT_CATEGORY_IDS = frozenset({1}) +COCO_VALIDATION_IMAGE_IDENTITIES_SHA256 = ( + "f57f71ba25171a0fd99be8c425a91d4a6fdd43d25aadc7e5be51dd37a73281a7" +) + + +def _path_has_symlink_component(path: Path) -> bool: + """Return whether a path traversal or its normalized ancestors contain a symlink.""" + + expanded_path = path.expanduser() + traversal_path = ( + expanded_path if expanded_path.is_absolute() else Path.cwd() / expanded_path + ) + normalized_path = Path(os.path.abspath(expanded_path)) + candidates = ( + traversal_path, + *traversal_path.parents, + normalized_path, + *normalized_path.parents, + ) + return any(component.is_symlink() for component in candidates) + + +def _files_by_stem( + directory: Path, + suffixes: set[str], + *, + reject_symlinks: bool = False, +) -> dict[str, Path] | None: + """Collect direct child files with supported suffixes by stem. + + Args: + directory: Directory containing candidate files. + suffixes: Accepted lowercase file suffixes. + reject_symlinks: Whether any symlinked directory or entry invalidates + the file collection. + + Returns: + Files keyed by stem, an empty mapping for a missing directory, or + ``None`` for duplicate stems or rejected symlinks. + """ + + if (reject_symlinks and directory.is_symlink()) or not directory.is_dir(): + return {} + entries = list(directory.iterdir()) + if reject_symlinks and any(path.is_symlink() for path in entries): + return None + paths = [ + path for path in entries if path.is_file() and path.suffix.lower() in suffixes + ] + files = {path.stem: path for path in paths} + return files if len(files) == len(paths) else None + + +def _has_positive_polygon_area(polygon: list[int | float]) -> bool: + """Return whether a finite flat polygon encloses non-zero signed area.""" + + points = np.asarray(polygon, dtype=np.float64).reshape(-1, 2) + signed_double_area = np.dot(points[:, 0], np.roll(points[:, 1], -1)) - np.dot( + points[:, 1], np.roll(points[:, 0], -1) + ) + return bool(abs(signed_double_area) > 0) + + +def _canonicalize_quadrilateral( + coordinates: list[int | float] | tuple[int | float, ...], +) -> tuple[float, ...]: + """Return a quadrilateral key independent of its start vertex and winding.""" + + if len(coordinates) != 8: + raise ValueError( + "A quadrilateral must contain exactly four two-dimensional vertices." + ) + points = tuple( + (float(coordinates[index]), float(coordinates[index + 1])) + for index in range(0, len(coordinates), 2) + ) + candidates = [] + for winding in (points, tuple(reversed(points))): + candidates.extend(winding[index:] + winding[:index] for index in range(4)) + return tuple(coordinate for point in min(candidates) for coordinate in point) + + +def _polygon_has_positive_image_overlap( + polygon: list[int | float], image_shape: tuple[int, int] | None +) -> bool: + """Return whether a polygon covers non-zero area within an image rectangle.""" + + if image_shape is None: + return True + height, width = image_shape + if height <= 0 or width <= 0: + return False + points = [ + tuple(point) for point in np.asarray(polygon, dtype=np.float64).reshape(-1, 2) + ] + + def _clip( + vertices: list[tuple[float, float]], + inside: Any, + intersect: Any, + ) -> list[tuple[float, float]]: + clipped: list[tuple[float, float]] = [] + if not vertices: + return clipped + previous = vertices[-1] + previous_inside = bool(inside(previous)) + for current in vertices: + current_inside = bool(inside(current)) + if current_inside != previous_inside: + clipped.append(intersect(previous, current)) + if current_inside: + clipped.append(current) + previous = current + previous_inside = current_inside + return clipped + + def _vertical_intersection( + x: float, start: tuple[float, float], end: tuple[float, float] + ) -> tuple[float, float]: + delta_x = end[0] - start[0] + if delta_x == 0: + return x, start[1] + ratio = (x - start[0]) / delta_x + return x, start[1] + ratio * (end[1] - start[1]) + + def _horizontal_intersection( + y: float, start: tuple[float, float], end: tuple[float, float] + ) -> tuple[float, float]: + delta_y = end[1] - start[1] + if delta_y == 0: + return start[0], y + ratio = (y - start[1]) / delta_y + return start[0] + ratio * (end[0] - start[0]), y + + points = _clip( + points, + lambda point: point[0] >= 0, + lambda a, b: _vertical_intersection(0, a, b), + ) + points = _clip( + points, + lambda point: point[0] <= width, + lambda a, b: _vertical_intersection(width, a, b), + ) + points = _clip( + points, + lambda point: point[1] >= 0, + lambda a, b: _horizontal_intersection(0, a, b), + ) + points = _clip( + points, + lambda point: point[1] <= height, + lambda a, b: _horizontal_intersection(height, a, b), + ) + if len(points) < 3: + return False + clipped = [coordinate for point in points for coordinate in point] + return _has_positive_polygon_area(clipped) + + +def _imagenet_ready(root: Path) -> bool: + """Check the organizer's complete ImageNet-1k validation class tree.""" + + if not root.is_dir(): + return False + class_dirs = [path for path in root.iterdir() if path.is_dir()] + if {path.name for path in class_dirs} != IMAGENET_SYNSETS: + return False + image_names: set[str] = set() + for class_dir in class_dirs: + if IMAGENET_CLASS_PATTERN.fullmatch(class_dir.name) is None: + return False + images = [ + path + for path in class_dir.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ] + if len(images) != IMAGENET_IMAGES_PER_CLASS or any( + path.suffix != ".JPEG" + or IMAGENET_IMAGE_PATTERN.fullmatch(path.stem) is None + for path in images + ): + return False + image_names.update(path.name for path in images) + return len(image_names) == IMAGENET_CLASS_COUNT * IMAGENET_IMAGES_PER_CLASS + + +def _load_coco_image_names( + annotation_path: Path, task: str = "object_detection" +) -> set[str] | None: + """Load unique validation image filenames from a COCO annotation file.""" + + try: + annotation: Any = json.loads(annotation_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError, UnicodeError): + return None + if not isinstance(annotation, dict): + return None + image_records = annotation.get("images") + if ( + not isinstance(image_records, list) + or len(image_records) != COCO_VALIDATION_SAMPLE_COUNT + ): + return None + names: list[str] = [] + image_ids: list[int] = [] + image_shapes: dict[int, tuple[int, int] | None] = {} + for record in image_records: + if not isinstance(record, dict): + continue + file_name = record.get("file_name") + image_id = record.get("id") + if ( + isinstance(file_name, str) + and isinstance(image_id, int) + and not isinstance(image_id, bool) + ): + names.append(file_name) + image_ids.append(image_id) + height, width = record.get("height"), record.get("width") + if ( + isinstance(height, int) + and not isinstance(height, bool) + and height > 0 + and isinstance(width, int) + and not isinstance(width, bool) + and width > 0 + ): + image_shapes[image_id] = (height, width) + else: + return None + if ( + len(names) != len(image_records) + or len(names) != len(set(names)) + or len(image_ids) != len(set(image_ids)) + ): + return None + if len(image_records) == 5000: + identity_payload = "".join( + f"{image_id}:{file_name}\n" + for image_id, file_name in sorted(zip(image_ids, names, strict=True)) + ).encode() + if ( + hashlib.sha256(identity_payload).hexdigest() + != COCO_VALIDATION_IMAGE_IDENTITIES_SHA256 + ): + return None + annotation_records = annotation.get("annotations") + categories = annotation.get("categories") + expected_annotations = COCO_ANNOTATION_COUNTS.get(annotation_path.name) + expected_categories = COCO_CATEGORY_COUNTS.get(annotation_path.name) + if ( + not isinstance(annotation_records, list) + or not isinstance(categories, list) + or len(annotation_records) != expected_annotations + or len(categories) != expected_categories + ): + return None + category_ids: list[int] = [] + for category in categories: + if not isinstance(category, dict): + return None + category_id = category.get("id") + if not isinstance(category_id, int) or isinstance(category_id, bool): + return None + category_ids.append(category_id) + if len(category_ids) != len(set(category_ids)): + return None + expected_category_ids = ( + COCO_PERSON_KEYPOINT_CATEGORY_IDS + if annotation_path.name == "person_keypoints_val2017.json" + else COCO_CATEGORY_IDS + ) + if set(category_ids) != expected_category_ids: + return None + if not _coco_task_annotations_valid( + annotation_records, + image_ids=set(image_ids), + category_ids=set(category_ids), + image_shapes=image_shapes, + task=task, + ): + return None + return set(names) + + +def _coco_task_annotations_valid( + annotation_records: list[Any], + *, + image_ids: set[int], + category_ids: set[int], + image_shapes: dict[int, tuple[int, int] | None], + task: str, +) -> bool: + """Validate task-specific COCO annotation payloads for readiness and APIs.""" + + annotation_ids: list[int] = [] + for record in annotation_records: + if not isinstance(record, dict): + return False + annotation_id = record.get("id") + image_id = record.get("image_id") + category_id = record.get("category_id") + if ( + not isinstance(annotation_id, int) + or isinstance(annotation_id, bool) + or not isinstance(image_id, int) + or isinstance(image_id, bool) + or image_id not in image_ids + or not isinstance(category_id, int) + or isinstance(category_id, bool) + or category_id not in category_ids + ): + return False + image_shape = image_shapes.get(image_id) + if image_shape is None: + return False + image_height, image_width = image_shape + bbox = record.get("bbox") + if ( + not isinstance(bbox, list) + or len(bbox) != 4 + or any( + not isinstance(value, (int, float)) + or isinstance(value, bool) + or not np.isfinite(value) + for value in bbox + ) + or bbox[2] <= 0 + or bbox[3] <= 0 + or bbox[0] >= image_width + or bbox[0] + bbox[2] <= 0 + or bbox[1] >= image_height + or bbox[1] + bbox[3] <= 0 + ): + return False + area = record.get("area") + iscrowd = record.get("iscrowd") + if ( + not isinstance(area, (int, float)) + or isinstance(area, bool) + or not np.isfinite(area) + or area <= 0 + or area > image_height * image_width + or not isinstance(iscrowd, int) + or isinstance(iscrowd, bool) + or iscrowd not in {0, 1} + ): + return False + if task == "pose_estimation" and area > bbox[2] * bbox[3]: + return False + if task == "instance_segmentation": + segmentation = record.get("segmentation") + if isinstance(segmentation, list): + if not segmentation or any( + not isinstance(polygon, list) + or len(polygon) < 6 + or len(polygon) % 2 + or any( + not isinstance(value, (int, float)) + or isinstance(value, bool) + or not np.isfinite(value) + for value in polygon + ) + or not _has_positive_polygon_area(polygon) + or not _polygon_has_positive_image_overlap( + polygon, image_shapes.get(image_id) + ) + or not _valid_coco_polygon(polygon, image_shapes.get(image_id)) + for polygon in segmentation + ): + return False + elif isinstance(segmentation, dict): + if not _valid_coco_rle(segmentation, image_shapes.get(image_id)): + return False + else: + return False + if task == "pose_estimation": + keypoints = record.get("keypoints") + num_keypoints = record.get("num_keypoints") + if ( + not isinstance(keypoints, list) + or len(keypoints) != 51 + or any( + not isinstance(value, (int, float)) + or isinstance(value, bool) + or not np.isfinite(value) + for value in keypoints + ) + or not isinstance(num_keypoints, int) + or isinstance(num_keypoints, bool) + or not 0 <= num_keypoints <= 17 + or any( + keypoints[index] not in {0, 1, 2} + for index in range(2, len(keypoints), 3) + ) + or num_keypoints + != sum(keypoints[index] > 0 for index in range(2, len(keypoints), 3)) + ): + return False + height, width = image_shape + if any( + keypoints[index + 2] > 0 + and not ( + 0 <= keypoints[index] < width and 0 <= keypoints[index + 1] < height + ) + for index in range(0, len(keypoints), 3) + ): + return False + annotation_ids.append(annotation_id) + return len(annotation_ids) == len(set(annotation_ids)) + + +def _decode_coco_rle_counts(counts: str) -> list[int] | None: + """Decode COCO's compact RLE run-length string without trusting its payload.""" + + run_counts: list[int] = [] + position = 0 + while position < len(counts): + value = 0 + shift = 0 + more = True + while more: + if position >= len(counts): + return None + code = ord(counts[position]) - 48 + position += 1 + if not 0 <= code <= 0x3F: + return None + value |= (code & 0x1F) << shift + more = bool(code & 0x20) + shift += 5 + if shift > 60: + return None + if not more and code & 0x10: + value |= -1 << shift + if len(run_counts) > 2: + value += run_counts[-2] + if value < 0: + return None + run_counts.append(value) + return run_counts + + +def _valid_coco_rle( + segmentation: dict[str, Any], image_shape: tuple[int, int] | None +) -> bool: + """Validate and decode an RLE mask against its referenced COCO image shape.""" + + counts = segmentation.get("counts") + size = segmentation.get("size") + if ( + image_shape is None + or not isinstance(size, list) + or len(size) != 2 + or any( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + for value in size + ) + or tuple(size) != image_shape + or not isinstance(counts, (str, list)) + ): + return False + if isinstance(counts, list): + if any( + not isinstance(value, int) or isinstance(value, bool) or value < 0 + for value in counts + ): + return False + run_counts = counts + else: + run_counts = _decode_coco_rle_counts(counts) + if run_counts is None: + return False + if sum(run_counts) != math.prod(size): + return False + try: + encoded = ( + coco_mask.frPyObjects(segmentation, size[0], size[1]) + if isinstance(counts, list) + else segmentation + ) + decoded = np.asarray(coco_mask.decode(encoded)) + except (RuntimeError, TypeError, ValueError): + return False + return decoded.shape == tuple(size) and bool(np.any(decoded)) + + +def _valid_coco_polygon( + polygon: list[int | float], image_shape: tuple[int, int] | None +) -> bool: + """Require a COCO polygon to rasterize to foreground in its image.""" + + if image_shape is None: + return True + height, width = image_shape + try: + encoded = coco_mask.frPyObjects([polygon], height, width) + decoded = np.asarray(coco_mask.decode(encoded)) + except (RuntimeError, TypeError, ValueError): + return False + return bool(np.any(decoded)) + + +def _coco_ready(root: Path, task: str) -> bool: + """Check the complete COCO 2017 image split and task annotation metadata.""" + + annotation_name = ( + "person_keypoints_val2017.json" + if task == "pose_estimation" + else "instances_val2017.json" + ) + annotation_names = _load_coco_image_names(root / annotation_name, task) + if annotation_names is None: + return False + image_dir = root / "val2017" + if not image_dir.is_dir(): + return False + image_paths = [ + path + for path in image_dir.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ] + if len(image_paths) != COCO_VALIDATION_SAMPLE_COUNT or any( + path.suffix.lower() != ".jpg" or COCO_IMAGE_PATTERN.fullmatch(path.stem) is None + for path in image_paths + ): + return False + if {path.name for path in image_paths} != annotation_names: + return False + try: + payload = json.loads((root / annotation_name).read_text(encoding="utf-8")) + image_records = payload["images"] + image_shapes = { + record["file_name"]: (record["height"], record["width"]) + for record in image_records + } + except (KeyError, TypeError, ValueError, json.JSONDecodeError, OSError): + return False + for image_path in image_paths: + expected_shape = image_shapes.get(image_path.name) + if ( + not isinstance(expected_shape, tuple) + or len(expected_shape) != 2 + or any( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + for value in expected_shape + ) + ): + return False + try: + with Image.open(image_path) as image: + image.load() + width, height = image.size + except OSError: + return False + if (height, width) != expected_shape: + return False + return True + + +def _dotav1_ready(root: Path) -> bool: + """Check complete paired DOTAv1 validation images and labels.""" + + flat_image_dir = root / "images" + flat_images = _files_by_stem(flat_image_dir, IMAGE_SUFFIXES) + if flat_images is None: + return False + image_dir = flat_image_dir if flat_images else flat_image_dir / "val" + images = flat_images if flat_images else _files_by_stem(image_dir, IMAGE_SUFFIXES) + if images is None or len(images) != DOTAV1_VALIDATION_SAMPLE_COUNT: + return False + + normalized_labels = _files_by_stem(root / "labels" / "val", {".txt"}) + original_labels = _files_by_stem(root / "labels" / "val_original", {".txt"}) + if normalized_labels is None or original_labels is None: + return False + label_stems = normalized_labels.keys() | original_labels.keys() + return images.keys() == label_stems + + +def _widerface_ready(root: Path) -> bool: + """Check the complete WiderFace validation image tree and metadata files.""" + + required_files = ( + "wider_face_val.mat", + "wider_easy_val.mat", + "wider_medium_val.mat", + "wider_hard_val.mat", + ) + if not all((root / file_name).is_file() for file_name in required_files): + return False + expected_images = _load_widerface_image_names(root / "wider_face_val.mat") + if expected_images is None: + return False + image_root = root / "images" + if not image_root.is_dir(): + return False + event_dirs = [path for path in image_root.iterdir() if path.is_dir()] + if len(event_dirs) != WIDERFACE_EVENT_COUNT or any( + WIDERFACE_EVENT_PATTERN.fullmatch(path.name) is None for path in event_dirs + ): + return False + if {path.name for path in event_dirs} != expected_images.keys(): + return False + actual_images = { + event_dir.name: { + path.name + for path in event_dir.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + } + for event_dir in event_dirs + } + image_shapes = _widerface_image_shapes(root, expected_images) + return ( + actual_images == expected_images + and sum(len(image_names) for image_names in actual_images.values()) + == WIDERFACE_VALIDATION_SAMPLE_COUNT + and image_shapes is not None + and _widerface_difficulty_metadata_ready( + root, expected_images, image_shapes=image_shapes + ) + ) + + +def _widerface_difficulty_metadata_ready( + root: Path, + expected_images: dict[str, set[str]], + *, + image_shapes: list[list[tuple[int, int]]] | None = None, +) -> bool: + """Validate WiderFace difficulty metadata and its decoded-image geometry.""" + + try: + main = loadmat(root / "wider_face_val.mat") + face_boxes = main["face_bbx_list"] + difficulties = [ + loadmat(root / file_name)["gt_list"] + for file_name in ( + "wider_easy_val.mat", + "wider_medium_val.mat", + "wider_hard_val.mat", + ) + ] + except (IndexError, KeyError, MatReadError, OSError, TypeError, ValueError): + return False + if len(face_boxes) != len(expected_images) or any( + len(table) != len(expected_images) for table in difficulties + ): + return False + if image_shapes is not None and len(image_shapes) != len(expected_images): + return False + difficulty_has_eligible_face = [False] * len(difficulties) + for event_index, image_names in enumerate(expected_images.values()): + try: + event_faces = face_boxes[event_index][0] + except (IndexError, TypeError): + return False + if len(event_faces) != len(image_names): + return False + if image_shapes is not None and len(image_shapes[event_index]) != len( + image_names + ): + return False + face_counts: list[int] = [] + for image_index, face_entry in enumerate(event_faces): + try: + face_array = np.asarray(face_entry[0]) + except (IndexError, TypeError): + return False + if ( + face_array.ndim != 2 + or face_array.shape[1] != 4 + or len(face_array) == 0 + or not np.issubdtype(face_array.dtype, np.number) + or np.issubdtype(face_array.dtype, np.complexfloating) + or not np.isfinite(face_array).all() + or (face_array[:, 2:] <= 0).any() + or len(np.unique(face_array, axis=0)) != len(face_array) + ): + return False + if image_shapes is not None: + height, width = image_shapes[event_index][image_index] + if ( + height <= 0 + or width <= 0 + or not ( + (face_array[:, 0] < width) + & (face_array[:, 0] + face_array[:, 2] > 0) + & (face_array[:, 1] < height) + & (face_array[:, 1] + face_array[:, 3] > 0) + ).all() + ): + return False + face_counts.append(len(face_array)) + for difficulty_index, table in enumerate(difficulties): + try: + event_indices = table[event_index][0] + except (IndexError, TypeError): + return False + if len(event_indices) != len(image_names): + return False + for image_index in range(len(event_faces)): + try: + keep_indices = np.asarray(event_indices[image_index][0]) + except (IndexError, TypeError): + return False + if keep_indices.ndim == 0: + keep_indices = keep_indices.reshape(1) + elif keep_indices.ndim == 2: + if keep_indices.size == 0: + if keep_indices.shape[0] != 0: + return False + keep_indices = keep_indices.reshape(0) + elif keep_indices.shape[1] == 1: + keep_indices = keep_indices[:, 0] + else: + return False + elif keep_indices.ndim != 1: + return False + try: + valid_indices = ( + np.isfinite(keep_indices).all() + and np.equal(keep_indices, np.trunc(keep_indices)).all() + ) + except (TypeError, ValueError): + return False + if not valid_indices: + return False + if keep_indices.size and ( + int(keep_indices.min()) < 1 + or int(keep_indices.max()) > face_counts[image_index] + ): + return False + if len(np.unique(keep_indices)) != keep_indices.size: + return False + difficulty_has_eligible_face[difficulty_index] |= bool( + keep_indices.size + ) + if not all(difficulty_has_eligible_face): + return False + return True + + +def _flatten_matlab_strings(value: Any) -> list[str]: + """Flatten strings stored inside nested MATLAB cell arrays.""" + + if isinstance(value, (str, np.str_)): + return [str(value)] + if isinstance(value, np.ndarray): + strings: list[str] = [] + for item in value.flat: + strings.extend(_flatten_matlab_strings(item)) + return strings + return [] + + +def _load_widerface_image_names(annotation_path: Path) -> dict[str, set[str]] | None: + """Load exact event and image identities from WiderFace validation metadata.""" + + try: + annotation = loadmat(annotation_path) + event_list = annotation["event_list"] + file_list = annotation["file_list"] + except (IndexError, KeyError, MatReadError, OSError, TypeError, ValueError): + return None + if len(event_list) != len(file_list): + return None + + expected: dict[str, set[str]] = {} + for event_cell, file_cell in zip(event_list, file_list, strict=True): + event_names = _flatten_matlab_strings(event_cell) + image_stems = _flatten_matlab_strings(file_cell) + if ( + len(event_names) != 1 + or not image_stems + or len(image_stems) != len(set(image_stems)) + or WIDERFACE_EVENT_PATTERN.fullmatch(event_names[0]) is None + or any(not stem or Path(stem).name != stem for stem in image_stems) + or event_names[0] in expected + ): + return None + expected[event_names[0]] = {f"{stem}.jpg" for stem in image_stems} + if len(expected) != WIDERFACE_EVENT_COUNT: + return None + return expected + + +def _widerface_image_shapes( + root: Path, expected_images: dict[str, set[str]] +) -> list[list[tuple[int, int]]] | None: + """Load WiderFace image shapes in the annotation's event and file order.""" + + try: + annotation = loadmat(root / "wider_face_val.mat") + event_list = annotation["event_list"] + file_list = annotation["file_list"] + except (IndexError, KeyError, MatReadError, OSError, TypeError, ValueError): + return None + if len(event_list) != len(file_list) or len(event_list) != len(expected_images): + return None + + all_shapes: list[list[tuple[int, int]]] = [] + for event_cell, file_cell in zip(event_list, file_list, strict=True): + event_names = _flatten_matlab_strings(event_cell) + image_stems = _flatten_matlab_strings(file_cell) + if ( + len(event_names) != 1 + or not image_stems + or expected_images.get(event_names[0]) + != {f"{stem}.jpg" for stem in image_stems} + ): + return None + event_shapes: list[tuple[int, int]] = [] + for stem in image_stems: + image_path = root / "images" / event_names[0] / f"{stem}.jpg" + if image_path.is_symlink(): + return None + try: + with Image.open(image_path) as image: + image.load() + width, height = image.size + except OSError: + return None + event_shapes.append((height, width)) + all_shapes.append(event_shapes) + return all_shapes + + +def dense_dataset_ready(data_path: str | Path, dataset: str) -> bool: + """Return whether a dense dataset matches its taxonomy and full validation split. + + Args: + data_path: Organized dataset root. + dataset: Dense validation taxonomy. + + Returns: + Whether the dataset has the expected filename identity, matched targets, + and complete validation sample count. + """ + + root = Path(data_path).expanduser() + if _path_has_symlink_component(root) or not root.is_dir(): + return False + normalized = dataset.lower() + if normalized == "nyu-depth": + images = _files_by_stem( + root / "images", {".jpg", ".jpeg", ".png"}, reject_symlinks=True + ) + depths = _files_by_stem(root / "depth", {".npy"}, reject_symlinks=True) + if images is None or depths is None: + return False + if not ( + len(images) == NYU_DEPTH_VALIDATION_SAMPLE_COUNT + and len(depths) == NYU_DEPTH_VALIDATION_SAMPLE_COUNT + and images.keys() == depths.keys() + ): + return False + for stem, image_path in images.items(): + try: + with Image.open(image_path) as image: + image.load() + image_shape = (image.height, image.width) + raw_depth = np.load(depths[stem], allow_pickle=False) + except (OSError, ValueError): + return False + if not np.issubdtype(raw_depth.dtype, np.number) or np.issubdtype( + raw_depth.dtype, np.complexfloating + ): + return False + with np.errstate(over="ignore", invalid="ignore"): + depth = np.asarray(raw_depth, dtype=np.float32) + if ( + depth.ndim != 2 + or depth.shape != image_shape + or not np.isfinite(depth).all() + or bool((depth < 0).any()) + or not bool(((depth > 0.001) & (depth < 100.0)).any()) + ): + return False + return True + + images = _files_by_stem( + root / "images", {".jpg", ".jpeg", ".png"}, reject_symlinks=True + ) + annotations = _files_by_stem(root / "annotations", {".png"}, reject_symlinks=True) + if images is None or annotations is None or images.keys() != annotations.keys(): + return False + + for stem, image_path in images.items(): + try: + with Image.open(image_path) as image: + image.load() + image_shape = (image.height, image.width) + with Image.open(annotations[stem]) as annotation_image: + annotation = np.asarray(annotation_image) + except OSError: + return False + if normalized == "cityscapes" and annotation.ndim == 3: + if ( + annotation.shape[2] not in {3, 4} + or not np.array_equal(annotation[..., 0], annotation[..., 1]) + or not np.array_equal(annotation[..., 0], annotation[..., 2]) + ): + return False + annotation = annotation[..., 0] + if annotation.ndim != 2 or annotation.shape != image_shape: + return False + if normalized == "ade20k": + if annotation.dtype != np.uint8 or ( + annotation.size and int(annotation.max()) > 150 + ): + return False + if not bool((annotation > 0).any()): + return False + elif normalized == "cityscapes": + if annotation.size and ( + int(annotation.min()) < 0 or int(annotation.max()) > 255 + ): + return False + if not bool(np.all((annotation <= 33) | (annotation == 255))): + return False + train_ids = CITYSCAPES_SOURCE_TO_TRAIN_ID[annotation.astype(np.uint8)] + if not bool((train_ids != 255).any()): + return False + + if normalized == "ade20k": + return ( + len(images) == ADE20K_VALIDATION_SAMPLE_COUNT + and all(stem.startswith("ADE_val_") for stem in images) + and all( + path.suffix.lower() in {".jpg", ".jpeg"} for path in images.values() + ) + and all( + not (root / file_name).is_symlink() and (root / file_name).is_file() + for file_name in ADE20K_METADATA_FILES + ) + ) + if normalized == "cityscapes": + city_counts: dict[str, int] = {} + for stem in images: + match = CITYSCAPES_SAMPLE_ID_PATTERN.fullmatch(stem) + if match is not None: + city = match.group("city") + city_counts[city] = city_counts.get(city, 0) + 1 + return ( + len(images) == CITYSCAPES_VALIDATION_SAMPLE_COUNT + and all( + CITYSCAPES_SAMPLE_ID_PATTERN.fullmatch(stem) is not None + for stem in images + ) + and all(path.suffix.lower() == ".png" for path in images.values()) + and ( + CITYSCAPES_VALIDATION_SAMPLE_COUNT != 500 + or city_counts == CITYSCAPES_VALIDATION_CITY_COUNTS + ) + ) + return False + + +def dataset_ready(data_path: str | Path, task: str, dataset: str | None = None) -> bool: + """Return whether an organized dataset matches its task, taxonomy, and full validation split. + + Args: + data_path: Organized dataset root. + task: Canonical vision task. + dataset: Optional validation taxonomy. + + Returns: + Whether the dataset has the expected identity, metadata, and sample count. + """ + + root = Path(data_path).expanduser() + normalized_task = normalize_vision_task(task) + expected_dataset = { + "image_classification": "imagenet", + "object_detection": "coco", + "instance_segmentation": "coco", + "pose_estimation": "coco", + "face_detection": "widerface", + "obb": "dotav1", + "depth_estimation": "nyu-depth", + }.get(normalized_task) + normalized_dataset = (dataset or expected_dataset or "").lower() + + if normalized_task == "semantic_segmentation": + return dense_dataset_ready(root, normalized_dataset or "ade20k") + if expected_dataset is None or normalized_dataset != expected_dataset: + return False + if normalized_task == "image_classification": + return _imagenet_ready(root) + if normalized_task in { + "object_detection", + "instance_segmentation", + "pose_estimation", + }: + return _coco_ready(root, normalized_task) + if normalized_task == "face_detection": + return _widerface_ready(root) + if normalized_task == "obb": + return _dotav1_ready(root) + if normalized_task == "depth_estimation": + return dense_dataset_ready(root, normalized_dataset) + return False diff --git a/mblt_vision/utils/evaluation/__init__.py b/mblt_vision/utils/evaluation/__init__.py new file mode 100644 index 0000000..21b2916 --- /dev/null +++ b/mblt_vision/utils/evaluation/__init__.py @@ -0,0 +1,51 @@ +""" +Evaluation scripts for various datasets. +""" + +from __future__ import annotations + +from ._result import EvaluationResult +from .eval_ade20k import ( + ADE20KResult, + SemanticMetricAccumulator, + SemanticSegmentationResult, + calculate_semantic_metrics, + eval_ade20k, + eval_semantic_segmentation, +) +from .eval_cityscapes import eval_cityscapes +from .eval_coco import COCOResult, eval_coco, eval_coco_metrics +from .eval_dota import DOTAResult, eval_dota +from .eval_imagenet import ImageNetResult, eval_imagenet, eval_imagenet_metrics +from .eval_nyu_depth import ( + NYUDepthMetricAccumulator, + NYUDepthResult, + calculate_nyu_depth_metrics, + eval_nyu_depth, +) +from .eval_widerface import WiderFaceResult, eval_widerface + +__all__: list[str] = [ + "eval_coco", + "eval_coco_metrics", + "COCOResult", + "EvaluationResult", + "ADE20KResult", + "SemanticMetricAccumulator", + "SemanticSegmentationResult", + "calculate_semantic_metrics", + "eval_ade20k", + "eval_cityscapes", + "eval_semantic_segmentation", + "DOTAResult", + "eval_dota", + "ImageNetResult", + "eval_imagenet", + "eval_imagenet_metrics", + "NYUDepthResult", + "NYUDepthMetricAccumulator", + "calculate_nyu_depth_metrics", + "eval_nyu_depth", + "WiderFaceResult", + "eval_widerface", +] diff --git a/mblt_vision/utils/evaluation/_result.py b/mblt_vision/utils/evaluation/_result.py new file mode 100644 index 0000000..2e8c202 --- /dev/null +++ b/mblt_vision/utils/evaluation/_result.py @@ -0,0 +1,22 @@ +"""Shared structural contract for Vision evaluator results.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class EvaluationResult(Protocol): + """Result object exposing the primary and secondary benchmark scores.""" + + @property + def primary_score(self) -> float: + """Return the evaluator's primary score.""" + + ... + + @property + def secondary_score(self) -> float: + """Return the evaluator's secondary score.""" + + ... diff --git a/mblt_vision/utils/evaluation/eval_ade20k.py b/mblt_vision/utils/evaluation/eval_ade20k.py new file mode 100644 index 0000000..5b6897a --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_ade20k.py @@ -0,0 +1,302 @@ +"""ADE20K evaluation for semantic-segmentation models.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple + +import numpy as np +import torch +from tqdm import tqdm + +from ..datasets import ( + CustomADE20K, + CustomCityscapes, + get_ade20k_loader, + get_cityscapes_loader, +) + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + + +class SemanticSegmentationResult(NamedTuple): + """Generic semantic metrics ordered from primary to secondary.""" + + miou: float + pixel_accuracy: float + + @property + def primary_score(self) -> float: + """Return mean intersection-over-union.""" + + return self.miou + + @property + def secondary_score(self) -> float: + """Return overall valid-pixel accuracy.""" + + return self.pixel_accuracy + + +ADE20KResult = SemanticSegmentationResult + + +class SemanticMetricAccumulator: + """Accumulate an ignore-aware semantic confusion matrix.""" + + def __init__(self, nc: int, ignore_label: int = 255) -> None: + """Initialize an empty confusion matrix. + + Args: + nc: Number of semantic classes. + ignore_label: Target label excluded from metrics. + """ + + self.nc = nc + self.ignore_label = ignore_label + self.matrix = np.zeros((nc, nc), dtype=np.int64) + + def update(self, prediction: np.ndarray, target: np.ndarray) -> None: + """Accumulate one or more predicted and target class maps. + + Args: + prediction: Predicted class maps. + target: Target class maps with optional ignore labels. + + Raises: + ValueError: If prediction and target shapes differ. + """ + + prediction = np.asarray(prediction) + target = np.asarray(target) + if prediction.shape != target.shape: + raise ValueError( + f"Semantic prediction and target shapes must match, got {prediction.shape} and {target.shape}." + ) + valid_target = ( + np.isfinite(target) + & (target >= 0) + & (target < self.nc) + & (target == np.floor(target)) + ) + allowed_target = valid_target | (target == self.ignore_label) + if not bool(allowed_target.all()): + invalid_values = np.asarray(np.unique(target[~allowed_target])) + raise ValueError( + f"Semantic targets must be finite class IDs in [0, {self.nc - 1}] " + f"or ignore label {self.ignore_label}; got {invalid_values.tolist()}." + ) + valid_prediction = ( + np.isfinite(prediction) + & (prediction >= 0) + & (prediction < self.nc) + & (prediction == np.floor(prediction)) + ) + invalid_prediction = valid_target & ~valid_prediction + if invalid_prediction.any(): + invalid_values = np.asarray(np.unique(prediction[invalid_prediction])) + raise ValueError( + f"Semantic predictions at valid target pixels must be finite class IDs in [0, {self.nc - 1}], " + f"got {invalid_values.tolist()}." + ) + if valid_target.any(): + histogram = np.bincount( + self.nc * target[valid_target].astype(np.int64) + + prediction[valid_target].astype(np.int64), + minlength=self.nc**2, + ) + self.matrix += histogram.reshape(self.nc, self.nc) + + def result(self) -> SemanticSegmentationResult: + """Compute mIoU over present classes and overall pixel accuracy. + + Returns: + Pooled semantic-segmentation metrics. + + Raises: + ValueError: If no valid target pixels were accumulated. + """ + + ground_truth = self.matrix.sum(axis=1) + predicted = self.matrix.sum(axis=0) + intersection = np.diag(self.matrix) + union = ground_truth + predicted - intersection + present = ground_truth > 0 + if not present.any(): + raise ValueError("Semantic evaluation received no valid target pixels.") + iou = np.divide( + intersection, + union, + out=np.zeros(self.nc, dtype=np.float64), + where=union > 0, + ) + total = int(self.matrix.sum()) + return SemanticSegmentationResult( + miou=float(iou[present].mean()), + pixel_accuracy=float(intersection.sum() / total), + ) + + +def calculate_semantic_metrics( + prediction: np.ndarray, + target: np.ndarray, + nc: int = 150, + ignore_label: int = 255, +) -> ADE20KResult: + """Calculate semantic metrics for one batch of class maps. + + Args: + prediction: Predicted class maps. + target: Target class maps. + nc: Number of semantic classes. + ignore_label: Target label excluded from metrics. + + Returns: + Semantic metrics exposed through the ADE20K compatibility alias. + + Raises: + ValueError: If shapes differ or no valid target pixels are present. + """ + + accumulator = SemanticMetricAccumulator(nc=nc, ignore_label=ignore_label) + accumulator.update(prediction, target) + return accumulator.result() + + +def _evaluate_semantic_loader( + model: MBLT_Engine, + loader: torch.utils.data.DataLoader, + nc: int, + description: str = "Evaluating semantic segmentation", +) -> SemanticSegmentationResult: + """Evaluate input-space semantic maps from an organized validation loader. + + Args: + model: Initialized semantic-segmentation engine. + loader: Organized validation data loader. + nc: Number of semantic classes. + description: Progress-bar description. + + Returns: + Pooled semantic-segmentation metrics. + + Raises: + ValueError: If postprocessing returns no class maps or no valid targets exist. + """ + + accumulator = SemanticMetricAccumulator(nc=nc) + for inputs, targets, _shapes, _ratio_pads, _ in tqdm(loader, desc=description): + # TODO: Restore logits to original geometry using shapes and ratio_pads when + # Ultralytics adopts native-geometry semantic validation metrics. + result = model.postprocess(model(inputs)) + semantic_mask = result.semantic_mask + if semantic_mask is None: + raise ValueError("Semantic postprocessor returned no class maps.") + prediction = ( + semantic_mask.detach().cpu().numpy() + if isinstance(semantic_mask, torch.Tensor) + else semantic_mask + ) + accumulator.update(np.asarray(prediction), targets) + return accumulator.result() + + +def eval_semantic_segmentation( + model: MBLT_Engine, + data_path: str, + batch_size: int, + dataset: str | None = None, +) -> SemanticSegmentationResult: + """Evaluate a semantic model with the loader for its configured taxonomy. + + Args: + model: Initialized semantic-segmentation engine. + data_path: Organized dataset root. + batch_size: Number of validation samples per inference batch. + dataset: Optional taxonomy override. Defaults to ``model.post_cfg.dataset``. + + Returns: + Generic mIoU and pixel-accuracy result. + + Raises: + ValueError: If preprocessing metadata, image size, taxonomy, predictions, or targets are invalid. + """ + + configured_dataset = model.post_cfg.get("dataset") + if not isinstance(configured_dataset, str) or not configured_dataset: + raise ValueError( + "Semantic validation requires model.post_cfg.dataset to declare the model taxonomy." + ) + configured_taxonomy = configured_dataset.lower() + if dataset is not None: + if not isinstance(dataset, str) or not dataset: + raise ValueError( + "Semantic validation dataset overrides must be non-empty strings." + ) + requested_taxonomy = dataset.lower() + if requested_taxonomy != configured_taxonomy: + raise ValueError( + f"Requested semantic validation taxonomy {requested_taxonomy!r} conflicts with " + f"the model's configured taxonomy {configured_taxonomy!r}." + ) + taxonomy = configured_taxonomy + + letterbox_cfg = model.pre_cfg.get("LetterBox") + if not isinstance(letterbox_cfg, dict) or "img_size" not in letterbox_cfg: + raise ValueError( + "Semantic validation requires a LetterBox img_size in the model preprocessing config." + ) + image_size = letterbox_cfg["img_size"] + if not isinstance(image_size, list) or len(image_size) != 2: + raise ValueError( + "Semantic validation img_size must be a two-item [height, width] list." + ) + + image_size_tuple = (int(image_size[0]), int(image_size[1])) + if taxonomy == "ade20k": + validation_dataset = CustomADE20K(data_path) + loader = get_ade20k_loader( + validation_dataset, + batch_size, + model.preprocess_with_metadata, + image_size=image_size_tuple, + ) + default_nc = 150 + description = "Evaluating ADE20K" + elif taxonomy == "cityscapes": + validation_dataset = CustomCityscapes(data_path) + loader = get_cityscapes_loader( + validation_dataset, + batch_size, + model.preprocess_with_metadata, + image_size=image_size_tuple, + ) + default_nc = 19 + description = "Evaluating Cityscapes" + else: + raise ValueError(f"Unsupported semantic validation dataset: {taxonomy!r}.") + nc = int(getattr(model.postprocessor, "nc", default_nc)) + return _evaluate_semantic_loader(model, loader, nc, description=description) + + +def eval_ade20k(model: MBLT_Engine, data_path: str, batch_size: int) -> ADE20KResult: + """Evaluate a semantic-segmentation model on ADE20K validation masks. + + Args: + model: Initialized semantic-segmentation engine. + data_path: Organized ADE20K dataset root. + batch_size: Number of validation samples per inference batch. + + Returns: + ADE20K metrics through the generic semantic result type. + + Raises: + ValueError: If the dataset, preprocessing metadata, predictions, or targets are invalid. + """ + + return eval_semantic_segmentation( + model, + data_path, + batch_size, + dataset="ade20k", + ) diff --git a/mblt_vision/utils/evaluation/eval_cityscapes.py b/mblt_vision/utils/evaluation/eval_cityscapes.py new file mode 100644 index 0000000..d12f9e5 --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_cityscapes.py @@ -0,0 +1,32 @@ +"""Cityscapes evaluation for semantic-segmentation models.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .eval_ade20k import SemanticSegmentationResult, eval_semantic_segmentation + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + + +def eval_cityscapes( + model: MBLT_Engine, data_path: str, batch_size: int +) -> SemanticSegmentationResult: + """Evaluate a semantic-segmentation model on Cityscapes validation masks. + + Args: + model: Initialized semantic-segmentation engine. + data_path: Organized Cityscapes dataset root. + batch_size: Number of validation samples per inference batch. + + Returns: + Cityscapes mIoU and pixel-accuracy metrics. + + Raises: + ValueError: If the dataset, preprocessing metadata, predictions, or targets are invalid. + """ + + return eval_semantic_segmentation( + model, data_path, batch_size, dataset="cityscapes" + ) diff --git a/mblt_vision/utils/evaluation/eval_coco.py b/mblt_vision/utils/evaluation/eval_coco.py new file mode 100644 index 0000000..b1e00c7 --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_coco.py @@ -0,0 +1,465 @@ +"""Evaluation script for COCO dataset.""" + +from __future__ import annotations + +import logging +import math +import os +from time import time +from typing import TYPE_CHECKING, Any, NamedTuple + +from faster_coco_eval import COCO, COCOeval_faster +from tqdm import tqdm + +from ..._tasks import normalize_vision_task +from ...datasets import get_dataset_category_ids +from ..datasets import CustomCOCODataset, get_coco_loader +from ..datasets.readiness import _coco_task_annotations_valid + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + from ..results import Results + +logger = logging.getLogger(__name__) + + +class COCOResult(NamedTuple): + """COCO mAP metrics.""" + + map5095: float + map50: float + + @property + def primary_score(self) -> float: + """Return mAP50-95.""" + + return self.map5095 + + @property + def secondary_score(self) -> float: + """Return mAP50.""" + + return self.map50 + + +def _require_batch_cardinality(expected: int, **values: Any) -> None: + """Reject postprocessing metadata that cannot represent every input image.""" + + invalid = { + name: len(value) for name, value in values.items() if len(value) != expected + } + if invalid: + details = ", ".join(f"{name}={length}" for name, length in invalid.items()) + raise ValueError( + "COCO evaluation batch cardinality mismatch: " + f"expected {expected}, got {details}." + ) + + +def _validate_coco_dataset_taxonomy(dataset: CustomCOCODataset, task: str) -> None: + """Ensure direct COCO evaluation uses only task-compatible category IDs.""" + + raw_annotation = getattr(dataset, "raw_annotation", None) + if raw_annotation is not None and not isinstance(raw_annotation, dict): + raise ValueError("COCO evaluation dataset has invalid raw annotation data.") + raw_images = ( + raw_annotation.get("images") if isinstance(raw_annotation, dict) else None + ) + raw_categories = ( + raw_annotation.get("categories") if isinstance(raw_annotation, dict) else None + ) + raw_annotations = ( + raw_annotation.get("annotations") if isinstance(raw_annotation, dict) else None + ) + raw_image_records: list[Any] = [] + raw_category_records: list[Any] = [] + raw_annotation_records: list[Any] = [] + if raw_annotation is not None: + if not isinstance(raw_images, list): + raise ValueError("COCO evaluation dataset has malformed raw image table.") + if not isinstance(raw_categories, list): + raise ValueError( + "COCO evaluation dataset has malformed raw category table." + ) + if not isinstance(raw_annotations, list): + raise ValueError( + "COCO evaluation dataset has malformed raw annotation table." + ) + raw_image_records = raw_images + raw_category_records = raw_categories + raw_annotation_records = raw_annotations + + categories = getattr(dataset.coco, "cats", None) + if not isinstance(categories, dict) or not categories: + raise ValueError("COCO evaluation dataset must define at least one category.") + category_ids = set(categories) + if any( + not isinstance(category_id, int) or isinstance(category_id, bool) + for category_id in category_ids + ): + raise ValueError("COCO evaluation dataset contains invalid category IDs.") + expected_ids = ( + {1} if task == "pose_estimation" else set(get_dataset_category_ids("coco")) + ) + unsupported_ids = category_ids - expected_ids + if unsupported_ids: + raise ValueError( + "COCO evaluation dataset contains unsupported category IDs: " + f"{sorted(unsupported_ids)}." + ) + annotations = getattr(dataset.coco, "anns", None) + if not isinstance(annotations, dict): + raise ValueError("COCO evaluation dataset must define an annotation table.") + images = getattr(dataset.coco, "imgs", None) + if not isinstance(images, dict): + raise ValueError("COCO evaluation dataset must define an image table.") + image_shapes: dict[int, tuple[int, int] | None] = {} + for image_id, image in images.items(): + if ( + not isinstance(image_id, int) + or isinstance(image_id, bool) + or not isinstance(image, dict) + ): + raise ValueError("COCO evaluation dataset contains invalid image metadata.") + height, width = image.get("height"), image.get("width") + if not ( + isinstance(height, int) + and not isinstance(height, bool) + and height > 0 + and isinstance(width, int) + and not isinstance(width, bool) + and width > 0 + ): + raise ValueError( + "COCO evaluation dataset contains invalid image dimensions." + ) + image_shapes[image_id] = (height, width) + + annotation_records = ( + raw_annotation_records + if raw_annotation is not None + else list(annotations.values()) + ) + if not annotation_records: + raise ValueError("COCO evaluation dataset must define at least one annotation.") + if raw_annotation is not None: + raw_image_ids = [ + record.get("id") if isinstance(record, dict) else None + for record in raw_image_records + ] + raw_category_ids = [ + record.get("id") if isinstance(record, dict) else None + for record in raw_category_records + ] + raw_annotation_ids = [ + record.get("id") if isinstance(record, dict) else None + for record in raw_annotation_records + ] + if ( + any( + not isinstance(record_id, int) or isinstance(record_id, bool) + for record_id in ( + *raw_image_ids, + *raw_category_ids, + *raw_annotation_ids, + ) + ) + or len(raw_image_ids) != len(set(raw_image_ids)) + or len(raw_category_ids) != len(set(raw_category_ids)) + or len(raw_annotation_ids) != len(set(raw_annotation_ids)) + ): + raise ValueError( + "COCO evaluation dataset has duplicate or invalid raw IDs." + ) + if ( + set(raw_image_ids) != set(images) + or set(raw_category_ids) != category_ids + or set(raw_annotation_ids) != set(annotations) + ): + raise ValueError( + "COCO evaluation dataset raw and indexed records disagree." + ) + if not _coco_task_annotations_valid( + annotation_records, + image_ids=set(images), + category_ids=category_ids, + image_shapes=image_shapes, + task=task, + ): + raise ValueError( + "COCO evaluation dataset contains invalid task-specific annotations." + ) + + +def format_coco_results( + task: str, + nms_outs: Results, + input_shape: tuple[int, ...], + org_shape: tuple[int, ...], + ratio_pad: list[Any], + idx: list[int], + dataset_ids: list[int], + postprocess: Any, +) -> list[dict[str, Any]]: + """Format the results for COCO evaluation. + + Args: + task (str): The task to evaluate. + nms_outs (Results): The output of the postprocessing. + input_shape (tuple): The shape of the input tensor. + org_shape (tuple): The original shape of the image. + idx (list): The indices of the images in the batch. + dataset_ids (list): The list of image IDs in the dataset. + postprocess: The postprocessing instance. + Returns: + list: The formatted results. + """ + results = [] + if task == "object_detection": + labels_list, boxes_list, scores_list = postprocess.nmsout2eval( + nms_outs.output, + input_shape, + org_shape, + ratio_pad=ratio_pad, + ) + _require_batch_cardinality( + len(idx), + org_shape=org_shape, + ratio_pad=ratio_pad, + labels=labels_list, + boxes=boxes_list, + scores=scores_list, + ) + for i, labels, boxes, scores in zip( + idx, labels_list, boxes_list, scores_list, strict=True + ): + results.extend( + [ + { + "image_id": dataset_ids[i], + "category_id": label, + "bbox": box, + "score": score, + } + for box, score, label in zip(boxes, scores, labels, strict=True) + ] + ) + elif task == "instance_segmentation": + labels_list, boxes_list, scores_list, extra_list = postprocess.nmsout2eval( + nms_outs.output, + input_shape, + org_shape, + ratio_pad=ratio_pad, + ) + _require_batch_cardinality( + len(idx), + org_shape=org_shape, + ratio_pad=ratio_pad, + labels=labels_list, + boxes=boxes_list, + scores=scores_list, + extra=extra_list, + ) + for i, labels, boxes, scores, extra in zip( + idx, labels_list, boxes_list, scores_list, extra_list, strict=True + ): + results.extend( + [ + { + "image_id": dataset_ids[i], + "category_id": label, + "bbox": box, + "score": score, + "segmentation": extra, + } + for box, score, label, extra in zip( + boxes, scores, labels, extra, strict=True + ) + ] + ) + elif task == "pose_estimation": + labels_list, boxes_list, scores_list, extra_list = postprocess.nmsout2eval( + nms_outs.output, + input_shape, + org_shape, + ratio_pad=ratio_pad, + ) + _require_batch_cardinality( + len(idx), + org_shape=org_shape, + ratio_pad=ratio_pad, + labels=labels_list, + boxes=boxes_list, + scores=scores_list, + extra=extra_list, + ) + for i, labels, boxes, scores, extra in zip( + idx, labels_list, boxes_list, scores_list, extra_list, strict=True + ): + results.extend( + [ + { + "image_id": dataset_ids[i], + "category_id": label, + "bbox": box, + "score": score, + "keypoints": extra, + } + for box, score, label, extra in zip( + boxes, scores, labels, extra, strict=True + ) + ] + ) + else: + raise NotImplementedError( + f"Only object detection, instance segmentation, and pose estimation are supported, but we got {task}" + ) + return results + + +def eval_coco( + model: MBLT_Engine, + data_path: str, + batch_size: int, + conf_thres: float | None = None, + iou_thres: float | None = None, +) -> float: + """Evaluate a model on COCO and return the legacy numeric mAP50-95 score.""" + + return eval_coco_metrics( + model, data_path, batch_size, conf_thres, iou_thres + ).primary_score + + +def eval_coco_metrics( + model: MBLT_Engine, + data_path: str, + batch_size: int, + conf_thres: float | None = None, + iou_thres: float | None = None, +) -> COCOResult: + """Evaluate a model on COCO and return structured mAP metrics. + + Args: + model (MBLT_Engine): The model engine to evaluate. + data_path (str): Path to the COCO dataset. + batch_size (int): Batch size for evaluation. + conf_thres (float | None): Optional confidence threshold override. + iou_thres (float | None): Optional IoU threshold override. + + Returns: + Structured mAP50-95 primary and mAP50 secondary metrics. + """ + task = normalize_vision_task( + model.post_cfg["task"], + supported=("object_detection", "instance_segmentation", "pose_estimation"), + ) + dataset_name = model.post_cfg.get("dataset") + if not isinstance(dataset_name, str) or dataset_name.lower() != "coco": + raise ValueError( + "COCO evaluation requires model post_cfg.dataset to be 'coco', " + f"got {dataset_name!r}." + ) + if task in {"object_detection", "instance_segmentation"}: + dataset = CustomCOCODataset( + os.path.join(data_path, "val2017"), + os.path.join(data_path, "instances_val2017.json"), + ) + else: + dataset = CustomCOCODataset( + os.path.join(data_path, "val2017"), + os.path.join(data_path, "person_keypoints_val2017.json"), + ) + _validate_coco_dataset_taxonomy(dataset, task) + + dataloader = get_coco_loader(dataset, batch_size, model.preprocess_with_metadata) + model.set_postprocess_thresholds(conf_thres=conf_thres, iou_thres=iou_thres) + + results = [] + num_data = len(dataset) + total_iter = math.ceil(num_data / batch_size) + pbar = tqdm(dataloader, total=total_iter, desc="Evaluating COCO") + + inference_time = 0.0 + cum_num_data = 0 + + for input_npu, org_shape, ratio_pad, idx in pbar: + cum_num_data += len(idx) + tic = time() + out_npu = model(input_npu) + inference_time += time() - tic + + nms_outs = model.postprocess(out_npu, multi_label=True) + results.extend( + format_coco_results( + task, + nms_outs, + input_npu.shape[1:-1], + org_shape, + ratio_pad, + idx, + dataset.ids, + model.postprocessor, + ) + ) + + pbar.set_postfix_str(f"NPU FPS: {cum_num_data / inference_time:.3f}") + + pbar.close() + res = evaluate_predictions_on_coco(dataset.coco, results, task, img_ids=dataset.ids) + + print("COCO evaluation completed") + return COCOResult( + map5095=float(res.stats[0].item()), map50=float(res.stats[1].item()) + ) + + +def evaluate_predictions_on_coco( + coco_gt: COCO, + coco_results: list[dict[str, Any]], + task: str, + img_ids: list[int] | None = None, +) -> COCOeval_faster: + """Evaluates predictions using the COCO API. + + Args: + coco_gt (COCO): Ground truth COCO object. + coco_results (list): Predictions in COCO format. + task (str): Task type ('object_detection', 'instance_segmentation', or 'pose_estimation'). + img_ids: Optional image IDs to include in evaluation. + + Returns: + COCOeval_faster: The COCO evaluation object containing results. + """ + normalized_task = normalize_vision_task( + task, + supported=("object_detection", "instance_segmentation", "pose_estimation"), + ) + + if coco_results: + coco_dt = coco_gt.loadRes(coco_results) + else: + coco_dt = COCO() + + if normalized_task == "object_detection": + coco_eval = COCOeval_faster( + coco_gt, coco_dt, "bbox", print_function=logger.info + ) + elif normalized_task == "instance_segmentation": + coco_eval = COCOeval_faster( + coco_gt, coco_dt, "segm", print_function=logger.info + ) + elif normalized_task == "pose_estimation": + coco_eval = COCOeval_faster( + coco_gt, coco_dt, "keypoints", print_function=logger.info + ) + else: + raise RuntimeError(f"Unexpected validated COCO task: {normalized_task}") + + if img_ids is not None: + coco_eval.params.imgIds = img_ids + + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() + return coco_eval diff --git a/mblt_vision/utils/evaluation/eval_dota.py b/mblt_vision/utils/evaluation/eval_dota.py new file mode 100644 index 0000000..8db6392 --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_dota.py @@ -0,0 +1,878 @@ +"""Evaluation script for DOTAv1 OBB.""" + +from __future__ import annotations + +import math +import re +from collections import defaultdict +from pathlib import Path +from time import time +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +import numpy as np +import torch +from mblt_vision.utils.postprocess.common import ( + batch_probiou, + rotated_nms, + xywhr2xyxyxyxy, + xyxyxyxy2xywhr, +) +from tqdm import tqdm + +from ..._tasks import normalize_vision_task +from ..datasets import ( + CustomDOTAv1, + get_dota_loader, + get_dotav1_class_num, + get_dotav1_label, +) +from ..datasets.readiness import ( + _canonicalize_quadrilateral, + _polygon_has_positive_image_overlap, +) +from ..letterbox import RatioPad, resolve_ratio_pad + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + from ..results import Results + +DOTAV1_CLASS_TO_IDX = { + get_dotav1_label(index): index for index in range(get_dotav1_class_num()) +} + + +class DOTAResult(NamedTuple): + """DOTAv1 rotated detection metrics in the legacy tuple order.""" + + map50: float + map5095: float + + @property + def primary_score(self) -> float: + """Return the primary DOTAv1 validation metric.""" + return self.map5095 + + @property + def secondary_score(self) -> float: + """Return the secondary DOTAv1 validation metric.""" + return self.map50 + + +def _label_to_index(label: str) -> int: + """Convert a DOTAv1 class token to a class index.""" + try: + return int(label) + except ValueError: + return DOTAV1_CLASS_TO_IDX[label] + + +def _validate_polygon_area( + coords: torch.Tensor, annotation_path: Path, line_number: int +) -> None: + """Require a non-degenerate quadrilateral before converting it to an OBB.""" + + shifted = torch.roll(coords, shifts=-1, dims=0) + area = 0.5 * torch.abs( + torch.sum(coords[:, 0] * shifted[:, 1] - coords[:, 1] * shifted[:, 0]) + ) + if not bool(area > 0): + raise ValueError( + "DOTAv1 annotation polygon must have positive area at " + f"{annotation_path}:{line_number}." + ) + + +def _validate_polygon_vertices( + coords: torch.Tensor, annotation_path: Path, line_number: int +) -> None: + """Require four distinct, consistently ordered DOTAv1 quadrilateral vertices.""" + + if torch.unique(coords, dim=0).shape[0] != 4: + raise ValueError( + "DOTAv1 annotation polygon must contain four distinct vertices at " + f"{annotation_path}:{line_number}." + ) + edges = torch.roll(coords, shifts=-1, dims=0) - coords + next_edges = torch.roll(edges, shifts=-1, dims=0) + turns = edges[:, 0] * next_edges[:, 1] - edges[:, 1] * next_edges[:, 0] + if not bool(torch.all(turns > 0) or torch.all(turns < 0)): + raise ValueError( + "DOTAv1 annotation polygon vertices must be consistently ordered at " + f"{annotation_path}:{line_number}." + ) + + +def _validate_polygon_image_overlap( + coords: torch.Tensor, + image_shape: tuple[int, int], + annotation_path: Path, + line_number: int, +) -> None: + """Require a quadrilateral to cover non-zero area inside its source image.""" + + polygon = [coordinate for point in coords.tolist() for coordinate in point] + if not _polygon_has_positive_image_overlap(polygon, image_shape): + raise ValueError( + "DOTAv1 annotation polygon must overlap its source image at " + f"{annotation_path}:{line_number}." + ) + + +def _load_ground_truths( + data_path: str, dataset: CustomDOTAv1 +) -> dict[str, dict[str, torch.Tensor]]: + """Load DOTAv1 OBB ground-truth labels in original-image coordinates. + + Args: + data_path: DOTAv1 root directory. + dataset: Dataset containing image IDs and image paths. + + Returns: + Mapping from image ID to tensors for positive and ignored classes, polygons, + and ``xywhr`` boxes. + """ + label_dir = Path(data_path) / "labels" / "val" + original_label_dir = Path(data_path) / "labels" / "val_original" + image_ids = set(dataset.ids) + label_ids = {path.stem for path in label_dir.glob("*.txt")} | { + path.stem for path in original_label_dir.glob("*.txt") + } + orphan_label_ids = sorted(label_ids - image_ids) + if orphan_label_ids: + raise ValueError( + "DOTAv1 annotations have no corresponding validation image: " + f"{', '.join(orphan_label_ids[:5])}." + ) + ground_truths: dict[str, dict[str, torch.Tensor]] = {} + for image_id, image_path in zip(dataset.ids, dataset.image_paths): + image = dataset._load_image(image_path) + height, width = image.shape[:2] + label_path = label_dir / f"{image_id}.txt" + original_label_path = original_label_dir / f"{image_id}.txt" + classes = [] + polygons = [] + ignore_classes = [] + ignore_polygons = [] + seen_targets: set[tuple[int, tuple[float, ...]]] = set() + + # Organizer output keeps the official original annotation alongside the + # normalized convenience label. Prefer the authoritative original when + # both are present so a stale normalized file cannot hide ground truth. + if label_path.is_file() and not original_label_path.is_file(): + for line_number, line in enumerate( + label_path.read_text(encoding="utf-8").splitlines(), start=1 + ): + parts = line.split() + if not parts: + continue + if len(parts) < 9: + raise ValueError( + "Malformed normalized DOTAv1 annotation at " + f"{label_path}:{line_number}: expected at least 9 fields, " + f"got {len(parts)}." + ) + cls = _label_to_index(parts[0]) + if not 0 <= cls < get_dotav1_class_num(): + raise ValueError( + f"Unsupported DOTAv1 class index {cls} in {label_path}." + ) + coords = torch.tensor( + [float(value) for value in parts[1:9]], dtype=torch.float32 + ).reshape(4, 2) + if not bool(torch.isfinite(coords).all()): + raise ValueError( + "DOTAv1 annotation coordinates must be finite at " + f"{label_path}:{line_number}." + ) + coords[:, 0] *= width + coords[:, 1] *= height + _validate_polygon_area(coords, label_path, line_number) + _validate_polygon_vertices(coords, label_path, line_number) + _validate_polygon_image_overlap( + coords, (height, width), label_path, line_number + ) + if len(parts) >= 10 and parts[9] not in {"0", "1", "2"}: + raise ValueError( + f"Unsupported DOTAv1 difficulty flag {parts[9]!r} at " + f"{label_path}:{line_number}." + ) + target_key = ( + cls, + _canonicalize_quadrilateral(coords.flatten().tolist()), + ) + if target_key in seen_targets: + raise ValueError( + "Duplicate DOTAv1 annotation target at " + f"{label_path}:{line_number}." + ) + seen_targets.add(target_key) + if len(parts) >= 10 and parts[9] in {"1", "2"}: + ignore_classes.append(cls) + ignore_polygons.append(coords) + else: + classes.append(cls) + polygons.append(coords) + elif original_label_path.is_file(): + for line_number, line in enumerate( + original_label_path.read_text(encoding="utf-8").splitlines(), start=1 + ): + parts = line.split() + if parts and ( + parts[0].startswith("imagesource:") or parts[0].startswith("gsd:") + ): + continue + if len(parts) < 10: + raise ValueError( + "Malformed original DOTAv1 annotation at " + f"{original_label_path}:{line_number}: expected at least 10 fields, " + f"got {len(parts)}." + ) + cls = _label_to_index(parts[8]) + if not 0 <= cls < get_dotav1_class_num(): + raise ValueError( + f"Unsupported DOTAv1 class index {cls} in {original_label_path}." + ) + coords = torch.tensor( + [float(value) for value in parts[:8]], dtype=torch.float32 + ).reshape(4, 2) + if not bool(torch.isfinite(coords).all()): + raise ValueError( + "DOTAv1 annotation coordinates must be finite at " + f"{original_label_path}:{line_number}." + ) + _validate_polygon_area(coords, original_label_path, line_number) + _validate_polygon_vertices(coords, original_label_path, line_number) + _validate_polygon_image_overlap( + coords, (height, width), original_label_path, line_number + ) + if parts[9] not in {"0", "1", "2"}: + raise ValueError( + f"Unsupported DOTAv1 difficulty flag {parts[9]!r} at " + f"{original_label_path}:{line_number}." + ) + target_key = ( + cls, + _canonicalize_quadrilateral(coords.flatten().tolist()), + ) + if target_key in seen_targets: + raise ValueError( + "Duplicate DOTAv1 annotation target at " + f"{original_label_path}:{line_number}." + ) + seen_targets.add(target_key) + if parts[9] in {"1", "2"}: + ignore_classes.append(cls) + ignore_polygons.append(coords) + else: + classes.append(cls) + polygons.append(coords) + + else: + raise FileNotFoundError( + "DOTAv1 annotation not found for image " + f"{image_id!r}; expected {label_path} or {original_label_path}." + ) + + polygon_tensor = ( + torch.stack(polygons).to(torch.float32) + if polygons + else torch.zeros((0, 4, 2), dtype=torch.float32) + ) + if polygon_tensor.numel(): + boxes_xywhr = xyxyxyxy2xywhr(polygon_tensor) + boxes = cast(torch.Tensor, boxes_xywhr).to(torch.float32) + else: + boxes = torch.zeros((0, 5), dtype=torch.float32) + + ignore_polygon_tensor = ( + torch.stack(ignore_polygons).to(torch.float32) + if ignore_polygons + else torch.zeros((0, 4, 2), dtype=torch.float32) + ) + if ignore_polygon_tensor.numel(): + ignored_xywhr = xyxyxyxy2xywhr(ignore_polygon_tensor) + ignore_boxes = cast(torch.Tensor, ignored_xywhr).to(torch.float32) + else: + ignore_boxes = torch.zeros((0, 5), dtype=torch.float32) + + ground_truths[image_id] = { + "cls": torch.tensor(classes, dtype=torch.int64), + "polygons": polygon_tensor, + "bboxes": boxes, + "ignore_cls": torch.tensor(ignore_classes, dtype=torch.int64), + "ignore_polygons": ignore_polygon_tensor, + "ignore_bboxes": ignore_boxes, + } + return ground_truths + + +def format_dota_results( + nms_outs: Results, + input_shape: tuple[int, ...], + org_shape: list[tuple[int, int]], + ratio_pad: list[Any], + image_ids: tuple[str, ...], + postprocess: Any, +) -> list[dict[str, Any]]: + """Format model outputs for DOTAv1 evaluation and export. + + Args: + nms_outs: Postprocessed model results. + input_shape: Preprocessed image shape. + org_shape: Original image shapes. + ratio_pad: Letterbox metadata. + image_ids: DOTAv1 image IDs. + postprocess: Postprocessor instance. + + Returns: + List of formatted prediction dictionaries. + """ + labels_list, polygons_list, scores_list, xywhr_list = postprocess.nmsout2eval( + nms_outs.output, + input_shape, + org_shape, + ratio_pad=ratio_pad, + include_xywhr=True, + ) + outer_lengths = { + "image IDs": len(image_ids), + "labels": len(labels_list), + "polygons": len(polygons_list), + "scores": len(scores_list), + "rotated boxes": len(xywhr_list), + } + if len(set(outer_lengths.values())) != 1: + details = ", ".join( + f"{name}={length}" for name, length in outer_lengths.items() + ) + raise ValueError(f"DOTAv1 export batch length mismatch: {details}.") + results = [] + for image_id, labels, polygons, scores, xywhrs in zip( + image_ids, labels_list, polygons_list, scores_list, xywhr_list, strict=True + ): + detection_lengths = { + "labels": len(labels), + "polygons": len(polygons), + "scores": len(scores), + "rotated boxes": len(xywhrs), + } + if len(set(detection_lengths.values())) != 1: + details = ", ".join( + f"{name}={length}" for name, length in detection_lengths.items() + ) + raise ValueError( + f"DOTAv1 export detection length mismatch for {image_id}: {details}." + ) + for label, polygon, score, xywhr in zip( + labels, polygons, scores, xywhrs, strict=True + ): + results.append( + { + "image_id": image_id, + "category_id": DOTAV1_CLASS_TO_IDX[label], + "category_name": label, + "poly": polygon, + "score": score, + "rbox": xywhr, + } + ) + return results + + +def _compute_ap( + recall: np.ndarray, precision: np.ndarray +) -> tuple[float, np.ndarray, np.ndarray]: + """Compute AP from recall and precision curves with Ultralytics interpolation.""" + mrec = np.concatenate(([0.0], recall, [1.0])) + mpre = np.concatenate(([1.0], precision, [0.0])) + mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) + grid = np.linspace(0, 1, 101) + integrate = getattr(np, "trapezoid", None) + if integrate is None: + integrate = np.trapz + return float(integrate(np.interp(grid, mrec, mpre), grid)), mpre, mrec + + +def _ap_per_class( + tp: np.ndarray, + ignore: np.ndarray, + conf: np.ndarray, + pred_cls: np.ndarray, + target_cls: np.ndarray, + eps: float = 1e-16, +) -> np.ndarray: + """Compute AP per class using the Ultralytics object-detection metric policy.""" + if target_cls.size == 0: + niou = tp.shape[1] if tp.ndim == 2 else 10 + return np.zeros((0, niou), dtype=np.float64) + + order = np.argsort(-conf) + tp = tp[order] + ignore = ignore[order] + conf = conf[order] + pred_cls = pred_cls[order] + + unique_classes, target_count = np.unique(target_cls, return_counts=True) + ap = np.zeros((unique_classes.shape[0], tp.shape[1]), dtype=np.float64) + for class_index, class_id in enumerate(unique_classes): + pred_mask = pred_cls == class_id + num_labels = target_count[class_index] + if not pred_mask.any() or num_labels == 0: + continue + + class_tp = tp[pred_mask] + class_ignore = ignore[pred_mask] + for iou_index in range(tp.shape[1]): + keep = ~class_ignore[:, iou_index] + if not keep.any(): + continue + true_positive = class_tp[keep, iou_index].cumsum(0) + false_positive = (1 - class_tp[keep, iou_index]).cumsum(0) + recall = true_positive / (num_labels + eps) + precision = true_positive / (true_positive + false_positive + eps) + ap[class_index, iou_index], _, _ = _compute_ap(recall, precision) + return ap + + +def _match_predictions( + pred_classes: torch.Tensor, + true_classes: torch.Tensor, + iou: torch.Tensor, + iouv: torch.Tensor, +) -> np.ndarray: + """Match predictions to ground-truth boxes with Ultralytics one-to-one matching.""" + correct = np.zeros((pred_classes.shape[0], iouv.shape[0]), dtype=bool) + if pred_classes.numel() == 0 or true_classes.numel() == 0: + return correct + + correct_class = true_classes[:, None] == pred_classes + iou_np = (iou * correct_class).cpu().numpy() + for iou_index, threshold in enumerate(iouv.cpu().tolist()): + matches = np.array(np.nonzero(iou_np >= threshold)).T + if matches.shape[0] == 0: + continue + if matches.shape[0] > 1: + matches = matches[iou_np[matches[:, 0], matches[:, 1]].argsort()[::-1]] + # Preserve the IoU-ranked order while keeping the first match per prediction and target. + matches = matches[np.sort(np.unique(matches[:, 1], return_index=True)[1])] + matches = matches[np.sort(np.unique(matches[:, 0], return_index=True)[1])] + correct[matches[:, 1].astype(int), iou_index] = True + return correct + + +def _empty_stats() -> dict[str, list[np.ndarray]]: + """Create an empty DOTAv1 metric statistics accumulator.""" + return {"tp": [], "ignore": [], "conf": [], "pred_cls": [], "target_cls": []} + + +def _append_stats( + stats: dict[str, list[np.ndarray]], image_stats: dict[str, np.ndarray] +) -> None: + """Append one image's metric statistics to the accumulator.""" + for key, value in image_stats.items(): + stats[key].append(value) + + +def _nms_output_to_predictions(nms_out: torch.Tensor) -> dict[str, torch.Tensor]: + """Convert OBB NMS rows to the prediction dictionary used by metric matching.""" + if nms_out.numel() == 0: + return { + "bboxes": torch.zeros((0, 5), dtype=torch.float32), + "conf": torch.zeros(0, dtype=torch.float32), + "cls": torch.zeros(0, dtype=torch.int64), + } + + nms_out = nms_out.detach().cpu() + return { + "bboxes": torch.cat([nms_out[:, :4], nms_out[:, 6:7]], dim=-1).to( + torch.float32 + ), + "conf": nms_out[:, 4].to(torch.float32), + "cls": nms_out[:, 5].to(torch.int64), + } + + +def _ratio_pad_for_shape( + input_shape: tuple[int, ...], + org_shape: tuple[int, int], + ratio_pad: RatioPad | None, +) -> tuple[float, tuple[float, float]]: + """Return letterbox gain and padding for an image.""" + if len(input_shape) < 2: + raise ValueError(f"Expected at least 2 input dimensions, got {input_shape}.") + + ratio, pad = resolve_ratio_pad( + (input_shape[0], input_shape[1]), org_shape, ratio_pad + ) + return float(ratio[0]), (float(pad[0]), float(pad[1])) + + +def _ground_truth_to_input_space( + ground_truth: dict[str, torch.Tensor], + input_shape: tuple[int, ...], + org_shape: tuple[int, int], + ratio_pad: RatioPad | None, +) -> dict[str, torch.Tensor]: + """Transform original-image DOTAv1 polygons to letterboxed ``xywhr`` boxes.""" + gain, pad = _ratio_pad_for_shape(input_shape, org_shape, ratio_pad) + + def transform(polygons: torch.Tensor | None, boxes: torch.Tensor) -> torch.Tensor: + if polygons is None: + return boxes + if polygons.numel() == 0: + return torch.zeros((0, 5), dtype=torch.float32) + transformed = polygons.clone().to(torch.float32) + transformed[..., 0] = transformed[..., 0] * gain + pad[0] + transformed[..., 1] = transformed[..., 1] * gain + pad[1] + return cast(torch.Tensor, xyxyxyxy2xywhr(transformed)).to(torch.float32) + + transformed_boxes = transform(ground_truth.get("polygons"), ground_truth["bboxes"]) + ignore_classes = ground_truth.get("ignore_cls", torch.zeros(0, dtype=torch.int64)) + ignore_boxes = transform( + ground_truth.get("ignore_polygons"), + ground_truth.get("ignore_bboxes", torch.zeros((0, 5), dtype=torch.float32)), + ) + return { + "cls": ground_truth["cls"], + "bboxes": transformed_boxes, + "ignore_cls": ignore_classes, + "ignore_bboxes": ignore_boxes, + } + + +def _process_image_stats( + predictions: dict[str, torch.Tensor], + target: dict[str, torch.Tensor], + iouv: torch.Tensor, +) -> dict[str, np.ndarray]: + """Build one image's true-positive, confidence, class, and target arrays.""" + target_cls = target["cls"].cpu().numpy() + if target["cls"].numel() == 0 or predictions["cls"].numel() == 0: + true_positive = np.zeros( + (predictions["cls"].shape[0], iouv.numel()), dtype=bool + ) + else: + iou = batch_probiou(target["bboxes"], predictions["bboxes"]) + true_positive = _match_predictions(predictions["cls"], target["cls"], iou, iouv) + ignore = np.zeros_like(true_positive) + ignore_classes = target.get("ignore_cls", torch.zeros(0, dtype=torch.int64)) + ignore_boxes = target.get("ignore_bboxes", torch.zeros((0, 5), dtype=torch.float32)) + if ignore_classes.numel() and predictions["cls"].numel(): + ignored_iou = batch_probiou(ignore_boxes, predictions["bboxes"]) + same_class = ignore_classes[:, None] == predictions["cls"] + ignored_matches = ( + (ignored_iou[:, :, None] >= iouv[None, None, :]) & same_class[:, :, None] + ).any(dim=0) + ignore = ignored_matches.cpu().numpy() & ~true_positive + return { + "tp": true_positive, + "ignore": ignore, + "conf": predictions["conf"].cpu().numpy(), + "pred_cls": predictions["cls"].cpu().numpy(), + "target_cls": target_cls, + } + + +def _evaluate_stats(stats: dict[str, list[np.ndarray]], niou: int = 10) -> DOTAResult: + """Compute DOTAv1 metrics in legacy tuple order: mAP50, then mAP50-95.""" + target_cls = ( + np.concatenate(stats["target_cls"], 0) + if stats["target_cls"] + else np.zeros(0, dtype=np.float64) + ) + if target_cls.size == 0: + raise ValueError("DOTAv1 evaluation requires at least one non-ignored target.") + + tp = ( + np.concatenate(stats["tp"], 0) + if stats["tp"] + else np.zeros((0, niou), dtype=bool) + ) + ignore = ( + np.concatenate(stats["ignore"], 0) + if stats["ignore"] + else np.zeros((0, niou), dtype=bool) + ) + conf = ( + np.concatenate(stats["conf"], 0) + if stats["conf"] + else np.zeros(0, dtype=np.float64) + ) + pred_cls = ( + np.concatenate(stats["pred_cls"], 0) + if stats["pred_cls"] + else np.zeros(0, dtype=np.float64) + ) + ap = _ap_per_class(tp, ignore, conf, pred_cls, target_cls) + if ap.size == 0: + return DOTAResult(map5095=0.0, map50=0.0) + return DOTAResult(map5095=float(ap.mean()), map50=float(ap[:, 0].mean())) + + +def evaluate_dota_predictions( + ground_truths: dict[str, dict[str, torch.Tensor]], + predictions: list[dict[str, Any]], +) -> DOTAResult: + """Evaluate DOTAv1 predictions with local rotated mAP. + + Args: + ground_truths: Mapping of image IDs to class and OBB tensors. + predictions: Formatted prediction dictionaries. + + Returns: + Rotated mAP at IoU ``0.50`` followed by mAP averaged across ``0.50:0.95``. + The ``primary_score`` and ``secondary_score`` properties expose mAP50-95 + and mAP50, respectively. + """ + iouv = torch.linspace(0.5, 0.95, 10) + stats = _empty_stats() + predictions_by_image: dict[str, list[dict[str, Any]]] = defaultdict(list) + for prediction in predictions: + predictions_by_image[str(prediction["image_id"])].append(prediction) + + image_ids = set(ground_truths) | set(predictions_by_image) + for image_id in sorted(image_ids): + rows = predictions_by_image.get(image_id, []) + if rows: + pred_dict = { + "bboxes": torch.tensor( + [row["rbox"] for row in rows], dtype=torch.float32 + ), + "conf": torch.tensor( + [row["score"] for row in rows], dtype=torch.float32 + ), + "cls": torch.tensor( + [row["category_id"] for row in rows], dtype=torch.int64 + ), + } + else: + pred_dict = _nms_output_to_predictions( + torch.zeros((0, 7), dtype=torch.float32) + ) + target = ground_truths.get( + image_id, + { + "cls": torch.zeros(0, dtype=torch.int64), + "bboxes": torch.zeros((0, 5), dtype=torch.float32), + }, + ) + _append_stats(stats, _process_image_stats(pred_dict, target, iouv)) + return _evaluate_stats(stats, niou=iouv.numel()) + + +def save_dota_task1_predictions( + predictions: list[dict[str, Any]], save_dir: str +) -> tuple[Path, Path]: + """Save split and merged predictions in DOTA Task1 text format. + + Args: + predictions: Formatted prediction dictionaries. + save_dir: Directory where prediction folders are written. + + Returns: + Tuple of split and merged prediction directories. + """ + root = Path(save_dir) + pred_txt = root / "predictions_txt" + pred_merged_txt = root / "predictions_merged_txt" + pred_txt.mkdir(parents=True, exist_ok=True) + pred_merged_txt.mkdir(parents=True, exist_ok=True) + + for cls_idx in range(get_dotav1_class_num()): + (pred_txt / f"Task1_{get_dotav1_label(cls_idx)}.txt").write_text( + "", encoding="utf-8" + ) + (pred_merged_txt / f"Task1_{get_dotav1_label(cls_idx)}.txt").write_text( + "", encoding="utf-8" + ) + + for pred in predictions: + class_name = pred["category_name"] + polygon = pred["poly"] + with (pred_txt / f"Task1_{class_name}.txt").open("a", encoding="utf-8") as file: + file.write( + f"{pred['image_id']} {pred['score']} " + f"{polygon[0]} {polygon[1]} {polygon[2]} {polygon[3]} " + f"{polygon[4]} {polygon[5]} {polygon[6]} {polygon[7]}\n" + ) + + merged_results: dict[str, list[list[float]]] = defaultdict(list) + offset_pattern = re.compile(r"(\d+)___(\d+)") + for pred in predictions: + image_id = pred["image_id"] + base_image_id = image_id.split("__", 1)[0] + offset_match = offset_pattern.search(image_id) + x_offset, y_offset = (0, 0) + if offset_match is not None: + x_offset, y_offset = ( + int(offset_match.group(1)), + int(offset_match.group(2)), + ) + rbox = list(pred["rbox"]) + rbox[0] += x_offset + rbox[1] += y_offset + merged_results[base_image_id].append( + [*rbox, pred["score"], float(pred["category_id"])] + ) + + for image_id, rows in merged_results.items(): + bbox = torch.tensor(rows, dtype=torch.float32) + if bbox.numel() == 0: + continue + max_wh = max(float(torch.max(bbox[:, :2]).item() * 2), 1.0) + class_offsets = bbox[:, 6:7] * max_wh + boxes = bbox[:, :5].clone() + boxes[:, :2] += class_offsets + keep = rotated_nms(boxes, bbox[:, 5], 0.3) + bbox = bbox[keep] + polygons = xywhr2xyxyxyxy(bbox[:, :5]).reshape(-1, 8) + for polygon, score, cls in zip( + polygons.tolist(), bbox[:, 5].tolist(), bbox[:, 6].tolist() + ): + class_name = get_dotav1_label(int(cls)) + rounded_polygon = [round(float(value), 3) for value in polygon] + with (pred_merged_txt / f"Task1_{class_name}.txt").open( + "a", encoding="utf-8" + ) as file: + file.write( + f"{image_id} {round(float(score), 3)} " + f"{rounded_polygon[0]} {rounded_polygon[1]} {rounded_polygon[2]} {rounded_polygon[3]} " + f"{rounded_polygon[4]} {rounded_polygon[5]} {rounded_polygon[6]} {rounded_polygon[7]}\n" + ) + + return pred_txt, pred_merged_txt + + +def _nms_output_list(nms_output: Any) -> list[torch.Tensor]: + """Normalize postprocess output to a per-image list of OBB tensors.""" + if isinstance(nms_output, list): + return nms_output + if isinstance(nms_output, tuple): + return list(nms_output) + if isinstance(nms_output, torch.Tensor): + if nms_output.ndim == 3: + return [image[image[:, 4] > 0] for image in nms_output] + return [nms_output] + raise TypeError(f"Unsupported OBB NMS output type: {type(nms_output).__name__}.") + + +def _validate_evaluation_batch_lengths( + nms_outputs: list[torch.Tensor], + input_batch_size: int, + org_shape: Any, + ratio_pad: Any, + image_ids: Any, +) -> None: + """Reject batches whose output or loader metadata omits an image.""" + + batch_lengths = { + "model outputs": len(nms_outputs), + "input batch": input_batch_size, + "original shapes": len(org_shape), + "ratio pads": len(ratio_pad), + "image IDs": len(image_ids), + } + if len(set(batch_lengths.values())) != 1: + details = ", ".join( + f"{name}={length}" for name, length in batch_lengths.items() + ) + raise ValueError(f"DOTAv1 evaluation batch length mismatch: {details}.") + + +def eval_dota( + model: MBLT_Engine, + data_path: str, + batch_size: int, + conf_thres: float | None = None, + iou_thres: float | None = None, + save_dir: str | None = None, +) -> DOTAResult: + """Evaluate a model on DOTAv1 validation. + + Args: + model: Model engine to evaluate. + data_path: DOTAv1 dataset root. + batch_size: Batch size for evaluation. + conf_thres: Optional confidence threshold override. + iou_thres: Optional IoU threshold override. + save_dir: Optional directory for DOTA Task1 prediction files. + + Returns: + Local rotated mAP scores. + """ + if normalize_vision_task(model.post_cfg["task"]) != "obb": + raise NotImplementedError( + f"Task {model.post_cfg['task']} is not supported for DOTAv1 evaluation." + ) + dataset_name = model.post_cfg.get("dataset") + if not isinstance(dataset_name, str) or dataset_name.lower() != "dotav1": + raise ValueError( + "DOTAv1 evaluation requires model post_cfg.dataset to be 'dotav1', " + f"got {dataset_name!r}." + ) + + dataset = CustomDOTAv1(data_path) + dataloader = get_dota_loader(dataset, batch_size, model.preprocess_with_metadata) + model.set_postprocess_thresholds(conf_thres=conf_thres, iou_thres=iou_thres) + ground_truths = _load_ground_truths(data_path, dataset) + iouv = torch.linspace(0.5, 0.95, 10) + stats = _empty_stats() + + results = [] + num_data = len(dataset) + total_iter = math.ceil(num_data / batch_size) + pbar = tqdm(dataloader, total=total_iter, desc="Evaluating DOTAv1") + inference_time = 0.0 + cum_num_data = 0 + + for input_npu, org_shape, ratio_pad, image_ids in pbar: + cum_num_data += len(image_ids) + tic = time() + out_npu = model(input_npu) + inference_time += time() - tic + nms_outs = model.postprocess(out_npu) + input_shape = tuple(int(value) for value in input_npu.shape[1:-1]) + nms_outputs = _nms_output_list(nms_outs.output) + _validate_evaluation_batch_lengths( + nms_outputs, + int(input_npu.shape[0]), + org_shape, + ratio_pad, + image_ids, + ) + for nms_out, image_id, image_shape, image_ratio_pad in zip( + nms_outputs, + image_ids, + org_shape, + ratio_pad, + strict=True, + ): + target = _ground_truth_to_input_space( + ground_truths[image_id], + input_shape, + (int(image_shape[0]), int(image_shape[1])), + image_ratio_pad, + ) + _append_stats( + stats, + _process_image_stats(_nms_output_to_predictions(nms_out), target, iouv), + ) + if save_dir is not None: + results.extend( + format_dota_results( + nms_outs, + input_shape, + org_shape, + ratio_pad, + image_ids, + model.postprocessor, + ) + ) + pbar.set_postfix_str(f"NPU FPS: {cum_num_data / inference_time:.3f}") + + pbar.close() + map_score = _evaluate_stats(stats, niou=iouv.numel()) + if save_dir is not None: + save_dota_task1_predictions(results, save_dir) + print("DOTAv1 evaluation completed") + return map_score diff --git a/mblt_vision/utils/evaluation/eval_imagenet.py b/mblt_vision/utils/evaluation/eval_imagenet.py new file mode 100644 index 0000000..c0eb716 --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_imagenet.py @@ -0,0 +1,151 @@ +""" +Evaluation script for ImageNet dataset. +""" + +from __future__ import annotations + +import math +from time import time +from typing import TYPE_CHECKING, NamedTuple + +import numpy as np +import torch +from tqdm import tqdm + +from ..datasets import CustomImageFolder, get_imagenet_loader +from ..datasets.readiness import IMAGENET_SYNSET_ORDER, IMAGENET_SYNSETS + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + + +class ImageNetResult(NamedTuple): + """ImageNet metrics ordered from primary to secondary.""" + + top1: float + top5: float + + @property + def primary_score(self) -> float: + """Return the primary ImageNet validation metric.""" + return self.top1 + + @property + def secondary_score(self) -> float: + """Return the secondary ImageNet validation metric.""" + return self.top5 + + +def eval_imagenet_metrics( + model: MBLT_Engine, data_path: str, batch_size: int +) -> ImageNetResult: + """Evaluates a classification model on the ImageNet validation set. + + Computes Top-1 and Top-5 accuracy and inference speed (FPS) on the NPU. + + Args: + model (MBLT_Engine): The vision engine to evaluate. + data_path (str): Path to the ImageNet validation images. + batch_size (int): Number of images per inference batch. + + Returns: + ImageNetResult: Top-1 primary accuracy and Top-5 secondary accuracy. + """ + dataset_name = model.post_cfg.get("dataset") + if not isinstance(dataset_name, str) or dataset_name.lower() != "imagenet": + raise ValueError( + "ImageNet evaluation requires model post_cfg.dataset to be 'imagenet', " + f"got {dataset_name!r}." + ) + dataset = CustomImageFolder(data_path) + unknown_synsets = set(dataset.classes) - IMAGENET_SYNSETS + if unknown_synsets: + raise ValueError( + "ImageNet evaluation found non-canonical synset directories: " + f"{', '.join(sorted(unknown_synsets)[:5])}." + ) + dataset.class_to_idx = { + synset: index + for index, synset in enumerate(IMAGENET_SYNSET_ORDER) + if synset in dataset.class_to_idx + } + dataset.make_dataset() + num_data = len(dataset) + if num_data == 0: + raise ValueError( + f"ImageNet evaluation dataset contains no supported images: {data_path}." + ) + dataloader = get_imagenet_loader(dataset, batch_size, model.preprocess) + total_iter = math.ceil(num_data / batch_size) + pbar = tqdm(dataloader, total=total_iter, desc="Evaluating ImageNet") + inference_time = 0.0 + cum_num_data = 0 + cum_top1_correct = 0 + cum_top5_correct = 0 + top1_acc = 0.0 + top5_acc = 0.0 + for input_npu, label in pbar: + cum_num_data += len(label) + tic = time() + out_npu = model(input_npu) + inference_time += time() - tic + result = model.postprocess(out_npu) + output = result.output + label_array = np.asarray(label) + if not isinstance(output, (np.ndarray, torch.Tensor)): + raise TypeError( + f"Expected classification output to be a tensor or ndarray, got {type(output)}." + ) + if output.ndim != 2: + raise ValueError( + f"ImageNet classification output must have shape [B, C], got {tuple(output.shape)}." + ) + + output_batch_size = output.shape[0] + if output_batch_size != len(label_array): + raise ValueError( + "ImageNet classification output batch size does not match labels: " + f"got {output_batch_size} outputs for {len(label_array)} labels." + ) + if isinstance(output, np.ndarray): + prediction = output.argmax(-1) + else: + prediction = output.argmax(-1).cpu().numpy() + top_k = min(5, output.shape[-1]) + if isinstance(output, torch.Tensor): + top5_prediction = output.topk(top_k, dim=-1).indices.cpu().numpy() + else: + top5_prediction = np.argpartition(output, -top_k, axis=-1)[:, -top_k:] + cum_top1_correct += (prediction == label_array).sum().item() + cum_top5_correct += ( + np.any(top5_prediction == label_array[:, np.newaxis], axis=-1).sum().item() + ) + top1_acc = cum_top1_correct / cum_num_data + top5_acc = cum_top5_correct / cum_num_data + pbar.set_postfix_str( + f"Top 1 Acc.: {100 * top1_acc:.3f}%, Top 5 Acc.: {100 * top5_acc:.3f}%, " + f"NPU FPS: {cum_num_data / inference_time:.3f}" + ) + pbar.close() + print("ImageNet evaluation completed") + print( + f"Top 1 Acc.: {100 * top1_acc:.3f}%, " + f"Top 5 Acc.: {100 * top5_acc:.3f}%, " + f"NPU FPS: {cum_num_data / inference_time:.3f}" + ) + return ImageNetResult(top1=top1_acc, top5=top5_acc) + + +def eval_imagenet(model: MBLT_Engine, data_path: str, batch_size: int) -> float: + """Evaluate ImageNet and return Top-1 accuracy for numeric API compatibility. + + Args: + model: Vision engine to evaluate. + data_path: Path to the ImageNet validation images. + batch_size: Number of images per inference batch. + + Returns: + Top-1 accuracy in the range 0.0 to 1.0. + """ + + return eval_imagenet_metrics(model, data_path, batch_size).top1 diff --git a/mblt_vision/utils/evaluation/eval_nyu_depth.py b/mblt_vision/utils/evaluation/eval_nyu_depth.py new file mode 100644 index 0000000..6d0099a --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_nyu_depth.py @@ -0,0 +1,189 @@ +"""NYU Depth V2 evaluation for monocular depth-estimation models.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from mblt_vision.utils.preprocess import build_preprocess +from tqdm import tqdm + +from ..datasets import CustomNYUDepth, get_nyu_depth_loader + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + + +@dataclass(frozen=True) +class NYUDepthResult: + """Median-aligned NYU Depth V2 metrics.""" + + delta1: float + abs_rel: float + rmse: float + + @property + def primary_score(self) -> float: + """Return the primary NYU Depth validation metric.""" + + return self.delta1 + + @property + def secondary_score(self) -> float: + """Return the secondary NYU Depth validation metric.""" + + return self.abs_rel + + +class NYUDepthMetricAccumulator: + """Accumulate median-aligned metrics over every valid NYU depth pixel.""" + + MIN_DEPTH = 0.001 + MAX_DEPTH = 100.0 + + def __init__(self) -> None: + """Initialize zero-valued pixel sums.""" + + self.delta1_sum = 0.0 + self.abs_rel_sum = 0.0 + self.squared_error_sum = 0.0 + self.valid_pixel_count = 0 + + def update(self, prediction: np.ndarray, target: np.ndarray) -> None: + """Median-align one prediction and add its valid-pixel statistics.""" + + prediction = _as_real_float32(prediction, "prediction") + target = _as_real_float32(target, "target") + if prediction.shape != target.shape: + raise ValueError( + f"NYU Depth prediction and target shapes must match, got {prediction.shape} and {target.shape}." + ) + if not np.isfinite(target).all(): + raise ValueError("NYU Depth target contains non-finite values.") + if (target < 0).any(): + raise ValueError("NYU Depth target contains negative values.") + valid = ( + np.isfinite(target) & (target > self.MIN_DEPTH) & (target < self.MAX_DEPTH) + ) + valid_pixel_count = int(valid.sum()) + if valid_pixel_count == 0: + raise ValueError( + "NYU Depth sample has no valid pixels in the (0.001, 100.0) range." + ) + + predicted, actual = prediction[valid], target[valid] + invalid_prediction_count = int((~np.isfinite(predicted)).sum()) + if invalid_prediction_count: + raise ValueError( + f"NYU Depth prediction contains {invalid_prediction_count} non-finite value(s) at valid target pixels." + ) + median_prediction = np.median(np.maximum(predicted, self.MIN_DEPTH)) + median_target = np.median(actual) + aligned = predicted * (median_target / median_prediction) + aligned = np.clip(aligned, self.MIN_DEPTH, self.MAX_DEPTH) + ratio = np.maximum(actual / aligned, aligned / actual) + self.delta1_sum += float(np.sum(ratio < 1.25)) + self.abs_rel_sum += float(np.sum(np.abs(actual - aligned) / actual)) + self.squared_error_sum += float(np.sum((actual - aligned) ** 2)) + self.valid_pixel_count += valid_pixel_count + + def result(self) -> NYUDepthResult: + """Return metrics pooled over all accumulated valid pixels.""" + + if self.valid_pixel_count == 0: + raise ValueError("NYU Depth evaluation received no valid pixels.") + return NYUDepthResult( + delta1=self.delta1_sum / self.valid_pixel_count, + abs_rel=self.abs_rel_sum / self.valid_pixel_count, + rmse=float(np.sqrt(self.squared_error_sum / self.valid_pixel_count)), + ) + + +def _as_real_float32(values: np.ndarray, name: str) -> np.ndarray: + """Validate metric input dtype before converting it to float32.""" + + array = np.asarray(values) + if not np.issubdtype(array.dtype, np.number) or np.issubdtype( + array.dtype, np.complexfloating + ): + raise ValueError( + f"NYU Depth {name} must use a real numeric dtype, got {array.dtype}." + ) + return np.asarray(array, dtype=np.float32) + + +def calculate_nyu_depth_metrics( + prediction: np.ndarray, target: np.ndarray +) -> NYUDepthResult: + """Calculate official pooled metrics for one median-aligned NYU sample.""" + + accumulator = NYUDepthMetricAccumulator() + accumulator.update(prediction, target) + return accumulator.result() + + +def eval_nyu_depth( + model: MBLT_Engine, data_path: str, batch_size: int +) -> NYUDepthResult: + """Evaluate a depth model on paired NYU validation images and depth maps.""" + + dataset_name = model.post_cfg.get("dataset") + if not isinstance(dataset_name, str) or dataset_name.lower() != "nyu-depth": + raise ValueError( + "NYU Depth evaluation requires model post_cfg.dataset to be 'nyu-depth', " + f"got {dataset_name!r}." + ) + dataset = CustomNYUDepth(data_path) + letterbox_cfg = model.pre_cfg.get("LetterBox") + if not isinstance(letterbox_cfg, dict) or "img_size" not in letterbox_cfg: + raise ValueError( + "NYU Depth validation requires a LetterBox img_size in the model preprocessing config." + ) + image_size = letterbox_cfg["img_size"] + if not isinstance(image_size, list) or len(image_size) != 2: + raise ValueError( + "NYU Depth validation img_size must be a two-item [height, width] list." + ) + + validation_pre_cfg = { + name: config for name, config in model.pre_cfg.items() if name != "LetterBox" + } + validation_preprocessor = build_preprocess(validation_pre_cfg) + loader = get_nyu_depth_loader( + dataset, + batch_size, + validation_preprocessor, + image_size=(int(image_size[0]), int(image_size[1])), + ) + accumulator = NYUDepthMetricAccumulator() + for inputs, targets, _, _, _ in tqdm(loader, desc="Evaluating NYU Depth"): + output = model(inputs) + result = model.postprocess(output) + depth = result.depth + if depth is None: + raise ValueError("Depth postprocessor returned no depth maps.") + if isinstance(depth, list): + maps = depth + elif len(targets) == 1 and depth.ndim == 2: + maps = [depth] + else: + if depth.ndim < 3 or depth.shape[0] != len(targets): + raise ValueError( + "Depth postprocessor output batch length mismatch: " + f"maps={depth.shape[0] if depth.ndim else 0}, " + f"targets={len(targets)}." + ) + maps = [depth[index] for index in range(len(targets))] + if len(maps) != len(targets): + raise ValueError( + f"Depth postprocessor returned {len(maps)} maps for {len(targets)} targets." + ) + for prediction, target in zip(maps, targets): + array = ( + prediction.detach().cpu().numpy() + if hasattr(prediction, "detach") + else np.asarray(prediction) + ) + accumulator.update(array, target) + return accumulator.result() diff --git a/mblt_vision/utils/evaluation/eval_widerface.py b/mblt_vision/utils/evaluation/eval_widerface.py new file mode 100644 index 0000000..148e332 --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_widerface.py @@ -0,0 +1,440 @@ +"""Evaluation script for WiderFace face detection.""" + +from __future__ import annotations + +import math +import os +from pathlib import Path +from time import time +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +import numpy as np +from mblt_vision.utils.postprocess.base import YOLODetectionPostBase +from scipy.io import loadmat +from tqdm import tqdm + +from ..datasets import CustomWiderFaceDataset, get_widerface_loader +from ..datasets.readiness import ( + _load_widerface_image_names, + _widerface_image_shapes, + _widerface_difficulty_metadata_ready, +) + +if TYPE_CHECKING: + from ...wrapper import MBLT_Engine + +CustomWiderface = CustomWiderFaceDataset + + +class WiderFaceResult(NamedTuple): + """WiderFace AP metrics.""" + + easy_ap: float + medium_ap: float + hard_ap: float + + @property + def mean_ap(self) -> float: + """Return the mean AP across Easy, Medium, and Hard.""" + + return (self.easy_ap + self.medium_ap + self.hard_ap) / 3.0 + + @property + def primary_score(self) -> float: + """Return mean AP across the three validation splits.""" + + return self.mean_ap + + @property + def secondary_score(self) -> float: + """Return Hard-set AP.""" + + return self.hard_ap + + +def _empty_prediction() -> np.ndarray: + """Return an empty WiderFace prediction array.""" + + return np.zeros((0, 5), dtype=np.float32) + + +def _initialize_predictions( + dataset: CustomWiderFaceDataset, +) -> dict[str, dict[str, np.ndarray]]: + """Initialize empty predictions for every WiderFace sample.""" + + predictions: dict[str, dict[str, np.ndarray]] = {} + for _, event_name, file_name in dataset.samples: + predictions.setdefault(event_name, {})[os.path.splitext(file_name)[0]] = ( + _empty_prediction() + ) + return predictions + + +def _boxes_scores_to_prediction( + boxes: list[list[float]], scores: list[float] +) -> np.ndarray: + """Convert xywh boxes and scores to a WiderFace prediction array.""" + + if len(boxes) != len(scores): + raise ValueError( + "WiderFace postprocess returned unequal box and score counts: " + f"boxes={len(boxes)}, scores={len(scores)}." + ) + if not boxes: + return _empty_prediction() + prediction = np.zeros((len(boxes), 5), dtype=np.float32) + for index, (box, score) in enumerate(zip(boxes, scores)): + prediction[index, :4] = np.asarray(box, dtype=np.float32) + prediction[index, 4] = float(score) + return prediction + + +def eval_widerface( + model: MBLT_Engine, + data_path: str, + batch_size: int, + conf_thres: float | None = None, + iou_thres: float | None = None, +) -> WiderFaceResult: + """Evaluate a face-detection model on WiderFace validation data. + + Args: + model: The face-detection engine to evaluate. + data_path: Organized WiderFace dataset root. + batch_size: Validation batch size. + conf_thres: Optional confidence threshold override. + iou_thres: Optional IoU threshold override. + + Returns: + WiderFace Easy, Medium, and Hard AP metrics. + """ + + if model.post_cfg["task"] != "face_detection": + raise NotImplementedError( + f"Task {model.post_cfg['task']} is not supported for WiderFace evaluation." + ) + dataset_name = model.post_cfg.get("dataset") + if not isinstance(dataset_name, str) or dataset_name.lower() != "widerface": + raise ValueError( + "WiderFace evaluation requires model post_cfg.dataset to be 'widerface', " + f"got {dataset_name!r}." + ) + + dataset_root = Path(data_path) + expected_images = _load_widerface_image_names(dataset_root / "wider_face_val.mat") + image_shapes = ( + _widerface_image_shapes(dataset_root, expected_images) + if expected_images is not None + else None + ) + if ( + expected_images is None + or image_shapes is None + or not _widerface_difficulty_metadata_ready( + dataset_root, expected_images, image_shapes=image_shapes + ) + ): + raise ValueError( + "WiderFace evaluation metadata is malformed or has inconsistent difficulty indices." + ) + + dataset = CustomWiderface(os.path.join(data_path, "images")) + actual_images: dict[str, set[str]] = {} + for _, event_name, file_name in dataset.samples: + actual_images.setdefault(event_name, set()).add(file_name) + if actual_images != expected_images or sum( + len(file_names) for file_names in actual_images.values() + ) != len(dataset.samples): + raise ValueError( + "WiderFace image tree does not match the validation metadata identities." + ) + dataloader = get_widerface_loader( + dataset, batch_size, model.preprocess_with_metadata + ) + model.set_postprocess_thresholds(conf_thres=conf_thres, iou_thres=iou_thres) + + predictions = _initialize_predictions(dataset) + num_data = len(dataset) + total_iter = math.ceil(num_data / batch_size) + pbar = tqdm(dataloader, total=total_iter, desc="Evaluating WiderFace") + inference_time = 0.0 + cum_num_data = 0 + + for input_npu, org_shape, ratio_pad, target_classes, fnames in pbar: + cum_num_data += len(fnames) + tic = time() + out_npu = model(input_npu) + inference_time += time() - tic + nms_outs = model.postprocess(out_npu) + input_shape = (int(input_npu.shape[1]), int(input_npu.shape[2])) + img0_shapes = [(int(shape[0]), int(shape[1])) for shape in org_shape.tolist()] + postprocessor = cast(YOLODetectionPostBase, model.postprocessor) + _, boxes_list, scores_list = postprocessor.nmsout2eval( + nms_outs.output, + input_shape, + img0_shapes, + ratio_pad=ratio_pad, + ) + batch_lengths = { + "input batch": int(input_npu.shape[0]), + "original shapes": len(org_shape), + "ratio pads": len(ratio_pad), + "target classes": len(target_classes), + "file names": len(fnames), + "boxes": len(boxes_list), + "scores": len(scores_list), + } + if len(set(batch_lengths.values())) != 1: + details = ", ".join( + f"{name}={length}" for name, length in batch_lengths.items() + ) + raise ValueError(f"WiderFace evaluation batch length mismatch: {details}.") + + for event_name, file_name, boxes, scores in zip( + target_classes, fnames, boxes_list, scores_list, strict=True + ): + predictions[event_name][os.path.splitext(file_name)[0]] = ( + _boxes_scores_to_prediction(boxes, scores) + ) + + pbar.set_postfix_str(f"NPU FPS: {cum_num_data / inference_time:.3f}") + + pbar.close() + aps = evaluation(norm_score(predictions), data_path) + print("WiderFace evaluation completed") + return WiderFaceResult(*aps) + + +def bbox_overlaps(boxes: np.ndarray, query_boxes: np.ndarray) -> np.ndarray: + """Compute pairwise IoU overlaps between boxes and query boxes.""" + + boxes = boxes.astype(np.float32) + query_boxes = query_boxes.astype(np.float32) + + boxes = boxes[:, None, :] + query_boxes = query_boxes[None, :, :] + + iw = ( + np.minimum(boxes[..., 2], query_boxes[..., 2]) + - np.maximum(boxes[..., 0], query_boxes[..., 0]) + + 1 + ) + ih = ( + np.minimum(boxes[..., 3], query_boxes[..., 3]) + - np.maximum(boxes[..., 1], query_boxes[..., 1]) + + 1 + ) + iw = np.maximum(iw, 0) + ih = np.maximum(ih, 0) + inter = iw * ih + + box_area = (boxes[..., 2] - boxes[..., 0] + 1) * (boxes[..., 3] - boxes[..., 1] + 1) + query_area = (query_boxes[..., 2] - query_boxes[..., 0] + 1) * ( + query_boxes[..., 3] - query_boxes[..., 1] + 1 + ) + union = box_area + query_area - inter + + return inter / union + + +def get_gt_boxes(gt_dir: str) -> tuple[Any, ...]: + """Load WiderFace evaluation `.mat` files from the organized dataset.""" + + gt_mat = loadmat(os.path.join(gt_dir, "wider_face_val.mat")) + hard_mat = loadmat(os.path.join(gt_dir, "wider_hard_val.mat")) + medium_mat = loadmat(os.path.join(gt_dir, "wider_medium_val.mat")) + easy_mat = loadmat(os.path.join(gt_dir, "wider_easy_val.mat")) + + facebox_list = gt_mat["face_bbx_list"] + event_list = gt_mat["event_list"] + file_list = gt_mat["file_list"] + hard_gt_list = hard_mat["gt_list"] + medium_gt_list = medium_mat["gt_list"] + easy_gt_list = easy_mat["gt_list"] + + return ( + facebox_list, + event_list, + file_list, + hard_gt_list, + medium_gt_list, + easy_gt_list, + ) + + +def norm_score(pred: dict[str, Any]) -> dict[str, Any]: + """Normalize WiderFace prediction scores to ``[0, 1]``.""" + + max_score = -1e9 + min_score = 1e9 + found = False + + for _, event_predictions in pred.items(): + for _, image_predictions in event_predictions.items(): + if len(image_predictions) == 0: + continue + found = True + _min = float(np.min(image_predictions[:, -1])) + _max = float(np.max(image_predictions[:, -1])) + if _max > max_score: + max_score = _max + if _min < min_score: + min_score = _min + + if not found: + return pred + + diff = max_score - min_score + if diff <= 0: + return pred + + for _, event_predictions in pred.items(): + for _, image_predictions in event_predictions.items(): + if len(image_predictions) == 0: + continue + image_predictions[:, -1] = (image_predictions[:, -1] - min_score) / diff + + return pred + + +def image_eval( + pred: np.ndarray, gt: np.ndarray, ignore: np.ndarray, iou_thresh: float +) -> tuple[np.ndarray, np.ndarray]: + """Evaluate one image worth of WiderFace predictions.""" + + _pred = pred.copy() + _gt = gt.copy() + pred_recall = np.zeros(_pred.shape[0]) + recall_list = np.zeros(_gt.shape[0]) + proposal_list = np.ones(_pred.shape[0]) + + _pred[:, 2] = _pred[:, 2] + _pred[:, 0] + _pred[:, 3] = _pred[:, 3] + _pred[:, 1] + _gt[:, 2] = _gt[:, 2] + _gt[:, 0] + _gt[:, 3] = _gt[:, 3] + _gt[:, 1] + + overlaps = bbox_overlaps(_pred[:, :4], _gt) + + for prediction_index in range(_pred.shape[0]): + gt_overlap = overlaps[prediction_index] + max_overlap = np.max(gt_overlap) + max_idx = np.argmax(gt_overlap) + if max_overlap >= iou_thresh: + if ignore[max_idx] == 0: + recall_list[max_idx] = -1 + proposal_list[prediction_index] = -1 + elif recall_list[max_idx] == 0: + recall_list[max_idx] = 1 + + r_keep_index = np.where(recall_list == 1)[0] + pred_recall[prediction_index] = len(r_keep_index) + return pred_recall, proposal_list + + +def img_pr_info( + thresh_num: int, + pred_info: np.ndarray, + proposal_list: np.ndarray, + pred_recall: np.ndarray, +) -> np.ndarray: + """Compute precision and recall contributions for one image.""" + + pr_info = np.zeros((thresh_num, 2), dtype=np.float32) + for threshold_index in range(thresh_num): + thresh = 1 - (threshold_index + 1) / thresh_num + recall_index = np.where(pred_info[:, 4] >= thresh)[0] + if len(recall_index) == 0: + pr_info[threshold_index, 0] = 0 + pr_info[threshold_index, 1] = 0 + else: + last_index = recall_index[-1] + proposal_index = np.where(proposal_list[: last_index + 1] == 1)[0] + pr_info[threshold_index, 0] = len(proposal_index) + pr_info[threshold_index, 1] = pred_recall[last_index] + return pr_info + + +def dataset_pr_info( + thresh_num: int, pr_curve: np.ndarray, count_face: int +) -> np.ndarray: + """Normalize a WiderFace precision-recall accumulator.""" + + _pr_curve = np.zeros((thresh_num, 2), dtype=np.float32) + for threshold_index in range(thresh_num): + proposals = pr_curve[threshold_index, 0] + matched = pr_curve[threshold_index, 1] + _pr_curve[threshold_index, 0] = matched / proposals if proposals > 0 else 0.0 + _pr_curve[threshold_index, 1] = matched / count_face if count_face > 0 else 0.0 + return _pr_curve + + +def voc_ap(rec: np.ndarray, prec: np.ndarray) -> float: + """Compute VOC-style average precision.""" + + mrec = np.concatenate((np.array([0.0]), rec, np.array([1.0]))) + mpre = np.concatenate((np.array([0.0]), prec, np.array([0.0]))) + + for index in range(mpre.size - 1, 0, -1): + mpre[index - 1] = np.maximum(mpre[index - 1], mpre[index]) + + recall_change_index = np.where(mrec[1:] != mrec[:-1])[0] + ap = np.sum( + (mrec[recall_change_index + 1] - mrec[recall_change_index]) + * mpre[recall_change_index + 1] + ) + return float(ap) + + +def evaluation( + pred: dict[str, Any], gt_path: str, iou_thresh: float = 0.5 +) -> list[float]: + """Evaluate WiderFace predictions against Easy, Medium, and Hard settings.""" + + facebox_list, event_list, file_list, hard_gt_list, medium_gt_list, easy_gt_list = ( + get_gt_boxes(gt_path) + ) + event_num = len(event_list) + thresh_num = 1000 + settings = ["easy", "medium", "hard"] + setting_gts = [easy_gt_list, medium_gt_list, hard_gt_list] + aps = [] + for setting_id in range(3): + gt_list = setting_gts[setting_id] + count_face = 0 + pr_curve = np.zeros((thresh_num, 2), dtype=np.float32) + pbar = tqdm(range(event_num)) + for event_index in pbar: + pbar.set_description(f"Processing {settings[setting_id]}") + event_name = str(event_list[event_index][0][0]) + img_list = file_list[event_index][0] + pred_list = pred[event_name] + sub_gt_list = gt_list[event_index][0] + gt_bbx_list = facebox_list[event_index][0] + for image_index, img_info in enumerate(img_list): + pred_info = pred_list[str(img_info[0][0])] + gt_boxes = np.array(gt_bbx_list[image_index][0], dtype=np.float32) + keep_index = np.array(sub_gt_list[image_index][0], dtype=np.int64) + count_face += len(keep_index) + + if len(gt_boxes) == 0 or len(pred_info) == 0: + continue + ignore = np.zeros(gt_boxes.shape[0]) + if len(keep_index) != 0: + ignore[keep_index - 1] = 1 + pred_recall, proposal_list = image_eval( + pred_info, gt_boxes, ignore, iou_thresh + ) + pr_curve += img_pr_info( + thresh_num, pred_info, proposal_list, pred_recall + ) + pbar.close() + pr_curve = dataset_pr_info(thresh_num, pr_curve, count_face) + aps.append(voc_ap(pr_curve[:, 1], pr_curve[:, 0])) + + print("==================== Results ====================") + print(f"Easy Val AP: {aps[0]}") + print(f"Medium Val AP: {aps[1]}") + print(f"Hard Val AP: {aps[2]}") + print("=================================================") + return aps diff --git a/mblt_vision/utils/letterbox.py b/mblt_vision/utils/letterbox.py new file mode 100644 index 0000000..b2ac02c --- /dev/null +++ b/mblt_vision/utils/letterbox.py @@ -0,0 +1,118 @@ +"""Shared forward and inverse geometry for aspect-preserving letterboxing.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypeAlias + +RatioPad: TypeAlias = tuple[tuple[float, float], tuple[float, float]] + + +@dataclass(frozen=True) +class LetterBoxGeometry: + """Geometry shared by letterbox preprocessing and output restoration.""" + + input_shape: tuple[int, int] + original_shape: tuple[int, int] + ratio: float + resized_shape: tuple[int, int] + pad: tuple[int, int] + + @classmethod + def from_shapes( + cls, + input_shape: tuple[int, int], + original_shape: tuple[int, int], + ) -> LetterBoxGeometry: + """Calculate YOLO-style centered letterbox geometry. + + Args: + input_shape: Target shape as ``(height, width)``. + original_shape: Source shape as ``(height, width)``. + + Returns: + Calculated resize ratio, resized shape, and top-left padding. + """ + + input_height, input_width = input_shape + original_height, original_width = original_shape + ratio = min(input_height / original_height, input_width / original_width) + resized_height = int(round(original_height * ratio)) + resized_width = int(round(original_width * ratio)) + left = int(round((input_width - resized_width) / 2 - 0.1)) + top = int(round((input_height - resized_height) / 2 - 0.1)) + return cls( + input_shape=input_shape, + original_shape=original_shape, + ratio=ratio, + resized_shape=(resized_height, resized_width), + pad=(left, top), + ) + + @property + def ratio_pad(self) -> RatioPad: + """Return metadata consumed by inverse letterbox operations.""" + + return ((self.ratio, self.ratio), self.pad) + + @property + def borders(self) -> tuple[int, int, int, int]: + """Return OpenCV border widths as ``(top, bottom, left, right)``.""" + + input_height, input_width = self.input_shape + resized_height, resized_width = self.resized_shape + left, top = self.pad + return ( + top, + input_height - resized_height - top, + left, + input_width - resized_width - left, + ) + + def crop_bounds( + self, + output_shape: tuple[int, int], + pad: tuple[float, float] | None = None, + ) -> tuple[int, int, int, int]: + """Scale inverse-letterbox crop bounds to a dense output shape. + + Args: + output_shape: Dense output shape as ``(height, width)``. + pad: Optional exact top-left padding metadata as ``(x, y)``. + + Returns: + Crop bounds as ``(top, bottom, left, right)``. + """ + + output_height, output_width = output_shape + input_height, input_width = self.input_shape + scale_x = output_width / input_width + scale_y = output_height / input_height + pad_x, pad_y = self.pad if pad is None else pad + left = int(round(pad_x * scale_x)) + top = int(round(pad_y * scale_y)) + resized_height, resized_width = self.resized_shape + right = left + int(round(resized_width * scale_x)) + bottom = top + int(round(resized_height * scale_y)) + return top, bottom, left, right + + +def resolve_ratio_pad( + input_shape: tuple[int, int], + original_shape: tuple[int, int], + ratio_pad: RatioPad | None = None, +) -> RatioPad: + """Return supplied letterbox metadata or derive it from image shapes. + + Args: + input_shape: Letterboxed shape as ``(height, width)``. + original_shape: Source shape as ``(height, width)``. + ratio_pad: Optional metadata recorded during preprocessing. + + Returns: + Resize ratios and top-left padding as ``((ratio_x, ratio_y), (pad_x, pad_y))``. + """ + + if ratio_pad is not None: + return ratio_pad + return LetterBoxGeometry.from_shapes(input_shape, original_shape).ratio_pad diff --git a/mblt_vision/utils/postprocess/__init__.py b/mblt_vision/utils/postprocess/__init__.py new file mode 100644 index 0000000..e5090d0 --- /dev/null +++ b/mblt_vision/utils/postprocess/__init__.py @@ -0,0 +1,9 @@ +""" +Postprocessing utilities for vision models. +""" + +from .build_post import build_postprocess +from .depth_post import DepthPost +from .semantic_seg_post import SemanticSegPost + +__all__ = ["DepthPost", "SemanticSegPost", "build_postprocess"] diff --git a/mblt_vision/utils/postprocess/_letterbox.py b/mblt_vision/utils/postprocess/_letterbox.py new file mode 100644 index 0000000..eaaddaf --- /dev/null +++ b/mblt_vision/utils/postprocess/_letterbox.py @@ -0,0 +1,102 @@ +"""Private helpers shared by dense prediction postprocessors.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +from ..letterbox import LetterBoxGeometry, RatioPad, resolve_ratio_pad +from .common import normalize_ratio_pads + + +def get_letterbox_input_shape( + pre_cfg: dict[str, Any], + requirement_name: str, + size_name: str | None = None, +) -> tuple[int, int]: + """Validate and return a dense task's configured letterbox input shape. + + Args: + pre_cfg: Model preprocessing configuration. + requirement_name: Task name used when LetterBox is absent. + size_name: Optional shorter name used for invalid-size errors. + + Returns: + Configured input height and width. + + Raises: + ValueError: If LetterBox or its two-item image size is missing or invalid. + """ + + letterbox_cfg = pre_cfg.get("LetterBox") + if not isinstance(letterbox_cfg, dict) or "img_size" not in letterbox_cfg: + raise ValueError( + f"{requirement_name} requires a LetterBox configuration in pre_cfg." + ) + image_size = letterbox_cfg["img_size"] + if not isinstance(image_size, list) or len(image_size) != 2: + raise ValueError( + f"{size_name or requirement_name} LetterBox img_size must be a two-item [height, width] list." + ) + return int(image_size[0]), int(image_size[1]) + + +def resolve_ratio_pads( + ratio_pad: RatioPad | Sequence[RatioPad | None] | None, + batch_size: int, + shapes: Sequence[tuple[int, int]], + input_shape: tuple[int, int], +) -> list[RatioPad]: + """Normalize letterbox metadata and derive values missing from a dense task batch. + + Args: + ratio_pad: Shared or per-image letterbox metadata. + batch_size: Number of images in the output batch. + shapes: Original image shapes. + input_shape: Configured model input shape. + + Returns: + One resolved ratio/padding pair per batch item. + + Raises: + ValueError: If ratio/padding metadata is invalid for the batch. + """ + + pads = normalize_ratio_pads(ratio_pad, batch_size) + return [ + resolve_ratio_pad(input_shape, shape, pad) for pad, shape in zip(pads, shapes) + ] + + +def crop_letterbox( + output: torch.Tensor, + shape: tuple[int, int], + ratio_pad: RatioPad, + input_shape: tuple[int, int], + task_name: str, +) -> torch.Tensor: + """Crop letterbox padding from a dense two-dimensional output. + + Args: + output: Dense two-dimensional model output. + shape: Original image height and width. + ratio_pad: Resize ratio and padding applied during preprocessing. + input_shape: Configured model input height and width. + task_name: Task label used in validation errors. + + Returns: + Output with letterbox padding removed. + + Raises: + ValueError: If inverse letterboxing produces an empty crop. + """ + + geometry = LetterBoxGeometry.from_shapes(input_shape, shape) + output_shape = (int(output.shape[0]), int(output.shape[1])) + top, bottom, left, right = geometry.crop_bounds(output_shape, pad=ratio_pad[1]) + cropped = output[top:bottom, left:right] + if cropped.numel() == 0: + raise ValueError(f"{task_name} letterbox restoration produced an empty crop.") + return cropped diff --git a/mblt_vision/utils/postprocess/base.py b/mblt_vision/utils/postprocess/base.py new file mode 100644 index 0000000..4671e15 --- /dev/null +++ b/mblt_vision/utils/postprocess/base.py @@ -0,0 +1,743 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Sequence +from typing import Any, cast + +import numpy as np +import torch + +from ..._tasks import normalize_vision_task +from ..letterbox import RatioPad +from ..preprocess._validation import normalize_image_size +from ..types import ListTensorLike, TensorLike +from .common import nmsout2eval, process_mask_upsample + + +class PostBase(ABC): + """Abstract base class for postprocessing.""" + + def __init__(self) -> None: + """Initialize PostBase.""" + super().__init__() + self.device = torch.device("cpu") + + @abstractmethod + def __call__( + self, x: TensorLike | ListTensorLike, *args: Any, **kwargs: Any + ) -> Any: + """Executes postprocessing on the model output. + + Args: + x (TensorLike | ListTensorLike): Input tensor or list of tensors from the model. + *args (Any): Additional positional arguments depending on the specific task. + **kwargs (Any): Additional keyword arguments depending on the specific task. + + Returns: + Any: Postprocessed results, format depends on the specific task. + """ + pass + + def to(self, device: str | torch.device) -> None: + """Move the operations to the specified device. + Args: + device (str | torch.device): Device to move the operations to. + """ + if isinstance(device, str): + self.device = torch.device(device) + elif isinstance(device, torch.device): + self.device = device + else: + raise TypeError(f"Got unexpected type for device={type(device)}.") + for name, value in self.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(self, name, value.to(self.device)) + + +class YOLODetectionPostBase(PostBase): + """Base class for YOLO postprocessing.""" + + NC_BY_DATASET_TASK: dict[tuple[str, str], int] = { + ("coco", "object_detection"): 80, + ("coco", "instance_segmentation"): 80, + ("coco", "pose_estimation"): 1, + ("dotav1", "obb"): 15, + ("widerface", "face_detection"): 1, + } + DEFAULT_NC_BY_TASK: dict[str, int] = { + "object_detection": 80, + "instance_segmentation": 80, + "pose_estimation": 1, + "obb": 15, + "face_detection": 1, + } + + def __init__( + self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any], **kwargs + ) -> None: + """Initialize the common YOLO detection postprocessor. + + Args: + pre_cfg (dict): Preprocessing configuration. + post_cfg (dict): Postprocessing configuration. + **kwargs: Optional runtime overrides for postprocess behavior. + + Raises: + TypeError: If unsupported keyword overrides are provided. + """ + super().__init__() + letterbox_cfg = pre_cfg.get("LetterBox") + if letterbox_cfg is None: + raise ValueError("LetterBox configuration should be provided in pre_cfg") + img_size = letterbox_cfg["img_size"] + self.imh: int + self.imw: int + self.imh, self.imw = normalize_image_size( + img_size, name="pre_cfg.LetterBox.img_size" + ) + task = post_cfg.get("task") + if task is None: + raise ValueError("task should be provided in post_cfg") + self.task = normalize_vision_task(task) + task_key = self.task + dataset = post_cfg.get("dataset") + self.dataset = dataset.lower() if isinstance(dataset, str) else None + dataset_nc = ( + self.NC_BY_DATASET_TASK.get((self.dataset, task_key)) + if self.dataset is not None + else None + ) + configured_nc = kwargs.pop("nc", post_cfg.get("nc")) + if ( + configured_nc is not None + and dataset_nc is not None + and int(configured_nc) != dataset_nc + ): + raise ValueError( + f"nc={configured_nc} conflicts with dataset '{self.dataset}' and task '{self.task}', " + f"which require nc={dataset_nc}." + ) + default_nc = ( + dataset_nc + if dataset_nc is not None + else self.DEFAULT_NC_BY_TASK.get(task_key) + ) + nc = configured_nc if configured_nc is not None else default_nc + if nc is None: + raise ValueError( + f"nc should be provided in post_cfg or kwargs for task '{self.task}'." + ) + self.nc: int = int(nc) + self.anchors: list[Any] | torch.Tensor | None = post_cfg.get( + "anchors", None + ) # anchor coordinates + self.stride: list[int] | torch.Tensor + self.nl: int + self.na: int + self.conf_thres: float + self.iou_thres: float + self.inv_conf_thres: float + + self.e2e = bool(kwargs.pop("e2e", post_cfg.get("e2e", True))) + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected YOLO postprocess kwargs: {unexpected}") + + if self.anchors is None: + nl = post_cfg.get("nl") + if nl is None: + raise ValueError("nl should be provided in post_cfg") + self.nl = nl + if self.nl == 2: + self.stride = [2 ** (4 + i) for i in range(self.nl)] + else: + self.stride = [2 ** (3 + i) for i in range(self.nl)] + self.make_anchors() + else: + if not isinstance(self.anchors, list): + raise TypeError( + f"anchors must be a list, got {type(self.anchors).__name__}." + ) + self.nl = len(self.anchors) + self.na = len(self.anchors[0]) // 2 + self.n_extra: int = post_cfg.get("n_extra", 0) + self.conf_thres = float(post_cfg.get("conf_thres", 0.25)) + self.iou_thres = float(post_cfg.get("iou_thres", 0.7)) + self.set_threshold() + + def anchors_as_list(self) -> list[Any]: + """Return anchors as the configured anchor list.""" + if not isinstance(self.anchors, list): + raise TypeError( + "anchors should be a list for anchor-based YOLO postprocessing." + ) + return self.anchors + + def anchors_as_tensor(self) -> torch.Tensor: + """Return anchors as the generated anchor-point tensor.""" + if not isinstance(self.anchors, torch.Tensor): + raise TypeError( + "anchors should be a tensor for anchor-free YOLO postprocessing." + ) + return cast(torch.Tensor, self.anchors) + + def stride_as_tensor(self) -> torch.Tensor: + """Return strides as the generated stride tensor.""" + if not isinstance(self.stride, torch.Tensor): + raise TypeError( + "stride should be a tensor for anchor-free YOLO postprocessing." + ) + return cast(torch.Tensor, self.stride) + + def __call__( + self, + x: TensorLike | ListTensorLike, + conf_thres: float | None = None, + iou_thres: float | None = None, + multi_label: bool = False, + ) -> list[Any]: + """Executes YOLO postprocessing. + + Includes rearranging, decoding, and NMS. + + Args: + x (TensorLike | ListTensorLike): Raw model outputs. + conf_thres (float | None): Confidence threshold for detection. + iou_thres (float | None): IoU threshold for NMS. + multi_label: Whether to emit one candidate for every class above the + confidence threshold. Validation uses this to match Ultralytics. + + Returns: + list: List of detections per image. + """ + self.set_threshold(conf_thres, iou_thres) + final_detections, proto_outs = self.extract_final_outputs(x) + if final_detections is not None: + if proto_outs is not None: + return self.masking(final_detections, proto_outs) + return final_detections + checked_input = self.check_input(x) + + if not self.e2e: + return self.non_e2e(checked_input) + + predictions, proto_outs = self._pre_process(checked_input) + + nms_output = ( + self.nms_multilabel(predictions) if multi_label else self.nms(predictions) + ) + + if proto_outs is not None: + return self.masking(nms_output, proto_outs) + return nms_output + + def non_e2e(self, x: list[torch.Tensor]) -> Any: + """Return the export-style postprocess output when end-to-end mode is disabled. + + Args: + x: Checked raw model outputs. + + Returns: + Export-style tensors whose batch dimensions remain intact. + """ + if len(x) == 1: + return self.conversion(x) + return self.rearrange(x) + + def _pre_process( + self, + x: list[torch.Tensor], + ) -> tuple[Any, torch.Tensor | list[torch.Tensor] | None]: + """Protected method to preprocess inputs into (predictions, prototypes). + + Args: + x: List of input tensors. + + Returns: + Tuple of (predictions, prototypes). Prototypes may be None. + """ + if len(x) == 1: + converted = self.conversion(x) + if not isinstance(converted, torch.Tensor): + raise TypeError( + "conversion should return a tensor for single-output YOLO postprocessing." + ) + return self.filter_conversion(converted), None + rearranged = self.rearrange(x) + return self.decode(rearranged), None + + def nmsout2eval( + self, + nms_out: Any, + img1_shape: tuple[int, int], + img0_shape: tuple[int, int] | list[tuple[int, int]], + ratio_pad: RatioPad | list[RatioPad | None] | None = None, + ) -> tuple[Any, ...]: + """Converts NMS output to evaluation format (labels, boxes, scores). + + Args: + nms_out: NMS output (tensor or list of tensors). + img1_shape: Resized image shape (height, width). + img0_shape: Original image shape(s). + + Returns: + Tuple: task-specific results. + - Detection: (labels_list, boxes_list, scores_list) + - Segmentation/Pose: (labels_list, boxes_list, scores_list, extra_list) + """ + + return nmsout2eval(nms_out, img1_shape, img0_shape, ratio_pads=ratio_pad) + + def extract_final_outputs( + self, + x: TensorLike | ListTensorLike, + ) -> tuple[list[torch.Tensor] | None, torch.Tensor | None]: + """Extract already-decoded ONNX-style detections when present. + + Args: + x: Raw postprocess input. + + Returns: + A tuple of ``(detections, prototypes)`` when the input already contains + final detections, otherwise ``(None, None)``. + """ + + final_det_dim = 6 + self.n_extra + + if isinstance(x, Sequence): + if not x: + return None, None + + normalized_detections: np.ndarray | torch.Tensor | None = None + normalized_proto: torch.Tensor | None = None + invalid_proto_error: ValueError | None = None + for output in x: + if not isinstance(output, (np.ndarray, torch.Tensor)): + continue + if normalized_detections is None: + normalized_detections = self._normalize_final_detection_tensor( + output, final_det_dim + ) + if normalized_detections is not None: + continue + if normalized_proto is None: + try: + normalized_proto = self._normalize_proto_batch(output) + except ValueError as exc: + if self.task == "instance_segmentation" and output.ndim == 4: + # Defer this until detections are found so unrelated + # four-dimensional outputs do not prevent raw-head + # decoding. Once this is a decoded segmentation output, + # every candidate prototype must be valid regardless of + # its position in the backend output sequence. + invalid_proto_error = exc + continue + + if normalized_detections is not None: + if invalid_proto_error is not None: + raise invalid_proto_error + if self.task == "instance_segmentation" and normalized_proto is None: + raise ValueError( + "Decoded instance-segmentation outputs require a mask prototype tensor." + ) + return self._final_detection_batches( + normalized_detections + ), normalized_proto + return None, None + + normalized_x = self._normalize_final_detection_tensor(x, final_det_dim) + if normalized_x is not None: + if self.task == "instance_segmentation": + raise ValueError( + "Decoded instance-segmentation outputs require a mask prototype tensor." + ) + return self._final_detection_batches(normalized_x), None + + return None, None + + def _normalize_final_detection_tensor( + self, + x: TensorLike, + final_det_dim: int, + ) -> np.ndarray | torch.Tensor | None: + """Return a batched final-detection tensor when ``x`` already contains decoded rows.""" + + while x.ndim == 4 and 1 in (x.shape[0], x.shape[1]): + if x.shape[1] == 1: + x = x[:, 0] + elif x.shape[0] == 1: + x = x[0] + if x.ndim == 2 and x.shape[-1] == final_det_dim: + x = x[None] + if x.ndim == 3 and x.shape[-1] == final_det_dim: + return x + if x.ndim == 3 and x.shape[1] == final_det_dim: + if isinstance(x, np.ndarray): + return np.swapaxes(x, 1, 2) + return x.transpose(1, 2) + + return None + + def _final_detection_batches( + self, x: np.ndarray | torch.Tensor + ) -> list[torch.Tensor]: + """Convert batched final detections to the internal per-image tensor list.""" + + if isinstance(x, np.ndarray): + tensor = torch.from_numpy(x).to(self.device) + else: + tensor = x.to(self.device) + batches: list[torch.Tensor] = [] + for batch in tensor: + valid_rows = torch.isfinite(batch).all(dim=1) + if not bool(valid_rows.all()): + invalid_rows = ( + torch.nonzero(~valid_rows, as_tuple=False) + .flatten() + .detach() + .cpu() + .tolist() + ) + raise ValueError( + "Decoded detection rows must contain only finite values; " + f"invalid rows: {invalid_rows}." + ) + labels = batch[:, 5] + scores = batch[:, 4] + if not bool(((scores >= 0) & (scores <= 1)).all()): + invalid_scores = ( + scores[(scores < 0) | (scores > 1)].detach().cpu().tolist() + ) + raise ValueError( + "Decoded detection confidence values must be in [0, 1]; " + f"got {invalid_scores}." + ) + valid_labels = ( + torch.isfinite(labels) + & (labels == labels.round()) + & (labels >= 0) + & (labels < self.nc) + ) + if not bool(valid_labels.all()): + invalid_labels = labels[~valid_labels].detach().cpu().tolist() + raise ValueError( + "Decoded detection class IDs must be finite integral values in " + f"[0, {self.nc}); got {invalid_labels}." + ) + retained = batch[batch[:, 4] > self.conf_thres] + if getattr(self, "task", "object_detection") == "obb": + valid_geometry = (retained[:, 2] > 0) & (retained[:, 3] > 0) + geometry_description = "positive width and height" + else: + valid_geometry = (retained[:, 2] > retained[:, 0]) & ( + retained[:, 3] > retained[:, 1] + ) + geometry_description = "positive xyxy area" + if not bool(valid_geometry.all()): + raise ValueError( + "Decoded detection boxes must have " + f"{geometry_description} after confidence filtering." + ) + if getattr(self, "task", "object_detection") == "pose_estimation": + keypoint_confidences = retained[:, 8::3] + if not bool( + ((keypoint_confidences >= 0) & (keypoint_confidences <= 1)).all() + ): + raise ValueError( + "Decoded pose keypoint confidence values must be in [0, 1]." + ) + batches.append(retained) + return batches + + def _normalize_proto_batch( + self, proto_outs: np.ndarray | torch.Tensor + ) -> torch.Tensor: + """Normalize prototype masks to ``(B, H, W, C)`` layout.""" + + if isinstance(proto_outs, np.ndarray): + proto = torch.from_numpy(proto_outs).to(self.device) + else: + proto = proto_outs.to(self.device) + + if proto.ndim != 4: + raise ValueError( + f"Expected 4D prototype tensor, got shape {tuple(proto.shape)}." + ) + if not bool(torch.isfinite(proto).all()): + raise ValueError("Mask prototype tensor must contain only finite values.") + if proto.shape[-1] == self.n_extra: + return proto + if proto.shape[1] == self.n_extra: + return proto.permute(0, 2, 3, 1) + raise ValueError(f"Unsupported prototype tensor shape {tuple(proto.shape)}.") + + def make_anchors(self, offset: float = 0.5) -> None: + """ + Generate anchor points and stride tensors based on image size and strides. + Args: + offset (float, optional): Offset for anchor points. Defaults to 0.5. + """ + anchor_points, stride_tensor = [], [] + strides = [2 ** (3 + i) for i in range(self.nl)] + if self.nl == 2: + strides = [strd * 2 for strd in strides] + for strd in strides: + ny, nx = self.imh // strd, self.imw // strd + sy = torch.arange(ny, dtype=torch.float32, device=self.device) + offset + sx = torch.arange(nx, dtype=torch.float32, device=self.device) + offset + yv, xv = torch.meshgrid(sy, sx, indexing="ij") + anchor_points.append(torch.stack((xv, yv), -1).reshape(-1, 2)) + stride_tensor.append( + torch.full((ny * nx, 1), strd, dtype=torch.float32, device=self.device) + ) + self.anchors = torch.cat(anchor_points, dim=0).permute(1, 0) + self.stride = torch.cat(stride_tensor, dim=0).permute(1, 0) + + def set_threshold( + self, conf_thres: float | None = None, iou_thres: float | None = None + ) -> None: + """Set confidence and IoU thresholds. + Args: + conf_thres (float, optional): Confidence threshold. + iou_thres (float, optional): IoU threshold. + """ + conf_thres = self.conf_thres if conf_thres is None else conf_thres + iou_thres = self.iou_thres if iou_thres is None else iou_thres + if isinstance(conf_thres, bool) or not isinstance(conf_thres, (int, float)): + raise TypeError( + f"conf_thres must be numeric, got {type(conf_thres).__name__}." + ) + if isinstance(iou_thres, bool) or not isinstance(iou_thres, (int, float)): + raise TypeError( + f"iou_thres must be numeric, got {type(iou_thres).__name__}." + ) + if not 0 < conf_thres < 1: + raise ValueError(f"conf_thres must be in (0, 1), got {conf_thres}.") + if not 0 < iou_thres < 1: + raise ValueError(f"iou_thres must be in (0, 1), got {iou_thres}.") + self.conf_thres = float(conf_thres) + self.iou_thres = float(iou_thres) + self.inv_conf_thres = -np.log(1 / conf_thres - 1) + + def check_input(self, x: TensorLike | ListTensorLike) -> list[torch.Tensor]: + """Check and prepare input tensors. + Args: + x (TensorLike | ListTensorLike): Input tensor or list of tensors. + Returns: + list[torch.Tensor]: List of tensors on the correct device. + """ + if isinstance(x, np.ndarray): + tensors = [torch.from_numpy(x).to(self.device)] + elif isinstance(x, torch.Tensor): + tensor_input = cast(torch.Tensor, x) + tensors = [tensor_input.to(self.device)] + else: + if not isinstance(x, Sequence): + raise TypeError(f"Got unexpected type for x={type(x)}.") + if all(isinstance(xi, np.ndarray) for xi in x): + tensors = [torch.from_numpy(xi).to(self.device) for xi in x] + elif all(isinstance(xi, torch.Tensor) for xi in x): + torch_inputs = cast(Sequence[torch.Tensor], x) + tensors = [xi.to(self.device) for xi in torch_inputs] + else: + raise TypeError(f"Got unexpected element type for x[0]={type(x[0])}.") + if any(not bool(torch.isfinite(tensor).all()) for tensor in tensors): + raise ValueError( + "Detection output tensors must contain only finite values." + ) + return self.check_dim(tensors) + + def check_dim(self, x: list[torch.Tensor]) -> list[torch.Tensor]: + """Check tensor dimensions. + Args: + x (list[torch.Tensor]): List of tensors. + Returns: + list[torch.Tensor]: List of tensors with corrected dimensions. + """ + y = [] + for xi in x: + if xi.ndim == 3: + xi = xi.unsqueeze(0) + elif xi.ndim in (4, 5): + pass + else: + raise ValueError(f"Got unexpected dim for xi={xi.ndim}.") + y.append(xi) + return y + + def normalize_split_head( + self, x: torch.Tensor, expected_channels: set[int] + ) -> torch.Tensor: + """Normalize a split detection head to ``(B, C, H, W)`` layout. + + This accepts the channel-last tensors produced by ONNX export flows as + well as the channel-first tensors commonly returned by MXQ/NPU inference. + + Args: + x: Raw split-head tensor. + expected_channels: Valid channel sizes for the current head group. + + Returns: + The normalized tensor in ``(B, C, H, W)`` format. + + Raises: + ValueError: If the tensor shape cannot be interpreted. + """ + while x.ndim > 4: + singleton_dims = [idx for idx, size in enumerate(x.shape) if size == 1] + if not singleton_dims: + raise ValueError( + f"Expected up to 4D split-head tensor, got shape {tuple(x.shape)}." + ) + x = x.squeeze(singleton_dims[0]) + if x.ndim == 3: + x = x.unsqueeze(0) + if x.ndim != 4: + raise ValueError( + f"Expected 3D or 4D split-head tensor, got shape {tuple(x.shape)}." + ) + + if x.shape[1] in expected_channels and x.shape[-1] not in expected_channels: + return x + if x.shape[-1] in expected_channels and x.shape[1] not in expected_channels: + return x.permute(0, 3, 1, 2) + if x.shape[1] in expected_channels and x.shape[-1] in expected_channels: + return x + + raise ValueError( + f"Could not infer split-head layout for shape {tuple(x.shape)} with expected channels {expected_channels}." + ) + + @abstractmethod + def rearrange(self, x: list[torch.Tensor]) -> Any: + """Rearranges raw model outputs into a task-specific intermediate form. + + Args: + x: Raw output tensors from the model. + + Returns: + A task-specific intermediate representation used by ``decode``. + """ + + @abstractmethod + def decode(self, x: Any) -> Any: + """Decodes rearranged outputs into a family-specific batched representation. + + Args: + x: Rearranged output tensors. + + Returns: + Decoded detections in the canonical representation for that YOLO family. + """ + + def conversion( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Converts raw outputs into a task-specific intermediate form. + + Args: + x: Input tensors. + + Returns: + A converted detection tensor, or a ``(detections, prototypes)`` tuple + for segmentation-style subclasses. + """ + if len(x) != 1: + raise ValueError( + f"Expected exactly one converted model output, got {len(x)}." + ) + return x[0] + + @abstractmethod + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters converted outputs into per-image detections before NMS. + + Args: + x: Converted output tensor. + + Returns: + Filtered detections for each image in the batch. + """ + + @abstractmethod + def nms(self, x: Any) -> list[torch.Tensor]: + """Performs non-maximum suppression on decoded detections. + + Args: + x: Decoded detections for each image. + + Returns: + Detections after NMS for each image in the batch. + """ + + def nms_multilabel(self, x: Any) -> list[torch.Tensor]: + """Perform validation NMS with all above-threshold class candidates. + + Args: + x: Decoded detections for each image. + + Returns: + Detections after NMS for each image in the batch. + """ + return self.nms(x) + + def validate_split_head_counts(self, **head_groups: Sequence[object]) -> None: + """Require every raw split-output group to provide every detection head.""" + + expected_count = self.nl + counts = {name: len(heads) for name, heads in head_groups.items()} + if any(count != expected_count for count in counts.values()): + found_counts = ", ".join( + f"{name}={count}" for name, count in counts.items() + ) + raise ValueError( + "Incomplete split-head outputs: " + f"expected {expected_count} heads per group, got {found_counts}." + ) + + def masking( + self, x: list[torch.Tensor], proto_outs: torch.Tensor | list[torch.Tensor] + ) -> list[list[torch.Tensor]]: + """Apply prototype masks to detection results. + + Args: + x: Detection results. + proto_outs: Prototype outputs for masks. + + Returns: + list: Detection results with masks. + """ + if len(x) != len(proto_outs): + raise ValueError( + "Detection and prototype batch sizes must match for instance " + f"segmentation, got {len(x)} detections and {len(proto_outs)} prototypes." + ) + masks = [] + for pred, proto in zip(x, proto_outs): + if proto.ndim != 3: + raise ValueError( + f"Expected 3D prototype tensor, got shape {tuple(proto.shape)}." + ) + if proto.shape[-1] == self.n_extra: + proto = proto.permute(2, 0, 1) + elif proto.shape[0] != self.n_extra: + raise ValueError( + f"Unsupported prototype tensor shape {tuple(proto.shape)}." + ) + if len(pred) == 0: + masks.append( + torch.zeros( + (0, self.imh, self.imw), dtype=torch.float32, device=self.device + ) + ) + continue + masks.append( + process_mask_upsample( + proto, pred[:, 6:], pred[:, :4], [self.imh, self.imw] + ) + ) + return [[xi, mask] for xi, mask in zip(x, masks)] + + +# Name retained from the first standalone draft. +YOLOPostBase = YOLODetectionPostBase diff --git a/mblt_vision/utils/postprocess/build_post.py b/mblt_vision/utils/postprocess/build_post.py new file mode 100644 index 0000000..305ac44 --- /dev/null +++ b/mblt_vision/utils/postprocess/build_post.py @@ -0,0 +1,120 @@ +""" +Postprocessing builder. +""" + +from __future__ import annotations + +from ..._tasks import normalize_vision_task +from .base import PostBase +from .cls_post import ClsPost +from .depth_post import DepthPost +from .semantic_seg_post import SemanticSegPost +from .yolo_anchor_post import YOLOAnchorDetectionPost, YOLOAnchorSegPost +from .yolo_anchorless_post import ( + YOLOAnchorlessDetectionPost, + YOLOAnchorlessOBBPost, + YOLOAnchorlessPosePost, + YOLOAnchorlessSegPost, +) +from .yolo_dflfree_post import ( + YOLODFLFreeDetectionPost, + YOLODFLFreeOBBPost, + YOLODFLFreePosePost, + YOLODFLFreeSegPost, +) +from .yolo_nmsfree_post import YOLONMSFreeDetectionPost + + +def build_postprocess( + pre_cfg: dict, + post_cfg: dict, + **kwargs: object, +) -> PostBase: + """Builds a postprocessing object based on the model configuration. + + Args: + pre_cfg (dict): Preprocessing configuration from the model info. + post_cfg (dict): Postprocessing configuration from the model info. + Must contain "task" and relevant flags for the specific task. + **kwargs: Optional runtime overrides passed to the postprocessor. + + Returns: + PostBase: An instance of a postprocessing class tailored for the task. + + Raises: + NotImplementedError: If the specified task is not supported. + """ + task = normalize_vision_task(post_cfg["task"]) + if task == "image_classification": + return ClsPost(pre_cfg, post_cfg) + if task == "depth_estimation": + return DepthPost(pre_cfg, post_cfg) + if task == "semantic_segmentation": + return SemanticSegPost(pre_cfg, post_cfg) + if task in {"object_detection", "face_detection"}: + if post_cfg.get("anchors", False): + return YOLOAnchorDetectionPost( + pre_cfg, + post_cfg, + **kwargs, + ) + if post_cfg.get("dflfree", False): # nms free is only available for detection + return YOLODFLFreeDetectionPost( + pre_cfg, + post_cfg, + **kwargs, + ) + if post_cfg.get("nmsfree", False): + return YOLONMSFreeDetectionPost( + pre_cfg, + post_cfg, + **kwargs, + ) + return YOLOAnchorlessDetectionPost( + pre_cfg, + post_cfg, + **kwargs, + ) + if task == "instance_segmentation": + if post_cfg.get("anchors", False): + return YOLOAnchorSegPost( + pre_cfg, + post_cfg, + **kwargs, + ) + if post_cfg.get("dflfree", False): + return YOLODFLFreeSegPost( + pre_cfg, + post_cfg, + **kwargs, + ) + return YOLOAnchorlessSegPost( + pre_cfg, + post_cfg, + **kwargs, + ) + if task == "pose_estimation": + if post_cfg.get("dflfree", False): + return YOLODFLFreePosePost( + pre_cfg, + post_cfg, + **kwargs, + ) + return YOLOAnchorlessPosePost( + pre_cfg, + post_cfg, + **kwargs, + ) + if task == "obb": + if post_cfg.get("dflfree", False): + return YOLODFLFreeOBBPost( + pre_cfg, + post_cfg, + **kwargs, + ) + return YOLOAnchorlessOBBPost( + pre_cfg, + post_cfg, + **kwargs, + ) + raise NotImplementedError(f"Task {post_cfg['task']} is not implemented yet") diff --git a/mblt_vision/utils/postprocess/cls_post.py b/mblt_vision/utils/postprocess/cls_post.py new file mode 100644 index 0000000..48bd5bf --- /dev/null +++ b/mblt_vision/utils/postprocess/cls_post.py @@ -0,0 +1,117 @@ +""" +Classification postprocessing. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +import torch + +from ...datasets import get_dataset_class_names, get_dataset_config +from ..types import ListTensorLike, TensorLike +from .base import PostBase + + +class ClsPost(PostBase): + """Post-processing for image classification models. + + Typically applies softmax to logits and ensures correct output shape. + """ + + def __init__(self, pre_cfg: dict, post_cfg: dict) -> None: + """Initializes the classification post-processing. + + Args: + pre_cfg (dict): Preprocessing configuration. + post_cfg (dict): Postprocessing configuration. + """ + super().__init__() + self.softmax = post_cfg.get("softmax", False) + self.dataset: str | None = None + self.num_classes: int | None = None + dataset = post_cfg.get("dataset") + if dataset is not None: + if not isinstance(dataset, str) or not dataset: + raise ValueError( + "Classification postprocessing requires post_cfg.dataset to be a non-empty string." + ) + dataset = dataset.lower() + dataset_config = get_dataset_config(dataset) + if "image_classification" not in dataset_config["tasks"]: + raise ValueError( + f"Dataset '{dataset}' does not support image classification." + ) + self.dataset = dataset + self.num_classes = len(get_dataset_class_names(dataset)) + + def __call__(self, x: TensorLike | ListTensorLike) -> torch.Tensor: + """Executes classification post-processing. + + Typically applies softmax to convert logits to probabilities. + + Args: + x (TensorLike | ListTensorLike): Raw model outputs. + Expected to be pre-softmax logits. + + Returns: + torch.Tensor: Softmax probabilities of shape (N, C). + """ + if isinstance(x, Sequence): + if len(x) != 1: + raise ValueError( + f"Classification postprocessing expects one output tensor, got {len(x)}." + ) + x = x[0] + if isinstance(x, np.ndarray): + x = torch.from_numpy(x).to(self.device) + elif isinstance(x, torch.Tensor): + x = x.to(self.device) + else: + raise TypeError(f"Got unexpected type for x={type(x)}.") + if x.ndim == 2: + x = x.unsqueeze(-1).unsqueeze(-1) + elif x.ndim == 3: + if self.num_classes is None: + raise ValueError( + "Classification 3D outputs require a configured class count to disambiguate layout." + ) + if x.shape[1] == self.num_classes and x.shape[-1] == 1: + x = x.unsqueeze(-1) + elif x.shape[0] == self.num_classes: + x = x.unsqueeze(0) + else: + raise ValueError( + f"Unsupported 3D classification output shape {tuple(x.shape)} for {self.num_classes} classes." + ) + if x.ndim != 4: + raise ValueError( + f"Classification output must be convertible to NCHW, got shape {tuple(x.shape)}." + ) + x = x.flatten( + 1 + ) # Classification heads may retain singleton spatial dimensions. + if self.num_classes is not None and x.shape[1] != self.num_classes: + raise ValueError( + f"Classification output has {x.shape[1]} classes, but dataset " + f"'{self.dataset}' requires {self.num_classes}." + ) + if not torch.isfinite(x).all(): + raise ValueError("Classification output scores must all be finite.") + if self.softmax: + if not bool(((x >= 0) & (x <= 1)).all()): + raise ValueError( + "Classification probability outputs must be in [0, 1] when post_cfg.softmax is true." + ) + if not torch.allclose( + x.sum(dim=-1), + torch.ones(x.shape[0], dtype=x.dtype, device=x.device), + rtol=1e-4, + atol=1e-4, + ): + raise ValueError( + "Classification probability outputs must sum to 1 per sample when post_cfg.softmax is true." + ) + return x + return x.softmax(dim=-1) diff --git a/mblt_vision/utils/postprocess/common.py b/mblt_vision/utils/postprocess/common.py new file mode 100644 index 0000000..dc6ed5a --- /dev/null +++ b/mblt_vision/utils/postprocess/common.py @@ -0,0 +1,1571 @@ +"""Common postprocessing utility functions.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from typing import Any, TypeGuard, overload + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F + +from ..datasets import get_coco_inv, get_dotav1_label +from ..letterbox import RatioPad, resolve_ratio_pad + + +def _is_ratio_pad(value: object) -> TypeGuard[RatioPad]: + """Return whether a value has the nested numeric shape of one RatioPad.""" + + return ( + isinstance(value, tuple) + and len(value) == 2 + and all( + isinstance(pair, tuple) + and len(pair) == 2 + and all(isinstance(component, (int, float)) for component in pair) + for pair in value + ) + ) + + +def normalize_image_shapes( + image_shapes: tuple[int, int] | Sequence[tuple[int, int]], + batch_size: int | None = None, +) -> list[tuple[int, int]]: + """Normalize one or many image shapes to a list, optionally validating its batch size.""" + + if len(image_shapes) == 2 and isinstance(image_shapes[0], int): + shapes = [(int(image_shapes[0]), int(image_shapes[1]))] # type: ignore[index] + if batch_size is not None: + shapes *= batch_size + else: + shapes = [(int(shape[0]), int(shape[1])) for shape in image_shapes] # type: ignore[union-attr] + if batch_size is not None and len(shapes) != batch_size: + raise ValueError(f"Expected {batch_size} image shapes, got {len(shapes)}.") + return shapes + + +def normalize_ratio_pads( + ratio_pads: RatioPad | Sequence[RatioPad | None] | None, + batch_size: int, +) -> list[RatioPad | None]: + """Normalize optional letterbox metadata to a batch-sized list.""" + + if ratio_pads is None: + return [None] * batch_size + if _is_ratio_pad(ratio_pads): + return [ratio_pads] * batch_size + pads: list[RatioPad | None] = [] + for ratio_pad in ratio_pads: + if ratio_pad is None: + pads.append(None) + elif _is_ratio_pad(ratio_pad): + pads.append(ratio_pad) + else: + raise TypeError( + "Each ratio_pad must be a ((ratio_x, ratio_y), (pad_x, pad_y)) tuple or None." + ) + if len(pads) != batch_size: + raise ValueError(f"Expected {batch_size} ratio_pad values, got {len(pads)}.") + return pads + + +# --- Box Conversion Utilities --- +@overload +def xywh2xyxy(x: np.ndarray) -> np.ndarray: + """Converts numpy boxes from ``xywh`` to ``xyxy`` format.""" + + +@overload +def xywh2xyxy(x: torch.Tensor) -> torch.Tensor: + """Converts torch boxes from ``xywh`` to ``xyxy`` format.""" + + +def xywh2xyxy(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """Converts bounding box coordinates from (cx, cy, w, h) to (x1, y1, x2, y2). + + (x1, y1) is the top-left corner and (x2, y2) is the bottom-right corner. + + Args: + x: Input bounding boxes in (cx, cy, w, h) format. + + Returns: + Bounding boxes in (x1, y1, x2, y2) format. + """ + if isinstance(x, np.ndarray): + y = np.copy(x) + y[..., 0] = x[..., 0] - x[..., 2] / 2 + y[..., 1] = x[..., 1] - x[..., 3] / 2 + y[..., 2] = x[..., 0] + x[..., 2] / 2 + y[..., 3] = x[..., 1] + x[..., 3] / 2 + return y + + if isinstance(x, torch.Tensor): + y = torch.clone(x) + y[..., 0] = x[..., 0] - x[..., 2] / 2 + y[..., 1] = x[..., 1] - x[..., 3] / 2 + y[..., 2] = x[..., 0] + x[..., 2] / 2 + y[..., 3] = x[..., 1] + x[..., 3] / 2 + return y + + raise ValueError("x should be np.ndarray or torch.Tensor") + + +@overload +def xyxy2xywh(x: np.ndarray) -> np.ndarray: + """Converts numpy boxes from ``xyxy`` to ``xywh`` format.""" + + +@overload +def xyxy2xywh(x: torch.Tensor) -> torch.Tensor: + """Converts torch boxes from ``xyxy`` to ``xywh`` format.""" + + +def xyxy2xywh(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """Converts bounding box coordinates from (x1, y1, x2, y2) to (cx, cy, w, h). + + (x1, y1) is the top-left corner and (x2, y2) is the bottom-right corner. + (cx, cy) is the center of the bounding box. + + Args: + x: Input bounding boxes in (x1, y1, x2, y2) format. + + Returns: + Bounding boxes in (cx, cy, w, h) format. + """ + if isinstance(x, np.ndarray): + y = np.copy(x) + y[..., 0] = (x[..., 0] + x[..., 2]) / 2 + y[..., 1] = (x[..., 1] + x[..., 3]) / 2 + y[..., 2] = x[..., 2] - x[..., 0] + y[..., 3] = x[..., 3] - x[..., 1] + return y + + if isinstance(x, torch.Tensor): + y = torch.clone(x) + y[..., 0] = (x[..., 0] + x[..., 2]) / 2 + y[..., 1] = (x[..., 1] + x[..., 3]) / 2 + y[..., 2] = x[..., 2] - x[..., 0] + y[..., 3] = x[..., 3] - x[..., 1] + return y + + raise ValueError("x should be np.ndarray or torch.Tensor") + + +def dist2bbox( + distance: torch.Tensor, + anchor_points: torch.Tensor, + xywh: bool = True, + dim: int = -1, +) -> torch.Tensor: + """ + Transform distance (ltrb) to bounding box (xywh or xyxy). + Args: + distance (torch.Tensor): Distance from anchor points to box boundaries + (left, top, right, bottom). + anchor_points (torch.Tensor): Anchor points (center points). + xywh (bool, optional): If True, return boxes in (cx, cy, w, h) format. + If False, return in (x1, y1, x2, y2) format. Defaults to True. + dim (int, optional): Dimension along which to chunk the distance tensor. Defaults to -1. + Returns: + torch.Tensor: Transformed bounding boxes. + """ + lt, rb = distance.chunk(2, dim) + x1y1 = anchor_points - lt + x2y2 = anchor_points + rb + if xywh: + return torch.cat(((x1y1 + x2y2) / 2, x2y2 - x1y1), dim) # xywh bbox + else: + return torch.cat((x1y1, x2y2), dim) # xyxy bbox + + +def dist2rbox( + distance: torch.Tensor, + angle: torch.Tensor, + anchor_points: torch.Tensor, + dim: int = -1, +) -> torch.Tensor: + """Decode rotated boxes from anchor-relative distances and angles. + + Args: + distance: Distance tensor in ``ltrb`` format. + angle: Rotation angle tensor in radians. + anchor_points: Anchor center points. + dim: Dimension along which box channels are split. + + Returns: + Rotated boxes in ``cx, cy, w, h`` format. + """ + lt, rb = distance.split(2, dim=dim) + cos_value = torch.cos(angle) + sin_value = torch.sin(angle) + xf, yf = ((rb - lt) / 2).split(1, dim=dim) + x = xf * cos_value - yf * sin_value + y = xf * sin_value + yf * cos_value + xy = torch.cat([x, y], dim=dim) + anchor_points + return torch.cat([xy, lt + rb], dim=dim) + + +@overload +def xywhr2xyxyxyxy(x: np.ndarray) -> np.ndarray: + """Converts numpy OBBs from ``xywhr`` to polygon corners.""" + + +@overload +def xywhr2xyxyxyxy(x: torch.Tensor) -> torch.Tensor: + """Converts torch OBBs from ``xywhr`` to polygon corners.""" + + +def xywhr2xyxyxyxy(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """Converts oriented boxes from ``cx, cy, w, h, angle`` to four corner points. + + Args: + x: Oriented boxes with shape ``(..., 5)`` and angle in radians. + + Returns: + Corner points with shape ``(..., 4, 2)``. + """ + if isinstance(x, torch.Tensor): + ctr = x[..., :2] + w, h, angle = (x[..., i : i + 1] for i in range(2, 5)) + cos_value = torch.cos(angle) + sin_value = torch.sin(angle) + vec1 = torch.cat([w / 2 * cos_value, w / 2 * sin_value], dim=-1) + vec2 = torch.cat([-h / 2 * sin_value, h / 2 * cos_value], dim=-1) + return torch.stack( + [ + ctr + vec1 + vec2, + ctr + vec1 - vec2, + ctr - vec1 - vec2, + ctr - vec1 + vec2, + ], + dim=-2, + ) + + if isinstance(x, np.ndarray): + ctr = x[..., :2] + w, h, angle = (x[..., i : i + 1] for i in range(2, 5)) + cos_value = np.cos(angle) + sin_value = np.sin(angle) + vec1 = np.concatenate([w / 2 * cos_value, w / 2 * sin_value], axis=-1) + vec2 = np.concatenate([-h / 2 * sin_value, h / 2 * cos_value], axis=-1) + return np.stack( + [ + ctr + vec1 + vec2, + ctr + vec1 - vec2, + ctr - vec1 - vec2, + ctr - vec1 + vec2, + ], + axis=-2, + ) + + raise ValueError("x should be np.ndarray or torch.Tensor") + + +def xyxyxyxy2xywhr(points: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """Converts OBB corner points to regularized ``xywhr`` boxes. + + Args: + points: Corner points with shape ``(..., 4, 2)``. + + Returns: + Rotated boxes in ``cx, cy, w, h, angle`` format. + """ + is_torch = isinstance(points, torch.Tensor) + points_np = points.detach().cpu().numpy() if is_torch else np.asarray(points) + flat_points = points_np.reshape(-1, 4, 2).astype(np.float32) + rboxes = [] + for pts in flat_points: + (cx, cy), (w, h), angle = cv2.minAreaRect(pts) + theta = angle / 180 * np.pi + if w < h: + w, h = h, w + theta += np.pi / 2 + while theta >= 3 * np.pi / 4: + theta -= np.pi + while theta < -np.pi / 4: + theta += np.pi + rboxes.append([cx, cy, w, h, theta]) + result_np = np.asarray(rboxes, dtype=points_np.dtype).reshape( + *points_np.shape[:-2], 5 + ) + if is_torch: + return torch.tensor(result_np, device=points.device, dtype=points.dtype) + return result_np + + +def regularize_rboxes(rboxes: torch.Tensor) -> torch.Tensor: + """Regularize rotated boxes to the angle range ``[0, pi / 2)``. + + Args: + rboxes: Rotated boxes in ``xywhr`` format. + + Returns: + Regularized rotated boxes. + """ + x, y, w, h, angle = rboxes.unbind(dim=-1) + swap = angle % math.pi >= math.pi / 2 + regularized_w = torch.where(swap, h, w) + regularized_h = torch.where(swap, w, h) + regularized_angle = angle % (math.pi / 2) + return torch.stack([x, y, regularized_w, regularized_h, regularized_angle], dim=-1) + + +def _get_covariance_matrix( + boxes: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return Gaussian covariance components for probabilistic OBB IoU.""" + gbbs = torch.cat((boxes[:, 2:4].pow(2) / 12, boxes[:, 4:]), dim=-1) + a, b, c = gbbs.split(1, dim=-1) + cos_value = c.cos() + sin_value = c.sin() + cos2 = cos_value.pow(2) + sin2 = sin_value.pow(2) + return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos_value * sin_value + + +def batch_probiou( + obb1: torch.Tensor | np.ndarray, obb2: torch.Tensor | np.ndarray, eps: float = 1e-7 +) -> torch.Tensor: + """Calculate pairwise probabilistic IoU for oriented boxes. + + Args: + obb1: First set of OBBs in ``xywhr`` format with shape ``(N, 5)``. + obb2: Second set of OBBs in ``xywhr`` format with shape ``(M, 5)``. + eps: Small value used for numerical stability. + + Returns: + Pairwise OBB similarities with shape ``(N, M)``. + """ + obb1 = torch.from_numpy(obb1) if isinstance(obb1, np.ndarray) else obb1 + obb2 = torch.from_numpy(obb2) if isinstance(obb2, np.ndarray) else obb2 + obb2 = obb2.to(device=obb1.device, dtype=obb1.dtype) + + x1, y1 = obb1[..., :2].split(1, dim=-1) + x2, y2 = (x.squeeze(-1)[None] for x in obb2[..., :2].split(1, dim=-1)) + a1, b1, c1 = _get_covariance_matrix(obb1) + a2, b2, c2 = (x.squeeze(-1)[None] for x in _get_covariance_matrix(obb2)) + + denominator = (a1 + a2) * (b1 + b2) - (c1 + c2).pow(2) + eps + t1 = ( + ((a1 + a2) * (y1 - y2).pow(2) + (b1 + b2) * (x1 - x2).pow(2)) / denominator + ) * 0.25 + t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / denominator) * 0.5 + t3 = ( + ((a1 + a2) * (b1 + b2) - (c1 + c2).pow(2)) + / ( + 4 + * ((a1 * b1 - c1.pow(2)).clamp_(0) * (a2 * b2 - c2.pow(2)).clamp_(0)).sqrt() + + eps + ) + + eps + ).log() * 0.5 + bd = (t1 + t2 + t3).clamp(eps, 100.0) + hd = (1.0 - (-bd).exp() + eps).sqrt() + return 1 - hd + + +def rotated_nms( + boxes: torch.Tensor, + scores: torch.Tensor, + iou_threshold: float, + iou_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = batch_probiou, +) -> torch.Tensor: + """Apply fast rotated NMS using an upper-triangular pairwise IoU matrix. + + Args: + boxes: OBBs in ``xywhr`` format. + scores: Confidence scores. + iou_threshold: IoU threshold for suppression. + iou_func: Pairwise IoU function. + + Returns: + Kept indices into the original inputs. + """ + if boxes.numel() == 0: + return torch.empty((0,), dtype=torch.int64, device=boxes.device) + sorted_idx = torch.argsort(scores, descending=True) + sorted_boxes = boxes[sorted_idx] + ious = iou_func(sorted_boxes, sorted_boxes).triu_(diagonal=1) + keep = torch.nonzero((ious >= iou_threshold).sum(0) <= 0).squeeze_(-1) + return sorted_idx[keep] + + +# --- Detection Utilities --- +def non_max_suppression( + boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float, max_output: int +) -> list[int]: + """ + Modified non-maximum suppression (NMS) implemented with PyTorch. + Args: + boxes (torch.Tensor): Bounding boxes in (x1, y1, x2, y2) format. + scores (torch.Tensor): Confidence scores for each box (assumed to be sorted in + descending order). + iou_threshold (float): IoU threshold for suppression. + max_output (int): Maximum number of boxes to keep. + Returns: + list[int]: Indices of the boxes that have been kept after NMS. + """ + if boxes.numel() == 0: + return [] + # Coordinates of bounding boxes + start_x = boxes[:, 0] + start_y = boxes[:, 1] + end_x = boxes[:, 2] + end_y = boxes[:, 3] + picked_indices: list[int] = [] + # Compute areas of bounding boxes + areas = (end_x - start_x) * (end_y - start_y) + # Create an index order (assumed scores are already sorted in descending order) + order = torch.arange(scores.size(0)).to(boxes.device) + while order.numel() > 0 and len(picked_indices) < max_output: + # The index with the highest score + index = int(order[0].item()) + picked_indices.append(index) + order = order[1:] # Remove the index from the order + if order.numel() == 0 or len(picked_indices) >= max_output: + break + # Compute the coordinates of the intersection boxes + x1 = torch.maximum(start_x[index], start_x[order]) + y1 = torch.maximum(start_y[index], start_y[order]) + x2 = torch.minimum(end_x[index], end_x[order]) + y2 = torch.minimum(end_y[index], end_y[order]) + # Compute width and height of the intersection boxes + w = torch.clamp(x2 - x1, min=0.0) + h = torch.clamp(y2 - y1, min=0.0) + intersection = w * h + # Compute the IoU ratio + union = areas[index] + areas[order] - intersection + ratio = intersection / union + # Keep boxes with IoU less than or equal to the threshold + keep = (ratio <= iou_threshold).to(order.device) + order = order[keep] + return picked_indices + + +def dual_topk( + pre_topk: torch.Tensor, + nc: int, + n_extra: int, + max_det: int = 300, + conf_thres: float = 0.25, + score_is_logits: bool = False, +) -> torch.Tensor: + """ + Perform dual-stage topk selection for NMS-free models. + Args: + pre_topk (torch.Tensor): Input tensor of shape (*, 4 + nc + n_extra). + nc (int): Number of classes. + n_extra (int): Number of extra elements (e.g., masks, keypoints). + max_det (int): Maximum detections to keep. Defaults to 300. + conf_thres (float): Confidence threshold. Defaults to 0.25. + score_is_logits (bool): Whether class scores are logits. When true, apply + the confidence cutoff and both rankings before sigmoid, then convert + only selected scores to probabilities. Defaults to false. + Returns: + torch.Tensor: Filtered detections of shape (*, 6 + n_extra). + """ + score_start = 4 + score_end = 4 + nc + score_view = pre_topk[:, score_start:score_end] + threshold = ( + math.log(conf_thres / (1.0 - conf_thres)) if score_is_logits else conf_thres + ) + ic = score_view.amax(dim=-1) > threshold + pre_topk = pre_topk[ic] + + if pre_topk.shape[0] == 0: + return torch.zeros( + (0, 6 + n_extra), dtype=torch.float32, device=pre_topk.device + ) + max_det = min(pre_topk.shape[0], max_det) + + row_index = torch.topk( + pre_topk[:, score_start:score_end].amax(dim=-1), max_det, dim=0 + ).indices + selected = pre_topk[row_index] + top_scores, flat_index = torch.topk( + selected[:, score_start:score_end].reshape(-1), max_det + ) + keep = top_scores > threshold + if not torch.any(keep): + return torch.zeros( + (0, 6 + n_extra), dtype=torch.float32, device=pre_topk.device + ) + + top_scores = top_scores[keep] + flat_index = flat_index[keep] + box_index = flat_index // nc + labels = (flat_index % nc).to(selected.dtype).unsqueeze(-1) + + output = torch.empty( + (top_scores.shape[0], 6 + n_extra), dtype=selected.dtype, device=selected.device + ) + output[:, :4] = selected[box_index, :4] + output[:, 4] = top_scores.sigmoid() if score_is_logits else top_scores + output[:, 5:6] = labels + if n_extra > 0: + output[:, 6:] = selected[box_index, score_end:] + return output + + +def yolo_multilabel_candidates( + detections: torch.Tensor, + nc: int, + n_extra: int, + conf_thres: float, +) -> torch.Tensor: + """Expand YOLO rows into one detection per class score above threshold. + + Args: + detections: Row-major detections with columns ``box, class scores, extra``. + nc: Number of classes. + n_extra: Number of extra channels after class scores. + conf_thres: Confidence threshold. + + Returns: + Canonical detection rows with columns ``box, score, class, extra``. + """ + if detections.numel() == 0: + return torch.zeros( + (0, 6 + n_extra), dtype=torch.float32, device=detections.device + ) + + boxes = detections[:, :4] + scores = detections[:, 4 : 4 + nc] + extra = detections[:, 4 + nc :] + box_index, class_index = torch.where(scores > conf_thres) + if box_index.numel() == 0: + return torch.zeros( + (0, 6 + n_extra), dtype=torch.float32, device=detections.device + ) + + output = torch.empty( + (box_index.numel(), 6 + n_extra), + dtype=detections.dtype, + device=detections.device, + ) + output[:, :4] = boxes[box_index] + output[:, 4] = scores[box_index, class_index] + output[:, 5] = class_index.to(detections.dtype) + if n_extra > 0: + output[:, 6:] = extra[box_index] + return output + + +def normalize_converted_obb_part(x: torch.Tensor, channel_count: int) -> torch.Tensor: + """Normalize a converted OBB output part to ``(batch, anchors, channels)``. + + Args: + x: Converted output part from a model runtime. + channel_count: Expected feature-channel count for this part. + + Returns: + The normalized row-major tensor. + """ + while x.ndim > 3: + singleton_dims = [ + idx for idx, size in enumerate(x.shape) if idx != 0 and size == 1 + ] + if not singleton_dims: + raise ValueError( + f"Expected converted OBB part with up to 3 non-batch dimensions, got {tuple(x.shape)}." + ) + x = x.squeeze(singleton_dims[0]) + if x.ndim == 2: + x = x.unsqueeze(0) + if x.ndim != 3: + raise ValueError( + f"Expected 2D or 3D converted OBB part, got shape {tuple(x.shape)}." + ) + if x.shape[-1] == channel_count: + return x + if x.shape[1] == channel_count: + return x.transpose(1, 2) + raise ValueError( + f"Could not find channel count {channel_count} in converted OBB part with shape {tuple(x.shape)}." + ) + + +def concat_converted_obb_outputs( + x: list[torch.Tensor], nc: int, n_extra: int +) -> torch.Tensor: + """Concatenate converted OBB box, class, and angle outputs in canonical order. + + Args: + x: Converted OBB runtime outputs. + nc: Number of OBB classes. + n_extra: Number of extra OBB channels. + + Returns: + Detections in ``cx, cy, w, h, class scores..., angle`` format. + """ + if len(x) == 1: + return x[0] + if len(x) != 3: + raise ValueError(f"Expected 1 or 3 converted OBB outputs, got {len(x)}.") + + expected_parts = {"box": 4, "scores": nc, "angle": n_extra} + parts: dict[str, torch.Tensor] = {} + for xi in x: + matches: list[tuple[str, torch.Tensor]] = [] + for name, channel_count in expected_parts.items(): + try: + matches.append((name, normalize_converted_obb_part(xi, channel_count))) + except ValueError: + continue + if len(matches) != 1: + match_names = ", ".join(name for name, _ in matches) or "none" + raise ValueError( + f"Could not uniquely classify converted OBB output {tuple(xi.shape)}; matches: {match_names}." + ) + name, normalized = matches[0] + if name in parts: + raise ValueError(f"Duplicate converted OBB {name} output.") + parts[name] = normalized + + missing = [name for name in expected_parts if name not in parts] + if missing: + raise ValueError(f"Missing converted OBB outputs: {', '.join(missing)}.") + return torch.cat([parts["box"], parts["scores"], parts["angle"]], dim=-1) + + +def decode_split_converted_obb_outputs( + x: list[torch.Tensor], + nc: int, + n_extra: int, + anchors: torch.Tensor, + stride: torch.Tensor, +) -> torch.Tensor: + """Decode MXQ decode-true OBB outputs split into score, angle, and coordinate tensors. + + Args: + x: Five converted runtime outputs: class scores, rotation angle, and + coordinate tensors containing decoded ``wh`` and pre-rotated center offsets. + nc: Number of OBB classes. + n_extra: Number of extra OBB channels. + anchors: Anchor points in ``(2, anchors)`` format. + stride: Stride tensor in ``(1, anchors)`` format. + + Returns: + Detections in ``cx, cy, w, h, class scores..., angle`` format. + """ + if n_extra != 1: + raise ValueError(f"Expected one OBB angle channel, got n_extra={n_extra}.") + if len(x) != 5: + raise ValueError(f"Expected five split converted OBB outputs, got {len(x)}.") + + try: + scores = normalize_converted_obb_part(x[0], nc) + angle = normalize_converted_obb_part(x[1], n_extra) + except ValueError: + scores = normalize_converted_obb_part(x[1], nc) + angle = normalize_converted_obb_part(x[0], n_extra) + wh = normalize_converted_obb_part(x[2], 2) + x_offset = normalize_converted_obb_part(x[3], 1) + y_offset = normalize_converted_obb_part(x[4], 1) + cos_value = torch.cos(angle) + sin_value = torch.sin(angle) + center_offset = torch.cat( + [ + x_offset * cos_value - y_offset * sin_value, + x_offset * sin_value + y_offset * cos_value, + ], + dim=-1, + ) + anchors_t = ( + anchors.transpose(0, 1).unsqueeze(0).to(device=wh.device, dtype=wh.dtype) + ) + stride_t = stride.transpose(0, 1).unsqueeze(0).to(device=wh.device, dtype=wh.dtype) + if anchors_t.shape[1] < wh.shape[1]: + raise ValueError( + f"Got {wh.shape[1]} OBB coordinate rows but only {anchors_t.shape[1]} anchors." + ) + anchors_t = anchors_t[:, : wh.shape[1]] + stride_t = stride_t[:, : wh.shape[1]] + box = torch.cat([anchors_t + center_offset, wh], dim=-1) * stride_t + return torch.cat([box, scores, angle], dim=-1) + + +# --- Scaling & Clipping Utilities --- +@overload +def scale_boxes( + img1_shape: tuple[int, int], + boxes: np.ndarray, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> np.ndarray: ... + + +@overload +def scale_boxes( + img1_shape: tuple[int, int], + boxes: torch.Tensor, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> torch.Tensor: ... + + +def scale_boxes( + img1_shape: tuple[int, int], + boxes: np.ndarray | torch.Tensor, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> np.ndarray | torch.Tensor: + """ + Original Source: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/ops.py#L92 + Rescales bounding boxes (in the format of xyxy) from the shape of the image they + were originally specified in (img1_shape) to the shape of a different image (img0_shape). + Args: + img1_shape (tuple): The shape of the image that the bounding boxes are for, + in the format of (height, width). + boxes (np.ndarray | torch.Tensor): the bounding boxes of the objects in the image, + in the format of (x1, y1, x2, y2) + img0_shape (tuple): the shape of the target image, in the format of (height, width). + ratio_pad (tuple): a tuple of (ratio, pad) for scaling the boxes. + If not provided, the ratio and pad will be calculated based on the size + difference between the two images. + padding (bool): If True, assuming the boxes is based on image augmented by + yolo style. If False then do regular rescaling. + Returns: + np.ndarray | torch.Tensor: The scaled bounding boxes, in the format of (x1, y1, x2, y2) + """ + ratio, pad = resolve_ratio_pad(img1_shape, img0_shape, ratio_pad) + gain = ratio[0] + if isinstance(boxes, np.ndarray): + if padding: + boxes[..., [0, 2]] -= pad[0] # x padding + boxes[..., [1, 3]] -= pad[1] # y padding + boxes[..., :4] /= gain + return clip_boxes(boxes, img0_shape) + if padding: + boxes[..., [0, 2]] -= pad[0] # x padding + boxes[..., [1, 3]] -= pad[1] # y padding + boxes[..., :4] /= gain + return clip_boxes(boxes, img0_shape) + + +@overload +def scale_coords( + img1_shape: tuple[int, int], + coords: np.ndarray, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> np.ndarray: ... + + +@overload +def scale_coords( + img1_shape: tuple[int, int], + coords: torch.Tensor, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> torch.Tensor: ... + + +def scale_coords( + img1_shape: tuple[int, int], + coords: np.ndarray | torch.Tensor, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> np.ndarray | torch.Tensor: + """ + Original Source: + https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/ops.py#L756 + Args: + img1_shape (tuple): The shape of the image that the bounding boxes are for, in the format of (height, width). + coords (np.ndarray | torch.Tensor): The coordinates of the objects in the image, in the format of (x, y). + img0_shape (tuple): The shape of the target image, in the format of (height, width). + ratio_pad (tuple): a tuple of (ratio, pad) for scaling the boxes. If not provided, the ratio and pad will be + calculated based on the size difference between the two images. + padding (bool): If True, assuming the boxes is based on image augmented by yolo style. If False then do regular + rescaling. + Returns: + np.ndarray | torch.Tensor: The scaled coordinates, in the format of (x, y) + """ + ratio, pad = resolve_ratio_pad(img1_shape, img0_shape, ratio_pad) + gain = ratio[0] + if isinstance(coords, np.ndarray): + if padding: + coords[..., 0] -= pad[0] # x padding + coords[..., 1] -= pad[1] # y padding + coords[..., :2] /= gain + return clip_coords(coords, img0_shape) + if padding: + coords[..., 0] -= pad[0] # x padding + coords[..., 1] -= pad[1] # y padding + coords[..., :2] /= gain + return clip_coords(coords, img0_shape) + + +def scale_rboxes( + img1_shape: tuple[int, int], + rboxes: torch.Tensor, + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> torch.Tensor: + """Rescale rotated boxes from model input size to an original image size. + + Args: + img1_shape: Processed image shape. + rboxes: Rotated boxes in ``xywhr`` format. + img0_shape: Original image shape. + ratio_pad: Optional precomputed resize ratio and padding. + padding: Whether YOLO-style letterbox padding was applied. + + Returns: + Rescaled rotated boxes in ``xywhr`` format. + """ + ratio, pad = resolve_ratio_pad(img1_shape, img0_shape, ratio_pad) + gain = ratio[0] + scaled = rboxes.clone() + if padding: + scaled[..., 0] -= pad[0] + scaled[..., 1] -= pad[1] + scaled[..., :4] /= gain + return scaled + + +def compute_ratio_pad( + img1_shape: tuple[int, int], + img0_shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, +) -> tuple[float, tuple[float, float]]: + """Return letterbox gain and padding for compatibility with existing callers. + + Args: + img1_shape (tuple): The target shape (height, width). + img0_shape (tuple): The original shape (height, width). + ratio_pad (tuple, optional): Pre-calculated (ratio, pad) tuple. + If None, it will be calculated from the shapes. Defaults to None. + + Returns: + tuple: (gain, pad) where gain is the scaling factor and pad is the (x, y) padding. + """ + ratio, pad = resolve_ratio_pad(img1_shape, img0_shape, ratio_pad) + return ratio[0], pad + + +@overload +def clip_boxes(boxes: np.ndarray, shape: tuple[int, int]) -> np.ndarray: ... + + +@overload +def clip_boxes(boxes: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor: ... + + +def clip_boxes( + boxes: np.ndarray | torch.Tensor, shape: tuple[int, int] +) -> np.ndarray | torch.Tensor: + """ + Clip bounding boxes to image shape. + Args: + boxes (np.ndarray | torch.Tensor): Bounding boxes. + shape (tuple): Image shape (height, width). + Returns: + np.ndarray | torch.Tensor: Clipped bounding boxes. + """ + if isinstance(boxes, torch.Tensor): + boxes[..., 0] = boxes[..., 0].clamp(0, shape[1]) + boxes[..., 1] = boxes[..., 1].clamp(0, shape[0]) + boxes[..., 2] = boxes[..., 2].clamp(0, shape[1]) + boxes[..., 3] = boxes[..., 3].clamp(0, shape[0]) + else: + boxes[..., 0] = np.clip(boxes[..., 0], 0, shape[1]) + boxes[..., 1] = np.clip(boxes[..., 1], 0, shape[0]) + boxes[..., 2] = np.clip(boxes[..., 2], 0, shape[1]) + boxes[..., 3] = np.clip(boxes[..., 3], 0, shape[0]) + return boxes + + +@overload +def clip_coords(coords: np.ndarray, shape: tuple[int, int]) -> np.ndarray: ... + + +@overload +def clip_coords(coords: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor: ... + + +def clip_coords( + coords: np.ndarray | torch.Tensor, shape: tuple[int, int] +) -> np.ndarray | torch.Tensor: + """Clips coordinates to the image shape. + + Args: + coords (np.ndarray | torch.Tensor): Coordinates to clip. + shape (tuple): Image shape (height, width). + + Returns: + np.ndarray | torch.Tensor: Clipped coordinates. + """ + if isinstance(coords, torch.Tensor): + coords[..., 0] = coords[..., 0].clamp(0, shape[1]) + coords[..., 1] = coords[..., 1].clamp(0, shape[0]) + else: + coords[..., 0] = np.clip(coords[..., 0], 0, shape[1]) + coords[..., 1] = np.clip(coords[..., 1], 0, shape[0]) + return coords + + +# --- Segmentation Utilities --- +def process_mask( + protos: torch.Tensor, + masks_in: torch.Tensor, + bboxes: torch.Tensor, + shape: tuple[int, int], + upsample: bool = False, +) -> torch.Tensor: + """Processes masks by applying coefficients to prototypes and cropping. + + Ref: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/ops.py#L680 + + Args: + protos (torch.Tensor): Prototype masks of shape [mask_dim, mask_h, mask_w]. + masks_in (torch.Tensor): Mask coefficients of shape [n, mask_dim]. + bboxes (torch.Tensor): Bounding boxes of shape [n, 4]. + shape (tuple): Input image size (h, w). + upsample (bool, optional): Whether to upsample the masks to the original image size. + Defaults to False. + + Returns: + torch.Tensor: Processed binary masks. + """ + c, mh, mw = protos.shape # CHW + ih, iw = shape + masks = (masks_in @ protos.float().view(c, -1)).view(-1, mh, mw) # n, CHW + downsampled_bboxes = bboxes.clone() + downsampled_bboxes[:, 0] *= mw / iw + downsampled_bboxes[:, 2] *= mw / iw + downsampled_bboxes[:, 3] *= mh / ih + downsampled_bboxes[:, 1] *= mh / ih + masks = crop_mask(masks, downsampled_bboxes) # CHW + if upsample: + masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[ + 0 + ] # CHW + return masks.gt_(0.0) + + +def process_mask_upsample( + protos: torch.Tensor, + masks_in: torch.Tensor, + bboxes: torch.Tensor, + shape: tuple[int, int] | list[int], +) -> torch.Tensor: + """Applies masks to bounding boxes with upsampling for higher quality. + + Ref: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/ops.py#L713 + This produces higher quality masks than `process_mask` but is slower. + + Args: + protos (torch.Tensor): Prototype masks of shape [mask_dim, mask_h, mask_w]. + masks_in (torch.Tensor): Mask coefficients of shape [n, mask_dim]. + bboxes (torch.Tensor): Bounding boxes of shape [n, 4]. + shape (tuple): Target image size (h, w). + + Returns: + torch.Tensor: Upsampled and thresholded binary masks. + """ + target_shape = (int(shape[0]), int(shape[1])) + c, mh, mw = protos.shape # CHW + + # Evaluate only the prototype pixels that contribute to each retained ROI. + # The ROI interpolation preserves the original global bilinear sampling + # coordinates, so it produces the same binary mask as the full-mask path. + if _use_roi_prototype_masks(masks_in, bboxes, c, mh, mw, target_shape): + return _process_mask_upsample_roi(protos, masks_in, bboxes, target_shape) + + masks = (masks_in @ protos.float().view(c, -1)).view(-1, mh, mw) # n, CHW + masks = scale_masks(masks, target_shape) # CHW + masks = crop_mask(masks, bboxes) # CHW + return masks.gt_(0.0) + + +def _use_roi_prototype_masks( + masks_in: torch.Tensor, + bboxes: torch.Tensor, + channels: int, + proto_h: int, + proto_w: int, + shape: tuple[int, int], +) -> bool: + """Return whether exact low-resolution ROI masking is expected to be cheaper.""" + count = masks_in.shape[0] + if bboxes.numel() == 0: + return False + height, width = shape + clipped = bboxes[:, :4].clone() + clipped[:, 0::2].clamp_(0, width) + clipped[:, 1::2].clamp_(0, height) + roi_pixels = ( + (clipped[:, 2] - clipped[:, 0]).clamp_min_(0).ceil() + * (clipped[:, 3] - clipped[:, 1]).clamp_min_(0).ceil() + ).sum() + # This is a conservative upper-bound for ROI work: the actual dot product + # runs at prototype resolution, while interpolation touches only ROI pixels. + full_work = count * (channels * proto_h * proto_w + height * width) + roi_work = channels * height * width + channels * roi_pixels + return bool(roi_work < full_work) + + +def _process_mask_upsample_roi( + protos: torch.Tensor, + masks_in: torch.Tensor, + bboxes: torch.Tensor, + shape: tuple[int, int], +) -> torch.Tensor: + """Create exact cropped masks from low-resolution prototype ROIs. + + The interpolation grid uses global ``align_corners=False`` coordinates. + Therefore every target pixel samples the same prototype neighborhood as + ``scale_masks(coefficients @ protos)`` without evaluating pixels outside + its bounding box. + """ + protos = protos.float() + channels, proto_h, proto_w = protos.shape + height, width = shape + top, left, bottom, right = _mask_scale_crop_bounds((proto_h, proto_w), shape) + crop_h, crop_w = bottom - top, right - left + masks = torch.zeros( + (masks_in.shape[0], height, width), dtype=torch.float32, device=protos.device + ) + boxes = bboxes.to(protos.device) + for index, box in enumerate(boxes): + x1 = max(0, min(width, math.ceil(float(box[0])))) + y1 = max(0, min(height, math.ceil(float(box[1])))) + x2 = max(0, min(width, math.ceil(float(box[2])))) + y2 = max(0, min(height, math.ceil(float(box[3])))) + if x1 >= x2 or y1 >= y2: + continue + proto_x1 = max(0, math.floor((x1 + 0.5) * crop_w / width - 0.5)) + proto_y1 = max(0, math.floor((y1 + 0.5) * crop_h / height - 0.5)) + proto_x2 = min(crop_w, math.floor((x2 - 0.5) * crop_w / width - 0.5) + 2) + proto_y2 = min(crop_h, math.floor((y2 - 0.5) * crop_h / height - 0.5) + 2) + proto_x1, proto_x2 = left + proto_x1, left + proto_x2 + proto_y1, proto_y2 = top + proto_y1, top + proto_y2 + prototype_roi = protos[:, proto_y1:proto_y2, proto_x1:proto_x2] + lowres_mask = (masks_in[index] @ prototype_roi.reshape(channels, -1)).reshape( + 1, 1, proto_y2 - proto_y1, proto_x2 - proto_x1 + ) + ys = torch.arange(y1, y2, device=protos.device, dtype=torch.float32) + xs = torch.arange(x1, x2, device=protos.device, dtype=torch.float32) + global_y, global_x = torch.meshgrid(ys, xs, indexing="ij") + local_y = (global_y + 0.5) * crop_h / height - 0.5 - (proto_y1 - top) + local_x = (global_x + 0.5) * crop_w / width - 0.5 - (proto_x1 - left) + grid = torch.stack( + ( + (local_x + 0.5) * 2 / (proto_x2 - proto_x1) - 1, + (local_y + 0.5) * 2 / (proto_y2 - proto_y1) - 1, + ), + dim=-1, + ).unsqueeze(0) + masks[index, y1:y2, x1:x2] = F.grid_sample( + lowres_mask, + grid, + mode="bilinear", + padding_mode="border", + align_corners=False, + )[0, 0] + return masks.gt_(0.0) + + +def _mask_scale_crop_bounds( + mask_shape: tuple[int, int], target_shape: tuple[int, int] +) -> tuple[int, int, int, int]: + """Return the crop applied by :func:`scale_masks` before interpolation.""" + mask_h, mask_w = mask_shape + target_h, target_w = target_shape + gain = min(mask_h / target_h, mask_w / target_w) + pad_w = (mask_w - round(target_w * gain)) / 2 + pad_h = (mask_h - round(target_h * gain)) / 2 + top, left = round(pad_h - 0.1), round(pad_w - 0.1) + bottom, right = mask_h - round(pad_h + 0.1), mask_w - round(pad_w + 0.1) + return top, left, bottom, right + + +def crop_mask(masks: torch.Tensor, boxes: torch.Tensor) -> torch.Tensor: + """Crops masks to bounding boxes. + + Args: + masks (torch.Tensor): Masks of shape [n, h, w]. + boxes (torch.Tensor): Bounding boxes of shape [n, 4] in (x1, y1, x2, y2) format. + + Returns: + torch.Tensor: Cropped masks. + """ + if boxes.device != masks.device: + boxes = boxes.to(masks.device) + _, h, w = masks.shape + x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) + rows = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] + cols = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] + return masks * ((rows >= x1) * (rows < x2) * (cols >= y1) * (cols < y2)) + + +def scale_masks( + masks: torch.Tensor, + shape: tuple[int, int], + ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None, + padding: bool = True, +) -> torch.Tensor: + """Rescales segment masks to the target shape. + + Args: + masks (torch.Tensor): Input masks of shape (C, H, W). + shape (tuple): Target shape (height, width). + ratio_pad (tuple, optional): Pre-calculated (ratio, pad) tuple. + If None, it will be calculated from the shapes. Defaults to None. + padding (bool, optional): If True, assumes the masks were generated from + an image with YOLO-style padding. Defaults to True. + + Returns: + torch.Tensor: Rescaled masks of shape (C, target_h, target_w). + """ + im1_h, im1_w = masks.shape[1:] + im0_h, im0_w = shape[:2] + if masks.numel() == 0: + return torch.zeros((0, im0_h, im0_w), dtype=masks.dtype, device=masks.device) + if im1_h == im0_h and im1_w == im0_w: + return masks + if ratio_pad is None: # calculate from im0_shape + gain = min(im1_h / im0_h, im1_w / im0_w) # gain = old / new + pad_w, pad_h = ( + (im1_w - round(im0_w * gain)), + (im1_h - round(im0_h * gain)), + ) # wh padding + if padding: + pad_w /= 2 + pad_h /= 2 + else: + pad_w, pad_h = ratio_pad[1] + top, left = (round(pad_h - 0.1), round(pad_w - 0.1)) if padding else (0, 0) + bottom, right = im1_h - round(pad_h + 0.1), im1_w - round(pad_w + 0.1) + masks = masks[..., top:bottom, left:right] + if isinstance(masks, np.ndarray): + masks = torch.from_numpy(masks) + masks = F.interpolate( + masks[None], shape, mode="bilinear", align_corners=False + ) # 1NHW + return masks[0] + + +def to_string(counts: list[int]) -> str: + """Converts the RLE object into a compact string representation. + + Each count is delta-encoded and variable-length encoded as a string. + + Args: + counts (list[int]): List of RLE counts. + + Returns: + str: Compact string representation of the RLE object. + """ + result = [] + + for i, x in enumerate(counts): + x = int(x) + + # Apply delta encoding for all counts after the second entry + if i > 2: + x -= int(counts[i - 2]) + + # Variable-length encode the value + while True: + c = x & 0x1F # Take 5 bits + x >>= 5 + + # If the sign bit (0x10) is set, continue if x != -1; + # otherwise, continue if x != 0 + more = (x != -1) if (c & 0x10) else (x != 0) + if more: + c |= 0x20 # Set continuation bit + c += 48 # Shift to ASCII + result.append(chr(c)) + if not more: + break + + return "".join(result) + + +def multi_encode(pixels: torch.Tensor) -> list[list[int]]: + """Convert multiple binary masks using Run-Length Encoding (RLE). + + Args: + pixels (torch.Tensor): A 2D tensor where each row represents a flattened binary mask + with shape [N, H*W]. + + Returns: + list[list[int]]: A list of RLE counts for each mask. + """ + pixel_rows = pixels.detach().cpu().numpy().astype(np.uint8, copy=False) + width = pixel_rows.shape[1] + counts = [] + for i in range(pixel_rows.shape[0]): + pixel_row = pixel_rows[i] + positions = np.flatnonzero(pixel_row[1:] != pixel_row[:-1]) + 1 + if positions.size: + count = np.diff(positions).tolist() + count.insert(0, int(positions[0])) + count.append(int(width - positions[-1])) + else: + count = [width] + if pixel_row[0] == 1: + count = [0, *count] + counts.append(count) + + return counts + + +def _encode_segmentation_masks(seg_result: torch.Tensor) -> list[dict[str, Any]]: + """Threshold resized instance masks and encode them as COCO RLE objects.""" + + h, w = seg_result.shape[1:3] + binary_masks = seg_result > 0.5 + encoded_pixels = ( + binary_masks.permute(0, 2, 1) + .contiguous() + .view(binary_masks.shape[0], h * w) + .to(torch.uint8) + ) + counts = multi_encode(encoded_pixels) + if len(counts) != encoded_pixels.shape[0]: + raise RuntimeError( + f"Encoded {len(counts)} masks for a mask tensor batch of {encoded_pixels.shape[0]}." + ) + return [{"size": [h, w], "counts": to_string(count)} for count in counts] + + +def nmsout2eval( + nms_outs: list[torch.Tensor] | torch.Tensor, + img1_shape: tuple[int, int], + img0_shapes: tuple[int, int] | Sequence[tuple[int, int]], + ratio_pads: RatioPad | Sequence[RatioPad | None] | None = None, +) -> tuple[list[list[int]], list[list[list[float]]], list[list[float]]]: + """Converts NMS output to COCO evaluation format. + + Args: + nms_outs (list[torch.Tensor] | torch.Tensor): The output of the NMS + operation of shape (n, 6), where n is the number of objects. + img1_shape (tuple): Processed image shape (H, W). + img0_shapes (list[tuple]): Original image shapes [(H, W), ...]. + + Returns: + tuple: A tuple containing: + - labels (list[list]): The labels of the objects for each image. + - boxes (list[list]): The bounding boxes (xywh) for each image. + - scores (list[list]): The confidence scores for each image. + """ + + if not isinstance(nms_outs, list): + nms_outs = [nms_outs] + actual_img0_shapes = normalize_image_shapes(img0_shapes, len(nms_outs)) + actual_ratio_pads = normalize_ratio_pads(ratio_pads, len(nms_outs)) + labels_list: list[list[int]] = [] + boxes_list: list[list[list[float]]] = [] + scores_list: list[list[float]] = [] + for nms_out, img0_shape, ratio_pad in zip( + nms_outs, actual_img0_shapes, actual_ratio_pads + ): + boxes = nms_out[:, :4].clone() + scores = nms_out[:, 4] + labels = nms_out[:, 5] + valid_labels = ( + torch.isfinite(labels) + & (labels == labels.round()) + & (labels >= 0) + & (labels < 80) + ) + if not bool(valid_labels.all()): + invalid_labels = labels[~valid_labels].detach().cpu().tolist() + raise ValueError( + "COCO class IDs must be finite integral values in [0, 79]; " + f"got {invalid_labels}." + ) + boxes = scale_boxes( + img1_shape, boxes, img0_shape, ratio_pad=ratio_pad + ) # scale boxes to original image size + boxes[:, 2:] = boxes[:, 2:] - boxes[:, :2] # xyxy to xywh with corner xy + + boxes_tolist = [ + [round(float(value), 3) for value in box] for box in boxes.tolist() + ] + scores_tolist = [round(float(score), 5) for score in scores.tolist()] + labels_tolist = labels.tolist() + labels_res = [get_coco_inv(int(label)) for label in labels_tolist] + + labels_list.append(labels_res) + boxes_list.append(boxes_tolist) + scores_list.append(scores_tolist) + + return labels_list, boxes_list, scores_list + + +def nmsout2eval_seg( + nms_outs: Any, + img1_shape: tuple[int, int], + img0_shapes: tuple[int, int] | list[tuple[int, int]], + ratio_pads: RatioPad | list[RatioPad | None] | None = None, +) -> tuple[ + list[list[int]], + list[list[list[float]]], + list[list[float]], + list[list[dict[str, Any]]], +]: + """Converts segmentation NMS output to COCO evaluation format. + + Args: + nms_outs (Union[list, tuple]): Segmentation postprocess output in one of two forms: + `(det_result, seg_result)` for a single image or a list of those pairs for a batch. + img1_shape (tuple): Processed image shape (H, W). + img0_shapes (tuple | list[tuple]): Original image shape for a single image or + a list of original shapes for a batch. + + Returns: + tuple: A tuple containing: + - labels (list[list]): The labels of the objects for each image. + - boxes (list[list]): The bounding boxes (xywh) for each image. + - scores (list[list]): The confidence scores for each image. + - extra (list[list]): The encoded segmentation masks for each image. + """ + actual_img0_shapes = normalize_image_shapes(img0_shapes) + actual_ratio_pads = normalize_ratio_pads(ratio_pads, len(actual_img0_shapes)) + + if not isinstance(nms_outs[0], (list, tuple)): + actual_nms_outs = [nms_outs] + else: + actual_nms_outs = nms_outs + + det_results = [] + seg_results = [] + for nms_out in actual_nms_outs: + det_results.append(nms_out[0]) + seg_results.append(nms_out[1]) + + labels_list, boxes_list, scores_list = nmsout2eval( + det_results, + img1_shape, + actual_img0_shapes, + ratio_pads=actual_ratio_pads, + ) + + scaled_seg_results = [ + scale_masks( + seg_result.to(torch.float32), + (img0_shape[0], img0_shape[1]), + ratio_pad=ratio_pad, + ) + for seg_result, img0_shape, ratio_pad in zip( + seg_results, actual_img0_shapes, actual_ratio_pads + ) + ] + + extra_list = [ + _encode_segmentation_masks(seg_result) for seg_result in scaled_seg_results + ] + for labels, boxes, scores, extra in zip( + labels_list, boxes_list, scores_list, extra_list + ): + if not len(labels) == len(boxes) == len(scores) == len(extra): + raise RuntimeError( + "Segmentation evaluation produced mismatched label, box, score, and mask counts." + ) + return labels_list, boxes_list, scores_list, extra_list + + +def nmsout2eval_pose( + nms_outs: list[torch.Tensor] | torch.Tensor, + img1_shape: tuple[int, int], + img0_shapes: tuple[int, int] | list[tuple[int, int]], + ratio_pads: RatioPad | list[RatioPad | None] | None = None, +) -> tuple[ + list[list[int]], list[list[list[float]]], list[list[float]], list[list[list[float]]] +]: + """Converts pose estimation NMS output to COCO evaluation format. + + Args: + nms_outs (list): The output of the NMS operation. + img1_shape (tuple): Processed image shape (H, W). + img0_shapes (list[tuple]): Original image shapes [(H, W), ...]. + + Returns: + tuple: A tuple containing: + - labels (list[list]): The labels of the objects for each image. + - boxes (list[list]): The bounding boxes (xywh) for each image. + - scores (list[list]): The confidence scores for each image. + - keypoints (list[list]): The scaled keypoints for each image. + """ + actual_img0_shapes = normalize_image_shapes(img0_shapes) + actual_ratio_pads = normalize_ratio_pads(ratio_pads, len(actual_img0_shapes)) + if not isinstance(nms_outs, list): + actual_nms_outs = [nms_outs] + else: + actual_nms_outs = nms_outs + labels_list, boxes_list, scores_list = nmsout2eval( + actual_nms_outs, + img1_shape, + actual_img0_shapes, + ratio_pads=actual_ratio_pads, + ) + extra = [ + scale_coords( + img1_shape, + nms_out[:, 6:].reshape(-1, 17, 3), + img0_shape, + ratio_pad=ratio_pad, + ).reshape(-1, 51) + for nms_out, img0_shape, ratio_pad in zip( + actual_nms_outs, actual_img0_shapes, actual_ratio_pads + ) + ] + return labels_list, boxes_list, scores_list, [x.tolist() for x in extra] + + +def nmsout2eval_obb( + nms_outs: list[torch.Tensor] | torch.Tensor, + img1_shape: tuple[int, int], + img0_shapes: tuple[int, int] | list[tuple[int, int]], + ratio_pads: RatioPad | list[RatioPad | None] | None = None, + include_xywhr: bool = False, +) -> tuple[Any, ...]: + """Converts OBB NMS output to DOTAv1 evaluation format. + + Args: + nms_outs: Detections with rows ``cx, cy, w, h, score, cls, angle``. + img1_shape: Processed image shape. + img0_shapes: Original image shape or shapes. + ratio_pads: Optional letterbox metadata. + include_xywhr: Whether to include scaled ``xywhr`` boxes in the return value. + + Returns: + DOTAv1 labels, polygons, scores, and optionally scaled ``xywhr`` boxes. + """ + actual_img0_shapes = normalize_image_shapes(img0_shapes) + actual_ratio_pads = normalize_ratio_pads(ratio_pads, len(actual_img0_shapes)) + actual_nms_outs = [nms_outs] if not isinstance(nms_outs, list) else nms_outs + + labels_list: list[list[str]] = [] + polygons_list: list[list[list[float]]] = [] + scores_list: list[list[float]] = [] + xywhr_list: list[list[list[float]]] = [] + for nms_out, img0_shape, ratio_pad in zip( + actual_nms_outs, actual_img0_shapes, actual_ratio_pads + ): + if nms_out.numel() == 0: + labels_list.append([]) + polygons_list.append([]) + scores_list.append([]) + xywhr_list.append([]) + continue + + rboxes = torch.cat([nms_out[:, :4], nms_out[:, 6:7]], dim=-1) + rboxes = scale_rboxes(img1_shape, rboxes, img0_shape, ratio_pad=ratio_pad) + polygons = xywhr2xyxyxyxy(rboxes).reshape(-1, 8) + polygons = scale_coords( + img0_shape, polygons.reshape(-1, 4, 2), img0_shape + ).reshape(-1, 8) + + labels = [get_dotav1_label(int(label)) for label in nms_out[:, 5].tolist()] + scores = [round(float(score), 5) for score in nms_out[:, 4].tolist()] + polygons_tolist = [ + [round(float(value), 3) for value in polygon] + for polygon in polygons.tolist() + ] + xywhr_tolist = [ + [round(float(value), 3) for value in rbox] for rbox in rboxes.tolist() + ] + + labels_list.append(labels) + polygons_list.append(polygons_tolist) + scores_list.append(scores) + xywhr_list.append(xywhr_tolist) + + if include_xywhr: + return labels_list, polygons_list, scores_list, xywhr_list + return labels_list, polygons_list, scores_list + + +class YOLOSegPostMixin: + """Mixin class for YOLO segmentation postprocessing.""" + + def nmsout2eval( + self, + nms_out: Any, + img1_shape: tuple[int, int], + img0_shape: tuple[int, int] | list[tuple[int, int]], + ratio_pad: RatioPad | list[RatioPad | None] | None = None, + ) -> tuple[Any, ...]: + """Converts NMS output to evaluation format for segmentation. + + Args: + nms_out: NMS output (detections and prototypes). + img1_shape: Resized image shape. + img0_shape: List of original image shapes. + + Returns: + Tuple: (labels_list, boxes_list, scores_list, extra_list). + """ + return nmsout2eval_seg(nms_out, img1_shape, img0_shape, ratio_pads=ratio_pad) + + +class YOLOPosePostMixin: + """Mixin class for YOLO pose estimation postprocessing.""" + + def nmsout2eval( + self, + nms_out: Any, + img1_shape: tuple[int, int], + img0_shape: tuple[int, int] | list[tuple[int, int]], + ratio_pad: RatioPad | list[RatioPad | None] | None = None, + ) -> tuple[Any, ...]: + """Converts NMS output to evaluation format for pose estimation. + + Args: + nms_out: NMS output (detections with keypoints). + img1_shape: Resized image shape. + img0_shape: List of original image shapes. + + Returns: + Tuple: (labels_list, boxes_list, scores_list, extra_list). + """ + return nmsout2eval_pose(nms_out, img1_shape, img0_shape, ratio_pads=ratio_pad) + + +class YOLOOBBPostMixin: + """Mixin class for YOLO oriented-bounding-box postprocessing.""" + + def nmsout2eval( + self, + nms_out: Any, + img1_shape: tuple[int, int], + img0_shape: tuple[int, int] | list[tuple[int, int]], + ratio_pad: RatioPad | list[RatioPad | None] | None = None, + include_xywhr: bool = False, + ) -> tuple[Any, ...]: + """Converts OBB detections to DOTAv1 labels, polygons, and scores. + + Args: + nms_out: NMS output with rows ``cx, cy, w, h, score, cls, angle``. + img1_shape: Resized image shape. + img0_shape: Original image shape or shapes. + ratio_pad: Optional letterbox metadata. + include_xywhr: Whether to include scaled rotated boxes. + + Returns: + DOTAv1 labels, polygons, scores, and optionally scaled ``xywhr`` boxes. + """ + return nmsout2eval_obb( + nms_out, + img1_shape, + img0_shape, + ratio_pads=ratio_pad, + include_xywhr=include_xywhr, + ) diff --git a/mblt_vision/utils/postprocess/depth_post.py b/mblt_vision/utils/postprocess/depth_post.py new file mode 100644 index 0000000..ee7dee0 --- /dev/null +++ b/mblt_vision/utils/postprocess/depth_post.py @@ -0,0 +1,106 @@ +"""Postprocessing for monocular depth-estimation models.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch +import torch.nn.functional as functional + +from ..letterbox import RatioPad +from ..types import ListTensorLike, TensorLike +from ._letterbox import crop_letterbox, get_letterbox_input_shape, resolve_ratio_pads +from .base import PostBase +from .common import normalize_image_shapes + + +class DepthPost(PostBase): + """Normalize depth outputs and undo letterbox padding when metadata is available.""" + + def __init__(self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any]) -> None: + """Initialize depth restoration from the model letterbox configuration.""" + + super().__init__() + del post_cfg + self.input_shape = get_letterbox_input_shape( + pre_cfg, "Depth estimation", "Depth" + ) + + def __call__( + self, + x: TensorLike | ListTensorLike, + img0_shape: tuple[int, int] | Sequence[tuple[int, int]] | None = None, + ratio_pad: RatioPad | Sequence[RatioPad | None] | None = None, + **kwargs: Any, + ) -> torch.Tensor | list[torch.Tensor]: + """Return normalized depth maps, optionally restored to original image sizes.""" + + if kwargs: + raise TypeError( + f"Unexpected depth postprocess kwargs: {', '.join(sorted(kwargs))}" + ) + depth = self._normalize_output(x) + if img0_shape is None: + return depth + + shapes = normalize_image_shapes(img0_shape, depth.shape[0]) + pads = resolve_ratio_pads(ratio_pad, depth.shape[0], shapes, self.input_shape) + restored = [ + self._restore(depth[index], shapes[index], pads[index]) + for index in range(depth.shape[0]) + ] + return restored[0] if len(restored) == 1 else restored + + def _normalize_output(self, x: TensorLike | ListTensorLike) -> torch.Tensor: + """Validate dense output and normalize it to input-sized ``[B, H, W]`` tensors.""" + + if isinstance(x, (list, tuple)): + if len(x) != 1: + raise ValueError( + f"Depth estimation expects one output tensor, received {len(x)}." + ) + x = x[0] + depth = torch.as_tensor(x) + if depth.ndim == 4: + if depth.shape[1] == 1: + depth = depth[:, 0] + elif depth.shape[-1] == 1: + depth = depth[..., 0] + else: + raise ValueError( + f"Depth estimation expects [B, 1, H, W] or [B, H, W, 1], got {tuple(depth.shape)}." + ) + elif depth.ndim == 3 and depth.shape[-1] == 1: + depth = depth[..., 0].unsqueeze(0) + elif depth.ndim != 3: + raise ValueError( + "Depth estimation expects [B, H, W], [B, 1, H, W], [H, W, 1], or [B, H, W, 1], " + f"got {tuple(depth.shape)}." + ) + depth = depth.to(device=self.device, dtype=torch.float32) + if not bool(torch.isfinite(depth).all()): + raise ValueError("Depth estimation output must contain only finite values.") + if tuple(depth.shape[-2:]) == self.input_shape: + return depth + + quarter_shape = tuple(dimension // 4 for dimension in self.input_shape) + if tuple(depth.shape[-2:]) == quarter_shape: + return functional.interpolate( + depth[:, None], scale_factor=4.0, mode="bilinear", align_corners=False + )[:, 0] + + raise ValueError( + f"Depth estimation output spatial shape must be {self.input_shape} or quarter-resolution {quarter_shape}, " + f"got {tuple(depth.shape[-2:])}." + ) + + def _restore( + self, depth: torch.Tensor, shape: tuple[int, int], ratio_pad: RatioPad + ) -> torch.Tensor: + """Crop padded depth pixels and bilinearly resize to an original image shape.""" + + cropped = crop_letterbox(depth, shape, ratio_pad, self.input_shape, "Depth") + return functional.interpolate( + cropped[None, None], size=shape, mode="bilinear", align_corners=False + )[0, 0] diff --git a/mblt_vision/utils/postprocess/semantic_seg_post.py b/mblt_vision/utils/postprocess/semantic_seg_post.py new file mode 100644 index 0000000..03f43ef --- /dev/null +++ b/mblt_vision/utils/postprocess/semantic_seg_post.py @@ -0,0 +1,192 @@ +"""Postprocessing for semantic-segmentation models.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch +import torch.nn.functional as functional + +from ..letterbox import RatioPad +from ..types import ListTensorLike, TensorLike +from ._letterbox import crop_letterbox, get_letterbox_input_shape, resolve_ratio_pads +from .base import PostBase +from .common import normalize_image_shapes + + +class SemanticSegPost(PostBase): + """Convert semantic logits to class maps and undo letterbox padding.""" + + NC_BY_DATASET: dict[str, int] = { + "ade20k": 150, + "cityscapes": 19, + } + + def __init__(self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any]) -> None: + """Initialize semantic output handling for a configured dataset taxonomy.""" + + super().__init__() + self.input_shape = get_letterbox_input_shape( + pre_cfg, "Semantic segmentation", "Semantic" + ) + dataset = post_cfg.get("dataset") + if not isinstance(dataset, str): + raise ValueError( + "Semantic segmentation requires a string dataset in post_cfg." + ) + self.dataset = dataset.lower() + dataset_nc = self.NC_BY_DATASET.get(self.dataset) + configured_nc = post_cfg.get("nc") + if configured_nc is None: + if dataset_nc is None: + raise ValueError( + f"Semantic segmentation requires nc for unknown dataset '{self.dataset}'." + ) + self.nc = dataset_nc + else: + self.nc = int(configured_nc) + if dataset_nc is not None and self.nc != dataset_nc: + raise ValueError( + f"nc={configured_nc} conflicts with semantic dataset '{self.dataset}', which requires nc={dataset_nc}." + ) + + def __call__( + self, + x: TensorLike | ListTensorLike, + img0_shape: tuple[int, int] | Sequence[tuple[int, int]] | None = None, + ratio_pad: RatioPad | Sequence[RatioPad | None] | None = None, + **kwargs: Any, + ) -> torch.Tensor | list[torch.Tensor]: + """Return class maps, optionally restored to original image sizes.""" + + if kwargs: + raise TypeError( + f"Unexpected semantic postprocess kwargs: {', '.join(sorted(kwargs))}" + ) + output, is_logits = self._normalize_output(x) + if img0_shape is None: + return self._to_input_space(output, is_logits) + + shapes = normalize_image_shapes(img0_shape, output.shape[0]) + pads = resolve_ratio_pads(ratio_pad, output.shape[0], shapes, self.input_shape) + restored = [ + self._restore(output[index], is_logits, shapes[index], pads[index]) + for index in range(output.shape[0]) + ] + return restored[0] if len(restored) == 1 else restored + + def _normalize_output( + self, x: TensorLike | ListTensorLike + ) -> tuple[torch.Tensor, bool]: + """Validate logits in NCHW/NHWC layout or a baked class-map tensor.""" + + if isinstance(x, (list, tuple)): + if len(x) != 1: + raise ValueError( + f"Semantic segmentation expects one output tensor, received {len(x)}." + ) + x = x[0] + output = torch.as_tensor(x, device=self.device) + if output.ndim == 4: + if output.shape[1] == self.nc: + return self._validate_logits(output), True + if output.shape[-1] == self.nc: + return self._validate_logits(output.permute(0, 3, 1, 2)), True + raise ValueError( + f"Semantic segmentation for '{self.dataset}' expects [B, {self.nc}, H, W] or " + f"[B, H, W, {self.nc}] logits, got {tuple(output.shape)}." + ) + if output.ndim == 3: + if output.shape[-1] == self.nc and output.is_floating_point(): + return self._validate_logits(output.permute(2, 0, 1).unsqueeze(0)), True + if tuple(output.shape[:2]) == self.input_shape: + raise ValueError( + f"Semantic segmentation for '{self.dataset}' expects [H, W, {self.nc}] MXQ logits, " + f"got {tuple(output.shape)}." + ) + if output.is_complex(): + raise ValueError("Semantic class-map values must be finite integers.") + if output.is_floating_point(): + if not bool(torch.isfinite(output).all()): + raise ValueError("Semantic class-map values must be finite.") + if not bool(torch.eq(output, output.trunc()).all()): + raise ValueError( + "Semantic class-map values must be integer-valued." + ) + if output.numel() and ( + int(output.min()) < 0 or int(output.max()) >= self.nc + ): + raise ValueError( + f"Semantic class-map values must be in [0, {self.nc - 1}]." + ) + return output.to(dtype=torch.int64), False + raise ValueError( + f"Semantic segmentation expects [B, C, H, W] or [B, H, W, C] logits, or [B, H, W] class maps, " + f"got {tuple(output.shape)}." + ) + + @staticmethod + def _validate_logits(output: torch.Tensor) -> torch.Tensor: + """Convert semantic logits to float while rejecting invalid artifact output.""" + + logits = output.to(dtype=torch.float32) + if not bool(torch.isfinite(logits).all()): + raise ValueError("Semantic logits must contain only finite values.") + return logits + + def _to_input_space(self, output: torch.Tensor, is_logits: bool) -> torch.Tensor: + """Restore model output resolution to configured input space.""" + + if is_logits: + if tuple(output.shape[-2:]) != self.input_shape: + output = functional.interpolate( + output, size=self.input_shape, mode="bilinear", align_corners=False + ) + return output.argmax(dim=1).to(dtype=torch.int64) + if tuple(output.shape[-2:]) != self.input_shape: + output = functional.interpolate( + output[:, None].float(), size=self.input_shape, mode="nearest" + )[:, 0] + return output.to(dtype=torch.int64) + + def _restore( + self, + output: torch.Tensor, + is_logits: bool, + shape: tuple[int, int], + ratio_pad: RatioPad, + ) -> torch.Tensor: + """Undo letterboxing, preserving logits until after bilinear restoration.""" + + if is_logits: + if tuple(output.shape[-2:]) != self.input_shape: + output = functional.interpolate( + output[None], + size=self.input_shape, + mode="bilinear", + align_corners=False, + )[0] + channels = [ + crop_letterbox(channel, shape, ratio_pad, self.input_shape, "Semantic") + for channel in output + ] + cropped_logits = torch.stack(channels) + restored_logits = functional.interpolate( + cropped_logits[None], + size=shape, + mode="bilinear", + align_corners=False, + )[0] + return restored_logits.argmax(dim=0).to(dtype=torch.int64) + + if tuple(output.shape[-2:]) != self.input_shape: + output = functional.interpolate( + output[None, None].float(), + size=self.input_shape, + mode="nearest", + )[0, 0].to(dtype=torch.int64) + cropped = crop_letterbox(output, shape, ratio_pad, self.input_shape, "Semantic") + return functional.interpolate( + cropped[None, None].float(), size=shape, mode="nearest" + )[0, 0].to(torch.int64) diff --git a/mblt_vision/utils/postprocess/yolo_anchor_post.py b/mblt_vision/utils/postprocess/yolo_anchor_post.py new file mode 100644 index 0000000..3848cd4 --- /dev/null +++ b/mblt_vision/utils/postprocess/yolo_anchor_post.py @@ -0,0 +1,469 @@ +""" +YOLO anchor-based postprocessing. +""" + +from __future__ import annotations + +from typing import Any, cast + +import torch + +from .base import YOLODetectionPostBase +from .common import YOLOSegPostMixin, non_max_suppression + + +class YOLOAnchorDetectionPost(YOLODetectionPostBase): + """Postprocessing for YOLO models with anchors.""" + + def __init__( + self, pre_cfg: dict[str, Any], post_cfg: dict[str, Any], **kwargs: Any + ) -> None: + """Initialize anchor-based YOLO detection postprocessing. + + Args: + pre_cfg (dict): Preprocessing configuration. + post_cfg (dict): Postprocessing configuration. + **kwargs: Optional runtime overrides for postprocess behavior. + """ + super().__init__(pre_cfg, post_cfg, **kwargs) + self.no = self.nc + 5 + self.n_extra + self.grid: torch.Tensor + self.anchor_grid: torch.Tensor + self.make_anchor_grid() + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return the export-style output tensor for anchor-based YOLO models.""" + if len(x) == 1: + converted = self.conversion(x) + if isinstance(converted, torch.Tensor): + return self._converted_to_batch_output(converted) + det_out, proto_out = converted + return [self._converted_to_batch_output(det_out), proto_out] + + rearranged = self.rearrange(x) + if isinstance(rearranged, tuple): + det_out, proto_out = rearranged + return [self.decode_batch(det_out), proto_out.permute(0, 3, 1, 2)] + return self.decode_batch(rearranged) + + def _converted_to_batch_output(self, x: torch.Tensor) -> torch.Tensor: + """Normalize converted outputs to the export-style batched layout.""" + while x.ndim == 4 and 1 in (x.shape[0], x.shape[1]): + if x.shape[0] == 1: + x = x.squeeze(0) + elif x.shape[1] == 1: + x = x.squeeze(1) + if x.ndim != 3: + raise ValueError( + f"Expected 3D converted tensor, got shape {tuple(x.shape)}." + ) + if x.shape[-1] == self.no: + return x + if x.shape[1] == self.no: + return x.transpose(1, 2) + raise ValueError( + f"Unsupported converted tensor shape {tuple(x.shape)} for non-e2e output." + ) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every anchor without filtering and preserve batch shape.""" + batch_size = x.shape[0] + grid = self.grid.unsqueeze(0).expand(batch_size, -1, -1) + anchor_grid = self.anchor_grid.unsqueeze(0).expand(batch_size, -1, -1) + stride = self.stride_as_tensor().unsqueeze(0).expand(batch_size, -1, -1) + + decoded = x.clone() + decoded[..., :2] = ( + decoded[..., :2].sigmoid().mul(2.0).add(grid).add(-0.5).mul(stride) + ) + decoded[..., 2:4] = ( + decoded[..., 2:4].sigmoid().mul(2.0).pow(2.0).mul(anchor_grid) + ) + conf = decoded[..., 4:5].sigmoid() + decoded[..., 4:5] = conf + decoded[..., 5 : 5 + self.nc] = decoded[..., 5 : 5 + self.nc].sigmoid() + if self.task == "instance_segmentation" and self.n_extra > 0: + decoded[..., 5 + self.nc :] = decoded[..., 5 + self.nc :] * conf + return decoded + + def rearrange( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Rearranges raw model output tensors into a concatenated decode input. + + Args: + x (list[torch.Tensor]): Raw output tensors from the model detection heads. + + Returns: + torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Concatenated tensor in + ``(batch, anchors, no)`` format, optionally paired with prototype masks in + segmentation subclasses. + """ + if len(x) != self.nl: + raise ValueError(f"Expected {self.nl} detection heads, got {len(x)}.") + y = [] + for i in range(self.nl): + tmp = x[i] + if tmp.shape[3] == self.no * self.na: + y.append( + tmp.permute(0, 3, 1, 2) + ) # (b, 80, 80, 255) -> (b, 255, 80, 80) + else: + raise NotImplementedError( + f"Got unsupported shape for input: {tmp.shape}." + ) + # sort by image size descending + y = sorted(y, key=lambda x: x.numel(), reverse=True) + return torch.cat( + [ + xi.reshape(xi.shape[0], self.na, self.no, xi.shape[-2], xi.shape[-1]) + .permute(0, 1, 3, 4, 2) + .reshape(xi.shape[0], -1, self.no) + for xi in y + ], + dim=1, + ) + + def decode(self, x: torch.Tensor) -> list[torch.Tensor]: + """Decodes model outputs into box coordinates and class scores. + + Applies sigmoid to predictions and transforms boxes from anchor-relative + to image-relative coordinates. + + Args: + x (torch.Tensor): Concatenated output tensor from `rearrange`. + + Returns: + list[torch.Tensor]: Per-image decoded detections after confidence filtering. + """ + return [self.process_box_cls(box_cls) for box_cls in x] + + def process_box_cls(self, x: torch.Tensor) -> torch.Tensor: + """Processes a single image's detection tensor. + + Args: + x: Raw detections for one image. + + Returns: + Decoded boxes, confidence, and scores. + """ + ic = x[:, 4] > self.inv_conf_thres # candidates + box_cls = x[ic] # (n, 85) + if box_cls.numel() == 0: + return box_cls.new_zeros((0, 5 + self.nc + self.n_extra)) + + grid = self.grid[ic, :] # (n, 2) + anchor_grid = self.anchor_grid[ic, :] # (n, 2) + stride = self.stride_as_tensor()[ic, :] # (n, 2) + + # Advanced indexing above materializes ``box_cls``, so in-place decode avoids a second output allocation. + box_cls[:, :2] = ( + box_cls[:, :2].sigmoid_().mul_(2.0).add_(grid).add_(-0.5).mul_(stride) + ) + box_cls[:, 2:4] = ( + box_cls[:, 2:4].sigmoid_().mul_(2.0).pow_(2.0).mul_(anchor_grid) + ) + conf = box_cls[:, 4:5].sigmoid_() + box_cls[:, 5 : 5 + self.nc].sigmoid_() + if self.task == "instance_segmentation" and self.n_extra > 0: + box_cls[:, 5 + self.nc :] *= conf + return box_cls + + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters out low-confidence detections from a single concatenated output tensor. + + Args: + x (torch.Tensor): Concatenated output tensor from the model. + + Returns: + list[torch.Tensor]: Filtered detections for each image in the batch. + """ + x_list = torch.split( + self._converted_to_batch_output(x), 1, dim=0 + ) # [(1, 25200, 85), ...] + + def process_conversion(x: torch.Tensor) -> torch.Tensor: + x = x.squeeze(0) # (25200, 85) + ic = x[:, 4] > self.conf_thres # candidates + x = x[ic] # (n, 85) + if len(x) == 0: + return x.new_zeros((0, self.no)) + return x + + return [process_conversion(xi) for xi in x_list] + + def _nms_single( + self, + xi: torch.Tensor, + max_det: int, + max_nms: int, + max_wh: int, + *, + multi_label: bool, + ) -> torch.Tensor: + """Apply anchor-based NMS to a single decoded image tensor.""" + mi = 5 + self.nc # mask index + if xi.numel() == 0: + return xi.new_zeros((0, 6 + self.n_extra)) + + scores = xi[:, 5:mi] * xi[:, 4:5] + if multi_label: + match_index = (scores > self.conf_thres).nonzero(as_tuple=False) + if match_index.numel() == 0: + return xi.new_zeros((0, 6 + self.n_extra)) + i, j = match_index[:, 0], match_index[:, 1] + rows = xi[i] + row_scores = scores[i, j] + else: + row_scores, j = scores.max(dim=1) + keep = row_scores > self.conf_thres + if not bool(keep.any()): + return xi.new_zeros((0, 6 + self.n_extra)) + rows, row_scores, j = xi[keep], row_scores[keep], j[keep] + boxes_xywh = rows[:, :4] + out = torch.empty( + (rows.shape[0], 6 + self.n_extra), dtype=rows.dtype, device=rows.device + ) + out[:, 0] = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2 + out[:, 1] = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2 + out[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2 + out[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2 + out[:, 4] = row_scores + out[:, 5] = j.to(rows.dtype) + if self.n_extra > 0: + out[:, 6:] = rows[:, mi:] + out = out[out[:, 4].argsort(descending=True)[:max_nms]] + c = out[:, 5:6] * max_wh + boxes, score = out[:, :4] + c, out[:, 4] + i_idx = non_max_suppression(boxes, score, self.iou_thres, max_det) + return out[i_idx] + + def nms( + self, + x: torch.Tensor | list[torch.Tensor], + max_det: int = 300, + max_nms: int = 30000, + max_wh: int = 7680, + multi_label: bool = False, + ) -> list[torch.Tensor]: + """ + Perform Non-Maximum Suppression (NMS) on the decoded detections. + Args: + x (list[torch.Tensor]): Decoded detections for each image. + max_det (int, optional): Maximum number of detections to keep. Defaults to 300. + max_nms (int, optional): Maximum number of candidates to consider for NMS. + Defaults to 30000. + max_wh (int, optional): Maximum box width/height for offset calculation. + Defaults to 7680. + Returns: + list[torch.Tensor]: Post-NMS detections for each image. + """ + if isinstance(x, list): + return [ + self._nms_single( + xi, + max_det=max_det, + max_nms=max_nms, + max_wh=max_wh, + multi_label=multi_label, + ) + for xi in x + ] + return [ + self._nms_single( + xi, + max_det=max_det, + max_nms=max_nms, + max_wh=max_wh, + multi_label=multi_label, + ) + for xi in x + ] + + def nms_multilabel( + self, x: torch.Tensor | list[torch.Tensor] + ) -> list[torch.Tensor]: + """Perform Ultralytics-compatible multi-label NMS for validation.""" + + return self.nms(x, multi_label=True) + + def make_anchor_grid(self) -> None: + """ + Pre-calculate the anchor grid for decoding. + """ + grid_parts: list[torch.Tensor] = [] + anchor_grid_parts: list[torch.Tensor] = [] + stride_parts: list[torch.Tensor] = [] + strides = [2 ** (3 + i) for i in range(self.nl)] + if self.nl == 2: + strides = [strd * 2 for strd in strides] + out_sizes = [ + [self.imh // strd, self.imw // strd] for strd in strides + ] # (80, 80), (40, 40), (20, 20) + for anchr, (ny, nx), strd in zip(self.anchors_as_list(), out_sizes, strides): + yv, xv = torch.meshgrid( + torch.arange(ny, dtype=torch.float32, device=self.device), + torch.arange(nx, dtype=torch.float32, device=self.device), + indexing="ij", + ) + grid = torch.stack((xv, yv), 2).expand(self.na, ny, nx, 2) + grid_parts.append(grid) + anchr_tensor = torch.broadcast_to( + torch.tensor(anchr).reshape(self.na, 1, 1, 2), + (self.na, ny, nx, 2), + ) + anchor_grid_parts.append(anchr_tensor) + stride_parts.append(strd * torch.ones(self.na, ny, nx, 2)) + self.grid = torch.cat([grd.reshape(-1, 2) for grd in grid_parts], dim=0) + self.anchor_grid = torch.cat( + [anc.reshape(-1, 2) for anc in anchor_grid_parts], dim=0 + ) + self.stride = torch.cat([strd.reshape(-1, 2) for strd in stride_parts], dim=0) + + def chop(self, npu_out: torch.Tensor, idx: int = 0) -> tuple[torch.Tensor, ...]: + """Splits the detection tensor into individual components (xy, wh, conf, scores, extra). + + Args: + npu_out (torch.Tensor): Raw detection tensor from one detection head. + idx (int, optional): Detection head index. Defaults to 0. + + Returns: + tuple: (xy, wh, conf, scores, extra). + """ + xy, wh, conf, scores, extra = torch.split( + npu_out, [2, 2, 1, self.nc, self.n_extra], dim=-1 + ) + return xy, wh, conf, scores, extra + + +class YOLOAnchorSegPost(YOLOSegPostMixin, YOLOAnchorDetectionPost): + """Postprocessing for YOLO segmentation models with anchors.""" + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return the export-style output tensor for anchor-based YOLO segmentation models. + + Args: + x: Checked raw model outputs. + + Returns: + A detection tensor, or detections paired with prototype masks. + """ + if any(xi.ndim <= 4 and self.no in xi.shape[1:] for xi in x): + converted, proto_outs = self.conversion(x) + return [ + self._converted_to_batch_output(converted), + self._proto_to_nchw(proto_outs), + ] + return super().non_e2e(x) + + def _proto_to_nchw(self, proto: torch.Tensor) -> torch.Tensor: + """Convert prototype tensors to ``(B, C, H, W)`` if needed. + + Args: + proto: Prototype tensor from a model runtime. + + Returns: + Prototype tensor in channel-first batch layout. + """ + if proto.ndim == 4 and proto.shape[1] == self.n_extra: + return proto + if proto.ndim == 4 and proto.shape[-1] == self.n_extra: + return proto.permute(0, 3, 1, 2) + raise ValueError( + f"Unsupported proto tensor shape {tuple(proto.shape)} for non-e2e output." + ) + + def _pre_process(self, x: list[torch.Tensor]) -> tuple[Any, torch.Tensor | None]: + """Preprocesses intermediate inputs into (boxes, proto) format. + + Args: + x (list[torch.Tensor]): Raw model output tensors. + + Returns: + tuple: (decoded_detections, prototype_masks). + """ + if any(xi.ndim <= 4 and self.no in xi.shape[1:] for xi in x): + converted, proto_outs = cast( + tuple[torch.Tensor, torch.Tensor], self.conversion(x) + ) + return self.filter_conversion(converted), proto_outs + rearranged, proto_outs = self.rearrange(x) + return self.decode(rearranged), proto_outs + + def conversion(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Converts raw model output tensors into detections and prototypes. + + Args: + x (list[torch.Tensor]): List of raw output tensors. + + Returns: + tuple: (detections, prototypes) + """ + det_out: torch.Tensor | None = None + proto_out: torch.Tensor | None = None + for xi in x: + if xi.ndim <= 4 and self.no in xi.shape[1:]: + det_out = xi + elif xi.ndim == 4 and self.n_extra in xi.shape[1:]: + proto_out = xi + if det_out is None or proto_out is None: + shapes = ", ".join(str(tuple(xi.shape)) for xi in x) + raise NotImplementedError( + f"Input shapes not supported for anchor segmentation: {shapes}." + ) + return det_out, proto_out + + def rearrange(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Rearranges model output tensors for segmentation tasks. + + Args: + x (list[torch.Tensor]): Raw output tensors from detection and prototype heads. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Concatenated detections and prototype masks. + """ + proto: torch.Tensor | None = None + for i, xi in enumerate(x): + if self.n_extra == xi.shape[-1]: + proto = x.pop(i) + break + if proto is None: + raise ValueError("Proto output is missing.") + y = [] + for xi in x: + if xi.shape[-1] == self.no * self.nl: + y.append(xi.permute(0, 3, 1, 2)) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort by image size descending + y = sorted(y, key=lambda x: x.numel(), reverse=True) + return ( + torch.cat( + [ + xi.reshape( + xi.shape[0], self.na, self.no, xi.shape[-2], xi.shape[-1] + ) + .permute(0, 1, 3, 4, 2) + .reshape(xi.shape[0], -1, self.no) + for xi in y + ], + dim=1, + ), + proto, + ) + + def chop(self, npu_out: torch.Tensor, idx: int = 0) -> tuple[torch.Tensor, ...]: + """Splits the detection tensor for segmentation tasks. + + Args: + npu_out (torch.Tensor): Raw detection tensor. + idx (int, optional): Detection head index. Defaults to 0. + + Returns: + tuple: (xy, wh, conf, scores, masks). + """ + xy, wh, conf, scores, masks = torch.split( + npu_out, [2, 2, 1, self.nc, self.n_extra], dim=-1 + ) + masks = masks * conf.sigmoid() + return xy, wh, conf, scores, masks diff --git a/mblt_vision/utils/postprocess/yolo_anchorless_post.py b/mblt_vision/utils/postprocess/yolo_anchorless_post.py new file mode 100644 index 0000000..cecbd72 --- /dev/null +++ b/mblt_vision/utils/postprocess/yolo_anchorless_post.py @@ -0,0 +1,952 @@ +""" +YOLO anchorless postprocessing. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import torch + +from .base import YOLODetectionPostBase +from .common import ( + YOLOOBBPostMixin, + YOLOPosePostMixin, + YOLOSegPostMixin, + concat_converted_obb_outputs, + decode_split_converted_obb_outputs, + dist2bbox, + dist2rbox, + non_max_suppression, + rotated_nms, + xywh2xyxy, + yolo_multilabel_candidates, +) + +AnchorlessOutputLayout = Literal["channels_first", "candidates_first"] + + +@dataclass(frozen=True) +class _AnchorlessNMSInput: + """Decoded anchorless detections together with their source layout.""" + + detections: torch.Tensor | list[torch.Tensor] + layout: AnchorlessOutputLayout + + +class YOLOAnchorlessDetectionPost(YOLODetectionPostBase): + """Postprocessing for YOLO models without anchors.""" + + def __init__(self, pre_cfg: dict, post_cfg: dict, **kwargs: object) -> None: + """Initialize the anchorless YOLO postprocessor. + + Args: + pre_cfg: Preprocessing configuration. + post_cfg: Postprocessing configuration. + **kwargs: Optional runtime overrides for postprocess behavior. + """ + super().__init__(pre_cfg, post_cfg, **kwargs) + self.reg_max = post_cfg.get("reg_max", 0) # DFL channels + self.no = self.nc + self.reg_max * 4 # number of outputs per anchor (144) + self.dfl_weight = torch.arange( + self.reg_max, dtype=torch.float32, device=self.device + ).reshape(1, -1, 1, 1) + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return the export-style output tensor for anchorless YOLO models.""" + if len(x) == 1: + converted = self.conversion(x) + if isinstance(converted, torch.Tensor): + return self._converted_to_batch_output(converted) + det_out, proto_out = converted + return [self._converted_to_batch_output(det_out), proto_out] + + rearranged = self.rearrange(x) + if isinstance(rearranged, tuple): + det_out, proto_out = rearranged + return [self.decode_batch(det_out), proto_out.permute(0, 3, 1, 2)] + return self.decode_batch(rearranged) + + def _converted_to_batch_output(self, x: torch.Tensor) -> torch.Tensor: + """Normalize converted outputs to the export-style batched layout.""" + if x.ndim == 4 and x.shape[1] == 1: + x = x.squeeze(1) + if x.ndim != 3: + raise ValueError( + f"Expected 3D converted tensor, got shape {tuple(x.shape)}." + ) + if x.shape[1] == 4 + self.nc + self.n_extra: + return x + if x.shape[-1] == 4 + self.nc + self.n_extra: + return x.transpose(1, 2) + raise ValueError( + f"Unsupported converted tensor shape {tuple(x.shape)} for non-e2e output." + ) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every anchor without confidence filtering for export-style output.""" + box, scores, extra = torch.split( + x, [self.reg_max * 4, self.nc, self.n_extra], dim=1 + ) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + dbox = dist2bbox(self.dfl(box), anchors, xywh=False, dim=1) * stride + return torch.cat([dbox, scores.sigmoid(), extra], dim=1) + + def rearrange( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Rearranges raw model output tensors into a concatenated decode input. + + Args: + x (list[torch.Tensor]): List of raw output tensors from the model detection heads. + + Returns: + torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Concatenated tensor in + ``(batch, channels, anchors)`` format, optionally paired with prototype masks. + """ + y_det = [] + y_cls = [] + for xi in x: # list of bchw outputs + if xi.ndim == 3: + xi = xi[None] + elif xi.ndim == 4: + pass + else: + raise NotImplementedError(f"Got unsupported ndim for input: {xi.ndim}.") + if xi.shape[-1] == self.reg_max * 4: + y_det.append( + xi.permute(0, 3, 1, 2) + ) # (b, 64, 80, 80), (b, 64 ,40, 40), ... + elif xi.shape[-1] == self.nc: + y_cls.append( + xi.permute(0, 3, 1, 2) + ) # (b, 80, 80, 80), (b, 80, 40, 40), ... + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort as box, scores + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + self.validate_split_head_counts(detection=y_det, classification=y_cls) + return torch.cat( + [ + torch.cat((yi_det, yi_cls), dim=1).flatten(2) + for yi_det, yi_cls in zip(y_det, y_cls) + ], + dim=-1, + ) + + def decode(self, x: torch.Tensor) -> list[torch.Tensor]: + """Decodes model outputs into box coordinates and class scores. + + Args: + x (torch.Tensor): Concatenated output tensor from `rearrange`. + + Returns: + list[torch.Tensor]: Per-image decoded detections in ``(channels, anchors)`` format. + """ + return [self.process_box_cls(box_cls) for box_cls in x] + + def _pre_process( + self, + x: list[torch.Tensor], + ) -> tuple[_AnchorlessNMSInput, torch.Tensor | None]: + """Decode detections while retaining whether the source was raw or converted. + + Args: + x: Checked model output tensors. + + Returns: + Layout-aware detections and an optional prototype output. + """ + if len(x) == 1: + converted = self.conversion(x) + if not isinstance(converted, torch.Tensor): + raise TypeError( + "conversion should return a tensor for single-output YOLO postprocessing." + ) + detections = self.filter_conversion(converted) + return _AnchorlessNMSInput(detections, "candidates_first"), None + rearranged = self.rearrange(x) + if not isinstance(rearranged, torch.Tensor): + raise TypeError( + "rearrange should return a tensor for non-segmentation YOLO postprocessing." + ) + detections = self.decode(rearranged) + return _AnchorlessNMSInput(detections, "channels_first"), None + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes detection results for a single image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Decoded boxes, scores, and extra data. + """ + if self.n_extra == 0: + ic = torch.amax(box_cls[-self.nc :, :], dim=0) > self.inv_conf_thres + else: + ic = ( + torch.amax(box_cls[-self.nc - self.n_extra : -self.n_extra, :], dim=0) + > self.inv_conf_thres + ) + box_cls = box_cls[:, ic] # (144, *) + if box_cls.numel() == 0: + return torch.zeros( + (4 + self.nc + self.n_extra, 0), dtype=torch.float32 + ) # (84, 0) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores, extra = torch.split( + box_cls[None], [self.reg_max * 4, self.nc, self.n_extra], dim=1 + ) # (1, 64, *), (1, 80, *), (1, 32, *) + dbox = ( + dist2bbox( + self.dfl(box), + anchors[:, ic], + xywh=False, + dim=1, + ) + * stride[:, ic] + ) + return torch.cat([dbox, scores.sigmoid(), extra], dim=1).squeeze(0) + + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters out low-confidence detections from a single concatenated output tensor. + + Args: + x (torch.Tensor): Concatenated output tensor from the model. + + Returns: + list[torch.Tensor]: Filtered detections for each image in the batch. + """ + while x.ndim == 4 and 1 in (x.shape[0], x.shape[1]): + if x.shape[0] == 1: + x = x.squeeze(0) + elif x.shape[1] == 1: + x = x.squeeze(1) + if x.ndim != 3: + raise ValueError( + f"Expected 3D converted tensor, got shape {tuple(x.shape)}." + ) + expected_dim = 4 + self.nc + self.n_extra + if x.shape[-1] == expected_dim: + normalized = x + elif x.shape[1] == expected_dim: + normalized = x.transpose(1, 2) + else: + raise ValueError(f"Unsupported converted tensor shape {tuple(x.shape)}.") + x_list = torch.split( + normalized, 1, dim=0 + ) # [(1, 8400, 84), (1, 8400, 84), ...] + + def process_conversion(x: torch.Tensor) -> torch.Tensor: + x = x.squeeze(0) # (8400, 84) + if self.n_extra == 0: + ic = torch.amax(x[:, -self.nc :], dim=1) > self.conf_thres + else: + ic = ( + torch.amax(x[:, -self.nc - self.n_extra : -self.n_extra], dim=1) + > self.conf_thres + ) + x = x[ic] + if x.numel() == 0: + return torch.zeros((0, 4 + self.nc + self.n_extra), dtype=torch.float32) + x = xywh2xyxy(x) + return x + + return [process_conversion(xi) for xi in x_list] + + def _nms_single( + self, + xi: torch.Tensor, + max_det: int, + max_nms: int, + max_wh: int, + multi_label: bool, + ) -> torch.Tensor: + """Apply anchorless NMS to a single decoded image tensor.""" + return self._nms_single_legacy_rows( + self._normalize_nms_input(xi, "channels_first"), + max_det=max_det, + max_nms=max_nms, + max_wh=max_wh, + multi_label=multi_label, + ) + + def _normalize_nms_input( + self, + xi: torch.Tensor, + layout: AnchorlessOutputLayout | None = None, + ) -> torch.Tensor: + """Normalize one decoded tensor to ``(candidates, channels)``. + + Source provenance resolves square tensors. Shape inference remains available + for callers that pass decoded tensors directly, with raw channel-first + layout taking precedence when both dimensions match. + + Args: + xi: One image's decoded detections. + layout: Known source layout, if available. + + Returns: + Detections in canonical ``(candidates, channels)`` layout. + + Raises: + ValueError: If the tensor shape conflicts with the expected channel count + or with the supplied source layout. + """ + if xi.ndim != 2: + raise ValueError( + f"Expected 2D decoded tensor, got shape {tuple(xi.shape)}." + ) + + expected_dim = 4 + self.nc + self.n_extra + if layout == "channels_first": + if xi.shape[0] != expected_dim: + raise ValueError( + f"Expected channel-first decoded tensor with {expected_dim} channels, got shape {tuple(xi.shape)}." + ) + return xi.transpose(0, 1) + if layout == "candidates_first": + if xi.shape[1] != expected_dim: + raise ValueError( + "Expected candidates-first decoded tensor " + f"with {expected_dim} channels, got shape {tuple(xi.shape)}." + ) + return xi + + if xi.shape[0] == expected_dim: + return xi.transpose(0, 1) + if xi.shape[1] == expected_dim: + return xi + raise ValueError(f"Unsupported decoded tensor shape {tuple(xi.shape)}.") + + def _nms_single_legacy_rows( + self, + xi: torch.Tensor, + max_det: int, + max_nms: int, + max_wh: int, + multi_label: bool, + ) -> torch.Tensor: + """Apply anchorless NMS to a single decoded image in row-major ``(anchors, channels)`` form.""" + if xi.numel() == 0: + return torch.zeros( + (0, 6 + self.n_extra), dtype=torch.float32, device=self.device + ) + if multi_label: + xi_out = yolo_multilabel_candidates( + xi, self.nc, self.n_extra, self.conf_thres + ) + else: + box, score, extra = xi[:, :4], xi[:, 4 : 4 + self.nc], xi[:, 4 + self.nc :] + conf, cls_idx = score.max(dim=1) + filt = conf > self.conf_thres + if not torch.any(filt): + return torch.zeros( + (0, 6 + self.n_extra), dtype=torch.float32, device=self.device + ) + box = box[filt] + conf = conf[filt] + cls_idx = cls_idx[filt] + extra = extra[filt] + xi_out = torch.empty( + (box.shape[0], 6 + self.n_extra), dtype=xi.dtype, device=xi.device + ) + xi_out[:, :4] = box + xi_out[:, 4] = conf + xi_out[:, 5] = cls_idx.to(xi.dtype) + if self.n_extra > 0: + xi_out[:, 6:] = extra + xi_out = xi_out[torch.argsort(xi_out[:, 4], descending=True)[:max_nms]] + c = xi_out[:, 5:6] * max_wh + boxes, scores = xi_out[:, :4] + c, xi_out[:, 4] + i_idx = non_max_suppression(boxes, scores, self.iou_thres, max_det) + return xi_out[i_idx] + + def nms( + self, + x: _AnchorlessNMSInput | torch.Tensor | list[torch.Tensor], + max_det: int = 300, + max_nms: int = 30000, + max_wh: int = 7680, + multi_label: bool = False, + ) -> list[torch.Tensor]: + """Performs Non-Maximum Suppression (NMS) on the decoded detections. + + Args: + x: Decoded detections for each image, optionally with source-layout + provenance. + max_det (int, optional): Maximum number of detections to keep. Defaults to 300. + max_nms (int, optional): Maximum number of candidates to consider for NMS. + Defaults to 30000. + max_wh (int, optional): Maximum box width/height for offset calculation. + Defaults to 7680. + multi_label: Whether to retain every class candidate above threshold. + + Returns: + list[torch.Tensor]: Post-NMS detections for each image. + """ + layout: AnchorlessOutputLayout | None = None + if isinstance(x, _AnchorlessNMSInput): + layout = x.layout + x = x.detections + + tensors = x if isinstance(x, list) else list(x) + return [ + self._nms_single_legacy_rows( + self._normalize_nms_input(xi, layout), + max_det=max_det, + max_nms=max_nms, + max_wh=max_wh, + multi_label=multi_label, + ) + for xi in tensors + ] + + def nms_multilabel( + self, + x: _AnchorlessNMSInput | torch.Tensor | list[torch.Tensor], + ) -> list[torch.Tensor]: + """Perform Ultralytics-compatible multi-label NMS for validation.""" + return self.nms(x, multi_label=True) + + def dfl(self, x: torch.Tensor) -> torch.Tensor: + """Applies Distribution Focal Loss projection. + + Args: + x: Tensor with shape ``(B, 4 * reg_max, A)``. + + Returns: + Tensor with shape ``(B, 4, A)`` containing projected distances. + """ + if self.reg_max == 0: # skip dfl for yolov6 s, n models + return x + if x.ndim != 3: + raise ValueError( + f"DFL input must have shape (B, 4 * reg_max, A), got {tuple(x.shape)}." + ) + B, _, A = x.shape + x = x.view(B, 4, self.reg_max, A).softmax(dim=2) + return (x * self.dfl_weight.view(1, 1, self.reg_max, 1)).sum(dim=2) + + +class YOLOAnchorlessSegPost(YOLOSegPostMixin, YOLOAnchorlessDetectionPost): + """Postprocessing for YOLO segmentation models without anchors.""" + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return the export-style output tensor for anchorless YOLO segmentation models. + + Args: + x: Checked raw model outputs. + + Returns: + A detection tensor, or detections paired with prototype masks. + """ + if len(x) == 2: + converted, proto_outs = self.conversion(x) + return [ + self._converted_to_batch_output(converted), + self._proto_to_nchw(proto_outs), + ] + return super().non_e2e(x) + + def _proto_to_nchw(self, proto: torch.Tensor) -> torch.Tensor: + """Convert prototype tensors to ``(B, C, H, W)`` if needed. + + Args: + proto: Prototype tensor from a model runtime. + + Returns: + Prototype tensor in channel-first batch layout. + """ + if proto.ndim == 4 and proto.shape[1] == self.n_extra: + return proto + if proto.ndim == 4 and proto.shape[-1] == self.n_extra: + return proto.permute(0, 3, 1, 2) + raise ValueError( + f"Unsupported proto tensor shape {tuple(proto.shape)} for non-e2e output." + ) + + def _pre_process( + self, x: list[torch.Tensor] + ) -> tuple[_AnchorlessNMSInput, torch.Tensor]: + """Preprocesses intermediate inputs into (boxes, proto) format. + + Args: + x (list[torch.Tensor]): Raw model output tensors. + + Returns: + tuple: (decoded_detections, prototype_masks). + """ + if len(x) == 2: + converted, proto_outs = self.conversion(x) + detections = self.filter_conversion(converted) + return _AnchorlessNMSInput(detections, "candidates_first"), proto_outs + rearranged, proto_outs = self.rearrange(x) + return _AnchorlessNMSInput( + self.decode(rearranged), "channels_first" + ), proto_outs + + def conversion(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Converts raw model output tensors into detections and prototypes. + + Args: + x (list[torch.Tensor]): List of raw output tensors. + + Returns: + tuple: (detections, prototypes) + """ + if (self.nc + self.n_extra + 4) in x[0].shape[1:] and self.n_extra in x[ + 1 + ].shape[1:]: + return ( + x[0], + x[1], + ) + if (self.nc + self.n_extra + 4) in x[1].shape[1:] and self.n_extra in x[ + 0 + ].shape[1:]: + return ( + x[1], + x[0], + ) + raise ValueError(f"Wrong shape of input: {x[0].shape}, {x[1].shape}") + + def rearrange(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """ + Rearrange model output tensors for segmentation tasks. + Args: + x (list[torch.Tensor]): Raw output tensors. + Returns: + tuple: (concatenated_detections, prototype_masks) + """ + y_det: list[torch.Tensor] = [] + y_cls: list[torch.Tensor] = [] + y_ext: list[torch.Tensor] = [] + for xi in x: + if xi.shape[-1] == self.n_extra: + y_ext.append( + xi.permute(0, 3, 1, 2) + ) # (b, 32, 160, 160), (b, 32, 80, 80), ... + elif xi.shape[-1] == self.reg_max * 4: + y_det.append( + xi.permute(0, 3, 1, 2) + ) # (b, 64, 80, 80), (b, 64 ,40, 40), ... + elif xi.shape[-1] == self.nc: + y_cls.append( + xi.permute(0, 3, 1, 2) + ) # (b, 80, 80, 80), (b, 80, 40, 40), ... + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort as box, scores + y_ext = sorted(y_ext, key=lambda x: x.numel(), reverse=True) + proto = y_ext.pop(0).permute(0, 2, 3, 1) + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + self.validate_split_head_counts( + detection=y_det, classification=y_cls, extra=y_ext + ) + y = torch.cat( + [ + torch.cat((yi_det, yi_cls, yi_ext), dim=1).flatten(2) + for yi_det, yi_cls, yi_ext in zip(y_det, y_cls, y_ext) + ], + dim=-1, + ) + return y, proto + + +class YOLOAnchorlessPosePost(YOLOPosePostMixin, YOLOAnchorlessDetectionPost): + """Postprocessing for YOLO pose estimation models without anchors.""" + + def rearrange(self, x: list[torch.Tensor]) -> torch.Tensor: + """Rearranges model output tensors for pose estimation tasks. + + Args: + x (list[torch.Tensor]): Raw output tensors. + + Returns: + torch.Tensor: Concatenated tensor for decode. + """ + y_det = [] + y_cls = [] + y_kpt = [] + for xi in x: # list of bchw outputs + if xi.ndim == 3: + xi = xi[None] + elif xi.ndim == 4: + pass + else: + raise NotImplementedError(f"Got unsupported ndim for input: {xi.ndim}.") + if xi.shape[-1] == self.reg_max * 4: + y_det.append( + xi.permute(0, 3, 1, 2) + ) # (b, 64, 80, 80), (b, 64 ,40, 40), ... + elif xi.shape[-1] == self.nc: + y_cls.append( + xi.permute(0, 3, 1, 2) + ) # (b, 1, 80, 80), (b, 1, 40, 40), ... + elif xi.shape[-1] == self.n_extra: + y_kpt.append( + xi.permute(0, 3, 1, 2).flatten(2) + ) # (b, 51, 80, 80), (b, 1, 40, 40), ... + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort as box, scores + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + y_kpt = sorted( + y_kpt, key=lambda x: x.numel(), reverse=True + ) # (b, 51, 6400), (b, 51, 1600), (b, 51, 400) + self.validate_split_head_counts( + detection=y_det, classification=y_cls, keypoint=y_kpt + ) + y_tmp = [ + torch.cat((yi_det, yi_cls), dim=1).flatten(2) + for (yi_det, yi_cls) in zip( + y_det, y_cls + ) # (b, 65, 6400), (b, 65, 1600), (b, 65, 400) + ] + return torch.cat( + [ + torch.cat((yi_tmp, yi_kpt), dim=1) + for yi_tmp, yi_kpt in zip(y_tmp, y_kpt) + ], + dim=-1, + ) + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes pose estimation results for a single image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Decoded boxes, scores, and keypoints. + """ + ic = ( + torch.amax(box_cls[-self.nc - self.n_extra : -self.n_extra, :], dim=0) + > self.inv_conf_thres + ) + box_cls = box_cls[:, ic] # (116, *) + if box_cls.numel() == 0: + return torch.zeros( + (4 + self.nc + self.n_extra, 0), dtype=torch.float32 + ) # (56, 0) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores, keypoints = torch.split( + box_cls[None], [self.reg_max * 4, self.nc, self.n_extra], dim=1 + ) # (1, 64, *), (1, 1, *), (1, 51, *) + dbox = ( + dist2bbox( + self.dfl(box), + anchors[:, ic], + xywh=False, + dim=1, + ) + * stride[:, ic] + ) + keypoints = keypoints.view(1, 17, 3, -1) + key_coord, key_conf = torch.split( + keypoints, [2, 1], dim=2 + ) # (1, 17, 2, 8400), (1, 17, 1, 8400) + key_coord = (key_coord * 2 + anchors[:, ic] - 0.5) * stride[ + :, ic + ] # (1, 17, 2, *) + keypoints = torch.cat([key_coord, key_conf.sigmoid()], dim=2).view( + 1, self.n_extra, -1 + ) # (1, 51, *) + return torch.cat([dbox, scores.sigmoid(), keypoints], dim=1).squeeze( + 0 + ) # (56, *) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every anchor without confidence filtering for export-style pose output.""" + box, scores, keypoints = torch.split( + x, [self.reg_max * 4, self.nc, self.n_extra], dim=1 + ) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + dbox = dist2bbox(self.dfl(box), anchors, xywh=False, dim=1) * stride + keypoints = keypoints.view(x.shape[0], 17, 3, -1) + key_coord, key_conf = torch.split(keypoints, [2, 1], dim=2) + key_coord = (key_coord * 2 + anchors.unsqueeze(1) - 0.5) * stride.unsqueeze(1) + keypoints = torch.cat([key_coord, key_conf.sigmoid()], dim=2).view( + x.shape[0], self.n_extra, -1 + ) + return torch.cat([dbox, scores.sigmoid(), keypoints], dim=1) + + +class YOLOAnchorlessOBBPost(YOLOOBBPostMixin, YOLOAnchorlessDetectionPost): + """Postprocessing for anchorless YOLO OBB models.""" + + def _angle_from_raw(self, angle: torch.Tensor) -> torch.Tensor: + """Decode YOLOv8/YOLO11 raw angle logits to radians.""" + return (angle.sigmoid() - 0.25) * torch.pi + + def _pre_process( + self, x: list[torch.Tensor] + ) -> tuple[_AnchorlessNMSInput, torch.Tensor | None]: + """Preprocess OBB inputs into row-major detections. + + Args: + x: Raw model outputs. + + Returns: + A tuple of detections and no prototype output. + """ + if len(x) in {1, 3, 5}: + detections = self.filter_conversion(self.conversion(x)) + return _AnchorlessNMSInput(detections, "candidates_first"), None + return _AnchorlessNMSInput( + self.decode(self.rearrange(x)), "channels_first" + ), None + + def conversion(self, x: list[torch.Tensor]) -> torch.Tensor: + """Convert exported OBB outputs to one canonical tensor. + + Args: + x: Converted OBB runtime outputs. + + Returns: + Detections in ``cx, cy, w, h, class scores..., angle`` format. + """ + if len(x) == 5: + return decode_split_converted_obb_outputs( + x, + self.nc, + self.n_extra, + self.anchors_as_tensor(), + self.stride_as_tensor(), + ) + return concat_converted_obb_outputs(x, self.nc, self.n_extra) + + def rearrange(self, x: list[torch.Tensor]) -> torch.Tensor: + """Rearrange raw OBB heads into ``(batch, channels, anchors)`` format. + + Args: + x: Raw model output tensors. + + Returns: + Concatenated OBB detection tensor. + """ + target_count = len(x) // 3 + y_det: list[torch.Tensor] = [] + y_cls: list[torch.Tensor] = [] + y_angle: list[torch.Tensor] = [] + ambiguous: list[tuple[torch.Tensor, list[int]]] = [] + for xi in x: + if xi.ndim == 3: + xi = xi.unsqueeze(0) + elif xi.ndim > 4: + while xi.ndim > 4 and 1 in xi.shape: + xi = xi.squeeze( + next(idx for idx, size in enumerate(xi.shape) if size == 1) + ) + if xi.ndim == 3: + xi = xi.unsqueeze(0) + if xi.ndim != 4: + raise ValueError( + f"Expected 3D or 4D OBB head, got shape {tuple(xi.shape)}." + ) + + candidates: list[tuple[int, torch.Tensor]] = [] + if xi.shape[1] in {self.reg_max * 4, self.nc, self.n_extra}: + candidates.append((int(xi.shape[1]), xi)) + if xi.shape[-1] in {self.reg_max * 4, self.nc, self.n_extra}: + candidates.append((int(xi.shape[-1]), xi.permute(0, 3, 1, 2))) + + deduped: list[tuple[int, torch.Tensor]] = [] + seen_channels: set[int] = set() + for channel_count, candidate in candidates: + if channel_count not in seen_channels: + seen_channels.add(channel_count) + deduped.append((channel_count, candidate)) + + if len(candidates) == 2 and len(deduped) == 1: + channel_count, _ = deduped[0] + normalized = xi.permute(0, 3, 1, 2) + if channel_count == self.reg_max * 4: + y_det.append(normalized) + elif channel_count == self.nc: + y_cls.append(normalized) + elif channel_count == self.n_extra: + y_angle.append(normalized) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + elif len(deduped) == 1: + channel_count, normalized = deduped[0] + if channel_count == self.reg_max * 4: + y_det.append(normalized) + elif channel_count == self.nc: + y_cls.append(normalized) + elif channel_count == self.n_extra: + y_angle.append(normalized) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + elif len(deduped) > 1: + ambiguous.append((xi, [channel_count for channel_count, _ in deduped])) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + + for xi, channel_options in ambiguous: + if self.reg_max * 4 in channel_options and len(y_det) < target_count: + y_det.append( + xi if xi.shape[1] == self.reg_max * 4 else xi.permute(0, 3, 1, 2) + ) + continue + if self.nc in channel_options and len(y_cls) < target_count: + y_cls.append(xi if xi.shape[1] == self.nc else xi.permute(0, 3, 1, 2)) + continue + if self.n_extra in channel_options and len(y_angle) < target_count: + y_angle.append( + xi if xi.shape[1] == self.n_extra else xi.permute(0, 3, 1, 2) + ) + continue + raise ValueError(f"Wrong shape of input: {xi.shape}") + + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + y_angle = sorted(y_angle, key=lambda x: x.numel(), reverse=True) + self.validate_split_head_counts( + detection=y_det, classification=y_cls, angle=y_angle + ) + return torch.cat( + [ + torch.cat((yi_det, yi_cls, yi_angle), dim=1).flatten(2) + for yi_det, yi_cls, yi_angle in zip(y_det, y_cls, y_angle) + ], + dim=-1, + ) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every OBB anchor without confidence filtering.""" + box, scores, angle = torch.split( + x, [self.reg_max * 4, self.nc, self.n_extra], dim=1 + ) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + angle = self._angle_from_raw(angle) + rbox = dist2rbox(self.dfl(box), angle, anchors, dim=1) * stride + return torch.cat([rbox, scores.sigmoid(), angle], dim=1) + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes OBB results for a single image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Decoded rotated boxes, scores, and angle. + """ + ic = ( + torch.amax(box_cls[-self.nc - self.n_extra : -self.n_extra, :], dim=0) + > self.inv_conf_thres + ) + box_cls = box_cls[:, ic] + if box_cls.numel() == 0: + return torch.zeros( + (4 + self.nc + self.n_extra, 0), dtype=torch.float32, device=self.device + ) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores, angle = torch.split( + box_cls[None], [self.reg_max * 4, self.nc, self.n_extra], dim=1 + ) + angle = self._angle_from_raw(angle) + rbox = dist2rbox(self.dfl(box), angle, anchors[:, ic], dim=1) * stride[:, ic] + return torch.cat([rbox, scores.sigmoid(), angle], dim=1).squeeze(0) + + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters converted OBB outputs from a single ONNX tensor. + + Args: + x: Converted output tensor. + + Returns: + Per-image row-major tensors in ``cx, cy, w, h, scores..., angle`` format. + """ + while x.ndim == 4 and 1 in (x.shape[0], x.shape[1]): + if x.shape[0] == 1: + x = x.squeeze(0) + elif x.shape[1] == 1: + x = x.squeeze(1) + if x.ndim != 3: + raise ValueError( + f"Expected 3D converted tensor, got shape {tuple(x.shape)}." + ) + expected_dim = 4 + self.nc + self.n_extra + if x.shape[-1] == expected_dim: + normalized = x + elif x.shape[1] == expected_dim: + normalized = x.transpose(1, 2) + else: + raise ValueError(f"Unsupported converted tensor shape {tuple(x.shape)}.") + + outputs = [] + for xi in torch.split(normalized, 1, dim=0): + xi = xi.squeeze(0) + ic = torch.amax(xi[:, 4 : 4 + self.nc], dim=1) > self.conf_thres + if torch.any(ic): + outputs.append(xi[ic]) + else: + outputs.append( + torch.zeros((0, expected_dim), dtype=xi.dtype, device=xi.device) + ) + return outputs + + def _nms_single( + self, + xi: torch.Tensor, + max_det: int, + max_nms: int, + max_wh: int, + multi_label: bool, + ) -> torch.Tensor: + """Apply rotated NMS to a single decoded OBB image tensor.""" + if xi.numel() == 0: + return torch.zeros((0, 7), dtype=torch.float32, device=self.device) + xi_t = xi.transpose(0, 1) + return self._nms_single_legacy_rows( + xi_t, + max_det=max_det, + max_nms=max_nms, + max_wh=max_wh, + multi_label=multi_label, + ) + + def _nms_single_legacy_rows( + self, + xi: torch.Tensor, + max_det: int, + max_nms: int, + max_wh: int, + multi_label: bool, + ) -> torch.Tensor: + """Apply rotated NMS to row-major OBB detections.""" + del multi_label + if xi.numel() == 0: + return torch.zeros((0, 7), dtype=torch.float32, device=self.device) + xi_out = yolo_multilabel_candidates(xi, self.nc, self.n_extra, self.conf_thres) + if xi_out.numel() == 0: + return torch.zeros((0, 7), dtype=torch.float32, device=self.device) + xi_out = xi_out[torch.argsort(xi_out[:, 4], descending=True)[:max_nms]] + c = xi_out[:, 5:6] * max_wh + boxes = torch.cat([xi_out[:, :2] + c, xi_out[:, 2:4], xi_out[:, 6:7]], dim=-1) + keep = rotated_nms(boxes, xi_out[:, 4], self.iou_thres)[:max_det] + return xi_out[keep] + + def nms_multilabel( + self, + x: _AnchorlessNMSInput | torch.Tensor | list[torch.Tensor], + ) -> list[torch.Tensor]: + """Preserve existing OBB NMS behavior outside the COCO validation scope.""" + return self.nms(x) + + +YOLOAnchorlessPost = YOLOAnchorlessDetectionPost diff --git a/mblt_vision/utils/postprocess/yolo_dflfree_post.py b/mblt_vision/utils/postprocess/yolo_dflfree_post.py new file mode 100644 index 0000000..02b557e --- /dev/null +++ b/mblt_vision/utils/postprocess/yolo_dflfree_post.py @@ -0,0 +1,983 @@ +from __future__ import annotations + +from typing import Any, cast + +import torch + +from .base import YOLODetectionPostBase +from .common import ( + YOLOOBBPostMixin, + YOLOPosePostMixin, + YOLOSegPostMixin, + concat_converted_obb_outputs, + decode_split_converted_obb_outputs, + dist2bbox, + dist2rbox, + dual_topk, + rotated_nms, + yolo_multilabel_candidates, +) + + +class YOLODFLFreeDetectionPost(YOLODetectionPostBase): + """Postprocessing for YOLO DFL-free models.""" + + max_det = 300 + reducemax_rtol = 1e-3 + reducemax_atol = 5e-2 + + def __init__(self, pre_cfg: dict, post_cfg: dict, **kwargs: object) -> None: + """Initialize the DFL-free YOLO postprocessor. + + Args: + pre_cfg: Preprocessing configuration. + post_cfg: Postprocessing configuration. + **kwargs: Optional runtime overrides for postprocess behavior. + """ + super().__init__(pre_cfg, post_cfg, **kwargs) + + def _normalize_converted_part( + self, x: torch.Tensor, channel_count: int + ) -> torch.Tensor | None: + """Normalize a split decode-true part to ``(B, anchors, channels)``.""" + + while x.ndim > 3 and 1 in x.shape: + x = x.squeeze(next(idx for idx, size in enumerate(x.shape) if size == 1)) + + if x.ndim == 2: + if x.shape[-1] == channel_count: + return x.unsqueeze(0) + if x.shape[0] == channel_count: + return x.transpose(0, 1).unsqueeze(0) + return None + + if x.ndim == 3: + if x.shape[-1] == channel_count: + return x + if x.shape[1] == channel_count: + return x.transpose(1, 2) + + return None + + def _collect_converted_parts( + self, + x: list[torch.Tensor], + *, + require_extra: bool, + ) -> tuple[torch.Tensor, set[int]] | None: + """Collect decode-true box/class/extra parts while ignoring reducemax.""" + + part_by_role: dict[str, torch.Tensor] = {} + used_indices: set[int] = set() + required_parts: list[tuple[str, int]] = [("boxes", 4), ("scores", self.nc)] + if require_extra: + required_parts.append(("extra", self.n_extra)) + + score_candidates = [ + (idx, cast(torch.Tensor, normalized)) + for idx, xi in enumerate(x) + if (normalized := self._normalize_converted_part(xi, self.nc)) is not None + ] + reducemax_candidates = [ + (idx, cast(torch.Tensor, normalized)) + for idx, xi in enumerate(x) + if (normalized := self._normalize_converted_part(xi, 1)) is not None + ] + + def _matches_reducemax(candidate_idx: int, candidate: torch.Tensor) -> bool: + if candidate.shape[-1] != self.nc: + return False + reduced = candidate.max(dim=-1, keepdim=True).values + return any( + reduced.shape == reducemax.shape + and torch.allclose( + reduced, + reducemax, + rtol=self.reducemax_rtol, + atol=self.reducemax_atol, + ) + for reducemax_idx, reducemax in reducemax_candidates + if reducemax_idx != candidate_idx + ) + + preferred_single_class_score_idx: int | None = None + if self.nc == 1 and len(score_candidates) > 1: + matched_score_candidates = [ + (idx, candidate) + for idx, candidate in score_candidates + if _matches_reducemax(idx, candidate) + ] + if matched_score_candidates: + preferred_single_class_score_idx, _ = max( + matched_score_candidates, + key=lambda item: float(item[1].sum()), + ) + + for idx, xi in enumerate(x): + for role, channel_count in required_parts: + if role in part_by_role: + continue + normalized = self._normalize_converted_part(xi, channel_count) + if normalized is None: + continue + + if role == "scores": + if ( + preferred_single_class_score_idx is not None + and idx != preferred_single_class_score_idx + ): + continue + if not _matches_reducemax(idx, normalized): + continue + elif ( + channel_count == self.nc + and self.nc == 4 + and _matches_reducemax(idx, normalized) + ): + continue + + part_by_role[role] = normalized + used_indices.add(idx) + break + + if any(role not in part_by_role for role, _ in required_parts): + return None + + batch_size = part_by_role["boxes"].shape[0] + anchor_count = part_by_role["boxes"].shape[1] + for role, _channel_count in required_parts[1:]: + part = part_by_role[role] + if part.shape[0] != batch_size or part.shape[1] != anchor_count: + return None + + ordered_parts = [part_by_role["boxes"], part_by_role["scores"]] + if require_extra: + ordered_parts.append(part_by_role["extra"]) + return torch.cat(ordered_parts, dim=-1), used_indices + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return the export-style output tensor for DFL-free YOLO models.""" + if len(x) == 2: + converted = cast(torch.Tensor, self.conversion(x)) + return self._stack_topk_outputs(self.filter_conversion(converted)) + if len(x) == 4: + converted, proto_outs = cast( + tuple[torch.Tensor, torch.Tensor], self.conversion(x) + ) + return [ + self._stack_topk_outputs(self.filter_conversion(converted)), + self._proto_to_nchw(proto_outs), + ] + if len(x) == 3: + converted = cast(torch.Tensor, self.conversion(x)) + return self._stack_topk_outputs(self.filter_conversion(converted)) + + rearranged = self.rearrange(x) + if isinstance(rearranged, tuple): + det_out, proto_outs = rearranged + return [self.decode_batch(det_out), self._proto_to_nchw(proto_outs)] + return self.decode_batch(rearranged) + + def _proto_to_nchw(self, proto: torch.Tensor) -> torch.Tensor: + """Convert prototype tensors to ``(B, C, H, W)`` if needed.""" + if proto.ndim == 4 and proto.shape[1] == self.n_extra: + return proto + if proto.ndim == 4 and proto.shape[-1] == self.n_extra: + return proto.permute(0, 3, 1, 2) + raise ValueError( + f"Unsupported proto tensor shape {tuple(proto.shape)} for non-e2e output." + ) + + def _stack_topk_outputs(self, outputs: list[torch.Tensor]) -> torch.Tensor: + """Pad or trim per-image detections to a fixed batch tensor.""" + if not outputs: + raise ValueError("At least one output tensor is required.") + + output_dim = int(outputs[0].shape[1]) + padded_outputs = [] + for output in outputs: + if output.ndim != 2: + raise ValueError( + f"Expected 2D detection rows, got shape {tuple(output.shape)}." + ) + if output.shape[1] != output_dim: + raise ValueError( + f"Inconsistent detection row width {output.shape[1]}; expected {output_dim}." + ) + output = output[: self.max_det] + if output.shape[0] < self.max_det: + pad = torch.zeros( + (self.max_det - output.shape[0], output_dim), + dtype=output.dtype, + device=output.device, + ) + output = torch.cat([output, pad], dim=0) + padded_outputs.append(output) + return torch.stack(padded_outputs, dim=0) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every anchor, then apply batched top-k selection for export-style output.""" + box, scores, extra = torch.split(x, [4, self.nc, self.n_extra], dim=1) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + dbox = dist2bbox(box, anchors, xywh=False, dim=1) * stride + decoded = torch.cat([dbox, scores, extra], dim=1).transpose(1, 2) + return self._stack_topk_outputs( + [ + dual_topk( + image, + self.nc, + self.n_extra, + max_det=self.max_det, + conf_thres=self.conf_thres, + score_is_logits=True, + ) + for image in decoded + ] + ) + + def _pre_process(self, x: list[torch.Tensor]) -> tuple[Any, torch.Tensor | None]: + """Preprocesses inputs for DFL-free models. + + Args: + x (list[torch.Tensor]): Raw model outputs. + + Returns: + tuple: (processed detections, None). + """ + if len(x) in {2, 3}: + converted = cast(torch.Tensor, self.conversion(x)) + return self.filter_conversion(converted), None + rearranged = self.rearrange(x) + if not isinstance(rearranged, torch.Tensor): + raise TypeError( + "rearrange should return a tensor for DFL-free detection postprocessing." + ) + return self.decode(rearranged), None + + def conversion( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Converts raw model output tensors into a single concatenated tensor. + + Args: + x (list[torch.Tensor]): List of raw output tensors. + + Returns: + torch.Tensor: + Concatenated tensor of shape ``(batch, num_anchors, 4 + nc + n_extra)``. + """ + converted_parts = self._collect_converted_parts( + x, require_extra=self.n_extra > 0 + ) + if converted_parts is not None: + converted, _ = converted_parts + return converted + + # sort by element number + x = sorted(x, key=lambda x: x.size(), reverse=self.nc < 4) + return torch.cat(x, dim=-1).squeeze(1) # [b, 8400, 84] + + def rearrange( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Rearranges raw outputs into a task-specific intermediate representation. + + Args: + x: Raw model output tensors. + + Returns: + A concatenated intermediate representation used by ``decode``. + """ + y_det = [] + y_cls = [] + for xi in x: # list of bchw outputs + if xi.ndim == 3: + xi = xi[None] + elif xi.ndim == 4: + pass + else: + raise NotImplementedError(f"Got unsupported ndim for input: {xi.ndim}.") + if xi.shape[-1] == 4: + y_det.append( + xi.permute(0, 3, 1, 2) + ) # (b, 4, 80, 80), (b, 4, 40, 40), ... + elif xi.shape[-1] == self.nc: + y_cls.append( + xi.permute(0, 3, 1, 2) + ) # (b, 80, 80, 80), (b, 80, 40, 40), ... + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort as box, scores + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + self.validate_split_head_counts(detection=y_det, classification=y_cls) + return torch.cat( + [ + torch.cat((yi_det, yi_cls), dim=1).flatten(2) + for yi_det, yi_cls in zip(y_det, y_cls) + ], + dim=-1, + ) + + def decode(self, x: torch.Tensor) -> list[torch.Tensor]: + """Decodes model outputs into box coordinates and class scores. + + Args: + x (torch.Tensor): Concatenated output tensor from `rearrange`. + + Returns: + list[torch.Tensor]: Per-image decoded detections after filtering and top-k selection. + """ + return [self.process_box_cls(box_cls) for box_cls in x] + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes detection results for a single image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Decoded boxes, scores, and extra data. + """ + if self.n_extra == 0: + ic = torch.amax(box_cls[-self.nc :, :], dim=0) > self.inv_conf_thres + else: + ic = ( + torch.amax(box_cls[-self.nc - self.n_extra : -self.n_extra, :], dim=0) + > self.inv_conf_thres + ) + box_cls = box_cls[:, ic] # (84, *) + if box_cls.numel() == 0: + return box_cls.new_zeros((0, 4 + self.nc + self.n_extra)) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores, extra = torch.split( + box_cls[None], [4, self.nc, self.n_extra], dim=1 + ) # (*, 4), (*, 80), (*, 32) + dbox = ( + dist2bbox( + box, + anchors[:, ic], + xywh=False, + dim=1, + ) + * stride[:, ic] + ) + pre_topk = ( + torch.cat([dbox, scores, extra], dim=1).squeeze(0).transpose(0, 1) + ) # (*, 84) + return dual_topk( + pre_topk, + self.nc, + self.n_extra, + conf_thres=self.conf_thres, + score_is_logits=True, + ) + + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters out low-confidence detections from a single concatenated output tensor. + + Args: + x (torch.Tensor): Output tensor from the model. + + Returns: + list[torch.Tensor]: Filtered detections for each image in the batch. + """ + x_list = torch.split(x, 1, dim=0) # [(1, 8400, 84), (1, 8400, 84), ...] + + return [ + dual_topk(xi.squeeze(0), self.nc, self.n_extra, conf_thres=self.conf_thres) + for xi in x_list + ] + + def nms( + self, + x: torch.Tensor | list[torch.Tensor], + _max_det: int = 300, + _max_nms: int = 30000, + _max_wh: int = 7680, + ) -> list[torch.Tensor]: + """Performs Non-Maximum Suppression (no-op for NMS-free models). + + Args: + x (list[torch.Tensor]): Decoded detections. + _max_det (int, optional): Maximum number of detections to keep. Defaults to 300. + _max_nms (int, optional): Maximum candidates for NMS. Defaults to 30000. + _max_wh (int, optional): Maximum box width/height. Defaults to 7680. + + Returns: + list[torch.Tensor]: Per-image detections with padded zero rows removed. + """ + if isinstance(x, list): + return x + return [xi[xi[:, 4] > 0] for xi in x] + + +class YOLODFLFreeSegPost(YOLOSegPostMixin, YOLODFLFreeDetectionPost): + """Postprocessing for YOLO NMS-free segmentation models.""" + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return export-style segmentation outputs for converted or raw split heads.""" + + if len(x) in {4, 5}: + converted, proto_outs = cast( + tuple[torch.Tensor, torch.Tensor], self.conversion(x) + ) + return [ + self._stack_topk_outputs(self.filter_conversion(converted)), + self._proto_to_nchw(proto_outs), + ] + + rearranged, proto_outs = self.rearrange(x) + return [self.decode_batch(rearranged), self._proto_to_nchw(proto_outs)] + + def _pre_process( + self, x: list[torch.Tensor] + ) -> tuple[list[torch.Tensor], torch.Tensor]: + """Preprocesses intermediate inputs into (boxes, proto) format. + + Args: + x (list[torch.Tensor]): Raw model output tensors. + + Returns: + tuple: (decoded_detections, prototype_masks). + """ + if len(x) in {4, 5}: + converted, proto_outs = cast( + tuple[torch.Tensor, torch.Tensor], self.conversion(x) + ) + return self.filter_conversion(converted), proto_outs + rearranged, proto_outs = self.rearrange(x) + return self.decode(rearranged), proto_outs + + def conversion( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Converts raw outputs into detections and prototype masks. + + Args: + x: Input tensors. + + Returns: + A tuple of processed detections and prototype masks. + """ + + converted_parts = self._collect_converted_parts(x, require_extra=True) + if converted_parts is not None: + converted, used_indices = converted_parts + batch_size, anchor_count = converted.shape[:2] + reducemax_candidate_indices = { + idx + for idx, xi in enumerate(x) + if (normalized := self._normalize_converted_part(xi, 1)) is not None + and normalized.shape[0] == batch_size + and normalized.shape[1] == anchor_count + } + proto_candidates = [] + for idx, xi in enumerate(x): + if idx in used_indices or idx in reducemax_candidate_indices: + continue + proto = xi + if proto.ndim == 3: + proto = proto.unsqueeze(0) + if proto.ndim == 4 and ( + proto.shape[-1] == self.n_extra or proto.shape[1] == self.n_extra + ): + proto_candidates.append(proto) + if len(proto_candidates) == 1: + return converted, proto_candidates[0] + + x = sorted(x, key=lambda x: x.size(), reverse=self.nc < 4) + outputs: list[torch.Tensor] = [] + protos: list[torch.Tensor] = [] + for xi in x: + if xi.shape[-1] == self.n_extra: + protos.append(xi) + else: + outputs.append(xi) + proto = protos.pop(0 if self.nc < 4 else -1) + converted = torch.cat(outputs + protos, dim=-1).squeeze(1) + return converted, proto + + def rearrange(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Rearranges segmentation outputs into detections and prototype masks. + + Args: + x: Raw model output tensors. + + Returns: + A tuple of concatenated detections and prototype masks. + """ + y_det = [] + y_cls = [] + y_ext = [] + for xi in x: # list of bchw outputs + if xi.ndim == 3: + xi = xi[None] + elif xi.ndim == 4: + pass + else: + raise NotImplementedError(f"Got unsupported ndim for input: {xi.ndim}.") + if xi.shape[-1] == self.n_extra: + y_ext.append( + xi.permute(0, 3, 1, 2) + ) # (b, 32, 160, 160), (b, 32, 80, 80), ... + elif xi.shape[-1] == 4: + y_det.append( + xi.permute(0, 3, 1, 2) + ) # (b, 4, 80, 80), (b, 4 ,40, 40), ... + elif xi.shape[-1] == self.nc: + y_cls.append( + xi.permute(0, 3, 1, 2) + ) # (b, 80, 80, 80), (b, 80, 40, 40), ... + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort as box, scores + y_ext = sorted(y_ext, key=lambda x: x.numel(), reverse=True) + proto = y_ext.pop(0).permute(0, 2, 3, 1) + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + self.validate_split_head_counts( + detection=y_det, classification=y_cls, extra=y_ext + ) + y = torch.cat( + [ + torch.cat((yi_det, yi_cls, yi_ext), dim=1).flatten(2) + for yi_det, yi_cls, yi_ext in zip(y_det, y_cls, y_ext) + ], + dim=-1, + ) + return y, proto + + +class YOLODFLFreePosePost(YOLOPosePostMixin, YOLODFLFreeDetectionPost): + """Postprocessing for YOLO NMS-free pose estimation models.""" + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: + """Return export-style pose outputs for both converted and raw split heads.""" + + if len(x) in {3, 4}: + converted = cast(torch.Tensor, self.conversion(x)) + return self._stack_topk_outputs(self.filter_conversion(converted)) + + rearranged = self.rearrange(x) + return self.decode_batch(rearranged) + + def _pre_process( + self, x: list[torch.Tensor] + ) -> tuple[list[torch.Tensor], torch.Tensor | None]: + """Preprocesses inputs for pose estimation. + + Args: + x (list[torch.Tensor]): Raw model outputs. + + Returns: + tuple: (processed_detections, None). + """ + if len(x) in {3, 4}: + converted = cast(torch.Tensor, self.conversion(x)) + return self.filter_conversion(converted), None + rearranged = self.rearrange(x) + return self.decode(rearranged), None + + def conversion( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Convert input tensors. + Args: + x (list[torch.Tensor]): Input tensors. + Returns: + torch.Tensor: Converted tensor. + """ + converted_parts = self._collect_converted_parts(x, require_extra=True) + if converted_parts is not None: + converted, _ = converted_parts + return converted + + # sort by element number + x = sorted(x, key=lambda x: x.size(), reverse=True) + kpt: torch.Tensor = x.pop(0) + kpt = kpt.permute(0, 3, 1, 2).flatten(-2) + return torch.cat( + [torch.cat(x, dim=-1).squeeze(1), kpt], dim=-1 + ) # [b, 8400, 56] + + def rearrange(self, x: list[torch.Tensor]) -> torch.Tensor: + y_det = [] + y_cls = [] + y_kpt = [] + for xi in x: # list of bchw outputs + if xi.ndim == 3: + xi = xi[None] + elif xi.ndim == 4: + pass + else: + raise NotImplementedError(f"Got unsupported ndim for input: {xi.ndim}.") + if xi.shape[-1] == 4: + y_det.append( + xi.permute(0, 3, 1, 2) + ) # (b, 4, 80, 80), (b, 4 ,40, 40), ... + elif xi.shape[-1] == self.nc: + y_cls.append( + xi.permute(0, 3, 1, 2) + ) # (b, 1, 80, 80), (b, 1, 40, 40), ... + elif xi.shape[-1] == self.n_extra: + y_kpt.append( + xi.permute(0, 3, 1, 2).flatten(2) + ) # (b, 51, 80, 80), (b, 1, 40, 40), ... + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + # sort as box, scores + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + y_kpt = sorted( + y_kpt, key=lambda x: x.numel(), reverse=True + ) # (b, 51, 6400), (b, 51, 1600), (b, 51, 400) + self.validate_split_head_counts( + detection=y_det, classification=y_cls, keypoint=y_kpt + ) + y_tmp = [ + torch.cat((yi_det, yi_cls), dim=1).flatten(2) + for (yi_det, yi_cls) in zip( + y_det, y_cls + ) # (b, 65, 6400), (b, 65, 1600), (b, 65, 400) + ] + return torch.cat( + [ + torch.cat((yi_tmp, yi_kpt), dim=1) + for yi_tmp, yi_kpt in zip(y_tmp, y_kpt) + ], + dim=-1, + ) + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes pose estimation results for a single image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Decoded boxes, scores, and keypoints. + """ + ic = ( + torch.amax(box_cls[-self.nc - self.n_extra : -self.n_extra, :], dim=0) + > self.inv_conf_thres + ) + box_cls = box_cls[:, ic] # (116, *) + if box_cls.numel() == 0: + return box_cls.new_zeros((0, 4 + self.nc + self.n_extra)) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores, keypoints = torch.split( + box_cls[None], [4, self.nc, self.n_extra], dim=1 + ) # (1, 4, *), (1, 1, *), (1, 51, *) + dbox = ( + dist2bbox( + box, + anchors[:, ic], + xywh=False, + dim=1, + ) + * stride[:, ic] + ) + keypoints = keypoints.view(1, 17, 3, -1) + key_coord, key_conf = torch.split( + keypoints, [2, 1], dim=2 + ) # (1, 17, 2, 8400), (1, 17, 1, 8400) + key_coord = (key_coord + anchors[:, ic]) * stride[:, ic] # (1, 17, 2, *) + keypoints = torch.cat([key_coord, key_conf.sigmoid()], dim=2).view( + 1, self.n_extra, -1 + ) # (1, 51, *) + pre_topk = ( + torch.cat([dbox, scores, keypoints], dim=1).squeeze(0).transpose(0, 1) + ) # (*, 56) + return dual_topk( + pre_topk, + self.nc, + self.n_extra, + conf_thres=self.conf_thres, + score_is_logits=True, + ) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every anchor, then apply batched top-k selection for export-style pose output.""" + box, scores, keypoints = torch.split(x, [4, self.nc, self.n_extra], dim=1) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + dbox = dist2bbox(box, anchors, xywh=False, dim=1) * stride + keypoints = keypoints.view(x.shape[0], 17, 3, -1) + key_coord, key_conf = torch.split(keypoints, [2, 1], dim=2) + key_coord = (key_coord + anchors.unsqueeze(1)) * stride.unsqueeze(1) + keypoints = torch.cat([key_coord, key_conf.sigmoid()], dim=2).view( + x.shape[0], self.n_extra, -1 + ) + decoded = torch.cat([dbox, scores, keypoints], dim=1).transpose(1, 2) + return self._stack_topk_outputs( + [ + dual_topk( + image, + self.nc, + self.n_extra, + max_det=self.max_det, + conf_thres=self.conf_thres, + score_is_logits=True, + ) + for image in decoded + ] + ) + + +class YOLODFLFreeOBBPost(YOLOOBBPostMixin, YOLODFLFreeDetectionPost): + """Postprocessing for DFL-free YOLO OBB models.""" + + def _pre_process( + self, x: list[torch.Tensor] + ) -> tuple[list[torch.Tensor], torch.Tensor | None]: + """Preprocess OBB inputs into row-major detections. + + Args: + x: Raw model outputs. + + Returns: + A tuple of detections and no prototype output. + """ + if len(x) in {1, 3, 5}: + converted = cast(torch.Tensor, self.conversion(x)) + return self.filter_conversion(converted), None + rearranged = self.rearrange(x) + if not isinstance(rearranged, torch.Tensor): + raise TypeError( + "rearrange should return a tensor for DFL-free OBB postprocessing." + ) + return self.decode(rearranged), None + + def conversion( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Convert DFL-free OBB outputs to a single tensor. + + Args: + x: Input tensors. + + Returns: + Converted tensor with last dimension ``4 + nc + 1``. + """ + if len(x) == 5: + return decode_split_converted_obb_outputs( + x, + self.nc, + self.n_extra, + self.anchors_as_tensor(), + self.stride_as_tensor(), + ) + return concat_converted_obb_outputs(x, self.nc, self.n_extra) + + def rearrange(self, x: list[torch.Tensor]) -> torch.Tensor: + """Rearrange split raw DFL-free OBB heads. + + Args: + x: Raw model output tensors. + + Returns: + Concatenated tensor in ``(batch, channels, anchors)`` format. + """ + target_count = len(x) // 3 + y_det: list[torch.Tensor] = [] + y_cls: list[torch.Tensor] = [] + y_angle: list[torch.Tensor] = [] + ambiguous: list[tuple[torch.Tensor, list[int]]] = [] + for xi in x: + if xi.ndim == 3: + xi = xi.unsqueeze(0) + elif xi.ndim > 4: + while xi.ndim > 4 and 1 in xi.shape: + xi = xi.squeeze( + next(idx for idx, size in enumerate(xi.shape) if size == 1) + ) + if xi.ndim == 3: + xi = xi.unsqueeze(0) + if xi.ndim != 4: + raise ValueError( + f"Expected 3D or 4D OBB head, got shape {tuple(xi.shape)}." + ) + + candidates: list[tuple[int, torch.Tensor]] = [] + if xi.shape[1] in {4, self.nc, self.n_extra}: + candidates.append((int(xi.shape[1]), xi)) + if xi.shape[-1] in {4, self.nc, self.n_extra}: + candidates.append((int(xi.shape[-1]), xi.permute(0, 3, 1, 2))) + + deduped: list[tuple[int, torch.Tensor]] = [] + seen_channels: set[int] = set() + for channel_count, candidate in candidates: + if channel_count not in seen_channels: + seen_channels.add(channel_count) + deduped.append((channel_count, candidate)) + + if len(candidates) == 2 and len(deduped) == 1: + channel_count, _ = deduped[0] + normalized = xi.permute(0, 3, 1, 2) + if channel_count == 4: + y_det.append(normalized) + elif channel_count == self.nc: + y_cls.append(normalized) + elif channel_count == self.n_extra: + y_angle.append(normalized) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + elif len(deduped) == 1: + channel_count, normalized = deduped[0] + if channel_count == 4: + y_det.append(normalized) + elif channel_count == self.nc: + y_cls.append(normalized) + elif channel_count == self.n_extra: + y_angle.append(normalized) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + elif len(deduped) > 1: + ambiguous.append((xi, [channel_count for channel_count, _ in deduped])) + else: + raise ValueError(f"Wrong shape of input: {xi.shape}") + + for xi, channel_options in ambiguous: + if 4 in channel_options and len(y_det) < target_count: + y_det.append(xi if xi.shape[1] == 4 else xi.permute(0, 3, 1, 2)) + continue + if self.nc in channel_options and len(y_cls) < target_count: + y_cls.append(xi if xi.shape[1] == self.nc else xi.permute(0, 3, 1, 2)) + continue + if self.n_extra in channel_options and len(y_angle) < target_count: + y_angle.append( + xi if xi.shape[1] == self.n_extra else xi.permute(0, 3, 1, 2) + ) + continue + raise ValueError(f"Wrong shape of input: {xi.shape}") + + y_det = sorted(y_det, key=lambda x: x.numel(), reverse=True) + y_cls = sorted(y_cls, key=lambda x: x.numel(), reverse=True) + y_angle = sorted(y_angle, key=lambda x: x.numel(), reverse=True) + self.validate_split_head_counts( + detection=y_det, classification=y_cls, angle=y_angle + ) + return torch.cat( + [ + torch.cat((yi_det, yi_cls, yi_angle), dim=1).flatten(2) + for yi_det, yi_cls, yi_angle in zip(y_det, y_cls, y_angle) + ], + dim=-1, + ) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every OBB anchor for export-style output.""" + box, scores, angle = torch.split(x, [4, self.nc, self.n_extra], dim=1) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + rbox = dist2rbox(box, angle, anchors, dim=1) * stride + return torch.cat([rbox, scores.sigmoid(), angle], dim=1).transpose(1, 2) + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes raw DFL-free OBB results for one image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Raw OBB rows ``cx, cy, w, h, class scores, angle`` before NMS. + """ + ic = ( + torch.amax(box_cls[-self.nc - self.n_extra : -self.n_extra, :], dim=0) + > self.inv_conf_thres + ) + box_cls = box_cls[:, ic] + if box_cls.numel() == 0: + return box_cls.new_zeros((0, 4 + self.nc + self.n_extra)) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores, angle = torch.split( + box_cls[None], [4, self.nc, self.n_extra], dim=1 + ) + rbox = dist2rbox(box, angle, anchors[:, ic], dim=1) * stride[:, ic] + return ( + torch.cat([rbox, scores.sigmoid(), angle], dim=1).squeeze(0).transpose(0, 1) + ) + + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters converted DFL-free OBB outputs. + + Args: + x: Converted output tensor. + + Returns: + Per-image canonical OBB detection rows before rotated NMS. + """ + while x.ndim == 4 and 1 in (x.shape[0], x.shape[1]): + if x.shape[0] == 1: + x = x.squeeze(0) + elif x.shape[1] == 1: + x = x.squeeze(1) + if x.ndim != 3: + raise ValueError( + f"Expected 3D converted tensor, got shape {tuple(x.shape)}." + ) + expected_dim = 4 + self.nc + self.n_extra + if x.shape[-1] == expected_dim: + normalized = x + elif x.shape[1] == expected_dim: + normalized = x.transpose(1, 2) + else: + raise ValueError(f"Unsupported converted tensor shape {tuple(x.shape)}.") + outputs = [] + for xi in normalized: + keep = xi[:, 4 : 4 + self.nc].amax(dim=1) > self.conf_thres + if torch.any(keep): + outputs.append(xi[keep]) + else: + outputs.append(xi.new_zeros((0, expected_dim))) + return outputs + + def nms( + self, + x: torch.Tensor | list[torch.Tensor], + max_det: int = 300, + max_nms: int = 30000, + max_wh: int = 7680, + ) -> list[torch.Tensor]: + """Apply rotated NMS to DFL-free OBB detections. + + Args: + x: Decoded detections. + max_det: Maximum detections to keep. + max_nms: Maximum candidates to consider. + max_wh: Class offset size. + + Returns: + Per-image OBB detections after rotated NMS. + """ + detections = x if isinstance(x, list) else list(x) + output = [] + for xi in detections: + if xi.numel() == 0: + output.append(xi.new_zeros((0, 7))) + continue + if xi.shape[1] == 4 + self.nc + self.n_extra: + xi = yolo_multilabel_candidates( + xi, self.nc, self.n_extra, self.conf_thres + ) + elif xi.shape[1] == 6 + self.n_extra: + xi = xi[xi[:, 4] > self.conf_thres] + else: + raise ValueError(f"Unsupported OBB detection shape {tuple(xi.shape)}.") + if xi.numel() == 0: + output.append(xi.new_zeros((0, 7))) + continue + xi = xi[torch.argsort(xi[:, 4], descending=True)[:max_nms]] + c = xi[:, 5:6] * max_wh + boxes = torch.cat([xi[:, :2] + c, xi[:, 2:4], xi[:, 6:7]], dim=-1) + keep = rotated_nms(boxes, xi[:, 4], self.iou_thres)[:max_det] + output.append(xi[keep]) + return output + + +YOLODFLFreePost = YOLODFLFreeDetectionPost diff --git a/mblt_vision/utils/postprocess/yolo_nmsfree_post.py b/mblt_vision/utils/postprocess/yolo_nmsfree_post.py new file mode 100644 index 0000000..cd4f125 --- /dev/null +++ b/mblt_vision/utils/postprocess/yolo_nmsfree_post.py @@ -0,0 +1,176 @@ +""" +YOLO NMS-free postprocessing. +""" + +from __future__ import annotations + +from typing import Any, cast + +import torch + +from .common import dist2bbox, dual_topk +from .yolo_anchorless_post import YOLOAnchorlessDetectionPost, _AnchorlessNMSInput + + +class YOLONMSFreeDetectionPost(YOLOAnchorlessDetectionPost): + """Postprocessing for YOLO NMS-free models.""" + + max_det = 300 + + def non_e2e(self, x: list[torch.Tensor]) -> torch.Tensor: + """Return the export-style output tensor for NMS-free YOLO models.""" + if len(x) == 2: + converted = cast(torch.Tensor, self.conversion(x)) + return self._stack_topk_outputs(self.filter_conversion(converted)) + + rearranged = cast(torch.Tensor, self.rearrange(x)) + return self.decode_batch(rearranged) + + def _stack_topk_outputs(self, outputs: list[torch.Tensor]) -> torch.Tensor: + """Pad or trim per-image detections to a fixed batch tensor.""" + padded_outputs = [] + for output in outputs: + output = output[: self.max_det] + if output.shape[0] < self.max_det: + pad = torch.zeros( + (self.max_det - output.shape[0], 6), + dtype=output.dtype, + device=output.device, + ) + output = torch.cat([output, pad], dim=0) + padded_outputs.append(output) + return torch.stack(padded_outputs, dim=0) + + def decode_batch(self, x: torch.Tensor) -> torch.Tensor: + """Decode every anchor, then apply batched top-k selection for export-style output.""" + box, scores = torch.split(x, [self.reg_max * 4, self.nc], dim=1) + anchors = self.anchors_as_tensor().unsqueeze(0) + stride = self.stride_as_tensor().unsqueeze(0) + dbox = dist2bbox(self.dfl(box), anchors, xywh=False, dim=1) * stride + decoded = torch.cat([dbox, scores], dim=1).transpose(1, 2) + return self._stack_topk_outputs( + [ + dual_topk( + image, + self.nc, + self.n_extra, + max_det=self.max_det, + conf_thres=self.conf_thres, + score_is_logits=True, + ) + for image in decoded + ] + ) + + def _pre_process(self, x: list[torch.Tensor]) -> tuple[Any, torch.Tensor | None]: + """Preprocesses inputs for NMS-free models. + + Args: + x (list[torch.Tensor]): Raw model outputs. + + Returns: + tuple: (processed_detections, None). + """ + if len(x) == 2: + converted = cast(torch.Tensor, self.conversion(x)) + return self.filter_conversion(converted), None + rearranged = cast(torch.Tensor, self.rearrange(x)) + return self.decode(rearranged), None + + def conversion( + self, x: list[torch.Tensor] + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Convert input tensors. + Args: + x (list[torch.Tensor]): Input tensors. + Returns: + torch.Tensor: Converted tensor. + """ + # sort by element number + x = sorted(x, key=lambda x: x.size(), reverse=self.nc < 4) + return torch.cat(x, dim=-1).squeeze(1) # [b, 8400, 84] + + def filter_conversion(self, x: torch.Tensor) -> list[torch.Tensor]: + """Filters out low-confidence detections from a single output tensor. + + Args: + x (torch.Tensor): Model output tensor. + + Returns: + list[torch.Tensor]: Decoded and filtered outputs for each image. + """ + x_list = torch.split(x, 1, dim=0) # [(1, 8400, 84), (1, 8400, 84), ...] + + return [ + dual_topk(xi.squeeze(0), self.nc, self.n_extra, conf_thres=self.conf_thres) + for xi in x_list + ] + + def process_box_cls(self, box_cls: torch.Tensor) -> torch.Tensor: + """Processes detection results for a single image. + + Args: + box_cls: Raw detections for one image. + + Returns: + Decoded and top-k filtered detections. + """ + ic = torch.amax(box_cls[-self.nc :, :], dim=0) > self.inv_conf_thres + box_cls = box_cls[:, ic] # (144, *) + if box_cls.numel() == 0: + return box_cls.new_zeros((0, 6)) + anchors = self.anchors_as_tensor() + stride = self.stride_as_tensor() + box, scores = torch.split( + box_cls[None], [self.reg_max * 4, self.nc], dim=1 + ) # (1, 64, *), (1, 80, *) + dbox = ( + dist2bbox( + self.dfl(box), + anchors[:, ic], + xywh=False, + dim=1, + ) + * stride[:, ic] + ) + pre_topk = ( + torch.cat([dbox, scores], dim=1).squeeze(0).transpose(0, 1) + ) # (*, 84) + return dual_topk( + pre_topk, + self.nc, + self.n_extra, + conf_thres=self.conf_thres, + score_is_logits=True, + ) + + def nms( + self, + x: _AnchorlessNMSInput | torch.Tensor | list[torch.Tensor], + max_det: int = 300, + max_nms: int = 30000, + max_wh: int = 7680, + multi_label: bool = False, + ) -> list[torch.Tensor]: + """Perform Non-Maximum Suppression (no-op for NMS-free models). + + Args: + x: Decoded detections, optionally with source-layout provenance. + max_det (int, optional): Maximum number of detections to keep. Defaults to 300. + max_nms (int, optional): Maximum candidates for NMS. Defaults to 30000. + max_wh (int, optional): Maximum box width/height. Defaults to 7680. + multi_label: Ignored because NMS-free outputs already select one + class per candidate. + + Returns: + list[torch.Tensor]: Per-image detections with padded zero rows removed. + """ + del max_det, max_nms, max_wh, multi_label + if isinstance(x, _AnchorlessNMSInput): + x = x.detections + if isinstance(x, list): + return x + return [xi[xi[:, 4] > 0] for xi in x] + + +YOLONMSFreePost = YOLONMSFreeDetectionPost diff --git a/mblt_vision/utils/preprocess/__init__.py b/mblt_vision/utils/preprocess/__init__.py new file mode 100644 index 0000000..497dcea --- /dev/null +++ b/mblt_vision/utils/preprocess/__init__.py @@ -0,0 +1,27 @@ +""" +Preprocessing utilities for vision models. +""" + +from .base import PreBase, PreOps +from .build_pre import build_preprocess +from .center_crop import CenterCrop +from .letterbox import LetterBox, letterbox_semantic_mask +from .normalize import Normalize +from .order import SetOrder +from .reader import Reader +from .resize import Resize +from .yolo_pre import YoloPre + +__all__ = [ + "CenterCrop", + "LetterBox", + "Normalize", + "PreBase", + "PreOps", + "Reader", + "Resize", + "SetOrder", + "YoloPre", + "build_preprocess", + "letterbox_semantic_mask", +] diff --git a/mblt_vision/utils/preprocess/_validation.py b/mblt_vision/utils/preprocess/_validation.py new file mode 100644 index 0000000..5d393d9 --- /dev/null +++ b/mblt_vision/utils/preprocess/_validation.py @@ -0,0 +1,66 @@ +"""Validation helpers shared by image preprocessing operations.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + + +def normalize_uint8_rgb_array(image: np.ndarray, *, operation: str) -> np.ndarray: + """Return byte RGB data after validating or scaling a floating-point image. + + ``[0, 1]`` floating-point input is treated as normalized RGB and scaled to + ``[0, 255]``. Other floating-point input must already lie in ``[0, 255]``. + """ + + if image.dtype == np.uint8: + return image + if not np.issubdtype(image.dtype, np.floating): + raise TypeError( + f"{operation} accepts uint8 arrays or floating-point arrays with RGB " + f"values in [0, 1] or [0, 255]; got {image.dtype}." + ) + if not np.isfinite(image).all(): + raise ValueError( + f"{operation} requires floating-point image arrays to contain only " + "finite RGB values." + ) + + min_value = float(image.min()) + max_value = float(image.max()) + if min_value < 0.0 or max_value > 255.0: + raise ValueError( + f"{operation} accepts floating-point RGB values only in [0, 1] or " + f"[0, 255]; got range [{min_value}, {max_value}]." + ) + if max_value <= 1.0: + image = image * 255.0 + return np.rint(image).astype(np.uint8) + + +def normalize_image_size(size: int | Sequence[int], *, name: str = "size") -> list[int]: + """Normalize a positive scalar or two-dimensional image size to ``[height, width]``.""" + + if isinstance(size, bool): + raise TypeError( + f"{name} must be a positive integer or a two-item integer sequence, got bool." + ) + if isinstance(size, int): + if size <= 0: + raise ValueError(f"{name} must be positive, got {size}.") + return [size, size] + if isinstance(size, Sequence) and not isinstance(size, (str, bytes)): + if len(size) != 2: + raise ValueError(f"{name} must contain exactly two items, got {size!r}.") + if not all( + isinstance(value, int) and not isinstance(value, bool) for value in size + ): + raise TypeError(f"{name} items must be integers, got {size!r}.") + normalized = [int(size[0]), int(size[1])] + if any(value <= 0 for value in normalized): + raise ValueError(f"{name} items must be positive, got {size!r}.") + return normalized + raise TypeError( + f"{name} must be a positive integer or a two-item integer sequence, got {type(size).__name__}." + ) diff --git a/mblt_vision/utils/preprocess/base.py b/mblt_vision/utils/preprocess/base.py new file mode 100644 index 0000000..520c23e --- /dev/null +++ b/mblt_vision/utils/preprocess/base.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +import torch + + +def _spatial_shape(value: Any) -> tuple[int, int] | None: + """Return the height and width of image-like preprocessing input.""" + + shape = getattr(value, "shape", None) + if shape is None or len(shape) < 2: + return None + if len(shape) == 2: + return int(shape[0]), int(shape[1]) + if len(shape) == 3 and int(shape[-1]) in {1, 3, 4}: + return int(shape[0]), int(shape[1]) + return int(shape[-2]), int(shape[-1]) + + +class PreOps(ABC): + """Abstract base class for individual preprocessing operations. + + Attributes: + device: The torch device where tensors should be placed. + """ + + def __init__(self) -> None: + """Initializes the preprocessing operation.""" + super().__init__() + self.device = torch.device("cpu") + + @abstractmethod + def __call__( + self, + x: Any, + ) -> Any: + """Executes the preprocess operation. + + Args: + x: Input data to be processed. + + Returns: + Processed data. + """ + + def to( + self, + device: str | torch.device, + ) -> None: + """Move the operation to the specified device. + + Args: + device: Device to move the operation to. + """ + if isinstance(device, str): + self.device = torch.device(device) + elif isinstance(device, torch.device): + self.device = device + else: + raise TypeError(f"Got unexpected type for device={type(device)}.") + for name, value in self.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(self, name, value.to(self.device)) + + +class PreBase: + """Base class for orchestrating a series of preprocessing operations. + + Attributes: + Ops: List of ordered PreOps instances to be applied. + device: The torch device being used. + """ + + def __init__( + self, + Ops: list[PreOps], + ) -> None: + """Initializes the PreBase class with a list of operations. + + Args: + Ops: List of ordered PreOps instances to be applied. + """ + self.Ops = Ops + self._check_ops() + self.device = torch.device("cpu") + + def _check_ops(self) -> None: + """Check if the operations are valid.""" + for op in self.Ops: + if not isinstance(op, PreOps): + raise TypeError(f"Got unsupported type={type(op)}.") + + def __call__( + self, + x: Any, + ) -> Any: + """Applies the sequence of preprocessing operations to the input. + + Args: + x: Initial input data. + + Returns: + Fully processed data. + """ + for op in self.Ops: + x = op(x) + return x + + def with_metadata( + self, + x: Any, + ) -> tuple[Any, dict[str, Any]]: + """Apply preprocessing and return metadata produced by preprocessing operations. + + Args: + x: Initial input data. + + Returns: + A tuple of the processed data and collected metadata. + """ + metadata: dict[str, Any] = {} + img0_shape = _spatial_shape(x) + if img0_shape is not None: + metadata["img0_shape"] = img0_shape + for op in self.Ops: + x = op(x) + if "img0_shape" not in metadata: + img0_shape = _spatial_shape(x) + if img0_shape is not None: + metadata["img0_shape"] = img0_shape + ratio_pad = getattr(op, "ratio_pad", None) + if ratio_pad is not None: + metadata["ratio_pad"] = ratio_pad + return x, metadata + + def to( + self, + device: str | torch.device, + ) -> None: + """Move the operations to the specified device. + + Args: + device: Device to move the operations to. + """ + if isinstance(device, str): + self.device = torch.device(device) + elif isinstance(device, torch.device): + self.device = device + else: + raise TypeError(f"Got unexpected type for device={type(device)}.") + for name, value in self.__dict__.items(): + if isinstance(value, torch.Tensor): + setattr(self, name, value.to(self.device)) + for op in self.Ops: + op.to(self.device) diff --git a/mblt_vision/utils/preprocess/build_pre.py b/mblt_vision/utils/preprocess/build_pre.py new file mode 100644 index 0000000..3dcc432 --- /dev/null +++ b/mblt_vision/utils/preprocess/build_pre.py @@ -0,0 +1,48 @@ +""" +Preprocessing builder. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from .base import PreBase +from .center_crop import CenterCrop +from .letterbox import LetterBox +from .normalize import Normalize +from .order import SetOrder +from .reader import Reader +from .resize import Resize + + +def build_preprocess( + pre_cfg: Mapping[str, Mapping[str, Any]], +) -> PreBase: + """Builds a preprocessing pipeline based on the configuration. + + Args: + pre_cfg: Preprocessing configuration mapping operations to attributes. + + Returns: + An orchestrator for the sequence of preprocessing steps. + """ + res = [] + for pre_type, pre_attr in pre_cfg.items(): + pre_type_lower = pre_type.lower() + if pre_type_lower == Reader.__name__.lower(): + res.append(Reader(**pre_attr)) + elif pre_type_lower == Resize.__name__.lower(): + res.append(Resize(**pre_attr)) + elif pre_type_lower == CenterCrop.__name__.lower(): + res.append(CenterCrop(**pre_attr)) + elif pre_type_lower == SetOrder.__name__.lower(): + res.append(SetOrder(**pre_attr)) + elif pre_type_lower == LetterBox.__name__.lower(): + res.append(LetterBox(**pre_attr)) + elif pre_type_lower == Normalize.__name__.lower(): + res.append(Normalize(**pre_attr)) + else: + raise ValueError(f"Got unsupported pre_type={pre_type}.") + + return PreBase(res) diff --git a/mblt_vision/utils/preprocess/center_crop.py b/mblt_vision/utils/preprocess/center_crop.py new file mode 100644 index 0000000..b116b3c --- /dev/null +++ b/mblt_vision/utils/preprocess/center_crop.py @@ -0,0 +1,75 @@ +""" +Center crop preprocessing. +""" + +from __future__ import annotations + +import cv2 +import numpy as np +import torch +from PIL import Image + +from ..types import TensorLike +from ._validation import normalize_image_size +from .base import PreOps + + +class CenterCrop(PreOps): + """ + Center crop the image to a specified size. + """ + + def __init__(self, size: int | list[int] | tuple[int, int]) -> None: + """Initializes the CenterCrop operation. + + Args: + size (int | list[int]): Target size [h, w]. If int, size is [size, size]. + """ + super().__init__() + self.size = normalize_image_size(size) + + def __call__(self, x: TensorLike | Image.Image) -> np.ndarray: + """Applies center crop to the image. + + Args: + x (np.ndarray | torch.Tensor | Image.Image): Input image. + + Returns: + np.ndarray: Center-cropped image in HWC format. + """ + if isinstance(x, torch.Tensor): + image = x.detach().cpu().numpy() + elif isinstance(x, Image.Image): + image = np.array(x) + elif isinstance(x, np.ndarray): + image = x + else: + raise TypeError( + f"CenterCrop expects a NumPy array, tensor, or PIL image, got {type(x).__name__}." + ) + if image.ndim != 3: + raise ValueError( + f"CenterCrop expects a three-dimensional image, got shape {image.shape}." + ) + H, W = image.shape[:2] + if (self.size[0] == H) and (self.size[1] == W): + return image + elif (self.size[1] > W) or (self.size[0] > H): + image = cv2.copyMakeBorder( + image, + (self.size[0] - H) // 2 if self.size[0] > H else 0, + (self.size[0] - H + 1) // 2 if self.size[0] > H else 0, + (self.size[1] - W) // 2 if self.size[1] > W else 0, + (self.size[1] - W + 1) // 2 if self.size[1] > W else 0, + cv2.BORDER_CONSTANT, + value=(0.0,), + ) + H, W = image.shape[:2] + crop_top = round((H - self.size[0]) / 2.0) + crop_left = round((W - self.size[1]) / 2.0) + image = image[ + crop_top : crop_top + self.size[0], + crop_left : crop_left + self.size[1], + :, + ] + return image.astype(np.uint8) diff --git a/mblt_vision/utils/preprocess/letterbox.py b/mblt_vision/utils/preprocess/letterbox.py new file mode 100644 index 0000000..4d2c1fe --- /dev/null +++ b/mblt_vision/utils/preprocess/letterbox.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import cv2 +import numpy as np +import torch + +from ..letterbox import LetterBoxGeometry, RatioPad +from ..types import TensorLike +from ._validation import normalize_image_size, normalize_uint8_rgb_array +from .base import PreOps + + +def _apply_letterbox( + image: np.ndarray, + img_size: list[int], + interpolation: int, + padding_value: int | tuple[int, int, int], +) -> tuple[np.ndarray, RatioPad]: + """Resize and pad an array while preserving its aspect ratio. + + Args: + image: Image or two-dimensional semantic mask. + img_size: Target size as ``[height, width]``. + interpolation: OpenCV interpolation mode. + padding_value: Constant border value. + + Returns: + The letterboxed array and its resize/padding metadata. + """ + + input_shape = (int(img_size[0]), int(img_size[1])) + original_shape = (int(image.shape[0]), int(image.shape[1])) + geometry = LetterBoxGeometry.from_shapes(input_shape, original_shape) + resized_height, resized_width = geometry.resized_shape + if image.shape[:2] != geometry.resized_shape: + image = cv2.resize( + image, (resized_width, resized_height), interpolation=interpolation + ) + top, bottom, left, right = geometry.borders + image = cv2.copyMakeBorder( + image, + top, + bottom, + left, + right, + cv2.BORDER_CONSTANT, + value=padding_value, + ) + return image, geometry.ratio_pad + + +def letterbox_semantic_mask( + mask: np.ndarray, + img_size: list[int], + ignore_label: int = 255, +) -> tuple[np.ndarray, RatioPad]: + """Letterbox a semantic mask without interpolating class IDs. + + Args: + mask: Two-dimensional semantic class map. + img_size: Target size as ``[height, width]``. + ignore_label: Class value used for padded pixels. + + Returns: + The letterboxed mask and its resize/padding metadata. + + Raises: + ValueError: If the mask is not two-dimensional. + """ + + if mask.ndim != 2: + raise ValueError( + f"Semantic masks must be two-dimensional, got shape {mask.shape}." + ) + return _apply_letterbox(mask, img_size, cv2.INTER_NEAREST, ignore_label) + + +class LetterBox(PreOps): + """Preprocessing for YOLO models, implementing letterbox resizing. + + Resizes the image while maintaining aspect ratio, adding padding to meet + target dimensions. Floating-point RGB inputs in ``[0, 1]`` are scaled to + byte RGB; other floating-point values must be finite and in ``[0, 255]``. + Based on Ultralytics implementation. + + Ref: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/augment.py#L1535 + """ + + def __init__(self, img_size: list[int]) -> None: + """Initializes LetterBox with target image size. + + Args: + img_size (list[int]): Target image size [h, w]. + """ + super().__init__() + self.img_size = normalize_image_size(img_size, name="img_size") + self.ratio_pad: tuple[tuple[float, float], tuple[float, float]] | None = None + + def __call__(self, x: TensorLike) -> torch.Tensor: + """Executes YOLO preprocessing (letterbox resizing). + + Args: + x (TensorLike): Input image. + + Returns: + torch.Tensor: Preprocessed image in HWC format on the selected device. + """ + if isinstance(x, torch.Tensor): + x = x.detach().cpu().numpy() + elif not isinstance(x, np.ndarray): + raise TypeError( + f"LetterBox expects a NumPy array or tensor, got {type(x).__name__}." + ) + if x.ndim != 3: + raise ValueError(f"LetterBox expects an HWC image, got shape {x.shape}.") + x = normalize_uint8_rgb_array(x, operation="LetterBox") + img, self.ratio_pad = _apply_letterbox( + x, + self.img_size, + cv2.INTER_LINEAR, + (114, 114, 114), + ) + return torch.from_numpy(img).to(self.device).byte() diff --git a/mblt_vision/utils/preprocess/normalize.py b/mblt_vision/utils/preprocess/normalize.py new file mode 100644 index 0000000..be7dd1b --- /dev/null +++ b/mblt_vision/utils/preprocess/normalize.py @@ -0,0 +1,94 @@ +"""Normalization operation for image preprocessing.""" + +from __future__ import annotations + +import numpy as np +import torch +from PIL import Image + +from ..types import TensorLike +from .base import PreOps + +STYLE_PARAMS = { + "torch": ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + "tf": ([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), + "openai": ( + [0.48145466, 0.4578275, 0.40821073], + [0.26862954, 0.26130258, 0.27577711], + ), + "cv": ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]), +} +STYLE_LIST = list(STYLE_PARAMS.keys()) + + +class Normalize(PreOps): + """Normalization layer to scale and shift image data. + + Attributes: + style: Data source style (e.g., 'torch', 'tf', 'openai', 'cv'). + mean: Array of mean values for normalization. + std: Array of standard deviation values for normalization. + """ + + def __init__(self, style: str) -> None: + """Initializes the Normalize layer with a specific style. + + Args: + style: The preprocessing style to use. Must be one of STYLE_LIST. + """ + super().__init__() + + if not isinstance(style, str): + raise TypeError( + f"Normalize style must be a string, got {type(style).__name__}." + ) + if style.lower() not in STYLE_LIST: + raise ValueError( + f"Unsupported Normalize style {style!r}; expected one of {STYLE_LIST}." + ) + + self.style = style.lower() + mean, std = STYLE_PARAMS[self.style] + self.mean = np.array(mean) + self.std = np.array(std) + + def __call__(self, x: TensorLike | Image.Image) -> np.ndarray: + """Applies normalization to the input image or tensor. + + Args: + x (TensorLike | Image.Image): Input data as a torch.Tensor, PIL Image, or numpy-like array. + + Returns: + np.ndarray: The normalized image as a float32 numpy array. + """ + if isinstance(x, torch.Tensor): + x = x.detach().cpu().numpy() + elif isinstance(x, Image.Image): + x = np.array(x) + elif not isinstance(x, np.ndarray): + raise TypeError( + f"Normalize expects a NumPy array, tensor, or PIL image, got {type(x).__name__}." + ) + if x.ndim != 3: + raise ValueError( + f"Normalize expects a three-dimensional image, got shape {x.shape}." + ) + x = x.astype(np.float32) / 255.0 + channels_first = x.shape[0] == len(self.mean) + channels_last = x.shape[-1] == len(self.mean) + if channels_first and channels_last: + raise ValueError( + f"Normalize cannot infer channel order from ambiguous shape {x.shape}." + ) + if channels_last: + mean, std = self.mean, self.std + elif channels_first: + mean = self.mean[:, None, None] + std = self.std[:, None, None] + else: + raise ValueError( + f"Normalize expects HWC or CHW data with {len(self.mean)} channels, " + f"got shape {x.shape}." + ) + x = (x - mean) / std + return x.astype(np.float32) diff --git a/mblt_vision/utils/preprocess/order.py b/mblt_vision/utils/preprocess/order.py new file mode 100644 index 0000000..d4931bd --- /dev/null +++ b/mblt_vision/utils/preprocess/order.py @@ -0,0 +1,74 @@ +""" +Channel order preprocessing. +""" + +from __future__ import annotations + +import numpy as np +import torch + +from ..types import TensorLike +from .base import PreOps + + +class SetOrder(PreOps): + """Sets the channel order of the image to either HWC or CHW format.""" + + def __init__(self, shape: str = "HWC") -> None: + """Initializes the SetOrder operation. + + Args: + shape (str, optional): Target channel order, either "HWC" or "CHW". + Defaults to "HWC". + """ + super().__init__() + if not isinstance(shape, str): + raise TypeError( + f"SetOrder shape must be a string, got {type(shape).__name__}." + ) + if shape.lower() not in {"hwc", "chw"}: + raise ValueError( + f"Unsupported channel order {shape!r}; expected 'HWC' or 'CHW'." + ) + self.shape = shape.lower() + + def __call__(self, x: TensorLike) -> TensorLike: + """Reorders the dimensions of the input image. + + Args: + x (TensorLike): Input image of shape (3, H, W) or (H, W, 3). + + Returns: + TensorLike: Image with the specified channel order. + """ + if not isinstance(x, (np.ndarray, torch.Tensor)): + raise TypeError( + f"SetOrder expects a NumPy array or tensor, got {type(x).__name__}." + ) + if x.ndim != 3: + raise ValueError( + f"SetOrder expects a three-dimensional color image, got shape {x.shape}." + ) + channels_first = x.shape[0] == 3 + channels_last = x.shape[-1] == 3 + if channels_first and channels_last: + raise ValueError( + f"SetOrder cannot infer channel order from ambiguous shape {x.shape}." + ) + if channels_first: + cdim = 0 + elif channels_last: + cdim = 2 + else: + raise ValueError( + f"SetOrder expects HWC or CHW data with three channels, got shape {x.shape}." + ) + if cdim == 0 and self.shape == "hwc": + if isinstance(x, torch.Tensor): + return torch.permute(x, (1, 2, 0)) + return np.transpose(x, (1, 2, 0)) + elif cdim == 2 and self.shape == "chw": + if isinstance(x, torch.Tensor): + return torch.permute(x, (2, 0, 1)) + return np.transpose(x, (2, 0, 1)) + return x diff --git a/mblt_vision/utils/preprocess/reader.py b/mblt_vision/utils/preprocess/reader.py new file mode 100644 index 0000000..69c9f8d --- /dev/null +++ b/mblt_vision/utils/preprocess/reader.py @@ -0,0 +1,94 @@ +""" +Image reader preprocessing. +""" + +from __future__ import annotations + +from pathlib import Path + +import cv2 +import numpy as np +import torch +from PIL import Image + +from ..types import TensorLike +from ._validation import normalize_uint8_rgb_array +from .base import PreOps + + +class Reader(PreOps): + """ + Reader for loading images from file paths or converting existing objects. + Supports "pil" and "numpy" reading styles. + + For ``style="pil"``, arrays must be ``uint8`` RGB values or finite floating-point + RGB values. Floating-point arrays in ``[0, 1]`` are treated as normalized RGB and + scaled to ``[0, 255]``; other floating-point arrays must already be in ``[0, 255]``. + """ + + def __init__(self, style: str) -> None: + """Initializes the Reader operation. + + Args: + style (str): Reading style, either "pil" or "numpy". + """ + super().__init__() + if not isinstance(style, str): + raise TypeError( + f"Reader style must be a string, got {type(style).__name__}." + ) + if style.lower() not in {"pil", "numpy"}: + raise ValueError( + f"Unsupported Reader style {style!r}; expected 'pil' or 'numpy'." + ) + self.style = style.lower() + + def __call__( + self, x: str | Path | TensorLike | Image.Image + ) -> np.ndarray | Image.Image: + """Reads/converts the input into an image object. + + Args: + x (str | Path | TensorLike | Image.Image): Input image path or image object. + + Returns: + np.ndarray | Image.Image: Read image in the specified style. + """ + if self.style == "numpy": + if isinstance(x, np.ndarray): + return x + elif isinstance(x, torch.Tensor): + return x.detach().cpu().numpy() + elif isinstance(x, (str, Path)): + image = cv2.imread(str(x)) + if image is None: + raise FileNotFoundError(f"Image not found: {x}") + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + elif isinstance(x, Image.Image): + return np.array(x) + else: + raise TypeError( + f"Reader(style='numpy') does not support input type {type(x).__name__}." + ) + elif self.style == "pil": + if isinstance(x, np.ndarray): + return Image.fromarray( + normalize_uint8_rgb_array(x, operation="Reader(style='pil')") + ) + elif isinstance(x, torch.Tensor): + x = x.detach().cpu().numpy() + return Image.fromarray( + normalize_uint8_rgb_array(x, operation="Reader(style='pil')") + ) + elif isinstance(x, (str, Path)): + return Image.open(x).convert("RGB") + elif isinstance(x, Image.Image): + return x + else: + raise TypeError( + f"Reader(style='pil') does not support input type {type(x).__name__}." + ) + else: + raise RuntimeError( + f"Reader has an invalid validated style: {self.style!r}." + ) diff --git a/mblt_vision/utils/preprocess/resize.py b/mblt_vision/utils/preprocess/resize.py new file mode 100644 index 0000000..e1ed399 --- /dev/null +++ b/mblt_vision/utils/preprocess/resize.py @@ -0,0 +1,202 @@ +""" +Image resizing preprocessing. +""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from ..types import TensorLike +from ._validation import normalize_image_size +from .base import PreOps + +PIL_INTERP_CODES = { + "nearest": Image.Resampling.NEAREST, + "bilinear": Image.Resampling.BILINEAR, + "bicubic": Image.Resampling.BICUBIC, + "box": Image.Resampling.BOX, + "hamming": Image.Resampling.HAMMING, + "lanczos": Image.Resampling.LANCZOS, +} +TORCH_INTERPOLATION_MODES = frozenset({"nearest", "bilinear", "bicubic"}) + + +class Resize(PreOps): + """Resizes the image to a specified size using various interpolation modes. + + Supports PyTorch tensors in CHW or BCHW format, HWC NumPy arrays, and PIL images. + """ + + def __init__( + self, + size: int | list[int], + interpolation: str, + ) -> None: + """ + Initialize the Resize operation. + Args: + size (int | list[int]): Target size. If int, the shorter edge is resized to this size + maintaining aspect ratio. If [h, w], it is resized to exactly this size. + interpolation (str): Interpolation mode (e.g., "bilinear", "bicubic", "nearest"). + """ + # Note that this behaves different for npy image and PIL image + super().__init__() + self.size = ( + size + if isinstance(size, int) and not isinstance(size, bool) + else normalize_image_size(size) + ) + if isinstance(self.size, int) and self.size <= 0: + raise ValueError(f"size must be positive, got {self.size}.") + if interpolation not in PIL_INTERP_CODES: + raise ValueError( + f"Unsupported resize interpolation {interpolation!r}; expected one of {sorted(PIL_INTERP_CODES)}." + ) + self.interpolation = interpolation + + def __call__( + self, x: TensorLike | Image.Image + ) -> np.ndarray | torch.Tensor | Image.Image: + """Resizes the input image. + + Args: + x (TensorLike | Image.Image): Image to be resized. + + Returns: + np.ndarray | torch.Tensor | Image.Image: Resized image in the same format as input. + + Raises: + TypeError: If input type is not supported. + ValueError: If an input has an invalid layout or uses a PIL-only interpolation mode. + """ + if isinstance(x, np.ndarray): + self._validate_tensor_interpolation() + if x.ndim != 3: + raise ValueError( + f"Expected an HWC NumPy array, but got x.shape={x.shape}." + ) + img_h, img_w = x.shape[:2] + new_h, new_w = self._compute_resized_output_size(img_h, img_w) + if [img_h, img_w] == [new_h, new_w]: + return x + + tensor_x = torch.from_numpy(x).to(self.device) + tensor_x = tensor_x.permute(2, 0, 1) + tensor_x, need_cast, need_squeeze, out_dtype = self._cast_squeeze_in( + tensor_x, [torch.float32, torch.float64] + ) + tensor_x = F.interpolate( + tensor_x, + size=(new_h, new_w), + mode=self.interpolation, + align_corners=( + False if self.interpolation in ["bilinear", "bicubic"] else None + ), + antialias=self.interpolation in ["bilinear", "bicubic"], + ) + tensor_x = self._cast_squeeze_out( + tensor_x, need_cast, need_squeeze, out_dtype + ) + return tensor_x.permute(1, 2, 0).cpu().numpy() + elif isinstance(x, torch.Tensor): + self._validate_tensor_interpolation() + tensor_x = x.to(self.device) + elif isinstance(x, Image.Image): + img_w, img_h = x.size + new_h, new_w = self._compute_resized_output_size(img_h, img_w) + if [img_h, img_w] == [new_h, new_w]: + return x + return x.resize( + size=(new_w, new_h), + resample=PIL_INTERP_CODES[self.interpolation], + ) + else: + raise TypeError(f"Got unexpected type for x={type(x)}.") + + if tensor_x.ndim not in (3, 4): + raise ValueError( + f"Expected a CHW or BCHW tensor, but got x.shape={tensor_x.shape}." + ) + img_h, img_w = tensor_x.shape[-2:] + new_h, new_w = self._compute_resized_output_size(img_h, img_w) + if [img_h, img_w] == [new_h, new_w]: + return tensor_x + tensor_x, need_cast, need_squeeze, out_dtype = self._cast_squeeze_in( + tensor_x, [torch.float32, torch.float64] + ) + tensor_x = F.interpolate( + tensor_x, + size=(new_h, new_w), + mode=self.interpolation, + align_corners=( + False if self.interpolation in ["bilinear", "bicubic"] else None + ), + antialias=self.interpolation in ["bilinear", "bicubic"], + ) + tensor_x = self._cast_squeeze_out(tensor_x, need_cast, need_squeeze, out_dtype) + return tensor_x.to(self.device) + + def _validate_tensor_interpolation(self) -> None: + """Reject interpolation modes unsupported by PyTorch tensors.""" + + if self.interpolation not in TORCH_INTERPOLATION_MODES: + raise ValueError( + f"Resize interpolation {self.interpolation!r} is supported only for PIL images; " + f"NumPy arrays and tensors require one of {sorted(TORCH_INTERPOLATION_MODES)}." + ) + + def _compute_resized_output_size(self, img_h: int, img_w: int) -> list[int]: + if isinstance(self.size, int): + # to match the shortest side to self.size with the same ratio + if img_w <= img_h: + new_w = self.size + new_h = int(self.size * img_h / img_w) + else: + new_h = self.size + new_w = int(self.size * img_w / img_h) + elif isinstance(self.size, list): + new_h, new_w = self.size + else: + raise RuntimeError(f"Resize has an invalid validated size: {self.size!r}.") + return [new_h, new_w] + + def _cast_squeeze_in( + self, img: torch.Tensor, req_dtypes: list[torch.dtype] + ) -> tuple[torch.Tensor, bool, bool, torch.dtype]: + need_squeeze = False + # make image NCHW + if img.ndim < 4: + img = img.unsqueeze(dim=0) + need_squeeze = True + out_dtype = img.dtype + need_cast = False + if out_dtype not in req_dtypes: + need_cast = True + req_dtype = req_dtypes[0] + img = img.to(req_dtype) + return img, need_cast, need_squeeze, out_dtype + + def _cast_squeeze_out( + self, + img: torch.Tensor, + need_cast: bool, + need_squeeze: bool, + out_dtype: torch.dtype, + ) -> torch.Tensor: + if need_squeeze: + img = img.squeeze(dim=0) + if need_cast: + if out_dtype in ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ): + # it is better to round before cast + img = torch.round(img) + img = img.to(out_dtype) + return img diff --git a/mblt_vision/utils/preprocess/yolo_pre.py b/mblt_vision/utils/preprocess/yolo_pre.py new file mode 100644 index 0000000..f2fc151 --- /dev/null +++ b/mblt_vision/utils/preprocess/yolo_pre.py @@ -0,0 +1,7 @@ +"""Compatibility name for the extracted letterbox operation.""" + +from .letterbox import LetterBox + +YoloPre = LetterBox + +__all__ = ["YoloPre"] diff --git a/mblt_vision/utils/results.py b/mblt_vision/utils/results.py new file mode 100644 index 0000000..884a333 --- /dev/null +++ b/mblt_vision/utils/results.py @@ -0,0 +1,725 @@ +""" +Results processing and plotting. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import cast + +import cv2 +import numpy as np +import torch +from PIL import Image + +from .._tasks import normalize_vision_task +from .datasets import ( + get_ade20k_palette, + get_cityscapes_palette, + get_coco_det_palette, + get_coco_keypoint_palette, + get_coco_label, + get_coco_limb_palette, + get_coco_pose_skeleton, + get_dotav1_label, + get_dotav1_palette, + get_imagenet_label, +) +from .letterbox import LetterBoxGeometry +from mblt_vision.utils.postprocess.common import ( + crop_mask, + scale_boxes, + scale_coords, + scale_masks, + scale_rboxes, + xywhr2xyxyxyxy, +) +from .preprocess._validation import normalize_uint8_rgb_array +from .types import ListTensorLike, NestedListTensorLike, TensorLike + +LW = 2 # line width +RADIUS = 5 # circle radius +ALPHA = 0.3 # alpha for overlay +DENSE_OVERLAY_ALPHA = 0.6 + + +class Results: + """Handle, process, and plot model inference results.""" + + def __init__( + self, + pre_cfg: dict, + post_cfg: dict, + output: TensorLike | ListTensorLike | NestedListTensorLike, + **kwargs, + ) -> None: + """ + Initializes the Results object. + Args: + pre_cfg (dict): Preprocessing configuration. + post_cfg (dict): Postprocessing configuration. + output (TensorLike | ListTensorLike | NestedListTensorLike): Raw model output. + **kwargs: Additional arguments. + """ + self.pre_cfg = pre_cfg + self.post_cfg = post_cfg + self.task = normalize_vision_task(post_cfg["task"]) + self.conf_thres = kwargs.get("conf_thres", 0.25) + self.acc: torch.Tensor | np.ndarray | None = None + self.box_cls: torch.Tensor | np.ndarray | None = None + self.mask: torch.Tensor | np.ndarray | None = None + self.depth: torch.Tensor | np.ndarray | list[TensorLike] | None = None + self.semantic_mask: torch.Tensor | np.ndarray | list[TensorLike] | None = None + self.output: TensorLike | ListTensorLike | NestedListTensorLike | None = None + self.labels: torch.Tensor | None = None + self.scores: torch.Tensor | None = None + self.boxes: torch.Tensor | None = None + self.rboxes: torch.Tensor | None = None + self.kpts: torch.Tensor | None = None + self.set_output(output) + + def _read_image( + self, source_path: str | Path | np.ndarray | Image.Image + ) -> np.ndarray: + """ + Internal method to read an image from various input types and convert to BGR format. + Args: + source_path (str | np.ndarray | Image.Image): Path to image or image object. + Returns: + np.ndarray: Image in BGR format (cv2 style). + """ + source_img = None + if isinstance(source_path, Image.Image): # PIL image open + source_img = source_path.convert("RGB") + source_img = np.array(source_img) + source_img = cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR) + elif isinstance(source_path, np.ndarray): + source_img = np.array(source_path) + if source_img.ndim != 3 or source_img.shape[2] != 3: + raise ValueError( + f"Image arrays must have HWC shape with three channels, got {source_img.shape}." + ) + source_img = normalize_uint8_rgb_array(source_img, operation="Results.plot") + source_img = cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR) + elif isinstance(source_path, (str, Path)): + image_path = Path(source_path) + if not image_path.is_file(): + raise FileNotFoundError(f"Image file not found: {image_path}") + source_img = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + else: + raise TypeError( + f"Unsupported image source type: {type(source_path).__name__}." + ) + if source_img is None: + raise ValueError(f"Failed to decode image from {source_path!r}.") + return source_img + + @staticmethod + def _save_image(save_path: str | Path, image: np.ndarray) -> None: + """Save an image and report encoder or filesystem failures.""" + + path = Path(save_path) + if path.parent != Path("."): + path.parent.mkdir(parents=True, exist_ok=True) + if not cv2.imwrite(str(path), image): + raise OSError(f"Failed to write result image: {path}") + + def set_output( + self, output: TensorLike | ListTensorLike | NestedListTensorLike + ) -> None: + """ + Sets variables from the raw model output based on the task. + Args: + output (TensorLike | ListTensorLike | NestedListTensorLike): Raw model output. + Raises: + NotImplementedError: If the task is not supported. + """ + self.acc = None + self.box_cls = None + self.mask = None + self.depth = None + self.semantic_mask = None + if self.task == "image_classification": + if not isinstance(output, (np.ndarray, torch.Tensor)): + raise TypeError( + f"Expected tensor output for task {self.task}, got {type(output).__name__}." + ) + self.acc = cast(TensorLike, output) + elif self.task in { + "object_detection", + "face_detection", + "pose_estimation", + "obb", + }: + if not isinstance(output, Sequence): + raise TypeError( + f"Expected list output for task {self.task}, got {type(output).__name__}." + ) + if len(output) == 0: + raise ValueError( + f"Expected a non-empty output list for task {self.task}." + ) + if not isinstance(output[0], (np.ndarray, torch.Tensor)): + raise TypeError( + f"Expected a tensor as the first output for task {self.task}, got {type(output[0]).__name__}." + ) + self.box_cls = cast(TensorLike, output[0]) + elif self.task == "instance_segmentation": + if not isinstance(output, Sequence): + raise TypeError( + f"Expected nested list output for task {self.task}, got {type(output).__name__}." + ) + if len(output) == 0: + raise ValueError( + f"Expected a non-empty output list for task {self.task}." + ) + if not isinstance(output[0], Sequence): + raise TypeError( + f"Expected a nested output sequence for task {self.task}, got {type(output[0]).__name__}." + ) + if len(output[0]) < 2: + raise ValueError( + "Instance segmentation output must contain detections and masks." + ) + seg_output = cast(ListTensorLike, output[0]) + if not all( + isinstance(item, (np.ndarray, torch.Tensor)) for item in seg_output[:2] + ): + raise TypeError( + "Instance segmentation detections and masks must be tensors." + ) + self.box_cls = cast(TensorLike, seg_output[0]) + self.mask = cast(TensorLike, seg_output[1]) + elif self.task == "depth_estimation": + if isinstance(output, Sequence) and not isinstance( + output, (np.ndarray, torch.Tensor) + ): + if len(output) == 0: + raise ValueError("Expected at least one depth-map tensor.") + if not all( + isinstance(item, (np.ndarray, torch.Tensor)) for item in output + ): + raise TypeError( + f"Expected depth-map tensors for task {self.task}, got {type(output).__name__}." + ) + self.depth = [cast(TensorLike, item) for item in output] + elif isinstance(output, (np.ndarray, torch.Tensor)): + self.depth = output + else: + raise TypeError( + f"Expected tensor depth output for task {self.task}, got {type(output)}." + ) + elif self.task == "semantic_segmentation": + if isinstance(output, Sequence) and not isinstance( + output, (np.ndarray, torch.Tensor) + ): + if len(output) == 0: + raise ValueError("Expected at least one semantic-map tensor.") + if not all( + isinstance(item, (np.ndarray, torch.Tensor)) for item in output + ): + raise TypeError( + f"Expected semantic-map tensors for task {self.task}, got {type(output).__name__}." + ) + self.semantic_mask = [cast(TensorLike, item) for item in output] + elif isinstance(output, (np.ndarray, torch.Tensor)): + self.semantic_mask = output + else: + raise TypeError( + f"Expected tensor semantic output for task {self.task}, got {type(output)}." + ) + else: + raise NotImplementedError( + f"Task {self.task} is not supported for plotting results." + ) + self.output = output # store raw output + + def plot( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray | None: + """Plot inference results on the source image. + + Args: + source_path: Image path or object to plot on. + save_path: Optional output image path. + **kwargs: Additional task-specific plotting options (e.g., topk for classification). + + Returns: + Image with results visualized in BGR format, or ``None`` for classification without an output path. + + Raises: + NotImplementedError: If the task is not supported for plotting. + """ + if self.task == "image_classification": + return self._plot_image_classification(source_path, save_path, **kwargs) + elif self.task in {"object_detection", "face_detection"}: + return self._plot_object_detection(source_path, save_path, **kwargs) + elif self.task == "instance_segmentation": + return self._plot_instance_segmentation(source_path, save_path, **kwargs) + elif self.task == "depth_estimation": + return self._plot_depth_estimation(source_path, save_path, **kwargs) + elif self.task == "semantic_segmentation": + return self._plot_semantic_segmentation(source_path, save_path, **kwargs) + elif self.task == "pose_estimation": + return self._plot_pose_estimation(source_path, save_path, **kwargs) + elif self.task == "obb": + return self._plot_obb(source_path, save_path, **kwargs) + else: + raise NotImplementedError( + f"Task {self.task} is not supported for plotting results." + ) + + def _plot_depth_estimation( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + """Colorize the first depth map with near objects in red and blend it over the original image.""" + + del kwargs + if self.depth is None: + raise ValueError("No depth output found.") + depth_value = self.depth[0] if isinstance(self.depth, list) else self.depth + depth = ( + depth_value.detach().cpu().numpy() + if isinstance(depth_value, torch.Tensor) + else depth_value + ) + if depth.ndim == 3: + depth = depth[0] + if depth.ndim != 2: + raise ValueError( + f"Expected a 2D depth map or [B, H, W], got {depth.shape}." + ) + image = self._read_image(source_path) + image_shape = (int(image.shape[0]), int(image.shape[1])) + if tuple(depth.shape) != image_shape: + depth = self._restore_depth_map(depth, image_shape) + valid = np.isfinite(depth) & (depth > 0) + if not valid.any(): + raise ValueError("Depth output contains no positive finite values.") + disparity = np.zeros(depth.shape, dtype=np.float32) + disparity[valid] = 1.0 / depth[valid] + lower, upper = np.percentile(disparity[valid], (2, 98)) + if upper <= lower: + upper = lower + 1e-6 + normalized = np.zeros(depth.shape, dtype=np.uint8) + normalized[valid] = np.clip( + (disparity[valid] - lower) * 255 / (upper - lower), 0, 255 + ).astype(np.uint8) + overlay = cv2.applyColorMap(normalized, cv2.COLORMAP_JET) + overlay[~valid] = 0 + result = cv2.addWeighted( + image, 1.0 - DENSE_OVERLAY_ALPHA, overlay, DENSE_OVERLAY_ALPHA, 0 + ) + if save_path is not None: + self._save_image(save_path, result) + return result + + def _plot_semantic_segmentation( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + """Colorize a semantic class map and blend it over the original image.""" + + del kwargs + if self.semantic_mask is None: + raise ValueError("No semantic output found.") + semantic_value = ( + self.semantic_mask[0] + if isinstance(self.semantic_mask, list) + else self.semantic_mask + ) + class_map = ( + semantic_value.detach().cpu().numpy() + if isinstance(semantic_value, torch.Tensor) + else semantic_value + ) + if class_map.ndim == 3: + class_map = class_map[0] + if class_map.ndim != 2: + raise ValueError( + f"Expected a 2D semantic map or [B, H, W], got {class_map.shape}." + ) + image = self._read_image(source_path) + image_shape = (int(image.shape[0]), int(image.shape[1])) + if tuple(class_map.shape) != image_shape: + class_map = self._restore_semantic_map(class_map, image_shape) + dataset_value = self.post_cfg.get("dataset") + dataset = ( + dataset_value.lower() if isinstance(dataset_value, str) else dataset_value + ) + if dataset == "ade20k": + default_nc = 150 + palette_getter = get_ade20k_palette + elif dataset == "cityscapes": + default_nc = 19 + palette_getter = get_cityscapes_palette + else: + raise ValueError( + f"Unsupported semantic segmentation dataset palette: {dataset!r}." + ) + nc = int(self.post_cfg.get("nc", default_nc)) + valid = class_map != 255 + if valid.any() and ( + int(class_map[valid].min()) < 0 or int(class_map[valid].max()) >= nc + ): + raise ValueError( + f"Semantic class-map values must be in [0, {nc - 1}] or 255." + ) + palette = np.array( + [palette_getter(index) for index in range(nc)], dtype=np.uint8 + ) + overlay = np.zeros_like(image) + overlay[valid] = palette[class_map[valid].astype(np.int64)] + blended = cv2.addWeighted( + image, 1.0 - DENSE_OVERLAY_ALPHA, overlay, DENSE_OVERLAY_ALPHA, 0 + ) + result = image.copy() + result[valid] = blended[valid] + if save_path is not None: + self._save_image(save_path, result) + return result + + def _restore_semantic_map( + self, class_map: np.ndarray, image_shape: tuple[int, int] + ) -> np.ndarray: + """Undo the configured letterbox transform using nearest-neighbor interpolation.""" + + return self._restore_dense_map( + class_map, image_shape, cv2.INTER_NEAREST, "Semantic" + ) + + def _restore_depth_map( + self, depth: np.ndarray, image_shape: tuple[int, int] + ) -> np.ndarray: + """Undo the configured letterbox transform and resize a depth map to an image.""" + + return self._restore_dense_map(depth, image_shape, cv2.INTER_LINEAR, "Depth") + + def _restore_dense_map( + self, + output: np.ndarray, + image_shape: tuple[int, int], + interpolation: int, + task_name: str, + ) -> np.ndarray: + """Undo configured letterboxing for a dense two-dimensional output.""" + + letterbox_cfg = self.pre_cfg.get("LetterBox", {}) + input_shape = letterbox_cfg.get("img_size") + if not isinstance(input_shape, list) or len(input_shape) != 2: + return cv2.resize( + output, (image_shape[1], image_shape[0]), interpolation=interpolation + ) + geometry = LetterBoxGeometry.from_shapes( + (int(input_shape[0]), int(input_shape[1])), image_shape + ) + output_shape = (int(output.shape[0]), int(output.shape[1])) + top, bottom, left, right = geometry.crop_bounds(output_shape) + cropped = output[top:bottom, left:right] + if cropped.size == 0: + raise ValueError( + f"{task_name} letterbox restoration produced an empty crop." + ) + return cv2.resize( + cropped, (image_shape[1], image_shape[0]), interpolation=interpolation + ) + + def _plot_image_classification( + self, + source_path: str | Path | np.ndarray | Image.Image | None = None, + save_path: str | Path | None = None, + topk: int = 5, + **kwargs, + ) -> np.ndarray | None: + if self.acc is None: + raise ValueError("No accuracy output found.") + if isinstance(topk, bool) or not isinstance(topk, int): + raise TypeError(f"topk must be an integer, got {type(topk).__name__}.") + if topk <= 0: + raise ValueError(f"topk must be positive, got {topk}.") + if isinstance(self.acc, np.ndarray): + self.acc = torch.tensor(self.acc) + scores = self.acc.squeeze() + if scores.ndim != 1: + raise ValueError( + f"Classification plotting expects one class-score vector, got shape {tuple(self.acc.shape)}." + ) + topk = min(topk, int(scores.numel())) + topk_probs, topk_indices = torch.topk(scores, topk) + topk_probs = np.atleast_1d(topk_probs.squeeze().detach().cpu().numpy()) + topk_indices = np.atleast_1d(topk_indices.squeeze().detach().cpu().numpy()) + # load labels + labels = [get_imagenet_label(i) for i in topk_indices] + comments = [] + for i in range(topk): + comments.append(f"{labels[i]}: {topk_probs[i] * 100:.2f}%") + print(f"Label: {labels[i]}, Probability: {topk_probs[i] * 100:.2f}%") + if source_path is not None and save_path is not None: + comments_str = "\n".join(comments) + img = self._read_image(source_path) + avg_color = img.mean(axis=(0, 1)) + txt_color = ( + int(255 - avg_color[0]), + int(255 - avg_color[1]), + int(255 - avg_color[2]), + ) + for i, line in enumerate(comments_str.splitlines()): + (_, h), _ = cv2.getTextSize( + text=line, + fontFace=cv2.FONT_HERSHEY_SIMPLEX, + fontScale=0.5, + thickness=1, + ) + img = cv2.putText( + img, + line, + (15, 15 + int(1.5 * i * h)), # line spacing + fontFace=cv2.FONT_HERSHEY_SIMPLEX, + fontScale=0.5, + color=txt_color, + thickness=1, + lineType=cv2.LINE_AA, + ) + self._save_image(save_path, img) + return img + else: + return None + + def _plot_object_detection( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + box_cls = self._box_cls_tensor() + expected_columns = 6 + self.post_cfg.get("n_extra", 0) + if box_cls.ndim != 2 or box_cls.shape[1] != expected_columns: + raise ValueError( + f"Object detection output must have shape [N, {expected_columns}], got {tuple(box_cls.shape)}." + ) + img = self._read_image(source_path) + img1_shape = cast(tuple[int, int], self.pre_cfg["LetterBox"]["img_size"]) + img0_shape: tuple[int, int] = (img.shape[0], img.shape[1]) + self.labels = box_cls[:, 5].to(torch.int64) + self.scores = box_cls[:, 4] + self.boxes = scale_boxes( + img1_shape, + box_cls[:, :4].clone(), + img0_shape, + ) + boxes = self.boxes + scores = self.scores + labels = self.labels + contours: dict[int, list[np.ndarray]] = {} + for box, score, label in zip(boxes, scores, labels): + label_idx = int(label.item()) + palette = self._get_detection_palette(label_idx) + img = cv2.putText( + img, + f"{self._get_detection_label(label_idx)} {int(100 * score)}%", + (int(box[0]), int(box[1]) - 10), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + palette, + 1, + cv2.LINE_AA, + ) + contours.setdefault(label_idx, []).append( + np.array( + [ + [int(box[0]), int(box[1])], + [int(box[2]), int(box[1])], + [int(box[2]), int(box[3])], + [int(box[0]), int(box[3])], + ] + ) + ) + for label, contour in contours.items(): + if len(contour) > 0: + cv2.drawContours( + img, + contour, + -1, + self._get_detection_palette(label), + LW, + ) + if save_path is not None: + self._save_image(save_path, img) + return img + + def _plot_instance_segmentation( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + img = self._plot_object_detection(source_path, None, **kwargs) + if self.mask is None: + raise RuntimeError("Instance segmentation output has no mask tensor.") + if self.boxes is None: + raise RuntimeError("Instance segmentation boxes were not initialized.") + if self.labels is None: + raise RuntimeError("Instance segmentation labels were not initialized.") + mask = self._mask_tensor() + img0_shape: tuple[int, int] = (img.shape[0], img.shape[1]) + masks = ( + crop_mask(scale_masks(mask, img0_shape), self.boxes) + .gt_(0.0) + .permute(1, 2, 0) + .to(torch.float32) + .cpu() + .numpy() + ) + overlay = np.zeros((masks.shape[0], masks.shape[1], 3)) + for i, label in enumerate(self.labels): + label_idx = int(label.item()) + overlay = np.maximum( + overlay, + masks[:, :, i][:, :, np.newaxis] + * np.array(get_coco_det_palette(label_idx)).reshape(1, 1, 3), + ) + total_mask = overlay.max(axis=2, keepdims=True) + inv_mask = 1 - ALPHA * total_mask / 255 + img = (img * inv_mask + overlay * ALPHA).astype(np.uint8) + if save_path is not None: + self._save_image(save_path, img) + return img + + def _plot_pose_estimation( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + img = self._plot_object_detection(source_path, None, **kwargs) + box_cls = self._box_cls_tensor() + img0_shape: tuple[int, int] = (img.shape[0], img.shape[1]) + self.kpts = scale_coords( + self.pre_cfg["LetterBox"]["img_size"], + box_cls[:, 6:].reshape(-1, 17, 3).clone(), + img0_shape, + ) + kpts = self.kpts + if kpts is None: + raise ValueError("No keypoints output found.") + for kpt in kpts: + for i, (x, y, v) in enumerate(kpt): + color_k = get_coco_keypoint_palette(i) + if float(v) < self.conf_thres: + continue + cv2.circle( + img, + (int(x), int(y)), + RADIUS, + color_k, + -1, + lineType=cv2.LINE_AA, + ) + for j, sk in enumerate(get_coco_pose_skeleton()): + conf1 = float(kpt[sk[0] - 1, 2]) + conf2 = float(kpt[sk[1] - 1, 2]) + if conf1 < self.conf_thres or conf2 < self.conf_thres: + continue + pos1 = (int(kpt[sk[0] - 1, 0]), int(kpt[sk[0] - 1, 1])) + pos2 = (int(kpt[sk[1] - 1, 0]), int(kpt[sk[1] - 1, 1])) + cv2.line( + img, + pos1, + pos2, + get_coco_limb_palette(j), + thickness=int(np.ceil(LW / 2)), + lineType=cv2.LINE_AA, + ) + if save_path is not None: + self._save_image(save_path, img) + return img + + def _plot_obb( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + """Plot OBB detections on an image. + + Args: + source_path: Path or image object. + save_path: Optional path to save the plotted image. + **kwargs: Additional plotting arguments. + + Returns: + The plotted BGR image. + """ + del kwargs + box_cls = self._box_cls_tensor() + if box_cls.ndim != 2 or box_cls.shape[1] != 7: + raise ValueError( + f"OBB output must have shape [N, 7], got {tuple(box_cls.shape)}." + ) + img = self._read_image(source_path) + img0_shape: tuple[int, int] = (img.shape[0], img.shape[1]) + self.labels = box_cls[:, 5].to(torch.int64) + self.scores = box_cls[:, 4] + self.rboxes = scale_rboxes( + self.pre_cfg["LetterBox"]["img_size"], + torch.cat([box_cls[:, :4], box_cls[:, 6:7]], dim=-1), + img0_shape, + ) + polygons = xywhr2xyxyxyxy(self.rboxes).to(torch.int32).cpu().numpy() + for polygon, score, label in zip(polygons, self.scores, self.labels): + label_idx = int(label.item()) + color = get_dotav1_palette(label_idx) + text_anchor = polygon.min(axis=0) + img = cv2.putText( + img, + f"{get_dotav1_label(label_idx)} {int(100 * score)}%", + (int(text_anchor[0]), int(text_anchor[1]) - 10), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + color, + 1, + cv2.LINE_AA, + ) + cv2.drawContours(img, [polygon.reshape(-1, 1, 2)], -1, color, LW) + if save_path is not None: + self._save_image(save_path, img) + return img + + def _box_cls_tensor(self) -> torch.Tensor: + """Returns detection output as a torch tensor.""" + if self.box_cls is None: + raise ValueError("No box_cls output found.") + if isinstance(self.box_cls, np.ndarray): + return torch.from_numpy(self.box_cls) + return self.box_cls + + def _get_detection_label(self, label_idx: int) -> str: + """Return the display label for detection-style tasks.""" + if self.task == "face_detection": + if label_idx != 0: + raise ValueError(f"Unexpected face_detection class index: {label_idx}.") + return "face" + return get_coco_label(label_idx) + + def _get_detection_palette(self, label_idx: int) -> tuple[int, int, int]: + """Return the display color for detection-style tasks.""" + if self.task == "face_detection": + return get_coco_det_palette(0) + return get_coco_det_palette(label_idx) + + def _mask_tensor(self) -> torch.Tensor: + """Returns segmentation mask output as a torch tensor.""" + if self.mask is None: + raise ValueError("No mask output found.") + if isinstance(self.mask, np.ndarray): + return torch.from_numpy(self.mask) + return self.mask diff --git a/mblt_vision/utils/types.py b/mblt_vision/utils/types.py new file mode 100644 index 0000000..5963c94 --- /dev/null +++ b/mblt_vision/utils/types.py @@ -0,0 +1,15 @@ +""" +Type definitions for MBLT vision models. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TypeAlias + +import numpy as np +import torch + +TensorLike: TypeAlias = torch.Tensor | np.ndarray +ListTensorLike: TypeAlias = Sequence[TensorLike] +NestedListTensorLike: TypeAlias = Sequence[TensorLike | ListTensorLike] diff --git a/mblt_vision/wrapper.py b/mblt_vision/wrapper.py new file mode 100644 index 0000000..73552a3 --- /dev/null +++ b/mblt_vision/wrapper.py @@ -0,0 +1,1121 @@ +""" +Wrapper classes for MBLT model execution. +""" + +from __future__ import annotations + +import copy +import importlib +import os +import stat +import sys +import tempfile +from pathlib import Path +from types import TracebackType +from typing import TYPE_CHECKING, Any, Sequence, cast + +import numpy as np +import torch +import yaml +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import EntryNotFoundError + +from mblt_npu import MobilintNPUBackend, ONNXBackend, normalize_target_device +from ._model_paths import resolve_framework as _resolve_framework +from ._model_paths import split_model_paths as _split_model_paths +from ._model_paths import ( + uses_shifted_engine_model_path_layout as _uses_shifted_engine_model_path_layout, +) +from .utils.postprocess import build_postprocess +from .utils.preprocess import build_preprocess +from .utils.results import Results +from .utils.types import TensorLike + +if TYPE_CHECKING: + from qbruntime import Cluster, CoreId + +MODEL_CONFIG_DIR = Path(__file__).parent / "models" + +ONNXRUNTIME_INSTALL_GUIDE = ( + "onnxruntime is not installed. To use ONNX inference, install one of the optional extras:\n" + "pip install mblt-vision-python[onnxruntime]\n" + + ( + "or\npip install mblt-vision-python[onnxruntime-gpu]" + if sys.platform != "darwin" + else "" + ) +) +CoreMode = str +CORE_MODES: tuple[CoreMode, ...] = ("single", "multi", "global4", "global8") +REGULUS_TARGET_DEVICES = frozenset({"regulus-ra", "regulus-rb"}) + + +def core_modes_for_target_device(target_device: str) -> tuple[CoreMode, ...]: + """Return the Vision core modes supported by a normalized NPU board.""" + + if normalize_target_device(target_device) in REGULUS_TARGET_DEVICES: + return ("single",) + return CORE_MODES + + +def normalize_core_mode( + core_mode: str, *, target_device: str | None = None +) -> CoreMode: + """Validate a Vision engine core mode, including board compatibility.""" + + if core_mode not in CORE_MODES: + raise ValueError( + f"Invalid core mode '{core_mode}'. Expected one of {list(CORE_MODES)}." + ) + if target_device is not None: + normalized_target_device = normalize_target_device(target_device) + supported_modes = core_modes_for_target_device(normalized_target_device) + if core_mode not in supported_modes: + raise ValueError( + f"Core mode '{core_mode}' is not supported by " + f"{normalized_target_device}; expected one of {list(supported_modes)}." + ) + return core_mode + + +__all__ = [ + "CoreMode", + "core_modes_for_target_device", + "MOBILINT_CACHE_DIR", + "get_mobilint_cache_dir", + "normalize_core_mode", + "resolve_model_config", + "MBLT_Engine", +] + + +def _derive_onnx_filename(file_cfg: dict[str, Any]) -> str | None: + """Return the configured or MXQ-derived ONNX artifact filename. + + ``onnx_filename`` is only required when the Hub ONNX artifact does not share + the MXQ artifact stem. Otherwise every model configuration follows the same + ``.mxq`` to ``.onnx`` convention. + """ + + onnx_filename = file_cfg.get("onnx_filename") + if isinstance(onnx_filename, str) and onnx_filename: + return onnx_filename + + filename = file_cfg.get("filename") + if not isinstance(filename, str) or not filename: + return None + + return f"{Path(filename).stem}.onnx" + + +def _normalize_model_artifacts(model_config: dict[str, Any]) -> dict[str, Any]: + """Populate the derived ONNX artifact name in a resolved model configuration.""" + + normalized = copy.deepcopy(model_config) + file_cfg = normalized.get("file_cfg") + if not isinstance(file_cfg, dict): + return normalized + + onnx_filename = _derive_onnx_filename(file_cfg) + if onnx_filename is not None: + file_cfg["onnx_filename"] = onnx_filename + return normalized + + +def _default_cache_dir() -> str: + """Returns a writable cache directory for downloaded vision artifacts.""" + + preferred = Path(os.path.expanduser("~/.mblt_model_zoo")) + try: + preferred.mkdir(parents=True, mode=0o700, exist_ok=True) + preferred_stat = os.lstat(preferred) + if ( + not stat.S_ISDIR(preferred_stat.st_mode) + or stat.S_ISLNK(preferred_stat.st_mode) + or preferred_stat.st_uid != os.getuid() + or preferred_stat.st_mode & 0o077 + ): + return _fallback_cache_dir() + with tempfile.NamedTemporaryFile(prefix=".write_test-", dir=preferred): + pass + return str(preferred) + except OSError: + return _fallback_cache_dir() + + +def _fallback_cache_dir() -> str: + """Return a stable private fallback cache when the home cache is unavailable.""" + + uid = os.getuid() + fallback = Path(tempfile.gettempdir()) / f"mblt_model_zoo-{uid}" + try: + os.mkdir(fallback, mode=0o700) + except FileExistsError: + pass + except OSError as exc: + raise RuntimeError( + f"Unable to create fallback Mobilint cache directory: {fallback}" + ) from exc + + fallback_stat = os.lstat(fallback) + if ( + not stat.S_ISDIR(fallback_stat.st_mode) + or stat.S_ISLNK(fallback_stat.st_mode) + or fallback_stat.st_uid != uid + or fallback_stat.st_mode & 0o077 + ): + raise RuntimeError( + f"Fallback Mobilint cache directory is not a private directory owned by this user: {fallback}" + ) + + try: + with tempfile.NamedTemporaryFile(prefix=".write_test-", dir=fallback): + pass + except OSError as exc: + raise RuntimeError( + f"Fallback Mobilint cache directory is not writable: {fallback}" + ) from exc + return str(fallback) + + +MOBILINT_CACHE_DIR = os.path.expanduser("~/.mblt_model_zoo") +"""Preferred artifact cache root, without creating it during import.""" + +_resolved_cache_dir: str | None = None + + +def get_mobilint_cache_dir() -> str: + """Return a writable artifact cache directory, creating it only when needed.""" + + global _resolved_cache_dir + if _resolved_cache_dir is None: + _resolved_cache_dir = _default_cache_dir() + return _resolved_cache_dir + + +def _load_onnxruntime() -> Any: + """Loads ``onnxruntime`` only when ONNX inference is requested. + + Returns: + The imported ``onnxruntime`` module. + + Raises: + ImportError: If ``onnxruntime`` is unavailable in the current environment. + """ + + try: + module = importlib.import_module("onnxruntime") + except ImportError as exc: + raise ImportError(ONNXRUNTIME_INSTALL_GUIDE) from exc + + if not hasattr(module, "InferenceSession"): + module_path = getattr(module, "__file__", None) or "" + msg = ( + "onnxruntime is installed, but the package is incomplete or broken and does not expose " + f"`InferenceSession` (resolved from {module_path}). " + f"{ONNXRUNTIME_INSTALL_GUIDE.replace('is not installed. To use ONNX inference, install', 'Reinstall')}" + ) + raise ImportError(msg) + + return module + + +def _resolve_onnx_providers( + ort_module: Any, requested_providers: Sequence[str] | None = None +) -> list[str]: + """Selects ONNX Runtime execution providers. + + Args: + ort_module: Imported ``onnxruntime`` module or compatible test double. + requested_providers: Optional provider order requested by the caller. + + Returns: + The provider list passed to ``InferenceSession``. + """ + + return ( + list(requested_providers) + if requested_providers is not None + else ["CPUExecutionProvider"] + ) + + +def _model_name_aliasing(model_name: str) -> str: + """Find the YAML filename matching a model name. + + Args: + model_name: Model identifier provided by the caller. + + Returns: + Exact YAML filename stored in ``MODEL_CONFIG_DIR``. + + Raises: + ValueError: If no YAML file matches or the normalized name is ambiguous. + """ + + def _stem(name: str) -> str: + return name[: -len(".yaml")] if name.lower().endswith(".yaml") else name + + def _normalize_separators(name: str) -> str: + return "_".join( + part + for part in _stem(name) + .replace("-", "_") + .replace(" ", "_") + .lower() + .split("_") + if part + ) + + requested = _normalize_separators(model_name) + config_names = sorted(path.name for path in MODEL_CONFIG_DIR.glob("*.yaml")) + separator_matches = [ + name for name in config_names if _normalize_separators(name) == requested + ] + if len(separator_matches) == 1: + return separator_matches[0] + if len(separator_matches) > 1: + raise ValueError( + f"Ambiguous model name '{model_name}'. Matches: {separator_matches}." + ) + + compact_requested = requested.replace("_", "") + compact_matches = [ + name + for name in config_names + if _normalize_separators(name).replace("_", "") == compact_requested + ] + if len(compact_matches) == 1: + return compact_matches[0] + if len(compact_matches) > 1: + raise ValueError( + f"Ambiguous model name '{model_name}'. Matches: {compact_matches}." + ) + raise ValueError(f"Model name '{model_name}' is not supported.") + + +def resolve_model_config( + model_cls: str | dict[str, Any], model_type: str = "DEFAULT" +) -> dict[str, Any]: + """Resolve a vision model configuration without constructing a runtime. + + Args: + model_cls: Model name, YAML path, or direct model configuration. + model_type: Variant key within a YAML model definition. + + Returns: + A deep copy of the resolved ``file_cfg``, ``pre_cfg``, and ``post_cfg`` mapping. + ``file_cfg`` always includes an ONNX filename when an MXQ filename is available. + + Raises: + TypeError: If the YAML or resolved configuration is not a mapping. + ValueError: If a requested variant, alias, or update base is unavailable. + """ + + if isinstance(model_cls, dict): + return _normalize_model_artifacts(model_cls) + + config_path = Path(model_cls) + if not config_path.is_file(): + config_path = MODEL_CONFIG_DIR / _model_name_aliasing(model_cls) + + with config_path.open(encoding="utf-8") as config_file: + full_config = yaml.safe_load(config_file) + if not isinstance(full_config, dict): + raise TypeError( + f"Model configuration '{config_path}' should define a dictionary." + ) + + resolving: set[str] = set() + + def _resolve_variant(variant: str) -> dict[str, Any]: + if variant in resolving: + raise ValueError( + f"Circular model configuration reference detected for '{variant}'." + ) + resolving.add(variant) + try: + model_config_part = full_config.get(variant) + if model_config_part is None: + raise ValueError(f"Model type '{variant}' not found in configuration.") + if isinstance(model_config_part, str): + if model_config_part not in full_config: + raise ValueError( + f"Model alias '{model_config_part}' not found in configuration." + ) + return _resolve_variant(model_config_part) + if not isinstance(model_config_part, dict): + raise TypeError( + f"Resolved model configuration for '{variant}' is not a dictionary." + ) + + resolved = copy.deepcopy(model_config_part) + base_config_key = resolved.pop("update", None) + if base_config_key is None: + return resolved + if ( + not isinstance(base_config_key, str) + or base_config_key not in full_config + ): + raise ValueError( + f"Base configuration '{base_config_key}' not found for update." + ) + + merged_config = _resolve_variant(base_config_key) + for key, value in resolved.items(): + if ( + key in merged_config + and isinstance(merged_config[key], dict) + and isinstance(value, dict) + ): + merged_config[key].update(value) + else: + merged_config[key] = value + return merged_config + finally: + resolving.remove(variant) + + return _normalize_model_artifacts(_resolve_variant(model_type)) + + +class MBLT_Engine: + """Main engine class for running vision models from the MBLT zoo. + + Handles the full pipeline: Preprocessing -> Inference -> Postprocessing. + + Attributes: + file_cfg: Model configuration. + pre_cfg: Preprocessing configuration. + post_cfg: Postprocessing configuration. + model: The underlying MXQ_Model. + device: The torch device being used. + """ + + def __init__( + self, + model_cls: str | dict[str, Any], + model_type: str = "DEFAULT", + mxq_path: str = "", + onnx_path: str = "", + dev_no: int | None = None, + core_mode: CoreMode | None = None, + target_cores: Sequence[str | CoreId] | None = None, + target_clusters: Sequence[int | Cluster] | None = None, + postprocess_kwargs: dict[str, Any] | None = None, + framework: str | None = None, + onnx_providers: Sequence[str] | None = None, + model_path: str = "", + target_device: str | None = None, + ) -> None: + """Initializes the MBLT_Engine. + + Args: + model_cls(if dict): + file_cfg: Model configuration. + mxq_path: path to mxq file + onnx_path: path to onnx file + model_path: generic path to local model file + dev_no: Accelerator No. + core_mode: single, multi, global4, global8 + target_cores: single mode + target_clusters: multi, global modes + target_device: NPU board identifier; defaults to ``aries-rb``. + pre_cfg: Preprocessing configuration. + post_cfg: Postprocessing configuration. + model_cls(not dict): model name or yaml path + postprocess_kwargs: Optional runtime overrides passed to the postprocessor builder. + framework: Execution framework, either "mxq" or "onnx". When omitted, + ``model_path`` suffix is used first, then MXQ is the fallback. + onnx_providers: Optional ONNX Runtime execution provider order. + target_device: NPU board identifier. An explicit value takes precedence + over ``file_cfg.target_device``; otherwise defaults to ``aries-rb``. + """ + + if _uses_shifted_engine_model_path_layout( + model_path, + mxq_path, + dev_no, + core_mode, + target_cores, + postprocess_kwargs, + framework, + onnx_providers, + ): + ( + model_path, + mxq_path, + onnx_path, + dev_no, + core_mode, + target_cores, + target_clusters, + postprocess_kwargs, + framework, + onnx_providers, + ) = ( + mxq_path, + onnx_path, + cast(str, dev_no or ""), + cast(int | None, core_mode), + cast(CoreMode | None, target_cores), + cast(Sequence[str | Any] | None, target_clusters), + cast(Sequence[int | Any] | None, postprocess_kwargs), + cast(dict[str, Any] | None, framework), + cast(str | None, onnx_providers), + cast(Sequence[str] | None, model_path or None), + ) + model_config_part = resolve_model_config(model_cls, model_type) + for section in ("file_cfg", "pre_cfg", "post_cfg"): + if not isinstance(model_config_part.get(section), dict): + raise ValueError( + f"Model configuration section '{section}' must be a mapping." + ) + + if mxq_path and Path(mxq_path).suffix.lower() != ".mxq": + raise ValueError(f"Explicit mxq_path must end in '.mxq', got {mxq_path!r}.") + if onnx_path and Path(onnx_path).suffix.lower() != ".onnx": + raise ValueError( + f"Explicit onnx_path must end in '.onnx', got {onnx_path!r}." + ) + for key, suffix in (("mxq_path", ".mxq"), ("onnx_path", ".onnx")): + configured_path = model_config_part["file_cfg"].get(key) + if configured_path and ( + not isinstance(configured_path, str) + or Path(configured_path).suffix.lower() != suffix + ): + raise ValueError( + f"Configured file_cfg.{key} must end in {suffix!r}, " + f"got {configured_path!r}." + ) + + file_cfg_model_path = str(model_config_part["file_cfg"].get("model_path", "")) + file_cfg_onnx_path = str(model_config_part["file_cfg"].get("onnx_path", "")) + framework_model_path = model_path or file_cfg_model_path + if not framework_model_path and not mxq_path: + framework_model_path = onnx_path or file_cfg_onnx_path + self.framework = _resolve_framework(framework, framework_model_path) + mxq_path, onnx_path = _split_model_paths( + framework=self.framework, + model_path=model_path, + mxq_path=mxq_path, + onnx_path=onnx_path, + ) + + _mxq_path_passed = bool(mxq_path) + _onnx_path_passed = bool(onnx_path) + if _mxq_path_passed and not os.path.isfile(mxq_path): + raise FileNotFoundError( + "Explicit MXQ model path does not exist: " + f"{mxq_path}. Remove model_path/mxq_path to download the configured artifact." + ) + if _onnx_path_passed and not os.path.isfile(onnx_path): + raise FileNotFoundError( + "Explicit ONNX model path does not exist: " + f"{onnx_path}. Remove model_path/onnx_path to download the configured artifact." + ) + _dev_no_passed = dev_no is not None + _core_mode_passed = core_mode is not None + _target_cores_passed = target_cores is not None + _target_clusters_passed = target_clusters is not None + + if dev_no is None: + dev_no = 0 + if target_device is None: + target_device = model_config_part["file_cfg"].get( + "target_device", "aries-rb" + ) + target_device = normalize_target_device(target_device) + is_regulus = target_device in {"regulus-ra", "regulus-rb"} + if core_mode is None: + core_mode = "single" if is_regulus else "global8" + else: + core_mode = normalize_core_mode(core_mode, target_device=target_device) + if target_cores is None: + target_cores = ( + [] + if is_regulus + else ["0:0", "0:1", "0:2", "0:3", "1:0", "1:1", "1:2", "1:3"] + ) + if target_clusters is None: + target_clusters = [] if is_regulus else [0, 1] + + self.file_cfg = copy.deepcopy(model_config_part["file_cfg"]) + file_cfg_model_path = self.file_cfg.pop("model_path", "") + if file_cfg_model_path: + yaml_mxq_path, yaml_onnx_path = _split_model_paths( + framework=self.framework, + model_path=file_cfg_model_path, + mxq_path=str(self.file_cfg.get("mxq_path", "")), + onnx_path=str(self.file_cfg.get("onnx_path", "")), + ) + self.file_cfg["mxq_path"] = yaml_mxq_path + self.file_cfg["onnx_path"] = yaml_onnx_path + if _mxq_path_passed or "mxq_path" not in self.file_cfg: + self.file_cfg["mxq_path"] = mxq_path + if _onnx_path_passed or "onnx_path" not in self.file_cfg: + self.file_cfg["onnx_path"] = onnx_path + if _core_mode_passed or "core_mode" not in self.file_cfg or is_regulus: + self.file_cfg["core_mode"] = core_mode + if _target_cores_passed or "target_cores" not in self.file_cfg: + self.file_cfg["target_cores"] = target_cores + if _target_clusters_passed or "target_clusters" not in self.file_cfg: + self.file_cfg["target_clusters"] = target_clusters + if _dev_no_passed or "dev_no" not in self.file_cfg: + self.file_cfg["dev_no"] = dev_no + self.file_cfg["core_mode"] = normalize_core_mode( + self.file_cfg["core_mode"], target_device=target_device + ) + self.file_cfg["target_device"] = target_device + + self.pre_cfg = copy.deepcopy(model_config_part["pre_cfg"]) + self.post_cfg = copy.deepcopy(model_config_part["post_cfg"]) + self.postprocess_kwargs = ( + {} if postprocess_kwargs is None else dict(postprocess_kwargs) + ) + self.file_config_cleansing() + + self.model: Any + self._mxq_model: MobilintNPUBackend | None = None + self._onnx_model: ONNXBackend | None = None + self._onnx_session: Any = None + self._closed = False + + try: + if self.framework == "onnx": + ort = _load_onnxruntime() + resolved_onnx_path = self.file_cfg.get("onnx_path") + if not resolved_onnx_path: + raise RuntimeError( + f"ONNX path not resolved for model {model_cls}. Make sure the model repository has an ONNX file." + ) + if not os.path.isfile(resolved_onnx_path): + raise FileNotFoundError( + f"ONNX file not found at: {resolved_onnx_path}" + ) + + providers = _resolve_onnx_providers(ort, onnx_providers) + onnx_model = ONNXBackend( + resolved_onnx_path, providers=providers, ort_module=ort + ) + self._onnx_model = onnx_model + onnx_model.create() + self._onnx_session = onnx_model.session + self.model = self._onnx_session + onnx_inputs = self._onnx_session.get_inputs() + if len(onnx_inputs) != 1: + raise ValueError( + "ONNX models must declare exactly one input because " + "MBLT_Engine accepts one preprocessed tensor; got " + f"{len(onnx_inputs)} inputs." + ) + self.input_name = onnx_inputs[0].name + self.output_names = [o.name for o in self._onnx_session.get_outputs()] + else: + mxq_model = MobilintNPUBackend(**self._mxq_backend_kwargs()) + self._mxq_model = mxq_model + self.model = mxq_model + mxq_model.create() + mxq_model.launch() + + if mxq_model.get_dtype() == "DataType.Uint8": + self.pre_cfg.pop("Normalize", None) + + self.preprocessor = build_preprocess(self.pre_cfg) + self.postprocessor = build_postprocess( + self.pre_cfg, self.post_cfg, **self.postprocess_kwargs + ) + self.device = torch.device("cpu") + except Exception: + self._close(suppress_errors=True) + raise + + def _mxq_backend_kwargs(self) -> dict[str, Any]: + """Builds the MXQ backend kwargs from the resolved file config.""" + + excluded_keys = { + "repo_id", + "filename", + "revision", + "onnx_filename", + "onnx_path", + } + return { + key: value + for key, value in self.file_cfg.items() + if key not in excluded_keys + } + + def _derive_onnx_filename(self) -> str | None: + """Returns the ONNX filename associated with the configured MXQ artifact.""" + + onnx_filename = _derive_onnx_filename(self.file_cfg) + if onnx_filename is not None: + self.file_cfg["onnx_filename"] = onnx_filename + return onnx_filename + + def _resolve_local_onnx_path(self, mxq_path: str) -> str | None: + """Tries to resolve a sibling ONNX file next to a local MXQ artifact.""" + + onnx_path = self.file_cfg.get("onnx_path", "") + if onnx_path and os.path.isfile(onnx_path): + return onnx_path + + onnx_filename = self._derive_onnx_filename() + if onnx_filename: + sibling_path = Path(mxq_path).with_name(onnx_filename) + if sibling_path.is_file(): + return str(sibling_path) + + if mxq_path.endswith(".mxq"): + suffix_swapped = f"{mxq_path[:-4]}.onnx" + if os.path.isfile(suffix_swapped): + return suffix_swapped + + return None + + def _download_hub_artifact( + self, + *, + repo_id: str, + filename: str, + revision: str, + subfolders: Sequence[str] | None = None, + ) -> str: + """Downloads a model artifact from Hugging Face Hub and returns its cache path.""" + + last_error: Exception | None = None + normalized_subfolders = [""] if subfolders is None else list(subfolders) + for subfolder in normalized_subfolders: + kwargs: dict[str, Any] = { + "repo_id": repo_id, + "filename": filename, + "revision": revision, + "local_dir": get_mobilint_cache_dir(), + } + if subfolder: + kwargs["subfolder"] = subfolder + try: + return hf_hub_download(**kwargs) + except EntryNotFoundError as exc: + last_error = exc + + attempted_paths = ", ".join( + f"{subfolder}/{filename}" if subfolder else filename + for subfolder in normalized_subfolders + ) + raise RuntimeError( + f"Failed to download model from Hugging Face. Tried repo '{repo_id}' at: {attempted_paths}." + ) from last_error + + def file_config_cleansing(self) -> None: + """Validates and resolves the MXQ and ONNX model file paths in ``self.file_cfg``.""" + framework = getattr(self, "framework", "mxq") + model_path = self.file_cfg.pop("model_path", "") + if model_path: + mxq_path, onnx_path = _split_model_paths( + framework=framework, + model_path=model_path, + mxq_path=str(self.file_cfg.get("mxq_path", "")), + onnx_path=str(self.file_cfg.get("onnx_path", "")), + ) + self.file_cfg["mxq_path"] = mxq_path + self.file_cfg["onnx_path"] = onnx_path + mxq_path = self.file_cfg.get("mxq_path", "") + onnx_path = self.file_cfg.get("onnx_path", "") + onnx_filename = self._derive_onnx_filename() + + if onnx_path and os.path.isfile(onnx_path): + self.file_cfg["onnx_path"] = onnx_path + if framework == "onnx": + return + + if mxq_path and os.path.isfile(mxq_path): + resolved_local_onnx = self._resolve_local_onnx_path(mxq_path) + if resolved_local_onnx is not None: + self.file_cfg["onnx_path"] = resolved_local_onnx + self.file_cfg.pop("repo_id", None) + self.file_cfg.pop("filename", None) + self.file_cfg.pop("revision", None) + elif framework == "mxq": + self.file_cfg.pop("repo_id", None) + self.file_cfg.pop("filename", None) + self.file_cfg.pop("revision", None) + if framework == "mxq": + return + + repo_id = self.file_cfg.pop("repo_id", None) + filename = self.file_cfg.pop("filename", None) + revision = self.file_cfg.pop("revision", None) + if not repo_id or not revision: + return + + if filename and framework == "mxq": + target_device = self.file_cfg.get("target_device", "aries-rb") + self.file_cfg["mxq_path"] = self._download_hub_artifact( + repo_id=repo_id, + filename=filename, + revision=revision, + subfolders=[target_device], + ) + + if onnx_filename and framework == "onnx" and not self.file_cfg.get("onnx_path"): + self.file_cfg["onnx_path"] = self._download_hub_artifact( + repo_id=repo_id, + filename=onnx_filename, + revision=revision, + ) + + def _prepare_onnx_inputs(self, x: TensorLike) -> dict[str, np.ndarray]: + """Normalizes runtime inputs to match the ONNX session contract.""" + + if isinstance(x, torch.Tensor): + x_np = x.detach().cpu().numpy() + elif isinstance(x, np.ndarray): + x_np = x + else: + raise TypeError(f"Got unexpected type for ONNX input x={type(x)}.") + + if x_np.dtype == np.float64: + x_np = x_np.astype(np.float32) + + expected_shape = self._require_onnx_session().get_inputs()[0].shape + if len(expected_shape) == 4: + expected_second_dim = expected_shape[1] + expected_last_dim = expected_shape[-1] + expected_layout = None + expected_channels = None + # ONNX layout is encoded by the channel axis, not by a spatial + # dimension that happens to equal a channel count. This is crucial + # for square static inputs such as [1, 3, 224, 224]. + if isinstance(expected_second_dim, int) and expected_second_dim in { + 1, + 2, + 3, + 4, + }: + expected_layout = "nchw" + expected_channels = expected_second_dim + elif isinstance(expected_last_dim, int) and expected_last_dim in { + 1, + 2, + 3, + 4, + }: + expected_layout = "nhwc" + expected_channels = expected_last_dim + + if x_np.ndim == 3: + if ( + expected_layout == "nchw" + and expected_channels is not None + and x_np.shape[0] == expected_channels + ): + x_np = np.expand_dims(x_np, axis=0) + elif ( + expected_layout == "nchw" + and expected_channels is not None + and x_np.shape[-1] == expected_channels + ): + x_np = np.transpose(x_np, (2, 0, 1)) + x_np = np.expand_dims(x_np, axis=0) + elif ( + expected_layout == "nhwc" + and expected_channels is not None + and x_np.shape[-1] == expected_channels + ): + x_np = np.expand_dims(x_np, axis=0) + elif ( + expected_layout == "nhwc" + and expected_channels is not None + and x_np.shape[0] == expected_channels + ): + x_np = np.transpose(x_np, (1, 2, 0)) + x_np = np.expand_dims(x_np, axis=0) + elif x_np.ndim == 4: + if expected_layout == "nchw" and expected_channels is not None: + if x_np.shape[1] == expected_channels: + pass + elif x_np.shape[-1] == expected_channels: + x_np = np.transpose(x_np, (0, 3, 1, 2)) + elif expected_layout == "nhwc" and expected_channels is not None: + if x_np.shape[-1] == expected_channels: + pass + elif x_np.shape[1] == expected_channels: + x_np = np.transpose(x_np, (0, 2, 3, 1)) + + return {self.input_name: x_np} + + def _require_onnx_session(self) -> Any: + """Return the active ONNX session.""" + + session = getattr(self, "_onnx_session", None) + if session is None: + fallback = getattr(self, "model", None) + if ( + fallback is not None + and hasattr(fallback, "get_inputs") + and hasattr(fallback, "run") + ): + return fallback + raise RuntimeError("ONNX session is not initialized.") + return session + + def _require_mxq_model(self) -> MobilintNPUBackend: + """Return the active MXQ backend.""" + + model = getattr(self, "_mxq_model", None) + if model is None: + fallback = getattr(self, "model", None) + if ( + fallback is not None + and hasattr(fallback, "create") + and hasattr(fallback, "launch") + ): + return cast(MobilintNPUBackend, fallback) + raise RuntimeError("MXQ backend is not initialized.") + return model + + def __call__( + self, + x: TensorLike, + ) -> Any: + """Runs raw model inference on the input. + + Note: + This does NOT include preprocessing or postprocessing. + + Args: + x: Input tensor for the model. + + Returns: + Raw model output. + """ + self._ensure_open() + if self.framework == "onnx": + outputs = self._require_onnx_session().run( + self.output_names, self._prepare_onnx_inputs(x) + ) + if len(outputs) == 1: + return outputs[0] + return outputs + return cast(Any, self._require_mxq_model())(x) + + def preprocess( + self, + x: Any, + **kwargs: Any, + ) -> Any: + """Runs preprocessing on the input. + + Args: + x: Input data. + **kwargs: Additional arguments for preprocessing. + + Returns: + Preprocessed data. + """ + return self.preprocessor(x, **kwargs) + + def preprocess_with_metadata( + self, + x: Any, + ) -> tuple[Any, dict[str, Any]]: + """Runs preprocessing and returns metadata needed for exact postprocess scaling. + + Args: + x: Input data. + + Returns: + A tuple of preprocessed data and metadata such as ``ratio_pad``. + """ + return self.preprocessor.with_metadata(x) + + def postprocess( + self, + x: Any, + **kwargs: Any, + ) -> Results: + """Runs postprocessing on the input. + + Args: + x: Input data. + **kwargs: Additional arguments for postprocessing. + + Returns: + Postprocessed results. + """ + pre_result = self.postprocessor(x, **kwargs) + result_kwargs = dict(kwargs) + conf_thres = getattr(self.postprocessor, "conf_thres", None) + iou_thres = getattr(self.postprocessor, "iou_thres", None) + if "conf_thres" not in result_kwargs and conf_thres is not None: + result_kwargs["conf_thres"] = conf_thres + if "iou_thres" not in result_kwargs and iou_thres is not None: + result_kwargs["iou_thres"] = iou_thres + return Results(self.pre_cfg, self.post_cfg, pre_result, **result_kwargs) + + def set_postprocess_thresholds( + self, conf_thres: float | None = None, iou_thres: float | None = None + ) -> None: + """Updates configurable postprocess thresholds for the current model. + + Args: + conf_thres: Optional confidence threshold override. + iou_thres: Optional IoU threshold override. + + Raises: + NotImplementedError: If the current postprocessor does not support thresholds. + """ + set_threshold = getattr(self.postprocessor, "set_threshold", None) + if set_threshold is None: + raise NotImplementedError( + f"Threshold overrides are not supported for task `{self.post_cfg.get('task', 'unknown')}`." + ) + set_threshold(conf_thres=conf_thres, iou_thres=iou_thres) + + def to( + self, + device: str | torch.device, + ) -> None: + """Moves the engine and its components to the specified device. + + Args: + device: Target device. + + Raises: + TypeError: If device type is unexpected. + """ + self.preprocessor.to(device) + self.postprocessor.to(device) + + if isinstance(device, str): + self.device = torch.device(device) + elif isinstance(device, torch.device): + self.device = device + else: + raise TypeError(f"Got unexpected type for device={type(device)}.") + + def cpu(self) -> None: + """Moves the engine to CPU.""" + self.to(device="cpu") + + def gpu(self) -> None: + """Moves the engine to GPU (CUDA).""" + self.to(device="cuda") + + def cuda( + self, + device: str | int = 0, + ) -> None: + """Moves the engine to CUDA device. + + Args: + device: CUDA device identifier. Defaults to 0. + + Raises: + ValueError: If device string is invalid. + RuntimeError: If CUDA is not available. + """ + if isinstance(device, int): + device = f"cuda:{device}" + elif isinstance(device, str): + if not device.startswith("cuda:"): + raise ValueError("Invalid device string. It should start with 'cuda:'.") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available. Please check your environment.") + self.to(device=device) + + def launch(self) -> None: + """Launches the underlying model.""" + self._ensure_open() + if self.framework == "mxq": + self._require_mxq_model().launch() + + def dispose(self) -> None: + """Compatibility alias for :meth:`close`.""" + + self.close() + + def close(self) -> None: + """Release backend resources. Safe to call more than once.""" + + self._close(suppress_errors=False) + + def __enter__(self) -> MBLT_Engine: + """Return this engine for use in a context manager.""" + + self._ensure_open() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> bool: + """Release resources when leaving a context manager block.""" + + del exc_value, traceback + self._close(suppress_errors=exc_type is not None) + return False + + def __del__(self) -> None: + """Best-effort cleanup for engines not explicitly closed by callers.""" + + try: + self._close(suppress_errors=True) + except Exception: + pass + + def _ensure_open(self) -> None: + """Raise when inference is attempted after backend disposal.""" + + if getattr(self, "_closed", False): + raise RuntimeError("MBLT_Engine is closed.") + + def _close(self, *, suppress_errors: bool) -> None: + """Dispose acquired backends once, optionally suppressing cleanup failures.""" + + if getattr(self, "_closed", False): + return + self._closed = True + first_error: Exception | None = None + for backend in ( + getattr(self, "_mxq_model", None), + getattr(self, "_onnx_model", None), + ): + if backend is None: + continue + try: + backend.dispose() + except Exception as exc: + if first_error is None: + first_error = exc + self._onnx_session = None + if first_error is not None and not suppress_errors: + raise first_error + + def model_name_aliasing(self, model_name: str) -> str: + """Finds the YAML config filename that matches the given model name. + + Matching is case-insensitive and first preserves separator boundaries + so names such as ``regnet_x_16gf`` do not collide with + ``regnet_x_1_6gf``. A separator-stripped fallback is used only when it + resolves to a single unique configuration, so inputs like ``resnet50``, + ``ResNet50``, ``Resnet_50``, and ``resnet-50`` all resolve to + ``ResNet50.yaml``. + + Args: + model_name: The model identifier provided by the caller. + + Returns: + The exact YAML filename (basename only) stored in + ``MODEL_CONFIG_DIR`` that corresponds to ``model_name``. + + Raises: + ValueError: If no YAML file matches, or if the name is ambiguous + (i.e., multiple files match after normalization). + """ + + return _model_name_aliasing(model_name) diff --git a/pyproject.toml b/pyproject.toml index c94d3ed..5514221 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,6 +2,85 @@ requires = ["setuptools>=70.0.0", "wheel"] build-backend = "setuptools.build_meta" +[project] +name = "mblt-vision-python" +dynamic = ["version"] +description = "Mobilint NPU vision models: classification, detection, segmentation, pose" +readme = "README.md" +dependencies = [ + "mblt-npu-python>=0.0.0", + "numpy>=1.26.0", + "torch>=2.4.1", + "opencv-python>=4.11.0.86", + "pillow>=11.1.0", + "faster-coco-eval", + "gdown>=5.2.0", + "huggingface-hub", + "matplotlib", + "requests>=2.32.0", + "scipy", + "tqdm", + "PyYAML", +] +requires-python = ">=3.10,<3.13" +license = { text = "BSD-3-Clause" } +authors = [{ name = "Mobilint" }] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: BSD License", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] + [project.urls] Home = "https://www.mobilint.com/" -Repository = "https://github.com/mobilint/mblt-vision-python" \ No newline at end of file +Repository = "https://github.com/mobilint/mblt-vision-python" + +[project.scripts] +mblt-vision = "mblt_vision.cli:main" + +[project.optional-dependencies] +onnxruntime = [ + "onnx", + "onnxruntime", +] +onnxruntime-gpu = [ + "onnx; sys_platform != 'darwin'", + "onnxruntime-gpu; sys_platform != 'darwin'", +] +qbcompiler = [ + "qbcompiler>=1.2.0", + "onnxruntime", +] + +[tool.setuptools] +packages = { find = { where = ["."], include = ["mblt_vision", "mblt_vision.*"] } } + +[tool.setuptools.dynamic] +version = { attr = "mblt_vision.__version__" } + +[tool.setuptools.package-data] +"mblt_vision" = [ + "py.typed", + "models/*.yaml", + "datasets/*.yaml", + "datasets/*.txt", +] + +[tool.pytest.ini_options] +markers = [ + "requires_network: test downloads model artifacts from Hugging Face Hub.", + "requires_npu: test requires a configured Mobilint NPU runtime and device.", +] + +[dependency-groups] +dev = [ + "pre-commit", + "pytest", + "pytest-timeout", +] diff --git a/tests/TEST.md b/tests/TEST.md new file mode 100644 index 0000000..ffd156d --- /dev/null +++ b/tests/TEST.md @@ -0,0 +1,42 @@ +# Test `vision` + +You can validate Mobilint's Vision API with [`pytest`](https://docs.pytest.org/en/stable/). The snippets below assume your virtual environment is already activated. + +## Install Packages + +Install the runtime extras plus the developer tooling required by the test suite: + +```bash +pip install -e . +pip install pytest +``` + +## Run All Tests + +Execute the complete standalone Vision test matrix: + +```bash +pytest tests +``` + +## Run Offline Unit Tests + +Exclude Hugging Face downloads and NPU hardware: + +```bash +pytest tests -m "not requires_network and not requires_npu" +``` + +## Run Optional Integration Tests + +After authenticating with Hugging Face Hub, exercise representative ONNX models: + +```bash +pytest tests/test_onnx_classification.py -m requires_network +``` + +To run MXQ inference, add a configured NPU and the shared runtime options: + +```bash +pytest tests/test_mxq_inference.py -m requires_npu --mxq-path /path/to/model.mxq +``` diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..671929d --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,19 @@ +# The shared options and fixtures come from mblt_npu. Installed, they arrive +# through its pytest11 entry point and this file is unnecessary; from a source +# checkout the star import is what puts pytest_addoption in the root conftest +# namespace, which is the only place pytest looks for it. +from pathlib import Path + +import pytest +from PIL import Image + +from mblt_npu.pytest_plugin import * # noqa: F401,F403 + + +@pytest.fixture +def synthetic_image_path(tmp_path: Path) -> Path: + """Create a small RGB image for inference tests without committing assets.""" + + image_path = tmp_path / "synthetic-image.jpg" + Image.new("RGB", (64, 48), color=(64, 128, 192)).save(image_path) + return image_path diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..40b4182 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,45 @@ +"""Tests for the standalone public Vision API.""" + +from __future__ import annotations + +import pytest + +import mblt_vision +from mblt_vision import list_models, list_tasks +from mblt_vision._tasks import normalize_vision_task + + +def test_public_discovery_exposes_all_supported_tasks() -> None: + """Expose every standalone Vision task through the public API.""" + + assert list_tasks() == [ + "image_classification", + "depth_estimation", + "object_detection", + "instance_segmentation", + "semantic_segmentation", + "obb", + "pose_estimation", + "face_detection", + ] + assert list_models("obb")["obb"] + + +def test_model_exports_are_discoverable_from_task_and_top_level_namespaces() -> None: + """Keep task exports synchronized with lazy top-level compatibility exports.""" + + from mblt_vision.image_classification import ResNet50 + from mblt_vision.object_detection import YOLO11m + + assert mblt_vision.ResNet50 is ResNet50 + assert mblt_vision.YOLO11m is YOLO11m + assert "ResNet50" in list_models("image_classification")["image_classification"] + assert "YOLO11m" in list_models("object_detection")["object_detection"] + + +@pytest.mark.parametrize("task", [None, 1, object()]) +def test_task_normalization_rejects_non_strings(task: object) -> None: + """Reject invalid task inputs with a stable, actionable error.""" + + with pytest.raises(TypeError, match="must be a string"): + normalize_vision_task(task) # type: ignore[arg-type] diff --git a/tests/test_benchmark_vision.py b/tests/test_benchmark_vision.py new file mode 100644 index 0000000..74369a1 --- /dev/null +++ b/tests/test_benchmark_vision.py @@ -0,0 +1,505 @@ +"""Tests for the standardized multi-model vision benchmark tools.""" + +from __future__ import annotations + +import argparse +import csv +import runpy +from pathlib import Path + +import pytest + +from benchmark import benchmark_vision_models, compare_benchmark_results +from mblt_vision.datasets import get_dataset_config +from mblt_vision.utils.evaluation import ( + DOTAResult, + ImageNetResult, + NYUDepthResult, + SemanticSegmentationResult, +) + + +def test_direct_coco_organizer_uses_registry_download_defaults() -> None: + """Keep documented COCO organizer defaults aligned with accepted registry URLs.""" + + script = Path(__file__).parents[1] / "benchmark" / "organize_coco.py" + namespace = runpy.run_path(str(script)) + download_config = get_dataset_config("coco")["download"] + + assert namespace["DEFAULT_COCO_IMAGE_SOURCE"] == download_config["images"] + assert namespace["DEFAULT_COCO_ANNOTATION_SOURCE"] == download_config["annotations"] + assert namespace["DEFAULT_COCO_IMAGE_SOURCE"].startswith("https://") + assert namespace["DEFAULT_COCO_ANNOTATION_SOURCE"].startswith("https://") + + +def test_benchmark_records_imagenet_metrics_in_primary_order( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Use Top-1 as the score while retaining Top-5 in benchmark metrics.""" + + class FakeModel: + """Minimal classification model double.""" + + post_cfg = {"task": "image_classification"} + + import mblt_vision.utils.evaluation as evaluation_module + + monkeypatch.setattr( + evaluation_module, + "eval_imagenet_metrics", + lambda *args, **kwargs: ImageNetResult(top1=0.75, top5=0.95), + ) + args = argparse.Namespace( + task="image_classification", + data_path=str(tmp_path), + batch_size=1, + ) + + score, score_name, metrics = benchmark_vision_models._evaluate( + FakeModel(), args, tmp_path + ) + + assert score == 0.75 + assert score_name == "top1_accuracy" + assert metrics == {"top1_accuracy": 0.75, "top5_accuracy": 0.95} + + +def test_benchmark_records_depth_metrics( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Record all NYU depth metrics with delta1 as primary.""" + + class FakeModel: + post_cfg = {"task": "depth_estimation", "dataset": "nyu-depth"} + + import mblt_vision.utils.evaluation as evaluation_module + + monkeypatch.setattr( + evaluation_module, + "eval_nyu_depth", + lambda *args, **kwargs: NYUDepthResult(delta1=0.8, abs_rel=0.1, rmse=0.2), + ) + args = argparse.Namespace( + task="depth_estimation", data_path=str(tmp_path), batch_size=1 + ) + + score, score_name, metrics = benchmark_vision_models._evaluate( + FakeModel(), args, tmp_path + ) + + assert (score, score_name) == (0.8, "delta1") + assert metrics == {"delta1": 0.8, "abs_rel": 0.1, "rmse": 0.2} + + +@pytest.mark.parametrize( + ("dataset", "evaluator_name"), + [("ade20k", "eval_ade20k"), ("cityscapes", "eval_cityscapes")], +) +def test_benchmark_dispatches_semantic_taxonomy( + dataset: str, + evaluator_name: str, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Dispatch semantic evaluation explicitly from the configured taxonomy.""" + + class FakeModel: + post_cfg = {"task": "semantic_segmentation", "dataset": dataset} + + import mblt_vision.utils.evaluation as evaluation_module + + monkeypatch.setattr( + evaluation_module, + evaluator_name, + lambda *args, **kwargs: SemanticSegmentationResult( + miou=0.6, pixel_accuracy=0.9 + ), + ) + args = argparse.Namespace( + task="semantic_segmentation", data_path=str(tmp_path), batch_size=1 + ) + + score, score_name, metrics = benchmark_vision_models._evaluate( + FakeModel(), args, tmp_path + ) + + assert (score, score_name) == (0.6, "miou") + assert metrics == {"miou": 0.6, "pixel_accuracy": 0.9} + + +def test_benchmark_accepts_obb_model_task( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Dispatch an OBB model through DOTAv1 evaluation.""" + + class FakeModel: + post_cfg = {"task": "obb", "dataset": "dotav1"} + + import mblt_vision.utils.evaluation as evaluation_module + + monkeypatch.setattr( + evaluation_module, + "eval_dota", + lambda *args, **kwargs: DOTAResult(map50=0.7, map5095=0.5), + ) + args = argparse.Namespace( + task="obb", + data_path=str(tmp_path), + batch_size=1, + conf_thres=None, + iou_thres=None, + ) + + score, score_name, metrics = benchmark_vision_models._evaluate( + FakeModel(), args, tmp_path + ) + + assert (score, score_name) == (0.5, "map50_95") + assert metrics == {"map50_95": 0.5, "map50": 0.7} + + +def test_benchmark_continues_after_evaluator_type_error( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Record an unsupported evaluator output without aborting later targets.""" + + class FakeEngine: + """Minimal engine used to exercise per-target error handling.""" + + def __init__(self, *, model_cls: str, **kwargs: object) -> None: + self.model_cls = model_cls + + def dispose(self) -> None: + """Release the fake benchmark engine.""" + + def fake_evaluate( + model: FakeEngine, + args: object, + run_dir: Path, + ) -> tuple[float, str, dict[str, float]]: + if model.model_cls == "invalid-output": + raise TypeError("Unsupported model output") + return 0.9, "top1_accuracy", {"top1_accuracy": 0.9} + + import mblt_vision as vision + + monkeypatch.setattr(vision, "MBLT_Engine", FakeEngine) + monkeypatch.setattr(benchmark_vision_models, "_evaluate", fake_evaluate) + + result = benchmark_vision_models.main( + [ + "--models", + "invalid-output", + "valid-output", + "--task", + "image_classification", + "--data-path", + str(tmp_path / "dataset"), + "--results-dir", + str(tmp_path / "results"), + "--no-plot", + ] + ) + + with (tmp_path / "results" / "results.csv").open( + newline="", encoding="utf-8" + ) as results_file: + rows = list(csv.DictReader(results_file)) + + assert result == 1 + assert [row["status"] for row in rows] == ["error", "ok"] + assert rows[0]["error"] == "TypeError: Unsupported model output" + + +def test_benchmark_records_dispose_failures_without_losing_other_results( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Write artifacts and continue when a backend fails during cleanup.""" + + class FakeEngine: + """Minimal engine that can fail evaluation or disposal by model name.""" + + def __init__(self, *, model_cls: str, **kwargs: object) -> None: + self.model_cls = model_cls + + def dispose(self) -> None: + if self.model_cls in {"evaluation-failure", "cleanup-failure"}: + raise RuntimeError("Backend cleanup failed") + + def fake_evaluate( + model: FakeEngine, + args: object, + run_dir: Path, + ) -> tuple[float, str, dict[str, float]]: + if model.model_cls == "evaluation-failure": + raise ValueError("Evaluation failed") + return 0.9, "top1_accuracy", {"top1_accuracy": 0.9} + + import mblt_vision as vision + + monkeypatch.setattr(vision, "MBLT_Engine", FakeEngine) + monkeypatch.setattr(benchmark_vision_models, "_evaluate", fake_evaluate) + + result = benchmark_vision_models.main( + [ + "--models", + "evaluation-failure", + "cleanup-failure", + "valid-output", + "--task", + "image_classification", + "--data-path", + str(tmp_path / "dataset"), + "--results-dir", + str(tmp_path / "results"), + "--no-plot", + ] + ) + + with (tmp_path / "results" / "results.csv").open( + newline="", encoding="utf-8" + ) as results_file: + rows = list(csv.DictReader(results_file)) + + assert result == 1 + assert [row["status"] for row in rows] == ["error", "error", "ok"] + assert rows[0]["error"] == "ValueError: Evaluation failed" + assert rows[0]["cleanup_error"] == "RuntimeError: Backend cleanup failed" + assert rows[1]["error"] == "RuntimeError: Backend cleanup failed" + + +def test_benchmark_forwards_normalized_target_device( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Construct each benchmark engine with the requested board identifier.""" + + engine_kwargs: list[dict[str, object]] = [] + + class FakeEngine: + """Minimal benchmark engine that records its construction options.""" + + def __init__(self, **kwargs: object) -> None: + engine_kwargs.append(kwargs) + + def dispose(self) -> None: + """Release the fake benchmark engine.""" + + import mblt_vision as vision + + monkeypatch.setattr(vision, "MBLT_Engine", FakeEngine) + monkeypatch.setattr( + benchmark_vision_models, + "_evaluate", + lambda *args: (0.9, "top1_accuracy", {"top1_accuracy": 0.9}), + ) + + assert ( + benchmark_vision_models.main( + [ + "--models", + "model-a", + "--task", + "image_classification", + "--target-device", + "regulus", + "--data-path", + str(tmp_path / "dataset"), + "--results-dir", + str(tmp_path / "results"), + "--no-plot", + ] + ) + == 0 + ) + assert engine_kwargs == [ + { + "model_cls": "model-a", + "model_type": "DEFAULT", + "model_path": "", + "mxq_path": "", + "onnx_path": "", + "framework": None, + "dev_no": 0, + "target_device": "regulus-ra", + "core_mode": "single", + } + ] + + +@pytest.mark.parametrize( + ("target_device", "core_mode", "expected_modes"), + [ + ("aries-rb", None, ("global8",)), + ("aries-rb", "all", ("single", "multi", "global4", "global8")), + ("regulus-ra", None, ("single",)), + ("regulus-rb", "all", ("single",)), + ], +) +def test_benchmark_core_modes_follow_the_selected_target_device( + target_device: str, core_mode: str | None, expected_modes: tuple[str, ...] +) -> None: + """Use only execution modes supported by the selected NPU board.""" + + assert ( + benchmark_vision_models._core_modes(core_mode, target_device=target_device) + == expected_modes + ) + + +def test_benchmark_rejects_explicit_unsupported_regulus_core_mode( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Fail before creating an invalid Regulus benchmark engine.""" + + monkeypatch.setattr( + benchmark_vision_models, + "_run_target", + lambda *args: pytest.fail("unsupported mode must not start a benchmark"), + ) + + assert ( + benchmark_vision_models.main( + [ + "--models", + "model-a", + "--task", + "image_classification", + "--target-device", + "regulus-ra", + "--core-mode", + "global8", + "--data-path", + str(tmp_path / "dataset"), + "--results-dir", + str(tmp_path / "results"), + "--no-plot", + ] + ) + == 2 + ) + + +def test_comparison_rejects_matching_metrics_from_different_tasks( + tmp_path: Path, +) -> None: + """Reject task-incompatible inputs even when their score metric matches.""" + + for name, task in ( + ("detection", "object_detection"), + ("segmentation", "instance_segmentation"), + ): + results_path = tmp_path / name / "results.csv" + results_path.parent.mkdir() + results_path.write_text( + f"model,core_mode,task,status,score_name,score\nmodel-a,global8,{task},ok,map50_95,0.5\n", + encoding="utf-8", + ) + + with pytest.raises(SystemExit, match="incompatible benchmark tasks"): + compare_benchmark_results.main( + [str(tmp_path / "detection"), str(tmp_path / "segmentation")] + ) + + +@pytest.mark.parametrize( + "framework_args", + [ + ["--framework", "onnx"], + ["--model-path", "model.onnx"], + ["--model-path", "MODEL.ONNX"], + ["--onnx-path", "model.onnx"], + ], +) +def test_onnx_benchmark_uses_one_neutral_runtime_target( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + framework_args: list[str], +) -> None: + """Avoid recording repeated NPU core-mode runs for ONNX inference.""" + + captured_modes: list[str] = [] + + def fake_run_target( + model_name: str, core_mode: str, args: object, results_dir: Path + ) -> dict[str, object]: + captured_modes.append(core_mode) + return { + "model": model_name, + "core_mode": core_mode, + "task": "image_classification", + "batch_size": 1, + "status": "ok", + "score": 0.9, + "score_name": "top1", + "elapsed_s": 0.0, + } + + monkeypatch.setattr(benchmark_vision_models, "_run_target", fake_run_target) + + assert ( + benchmark_vision_models.main( + [ + "--models", + "model-a", + "--task", + "image_classification", + *framework_args, + "--core-mode", + "all", + "--data-path", + str(tmp_path / "dataset"), + "--results-dir", + str(tmp_path / "results"), + "--no-plot", + ] + ) + == 0 + ) + assert captured_modes == ["onnx"] + + +def test_comparison_uses_result_directory_names( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Derive default chart paths and legend labels from result directories.""" + + for name in ("baseline", "candidate"): + results_path = tmp_path / name / "results.csv" + results_path.parent.mkdir() + results_path.write_text( + "model,core_mode,task,status,score_name,score\nmodel-a,global8,object_detection,ok,map50_95,0.5\n", + encoding="utf-8", + ) + import mblt_vision.benchmark.chart_utils as chart_utils + + captured: dict[str, object] = {} + + def fake_default_charts_dir( + script_dir: Path, sources: list[Path], **kwargs: object + ) -> Path: + captured["sources"] = sources + return tmp_path / "charts" + + monkeypatch.setattr(chart_utils, "default_charts_dir", fake_default_charts_dir) + monkeypatch.setattr( + chart_utils, + "plot_grouped_scalar_barh", + lambda **kwargs: captured.update(kwargs), + ) + + assert ( + compare_benchmark_results.main( + [str(tmp_path / "baseline"), str(tmp_path / "candidate")] + ) + == 0 + ) + sources = captured["sources"] + assert isinstance(sources, list) + source_paths: list[Path] = [] + for source in sources: + assert isinstance(source, Path) + source_paths.append(source) + assert [path.name for path in source_paths] == ["baseline", "candidate"] + assert captured["group_labels"] == ["baseline", "candidate"] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..62250fe --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,107 @@ +"""Tests for the standalone Vision CLI parser.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from mblt_vision.cli import build_parser +from mblt_vision.cli.compile import _run_compile +from mblt_vision.cli.predict import _cmd_predict +from mblt_vision.cli import val as val_module +from mblt_vision.cli.val import _cmd_val + + +def test_standalone_cli_registers_vision_commands() -> None: + """Expose the supported Vision commands from the standalone distribution.""" + + parser = build_parser() + + predict_args = parser.parse_args( + ["predict", "--source", "image.jpg", "--model", "resnet50"] + ) + val_args = parser.parse_args(["val", "--model", "resnet50"]) + compile_args = parser.parse_args( + ["compile", "--model-cls", "resnet50", "--target-device", "aries-rb"] + ) + + assert predict_args._handler is _cmd_predict + assert val_args._handler is _cmd_val + assert compile_args._handler is _run_compile + assert compile_args.target_device == "aries-rb" + assert predict_args.target_device == "aries-rb" + assert val_args.target_device == "aries-rb" + + +def test_vision_cli_uses_single_core_mode_by_default_on_regulus() -> None: + """Keep the CLI's implicit core mode compatible with a Regulus target.""" + + parser = build_parser() + args = parser.parse_args( + [ + "predict", + "--source", + "image.jpg", + "--model", + "resnet50", + "--target-device", + "regulus-ra", + ] + ) + + assert args.target_device == "regulus-ra" + assert args.core_mode is None + + +@pytest.mark.parametrize("batch_size", ["0", "-1"]) +def test_val_rejects_nonpositive_batch_sizes(batch_size: str) -> None: + """Fail argument parsing before validation can construct a model.""" + + with pytest.raises(SystemExit): + build_parser().parse_args( + ["val", "--model", "resnet50", "--batch-size", batch_size] + ) + + +@pytest.mark.parametrize("removed_alias", ["classify", "detect", "pose", "segment"]) +def test_predict_command_has_no_task_aliases(removed_alias: str) -> None: + """Keep prediction discoverable through one task-agnostic command.""" + + with pytest.raises(SystemExit): + build_parser().parse_args([removed_alias]) + + +def test_predict_help_explains_supported_workflows( + capsys: pytest.CaptureFixture[str], +) -> None: + """Describe tasks, output behavior, framework choice, and examples in CLI help.""" + + with pytest.raises(SystemExit, match="0"): + build_parser().parse_args(["predict", "--help"]) + help_text = capsys.readouterr().out + + assert "Supported tasks:" in help_text + assert "--framework onnx" in help_text + assert "--target-device regulus-ra" in help_text + + +def test_validation_default_dataset_path_uses_resolved_cache_root( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Keep default organization alongside artifacts when the home cache falls back.""" + + configured_path = Path.home() / ".mblt_model_zoo" / "datasets" / "coco" + fallback_cache = tmp_path / "cache" + monkeypatch.setattr( + val_module, + "get_dataset_config_for_task", + lambda _task, _dataset: {"path": str(configured_path)}, + ) + monkeypatch.setattr( + val_module, "get_mobilint_cache_dir", lambda: str(fallback_cache) + ) + + assert val_module._default_data_path_for_task("object_detection") == str( + fallback_cache / "datasets" / "coco" + ) diff --git a/tests/test_compile_vision.py b/tests/test_compile_vision.py new file mode 100644 index 0000000..467f55b --- /dev/null +++ b/tests/test_compile_vision.py @@ -0,0 +1,1055 @@ +"""Tests for packaged vision compilation and calibration preparation.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from PIL import Image + +from mblt_vision.compile import vision as compile_module +from mblt_vision.compile.vision import ( + copy_calibration_subset, + prepare_calibration_arrays, + resolve_quantization_values, + select_calibration_images, + validate_calibration_dataset, +) +from mblt_vision.utils import datasets as dataset_utils +from mblt_vision.utils.datasets import readiness as readiness_module +from mblt_vision.wrapper import resolve_model_config + + +def _write_images(directory: Path, names: list[str]) -> list[Path]: + """Create placeholder image files for selection tests. + + Args: + directory: Destination directory. + names: Relative filenames to create. + + Returns: + Created image paths. + """ + + created: list[Path] = [] + for name in names: + path = directory / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(name.encode()) + created.append(path) + return created + + +def _write_image(path: Path, size: tuple[int, int] = (2, 2)) -> None: + """Write a minimal valid RGB image for dataset-readiness tests.""" + + path.parent.mkdir(parents=True, exist_ok=True) + Image.new("RGB", size).save(path) + + +def test_resolve_model_config_derives_onnx_filename() -> None: + """Derive ONNX artifact names consistently from MXQ artifact names.""" + + hub_config = resolve_model_config("alexnet") + direct_config = resolve_model_config( + { + "file_cfg": { + "filename": "example.mxq", + }, + "pre_cfg": {}, + "post_cfg": {}, + } + ) + custom_config = resolve_model_config( + { + "file_cfg": { + "filename": "example.mxq", + "onnx_filename": "exported-model.onnx", + }, + "pre_cfg": {}, + "post_cfg": {}, + } + ) + + assert hub_config["file_cfg"]["onnx_filename"] == "alexnet_IMAGENET1K_V1.onnx" + assert direct_config["file_cfg"]["onnx_filename"] == "example.onnx" + assert custom_config["file_cfg"]["onnx_filename"] == "exported-model.onnx" + + +@pytest.mark.parametrize( + ("task", "relative_dir"), + [ + ("depth_estimation", "images"), + ("object_detection", "val2017"), + ("semantic_segmentation", "images"), + ("obb", "images"), + ], +) +def test_non_imagenet_selection_is_deterministic( + task: str, relative_dir: str, tmp_path: Path +) -> None: + """Select deterministic total-count COCO and DOTAv1 samples.""" + + _write_images(tmp_path / relative_dir, [f"image-{index}.jpg" for index in range(5)]) + + first = select_calibration_images(task, tmp_path, subset_size=3, seed=7) + second = select_calibration_images(task, tmp_path, subset_size=3, seed=7) + + assert first == second + assert len(first) == 3 + + +def test_obb_selects_calibration_images(tmp_path: Path) -> None: + """Select OBB calibration images through exported sampling helpers.""" + + images = _write_images(tmp_path / "images", ["P0001.png", "P0002.png"]) + + assert set(select_calibration_images("obb", tmp_path, subset_size=2)) == set(images) + + +def test_calibration_default_dataset_path_uses_resolved_cache_root( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Keep implicit compilation datasets alongside the active artifact cache.""" + + configured_path = Path.home() / ".mblt_model_zoo" / "datasets" / "coco" + fallback_cache = tmp_path / "cache" + monkeypatch.setattr( + compile_module, + "get_dataset_config_for_task", + lambda *_: {"name": "coco", "path": str(configured_path)}, + ) + monkeypatch.setattr( + compile_module, "get_mobilint_cache_dir", lambda: str(fallback_cache) + ) + monkeypatch.setattr(compile_module, "_dataset_ready", lambda *_: True) + + assert compile_module.ensure_calibration_dataset("object_detection") == ( + fallback_cache / "datasets" / "coco" + ) + + +@pytest.mark.parametrize( + ("task", "relative_dir"), + [ + ("depth_estimation", "images"), + ("object_detection", "val2017"), + ("instance_segmentation", "val2017"), + ("semantic_segmentation", "images"), + ("pose_estimation", "val2017"), + ("obb", "images"), + ], +) +def test_selection_defaults_to_seed_zero( + task: str, relative_dir: str, tmp_path: Path +) -> None: + """Use seed zero whenever callers omit the vision sampling seed.""" + + _write_images(tmp_path / relative_dir, [f"image-{index}.jpg" for index in range(5)]) + + assert select_calibration_images( + task, tmp_path, subset_size=3 + ) == select_calibration_images( + task, + tmp_path, + subset_size=3, + seed=0, + ) + + +def test_widerface_selection_uses_per_category_size(tmp_path: Path) -> None: + """Select the requested number of images from every WiderFace category.""" + + _write_images( + tmp_path / "images" / "0--Parade", [f"parade-{index}.jpg" for index in range(3)] + ) + _write_images( + tmp_path / "images" / "1--Handshaking", + [f"handshake-{index}.jpg" for index in range(3)], + ) + + first = select_calibration_images("face_detection", tmp_path, subset_size=2, seed=7) + second = select_calibration_images( + "face_detection", tmp_path, subset_size=2, seed=7 + ) + default = select_calibration_images("face_detection", tmp_path) + seed_zero = select_calibration_images("face_detection", tmp_path, seed=0) + + assert first == second + assert len(first) == 4 + assert len(default) == 2 + assert default == seed_zero + assert {path.parent.name for path in first} == {"0--Parade", "1--Handshaking"} + + +def test_imagenet_selection_uses_per_class_size(tmp_path: Path) -> None: + """Select the requested number of images independently from each class.""" + + _write_images(tmp_path / "class-a", ["same.jpg", "a.jpg"]) + _write_images(tmp_path / "class-b", ["same.jpg", "b.jpg"]) + + selected = select_calibration_images( + "image_classification", tmp_path, subset_size=1 + ) + seed_zero = select_calibration_images( + "image_classification", tmp_path, subset_size=1, seed=0 + ) + + assert len(selected) == 2 + assert selected == seed_zero + assert {path.parent.name for path in selected} == {"class-a", "class-b"} + + +@pytest.mark.parametrize("subset_size", [0, -1, 4]) +def test_selection_rejects_invalid_sizes(subset_size: int, tmp_path: Path) -> None: + """Reject non-positive and unavailable calibration subset sizes.""" + + _write_images(tmp_path / "val2017", ["one.jpg", "two.jpg", "three.jpg"]) + + with pytest.raises(ValueError, match="subset_size"): + select_calibration_images("object_detection", tmp_path, subset_size=subset_size) + + +def test_selection_rejects_symlinked_images(tmp_path: Path) -> None: + """Reject image-looking symlinks instead of sampling their targets.""" + + external_image = _write_images(tmp_path / "external", ["secret.jpg"])[0] + image_dir = tmp_path / "dataset" / "val2017" + image_dir.mkdir(parents=True) + (image_dir / "stolen.jpg").symlink_to(external_image) + + with pytest.raises(ValueError, match="must not be a symlink"): + select_calibration_images( + "object_detection", tmp_path / "dataset", subset_size=1 + ) + + +def test_copy_rejects_symlinked_or_outside_images(tmp_path: Path) -> None: + """Defend direct copy callers against symlink and outside-root sources.""" + + dataset = tmp_path / "dataset" + external_image = _write_images(tmp_path / "external", ["secret.jpg"])[0] + symlink = dataset / "stolen.jpg" + dataset.mkdir() + symlink.symlink_to(external_image) + + with pytest.raises(ValueError, match="must not be a symlink"): + copy_calibration_subset([symlink], dataset, tmp_path / "subset") + with pytest.raises(ValueError, match="must remain within dataset root"): + copy_calibration_subset([external_image], dataset, tmp_path / "subset") + + +def test_flat_subset_names_are_collision_safe(tmp_path: Path) -> None: + """Preserve images with duplicate basenames from nested source directories.""" + + images = _write_images( + tmp_path / "dataset", ["images/a/same.jpg", "images/b/same.jpg"] + ) + + copied = copy_calibration_subset(images, tmp_path / "dataset", tmp_path / "subset") + + assert len(copied) == 2 + assert copied[0].name != copied[1].name + assert all(path.is_file() for path in copied) + + +@pytest.mark.parametrize("output_name", [".", "parent-output", "val2017/subset"]) +def test_subset_rejects_overlapping_output_paths( + tmp_path: Path, output_name: str +) -> None: + """Protect the organized source dataset from destructive replacement.""" + + data_path = tmp_path / "dataset" + image_path = _write_images(data_path / "val2017", ["one.jpg"])[0] + output_path = tmp_path if output_name == "." else data_path / output_name + + with pytest.raises(ValueError, match="must not overlap"): + compile_module.make_calibration_subset( + "object_detection", data_path, output_path, subset_size=1 + ) + + assert image_path.is_file() + + +def test_subset_keeps_existing_output_when_selection_fails(tmp_path: Path) -> None: + """Leave an existing subset untouched when a replacement cannot be selected.""" + + data_path = tmp_path / "dataset" + output_path = tmp_path / "subset" + _write_images(data_path / "val2017", ["one.jpg"]) + existing = _write_images(output_path, ["existing.jpg"])[0] + + with pytest.raises(ValueError, match="subset_size"): + compile_module.make_calibration_subset( + "object_detection", data_path, output_path, subset_size=2 + ) + + assert existing.read_bytes() == b"existing.jpg" + + +@pytest.mark.parametrize("symlinked_parent", [False, True]) +def test_subset_rejects_symlinked_output_path_before_replacement( + tmp_path: Path, symlinked_parent: bool +) -> None: + """Do not resolve a caller-controlled subset output through a symlink.""" + + data_path = tmp_path / "dataset" + _write_images(data_path / "val2017", ["one.jpg"]) + external_output = tmp_path / "external-parent" / "subset" + existing = _write_images(external_output, ["keep.jpg"])[0] + if symlinked_parent: + output_path = tmp_path / "linked-parent" / "subset" + (tmp_path / "linked-parent").symlink_to( + external_output.parent, target_is_directory=True + ) + else: + output_path = tmp_path / "subset" + output_path.symlink_to(external_output, target_is_directory=True) + + with pytest.raises(ValueError, match="output_dir must not be or contain a symlink"): + compile_module.make_calibration_subset( + "object_detection", data_path, output_path, subset_size=1 + ) + + assert existing.read_bytes() == b"keep.jpg" + + +def test_imagenet_readiness_rejects_partial_class_tree(tmp_path: Path) -> None: + """Organize ImageNet again when the existing class layout is incomplete.""" + + _write_images(tmp_path / "class-a", ["image.jpg"]) + + assert not compile_module._dataset_ready("image_classification", tmp_path) + + +def test_depth_dataset_readiness_requires_complete_pairs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Require all paired NYU samples before sampling calibration images.""" + + monkeypatch.setattr(readiness_module, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + image_path = tmp_path / "images" / "image.png" + _write_image(image_path) + (tmp_path / "depth").mkdir() + assert not compile_module._dataset_ready("depth_estimation", tmp_path, "nyu-depth") + + np.save(tmp_path / "depth" / "image.npy", np.ones((2, 2), dtype=np.float32)) + assert compile_module._dataset_ready("depth_estimation", tmp_path, "nyu-depth") + + +@pytest.mark.parametrize( + ("dataset", "image_name", "other_dataset"), + [ + ("ade20k", "ADE_val_00000001.jpg", "cityscapes"), + ("cityscapes", "frankfurt_000000_000001.png", "ade20k"), + ], +) +def test_semantic_dataset_readiness_checks_taxonomy( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + dataset: str, + image_name: str, + other_dataset: str, +) -> None: + """Require complete pairs with filenames belonging to the selected taxonomy.""" + + monkeypatch.setattr(readiness_module, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(readiness_module, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 1) + image_path = tmp_path / "images" / image_name + _write_image(image_path) + (tmp_path / "annotations").mkdir() + assert not compile_module._dataset_ready("semantic_segmentation", tmp_path, dataset) + + annotation_value = 1 if dataset == "ade20k" else 7 + Image.new("L", (2, 2), color=annotation_value).save( + tmp_path / "annotations" / f"{image_path.stem}.png" + ) + if dataset == "ade20k": + (tmp_path / "objectInfo150.txt").write_bytes(b"labels") + (tmp_path / "sceneCategories.txt").write_bytes(b"scenes") + assert compile_module._dataset_ready("semantic_segmentation", tmp_path, dataset) + assert not compile_module._dataset_ready( + "semantic_segmentation", tmp_path, other_dataset + ) + + +def test_obb_dataset_readiness_accepts_flat_images( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Accept the flat image layout produced by the DOTAv1 organizer.""" + + data_path = tmp_path / "dotav1" + monkeypatch.setattr( + compile_module, + "get_dataset_config_for_task", + lambda task, dataset=None: { + "name": "dotav1", + "path": str(data_path), + "download": {"url": "https://example.test/dotav1.zip"}, + }, + ) + + def _organize_dotav1(**kwargs: str) -> None: + assert kwargs["output_dir"] == str(data_path) + _write_images(data_path / "images", ["P0001.png"]) + (data_path / "labels" / "val_original").mkdir(parents=True) + (data_path / "labels" / "val_original" / "P0001.txt").write_bytes(b"label") + + monkeypatch.setattr(dataset_utils, "organize_dotav1", _organize_dotav1) + monkeypatch.setattr(readiness_module, "DOTAV1_VALIDATION_SAMPLE_COUNT", 1) + + assert ( + compile_module.ensure_calibration_dataset("obb", dataset="dotav1") == data_path + ) + + +def test_dense_dataset_organizers_follow_model_taxonomy( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Route dense compilation through the matching registry-backed organizer.""" + + data_path = tmp_path / "dataset" + cityscapes_images = tmp_path / "leftImg8bit_trainvaltest.zip" + cityscapes_annotations = tmp_path / "gtFine_trainvaltest.zip" + cityscapes_images.write_bytes(b"images") + cityscapes_annotations.write_bytes(b"annotations") + configs = { + "nyu-depth": { + "name": "nyu-depth", + "download": {"url": "https://example.test/nyu.zip"}, + }, + "ade20k": { + "name": "ade20k", + "download": {"url": "https://example.test/ade.zip"}, + }, + "cityscapes": { + "name": "cityscapes", + "download": { + "images_archive": cityscapes_images.name, + "annotations_archive": cityscapes_annotations.name, + }, + }, + } + calls: list[tuple[str, dict[str, str]]] = [] + monkeypatch.setattr( + compile_module, + "get_dataset_config_for_task", + lambda task, dataset=None: configs[str(dataset)], + ) + monkeypatch.setattr( + dataset_utils, + "organize_nyu_depth", + lambda **kwargs: calls.append(("nyu-depth", kwargs)), + ) + monkeypatch.setattr( + dataset_utils, + "organize_ade20k", + lambda **kwargs: calls.append(("ade20k", kwargs)), + ) + monkeypatch.setattr( + dataset_utils, + "organize_cityscapes", + lambda **kwargs: calls.append(("cityscapes", kwargs)), + ) + + compile_module._organize_dataset("depth_estimation", data_path, "nyu-depth") + compile_module._organize_dataset("semantic_segmentation", data_path, "ade20k") + compile_module._organize_dataset("semantic_segmentation", data_path, "cityscapes") + + assert calls == [ + ( + "nyu-depth", + { + "dataset_path": "https://example.test/nyu.zip", + "output_dir": str(data_path), + }, + ), + ( + "ade20k", + { + "dataset_path": "https://example.test/ade.zip", + "output_dir": str(data_path), + }, + ), + ( + "cityscapes", + { + "image_dir": str(cityscapes_images), + "annotation_dir": str(cityscapes_annotations), + "output_dir": str(data_path), + }, + ), + ] + + +def test_prepare_calibration_arrays_preserves_preprocess_output(tmp_path: Path) -> None: + """Save contiguous HWC float32 arrays matching engine preprocessing.""" + + source = _write_images(tmp_path, ["image.jpg"])[0] + expected = np.arange(18, dtype=np.float64).reshape(2, 3, 3) + + class _Engine: + def preprocess(self, image_path: str) -> np.ndarray: + assert image_path == str(source) + return expected + + saved = prepare_calibration_arrays(_Engine(), [source], tmp_path / "arrays") # type: ignore[arg-type] + actual = np.load(saved[0]) + + assert actual.dtype == np.float32 + assert actual.flags.c_contiguous + np.testing.assert_array_equal(actual, expected.astype(np.float32)) + + +def test_prepare_calibration_arrays_rejects_chw(tmp_path: Path) -> None: + """Reject preprocessing output that is not three-channel HWC data.""" + + source = _write_images(tmp_path, ["image.jpg"])[0] + + class _Engine: + def preprocess(self, image_path: str) -> np.ndarray: + return np.zeros((3, 2, 2), dtype=np.float32) + + with pytest.raises(ValueError, match="HWC three-channel"): + prepare_calibration_arrays(_Engine(), [source], tmp_path / "arrays") # type: ignore[arg-type] + + +@pytest.mark.parametrize("invalid_value", [np.nan, np.inf, -np.inf]) +def test_validate_calibration_dataset_rejects_nonfinite_values( + tmp_path: Path, invalid_value: float +) -> None: + """Reject non-finite ready calibration tensors before compilation.""" + + array = np.ones((2, 2, 3), dtype=np.float32) + array[0, 0, 0] = invalid_value + np.save(tmp_path / "selected.npy", array) + + with pytest.raises(ValueError, match="must contain only finite values"): + validate_calibration_dataset(tmp_path) + + +def test_validate_calibration_dataset_rejects_unexpected_model_image_shape( + tmp_path: Path, +) -> None: + """Ready calibration arrays must retain the configured preprocessor geometry.""" + + np.save(tmp_path / "selected.npy", np.ones((2, 2, 3), dtype=np.float32)) + + with pytest.raises(ValueError, match=r"pre_cfg image shape \(3, 2\)"): + validate_calibration_dataset(tmp_path, image_shape=(3, 2)) + + +def test_quantization_explicit_values_do_not_fetch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Give explicit quantization values precedence over hosted metadata.""" + + monkeypatch.setattr( + compile_module, + "_fetch_quantization_config", + lambda *args: pytest.fail("hosted metadata should not be fetched"), + ) + + assert resolve_quantization_values({"repo_id": "owner/model"}, 0.98, 0.03) == ( + 0.98, + 0.03, + ) + + +@pytest.mark.parametrize("topk_key", ["topk", "topk_ratio"]) +def test_quantization_reads_both_hosted_topk_spellings( + topk_key: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Read hosted percentile conversion and both supported top-k field names.""" + + monkeypatch.setattr( + compile_module, + "_fetch_quantization_config", + lambda repo_id, revision: {"percentile": 0.002, topk_key: 0.04}, + ) + + percentile, topk = resolve_quantization_values( + {"repo_id": "owner/model", "revision": "v1"}, None, None + ) + + assert percentile == pytest.approx(0.998) + assert topk == 0.04 + + +def test_quantization_resolves_missing_values_independently( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Retain an explicit percentile while obtaining only missing top-k metadata.""" + + monkeypatch.setattr( + compile_module, + "_fetch_quantization_config", + lambda repo_id, revision: {"percentile": 0.25, "topk_ratio": 0.08}, + ) + + assert resolve_quantization_values({"repo_id": "owner/model"}, 0.97, None) == ( + 0.97, + 0.08, + ) + + +def test_quantization_warns_and_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + """Use documented defaults when optional hosted values are unavailable.""" + + monkeypatch.setattr( + compile_module, "_fetch_quantization_config", lambda repo_id, revision: None + ) + + with pytest.warns(UserWarning) as warning_records: + values = resolve_quantization_values({"repo_id": "owner/model"}, None, None) + + assert values == ( + compile_module.DEFAULT_PERCENTILE, + compile_module.DEFAULT_TOPK_RATIO, + ) + assert len(warning_records) == 2 + + +def test_quantization_rejects_malformed_hosted_json( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Raise a contextual error for malformed hosted JSON.""" + + metadata_path = tmp_path / "best_result.json" + metadata_path.write_text("{bad json", encoding="utf-8") + monkeypatch.setattr( + compile_module, "hf_hub_download", lambda **kwargs: str(metadata_path) + ) + + with pytest.raises(ValueError, match="Malformed quantization metadata JSON"): + compile_module._fetch_quantization_config("owner/model", "main") + + +def test_quantization_rejects_malformed_config_value( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Reject non-numeric hosted quantization fields.""" + + metadata_path = tmp_path / "best_result.json" + metadata_path.write_text( + json.dumps({"config": {"percentile": "invalid"}}), encoding="utf-8" + ) + monkeypatch.setattr( + compile_module, "hf_hub_download", lambda **kwargs: str(metadata_path) + ) + + with pytest.raises(ValueError, match="must be numeric"): + resolve_quantization_values({"repo_id": "owner/model"}, None, 0.01) + + +def _run_fake_compile( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + *, + task: str, + model_path: str | Path | None, + onnx_path: str | Path | None = None, + entry_level: str = "data", + dataset: str | None = None, + target_device: str = "aries-rb", + save_path: Path | None = None, + fail: bool = False, + calls: dict[str, Any] | None = None, +) -> tuple[dict[str, Any], Path]: + """Run compilation with fake engine and qbcompiler dependencies. + + Args: + monkeypatch: Pytest monkeypatch fixture. + tmp_path: Temporary test root. + task: Fake model task. + model_path: Optional user-supplied model path. + entry_level: Calibration pipeline level to supply. + dataset: Optional fake model dataset taxonomy. + fail: Whether the fake compiler should fail. + calls: Optional mapping populated even when compilation fails. + + Returns: + Captured calls and resolved hosted ONNX path. + """ + + calls = {} if calls is None else calls + hosted_onnx = tmp_path / "hosted-model.onnx" + hosted_onnx.write_bytes(b"onnx") + dataset_path = tmp_path / "dataset" + dataset = ( + dataset + or { + "image_classification": "imagenet", + "depth_estimation": "nyu-depth", + "object_detection": "coco", + "instance_segmentation": "coco", + "semantic_segmentation": "ade20k", + "pose_estimation": "coco", + "face_detection": "widerface", + "obb": "dotav1", + }[task] + ) + if task == "image_classification": + _write_images(dataset_path / "class-a", ["one.jpg"]) + elif task in {"depth_estimation", "semantic_segmentation"}: + _write_images(dataset_path / "images", ["one.jpg"]) + else: + _write_images(dataset_path / "val2017", ["one.jpg"]) + + class _Engine: + def __init__(self, **kwargs: Any) -> None: + calls["engine_kwargs"] = kwargs + resolved_path = Path(str(kwargs.get("model_path", hosted_onnx))) + self.file_cfg = {"onnx_path": str(resolved_path)} + + def preprocess(self, image_path: str) -> np.ndarray: + if entry_level == "calibration": + pytest.fail("ready calibration tensors must skip preprocessing") + calls.setdefault("preprocessed", []).append(image_path) + return np.arange(12, dtype=np.float64).reshape(2, 2, 3) + + def dispose(self) -> None: + calls["disposed"] = True + + class _CalibrationConfig: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + def _compile(**kwargs: Any) -> None: + calls["compile_kwargs"] = kwargs + array_dir = Path(kwargs["calib_data_path"]) + calls["temporary_root"] = array_dir.parent + calls["array"] = np.load(next(array_dir.glob("*.npy"))) + calls["calibration_kwargs"] = kwargs["calibration_config"].kwargs + if fail: + raise RuntimeError("compile failed") + + monkeypatch.setattr(compile_module, "MBLT_Engine", _Engine) + monkeypatch.setattr( + compile_module, "_load_qbcompiler", lambda: (_compile, _CalibrationConfig) + ) + monkeypatch.setattr( + compile_module, + "resolve_model_config", + lambda model_cls, model_type: { + "file_cfg": { + "repo_id": "owner/model", + "revision": "main", + "filename": "hosted-model.mxq", + "onnx_path": str(hosted_onnx), + }, + "pre_cfg": {"LetterBox": {"img_size": [2, 2]}}, + "post_cfg": {"task": task, "dataset": dataset}, + }, + ) + monkeypatch.setattr( + compile_module, "resolve_quantization_values", lambda *args: (0.99, 0.02) + ) + subset_path: Path | None = None + calib_data_path: Path | None = None + original_data_path: Path | None = dataset_path + if entry_level == "subset": + original_data_path = None + subset_path = tmp_path / "provided-subset" + _write_images(subset_path, ["nested/selected.jpg"]) + elif entry_level == "calibration": + original_data_path = None + calib_data_path = tmp_path / "provided-calibration" + calib_data_path.mkdir() + np.save(calib_data_path / "selected.npy", np.ones((2, 2, 3), dtype=np.float32)) + elif entry_level != "data": + raise ValueError(f"Unsupported test entry level: {entry_level}") + + def _ensure_dataset( + task: str, data_path: str | Path | None, dataset: str | None + ) -> Path: + calls["ensure_dataset"] = (task, dataset) + if entry_level != "data": + pytest.fail(f"{entry_level} input must skip original dataset preparation") + return dataset_path + + monkeypatch.setattr(compile_module, "ensure_calibration_dataset", _ensure_dataset) + + compile_module.compile_vision_model( + "fake-model", + target_device=target_device, + model_type="VARIANT", + model_path=model_path, + onnx_path=onnx_path, + data_path=original_data_path, + subset_path=subset_path, + calib_data_path=calib_data_path, + save_path=save_path, + subset_size=1, + ) + return calls, hosted_onnx + + +def test_compile_uses_local_onnx_and_exact_options( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Prefer a local ONNX and pass exact Aries/GPU compiler options.""" + + local_onnx = tmp_path / "local.onnx" + local_onnx.write_bytes(b"onnx") + model_dir = tmp_path / ".mblt_model_zoo" + monkeypatch.setattr(compile_module, "DEFAULT_MODEL_DIR", model_dir) + + calls, _ = _run_fake_compile( + monkeypatch, tmp_path, task="image_classification", model_path=local_onnx + ) + + assert calls["engine_kwargs"]["model_path"] == str(local_onnx) + assert calls["engine_kwargs"]["model_type"] == "VARIANT" + assert calls["compile_kwargs"] | {"calibration_config": None} == { + "model": str(local_onnx), + "calib_data_path": calls["compile_kwargs"]["calib_data_path"], + "save_path": str(model_dir / "local.mxq"), + "image_channels": 3, + "backend": "onnx", + "device": "gpu", + "target_device": "aries-rb", + "inference_scheme": "all", + "calibration_config": None, + } + assert calls["calibration_kwargs"]["output"] == 0 + assert calls["calibration_kwargs"]["method"] == 1 + assert calls["calibration_kwargs"]["mode"] == 1 + assert calls["array"].dtype == np.float32 + assert calls["disposed"] is True + assert not calls["temporary_root"].exists() + + +def test_compile_forwards_required_target_device( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Forward the selected board to qbcompiler rather than hard-coding Aries.""" + + calls, _ = _run_fake_compile( + monkeypatch, + tmp_path, + task="image_classification", + model_path=None, + target_device="regulus-rb", + ) + + assert calls["compile_kwargs"]["target_device"] == "regulus-rb" + + +def test_compile_uses_obb_task_from_model_config( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Compile custom model metadata using the canonical OBB task.""" + + calls, _ = _run_fake_compile( + monkeypatch, + tmp_path, + task="obb", + model_path=None, + dataset="dotav1", + entry_level="calibration", + ) + + assert calls["calibration_kwargs"]["output"] == 1 + + +def test_compile_rejects_non_onnx_local_model_with_ready_calibration_data( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Fail clearly before sending an MXQ file to the ONNX compiler backend.""" + + mxq_path = tmp_path / "local.mxq" + mxq_path.write_bytes(b"mxq") + + with pytest.raises(ValueError, match=r"requires an ONNX model.*local\.mxq"): + _run_fake_compile( + monkeypatch, + tmp_path, + task="image_classification", + model_path=mxq_path, + entry_level="calibration", + ) + + +@pytest.mark.parametrize("entry_level", ["data", "calibration"]) +def test_compile_uses_onnx_path_when_model_path_is_empty( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, entry_level: str +) -> None: + """Treat an adapter-provided empty model path as absent, not preferred.""" + + local_onnx = tmp_path / "alias.onnx" + local_onnx.write_bytes(b"onnx") + + calls, _ = _run_fake_compile( + monkeypatch, + tmp_path, + task="image_classification", + model_path="", + onnx_path=local_onnx, + entry_level=entry_level, + ) + + assert calls["compile_kwargs"]["model"] == str(local_onnx) + + +@pytest.mark.parametrize("entry_level", ["data", "calibration"]) +def test_compile_requires_an_mxq_output_path( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, entry_level: str +) -> None: + """Reject mislabeled compiler output paths in both compilation branches.""" + + with pytest.raises(ValueError, match=r"output path must end with `.mxq`"): + _run_fake_compile( + monkeypatch, + tmp_path, + task="image_classification", + model_path=None, + entry_level=entry_level, + save_path=tmp_path / "compiled.onnx", + ) + + +@pytest.mark.parametrize("entry_level", ["data", "calibration"]) +def test_compile_rejects_output_path_resolving_to_source_onnx( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, entry_level: str +) -> None: + """Do not let an MXQ-named symlink overwrite the input ONNX model.""" + + source_onnx = tmp_path / "source.onnx" + source_onnx.write_bytes(b"onnx") + output_alias = tmp_path / "compiled.mxq" + output_alias.symlink_to(source_onnx) + + with pytest.raises(ValueError, match="must not be the same as the input ONNX path"): + _run_fake_compile( + monkeypatch, + tmp_path, + task="image_classification", + model_path=source_onnx, + entry_level=entry_level, + save_path=output_alias, + ) + + +def test_compile_ignores_missing_local_path_and_uses_hosted_onnx( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Construct the engine without a nonexistent local path so Hub fallback applies.""" + + model_dir = tmp_path / ".mblt_model_zoo" + monkeypatch.setattr(compile_module, "DEFAULT_MODEL_DIR", model_dir) + calls, hosted_onnx = _run_fake_compile( + monkeypatch, + tmp_path, + task="object_detection", + model_path=tmp_path / "missing.onnx", + ) + + assert "model_path" not in calls["engine_kwargs"] + assert calls["compile_kwargs"]["model"] == str(hosted_onnx) + assert calls["compile_kwargs"]["save_path"] == str(model_dir / "hosted-model.mxq") + assert calls["calibration_kwargs"]["output"] == 1 + + +def test_compile_routes_semantic_calibration_by_model_dataset( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Use postprocess taxonomy to distinguish ADE20K and Cityscapes compilation.""" + + calls, _ = _run_fake_compile( + monkeypatch, + tmp_path, + task="semantic_segmentation", + dataset="cityscapes", + model_path=None, + ) + + assert calls["ensure_dataset"] == ("semantic_segmentation", "cityscapes") + assert calls["preprocessed"] + + +def test_compile_starts_from_provided_image_subset( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Skip original dataset preparation and sampling for a supplied image subset.""" + + calls, _ = _run_fake_compile( + monkeypatch, + tmp_path, + task="object_detection", + model_path=None, + entry_level="subset", + ) + + assert "ensure_dataset" not in calls + assert calls["preprocessed"] == [ + str(tmp_path / "provided-subset" / "nested" / "selected.jpg") + ] + assert not calls["temporary_root"].exists() + + +def test_compile_uses_provided_calibration_dataset_directly( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Pass ready NumPy tensors directly without image dataset processing.""" + + calls, _ = _run_fake_compile( + monkeypatch, + tmp_path, + task="object_detection", + model_path=None, + entry_level="calibration", + ) + + calibration_path = tmp_path / "provided-calibration" + assert "ensure_dataset" not in calls + assert "engine_kwargs" not in calls + assert "preprocessed" not in calls + assert calls["compile_kwargs"]["calib_data_path"] == str(calibration_path) + assert (calibration_path / "selected.npy").is_file() + + +def test_compile_rejects_multiple_data_levels() -> None: + """Reject ambiguous original, subset, and calibration path combinations.""" + + with pytest.raises(ValueError, match="Provide only one calibration pipeline input"): + compile_module.compile_vision_model( + "alexnet", + target_device="aries-rb", + data_path="original", + subset_path="subset", + ) + + +def test_compile_cleans_up_and_disposes_after_failure( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Dispose the engine and delete temporary images and arrays after compiler failure.""" + + calls: dict[str, Any] = {} + with pytest.raises(RuntimeError, match="compile failed"): + _run_fake_compile( + monkeypatch, + tmp_path, + task="image_classification", + model_path=None, + fail=True, + calls=calls, + ) + + assert calls["disposed"] is True + assert not calls["temporary_root"].exists() + + +def test_cli_compile_parser_and_dispatch( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Parse compile options and dispatch them to the packaged API.""" diff --git a/tests/test_dataloader.py b/tests/test_dataloader.py new file mode 100644 index 0000000..6ca39c8 --- /dev/null +++ b/tests/test_dataloader.py @@ -0,0 +1,206 @@ +"""Tests for vision dataset sample discovery.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import cv2 +import numpy as np +import pytest +from PIL import Image + +from mblt_vision.utils.datasets.dataloader import ( + CustomCOCODataset, + CustomADE20K, + CustomCityscapes, + CustomDOTAv1, + CustomImageFolder, + CustomNYUDepth, + CustomWiderFaceDataset, +) + + +@pytest.mark.parametrize( + ("dataset_class", "class_name", "expected_count", "expected_suffixes"), + [ + (CustomImageFolder, "class-a", 1, {".jpg"}), + (CustomWiderFaceDataset, "0--Parade", 1, {".jpg"}), + ], +) +def test_class_based_datasets_ignore_non_image_files( + tmp_path: Path, + dataset_class: type[CustomImageFolder] | type[CustomWiderFaceDataset], + class_name: str, + expected_count: int, + expected_suffixes: set[str], +) -> None: + """Keep incidental files out of ImageNet and WiderFace sample lists.""" + + class_dir = tmp_path / class_name + nested_dir = class_dir / "nested" + nested_dir.mkdir(parents=True) + (class_dir / "image.JPG").write_bytes(b"image") + (nested_dir / "image.png").write_bytes(b"image") + (class_dir / ".DS_Store").write_bytes(b"metadata") + (nested_dir / "labels.txt").write_text("metadata", encoding="utf-8") + + dataset = dataset_class(str(tmp_path)) + + assert len(dataset) == expected_count + assert { + Path(sample[0]).suffix.lower() for sample in dataset.samples + } == expected_suffixes + + +def test_coco_dataset_rejects_annotation_geometry_mismatching_image() -> None: + """Do not evaluate COCO labels with geometry different from their decoded image.""" + + dataset = cast( + CustomCOCODataset, + SimpleNamespace( + ids=[1], + coco=SimpleNamespace(imgs={1: {"height": 4, "width": 5}}), + _load_image=lambda _: np.zeros((3, 5, 3), dtype=np.uint8), + ), + ) + + with pytest.raises( + ValueError, match=r"image ID 1: annotation \(4, 5\), image \(3, 5\)" + ): + CustomCOCODataset.__getitem__(dataset, 0) + + +@pytest.mark.parametrize("file_name", ["/tmp/outside.jpg", "../outside.jpg"]) +def test_coco_dataset_rejects_unsafe_annotation_file_names(file_name: str) -> None: + """Do not let COCO metadata access images outside the configured root.""" + + dataset = cast( + CustomCOCODataset, + SimpleNamespace( + root="/dataset", + coco=SimpleNamespace( + loadImgs=lambda _: [{"file_name": file_name}], + ), + ), + ) + + with pytest.raises(ValueError, match="unsafe file_name"): + CustomCOCODataset._load_image(dataset, 1) + + +def test_dota_dataset_rejects_duplicate_image_stems(tmp_path: Path) -> None: + """Direct DOTAv1 loading must not silently collapse same-stem images.""" + + image_dir = tmp_path / "images" + image_dir.mkdir() + (image_dir / "sample.jpg").write_bytes(b"jpg") + (image_dir / "sample.png").write_bytes(b"png") + + with pytest.raises(ValueError, match="duplicate filename stems"): + CustomDOTAv1(str(tmp_path)) + + +@pytest.mark.parametrize( + ("dataset_class", "source_id"), + [(CustomADE20K, 0), (CustomCityscapes, 255)], +) +def test_dense_semantic_datasets_reject_all_ignored_targets( + tmp_path: Path, + dataset_class: type[CustomADE20K] | type[CustomCityscapes], + source_id: int, +) -> None: + """Reject direct semantic evaluation samples that contain no valid labels.""" + + image_dir = tmp_path / "images" + annotation_dir = tmp_path / "annotations" + image_dir.mkdir() + annotation_dir.mkdir() + Image.new("RGB", (2, 2)).save(image_dir / "sample.png") + Image.fromarray(np.full((2, 2), source_id, dtype=np.uint8)).save( + annotation_dir / "sample.png" + ) + dataset = dataset_class(str(tmp_path)) + + with pytest.raises(ValueError, match="contains no evaluable class IDs"): + dataset[0] + + +@pytest.mark.parametrize( + ("dataset_class", "target_dir", "target_names"), + [ + (CustomNYUDepth, "depth", ("sample.npy", "sample.NPY")), + (CustomADE20K, "annotations", ("sample.png", "sample.PNG")), + (CustomCityscapes, "annotations", ("sample.png", "sample.PNG")), + ], +) +def test_dense_datasets_reject_duplicate_target_stems( + tmp_path: Path, + dataset_class: type[CustomNYUDepth] | type[CustomADE20K] | type[CustomCityscapes], + target_dir: str, + target_names: tuple[str, str], +) -> None: + """Do not let case variants silently overwrite a dense target mapping.""" + + (tmp_path / "images").mkdir() + (tmp_path / target_dir).mkdir() + (tmp_path / "images" / "sample.jpg").write_bytes(b"image") + for target_name in target_names: + (tmp_path / target_dir / target_name).write_bytes(b"target") + + with pytest.raises(ValueError, match="duplicate filename stems"): + dataset_class(str(tmp_path)) + + +def test_nyu_dataset_rejects_negative_depth_targets(tmp_path: Path) -> None: + """Do not silently exclude corrupted negative depths during evaluation.""" + + image_dir = tmp_path / "images" + depth_dir = tmp_path / "depth" + image_dir.mkdir() + depth_dir.mkdir() + assert cv2.imwrite(str(image_dir / "sample.jpg"), np.zeros((1, 1, 3), np.uint8)) + np.save(depth_dir / "sample.npy", np.array([[-1]], dtype=np.float32)) + + with pytest.raises(ValueError, match="must not contain negative values"): + CustomNYUDepth(str(tmp_path))[0] + + +def test_cityscapes_dataset_rejects_unknown_source_ids(tmp_path: Path) -> None: + """Reject corrupted Cityscapes labels instead of remapping them to ignore.""" + + image_dir = tmp_path / "images" + annotation_dir = tmp_path / "annotations" + image_dir.mkdir() + annotation_dir.mkdir() + assert cv2.imwrite( + str(image_dir / "sample.png"), np.zeros((2, 2, 3), dtype=np.uint8) + ) + Image.fromarray(np.full((2, 2), 200, dtype=np.uint8)).save( + annotation_dir / "sample.png" + ) + + with pytest.raises(ValueError, match=r"unsupported source IDs \[200\]"): + CustomCityscapes(str(tmp_path))[0] + + +@pytest.mark.parametrize( + "mask", [np.zeros((2, 2, 3), dtype=np.uint8), np.zeros((2, 2), dtype=np.uint16)] +) +def test_ade20k_dataset_rejects_noncanonical_mask_encodings( + tmp_path: Path, mask: np.ndarray +) -> None: + """Reject color and high-bit-depth ADE20K masks before source-ID conversion.""" + + image_dir = tmp_path / "images" + annotation_dir = tmp_path / "annotations" + image_dir.mkdir() + annotation_dir.mkdir() + assert cv2.imwrite( + str(image_dir / "sample.png"), np.zeros((2, 2, 3), dtype=np.uint8) + ) + Image.fromarray(mask).save(annotation_dir / "sample.png") + + with pytest.raises(ValueError, match="single-channel 8-bit PNG masks"): + CustomADE20K(str(tmp_path))[0] diff --git a/tests/test_dataset_organizer.py b/tests/test_dataset_organizer.py new file mode 100644 index 0000000..94a4098 --- /dev/null +++ b/tests/test_dataset_organizer.py @@ -0,0 +1,1682 @@ +"""Tests for dataset download organization helpers.""" + +from __future__ import annotations + +import hashlib +import io +import inspect +import json +import os +import shutil +import tarfile +from collections.abc import Callable +from pathlib import Path +from typing import Any +from urllib.parse import urlparse +from zipfile import ZipFile + +import pytest +import requests +import numpy as np +from PIL import Image + +import mblt_vision.utils.datasets.organizer as organizer +from mblt_vision.utils.datasets import readiness as readiness_module + + +class _DummyTqdm: + """Minimal tqdm stub for download tests.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.updated = 0 + + def __enter__(self) -> _DummyTqdm: + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + return None + + def update(self, value: int) -> None: + self.updated += value + + +def test_dotav1_normalized_labels_preserve_difficult_metadata(tmp_path: Path) -> None: + """Keep official difficult regions in the normalized label export.""" + + image_path = tmp_path / "image.png" + Image.new("RGB", (100, 50)).save(image_path) + source_path = tmp_path / "source.txt" + source_path.write_text( + "imagesource:GoogleEarth\n" + "0 0 20 0 20 20 0 20 plane 0\n" + "30 0 50 0 50 20 30 20 plane 1\n", + encoding="utf-8", + ) + output_path = tmp_path / "normalized.txt" + + organizer._write_dotav1_yolo_labels( + str(image_path), str(source_path), str(output_path) + ) + + assert output_path.read_text(encoding="utf-8").splitlines() == [ + "0 0 0 0.2 0 0.2 0.4 0 0.4 0", + "0 0.3 0 0.5 0 0.5 0.4 0.3 0.4 1", + ] + + +def test_dotav1_organizer_rejects_truncated_raw_annotation_rows( + tmp_path: Path, +) -> None: + """Do not silently remove malformed source annotations during organization.""" + + image_path = tmp_path / "image.png" + Image.new("RGB", (100, 50)).save(image_path) + source_path = tmp_path / "source.txt" + source_path.write_text( + "imagesource:GoogleEarth\ngsd:0.5\n0 0 20 0 20 20 0 plane 0\n", + encoding="utf-8", + ) + + with pytest.raises( + ValueError, + match=r"source\.txt at line 3: expected at least 10 fields, got 9", + ): + organizer._write_dotav1_yolo_labels( + str(image_path), str(source_path), str(tmp_path / "normalized.txt") + ) + + +def test_dotav1_organizer_rejects_repeated_polygon_vertices(tmp_path: Path) -> None: + """Keep malformed triangle-like rows out of the staged normalized labels.""" + + image_path = tmp_path / "image.png" + Image.new("RGB", (100, 50)).save(image_path) + source_path = tmp_path / "source.txt" + source_path.write_text("0 0 20 0 20 20 0 0 plane 0\n", encoding="utf-8") + + with pytest.raises(ValueError, match="four distinct vertices"): + organizer._write_dotav1_yolo_labels( + str(image_path), str(source_path), str(tmp_path / "normalized.txt") + ) + + +def test_dotav1_organizer_rejects_duplicate_targets(tmp_path: Path) -> None: + """Reject duplicate raw targets before creating normalized labels.""" + + image_path = tmp_path / "image.png" + Image.new("RGB", (100, 50)).save(image_path) + source_path = tmp_path / "source.txt" + target = "0 0 20 0 20 20 0 20 plane 0\n" + source_path.write_text(target * 2, encoding="utf-8") + + with pytest.raises(ValueError, match="Duplicate DOTAv1 annotation target"): + organizer._write_dotav1_yolo_labels( + str(image_path), str(source_path), str(tmp_path / "normalized.txt") + ) + + +def test_dotav1_organizer_rejects_reordered_duplicate_targets( + tmp_path: Path, +) -> None: + """Raw DOTAv1 labels cannot duplicate a target with another start vertex.""" + + image_path = tmp_path / "image.png" + Image.new("RGB", (100, 50)).save(image_path) + source_path = tmp_path / "source.txt" + source_path.write_text( + "0 0 20 0 20 20 0 20 plane 0\n20 0 20 20 0 20 0 0 plane 0\n", + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="Duplicate DOTAv1 annotation target"): + organizer._write_dotav1_yolo_labels( + str(image_path), str(source_path), str(tmp_path / "normalized.txt") + ) + + +def test_dotav1_stage_requires_a_non_difficult_target(tmp_path: Path) -> None: + """Do not replace a usable DOTAv1 cache with an unscorable staged split.""" + + image_path = tmp_path / "images" / "P0001.png" + image_path.parent.mkdir() + Image.new("RGB", (100, 50)).save(image_path) + label_path = tmp_path / "labels" / "val_original" / "P0001.txt" + label_path.parent.mkdir(parents=True) + label_path.write_text("0 0 20 0 20 20 0 20 plane 1\n", encoding="utf-8") + + with pytest.raises(ValueError, match="at least one non-difficult target"): + organizer._validate_staged_dotav1_labels(tmp_path) + + +@pytest.mark.parametrize( + ("organize_dataset", "dataset_name"), + [ + (organizer.organize_imagenet, "imagenet"), + (organizer.organize_coco, "coco"), + (organizer.organize_widerface, "widerface"), + (organizer.organize_nyu_depth, "nyu-depth"), + (organizer.organize_ade20k, "ade20k"), + (organizer.organize_cityscapes, "cityscapes"), + (organizer.organize_dotav1, "dotav1"), + ], +) +def test_organizer_defaults_use_the_lazy_cache_resolver( + monkeypatch: pytest.MonkeyPatch, + organize_dataset: Callable[..., Any], + dataset_name: str, + tmp_path: Path, +) -> None: + """Keep direct organizer APIs consistent with artifact cache fallback behavior.""" + + output_dir = inspect.signature(organize_dataset).parameters["output_dir"].default + assert output_dir is None + + import mblt_vision.wrapper as wrapper + + cache_dir = tmp_path / "cache" + monkeypatch.setattr(wrapper, "get_mobilint_cache_dir", lambda: str(cache_dir)) + assert organizer._resolve_organizer_output_dir(None, dataset_name) == str( + cache_dir / "datasets" / dataset_name + ) + + +class _FakeResponse: + """Simple streaming response test double.""" + + def __init__( + self, status_code: int, headers: dict[str, str], chunks: list[bytes | Exception] + ) -> None: + self.status_code = status_code + self.headers = headers + self._chunks = chunks + + def __enter__(self) -> _FakeResponse: + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + return None + + def raise_for_status(self) -> None: + if self.status_code >= 400: + response = requests.Response() + response.status_code = self.status_code + raise requests.HTTPError(response=response) + + def iter_content(self, chunk_size: int) -> Any: + del chunk_size + for chunk in self._chunks: + if isinstance(chunk, Exception): + raise chunk + yield chunk + + +def test_download_url_retries_and_resumes_partial_file( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Resume a partial archive download after a transient connection failure.""" + + first_chunk = b"abc" + second_chunk = b"def" + calls: list[dict[str, str]] = [] + responses = [ + _FakeResponse( + status_code=200, + headers={"Content-Length": str(len(first_chunk) + len(second_chunk))}, + chunks=[first_chunk, requests.ConnectionError("interrupted")], + ), + _FakeResponse( + status_code=206, + headers={ + "Content-Length": str(len(second_chunk)), + "Content-Range": "bytes 3-5/6", + }, + chunks=[second_chunk], + ), + ] + + def _fake_get( + url: str, stream: bool, timeout: tuple[int, int], headers: dict[str, str] + ) -> _FakeResponse: + del url, stream, timeout + calls.append(dict(headers)) + return responses.pop(0) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + monkeypatch.setattr(organizer, "tqdm", _DummyTqdm) + monkeypatch.setattr(organizer, "sleep", lambda _: None) + + local_path = tmp_path / "archive.tar" + result = organizer._download_url("https://example.com/archive.tar", str(local_path)) + + assert result == str(local_path) + assert local_path.read_bytes() == first_chunk + second_chunk + assert calls == [{}, {"Range": "bytes=3-"}] + + +def test_download_url_restarts_when_resume_offset_does_not_match( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Discard a partial file when a range response starts at another offset.""" + + calls: list[dict[str, str]] = [] + responses = [ + _FakeResponse( + status_code=206, + headers={"Content-Length": "3", "Content-Range": "bytes 0-2/6"}, + chunks=[b"bad"], + ), + _FakeResponse( + status_code=200, + headers={"Content-Length": "6"}, + chunks=[b"fresh!"], + ), + ] + + def _fake_get( + url: str, stream: bool, timeout: tuple[int, int], headers: dict[str, str] + ) -> _FakeResponse: + del url, stream, timeout + calls.append(dict(headers)) + return responses.pop(0) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + monkeypatch.setattr(organizer, "tqdm", _DummyTqdm) + local_path = tmp_path / "archive.tar" + local_path.write_bytes(b"old") + + assert organizer._download_url( + "https://example.com/archive.tar", str(local_path) + ) == str(local_path) + assert local_path.read_bytes() == b"fresh!" + assert calls == [{"Range": "bytes=3-"}, {}] + + +def test_download_url_accepts_completed_archive_on_range_not_satisfiable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Treat a matching 416 total as confirmation that the local file is complete.""" + + payload = b"complete" + calls: list[dict[str, str]] = [] + + def _fake_get( + url: str, stream: bool, timeout: tuple[int, int], headers: dict[str, str] + ) -> _FakeResponse: + del url, stream, timeout + calls.append(dict(headers)) + return _FakeResponse( + status_code=416, + headers={"Content-Range": f"bytes */{len(payload)}"}, + chunks=[], + ) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + local_path = tmp_path / "archive.tar" + local_path.write_bytes(payload) + + assert organizer._download_url( + "https://example.com/archive.tar", + str(local_path), + expected_sha256=hashlib.sha256(payload).hexdigest(), + ) == str(local_path) + assert local_path.read_bytes() == payload + assert calls == [{"Range": f"bytes={len(payload)}-"}] + + +def test_download_url_restarts_after_incomplete_range_not_satisfiable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Do not accept a 416 when the local partial file has the wrong size.""" + + responses = [ + _FakeResponse( + status_code=416, + headers={"Content-Range": "bytes */6"}, + chunks=[], + ), + _FakeResponse( + status_code=200, + headers={"Content-Length": "6"}, + chunks=[b"fresh!"], + ), + ] + + def _fake_get(url: str, **kwargs: Any) -> _FakeResponse: + del url, kwargs + return responses.pop(0) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + monkeypatch.setattr(organizer, "tqdm", _DummyTqdm) + local_path = tmp_path / "archive.tar" + local_path.write_bytes(b"old") + + assert organizer._download_url( + "https://example.com/archive.tar", str(local_path) + ) == str(local_path) + assert local_path.read_bytes() == b"fresh!" + + +@pytest.mark.parametrize("status_code", [408, 429, 503]) +def test_download_url_retries_transient_http_responses( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + status_code: int, +) -> None: + """Retry transient HTTP responses before accepting a successful download.""" + + calls = 0 + responses = [ + _FakeResponse(status_code=status_code, headers={}, chunks=[]), + _FakeResponse(status_code=200, headers={"Content-Length": "2"}, chunks=[b"ok"]), + ] + + def _fake_get(url: str, **kwargs: Any) -> _FakeResponse: + nonlocal calls + del url, kwargs + calls += 1 + return responses.pop(0) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + monkeypatch.setattr(organizer, "tqdm", _DummyTqdm) + monkeypatch.setattr(organizer, "sleep", lambda _: None) + + local_path = tmp_path / "archive.tar" + assert organizer._download_url( + "https://example.com/archive.tar", str(local_path) + ) == str(local_path) + assert local_path.read_bytes() == b"ok" + assert calls == 2 + + +def test_download_url_does_not_retry_permanent_http_client_errors( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Fail immediately for permanent client errors such as HTTP 404.""" + + calls = 0 + + def _fake_get(url: str, **kwargs: Any) -> _FakeResponse: + nonlocal calls + del url, kwargs + calls += 1 + return _FakeResponse(status_code=404, headers={}, chunks=[]) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + + with pytest.raises(requests.HTTPError): + organizer._download_url( + "https://example.com/archive.tar", str(tmp_path / "archive.tar") + ) + + assert calls == 1 + + +def test_download_url_verifies_pinned_archive_sha256( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Accept only downloaded archive bytes matching the configured digest.""" + + payload = b"verified archive" + + def _fake_get(url: str, **kwargs: Any) -> _FakeResponse: + del url, kwargs + return _FakeResponse( + status_code=200, + headers={"Content-Length": str(len(payload))}, + chunks=[payload], + ) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + monkeypatch.setattr(organizer, "tqdm", _DummyTqdm) + local_path = tmp_path / "archive.zip" + + assert organizer._download_url( + "https://example.com/archive.zip", + str(local_path), + expected_sha256=hashlib.sha256(payload).hexdigest(), + ) == str(local_path) + assert local_path.read_bytes() == payload + + +def test_download_url_rejects_and_removes_digest_mismatches( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Never leave an unverified archive available for later extraction.""" + + def _fake_get(url: str, **kwargs: Any) -> _FakeResponse: + del url, kwargs + return _FakeResponse( + status_code=200, + headers={"Content-Length": "9"}, + chunks=[b"malicious"], + ) + + monkeypatch.setattr("mblt_vision.utils.datasets.organizer.requests.get", _fake_get) + monkeypatch.setattr(organizer, "tqdm", _DummyTqdm) + local_path = tmp_path / "archive.zip" + + with pytest.raises(ValueError, match="SHA-256 mismatch"): + organizer._download_url( + "https://example.com/archive.zip", + str(local_path), + expected_sha256=hashlib.sha256(b"expected").hexdigest(), + ) + + assert not local_path.exists() + + +def test_download_if_url_passes_registered_archive_digest( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Use the package-pinned digest whenever a default archive is downloaded.""" + + source_url, expected_sha256 = next(iter(organizer.PINNED_ARCHIVE_SHA256.items())) + recorded: dict[str, str | None] = {} + + def _fake_download( + url: str, local_path: str, expected_sha256: str | None = None + ) -> str: + recorded.update(url=url, local_path=local_path, expected_sha256=expected_sha256) + return local_path + + monkeypatch.setattr(organizer, "_download_url", _fake_download) + + assert organizer._download_if_url(source_url, str(tmp_path)) == str( + tmp_path / Path(urlparse(source_url).path).name + ) + assert recorded["url"] == source_url + assert recorded["expected_sha256"] == expected_sha256 + + +def test_download_if_url_rejects_plaintext_http( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not fetch benchmark inputs over an unauthenticated transport.""" + + monkeypatch.setattr( + "mblt_vision.utils.datasets.organizer.requests.get", + lambda *args, **kwargs: pytest.fail("plaintext URL must not be requested"), + ) + + with pytest.raises(ValueError, match="must use HTTPS"): + organizer._download_if_url("http://example.com/archive.zip", str(tmp_path)) + + +def test_coco_and_ade20k_downloads_use_pinned_https_archives() -> None: + """Keep default benchmark datasets on authenticated, integrity-checked URLs.""" + + assert all(url.startswith("https://") for url in organizer.PINNED_ARCHIVE_SHA256) + assert all( + len(digest) == 64 + and all(character in "0123456789abcdef" for character in digest) + for digest in organizer.PINNED_ARCHIVE_SHA256.values() + ) + + +def test_should_download_serially_for_same_host_urls() -> None: + """Serialize same-host dataset archive downloads to avoid throttling.""" + + assert organizer._should_download_serially( + [ + "https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar", + "https://image-net.org/data/ILSVRC/2012/ILSVRC2012_bbox_val_v3.tgz", + ] + ) + + assert not organizer._should_download_serially( + [ + "https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar", + "https://example.com/data/annotations.tgz", + ] + ) + + +def _create_imagenet_source(tmp_path: Path) -> tuple[Path, Path]: + """Create a small structurally valid ImageNet organizer source.""" + + image_dir = tmp_path / "source-images" + xml_dir = tmp_path / "source-xml" / "val" + image_dir.mkdir() + xml_dir.mkdir(parents=True) + for index in range(50): + stem = f"ILSVRC2012_val_{index + 1:08d}" + (image_dir / f"{stem}.JPEG").write_bytes(b"image") + (xml_dir / f"{stem}.xml").write_text( + "n00000001", + encoding="utf-8", + ) + return image_dir, xml_dir.parent + + +def _create_coco_source(tmp_path: Path) -> tuple[Path, Path]: + """Create a small COCO organizer source.""" + + image_dir = tmp_path / "source-images" + annotation_dir = tmp_path / "source-annotations" / "annotations" + image_dir.mkdir() + annotation_dir.mkdir(parents=True) + (image_dir / "000000000001.jpg").write_bytes(b"image") + (annotation_dir / "instances_val2017.json").write_text("{}", encoding="utf-8") + return image_dir, annotation_dir.parent + + +def _create_widerface_source(tmp_path: Path) -> tuple[Path, Path]: + """Create a small WiderFace organizer source.""" + + image_dir = tmp_path / "source-images" / "images" / "0--Parade" + annotation_dir = tmp_path / "source-annotations" + image_dir.mkdir(parents=True) + annotation_dir.mkdir() + (image_dir / "sample.jpg").write_bytes(b"image") + (annotation_dir / "wider_face_val.mat").write_bytes(b"metadata") + return image_dir.parent.parent, annotation_dir + + +def test_construct_imagenet_replaces_stale_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Replace a managed ImageNet root only after staged readiness succeeds.""" + + image_dir, xml_dir = _create_imagenet_source(tmp_path) + readiness_calls: list[tuple[str, str]] = [] + + def _dataset_ready(_: str, task: str, dataset: str) -> bool: + readiness_calls.append((task, dataset)) + return True + + monkeypatch.setattr(organizer, "dataset_ready", _dataset_ready) + monkeypatch.setattr(organizer, "_validate_staged_payloads", lambda *_: None) + + output_dir = tmp_path / "imagenet" + (output_dir / "stale-class").mkdir(parents=True) + (output_dir / "stale-class" / "stale.JPEG").write_bytes(b"stale") + + organizer.construct_imagenet(str(image_dir), str(xml_dir), str(output_dir)) + + assert {path.name for path in output_dir.iterdir()} == {"n00000001"} + assert len(list((output_dir / "n00000001").iterdir())) == 50 + assert readiness_calls == [("image_classification", "imagenet")] + + +@pytest.mark.parametrize("class_name", ["../../escaped", "/tmp/escaped"]) +def test_construct_imagenet_rejects_unsafe_xml_class_name( + tmp_path: Path, class_name: str +) -> None: + """Reject XML class names that could escape the ImageNet staging root.""" + + image_dir, xml_dir = _create_imagenet_source(tmp_path) + xml_path = next((xml_dir / "val").glob("*.xml")) + xml_path.write_text( + f"{class_name}", + encoding="utf-8", + ) + output_dir = tmp_path / "imagenet" + + with pytest.raises(ValueError, match="invalid ImageNet synset name"): + organizer.construct_imagenet(str(image_dir), str(xml_dir), str(output_dir)) + + assert not (tmp_path / "escaped").exists() + assert not output_dir.exists() + + +def test_construct_coco_replaces_stale_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Replace a managed COCO root only after staged readiness succeeds.""" + + image_dir, annotation_dir = _create_coco_source(tmp_path) + readiness_calls: list[tuple[str, str]] = [] + + def _dataset_ready(_: str, task: str, dataset: str) -> bool: + readiness_calls.append((task, dataset)) + return True + + monkeypatch.setattr(organizer, "dataset_ready", _dataset_ready) + monkeypatch.setattr(organizer, "_validate_staged_payloads", lambda *_: None) + + output_dir = tmp_path / "coco" + (output_dir / "val2017").mkdir(parents=True) + (output_dir / "val2017" / "stale.jpg").write_bytes(b"stale") + (output_dir / "stale_val2017.json").write_text("{}", encoding="utf-8") + + organizer.construct_coco(str(image_dir), str(annotation_dir), str(output_dir)) + + assert {path.name for path in (output_dir / "val2017").iterdir()} == { + "000000000001.jpg" + } + assert {path.name for path in output_dir.iterdir()} == { + "val2017", + "instances_val2017.json", + } + assert readiness_calls == [ + ("object_detection", "coco"), + ("instance_segmentation", "coco"), + ("pose_estimation", "coco"), + ] + + +def test_construct_widerface_replaces_stale_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Replace a managed WiderFace root only after staged readiness succeeds.""" + + image_dir, annotation_dir = _create_widerface_source(tmp_path) + readiness_calls: list[tuple[str, str]] = [] + + def _dataset_ready(_: str, task: str, dataset: str) -> bool: + readiness_calls.append((task, dataset)) + return True + + monkeypatch.setattr(organizer, "dataset_ready", _dataset_ready) + monkeypatch.setattr(organizer, "_validate_staged_payloads", lambda *_: None) + + output_dir = tmp_path / "widerface" + (output_dir / "images" / "stale-event").mkdir(parents=True) + (output_dir / "images" / "stale-event" / "stale.jpg").write_bytes(b"stale") + (output_dir / "stale_val.mat").write_bytes(b"stale") + + organizer.construct_widerface(str(image_dir), str(annotation_dir), str(output_dir)) + + assert {path.name for path in (output_dir / "images").iterdir()} == {"0--Parade"} + assert {path.name for path in output_dir.iterdir()} == { + "images", + "wider_face_val.mat", + } + assert readiness_calls == [("face_detection", "widerface")] + + +@pytest.mark.parametrize("dataset", ["imagenet", "coco", "widerface"]) +def test_incomplete_staged_dataset_preserves_existing_cache( + tmp_path: Path, + dataset: str, +) -> None: + """Reject incomplete staged roots before replacing any existing cache.""" + + source_builders = { + "imagenet": _create_imagenet_source, + "coco": _create_coco_source, + "widerface": _create_widerface_source, + } + constructors = { + "imagenet": organizer.construct_imagenet, + "coco": organizer.construct_coco, + "widerface": organizer.construct_widerface, + } + first_source, second_source = source_builders[dataset](tmp_path) + + output_dir = tmp_path / dataset + output_dir.mkdir() + (output_dir / "valid-cache-marker").write_bytes(b"existing") + + with pytest.raises(ValueError, match="existing dataset cache was not replaced"): + constructors[dataset](str(first_source), str(second_source), str(output_dir)) + + assert (output_dir / "valid-cache-marker").read_bytes() == b"existing" + + +def test_safe_unpack_archive_preserves_regular_tar_layout(tmp_path: Path) -> None: + """Extract regular files and directories from a supported tar archive.""" + + archive_path = tmp_path / "dataset.tar" + with tarfile.open(archive_path, "w") as archive: + directory = tarfile.TarInfo("dataset/") + directory.type = tarfile.DIRTYPE + archive.addfile(directory) + payload = b"dataset contents" + member = tarfile.TarInfo("dataset/sample.txt") + member.size = len(payload) + archive.addfile(member, io.BytesIO(payload)) + + extract_dir = tmp_path / "extracted" + organizer._safe_unpack_archive(str(archive_path), str(extract_dir)) + + assert (extract_dir / "dataset" / "sample.txt").read_bytes() == b"dataset contents" + + +@pytest.mark.parametrize("suffix", [".zip", ".tar"]) +def test_safe_unpack_archive_rejects_duplicate_member_paths( + tmp_path: Path, suffix: str +) -> None: + """Reject duplicate archive members before any staged file can be replaced.""" + + archive_path = tmp_path / f"dataset{suffix}" + if suffix == ".zip": + with ZipFile(archive_path, "w") as archive: + archive.writestr("dataset/sample.txt", b"first") + archive.writestr("dataset/sample.txt", b"second") + else: + with tarfile.open(archive_path, "w") as archive: + for payload in (b"first", b"second"): + member = tarfile.TarInfo("dataset/sample.txt") + member.size = len(payload) + archive.addfile(member, io.BytesIO(payload)) + + with pytest.raises(ValueError, match="Duplicate archive member path"): + organizer._safe_unpack_archive(str(archive_path), str(tmp_path / "extracted")) + + +def test_safe_unpack_archive_rejects_tar_traversal_before_writing( + tmp_path: Path, +) -> None: + """Reject a traversal member without extracting earlier valid members.""" + + archive_path = tmp_path / "dataset.tar" + outside_path = tmp_path / "outside.txt" + with tarfile.open(archive_path, "w") as archive: + safe_payload = b"safe" + safe_member = tarfile.TarInfo("dataset/safe.txt") + safe_member.size = len(safe_payload) + archive.addfile(safe_member, io.BytesIO(safe_payload)) + outside_payload = b"outside" + outside_member = tarfile.TarInfo("../outside.txt") + outside_member.size = len(outside_payload) + archive.addfile(outside_member, io.BytesIO(outside_payload)) + + extract_dir = tmp_path / "extracted" + with pytest.raises(ValueError, match="Unsafe archive member path"): + organizer._safe_unpack_archive(str(archive_path), str(extract_dir)) + + assert not outside_path.exists() + assert not (extract_dir / "dataset" / "safe.txt").exists() + + +def test_organize_imagenet_rejects_tar_symlink_escape(tmp_path: Path) -> None: + """Reject a tar symlink before ImageNet extraction can write through it.""" + + image_archive = tmp_path / "images.tar" + xml_archive = tmp_path / "annotations.tgz" + outside_path = tmp_path / "outside.txt" + with tarfile.open(image_archive, "w") as archive: + link = tarfile.TarInfo("redirect") + link.type = tarfile.SYMTYPE + link.linkname = str(tmp_path) + archive.addfile(link) + payload = b"outside" + member = tarfile.TarInfo("redirect/outside.txt") + member.size = len(payload) + archive.addfile(member, io.BytesIO(payload)) + with tarfile.open(xml_archive, "w:gz"): + pass + + with pytest.raises(ValueError, match="Unsafe archive member type"): + organizer.organize_imagenet( + str(image_archive), str(xml_archive), str(tmp_path / "imagenet") + ) + + assert not outside_path.exists() + + +def test_organize_nyu_depth_extracts_only_validation_layout( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Install only NYU Depth validation image/depth pairs from an archive.""" + + monkeypatch.setattr(organizer, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(organizer, "_validate_staged_nyu_depth", lambda _: None) + archive_path = tmp_path / "nyu-depth.zip" + with ZipFile(archive_path, "w") as archive: + archive.writestr("nyu-depth/images/train/nyu_train.jpg", b"training image") + archive.writestr("nyu-depth/depth/train/nyu_train.npy", b"training depth") + archive.writestr("nyu-depth/images/val/nyu_0000.jpg", b"validation image") + archive.writestr("nyu-depth/depth/val/nyu_0000.npy", b"validation depth") + + output_dir = tmp_path / "organized" + organizer.organize_nyu_depth(str(archive_path), str(output_dir)) + + assert archive_path.is_file() + assert (output_dir / "images" / "nyu_0000.jpg").read_bytes() == b"validation image" + assert (output_dir / "depth" / "nyu_0000.npy").read_bytes() == b"validation depth" + assert not (output_dir / "images" / "train").exists() + assert not (output_dir / "depth" / "train").exists() + + +@pytest.mark.parametrize("relative_path", ["images/sample.jpg", "depth/sample.npy"]) +def test_construct_nyu_depth_rejects_symlinked_data_files( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + relative_path: str, +) -> None: + """Reject NYU data symlinks without replacing an existing managed cache.""" + + monkeypatch.setattr(organizer, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + dataset_dir = tmp_path / "source" + (dataset_dir / "images").mkdir(parents=True) + (dataset_dir / "depth").mkdir() + (dataset_dir / "images" / "sample.jpg").write_bytes(b"image") + (dataset_dir / "depth" / "sample.npy").write_bytes(b"depth") + external_file = tmp_path / "secret" + external_file.write_bytes(b"outside dataset") + source_path = dataset_dir / relative_path + source_path.unlink() + source_path.symlink_to(external_file) + output_dir = tmp_path / "organized" + output_dir.mkdir() + marker = output_dir / "valid-cache-marker" + marker.write_bytes(b"existing") + + with pytest.raises(ValueError, match="must not be a symlink"): + organizer.construct_nyu_depth(str(dataset_dir), str(output_dir)) + + assert marker.read_bytes() == b"existing" + assert not (output_dir / Path(relative_path).name).exists() + + +def test_construct_nyu_depth_rejects_source_outside_resolved_root( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Reject regular files reached through a directory symlink escaping the dataset root.""" + + monkeypatch.setattr(organizer, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + dataset_dir = tmp_path / "source" + selected_root = dataset_dir / "nyu-depth" + selected_root.mkdir(parents=True) + external_images = dataset_dir / "external-images" + external_images.mkdir() + (external_images / "sample.jpg").write_bytes(b"outside dataset") + (selected_root / "images").symlink_to(external_images, target_is_directory=True) + (selected_root / "depth").mkdir() + (selected_root / "depth" / "sample.npy").write_bytes(b"depth") + + with pytest.raises(ValueError, match="must remain within dataset root"): + organizer.construct_nyu_depth(str(dataset_dir), str(tmp_path / "organized")) + + +def test_construct_nyu_depth_rejects_duplicate_nested_source_stems( + tmp_path: Path, +) -> None: + """Do not let recursive NYU source traversal overwrite a flattened sample.""" + + dataset_dir = tmp_path / "source" + for relative_path in ( + "images/first/sample.jpg", + "images/second/sample.png", + "depth/sample.npy", + ): + path = dataset_dir / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"source") + + with pytest.raises( + ValueError, match="NYU Depth images contain duplicate filename stem" + ): + organizer._collect_nyu_depth_validation_files( + str(dataset_dir / "images"), + str(dataset_dir / "depth"), + dataset_dir.resolve(), + ) + + +def test_nyu_depth_install_preserves_backups_when_rollback_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Leave recoverable NYU backups outside staging when rollback fails.""" + + monkeypatch.setattr(organizer, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(organizer, "_validate_staged_nyu_depth", lambda _: None) + dataset_dir = tmp_path / "source" + (dataset_dir / "images").mkdir(parents=True) + (dataset_dir / "depth").mkdir() + (dataset_dir / "images" / "sample.jpg").write_bytes(b"new image") + (dataset_dir / "depth" / "sample.npy").write_bytes(b"new depth") + + output_dir = tmp_path / "organized" + (output_dir / "images").mkdir(parents=True) + (output_dir / "depth").mkdir() + (output_dir / "images" / "keep.jpg").write_bytes(b"old image") + (output_dir / "depth" / "keep.npy").write_bytes(b"old depth") + + real_replace = os.replace + + def _fail_install_and_rollback(source: str, destination: str) -> None: + source_path = Path(source) + destination_path = Path(destination) + if ( + source_path.parent.name.startswith(".nyu-depth-staging-") + and source_path.name == "depth" + ): + raise OSError("simulated install failure") + if ( + source_path.parent.name.startswith(".nyu-depth-backup-") + and destination_path == output_dir / "images" + ): + raise OSError("simulated rollback failure") + real_replace(source, destination) + + monkeypatch.setattr( + "mblt_vision.utils.datasets.organizer.os.replace", _fail_install_and_rollback + ) + + with pytest.raises(OSError, match="backups are preserved"): + organizer.construct_nyu_depth(str(dataset_dir), str(output_dir)) + + backup_dirs = list(tmp_path.glob(".nyu-depth-backup-*")) + assert len(backup_dirs) == 1 + assert (backup_dirs[0] / "images" / "keep.jpg").read_bytes() == b"old image" + assert (backup_dirs[0] / "depth" / "keep.npy").read_bytes() == b"old depth" + + +@pytest.mark.parametrize( + "depth", + [ + np.array([[np.nan]], dtype=np.float32), + np.array([[1]], dtype=np.complex64), + np.array([[-1]], dtype=np.float32), + np.array([[0]], dtype=np.float32), + np.array([[100]], dtype=np.float32), + ], +) +def test_staged_nyu_depth_validation_rejects_malformed_payloads( + tmp_path: Path, depth: np.ndarray +) -> None: + """Reject corrupt staged NYU targets before cache replacement can begin.""" + + image_dir = tmp_path / "images" + depth_dir = tmp_path / "depth" + image_dir.mkdir() + depth_dir.mkdir() + Image.new("RGB", (1, 1)).save(image_dir / "sample.png") + np.save(depth_dir / "sample.npy", depth) + + with pytest.raises( + ValueError, + match="real numeric dtype|finite values|negative values|valid metric depth", + ): + organizer._validate_staged_nyu_depth(str(tmp_path)) + + +@pytest.mark.parametrize("dataset", ["ade20k", "dotav1"]) +def test_staged_payload_validation_rejects_corrupt_files( + tmp_path: Path, dataset: str +) -> None: + """Do not replace a valid cache with structurally complete corrupt data.""" + + image_dir = tmp_path / "images" + image_dir.mkdir() + if dataset == "ade20k": + annotation_dir = tmp_path / "annotations" + annotation_dir.mkdir() + (image_dir / "ADE_val_00000001.jpg").write_bytes(b"not an image") + (annotation_dir / "ADE_val_00000001.png").write_bytes(b"not a PNG") + else: + Image.new("RGB", (1, 1)).save(image_dir / "P0001.png") + label_dir = tmp_path / "labels" / "val" + original_label_dir = tmp_path / "labels" / "val_original" + label_dir.mkdir(parents=True) + original_label_dir.mkdir() + (label_dir / "P0001.txt").write_text("bad label", encoding="utf-8") + (original_label_dir / "P0001.txt").write_text("bad label", encoding="utf-8") + + with pytest.raises(ValueError, match="unreadable|Malformed staged DOTAv1"): + organizer._validate_staged_payloads(tmp_path, dataset) + + +def test_staged_cityscapes_payload_validation_rejects_unknown_source_ids( + tmp_path: Path, +) -> None: + """Reject source IDs the Cityscapes loader cannot classify consistently.""" + + image_dir = tmp_path / "images" + annotation_dir = tmp_path / "annotations" + image_dir.mkdir() + annotation_dir.mkdir() + Image.new("RGB", (2, 2)).save(image_dir / "sample.png") + Image.fromarray(np.array([[0, 34], [255, 7]], dtype=np.uint8)).save( + annotation_dir / "sample.png" + ) + + with pytest.raises(ValueError, match="unsupported source IDs.*34"): + organizer._validate_staged_payloads(tmp_path, "cityscapes") + + +@pytest.mark.parametrize( + ("dataset", "source_id"), + [("ade20k", 0), ("cityscapes", 255)], +) +def test_staged_semantic_payload_validation_rejects_all_ignored_targets( + tmp_path: Path, dataset: str, source_id: int +) -> None: + """Do not install semantic masks that cannot contribute to validation metrics.""" + + image_dir = tmp_path / "images" + annotation_dir = tmp_path / "annotations" + image_dir.mkdir() + annotation_dir.mkdir() + Image.new("RGB", (2, 2)).save(image_dir / "sample.png") + Image.fromarray(np.full((2, 2), source_id, dtype=np.uint8)).save( + annotation_dir / "sample.png" + ) + + with pytest.raises(ValueError, match="contains no evaluable class IDs"): + organizer._validate_staged_payloads(tmp_path, dataset) + + +def test_staged_coco_payload_validation_rejects_annotation_geometry_mismatch( + tmp_path: Path, +) -> None: + """Preserve a usable cache when JSON dimensions differ from staged image bytes.""" + + image_dir = tmp_path / "val2017" + image_dir.mkdir() + Image.new("RGB", (2, 1)).save(image_dir / "000000000001.jpg") + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [ + { + "id": 1, + "file_name": "000000000001.jpg", + "height": 2, + "width": 1, + } + ] + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="geometry does not match"): + organizer._validate_staged_payloads(tmp_path, "coco") + + +def _create_ade20k_source(tmp_path: Path) -> Path: + """Create a compact extracted ADE20K validation source.""" + + dataset_dir = tmp_path / "ADEChallengeData2016" + (dataset_dir / "images").mkdir(parents=True) + (dataset_dir / "annotations").mkdir() + (dataset_dir / "images" / "ADE_val_00000001.jpg").write_bytes(b"validation") + (dataset_dir / "annotations" / "ADE_val_00000001.png").write_bytes(b"annotation") + (dataset_dir / "objectInfo150.txt").write_bytes(b"labels") + (dataset_dir / "sceneCategories.txt").write_bytes(b"scenes") + return dataset_dir + + +def test_organize_ade20k_extracts_flat_validation_layout( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Install only ADE20K validation image/mask pairs in the reference layout.""" + + monkeypatch.setattr(organizer, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(readiness_module, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(organizer, "_validate_staged_payloads", lambda *_: None) + archive_path = tmp_path / "ADEChallengeData2016.zip" + image_buffer = io.BytesIO() + Image.new("RGB", (1, 1)).save(image_buffer, format="JPEG") + annotation_buffer = io.BytesIO() + Image.new("L", (1, 1), color=1).save(annotation_buffer, format="PNG") + with ZipFile(archive_path, "w") as archive: + archive.writestr( + "ADEChallengeData2016/images/training/ADE_train_00000001.jpg", b"training" + ) + archive.writestr( + "ADEChallengeData2016/annotations/training/ADE_train_00000001.png", + b"training", + ) + archive.writestr( + "ADEChallengeData2016/images/validation/ADE_val_00000001.jpg", + image_buffer.getvalue(), + ) + archive.writestr( + "ADEChallengeData2016/annotations/validation/ADE_val_00000001.png", + annotation_buffer.getvalue(), + ) + archive.writestr("ADEChallengeData2016/objectInfo150.txt", b"labels") + archive.writestr("ADEChallengeData2016/sceneCategories.txt", b"scenes") + + output_dir = tmp_path / "organized" + organizer.organize_ade20k(str(archive_path), str(output_dir)) + + assert ( + output_dir / "images" / "ADE_val_00000001.jpg" + ).read_bytes() == image_buffer.getvalue() + assert ( + output_dir / "annotations" / "ADE_val_00000001.png" + ).read_bytes() == annotation_buffer.getvalue() + assert (output_dir / "objectInfo150.txt").read_bytes() == b"labels" + assert (output_dir / "sceneCategories.txt").read_bytes() == b"scenes" + assert not (output_dir / "images" / "training").exists() + + +def test_construct_ade20k_requires_metadata_before_replacing_cache( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Preserve the cache when an ADE20K source omits required metadata.""" + + monkeypatch.setattr(organizer, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + dataset_dir = _create_ade20k_source(tmp_path) + (dataset_dir / "sceneCategories.txt").unlink() + output_dir = tmp_path / "organized" + output_dir.mkdir() + (output_dir / "valid-cache-marker").write_bytes(b"existing") + + with pytest.raises(ValueError, match="sceneCategories.txt"): + organizer.construct_ade20k(str(dataset_dir), str(output_dir)) + + assert (output_dir / "valid-cache-marker").read_bytes() == b"existing" + + +def test_construct_ade20k_rejects_duplicate_source_stems( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Reject flat ADE20K source names that would overwrite during staging.""" + + monkeypatch.setattr(organizer, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + dataset_dir = _create_ade20k_source(tmp_path) + image_path = dataset_dir / "images" / "ADE_val_00000001.jpg" + shutil.copy2(image_path, image_path.with_suffix(".JPG")) + + with pytest.raises( + ValueError, match="ADE20K images contain duplicate filename stem" + ): + organizer.construct_ade20k(str(dataset_dir), str(tmp_path / "organized")) + + +@pytest.mark.parametrize( + "relative_path", + [ + "images/ADE_val_00000001.jpg", + "annotations/ADE_val_00000001.png", + "objectInfo150.txt", + "sceneCategories.txt", + ], +) +def test_construct_ade20k_rejects_symlinked_sources( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + relative_path: str, +) -> None: + """Reject ADE20K data and metadata symlinks before replacing a managed cache.""" + + monkeypatch.setattr(organizer, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + dataset_dir = _create_ade20k_source(tmp_path) + external_file = tmp_path / "secret" + external_file.write_bytes(b"outside dataset") + source_path = dataset_dir / relative_path + source_path.unlink() + source_path.symlink_to(external_file) + output_dir = tmp_path / "organized" + output_dir.mkdir() + marker = output_dir / "valid-cache-marker" + marker.write_bytes(b"existing") + + with pytest.raises(ValueError, match="must not be a symlink"): + organizer.construct_ade20k(str(dataset_dir), str(output_dir)) + + assert marker.read_bytes() == b"existing" + + +def test_construct_ade20k_preserves_cache_when_metadata_staging_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Preserve the cache when copying required ADE20K metadata into staging fails.""" + + monkeypatch.setattr(organizer, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + dataset_dir = _create_ade20k_source(tmp_path) + output_dir = tmp_path / "organized" + output_dir.mkdir() + (output_dir / "valid-cache-marker").write_bytes(b"existing") + real_copy2 = shutil.copy2 + + def _fail_scene_metadata_copy(source: str, destination: str) -> str: + if Path(source).name == "sceneCategories.txt": + raise OSError("simulated metadata copy failure") + return real_copy2(source, destination) + + monkeypatch.setattr( + "mblt_vision.utils.datasets.organizer.shutil.copy2", + _fail_scene_metadata_copy, + ) + + with pytest.raises(OSError, match="simulated metadata copy failure"): + organizer.construct_ade20k(str(dataset_dir), str(output_dir)) + + assert (output_dir / "valid-cache-marker").read_bytes() == b"existing" + + +def _write_cityscapes_archives( + tmp_path: Path, sample_ids: list[str] +) -> tuple[Path, Path]: + """Create compact official-layout Cityscapes ZIP fixtures.""" + + image_archive = tmp_path / "leftImg8bit_trainvaltest.zip" + annotation_archive = tmp_path / "gtFine_trainvaltest.zip" + with ZipFile(image_archive, "w") as archive: + archive.writestr( + "leftImg8bit/train/aachen/aachen_000000_000001_leftImg8bit.png", + b"train image", + ) + archive.writestr( + "leftImg8bit/test/berlin/berlin_000000_000001_leftImg8bit.png", + b"test image", + ) + for sample_id in sample_ids: + city = sample_id.split("_", maxsplit=1)[0] + archive.writestr( + f"leftImg8bit/val/{city}/{sample_id}_leftImg8bit.png", + f"image:{sample_id}".encode(), + ) + with ZipFile(annotation_archive, "w") as archive: + archive.writestr( + "gtFine/train/aachen/aachen_000000_000001_gtFine_labelIds.png", + b"train mask", + ) + for sample_id in sample_ids: + city = sample_id.split("_", maxsplit=1)[0] + archive.writestr( + f"gtFine/val/{city}/{sample_id}_gtFine_labelIds.png", + f"mask:{sample_id}".encode(), + ) + archive.writestr( + f"gtFine/val/{city}/{sample_id}_gtFine_color.png", b"color" + ) + archive.writestr( + f"gtFine/val/{city}/{sample_id}_gtFine_instanceIds.png", b"instance" + ) + archive.writestr( + f"gtFine/val/{city}/{sample_id}_gtFine_polygons.json", b"{}" + ) + archive.writestr( + f"gtFine/val/{city}/{sample_id}_gtFine_trainIds.png", b"train IDs" + ) + return image_archive, annotation_archive + + +@pytest.mark.parametrize("dataset", ["nyu-depth", "ade20k", "cityscapes"]) +def test_dense_organizers_reject_symlinked_output_roots( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + dataset: str, +) -> None: + """Fail before organizing through a symlinked managed root or ancestor.""" + + if dataset == "nyu-depth": + monkeypatch.setattr(organizer, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + source_dir = tmp_path / "nyu-source" + (source_dir / "images").mkdir(parents=True) + (source_dir / "depth").mkdir() + (source_dir / "images" / "sample.jpg").write_bytes(b"image") + (source_dir / "depth" / "sample.npy").write_bytes(b"depth") + + def organizer_fn(output_dir: str) -> None: + organizer.organize_nyu_depth(str(source_dir), output_dir) + elif dataset == "ade20k": + monkeypatch.setattr(organizer, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + source_dir = _create_ade20k_source(tmp_path) + + def organizer_fn(output_dir: str) -> None: + organizer.organize_ade20k(str(source_dir), output_dir) + else: + monkeypatch.setattr(organizer, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 1) + image_archive, annotation_archive = _write_cityscapes_archives( + tmp_path, + ["lindau_000000_000019"], + ) + + def organizer_fn(output_dir: str) -> None: + organizer.organize_cityscapes( + str(image_archive), + str(annotation_archive), + output_dir, + ) + + for topology in ( + "root", + "ancestor", + "normalized-ancestor", + "symlink-parent-traversal", + ): + if topology == "root": + protected_dir = tmp_path / "root-target" + protected_dir.mkdir() + output_dir = tmp_path / "managed" + output_dir.symlink_to(protected_dir, target_is_directory=True) + elif topology == "ancestor": + target_parent = tmp_path / "ancestor-target" + protected_dir = target_parent / "managed" + protected_dir.mkdir(parents=True) + symlinked_parent = tmp_path / "datasets-link" + symlinked_parent.symlink_to(target_parent, target_is_directory=True) + output_dir = symlinked_parent / "managed" + elif topology == "normalized-ancestor": + target_parent = tmp_path / "normalized-ancestor-target" + protected_dir = target_parent / "managed" + protected_dir.mkdir(parents=True) + symlinked_parent = tmp_path / "normalized-datasets-link" + symlinked_parent.symlink_to(target_parent, target_is_directory=True) + output_dir = tmp_path / "missing" / ".." / symlinked_parent.name / "managed" + else: + target_parent = tmp_path / "traversal-target" + target_child = target_parent / "child" + target_child.mkdir(parents=True) + protected_dir = target_parent / "traversed-managed" + protected_dir.mkdir() + symlinked_parent = tmp_path / "traversal-link" + symlinked_parent.symlink_to(target_child, target_is_directory=True) + output_dir = symlinked_parent / ".." / protected_dir.name + marker = protected_dir / "keep" + marker.write_bytes(b"existing") + + with pytest.raises(ValueError, match="existing parents must not be symlinks"): + organizer_fn(str(output_dir)) + + assert marker.read_bytes() == b"existing" + assert list(protected_dir.iterdir()) == [marker] + + +@pytest.mark.parametrize( + ("dataset", "layout_name"), + [ + ("nyu-depth", "images"), + ("nyu-depth", "depth"), + ("ade20k", "images"), + ("ade20k", "annotations"), + ("cityscapes", "images"), + ("cityscapes", "annotations"), + ], +) +def test_dense_organizers_reject_symlinked_output_layout_directories( + tmp_path: Path, + dataset: str, + layout_name: str, +) -> None: + """Reject symlinked managed layout children without modifying their targets.""" + + output_dir = tmp_path / "managed" + output_dir.mkdir() + protected_dir = tmp_path / "protected-layout" + protected_dir.mkdir() + marker = protected_dir / "keep" + marker.write_bytes(b"existing") + (output_dir / layout_name).symlink_to(protected_dir, target_is_directory=True) + + with pytest.raises(ValueError, match="layout directories must not be symlinks"): + if dataset == "nyu-depth": + organizer.organize_nyu_depth("unused", str(output_dir)) + elif dataset == "ade20k": + organizer.organize_ade20k("unused", str(output_dir)) + else: + organizer.organize_cityscapes( + "unused-images", "unused-annotations", str(output_dir) + ) + + assert marker.read_bytes() == b"existing" + assert list(protected_dir.iterdir()) == [marker] + assert (output_dir / layout_name).is_symlink() + + +def test_organize_cityscapes_materializes_lossless_validation_pairs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Copy only exactly paired validation images and label-ID masks without transcoding.""" + + monkeypatch.setattr(organizer, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 2) + monkeypatch.setattr(organizer, "dataset_ready", lambda *_: True) + monkeypatch.setattr(organizer, "_validate_staged_payloads", lambda *_: None) + sample_ids = ["frankfurt_000000_000294", "munster_000001_000019"] + image_archive, annotation_archive = _write_cityscapes_archives(tmp_path, sample_ids) + output_dir = tmp_path / "cityscapes" + (output_dir / "images").mkdir(parents=True) + (output_dir / "annotations").mkdir() + (output_dir / "images" / "stale.png").write_bytes(b"stale") + (output_dir / "annotations" / "stale.png").write_bytes(b"stale") + + organizer.organize_cityscapes( + str(image_archive), str(annotation_archive), str(output_dir) + ) + + image_paths = sorted((output_dir / "images").glob("*.png")) + annotation_paths = sorted((output_dir / "annotations").glob("*.png")) + assert [path.name for path in image_paths] == [ + f"{sample_id}.png" for sample_id in sample_ids + ] + assert [path.name for path in annotation_paths] == [ + f"{sample_id}.png" for sample_id in sample_ids + ] + assert image_paths[0].read_bytes() == f"image:{sample_ids[0]}".encode() + assert annotation_paths[0].read_bytes() == f"mask:{sample_ids[0]}".encode() + assert not list(output_dir.rglob("*train*")) + assert not list(output_dir.rglob("*color*")) + assert not list(output_dir.rglob("*instance*")) + assert not list(output_dir.rglob("*.json")) + + +def test_organize_cityscapes_enforces_validation_pair_count( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Reject incomplete validation sources before replacing existing data.""" + + monkeypatch.setattr(organizer, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 2) + image_archive, annotation_archive = _write_cityscapes_archives( + tmp_path, ["lindau_000000_000019"] + ) + output_dir = tmp_path / "cityscapes" + (output_dir / "images").mkdir(parents=True) + marker = output_dir / "images" / "keep.png" + marker.write_bytes(b"keep") + + with pytest.raises(ValueError, match="must contain 2 pairs"): + organizer.organize_cityscapes( + str(image_archive), str(annotation_archive), str(output_dir) + ) + + assert marker.read_bytes() == b"keep" + + +def test_organize_cityscapes_rejects_mismatched_and_malformed_sources( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Require exact official stems and reject malformed validation candidates.""" + + monkeypatch.setattr(organizer, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 1) + image_archive, annotation_archive = _write_cityscapes_archives( + tmp_path, ["lindau_000000_000019"] + ) + with ZipFile(annotation_archive, "w") as archive: + archive.writestr( + "gtFine/val/lindau/lindau_000000_000020_gtFine_labelIds.png", b"mask" + ) + with pytest.raises(ValueError, match="mismatch"): + organizer.organize_cityscapes( + str(image_archive), str(annotation_archive), str(tmp_path / "mismatched") + ) + + with ZipFile(annotation_archive, "w") as archive: + archive.writestr( + "gtFine/val/lindau/frankfurt_000000_000019_gtFine_labelIds.png", b"mask" + ) + with pytest.raises(ValueError, match="Malformed Cityscapes annotation filename"): + organizer.organize_cityscapes( + str(image_archive), str(annotation_archive), str(tmp_path / "malformed") + ) + + +def test_organize_cityscapes_rejects_non_zip_duplicate_and_unsafe_inputs( + tmp_path: Path, +) -> None: + """Reject invalid ZIPs, duplicate members, and traversal paths before installation.""" + + invalid_archive = tmp_path / "invalid.zip" + invalid_archive.write_bytes(b"not a zip") + valid_image_archive, valid_annotation_archive = _write_cityscapes_archives( + tmp_path, + ["lindau_000000_000019"], + ) + with pytest.raises(ValueError, match="does not exist"): + organizer.organize_cityscapes( + str(tmp_path / "missing.zip"), + str(valid_annotation_archive), + str(tmp_path / "missing"), + ) + with pytest.raises(ValueError, match="must be a valid ZIP"): + organizer.organize_cityscapes( + str(invalid_archive), + str(valid_annotation_archive), + str(tmp_path / "invalid"), + ) + + duplicate_archive = tmp_path / "duplicate.zip" + duplicate_member = "leftImg8bit/val/lindau/lindau_000000_000019_leftImg8bit.png" + with pytest.warns(UserWarning, match="Duplicate name"): + with ZipFile(duplicate_archive, "w") as archive: + archive.writestr(duplicate_member, b"first") + archive.writestr(duplicate_member, b"second") + with pytest.raises(ValueError, match="duplicate members"): + organizer.organize_cityscapes( + str(duplicate_archive), + str(valid_annotation_archive), + str(tmp_path / "duplicate"), + ) + + unsafe_archive = tmp_path / "unsafe.zip" + outside_marker = tmp_path / "outside.png" + with ZipFile(unsafe_archive, "w") as archive: + archive.writestr("../outside.png", b"outside") + with pytest.raises(ValueError, match="Unsafe archive member path"): + organizer.organize_cityscapes( + str(unsafe_archive), str(valid_annotation_archive), str(tmp_path / "unsafe") + ) + assert not outside_marker.exists() + + +def test_organize_cityscapes_rolls_back_failed_atomic_replacement( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Restore both previous directories if the staged installation fails.""" + + monkeypatch.setattr(organizer, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(organizer, "dataset_ready", lambda *_: True) + monkeypatch.setattr(organizer, "_validate_staged_payloads", lambda *_: None) + image_archive, annotation_archive = _write_cityscapes_archives( + tmp_path, ["lindau_000000_000019"] + ) + output_dir = tmp_path / "cityscapes" + (output_dir / "images").mkdir(parents=True) + (output_dir / "annotations").mkdir() + (output_dir / "images" / "keep.png").write_bytes(b"old image") + (output_dir / "annotations" / "keep.png").write_bytes(b"old annotation") + real_replace = os.replace + failed = False + + def _fail_annotation_install(source: str, destination: str) -> None: + nonlocal failed + if ( + not failed + and Path(source).name == "annotations" + and Path(destination) == output_dir / "annotations" + ): + failed = True + raise OSError("simulated install failure") + real_replace(source, destination) + + monkeypatch.setattr( + "mblt_vision.utils.datasets.organizer.os.replace", _fail_annotation_install + ) + + with pytest.raises(OSError, match="simulated"): + organizer.organize_cityscapes( + str(image_archive), str(annotation_archive), str(output_dir) + ) + + assert (output_dir / "images" / "keep.png").read_bytes() == b"old image" + assert (output_dir / "annotations" / "keep.png").read_bytes() == b"old annotation" + + +def test_organize_cityscapes_preserves_cache_when_staging_fails_readiness( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Validate the staged tree before it can replace a usable Cityscapes cache.""" + + monkeypatch.setattr(organizer, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(organizer, "dataset_ready", lambda *_: False) + image_archive, annotation_archive = _write_cityscapes_archives( + tmp_path, ["lindau_000000_000019"] + ) + output_dir = tmp_path / "cityscapes" + (output_dir / "images").mkdir(parents=True) + (output_dir / "annotations").mkdir() + (output_dir / "images" / "keep.png").write_bytes(b"old image") + (output_dir / "annotations" / "keep.png").write_bytes(b"old annotation") + + with pytest.raises(ValueError, match="failed identity and completeness checks"): + organizer.organize_cityscapes( + str(image_archive), str(annotation_archive), str(output_dir) + ) + + assert (output_dir / "images" / "keep.png").read_bytes() == b"old image" + assert (output_dir / "annotations" / "keep.png").read_bytes() == b"old annotation" + + +def test_dense_install_preserves_backups_when_rollback_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Keep recoverable backups outside staging when rollback cannot finish.""" + + staging_dir = tmp_path / ".staging" + staged_images = staging_dir / "images" + staged_annotations = staging_dir / "annotations" + staged_images.mkdir(parents=True) + staged_annotations.mkdir() + (staged_images / "new.png").write_bytes(b"new image") + (staged_annotations / "new.png").write_bytes(b"new annotation") + + output_dir = tmp_path / "organized" + (output_dir / "images").mkdir(parents=True) + (output_dir / "annotations").mkdir() + (output_dir / "images" / "keep.png").write_bytes(b"old image") + (output_dir / "annotations" / "keep.png").write_bytes(b"old annotation") + + real_replace = os.replace + + def _fail_install_and_rollback(source: str, destination: str) -> None: + source_path = Path(source) + destination_path = Path(destination) + if source_path == staged_annotations: + raise OSError("simulated install failure") + if ( + source_path.parent.name.startswith(".dense-backup-") + and destination_path == output_dir / "images" + ): + raise OSError("simulated rollback failure") + real_replace(source, destination) + + monkeypatch.setattr( + "mblt_vision.utils.datasets.organizer.os.replace", _fail_install_and_rollback + ) + + replacements = ( + (str(staged_images), str(output_dir / "images")), + (str(staged_annotations), str(output_dir / "annotations")), + ) + with pytest.raises(OSError, match="backups are preserved"): + organizer._replace_staged_directories( + replacements, str(tmp_path), ".dense-backup-" + ) + + backup_dirs = list(tmp_path.glob(".dense-backup-*")) + assert len(backup_dirs) == 1 + assert (backup_dirs[0] / "images" / "keep.png").read_bytes() == b"old image" + assert ( + backup_dirs[0] / "annotations" / "keep.png" + ).read_bytes() == b"old annotation" diff --git a/tests/test_dataset_readiness.py b/tests/test_dataset_readiness.py new file mode 100644 index 0000000..bf2737f --- /dev/null +++ b/tests/test_dataset_readiness.py @@ -0,0 +1,981 @@ +"""Tests for organized vision dataset identity and completeness checks.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image +from scipy.io import savemat + +from mblt_vision.utils.datasets import readiness + + +def _write_file(path: Path) -> None: + """Create a placeholder dataset file and its parent directories.""" + + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"data") + + +def _write_image(path: Path, size: tuple[int, int]) -> None: + """Create a decodable RGB image with the requested ``(width, height)``.""" + + path.parent.mkdir(parents=True, exist_ok=True) + Image.new("RGB", size).save(path) + + +def _write_widerface_metadata(path: Path, event_images: dict[str, list[str]]) -> None: + """Write the WiderFace event and image-name cell arrays used by readiness.""" + + event_list = np.empty((len(event_images), 1), dtype=object) + file_list = np.empty((len(event_images), 1), dtype=object) + for index, (event_name, image_stems) in enumerate(event_images.items()): + event_list[index, 0] = event_name + file_list[index, 0] = np.array([[stem] for stem in image_stems], dtype=object) + savemat(path, {"event_list": event_list, "file_list": file_list}) + + +def test_imagenet_readiness_requires_complete_official_class_tree( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Reject incomplete or non-ImageNet classification directory trees.""" + + monkeypatch.setattr(readiness, "IMAGENET_CLASS_COUNT", 2) + monkeypatch.setattr(readiness, "IMAGENET_IMAGES_PER_CLASS", 2) + monkeypatch.setattr(readiness, "IMAGENET_SYNSETS", {"n00000000", "n00000001"}) + for class_index in range(2): + for image_index in range(2): + if (class_index, image_index) == (1, 1): + continue + _write_file( + tmp_path + / f"n{class_index:08d}" + / f"ILSVRC2012_val_{class_index * 2 + image_index + 1:08d}.JPEG" + ) + + assert not readiness.dataset_ready(tmp_path, "image_classification", "imagenet") + + _write_file(tmp_path / "n00000001" / "ILSVRC2012_val_00000004.JPEG") + + assert readiness.dataset_ready(tmp_path, "image_classification", "imagenet") + assert not readiness.dataset_ready(tmp_path, "image_classification", "coco") + + +@pytest.mark.parametrize( + ("task", "annotation_name"), + [ + ("object_detection", "instances_val2017.json"), + ("instance_segmentation", "instances_val2017.json"), + ("pose_estimation", "person_keypoints_val2017.json"), + ], +) +def test_coco_readiness_matches_images_to_task_annotations( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + task: str, + annotation_name: str, +) -> None: + """Require every official COCO image in the task-specific annotation file.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 2) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, annotation_name, 2) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, annotation_name, 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + image_names = ["000000000001.jpg", "000000000002.jpg"] + for image_name in image_names: + _write_image(tmp_path / "val2017" / image_name, (2, 2)) + (tmp_path / annotation_name).write_text( + json.dumps({"images": [{"id": 1, "file_name": image_names[0]}]}), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, task, "coco") + + (tmp_path / annotation_name).write_text( + json.dumps( + { + "images": [ + { + "id": index, + "file_name": image_name, + "height": 2, + "width": 2, + } + for index, image_name in enumerate(image_names, start=1) + ], + "categories": [{"id": 1}], + "annotations": [ + { + "id": index, + "image_id": index, + "category_id": 1, + "bbox": [0, 0, 2, 2], + "area": 4, + "iscrowd": 0, + **( + {"segmentation": [[0, 0, 2, 0, 2, 2]]} + if task == "instance_segmentation" + else {} + ), + **( + {"keypoints": [0, 0, 0] * 17, "num_keypoints": 0} + if task == "pose_estimation" + else {} + ), + } + for index in range(1, 3) + ], + } + ), + encoding="utf-8", + ) + + assert readiness.dataset_ready(tmp_path, task, "coco") + assert not readiness.dataset_ready(tmp_path, task, "imagenet") + + +@pytest.mark.parametrize( + ("task", "invalid_field"), + [ + ("instance_segmentation", "segmentation"), + ("pose_estimation", "keypoints"), + ("pose_estimation", "num_keypoints"), + ], +) +def test_coco_readiness_rejects_missing_task_specific_ground_truth( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + task: str, + invalid_field: str, +) -> None: + """Reject COCO metadata missing the payload required by its evaluation task.""" + + annotation_name = ( + "person_keypoints_val2017.json" + if task == "pose_estimation" + else "instances_val2017.json" + ) + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, annotation_name, 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, annotation_name, 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + _write_image(tmp_path / "val2017" / "000000000001.jpg", (2, 2)) + annotation = { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + "segmentation": [[0, 0, 1, 0, 1, 1]], + "keypoints": [0, 0, 0] * 17, + "num_keypoints": 0, + } + del annotation[invalid_field] + (tmp_path / annotation_name).write_text( + json.dumps( + { + "images": [{"id": 1, "file_name": "000000000001.jpg"}], + "categories": [{"id": 1}], + "annotations": [annotation], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, task, "coco") + + +def test_coco_readiness_rejects_noncanonical_category_ids( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Require the canonical COCO category-ID set, not merely its count.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + image_name = "000000000001.jpg" + _write_file(tmp_path / "val2017" / image_name) + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [{"id": 1, "file_name": image_name}], + "categories": [{"id": 2}], + "annotations": [{"id": 1, "image_id": 1, "category_id": 2}], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("segmentation", [[0, 0, 1, 1, 2, 2]]), + ("area", 0), + ("iscrowd", 2), + ], +) +def test_coco_readiness_rejects_invalid_evaluation_metadata( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + field: str, + value: object, +) -> None: + """Reject metadata that can alter COCO matching or area-range evaluation.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + _write_file(tmp_path / "val2017" / "000000000001.jpg") + annotation = { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + "segmentation": [[0, 0, 1, 0, 1, 1]], + } + annotation[field] = value + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [{"id": 1, "file_name": "000000000001.jpg"}], + "categories": [{"id": 1}], + "annotations": [annotation], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "instance_segmentation", "coco") + + +@pytest.mark.parametrize("bbox", [[0, 0, 0, 1], [0, 0, 1, float("nan")]]) +def test_coco_readiness_rejects_invalid_instance_box_geometry( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bbox: list[float] +) -> None: + """Reject invalid instance boxes before a malformed cache reaches COCOeval.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + _write_file(tmp_path / "val2017" / "000000000001.jpg") + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [{"id": 1, "file_name": "000000000001.jpg"}], + "categories": [{"id": 1}], + "annotations": [ + {"id": 1, "image_id": 1, "category_id": 1, "bbox": bbox} + ], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + +@pytest.mark.parametrize( + ("bbox", "area"), + [([3, 0, 1, 1], 1), ([0, 0, 1, 1], 5)], +) +def test_coco_readiness_rejects_unrealizable_box_or_area( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + bbox: list[int], + area: int, +) -> None: + """Boxes and areas must describe foreground that fits the source image.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + image_name = "000000000001.jpg" + _write_image(tmp_path / "val2017" / image_name, (2, 2)) + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [{"id": 1, "file_name": image_name, "height": 2, "width": 2}], + "categories": [{"id": 1}], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": bbox, + "area": area, + "iscrowd": 0, + } + ], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + +@pytest.mark.parametrize( + ("counts", "expected_ready"), + [([3, 1], True), ("31", True), ([4], False), ([3], False), ("4", False)], +) +def test_coco_readiness_decodes_and_validates_rle_segmentations( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + counts: list[int] | str, + expected_ready: bool, +) -> None: + """Require nonempty, complete RLE masks with the referenced image geometry.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + _write_image(tmp_path / "val2017" / "000000000001.jpg", (2, 2)) + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [ + { + "id": 1, + "file_name": "000000000001.jpg", + "height": 2, + "width": 2, + } + ], + "categories": [{"id": 1}], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + "segmentation": {"size": [2, 2], "counts": counts}, + } + ], + } + ), + encoding="utf-8", + ) + + assert ( + readiness.dataset_ready(tmp_path, "instance_segmentation", "coco") + is expected_ready + ) + + +def test_coco_readiness_rejects_corrupt_or_mismatched_images( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Cache readiness must decode COCO images and match annotation geometry.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 1) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + _write_image(tmp_path / "val2017" / "000000000001.jpg", (1, 1)) + annotation_path = tmp_path / "instances_val2017.json" + annotation_path.write_text( + json.dumps( + { + "images": [ + { + "id": 1, + "file_name": "000000000001.jpg", + "height": 2, + "width": 2, + } + ], + "categories": [{"id": 1}], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + } + ], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + _write_file(tmp_path / "val2017" / "000000000001.jpg") + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + +def test_nyu_readiness_matches_the_evaluator_depth_domain( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Cache reuse must reject depth targets that evaluation cannot score.""" + + monkeypatch.setattr(readiness, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + _write_image(tmp_path / "images" / "sample.png", (2, 2)) + depth_path = tmp_path / "depth" / "sample.npy" + depth_path.parent.mkdir() + + np.save(depth_path, np.full((2, 2), 0.001, dtype=np.float64)) + assert not readiness.dense_dataset_ready(tmp_path, "nyu-depth") + + np.save( + depth_path, + np.array([[1.0, np.finfo(np.float64).max], [1.0, 1.0]], dtype=np.float64), + ) + assert not readiness.dense_dataset_ready(tmp_path, "nyu-depth") + + np.save(depth_path, np.ones((2, 2), dtype=np.float32)) + assert readiness.dense_dataset_ready(tmp_path, "nyu-depth") + + +def test_cityscapes_readiness_rejects_unsupported_source_ids( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Source IDs outside the Cityscapes label taxonomy cannot hide as ignore.""" + + monkeypatch.setattr(readiness, "CITYSCAPES_VALIDATION_SAMPLE_COUNT", 1) + stem = "frankfurt_000000_000000" + _write_image(tmp_path / "images" / f"{stem}.png", (2, 2)) + annotation_path = tmp_path / "annotations" / f"{stem}.png" + annotation_path.parent.mkdir() + Image.new("L", (2, 2), color=34).save(annotation_path) + + assert not readiness.dense_dataset_ready(tmp_path, "cityscapes") + + Image.new("L", (2, 2), color=7).save(annotation_path) + assert readiness.dense_dataset_ready(tmp_path, "cityscapes") + + +def test_coco_annotation_identity_digest_rejects_wrong_validation_split( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Bind a complete COCO annotation table to its canonical image identities.""" + + images = [ + {"id": index, "file_name": f"{index:012d}.jpg", "height": 1, "width": 1} + for index in range(1, 5001) + ] + payload = "".join(f"{item['id']}:{item['file_name']}\n" for item in images).encode() + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 5000) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 0) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + monkeypatch.setattr( + readiness, + "COCO_VALIDATION_IMAGE_IDENTITIES_SHA256", + hashlib.sha256(payload).hexdigest(), + ) + annotation_path = tmp_path / "instances_val2017.json" + annotation_path.write_text( + json.dumps({"images": images, "categories": [{"id": 1}], "annotations": []}), + encoding="utf-8", + ) + + assert readiness._load_coco_image_names(annotation_path) is not None + images[0]["file_name"] = "999999999999.jpg" + annotation_path.write_text( + json.dumps({"images": images, "categories": [{"id": 1}], "annotations": []}), + encoding="utf-8", + ) + assert readiness._load_coco_image_names(annotation_path) is None + + +def test_coco_readiness_rejects_duplicate_image_ids( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Reject COCO metadata whose duplicate IDs would overwrite image records.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 2) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 2) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + image_names = ["000000000001.jpg", "000000000002.jpg"] + for image_name in image_names: + _write_file(tmp_path / "val2017" / image_name) + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [ + {"id": 1, "file_name": image_names[0]}, + {"id": 1, "file_name": image_names[1]}, + ], + "categories": [{"id": 1}], + "annotations": [ + {"id": 1, "image_id": 1, "category_id": 1}, + {"id": 2, "image_id": 1, "category_id": 1}, + ], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + +def test_coco_readiness_rejects_truncated_annotation_table( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not reuse complete-looking images with incomplete COCO ground truth.""" + + monkeypatch.setattr(readiness, "COCO_VALIDATION_SAMPLE_COUNT", 2) + monkeypatch.setitem(readiness.COCO_ANNOTATION_COUNTS, "instances_val2017.json", 2) + monkeypatch.setitem(readiness.COCO_CATEGORY_COUNTS, "instances_val2017.json", 1) + monkeypatch.setattr(readiness, "COCO_CATEGORY_IDS", frozenset({1})) + image_names = ["000000000001.jpg", "000000000002.jpg"] + for image_name in image_names: + _write_file(tmp_path / "val2017" / image_name) + (tmp_path / "instances_val2017.json").write_text( + json.dumps( + { + "images": [ + {"id": index, "file_name": image_name} + for index, image_name in enumerate(image_names, start=1) + ], + "categories": [{"id": 1}], + "annotations": [{"id": 1, "image_id": 1, "category_id": 1}], + } + ), + encoding="utf-8", + ) + + assert not readiness.dataset_ready(tmp_path, "object_detection", "coco") + + +def test_widerface_readiness_rejects_invalid_difficulty_metadata( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Reject present-but-malformed difficulty files before evaluation indexes them.""" + + monkeypatch.setattr(readiness, "WIDERFACE_EVENT_COUNT", 1) + monkeypatch.setattr(readiness, "WIDERFACE_VALIDATION_SAMPLE_COUNT", 1) + _write_widerface_metadata( + tmp_path / "wider_face_val.mat", {"0--Parade": ["sample"]} + ) + for file_name in ( + "wider_easy_val.mat", + "wider_medium_val.mat", + "wider_hard_val.mat", + ): + _write_file(tmp_path / file_name) + _write_file(tmp_path / "images" / "0--Parade" / "sample.jpg") + + assert not readiness.dataset_ready(tmp_path, "face_detection", "widerface") + + +def test_widerface_readiness_rejects_empty_face_ground_truth( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not accept empty face arrays that evaluation would silently skip.""" + + face_boxes = np.empty((1, 1), dtype=object) + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.empty((0, 4), dtype=np.float64) + face_boxes[0, 0] = event_faces + difficulty = np.empty((1, 1), dtype=object) + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = np.empty((0,), dtype=np.int64) + difficulty[0, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert not readiness._widerface_difficulty_metadata_ready( + tmp_path, {"0--Parade": {"sample.jpg"}} + ) + + +def test_widerface_readiness_rejects_duplicate_difficulty_indices( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Difficulty masks must identify each eligible face at most once.""" + + face_boxes = np.empty((1, 1), dtype=object) + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.array([[0, 0, 1, 1]], dtype=np.float64) + face_boxes[0, 0] = event_faces + difficulty = np.empty((1, 1), dtype=object) + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = np.array([1, 1], dtype=np.int64) + difficulty[0, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert not readiness._widerface_difficulty_metadata_ready( + tmp_path, {"0--Parade": {"sample.jpg"}} + ) + + +def test_widerface_readiness_rejects_row_vector_difficulty_indices( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Reject row vectors that the evaluator would count as one eligible face.""" + + face_boxes = np.empty((1, 1), dtype=object) + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.array([[0, 0, 1, 1], [1, 1, 1, 1]], dtype=np.float64) + face_boxes[0, 0] = event_faces + difficulty = np.empty((1, 1), dtype=object) + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = np.array([[1, 2]], dtype=np.int64) + difficulty[0, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert not readiness._widerface_difficulty_metadata_ready( + tmp_path, {"0--Parade": {"sample.jpg"}} + ) + + +def test_widerface_readiness_rejects_all_empty_difficulty_indices( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Every complete WiderFace difficulty split needs an eligible face.""" + + face_boxes = np.empty((1, 1), dtype=object) + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.array([[0, 0, 1, 1]], dtype=np.float64) + face_boxes[0, 0] = event_faces + difficulty = np.empty((1, 1), dtype=object) + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = np.empty((0, 0), dtype=np.int64) + difficulty[0, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert not readiness._widerface_difficulty_metadata_ready( + tmp_path, {"0--Parade": {"sample.jpg"}} + ) + + +def test_widerface_readiness_rejects_face_boxes_outside_images( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Face ground truth must retain foreground in its decoded source image.""" + + face_boxes = np.empty((1, 1), dtype=object) + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.array([[11, 0, 1, 1]], dtype=np.float64) + face_boxes[0, 0] = event_faces + difficulty = np.empty((1, 1), dtype=object) + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = np.array([1], dtype=np.int64) + difficulty[0, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert not readiness._widerface_difficulty_metadata_ready( + tmp_path, + {"0--Parade": {"sample.jpg"}}, + image_shapes=[[(10, 10)]], + ) + + +def test_widerface_readiness_allows_empty_event_difficulty_contributions( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A difficulty split may have no eligible faces in one event but not all.""" + + face_boxes = np.empty((2, 1), dtype=object) + difficulty = np.empty((2, 1), dtype=object) + for event_index in range(2): + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.array([[0, 0, 1, 1]], dtype=np.float64) + face_boxes[event_index, 0] = event_faces + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = ( + np.empty((0, 0), dtype=np.int64) + if event_index == 0 + else np.array([1], dtype=np.int64) + ) + difficulty[event_index, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert readiness._widerface_difficulty_metadata_ready( + tmp_path, + {"0--Parade": {"first.jpg"}, "1--Handshaking": {"second.jpg"}}, + ) + + +def test_widerface_readiness_rejects_duplicate_face_boxes( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Duplicate ground-truth rows cannot alter the WiderFace recall denominator.""" + + face_boxes = np.empty((1, 1), dtype=object) + event_faces = np.empty((1, 1), dtype=object) + event_faces[0, 0] = np.array([[0, 0, 1, 1], [0, 0, 1, 1]], dtype=np.float64) + face_boxes[0, 0] = event_faces + difficulty = np.empty((1, 1), dtype=object) + event_indices = np.empty((1, 1), dtype=object) + event_indices[0, 0] = np.array([1, 2], dtype=np.int64) + difficulty[0, 0] = event_indices + + def _loadmat(path: Path) -> dict[str, np.ndarray]: + if path.name == "wider_face_val.mat": + return {"face_bbx_list": face_boxes} + return {"gt_list": difficulty} + + monkeypatch.setattr(readiness, "loadmat", _loadmat) + + assert not readiness._widerface_difficulty_metadata_ready( + tmp_path, {"0--Parade": {"sample.jpg"}} + ) + + +@pytest.mark.parametrize("relative_image_dir", ["images", "images/val"]) +def test_dotav1_readiness_requires_complete_image_label_pairs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + relative_image_dir: str, +) -> None: + """Accept flat and legacy DOTA images only when every image has a label.""" + + monkeypatch.setattr(readiness, "DOTAV1_VALIDATION_SAMPLE_COUNT", 2) + for stem in ("P0001", "P0002"): + _write_file(tmp_path / relative_image_dir / f"{stem}.png") + _write_file(tmp_path / "labels" / "val_original" / "P0001.txt") + + assert not readiness.dataset_ready(tmp_path, "obb", "dotav1") + + _write_file(tmp_path / "labels" / "val" / "P0002.txt") + + assert readiness.dataset_ready(tmp_path, "obb", "dotav1") + + external_image = tmp_path / "external.png" + external_label = tmp_path / "external.txt" + _write_file(external_image) + _write_file(external_label) + image_path = tmp_path / relative_image_dir / "P0001.png" + label_path = tmp_path / "labels" / "val" / "P0002.txt" + image_path.unlink() + label_path.unlink() + image_path.symlink_to(external_image) + label_path.symlink_to(external_label) + + assert readiness.dataset_ready(tmp_path, "obb", "dotav1") + + +def test_widerface_readiness_requires_complete_event_tree_and_metadata( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Require all WiderFace validation events, images, and evaluation metadata.""" + + monkeypatch.setattr(readiness, "WIDERFACE_EVENT_COUNT", 2) + monkeypatch.setattr(readiness, "WIDERFACE_VALIDATION_SAMPLE_COUNT", 2) + monkeypatch.setattr( + readiness, "_widerface_difficulty_metadata_ready", lambda *_, **__: True + ) + monkeypatch.setattr( + readiness, "_widerface_image_shapes", lambda *_: [[(1, 1)], [(1, 1)]] + ) + _write_widerface_metadata( + tmp_path / "wider_face_val.mat", + {"0--Parade": ["sample-0"], "1--Handshaking": ["sample-1"]}, + ) + for file_name in ( + "wider_easy_val.mat", + "wider_medium_val.mat", + "wider_hard_val.mat", + ): + _write_file(tmp_path / file_name) + _write_file(tmp_path / "images" / "0--Parade" / "sample-0.jpg") + (tmp_path / "images" / "1--Handshaking").mkdir() + + assert not readiness.dataset_ready(tmp_path, "face_detection", "widerface") + + _write_file(tmp_path / "images" / "1--Handshaking" / "sample-1.jpg") + + assert readiness.dataset_ready(tmp_path, "face_detection", "widerface") + assert not readiness.dataset_ready(tmp_path, "face_detection", "coco") + + +@pytest.mark.parametrize( + ("event_name", "image_name"), + [ + ("0--Parade", "stale"), + ("1--Handshaking", "expected"), + ], +) +def test_widerface_readiness_rejects_tree_not_named_by_metadata( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + event_name: str, + image_name: str, +) -> None: + """Reject complete-looking trees whose event or image identity differs from metadata.""" + + monkeypatch.setattr(readiness, "WIDERFACE_EVENT_COUNT", 1) + monkeypatch.setattr(readiness, "WIDERFACE_VALIDATION_SAMPLE_COUNT", 1) + _write_widerface_metadata( + tmp_path / "wider_face_val.mat", {"0--Parade": ["expected"]} + ) + for file_name in ( + "wider_easy_val.mat", + "wider_medium_val.mat", + "wider_hard_val.mat", + ): + _write_file(tmp_path / file_name) + _write_file(tmp_path / "images" / event_name / f"{image_name}.jpg") + + assert not readiness.dataset_ready(tmp_path, "face_detection", "widerface") + + +def test_ade20k_readiness_requires_source_metadata( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Require both ADE20K metadata files before reusing an organized cache.""" + + monkeypatch.setattr(readiness, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + _write_image(tmp_path / "images" / "ADE_val_00000001.jpg", (1, 1)) + annotation_path = tmp_path / "annotations" / "ADE_val_00000001.png" + annotation_path.parent.mkdir(parents=True, exist_ok=True) + Image.new("L", (1, 1), color=1).save(annotation_path) + + assert not readiness.dataset_ready(tmp_path, "semantic_segmentation", "ade20k") + + _write_file(tmp_path / "objectInfo150.txt") + assert not readiness.dataset_ready(tmp_path, "semantic_segmentation", "ade20k") + + _write_file(tmp_path / "sceneCategories.txt") + assert readiness.dataset_ready(tmp_path, "semantic_segmentation", "ade20k") + + +@pytest.mark.parametrize( + ("dataset", "task", "relative_path"), + [ + ("nyu-depth", "depth_estimation", "images/sample.jpg"), + ("nyu-depth", "depth_estimation", "depth/sample.npy"), + ("nyu-depth", "depth_estimation", "images/extra.jpg"), + ("ade20k", "semantic_segmentation", "images/ADE_val_00000001.jpg"), + ("ade20k", "semantic_segmentation", "annotations/ADE_val_00000001.png"), + ("ade20k", "semantic_segmentation", "objectInfo150.txt"), + ], +) +def test_dense_readiness_rejects_symlinked_files( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + dataset: str, + task: str, + relative_path: str, +) -> None: + """Do not reuse a complete-looking dense cache containing symlinked files.""" + + monkeypatch.setattr(readiness, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(readiness, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + if dataset == "nyu-depth": + _write_file(tmp_path / "images" / "sample.jpg") + _write_file(tmp_path / "depth" / "sample.npy") + else: + _write_file(tmp_path / "images" / "ADE_val_00000001.jpg") + _write_file(tmp_path / "annotations" / "ADE_val_00000001.png") + for file_name in readiness.ADE20K_METADATA_FILES: + _write_file(tmp_path / file_name) + external_file = tmp_path.parent / f"{tmp_path.name}-outside" + external_file.write_bytes(b"outside dataset") + source_path = tmp_path / relative_path + if source_path.exists(): + source_path.unlink() + source_path.symlink_to(external_file) + + assert not readiness.dataset_ready(tmp_path, task, dataset) + + +def test_dense_readiness_rejects_symlinked_root_ancestors( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Do not reuse complete dense roots reached through a symlinked parent.""" + + monkeypatch.setattr(readiness, "NYU_DEPTH_VALIDATION_SAMPLE_COUNT", 1) + monkeypatch.setattr(readiness, "ADE20K_VALIDATION_SAMPLE_COUNT", 1) + target_parent = tmp_path / "target" + nyu_root = target_parent / "nyu-depth" + _write_file(nyu_root / "images" / "sample.jpg") + _write_file(nyu_root / "depth" / "sample.npy") + ade20k_root = target_parent / "ade20k" + _write_file(ade20k_root / "images" / "ADE_val_00000001.jpg") + _write_file(ade20k_root / "annotations" / "ADE_val_00000001.png") + for file_name in readiness.ADE20K_METADATA_FILES: + _write_file(ade20k_root / file_name) + symlinked_parent = tmp_path / "datasets" + symlinked_parent.symlink_to(target_parent, target_is_directory=True) + + assert not readiness.dataset_ready( + symlinked_parent / "nyu-depth", "depth_estimation", "nyu-depth" + ) + assert not readiness.dataset_ready( + symlinked_parent / "ade20k", "semantic_segmentation", "ade20k" + ) + + existing_dir = tmp_path / "existing" + existing_dir.mkdir() + traversed_parent = existing_dir / ".." / symlinked_parent.name + assert not readiness.dataset_ready( + traversed_parent / "nyu-depth", "depth_estimation", "nyu-depth" + ) + assert not readiness.dataset_ready( + traversed_parent / "ade20k", "semantic_segmentation", "ade20k" + ) + + target_child = target_parent / "child" + target_child.mkdir() + traversal_link = tmp_path / "traversal-link" + traversal_link.symlink_to(target_child, target_is_directory=True) + symlink_traversed_parent = traversal_link / ".." + assert not readiness.dataset_ready( + symlink_traversed_parent / "nyu-depth", "depth_estimation", "nyu-depth" + ) + assert not readiness.dataset_ready( + symlink_traversed_parent / "ade20k", "semantic_segmentation", "ade20k" + ) diff --git a/tests/test_depth_estimation.py b/tests/test_depth_estimation.py new file mode 100644 index 0000000..5bd9087 --- /dev/null +++ b/tests/test_depth_estimation.py @@ -0,0 +1,182 @@ +"""Unit tests for YOLO26 depth-estimation support.""" + +from __future__ import annotations + +import cv2 +import numpy as np +import pytest +import torch + +from mblt_vision.utils.datasets import CustomNYUDepth +from mblt_vision.utils.postprocess import DepthPost +from mblt_vision.wrapper import resolve_model_config + + +@pytest.mark.parametrize( + ("model_name", "task"), + [ + ("yolo26s", "object_detection"), + ("yolo26s-seg", "instance_segmentation"), + ("yolo26s-pose", "pose_estimation"), + ("yolo26s-obb", "obb"), + ], +) +def test_other_yolo_validation_tasks_keep_letterbox(model_name: str, task: str) -> None: + """Keep aspect-preserving letterbox preprocessing for non-depth YOLO tasks.""" + + config = resolve_model_config(model_name) + assert "LetterBox" in config["pre_cfg"] + assert "Resize" not in config["pre_cfg"] + assert config["post_cfg"]["task"] == task + + +def test_depth_post_restores_letterbox_padding() -> None: + """Crop letterbox padding before bilinearly restoring an original image shape.""" + + post = DepthPost({"LetterBox": {"img_size": [8, 8]}}, {}) + output = torch.zeros((1, 1, 8, 8)) + output[:, :, 2:6, :] = 2.0 + restored = post(output, img0_shape=(2, 4), ratio_pad=((2.0, 2.0), (0.0, 2.0))) + assert isinstance(restored, torch.Tensor) + assert restored.shape == (2, 4) + assert torch.allclose(restored, torch.full((2, 4), 2.0)) + with pytest.raises(ValueError, match=r"expects \[B, 1, H, W\] or \[B, H, W, 1\]"): + post(torch.zeros((1, 2, 8, 8))) + + +def test_depth_post_normalizes_quarter_resolution_mxq_before_restoring() -> None: + """Upsample the MXQ depth layout before undoing letterbox padding.""" + + post = DepthPost({"LetterBox": {"img_size": [8, 8]}}, {}) + mxq_depth = torch.arange(4, dtype=torch.float32).reshape(1, 2, 2) + expected = torch.nn.functional.interpolate( + mxq_depth[:, None], scale_factor=4.0, mode="bilinear", align_corners=False + )[:, 0] + + normalized = post(mxq_depth) + assert isinstance(normalized, torch.Tensor) + assert torch.equal(normalized, expected) + + restored = post(mxq_depth, img0_shape=(2, 4), ratio_pad=((2.0, 2.0), (0.0, 2.0))) + assert isinstance(restored, torch.Tensor) + assert restored.shape == (2, 4) + expected_restored = torch.nn.functional.interpolate( + expected[0, 2:6][None, None], size=(2, 4), mode="bilinear", align_corners=False + )[0, 0] + assert torch.equal(restored, expected_restored) + + +def test_depth_post_keeps_full_resolution_onnx_output_and_rejects_other_scales() -> ( + None +): + """Keep ONNX depth unchanged while rejecting unsupported dense output scales.""" + + post = DepthPost({"LetterBox": {"img_size": [8, 8]}}, {}) + full_resolution = torch.arange(64, dtype=torch.float32).reshape(1, 1, 8, 8) + normalized = post(full_resolution) + assert isinstance(normalized, torch.Tensor) + assert torch.equal(normalized, full_resolution[:, 0]) + + with pytest.raises(ValueError, match="spatial shape must be"): + post(torch.zeros((1, 3, 3))) + + +def test_depth_post_normalizes_full_resolution_channel_last_mxq_output() -> None: + """Normalize batched and single-image baked MXQ resize outputs without resizing again.""" + + post = DepthPost({"LetterBox": {"img_size": [8, 8]}}, {}) + full_resolution = torch.arange(64, dtype=torch.float32).reshape(1, 8, 8, 1) + normalized = post(full_resolution) + assert isinstance(normalized, torch.Tensor) + assert torch.equal(normalized, full_resolution[..., 0]) + + single_image = full_resolution[0] + normalized_single_image = post(single_image) + assert isinstance(normalized_single_image, torch.Tensor) + assert torch.equal(normalized_single_image, full_resolution[..., 0]) + + +def test_depth_post_normalizes_quarter_resolution_channel_last_output() -> None: + """Remove a single-image channel-last axis before checking quarter resolution.""" + + post = DepthPost({"LetterBox": {"img_size": [8, 8]}}, {}) + channel_last_depth = torch.arange(4, dtype=torch.float32).reshape(2, 2, 1) + expected = torch.nn.functional.interpolate( + channel_last_depth[..., 0][None, None], + scale_factor=4.0, + mode="bilinear", + align_corners=False, + )[:, 0] + + normalized = post(channel_last_depth) + + assert isinstance(normalized, torch.Tensor) + assert torch.equal(normalized, expected) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")]) +def test_depth_post_rejects_nonfinite_outputs(invalid_value: float) -> None: + """Reject invalid local depth tensors before resizing or visualization.""" + + post = DepthPost({"LetterBox": {"img_size": [8, 8]}}, {}) + depth = torch.zeros((1, 8, 8)) + depth[0, 0, 0] = invalid_value + + with pytest.raises(ValueError, match="must contain only finite values"): + post(depth) + + +def test_nyu_depth_dataset_rejects_mismatched_image_and_target_shapes(tmp_path) -> None: + """Reject paired NYU files whose pixels cannot be compared one-to-one.""" + + image_dir = tmp_path / "images" + depth_dir = tmp_path / "depth" + image_dir.mkdir() + depth_dir.mkdir() + assert cv2.imwrite( + str(image_dir / "sample.png"), np.zeros((4, 5, 3), dtype=np.uint8) + ) + np.save(depth_dir / "sample.npy", np.zeros((3, 5), dtype=np.float32)) + + with pytest.raises(ValueError, match=r"sample: image \(4, 5\), depth \(3, 5\)"): + CustomNYUDepth(str(tmp_path))[0] + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")]) +def test_nyu_depth_dataset_rejects_nonfinite_targets( + tmp_path, invalid_value: float +) -> None: + """Do not silently remove invalid NYU ground-truth pixels from metrics.""" + + image_dir = tmp_path / "images" + depth_dir = tmp_path / "depth" + image_dir.mkdir() + depth_dir.mkdir() + assert cv2.imwrite( + str(image_dir / "sample.png"), np.zeros((4, 5, 3), dtype=np.uint8) + ) + depth = np.ones((4, 5), dtype=np.float32) + depth[0, 0] = invalid_value + np.save(depth_dir / "sample.npy", depth) + + with pytest.raises(ValueError, match=r"finite values: .*sample\.npy"): + CustomNYUDepth(str(tmp_path))[0] + + +@pytest.mark.parametrize("dtype", [np.complex64, np.dtype("U4")]) +def test_nyu_depth_dataset_rejects_nonreal_or_nonnumeric_targets( + tmp_path, dtype: np.dtype +) -> None: + """Reject target arrays before lossy conversion to float32.""" + + image_dir = tmp_path / "images" + depth_dir = tmp_path / "depth" + image_dir.mkdir() + depth_dir.mkdir() + assert cv2.imwrite( + str(image_dir / "sample.png"), np.zeros((4, 5, 3), dtype=np.uint8) + ) + np.save(depth_dir / "sample.npy", np.ones((4, 5), dtype=dtype)) + + with pytest.raises(ValueError, match="real numeric dtype"): + CustomNYUDepth(str(tmp_path))[0] diff --git a/tests/test_eval_coco.py b/tests/test_eval_coco.py new file mode 100644 index 0000000..98b0b68 --- /dev/null +++ b/tests/test_eval_coco.py @@ -0,0 +1,373 @@ +"""Focused COCO evaluator regression tests.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace +from typing import Any + +import pytest + +eval_coco_module = importlib.import_module("mblt_vision.utils.evaluation.eval_coco") + + +def test_pose_evaluation_uses_all_keypoints_annotation_images( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not omit COCO images that lack visible-person keypoint annotations.""" + + constructed_kwargs: dict[str, Any] = {} + evaluated_image_ids: list[int] | None = None + + class _Dataset: + def __init__(self, root: str, annotation_path: str, **kwargs: Any) -> None: + constructed_kwargs.update( + root=root, annotation_path=annotation_path, **kwargs + ) + self.ids = [1, 2] + self.coco = SimpleNamespace( + cats={1: {}}, + anns={ + 1: { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + "keypoints": [0, 0, 0] * 17, + "num_keypoints": 0, + } + }, + imgs={1: {"height": 1, "width": 1}, 2: {"height": 1, "width": 1}}, + ) + + def __len__(self) -> int: + return len(self.ids) + + class _Model: + post_cfg = {"task": "pose_estimation", "dataset": "coco"} + + def set_postprocess_thresholds( + self, *, conf_thres: float | None, iou_thres: float | None + ) -> None: + del conf_thres, iou_thres + + def preprocess_with_metadata(self, image: object) -> object: + return image + + def _evaluate( + coco_gt: object, + coco_results: list[dict[str, Any]], + task: str, + img_ids: list[int] | None = None, + ) -> SimpleNamespace: + nonlocal evaluated_image_ids + del coco_gt, coco_results + assert task == "pose_estimation" + evaluated_image_ids = img_ids + return SimpleNamespace(stats=[SimpleNamespace(item=lambda: 0.1)] * 2) + + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", _Dataset) + monkeypatch.setattr(eval_coco_module, "get_coco_loader", lambda *_: []) + monkeypatch.setattr(eval_coco_module, "evaluate_predictions_on_coco", _evaluate) + + result = eval_coco_module.eval_coco_metrics(_Model(), "/dataset", batch_size=2) + + assert "min_keypoints" not in constructed_kwargs + assert constructed_kwargs["annotation_path"].endswith( + "person_keypoints_val2017.json" + ) + assert evaluated_image_ids == [1, 2] + assert result.map5095 == result.map50 == 0.1 + + +def test_coco_evaluation_rejects_non_coco_model_taxonomy() -> None: + """Do not evaluate another taxonomy using the hard-coded COCO ID mapping.""" + + model = SimpleNamespace(post_cfg={"task": "object_detection", "dataset": "dotav1"}) + + with pytest.raises(ValueError, match="post_cfg.dataset to be 'coco'"): + eval_coco_module.eval_coco_metrics(model, "/dataset", batch_size=1) + + +def test_coco_evaluation_rejects_noncanonical_artifact_categories( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not score direct COCO artifacts with an incompatible category taxonomy.""" + + dataset = SimpleNamespace(coco=SimpleNamespace(cats={999: {}}, anns={})) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match=r"unsupported category IDs: \[999\]"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace(post_cfg={"task": "object_detection", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_duplicate_raw_annotation_ids( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Validate raw COCO IDs before the backend's index can overwrite one row.""" + + annotation = { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + } + dataset = SimpleNamespace( + raw_annotation={ + "images": [{"id": 1, "height": 1, "width": 1}], + "categories": [{"id": 1}], + "annotations": [annotation, annotation.copy()], + }, + coco=SimpleNamespace( + cats={1: {}}, + imgs={1: {"height": 1, "width": 1}}, + anns={1: annotation}, + ), + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="duplicate or invalid raw IDs"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace(post_cfg={"task": "object_detection", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_empty_ground_truth( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A direct COCO run must not present undefined AP as a normal metric.""" + + dataset = SimpleNamespace( + coco=SimpleNamespace(cats={1: {}}, imgs={1: {"height": 1, "width": 1}}, anns={}) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="at least one annotation"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace(post_cfg={"task": "object_detection", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_undeclared_annotation_categories( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not let malformed annotation categories be omitted from direct AP.""" + + dataset = SimpleNamespace( + coco=SimpleNamespace( + cats={1: {}}, + anns={7: {"category_id": 999}}, + imgs={1: {"height": 1, "width": 1}}, + ) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="invalid task-specific annotations"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace(post_cfg={"task": "object_detection", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_invalid_task_payloads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Use readiness-equivalent checks before direct COCO evaluation starts.""" + + dataset = SimpleNamespace( + coco=SimpleNamespace( + cats={1: {}}, + imgs={1: {"height": 2, "width": 2}}, + anns={ + 7: { + "id": 7, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": -1, + "iscrowd": 0, + "segmentation": [[0, 0, 1, 0, 1, 1]], + } + }, + ) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="invalid task-specific annotations"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace( + post_cfg={"task": "instance_segmentation", "dataset": "coco"} + ), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_polygons_outside_image_bounds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not let a nonempty off-image polygon become an empty COCO mask.""" + + dataset = SimpleNamespace( + coco=SimpleNamespace( + cats={1: {}}, + imgs={1: {"height": 10, "width": 10}}, + anns={ + 7: { + "id": 7, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + "segmentation": [[11, 0, 12, 0, 12, 1, 11, 1]], + } + }, + ) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="invalid task-specific annotations"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace( + post_cfg={"task": "instance_segmentation", "dataset": "coco"} + ), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_polygons_without_rasterized_foreground( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Continuous subpixel overlap must not admit an empty raster mask.""" + + dataset = SimpleNamespace( + coco=SimpleNamespace( + cats={1: {}}, + imgs={1: {"height": 10, "width": 10}}, + anns={ + 7: { + "id": 7, + "image_id": 1, + "category_id": 1, + "bbox": [9, 9, 1, 1], + "area": 1, + "iscrowd": 0, + "segmentation": [[9.9, 9.9, 9.99, 9.9, 9.99, 9.99]], + } + }, + ) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="invalid task-specific annotations"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace( + post_cfg={"task": "instance_segmentation", "dataset": "coco"} + ), + "/dataset", + batch_size=1, + ) + + +@pytest.mark.parametrize("visibility", [1, 2]) +def test_coco_evaluation_rejects_labeled_keypoints_outside_images( + monkeypatch: pytest.MonkeyPatch, visibility: int +) -> None: + """Both occluded and visible pose coordinates must be in the source image.""" + + keypoints = [0.0, 0.0, 0.0] * 17 + keypoints[:3] = [11.0, 1.0, visibility] + dataset = SimpleNamespace( + coco=SimpleNamespace( + cats={1: {}}, + imgs={1: {"height": 10, "width": 10}}, + anns={ + 7: { + "id": 7, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 1, + "iscrowd": 0, + "keypoints": keypoints, + "num_keypoints": 1, + } + }, + ) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="invalid task-specific annotations"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace(post_cfg={"task": "pose_estimation", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_coco_evaluation_rejects_pose_area_larger_than_its_box( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not let inflated pose areas weaken OKS matching thresholds.""" + + dataset = SimpleNamespace( + coco=SimpleNamespace( + cats={1: {}}, + imgs={1: {"height": 10, "width": 10}}, + anns={ + 7: { + "id": 7, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 1, 1], + "area": 100, + "iscrowd": 0, + "keypoints": [0, 0, 0] * 17, + "num_keypoints": 0, + } + }, + ) + ) + monkeypatch.setattr(eval_coco_module, "CustomCOCODataset", lambda *_: dataset) + + with pytest.raises(ValueError, match="invalid task-specific annotations"): + eval_coco_module.eval_coco_metrics( + SimpleNamespace(post_cfg={"task": "pose_estimation", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_coco_result_formatter_rejects_truncated_postprocess_batch() -> None: + """Require one decoded result for every submitted COCO image.""" + + postprocess = SimpleNamespace( + nmsout2eval=lambda *_args, **_kwargs: ([[1]], [[[0, 0, 1, 1]]], [[0.9]]) + ) + + with pytest.raises(ValueError, match="batch cardinality mismatch"): + eval_coco_module.format_coco_results( + "object_detection", + SimpleNamespace(output=[]), + (640, 640), + [(640, 640), (640, 640)], + [None, None], + [0, 1], + [101, 102], + postprocess, + ) diff --git a/tests/test_eval_dota.py b/tests/test_eval_dota.py new file mode 100644 index 0000000..759054c --- /dev/null +++ b/tests/test_eval_dota.py @@ -0,0 +1,528 @@ +"""Tests for DOTAv1 difficult-region evaluation.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest +import torch + +from mblt_vision.utils.datasets import CustomDOTAv1 +from mblt_vision.utils.evaluation.eval_dota import ( + _load_ground_truths, + evaluate_dota_predictions, + format_dota_results, +) +from mblt_vision.utils.results import Results + +eval_dota_module = importlib.import_module("mblt_vision.utils.evaluation.eval_dota") + + +def test_eval_dota_rejects_truncated_postprocess_batches( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + """Do not silently omit ground truths when a batched model output is short.""" + + class _FakeProgress: + def __init__(self, iterable, **kwargs) -> None: + del kwargs + self._iterable = iterable + + def __iter__(self): + return iter(self._iterable) + + def set_postfix_str(self, value: str) -> None: + del value + + def close(self) -> None: + return None + + class _FakeModel: + post_cfg = {"task": "obb", "dataset": "dotav1"} + preprocess_with_metadata = object() + + def set_postprocess_thresholds(self, **kwargs) -> None: + del kwargs + + def __call__(self, inputs: torch.Tensor) -> torch.Tensor: + return inputs + + def postprocess(self, outputs: torch.Tensor) -> SimpleNamespace: + del outputs + return SimpleNamespace(output=[torch.zeros((0, 7), dtype=torch.float32)]) + + class _FakeDataset: + def __len__(self) -> int: + return 2 + + batch = ( + torch.zeros((2, 8, 8, 3), dtype=torch.float32), + [(8, 8), (8, 8)], + [None, None], + ("first", "second"), + ) + monkeypatch.setattr(eval_dota_module, "CustomDOTAv1", lambda _: _FakeDataset()) + monkeypatch.setattr(eval_dota_module, "get_dota_loader", lambda *args: [batch]) + monkeypatch.setattr(eval_dota_module, "_load_ground_truths", lambda *args: {}) + monkeypatch.setattr(eval_dota_module, "tqdm", _FakeProgress) + + with pytest.raises( + ValueError, + match=r"DOTAv1 evaluation batch length mismatch: model outputs=1, input batch=2", + ): + eval_dota_module.eval_dota(_FakeModel(), str(tmp_path), batch_size=2) + + +def test_eval_dota_rejects_wrong_model_taxonomy(tmp_path) -> None: + """Do not score another OBB taxonomy against DOTAv1 ground truth.""" + + with pytest.raises(ValueError, match="post_cfg.dataset to be 'dotav1'"): + eval_dota_module.eval_dota( + SimpleNamespace(post_cfg={"task": "obb", "dataset": "coco"}), + str(tmp_path), + batch_size=1, + ) + + +def test_dota_ground_truth_requires_annotation_for_every_image(tmp_path) -> None: + """Do not silently manufacture empty ground truth for unlabeled DOTAv1 samples.""" + + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises( + FileNotFoundError, match="annotation not found for image 'image'" + ): + _load_ground_truths(str(tmp_path), dataset) + + +def test_normalized_difficult_flag_loads_as_an_ignored_region(tmp_path) -> None: + """Read organizer-produced difficult metadata before selecting evaluation targets.""" + + label_dir = tmp_path / "labels" / "val" + label_dir.mkdir(parents=True) + (label_dir / "image.txt").write_text( + "0 0 0 0.2 0 0.2 0.2 0 0.2 0\n0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6 1\n", + encoding="utf-8", + ) + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + ground_truth = _load_ground_truths(str(tmp_path), dataset)["image"] + + assert ground_truth["cls"].tolist() == [0] + assert ground_truth["ignore_cls"].tolist() == [0] + + +def test_normalized_truncated_annotation_raises_with_file_and_line(tmp_path) -> None: + """Reject nonblank normalized annotations without a complete OBB polygon.""" + + label_dir = tmp_path / "labels" / "val" + label_dir.mkdir(parents=True) + label_path = label_dir / "image.txt" + label_path.write_text("\n0 0 0 0.2\n", encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises( + ValueError, + match=r"image\.txt:2: expected at least 9 fields, got 4", + ): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "0 0 0 0.2 0 0.2 0.2 0 0.2"), + ("labels/val_original/image.txt", "0 0 20 0 20 20 0 20 plane 0"), + ], +) +def test_dota_annotations_reject_duplicate_targets( + tmp_path, label_path: str, annotation: str +) -> None: + """A target may be represented only once within one source image.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(f"{annotation}\n{annotation}\n", encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="Duplicate DOTAv1 annotation target"): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotations"), + [ + ( + "labels/val/image.txt", + "0 0 0 0.2 0 0.2 0.2 0 0.2\n0 0.2 0 0.2 0.2 0 0.2 0 0", + ), + ( + "labels/val_original/image.txt", + "0 0 20 0 20 20 0 20 plane 0\n20 0 20 20 0 20 0 0 plane 0", + ), + ], +) +def test_dota_annotations_reject_equivalent_reordered_targets( + tmp_path, label_path: str, annotations: str +) -> None: + """Duplicate quadrilaterals remain duplicates after cyclic reordering.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotations, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="Duplicate DOTAv1 annotation target"): + _load_ground_truths(str(tmp_path), dataset) + + +def test_original_truncated_annotation_raises_with_file_and_line(tmp_path) -> None: + """Require a difficulty flag when loading original DOTAv1 labels directly.""" + + original_label_dir = tmp_path / "labels" / "val_original" + original_label_dir.mkdir(parents=True) + label_path = original_label_dir / "image.txt" + label_path.write_text( + "imagesource:GoogleEarth\n0 0 20 0 20 20 0 20 plane\n", encoding="utf-8" + ) + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises( + ValueError, + match=r"image\.txt:2: expected at least 10 fields, got 9", + ): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "0 0 0 1 0 2 0 3 0"), + ("labels/val_original/image.txt", "0 0 1 0 2 0 3 0 plane 0"), + ], +) +def test_dota_annotations_reject_degenerate_polygons( + tmp_path, label_path: str, annotation: str +) -> None: + """Reject line-like quadrilaterals before deriving rotated boxes.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotation, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="polygon must have positive area"): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "0 0 0 0.2 0 0.2 0.2 0 0"), + ("labels/val_original/image.txt", "0 0 20 0 20 20 0 0 plane 0"), + ], +) +def test_dota_annotations_reject_repeated_vertices( + tmp_path, label_path: str, annotation: str +) -> None: + """DOTAv1 quadrilaterals cannot silently degrade into triangles.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotation, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="four distinct vertices"): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "0 1.1 0 1.2 0 1.2 0.1 1.1 0.1"), + ("labels/val_original/image.txt", "110 0 120 0 120 10 110 10 plane 0"), + ], +) +def test_dota_annotations_reject_polygons_outside_images( + tmp_path, label_path: str, annotation: str +) -> None: + """Ground-truth polygons must retain foreground within their source images.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotation, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="must overlap its source image"): + _load_ground_truths(str(tmp_path), dataset) + + +def test_dota_normalized_coordinates_are_scaled_even_when_partially_clipped( + tmp_path, +) -> None: + """The normalized-label directory always uses normalized image coordinates.""" + + label_path = tmp_path / "labels" / "val" / "image.txt" + label_path.parent.mkdir(parents=True) + label_path.write_text("0 0.9 0 1.6 0 1.6 0.1 0.9 0.1", encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + ground_truth = _load_ground_truths(str(tmp_path), dataset)["image"] + + assert ground_truth["polygons"][0, :, 0].tolist() == [90.0, 160.0, 160.0, 90.0] + + +def test_dota_ground_truth_rejects_orphan_label_files(tmp_path) -> None: + """Direct evaluation must not ignore labels without a selected image.""" + + label_path = tmp_path / "labels" / "val" / "orphan.txt" + label_path.parent.mkdir(parents=True) + label_path.write_text("0 0 0 0.2 0 0.2 0.2 0 0.2", encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace(ids=[], image_paths=[], _load_image=lambda _: None), + ) + + with pytest.raises(ValueError, match="no corresponding validation image"): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "-1 0 0 0.2 0 0.2 0.2 0 0.2"), + ("labels/val_original/image.txt", "0 0 20 0 20 20 0 20 -1 0"), + ], +) +def test_dota_annotations_reject_negative_class_indices( + tmp_path, label_path: str, annotation: str +) -> None: + """Reject invalid negative classes in normalized and original labels.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotation, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="Unsupported DOTAv1 class index -1"): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "0 nan 0 0.2 0 0.2 0.2 0 0.2"), + ("labels/val_original/image.txt", "inf 0 20 0 20 20 0 20 plane 0"), + ], +) +def test_dota_annotations_reject_nonfinite_coordinates( + tmp_path, label_path: str, annotation: str +) -> None: + """Reject poisoned normalized and original DOTA polygon coordinates.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotation, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match=r"coordinates must be finite.*image\.txt:1"): + _load_ground_truths(str(tmp_path), dataset) + + +@pytest.mark.parametrize( + ("label_path", "annotation"), + [ + ("labels/val/image.txt", "0 0 0 0.2 0 0.2 0.2 0 0.2 3"), + ("labels/val_original/image.txt", "0 0 20 0 20 20 0 20 plane invalid"), + ], +) +def test_dota_annotations_reject_unknown_difficulty_flags( + tmp_path, label_path: str, annotation: str +) -> None: + """Keep invalid difficult-region metadata from changing positive counts.""" + + path = tmp_path / label_path + path.parent.mkdir(parents=True) + path.write_text(annotation, encoding="utf-8") + dataset = cast( + CustomDOTAv1, + SimpleNamespace( + ids=["image"], + image_paths=["unused"], + _load_image=lambda _: np.zeros((100, 100, 3), dtype=np.uint8), + ), + ) + + with pytest.raises(ValueError, match="Unsupported DOTAv1 difficulty flag"): + _load_ground_truths(str(tmp_path), dataset) + + +def test_difficult_regions_do_not_count_as_positive_or_false_positive() -> None: + """Ignore a detection on a difficult region while retaining positive matching.""" + + ground_truths = { + "image": { + "cls": torch.tensor([0]), + "bboxes": torch.tensor([[10.0, 10.0, 4.0, 4.0, 0.0]]), + "ignore_cls": torch.tensor([0]), + "ignore_bboxes": torch.tensor([[30.0, 30.0, 4.0, 4.0, 0.0]]), + } + } + predictions = [ + { + "image_id": "image", + "category_id": 0, + "score": 0.99, + "rbox": [30.0, 30.0, 4.0, 4.0, 0.0], + }, + { + "image_id": "image", + "category_id": 0, + "score": 0.90, + "rbox": [10.0, 10.0, 4.0, 4.0, 0.0], + }, + ] + + result = evaluate_dota_predictions(ground_truths, predictions) + baseline = evaluate_dota_predictions(ground_truths, predictions[1:]) + + assert result == baseline + + +def test_dota_ap_interpolation_uses_terminal_recall_sentinel() -> None: + """Preserve the reference AP curve after the final observed recall point.""" + + ap, _, recall_curve = eval_dota_module._compute_ap(np.array([0.5]), np.array([1.0])) + + assert recall_curve.tolist() == [0.0, 0.5, 1.0] + assert ap == pytest.approx(0.75) + + +def test_dota_metrics_require_non_ignored_ground_truth() -> None: + """Empty or difficult-only labels must not yield a normal AP result.""" + + with pytest.raises(ValueError, match="at least one non-ignored target"): + eval_dota_module._evaluate_stats(eval_dota_module._empty_stats()) + + +def test_dota_export_rejects_truncated_converted_batches() -> None: + """Do not silently omit Task1 files when converted output batches are short.""" + + postprocess = SimpleNamespace( + nmsout2eval=lambda *_args, **_kwargs: ([[]], [[]], [[]], [[]]) + ) + + with pytest.raises(ValueError, match="export batch length mismatch"): + format_dota_results( + cast(Results, SimpleNamespace(output=[])), + (640, 640), + [(640, 640), (640, 640)], + [None, None], + ("first", "second"), + postprocess, + ) + + +def test_dota_export_rejects_mismatched_detection_fields() -> None: + """Require a polygon, score, and rotated box for every exported label.""" + + postprocess = SimpleNamespace( + nmsout2eval=lambda *_args, **_kwargs: ([["plane"]], [[]], [[0.9]], [[]]) + ) + + with pytest.raises(ValueError, match="export detection length mismatch"): + format_dota_results( + cast(Results, SimpleNamespace(output=[])), + (640, 640), + [(640, 640)], + [None], + ("first",), + postprocess, + ) diff --git a/tests/test_eval_imagenet.py b/tests/test_eval_imagenet.py new file mode 100644 index 0000000..4520575 --- /dev/null +++ b/tests/test_eval_imagenet.py @@ -0,0 +1,167 @@ +"""Focused ImageNet evaluator regression tests.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import pytest +import torch + +eval_imagenet_module = importlib.import_module( + "mblt_vision.utils.evaluation.eval_imagenet" +) + + +def test_imagenet_evaluation_rejects_truncated_classification_batches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not broadcast one classification row across a multi-image label batch.""" + + class _FakeDataset: + classes = ["n00000000", "n00000001"] + class_to_idx = {"n00000000": 0, "n00000001": 1} + + def make_dataset(self) -> None: + return None + + def __len__(self) -> int: + return 2 + + class _FakeProgress: + def __init__(self, iterable, **kwargs) -> None: + del kwargs + self._iterable = iterable + + def __iter__(self): + return iter(self._iterable) + + def set_postfix_str(self, value: str) -> None: + del value + + def close(self) -> None: + return None + + class _FakeModel: + post_cfg = {"dataset": "imagenet"} + + def preprocess(self, value: object) -> object: + return value + + def __call__(self, inputs: torch.Tensor) -> torch.Tensor: + return inputs + + def postprocess(self, outputs: torch.Tensor) -> SimpleNamespace: + del outputs + return SimpleNamespace(output=torch.zeros((1, 5), dtype=torch.float32)) + + batch = (torch.zeros((2, 3, 8, 8)), torch.tensor([0, 1])) + monkeypatch.setattr( + eval_imagenet_module, "CustomImageFolder", lambda _: _FakeDataset() + ) + monkeypatch.setattr( + eval_imagenet_module, "get_imagenet_loader", lambda *args: [batch] + ) + monkeypatch.setattr(eval_imagenet_module, "tqdm", _FakeProgress) + monkeypatch.setattr( + eval_imagenet_module, "IMAGENET_SYNSET_ORDER", ("n00000000", "n00000001") + ) + monkeypatch.setattr( + eval_imagenet_module, "IMAGENET_SYNSETS", frozenset({"n00000000", "n00000001"}) + ) + + with pytest.raises( + ValueError, + match="got 1 outputs for 2 labels", + ): + eval_imagenet_module.eval_imagenet_metrics( + _FakeModel(), "/dataset", batch_size=2 + ) + + +def test_imagenet_evaluation_preserves_canonical_synset_indices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not renumber a partial ImageNet directory tree by local sort order.""" + + class _Dataset(SimpleNamespace): + def __len__(self) -> int: + return 1 + + dataset = _Dataset( + classes=["n00000001"], + class_to_idx={"n00000001": 0}, + make_dataset=lambda: None, + ) + monkeypatch.setattr(eval_imagenet_module, "CustomImageFolder", lambda _: dataset) + monkeypatch.setattr( + eval_imagenet_module, "IMAGENET_SYNSET_ORDER", ("n00000000", "n00000001") + ) + monkeypatch.setattr( + eval_imagenet_module, "IMAGENET_SYNSETS", frozenset({"n00000000", "n00000001"}) + ) + + class _StopAfterClassMapping(Exception): + pass + + def _stop_after_mapping(*args: object) -> object: + del args + raise _StopAfterClassMapping + + monkeypatch.setattr( + eval_imagenet_module, "get_imagenet_loader", _stop_after_mapping + ) + + with pytest.raises(_StopAfterClassMapping): + eval_imagenet_module.eval_imagenet_metrics( + SimpleNamespace(post_cfg={"dataset": "imagenet"}, preprocess=object()), + "/dataset", + batch_size=1, + ) + + assert dataset.class_to_idx == {"n00000001": 1} + + +def test_imagenet_evaluation_rejects_empty_dataset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Fail before loader construction when no supported samples are available.""" + + class _EmptyDataset: + classes = ["n00000000"] + class_to_idx = {"n00000000": 0} + + def make_dataset(self) -> None: + return None + + def __len__(self) -> int: + return 0 + + monkeypatch.setattr( + eval_imagenet_module, "CustomImageFolder", lambda _: _EmptyDataset() + ) + monkeypatch.setattr(eval_imagenet_module, "IMAGENET_SYNSET_ORDER", ("n00000000",)) + monkeypatch.setattr( + eval_imagenet_module, "IMAGENET_SYNSETS", frozenset({"n00000000"}) + ) + monkeypatch.setattr( + eval_imagenet_module, + "get_imagenet_loader", + lambda *_: pytest.fail("empty datasets must not create an ImageNet loader"), + ) + + with pytest.raises(ValueError, match="contains no supported images"): + eval_imagenet_module.eval_imagenet_metrics( + SimpleNamespace(post_cfg={"dataset": "imagenet"}, preprocess=object()), + "/empty", + batch_size=1, + ) + + +def test_imagenet_evaluation_rejects_wrong_model_taxonomy() -> None: + """Do not score another classification taxonomy against ImageNet labels.""" + + with pytest.raises(ValueError, match="post_cfg.dataset to be 'imagenet'"): + eval_imagenet_module.eval_imagenet_metrics( + SimpleNamespace(post_cfg={"dataset": "coco"}), "/dataset", batch_size=1 + ) diff --git a/tests/test_eval_nyu_depth.py b/tests/test_eval_nyu_depth.py new file mode 100644 index 0000000..e635938 --- /dev/null +++ b/tests/test_eval_nyu_depth.py @@ -0,0 +1,77 @@ +"""Regression tests for NYU depth-evaluation output validation.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +eval_nyu_depth_module = importlib.import_module( + "mblt_vision.utils.evaluation.eval_nyu_depth" +) + + +def test_nyu_depth_evaluation_rejects_surplus_output_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not silently discard depth maps from a malformed backend batch.""" + + class _Model: + post_cfg = {"dataset": "nyu-depth"} + pre_cfg = {"LetterBox": {"img_size": [4, 4]}} + + def __call__(self, inputs: torch.Tensor) -> torch.Tensor: + return inputs + + def postprocess(self, output: torch.Tensor) -> SimpleNamespace: + del output + return SimpleNamespace(depth=torch.zeros((2, 4, 4))) + + batch = ( + torch.zeros((1, 4, 4, 3)), + np.ones((1, 4, 4), dtype=np.float32), + [(4, 4)], + [None], + ("sample",), + ) + monkeypatch.setattr(eval_nyu_depth_module, "CustomNYUDepth", lambda _: object()) + monkeypatch.setattr( + eval_nyu_depth_module, "get_nyu_depth_loader", lambda *_args, **_kwargs: [batch] + ) + monkeypatch.setattr(eval_nyu_depth_module, "build_preprocess", lambda _: object()) + + with pytest.raises( + ValueError, match=r"output batch length mismatch: maps=2, targets=1" + ): + eval_nyu_depth_module.eval_nyu_depth(_Model(), "/dataset", batch_size=1) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf"), -1.0]) +def test_nyu_depth_metrics_reject_invalid_targets(invalid_value: float) -> None: + """Keep direct metric callers from silently excluding corrupt depth targets.""" + + target = np.array([[1.0, invalid_value]], dtype=np.float32) + with pytest.raises(ValueError, match="target contains non-finite|negative values"): + eval_nyu_depth_module.calculate_nyu_depth_metrics(np.ones((1, 2)), target) + + +@pytest.mark.parametrize("dtype", [np.complex64, np.dtype("U4")]) +def test_nyu_depth_metrics_reject_nonreal_or_nonnumeric_inputs(dtype: np.dtype) -> None: + """Direct metric callers must not lose source-dtype corruption during casting.""" + + values = np.ones((1, 2), dtype=dtype) + + with pytest.raises(ValueError, match="real numeric dtype"): + eval_nyu_depth_module.calculate_nyu_depth_metrics(values, values) + + +def test_nyu_depth_evaluation_rejects_wrong_model_taxonomy() -> None: + """Require the model's declared depth taxonomy before loading a dataset.""" + + with pytest.raises(ValueError, match="post_cfg.dataset to be 'nyu-depth'"): + eval_nyu_depth_module.eval_nyu_depth( + SimpleNamespace(post_cfg={"dataset": "ade20k"}), "/dataset", batch_size=1 + ) diff --git a/tests/test_eval_widerface.py b/tests/test_eval_widerface.py new file mode 100644 index 0000000..d86f651 --- /dev/null +++ b/tests/test_eval_widerface.py @@ -0,0 +1,167 @@ +"""Focused WiderFace evaluator validation tests.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +eval_widerface_module = importlib.import_module( + "mblt_vision.utils.evaluation.eval_widerface" +) + + +def test_widerface_evaluation_rejects_wrong_model_taxonomy() -> None: + """Do not evaluate another detector taxonomy using WiderFace metadata.""" + + with pytest.raises(ValueError, match="post_cfg.dataset to be 'widerface'"): + eval_widerface_module.eval_widerface( + SimpleNamespace(post_cfg={"task": "face_detection", "dataset": "coco"}), + "/dataset", + batch_size=1, + ) + + +def test_widerface_evaluation_rejects_truncated_postprocess_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Require one decoded face prediction list for every submitted image.""" + + class _Dataset: + samples = [("unused", "event", "first.jpg"), ("unused", "event", "second.jpg")] + + def __len__(self) -> int: + return 2 + + class _Progress: + def __init__(self, iterable, **kwargs) -> None: + del kwargs + self.iterable = iterable + + def __iter__(self): + return iter(self.iterable) + + def set_postfix_str(self, value: str) -> None: + del value + + def close(self) -> None: + return None + + class _Model: + post_cfg = {"task": "face_detection", "dataset": "widerface"} + postprocessor = SimpleNamespace( + nmsout2eval=lambda *_args, **_kwargs: ([], [[]], [[]]) + ) + preprocess_with_metadata = object() + + def set_postprocess_thresholds(self, **kwargs) -> None: + del kwargs + + def __call__(self, inputs: torch.Tensor) -> torch.Tensor: + return inputs + + def postprocess(self, outputs: torch.Tensor) -> SimpleNamespace: + del outputs + return SimpleNamespace(output=[]) + + batch = ( + torch.zeros((2, 8, 8, 3)), + np.array([[8, 8], [8, 8]]), + [None, None], + ("event", "event"), + ("first.jpg", "second.jpg"), + ) + monkeypatch.setattr(eval_widerface_module, "CustomWiderface", lambda _: _Dataset()) + monkeypatch.setattr( + eval_widerface_module, + "_load_widerface_image_names", + lambda _: {"event": {"first.jpg", "second.jpg"}}, + ) + monkeypatch.setattr( + eval_widerface_module, + "_widerface_difficulty_metadata_ready", + lambda *_, **__: True, + ) + monkeypatch.setattr( + eval_widerface_module, + "_widerface_image_shapes", + lambda *_: [[(8, 8), (8, 8)]], + ) + monkeypatch.setattr( + eval_widerface_module, "get_widerface_loader", lambda *_: [batch] + ) + monkeypatch.setattr(eval_widerface_module, "tqdm", _Progress) + + with pytest.raises(ValueError, match="WiderFace evaluation batch length mismatch"): + eval_widerface_module.eval_widerface(_Model(), "/dataset", batch_size=2) + + +def test_widerface_evaluation_rejects_malformed_difficulty_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Validate metadata before direct evaluation constructs a dataset or loader.""" + + model = SimpleNamespace(post_cfg={"task": "face_detection", "dataset": "widerface"}) + monkeypatch.setattr( + eval_widerface_module, + "_load_widerface_image_names", + lambda _: {"0--Parade": {"sample.jpg"}}, + ) + monkeypatch.setattr( + eval_widerface_module, + "_widerface_difficulty_metadata_ready", + lambda *_, **__: False, + ) + monkeypatch.setattr( + eval_widerface_module, + "CustomWiderface", + lambda _: pytest.fail( + "invalid metadata must not construct a WiderFace dataset" + ), + ) + + with pytest.raises(ValueError, match="metadata is malformed"): + eval_widerface_module.eval_widerface(model, "/dataset", batch_size=1) + + +def test_widerface_evaluation_rejects_image_tree_mismatching_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Require exact event and filename identities before direct inference.""" + + class _Dataset: + samples = [("unused", "0--Parade", "sample.png")] + + model = SimpleNamespace(post_cfg={"task": "face_detection", "dataset": "widerface"}) + monkeypatch.setattr( + eval_widerface_module, + "_load_widerface_image_names", + lambda _: {"0--Parade": {"sample.jpg"}}, + ) + monkeypatch.setattr( + eval_widerface_module, + "_widerface_difficulty_metadata_ready", + lambda *_, **__: True, + ) + monkeypatch.setattr( + eval_widerface_module, "_widerface_image_shapes", lambda *_: [[(8, 8)]] + ) + monkeypatch.setattr(eval_widerface_module, "CustomWiderface", lambda _: _Dataset()) + monkeypatch.setattr( + eval_widerface_module, + "get_widerface_loader", + lambda *_: pytest.fail("mismatched image trees must not create a loader"), + ) + + with pytest.raises(ValueError, match="does not match the validation metadata"): + eval_widerface_module.eval_widerface(model, "/dataset", batch_size=1) + + +def test_widerface_evaluation_rejects_unequal_box_and_score_counts() -> None: + """Do not fabricate or drop face detections during result conversion.""" + + with pytest.raises(ValueError, match="unequal box and score counts"): + eval_widerface_module._boxes_scores_to_prediction([[0, 0, 1, 1]], []) diff --git a/tests/test_face_detection.py b/tests/test_face_detection.py new file mode 100644 index 0000000..58a8091 --- /dev/null +++ b/tests/test_face_detection.py @@ -0,0 +1,123 @@ +"""CPU regression tests for face-detection postprocessing and exports.""" + +from __future__ import annotations + +from typing import Any, cast + +import cv2 +import numpy as np +import pytest +import torch +from mblt_vision.utils.postprocess import build_postprocess +from mblt_vision.utils.postprocess.base import YOLODetectionPostBase +from mblt_vision.utils.postprocess.yolo_anchor_post import YOLOAnchorDetectionPost +from mblt_vision.utils.postprocess.yolo_anchorless_post import ( + YOLOAnchorlessDetectionPost, +) +from mblt_vision.utils.postprocess.yolo_dflfree_post import YOLODFLFreeDetectionPost +from mblt_vision.utils.postprocess.yolo_nmsfree_post import YOLONMSFreeDetectionPost + +from mblt_vision import YOLO11m_face, list_models +from mblt_vision.face_detection import YOLO11m_face as FaceDetectionYOLO11mFace +from mblt_vision.utils.results import Results + + +def _pre_cfg() -> dict[str, Any]: + """Return a representative face preprocessing configuration.""" + + return {"LetterBox": {"img_size": [640, 640]}} + + +def _post_cfg(**overrides: Any) -> dict[str, Any]: + """Return a representative face postprocessing configuration.""" + + return { + "task": "face_detection", + "nl": 3, + "reg_max": 16, + "conf_thres": 0.25, + **overrides, + } + + +@pytest.mark.parametrize( + ("post_cfg", "expected_type"), + [ + ({"nl": 3, "reg_max": 16}, YOLOAnchorlessDetectionPost), + ({"nl": 3, "dflfree": True}, YOLODFLFreeDetectionPost), + ({"nl": 3, "nmsfree": True}, YOLONMSFreeDetectionPost), + ({"anchors": [[10, 13, 16, 30, 33, 23]]}, YOLOAnchorDetectionPost), + ], +) +def test_face_detection_routes_postprocessors( + post_cfg: dict[str, Any], expected_type: type[YOLODetectionPostBase] +) -> None: + """Route every supported face head family to its YOLO postprocessor.""" + + postprocessor = build_postprocess(_pre_cfg(), _post_cfg(**post_cfg)) + + assert isinstance(postprocessor, expected_type) + assert cast(YOLODetectionPostBase, postprocessor).nc == 1 + + +def test_face_detection_exports_and_plot_label( + monkeypatch: pytest.MonkeyPatch, tmp_path: Any +) -> None: + """Preserve legacy exports and render a face-specific detection label.""" + + assert "YOLO11m_face" in list_models("face_detection")["face_detection"] + assert YOLO11m_face is FaceDetectionYOLO11mFace + source_path = tmp_path / "source.jpg" + output_path = tmp_path / "face.jpg" + cv2.imwrite(str(source_path), np.full((32, 32, 3), 255, dtype=np.uint8)) + labels: list[str] = [] + original_put_text = cv2.putText + monkeypatch.setattr( + cv2, + "putText", + lambda *args, **kwargs: ( + labels.append(args[1]), + original_put_text(*args, **kwargs), + )[1], + ) + + Results( + _pre_cfg(), + {"task": "face_detection"}, + [torch.tensor([[1, 1, 16, 16, 0.95, 0]])], + ).plot(str(source_path), save_path=str(output_path)) + + assert output_path.is_file() + assert labels == ["face 95%"] + + +def test_face_detection_non_e2e_converted_and_raw_outputs() -> None: + """Keep converted and raw face heads available through the legacy non-E2E contract.""" + + postprocessor = build_postprocess(_pre_cfg(), _post_cfg(e2e=False)) + converted = torch.tensor( + [ + [ + [10.0, 12.0, 0.9], + [20.0, 18.0, 0.1], + [8.0, 7.0, 0.2], + [9.0, 6.0, 0.3], + [0.8, 0.5, 0.1], + ] + ] + ) + converted_result = postprocessor(converted) + raw_heads = [torch.zeros((1, size, size, 64)) for size in (80, 40, 20)] + raw_heads = [ + tensor + for pair in zip( + raw_heads, [torch.zeros((*tensor.shape[:3], 1)) for tensor in raw_heads] + ) + for tensor in pair + ] + raw_result = postprocessor(raw_heads) + + assert isinstance(converted_result, torch.Tensor) + assert converted_result.shape == (1, 5, 3) + assert isinstance(raw_result, torch.Tensor) + assert raw_result.shape == (1, 5, 8400) diff --git a/tests/test_model_aliasing.py b/tests/test_model_aliasing.py new file mode 100644 index 0000000..80a0e2e --- /dev/null +++ b/tests/test_model_aliasing.py @@ -0,0 +1,40 @@ +"""Tests for vision model name aliasing.""" + +from __future__ import annotations + +import pytest + +from mblt_vision.wrapper import MBLT_Engine + + +@pytest.mark.parametrize( + ("model_name", "expected_yaml"), + [ + ("regnet_x_16gf", "RegNet_X_16GF.yaml"), + ("regnet-x-16gf", "RegNet_X_16GF.yaml"), + ("RegNet_X_16GF.yaml", "RegNet_X_16GF.yaml"), + ("regnet_x_1_6gf", "RegNet_X_1_6GF.yaml"), + ("regnet-x-1-6gf", "RegNet_X_1_6GF.yaml"), + ("resnet50", "ResNet50.yaml"), + ("resnet-50", "ResNet50.yaml"), + ("resnet_50", "ResNet50.yaml"), + ], +) +def test_model_name_aliasing_resolves_precise_separator_matches( + model_name: str, + expected_yaml: str, +) -> None: + """Resolve aliases without collapsing distinct separator boundaries too early.""" + + engine = MBLT_Engine.__new__(MBLT_Engine) + + assert engine.model_name_aliasing(model_name) == expected_yaml + + +def test_model_name_aliasing_reports_compact_ambiguity() -> None: + """Keep compact ambiguous names explicit.""" + + engine = MBLT_Engine.__new__(MBLT_Engine) + + with pytest.raises(ValueError, match="Ambiguous model name"): + engine.model_name_aliasing("regnetx16gf") diff --git a/tests/test_model_metadata.py b/tests/test_model_metadata.py new file mode 100644 index 0000000..de3fb68 --- /dev/null +++ b/tests/test_model_metadata.py @@ -0,0 +1,42 @@ +"""Tests for model dataset metadata and dataset-aware postprocessing.""" + +from __future__ import annotations + +from pathlib import Path + +import mblt_vision +import yaml + +from mblt_vision.wrapper import resolve_model_config + +MODEL_CONFIG_DIR = Path(mblt_vision.__file__).parent / "models" + +DATASETS_BY_TASK = { + "depth_estimation": {"nyu-depth"}, + "face_detection": {"widerface"}, + "image_classification": {"imagenet"}, + "instance_segmentation": {"coco"}, + "object_detection": {"coco"}, + "obb": {"dotav1"}, + "pose_estimation": {"coco"}, + "semantic_segmentation": {"ade20k", "cityscapes"}, +} + + +def test_all_model_variants_declare_a_supported_dataset() -> None: + """Require every resolved model variant to identify its output taxonomy.""" + + checked = 0 + for config_path in sorted(MODEL_CONFIG_DIR.glob("*.yaml")): + full_config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert isinstance(full_config, dict) + for variant in full_config: + config = resolve_model_config(str(config_path), variant) + post_cfg = config["post_cfg"] + task = post_cfg["task"] + assert ( + post_cfg["dataset"] in DATASETS_BY_TASK[task] + ), f"{config_path.name}:{variant}" + checked += 1 + + assert checked > 0 diff --git a/tests/test_mxq_inference.py b/tests/test_mxq_inference.py new file mode 100644 index 0000000..a99a526 --- /dev/null +++ b/tests/test_mxq_inference.py @@ -0,0 +1,24 @@ +"""Optional end-to-end MXQ inference coverage for the standalone package.""" + +from __future__ import annotations + +import pytest + +from mblt_npu.pytest_plugin import NpuParams +from mblt_vision import MBLT_Engine + +pytestmark = [pytest.mark.requires_network, pytest.mark.requires_npu] + + +def test_mxq_classification_runs_with_shared_npu_options( + npu_params: NpuParams, synthetic_image_path +) -> None: + """Load a representative MXQ model using the shared NPU test options.""" + + model = MBLT_Engine(model_cls="resnet50", **npu_params.base) + try: + result = model.postprocess(model(model.preprocess(str(synthetic_image_path)))) + assert result.task == "image_classification" + assert result.acc is not None + finally: + model.dispose() diff --git a/tests/test_onnx_classification.py b/tests/test_onnx_classification.py new file mode 100644 index 0000000..d05d31f --- /dev/null +++ b/tests/test_onnx_classification.py @@ -0,0 +1,35 @@ +"""Tests for vision model ONNX inference on image classification.""" + +from __future__ import annotations + +import pytest + +from mblt_vision.image_classification import AlexNet, CAFormer_B36, YOLO26sCls + +pytestmark = pytest.mark.requires_network + + +@pytest.mark.parametrize( + "model_cls", + [ + AlexNet, + CAFormer_B36, + YOLO26sCls, + ], +) +def test_onnx_classification(model_cls, synthetic_image_path) -> None: + """Run ONNX inference for representative classification models.""" + + model = model_cls(framework="onnx") + + try: + input_img = model.preprocess(str(synthetic_image_path)) + output = model(input_img) + result = model.postprocess(output) + + assert result is not None + assert result.task == "image_classification" + assert result.acc is not None + assert result.output is not None + finally: + model.dispose() diff --git a/tests/test_onnx_yolo.py b/tests/test_onnx_yolo.py new file mode 100644 index 0000000..6b85592 --- /dev/null +++ b/tests/test_onnx_yolo.py @@ -0,0 +1,50 @@ +"""Tests for ONNX inference across representative YOLO postprocess families.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from mblt_vision import MBLT_Engine + +pytestmark = pytest.mark.requires_network + + +@pytest.mark.parametrize( + ("model_cls", "task"), + [ + ("yolov5m", "object_detection"), + ("yolov5m-seg", "instance_segmentation"), + ("yolo11m", "object_detection"), + ("yolo11m-seg", "instance_segmentation"), + ("yolo11m-pose", "pose_estimation"), + ("yolov10m", "object_detection"), + ("yolo26m", "object_detection"), + ("yolo26m-seg", "instance_segmentation"), + ("yolo26m-pose", "pose_estimation"), + ("yolov8m-obb", "obb"), + ("yolo11m-obb", "obb"), + ("yolo26m-obb", "obb"), + ], +) +def test_onnx_yolo_inference( + model_cls: str, task: str, tmp_path: Path, synthetic_image_path: Path +) -> None: + """Run ONNX inference for representative YOLO postprocess families.""" + + save_path = tmp_path / f"{model_cls}_visualization.jpg" + model = MBLT_Engine(model_cls=model_cls, framework="onnx") + + try: + input_img = model.preprocess(str(synthetic_image_path)) + output = model(input_img) + result = model.postprocess(output) + + assert result is not None + assert result.task == task + assert result.output is not None + result.plot(str(synthetic_image_path), save_path=str(save_path)) + assert save_path.is_file() + finally: + model.dispose() diff --git a/tests/test_postprocess_hierarchy.py b/tests/test_postprocess_hierarchy.py new file mode 100644 index 0000000..a03d490 --- /dev/null +++ b/tests/test_postprocess_hierarchy.py @@ -0,0 +1,514 @@ +"""Focused tests for the vision postprocessor class hierarchy and builder.""" + +from __future__ import annotations + +import warnings +from typing import Any, cast + +import pytest +import torch +from mblt_vision.utils.postprocess import build_postprocess +from mblt_vision.utils.postprocess import common as common_module +from mblt_vision.utils.postprocess.base import PostBase, YOLODetectionPostBase +from mblt_vision.utils.postprocess.cls_post import ClsPost +from mblt_vision.utils.postprocess.depth_post import DepthPost +from mblt_vision.utils.postprocess.semantic_seg_post import SemanticSegPost +from mblt_vision.utils.postprocess.yolo_anchor_post import ( + YOLOAnchorDetectionPost, + YOLOAnchorSegPost, +) +from mblt_vision.utils.postprocess.yolo_anchorless_post import ( + YOLOAnchorlessDetectionPost, + YOLOAnchorlessOBBPost, + YOLOAnchorlessPosePost, + YOLOAnchorlessSegPost, +) +from mblt_vision.utils.postprocess.yolo_dflfree_post import ( + YOLODFLFreeDetectionPost, + YOLODFLFreeOBBPost, + YOLODFLFreePosePost, + YOLODFLFreeSegPost, +) +from mblt_vision.utils.postprocess.yolo_nmsfree_post import YOLONMSFreeDetectionPost + + +def test_segmentation_rle_encoding_thresholds_resized_masks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep fractional bilinear boundaries from being truncated during RLE export.""" + + resized_mask = common_module.scale_masks( + torch.tensor([[[1.0, 0.0], [1.0, 0.0]]]), (3, 4) + ) + captured_pixels: list[torch.Tensor] = [] + + def _capture_encode(pixels: torch.Tensor) -> list[list[int]]: + captured_pixels.append(pixels.clone()) + return [[pixels.numel()]] + + monkeypatch.setattr(common_module, "multi_encode", _capture_encode) + + common_module._encode_segmentation_masks(resized_mask) + + assert 0.5 < resized_mask[0, 0, 1] < 1.0 + assert captured_pixels[0].view(1, 4, 3).permute(0, 2, 1).tolist() == [ + [[1, 1, 0, 0]] * 3 + ] + + +def test_classification_postprocessor_rejects_wrong_taxonomy_width() -> None: + """Reject local classification artifacts whose heads do not match ImageNet.""" + + postprocessor = ClsPost({}, {"task": "image_classification", "dataset": "imagenet"}) + + with pytest.raises( + ValueError, + match="Classification output has 999 classes, but dataset 'imagenet' requires 1000", + ): + postprocessor(torch.zeros((1, 999), dtype=torch.float32)) + + +def test_classification_postprocessor_keeps_batched_singleton_outputs() -> None: + """Preserve batch size for local logits shaped [B, C, 1].""" + + postprocessor = ClsPost({}, {"task": "image_classification", "dataset": "imagenet"}) + output = postprocessor(torch.zeros((2, 1000, 1), dtype=torch.float32)) + + assert output.shape == (2, 1000) + + +@pytest.mark.parametrize( + "scores", + [ + torch.tensor([[1.1] + [0.0] * 999]), + torch.tensor([[0.5] + [0.0] * 999]), + ], +) +def test_classification_probability_postprocessor_rejects_invalid_probabilities( + scores: torch.Tensor, +) -> None: + """Validate local artifacts that declare already-softmaxed outputs.""" + + postprocessor = ClsPost( + {}, {"task": "image_classification", "dataset": "imagenet", "softmax": True} + ) + + with pytest.raises(ValueError, match="probability outputs"): + postprocessor(scores) + + +@pytest.mark.parametrize("invalid_score", [float("nan"), float("inf")]) +def test_classification_postprocessor_rejects_nonfinite_scores( + invalid_score: float, +) -> None: + """Do not convert malformed classification logits into plausible predictions.""" + + scores = torch.zeros((1, 1000), dtype=torch.float32) + scores[0, 0] = invalid_score + + with pytest.raises(ValueError, match="scores must all be finite"): + ClsPost({}, {"task": "image_classification", "dataset": "imagenet"})(scores) + + +@pytest.mark.parametrize("invalid_class", [-1.0, 1.5, 2.0]) +def test_already_decoded_detections_reject_invalid_task_class_ids( + invalid_class: float, +) -> None: + """Validate local decoded detection outputs against the configured taxonomy.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + detections = torch.tensor( + [[[0.0, 0.0, 1.0, 1.0, 0.9, invalid_class]]], dtype=torch.float32 + ) + + with pytest.raises(ValueError, match="Decoded detection class IDs"): + postprocessor._final_detection_batches(detections) + + +@pytest.mark.parametrize( + ("column", "invalid_value"), + [(0, float("nan")), (4, float("inf")), (5, float("-inf"))], +) +def test_already_decoded_detections_reject_nonfinite_rows( + column: int, invalid_value: float +) -> None: + """Reject malformed coordinates, scores, and labels before filtering them.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + detections = torch.tensor([[[0.0, 0.0, 1.0, 1.0, 0.9, 1.0]]], dtype=torch.float32) + detections[0, 0, column] = invalid_value + + with pytest.raises(ValueError, match="finite values"): + postprocessor._final_detection_batches(detections) + + +@pytest.mark.parametrize("invalid_score", [-0.1, 1.1]) +def test_already_decoded_detections_reject_invalid_confidence( + invalid_score: float, +) -> None: + """Reject finite but impossible confidence values before thresholding.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + detections = torch.tensor( + [[[0.0, 0.0, 1.0, 1.0, invalid_score, 1.0]]], dtype=torch.float32 + ) + + with pytest.raises(ValueError, match=r"confidence values must be in \[0, 1\]"): + postprocessor._final_detection_batches(detections) + + +@pytest.mark.parametrize("coordinates", [(1.0, 0.0, 1.0, 2.0), (0.0, 2.0, 1.0, 2.0)]) +def test_already_decoded_detections_reject_nonpositive_box_area( + coordinates: tuple[float, float, float, float], +) -> None: + """Reject degenerate decoded XYXY boxes before they reach evaluation.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + detections = torch.tensor([[[*coordinates, 0.9, 1.0]]], dtype=torch.float32) + + with pytest.raises(ValueError, match="positive xyxy area"): + postprocessor._final_detection_batches(detections) + + +def test_already_decoded_detections_ignore_degenerate_padding_rows() -> None: + """Validate geometry only after low-confidence fixed-size padding is removed.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + detections = torch.tensor([[[0.0, 0.0, 0.0, 0.0, 0.1, 1.0]]], dtype=torch.float32) + + assert postprocessor._final_detection_batches(detections)[0].shape == (0, 6) + + +def test_already_decoded_obb_uses_width_height_geometry() -> None: + """Do not interpret OBB center/size fields as XYXY corners.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + postprocessor.task = "obb" + detections = torch.tensor( + [[[500.0, 20.0, 100.0, 10.0, 0.9, 1.0]]], dtype=torch.float32 + ) + + assert postprocessor._final_detection_batches(detections)[0].shape == (1, 6) + + +@pytest.mark.parametrize("prototype_first", [False, True]) +def test_decoded_segmentation_propagates_invalid_mask_prototypes( + prototype_first: bool, +) -> None: + """Do not discard malformed required prototype tensors as unrelated outputs.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + postprocessor.n_extra = 2 + postprocessor.task = "instance_segmentation" + detections = torch.tensor( + [[[0.0, 0.0, 1.0, 1.0, 0.9, 1.0, 0.0, 0.0]]], dtype=torch.float32 + ) + invalid_proto = torch.full((1, 2, 2, 2), float("nan")) + outputs = ( + [invalid_proto, detections] if prototype_first else [detections, invalid_proto] + ) + + with pytest.raises( + ValueError, match="Mask prototype tensor must contain only finite" + ): + postprocessor.extract_final_outputs(outputs) + + +def test_decoded_segmentation_requires_mask_prototypes() -> None: + """Reject decoded segmentation artifacts that cannot produce instance masks.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + postprocessor.conf_thres = 0.25 + postprocessor.nc = 2 + postprocessor.n_extra = 2 + postprocessor.task = "instance_segmentation" + detections = torch.tensor( + [[[0.0, 0.0, 1.0, 1.0, 0.9, 1.0, 0.0, 0.0]]], dtype=torch.float32 + ) + + with pytest.raises(ValueError, match="require a mask prototype tensor"): + postprocessor.extract_final_outputs([detections]) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")]) +def test_detection_postprocessor_rejects_nonfinite_raw_heads( + invalid_value: float, +) -> None: + """Reject malformed raw detector heads before they reach decoding and NMS.""" + + postprocessor = cast( + Any, YOLOAnchorlessDetectionPost.__new__(YOLOAnchorlessDetectionPost) + ) + postprocessor.device = torch.device("cpu") + raw_head = torch.zeros((1, 6, 2, 2), dtype=torch.float32) + raw_head[0, 0, 0, 0] = invalid_value + + with pytest.raises( + ValueError, match="Detection output tensors must contain only finite" + ): + postprocessor.check_input(raw_head) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")]) +def test_detection_postprocessor_rejects_nonfinite_mask_prototypes( + invalid_value: float, +) -> None: + """Reject invalid mask prototypes from already-decoded segmentation artifacts.""" + + postprocessor = cast(Any, YOLOAnchorlessSegPost.__new__(YOLOAnchorlessSegPost)) + postprocessor.device = torch.device("cpu") + postprocessor.n_extra = 2 + prototype = torch.zeros((1, 2, 2, 2)) + prototype[0, 0, 0, 0] = invalid_value + + with pytest.raises( + ValueError, match="Mask prototype tensor must contain only finite" + ): + postprocessor._normalize_proto_batch(prototype) + + +@pytest.mark.parametrize( + ("post_cfg", "expected_type"), + [ + ( + {"task": "object_detection", "anchors": [[10, 13, 16, 30, 33, 23]]}, + YOLOAnchorDetectionPost, + ), + ( + {"task": "object_detection", "nl": 3, "reg_max": 16}, + YOLOAnchorlessDetectionPost, + ), + ( + {"task": "object_detection", "nl": 3, "dflfree": True}, + YOLODFLFreeDetectionPost, + ), + ( + {"task": "object_detection", "nl": 3, "reg_max": 16, "nmsfree": True}, + YOLONMSFreeDetectionPost, + ), + ( + { + "task": "instance_segmentation", + "anchors": [[10, 13, 16, 30, 33, 23]], + "n_extra": 32, + }, + YOLOAnchorSegPost, + ), + ( + {"task": "instance_segmentation", "nl": 3, "reg_max": 16, "n_extra": 32}, + YOLOAnchorlessSegPost, + ), + ( + {"task": "instance_segmentation", "nl": 3, "dflfree": True, "n_extra": 32}, + YOLODFLFreeSegPost, + ), + ( + {"task": "pose_estimation", "nl": 3, "reg_max": 16, "n_extra": 51}, + YOLOAnchorlessPosePost, + ), + ( + {"task": "pose_estimation", "nl": 3, "dflfree": True, "n_extra": 51}, + YOLODFLFreePosePost, + ), + ({"task": "obb", "nl": 3, "reg_max": 16, "n_extra": 1}, YOLOAnchorlessOBBPost), + ({"task": "obb", "nl": 3, "dflfree": True, "n_extra": 1}, YOLODFLFreeOBBPost), + ], +) +def test_builder_routes_detection_backends_without_warnings( + post_cfg: dict[str, Any], expected_type: type[YOLODetectionPostBase] +) -> None: + """Build every detection family through canonical warning-free imports.""" + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + postprocessor = build_postprocess( + {"LetterBox": {"img_size": [640, 640]}}, post_cfg + ) + assert type(postprocessor) is expected_type + + +@pytest.mark.parametrize( + ("pre_cfg", "post_cfg", "expected_type"), + [ + ({}, {"task": "image_classification"}, ClsPost), + ({"LetterBox": {"img_size": [8, 8]}}, {"task": "depth_estimation"}, DepthPost), + ( + {"LetterBox": {"img_size": [8, 8]}}, + {"task": "semantic_segmentation", "dataset": "ade20k"}, + SemanticSegPost, + ), + ], +) +def test_builder_keeps_non_detection_routing_warning_free( + pre_cfg: dict[str, Any], post_cfg: dict[str, Any], expected_type: type[PostBase] +) -> None: + """Preserve classification and dense prediction builder routes.""" + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + postprocessor = build_postprocess(pre_cfg, post_cfg) + assert type(postprocessor) is expected_type + + +@pytest.mark.parametrize( + ("postprocessor_type", "input_shape", "n_extra", "output_width"), + [ + (YOLONMSFreeDetectionPost, (5, 1), 0, 6), + (YOLODFLFreeDetectionPost, (5, 1), 0, 5), + (YOLODFLFreePosePost, (56, 1), 51, 56), + (YOLODFLFreeOBBPost, (6, 1), 1, 6), + ], +) +def test_empty_anchorless_outputs_preserve_input_device_and_dtype( + postprocessor_type: type[YOLODetectionPostBase], + input_shape: tuple[int, int], + n_extra: int, + output_width: int, +) -> None: + """Create empty decoded rows from the selected input tensor.""" + + postprocessor = cast(Any, postprocessor_type.__new__(postprocessor_type)) + postprocessor.nc = 1 + postprocessor.n_extra = n_extra + postprocessor.inv_conf_thres = 1.0 + input_tensor = torch.zeros(input_shape, dtype=torch.float64) + + output = postprocessor.process_box_cls(input_tensor) + + assert output.shape == (0, output_width) + assert output.device == input_tensor.device + assert output.dtype == input_tensor.dtype + + +def test_empty_anchor_outputs_preserve_input_device_and_dtype() -> None: + """Keep converted and decoded empty anchor rows on the input device.""" + + postprocessor = YOLOAnchorDetectionPost.__new__(YOLOAnchorDetectionPost) + postprocessor.nc = 1 + postprocessor.n_extra = 0 + postprocessor.no = 6 + postprocessor.inv_conf_thres = 1.0 + postprocessor.conf_thres = 1.0 + decoded_input = torch.zeros((1, 6), dtype=torch.float64) + converted_input = torch.zeros((1, 1, 6), dtype=torch.float64) + + decoded_output = postprocessor.process_box_cls(decoded_input) + converted_output = postprocessor.filter_conversion(converted_input)[0] + + for output, input_tensor in ( + (decoded_output, decoded_input), + (converted_output, converted_input), + ): + assert output.device == input_tensor.device + assert output.dtype == input_tensor.dtype + + +def test_detection_postprocessor_rejects_missing_head_count() -> None: + """Raise a stable runtime error when anchorless metadata omits nl.""" + + with pytest.raises(ValueError, match="nl should be provided in post_cfg"): + YOLOAnchorlessDetectionPost( + {"LetterBox": {"img_size": [640, 640]}}, + {"task": "object_detection", "dataset": "coco"}, + ) + + +@pytest.mark.parametrize( + ("post_cfg", "channels"), + [ + ({"task": "object_detection", "nl": 3, "reg_max": 16}, [64] * 3 + [80] * 2), + ( + {"task": "instance_segmentation", "nl": 3, "reg_max": 16, "n_extra": 32}, + [32] * 4 + [64] * 3 + [80] * 2, + ), + ( + {"task": "pose_estimation", "nl": 3, "reg_max": 16, "n_extra": 51}, + [64] * 3 + [1] * 2 + [51] * 3, + ), + ( + {"task": "obb", "nl": 3, "reg_max": 16, "n_extra": 1}, + [64] * 3 + [15] * 2 + [1] * 3, + ), + ({"task": "object_detection", "nl": 3, "dflfree": True}, [4] * 3 + [80] * 2), + ( + {"task": "instance_segmentation", "nl": 3, "dflfree": True, "n_extra": 32}, + [32] * 4 + [4] * 3 + [80] * 2, + ), + ( + {"task": "pose_estimation", "nl": 3, "dflfree": True, "n_extra": 51}, + [4] * 3 + [1] * 2 + [51] * 3, + ), + ( + {"task": "obb", "nl": 3, "dflfree": True, "n_extra": 1}, + [4] * 3 + [15] * 2 + [1] * 3, + ), + ], +) +def test_split_head_postprocessors_reject_incomplete_head_groups( + post_cfg: dict[str, Any], channels: list[int] +) -> None: + """Reject malformed raw output sets before zip can discard a detection scale.""" + + postprocessor = build_postprocess({"LetterBox": {"img_size": [64, 64]}}, post_cfg) + raw_outputs = [torch.zeros((1, 2, 2, channel)) for channel in channels] + + with pytest.raises( + ValueError, + match=r"Incomplete split-head outputs: expected 3 heads per group.*classification=2", + ): + postprocessor.rearrange(raw_outputs) # type: ignore[attr-defined] + + +def test_segmentation_postprocessor_rejects_mismatched_prototype_batch() -> None: + """Do not silently drop detections when prototypes have fewer images.""" + + postprocessor = YOLODFLFreeSegPost( + {"LetterBox": {"img_size": [640, 640]}}, + {"task": "instance_segmentation", "nl": 3, "n_extra": 32}, + ) + detections = [torch.empty((0, 38)), torch.empty((0, 38))] + prototypes = torch.empty((1, 32, 160, 160)) + + with pytest.raises( + ValueError, + match="Detection and prototype batch sizes must match.*2 detections and 1 prototypes", + ): + postprocessor.masking(detections, prototypes) diff --git a/tests/test_preprocess_validation.py b/tests/test_preprocess_validation.py new file mode 100644 index 0000000..9ed34d1 --- /dev/null +++ b/tests/test_preprocess_validation.py @@ -0,0 +1,220 @@ +"""Runtime validation tests for shared Vision preprocessing.""" + +from __future__ import annotations + +import subprocess +import sys + +import numpy as np +import pytest +import torch +from PIL import Image +from mblt_vision.utils.preprocess.center_crop import CenterCrop +from mblt_vision.utils.preprocess.letterbox import LetterBox +from mblt_vision.utils.preprocess.normalize import Normalize +from mblt_vision.utils.preprocess.order import SetOrder +from mblt_vision.utils.preprocess.reader import Reader +from mblt_vision.utils.preprocess.resize import Resize + + +@pytest.mark.parametrize( + "operation", [CenterCrop(2), Resize(2, "bilinear"), LetterBox([2, 2])] +) +def test_image_size_operations_reject_nonpositive_sizes(operation: object) -> None: + """Keep size validation active when Python assertions are optimized away.""" + + operation_type = type(operation) + if operation_type is Resize: + with pytest.raises(ValueError, match="positive"): + Resize(0, "bilinear") + elif operation_type is CenterCrop: + with pytest.raises(ValueError, match="positive"): + CenterCrop([0, 2]) + else: + with pytest.raises(ValueError, match="positive"): + LetterBox([0, 2]) + + +@pytest.mark.parametrize("size", [[1], [1, 2, 3], [1.5, 2], "2"]) +def test_center_crop_rejects_invalid_size_values(size: object) -> None: + """Reject invalid crop types and dimensions before processing.""" + + with pytest.raises((TypeError, ValueError)): + CenterCrop(size) # type: ignore[arg-type] + + +def test_set_order_rejects_ambiguous_channel_layout() -> None: + """Do not guess between CHW and HWC when both ends have three channels.""" + + with pytest.raises(ValueError, match="ambiguous"): + SetOrder("HWC")(np.zeros((3, 5, 3), dtype=np.uint8)) + + +def test_normalize_rejects_ambiguous_channel_layout() -> None: + """Do not silently apply HWC statistics to potentially CHW data.""" + + with pytest.raises(ValueError, match="ambiguous"): + Normalize("torch")(np.zeros((3, 640, 3), dtype=np.uint8)) + + +def test_normalize_supports_channel_first_set_order_output() -> None: + """Broadcast normalization statistics per channel after a CHW conversion.""" + + image = np.array( + [ + [[255, 0, 127], [0, 255, 127], [127, 127, 255], [64, 32, 16], [1, 2, 3]], + [[4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15], [16, 17, 18]], + [[19, 20, 21], [22, 23, 24], [25, 26, 27], [28, 29, 30], [31, 32, 33]], + [[34, 35, 36], [37, 38, 39], [40, 41, 42], [43, 44, 45], [46, 47, 48]], + ], + dtype=np.uint8, + ) + + normalized = Normalize("torch")(SetOrder("CHW")(image)) + expected = ( + image.transpose(2, 0, 1) / 255.0 + - np.array([0.485, 0.456, 0.406])[:, None, None] + ) / np.array([0.229, 0.224, 0.225])[:, None, None] + + assert normalized.shape == (3, 4, 5) + np.testing.assert_allclose( + normalized, expected.astype(np.float32), rtol=1e-6, atol=1e-6 + ) + + +@pytest.mark.parametrize( + ("operation", "output_type"), + [ + (Reader("numpy"), np.ndarray), + (Reader("pil"), Image.Image), + (Normalize("cv"), np.ndarray), + (CenterCrop([2, 3]), np.ndarray), + (LetterBox([2, 3]), torch.Tensor), + ], +) +def test_preprocessors_accept_grad_tracking_tensors( + operation: object, output_type: type[object] +) -> None: + """Detach tensors before preprocessing converts them to NumPy arrays.""" + + image = torch.ones((4, 5, 3), requires_grad=True) + + assert isinstance(operation(image), output_type) # type: ignore[operator] + + +def test_reader_pil_scales_normalized_float_arrays() -> None: + """Convert normalized float RGB inputs deliberately instead of truncating them.""" + + image = np.full((2, 3, 3), 0.5, dtype=np.float32) + + converted = np.asarray(Reader("pil")(image)) + + np.testing.assert_array_equal(converted, np.full((2, 3, 3), 128, dtype=np.uint8)) + + +def test_letterbox_scales_normalized_float_arrays() -> None: + """Do not truncate normalized float RGB input before LetterBox resizing.""" + + image = np.full((2, 2, 3), 0.5, dtype=np.float32) + + converted = LetterBox([4, 4])(image) + + assert converted.dtype == torch.uint8 + assert torch.equal(converted, torch.full((4, 4, 3), 128, dtype=torch.uint8)) + + +@pytest.mark.parametrize( + "image", + [ + np.full((2, 3, 3), -0.1, dtype=np.float32), + np.full((2, 3, 3), 256.0, dtype=np.float32), + np.full((2, 3, 3), np.nan, dtype=np.float32), + ], +) +def test_reader_pil_rejects_invalid_float_image_arrays(image: np.ndarray) -> None: + """Do not silently cast invalid float RGB values to byte images.""" + + with pytest.raises(ValueError, match=r"finite RGB values|\[0, 1\] or \[0, 255\]"): + Reader("pil")(image) + + +@pytest.mark.parametrize( + "image", + [ + np.full((2, 3, 3), -0.1, dtype=np.float32), + np.full((2, 3, 3), 256.0, dtype=np.float32), + np.full((2, 3, 3), np.nan, dtype=np.float32), + ], +) +def test_letterbox_rejects_invalid_float_image_arrays(image: np.ndarray) -> None: + """Do not wrap out-of-range or non-finite float pixels to byte RGB.""" + + with pytest.raises(ValueError, match=r"finite RGB values|\[0, 1\] or \[0, 255\]"): + LetterBox([4, 4])(image) + + +@pytest.mark.parametrize( + ("factory", "expected"), + [ + (lambda: Reader("unsupported"), ValueError), + (lambda: Normalize("unsupported"), ValueError), + (lambda: SetOrder("unsupported"), ValueError), + ], +) +def test_preprocessing_styles_raise_explicit_errors( + factory, expected: type[Exception] +) -> None: + """Use explicit configuration errors instead of runtime assertions.""" + + with pytest.raises(expected): + factory() + + +@pytest.mark.parametrize("interpolation", ["box", "hamming", "lanczos"]) +@pytest.mark.parametrize( + "image", [np.zeros((2, 3, 3), dtype=np.uint8), torch.zeros((3, 2, 3))] +) +def test_resize_rejects_pil_only_modes_for_tensor_backed_inputs( + interpolation: str, image: np.ndarray | torch.Tensor +) -> None: + """Fail clearly instead of passing unsupported modes to torch.interpolate.""" + + with pytest.raises(ValueError, match="supported only for PIL images"): + Resize([4, 6], interpolation)(image) + + +@pytest.mark.parametrize("interpolation", ["box", "hamming", "lanczos"]) +def test_resize_keeps_pil_only_modes_available_for_pil_images( + interpolation: str, +) -> None: + """Retain the documented PIL resize modes for PIL callers.""" + + image = Image.fromarray(np.zeros((2, 3, 3), dtype=np.uint8)) + + resized = Resize([4, 6], interpolation)(image) + + assert isinstance(resized, Image.Image) + assert resized.size == (6, 4) + + +def test_runtime_validation_survives_optimized_python() -> None: + """Keep configuration checks active under ``python -O``.""" + + code = """ +from mblt_vision.utils.postprocess import build_postprocess +try: + build_postprocess( + {"LetterBox": {"img_size": [640, 640]}}, + { + "task": "object_detection", + "dataset": "coco", + "nl": 3, + "reg_max": 16, + "conf_thres": 0, + }, + ) +except ValueError: + raise SystemExit(0) +raise SystemExit(1) +""" + subprocess.run([sys.executable, "-O", "-c", code], check=True) diff --git a/tests/test_results.py b/tests/test_results.py new file mode 100644 index 0000000..0b4d446 --- /dev/null +++ b/tests/test_results.py @@ -0,0 +1,286 @@ +"""Tests for vision plotting helpers.""" + +from __future__ import annotations + +from pathlib import Path + +import cv2 +import numpy as np +import pytest +import torch + +import mblt_vision.utils.results as results_module +from mblt_vision.utils.datasets import get_dotav1_palette +from mblt_vision.utils.results import Results +from mblt_vision.utils.types import NestedListTensorLike + + +def test_image_classification_plot_saves_without_gui_cleanup( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Save classification results without requiring OpenCV GUI support.""" + + source_path = tmp_path / "source.jpg" + save_path = tmp_path / "result.jpg" + image = np.full((32, 32, 3), 255, dtype=np.uint8) + cv2.imwrite(str(source_path), image) + + def _raise_destroy_all_windows() -> None: + raise cv2.error("cvDestroyAllWindows is unavailable") + + monkeypatch.setattr(cv2, "destroyAllWindows", _raise_destroy_all_windows) + + output = torch.zeros(1000) + output[980] = 0.9 + result = Results({}, {"task": "image_classification"}, output) + + plotted = result.plot(str(source_path), str(save_path), topk=1) + + assert plotted is not None + assert save_path.is_file() + + +def test_image_classification_plot_rejects_missing_output() -> None: + """Raise a runtime validation error when classification output is absent.""" + + result = Results({}, {"task": "image_classification"}, torch.zeros(1)) + result.acc = None + + with pytest.raises(ValueError, match="No accuracy output found"): + result.plot(np.zeros((8, 8, 3), dtype=np.uint8), topk=1) + + +def test_instance_segmentation_plot_supports_nonzero_coco_labels( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Plot segmentation results when detections use regular COCO class ids.""" + + source_path = tmp_path / "source.jpg" + save_path = tmp_path / "segmentation.jpg" + image = np.full((32, 32, 3), 255, dtype=np.uint8) + cv2.imwrite(str(source_path), image) + + monkeypatch.setattr(results_module, "scale_boxes", lambda *args, **kwargs: args[1]) + monkeypatch.setattr(results_module, "scale_masks", lambda mask, img0_shape: mask) + monkeypatch.setattr(results_module, "crop_mask", lambda mask, boxes: mask) + + box_cls = torch.tensor([[4.0, 6.0, 20.0, 24.0, 0.9, 45.0]], dtype=torch.float32) + mask = torch.ones((1, 32, 32), dtype=torch.float32) + result = Results( + {"LetterBox": {"img_size": (32, 32)}}, + {"task": "instance_segmentation"}, + [[box_cls, mask]], + ) + + plotted = result.plot(str(source_path), str(save_path)) + + assert plotted is not None + assert save_path.is_file() + + +def test_object_detection_plot_preserves_raw_output_coordinates() -> None: + """Inverse scaling for plotting must not mutate stored postprocess boxes.""" + + box_cls = torch.tensor([[40.0, 60.0, 160.0, 180.0, 0.9, 0.0]], dtype=torch.float32) + result = Results( + {"LetterBox": {"img_size": (200, 200)}}, + {"task": "object_detection"}, + [box_cls], + ) + expected = box_cls.clone() + image = np.zeros((100, 100, 3), dtype=np.uint8) + + first = result.plot(image) + second = result.plot(image) + + torch.testing.assert_close(box_cls, expected) + torch.testing.assert_close(result._box_cls_tensor(), expected) + assert first is not None + assert second is not None + assert np.array_equal(first, second) + + +def test_plot_converts_rgb_array_source_to_bgr() -> None: + """Return OpenCV-order output when callers provide an RGB NumPy image.""" + + result = Results( + {"LetterBox": {"img_size": (1, 1)}}, + {"task": "object_detection"}, + [torch.zeros((0, 6), dtype=torch.float32)], + ) + rgb = np.array([[[255, 0, 0]]], dtype=np.uint8) + + plotted = result.plot(rgb) + + assert plotted is not None + assert np.array_equal(plotted, np.array([[[0, 0, 255]]], dtype=np.uint8)) + assert np.array_equal(rgb, np.array([[[255, 0, 0]]], dtype=np.uint8)) + + +def test_plot_normalizes_float_rgb_array_source_to_uint8_bgr() -> None: + """Plot normalized RGB arrays without rendering a near-black background.""" + + result = Results( + {"LetterBox": {"img_size": (1, 1)}}, + {"task": "object_detection"}, + [torch.zeros((0, 6), dtype=torch.float32)], + ) + rgb = np.array([[[0.5, 0.0, 0.0]]], dtype=np.float32) + + plotted = result.plot(rgb) + + assert plotted is not None + assert np.array_equal(plotted, np.array([[[0, 0, 128]]], dtype=np.uint8)) + assert np.array_equal(rgb, np.array([[[0.5, 0.0, 0.0]]], dtype=np.float32)) + + +def test_pose_plot_hides_low_visibility_keypoints_and_limbs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Draw pose elements only when every required keypoint is visible.""" + + circles: list[tuple[int, int]] = [] + limbs: list[tuple[tuple[int, int], tuple[int, int]]] = [] + + def capture_circle( + image: np.ndarray, center: tuple[int, int], *args: object, **kwargs: object + ) -> np.ndarray: + del args, kwargs + circles.append(center) + return image + + def capture_line( + image: np.ndarray, + point1: tuple[int, int], + point2: tuple[int, int], + *args: object, + **kwargs: object, + ) -> np.ndarray: + del args, kwargs + limbs.append((point1, point2)) + return image + + monkeypatch.setattr(cv2, "circle", capture_circle) + monkeypatch.setattr(cv2, "line", capture_line) + keypoints = torch.zeros((17, 3), dtype=torch.float32) + keypoints[0] = torch.tensor([20.0, 20.0, 0.9]) + keypoints[1] = torch.tensor([40.0, 20.0, 0.9]) + box_cls = torch.cat( + [ + torch.tensor([[10.0, 10.0, 50.0, 50.0, 0.9, 0.0]]), + keypoints.reshape(1, -1), + ], + dim=1, + ) + result = Results( + {"LetterBox": {"img_size": (100, 100)}}, + {"task": "pose_estimation", "n_extra": 51}, + [box_cls], + conf_thres=0.5, + ) + + result.plot(np.zeros((100, 100, 3), dtype=np.uint8)) + + assert circles == [(20, 20), (40, 20)] + assert limbs == [((20, 20), (40, 20))] + + +def test_obb_plot_uses_dotav1_palette(monkeypatch: pytest.MonkeyPatch) -> None: + """Plot DOTAv1 boxes without consulting the COCO palette.""" + + def _reject_coco_palette(label_idx: int) -> tuple[int, int, int]: + raise AssertionError( + f"Unexpected COCO palette lookup for DOTAv1 class {label_idx}." + ) + + monkeypatch.setattr(results_module, "get_coco_det_palette", _reject_coco_palette) + box_cls = torch.tensor( + [[16.0, 16.0, 10.0, 10.0, 0.9, 2.0, 0.0]], dtype=torch.float32 + ) + result = Results( + {"LetterBox": {"img_size": (32, 32)}}, + {"task": "obb"}, + [box_cls], + ) + + plotted = result.plot(np.zeros((32, 32, 3), dtype=np.uint8)) + + assert plotted is not None + assert np.any(np.all(plotted == get_dotav1_palette(2), axis=2)) + + +def test_dotav1_palette_wraps_class_indices() -> None: + """Match the modulo behavior used by the other visualization palettes.""" + + assert get_dotav1_palette(15) == get_dotav1_palette(0) + + +def test_results_accept_path_and_basename_save_path( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Accept pathlib inputs and avoid creating an empty parent directory.""" + + source_path = tmp_path / "source.jpg" + assert cv2.imwrite(str(source_path), np.zeros((8, 8, 3), dtype=np.uint8)) + monkeypatch.chdir(tmp_path) + result = Results( + {}, {"task": "image_classification"}, torch.arange(3, dtype=torch.float32) + ) + + result.plot(source_path, Path("result.jpg"), topk=10) + + assert (tmp_path / "result.jpg").is_file() + + +def test_results_report_failed_image_write(monkeypatch: pytest.MonkeyPatch) -> None: + """Raise OSError when OpenCV cannot encode or write the requested image.""" + + monkeypatch.setattr(cv2, "imwrite", lambda *args, **kwargs: False) + result = Results({}, {"task": "image_classification"}, torch.ones(2)) + + with pytest.raises(OSError, match="Failed to write"): + result.plot(np.zeros((8, 8, 3), dtype=np.uint8), "result.jpg", topk=1) + + +@pytest.mark.parametrize("topk", [0, -1, 1.5, True]) +def test_results_validate_classification_topk(topk: object) -> None: + """Reject invalid classification Top-K values.""" + + result = Results({}, {"task": "image_classification"}, torch.ones(2)) + with pytest.raises((TypeError, ValueError)): + result.plot(np.zeros((8, 8, 3), dtype=np.uint8), topk=topk) + + +@pytest.mark.parametrize( + ("task", "output"), + [ + ("object_detection", []), + ("instance_segmentation", []), + ("instance_segmentation", [[]]), + ], +) +def test_results_reject_empty_structured_outputs( + task: str, output: NestedListTensorLike +) -> None: + """Validate structured result containers before indexing.""" + + with pytest.raises(ValueError): + Results({}, {"task": task}, output) + + +def test_results_normalize_obb_task_and_semantic_taxonomy_case() -> None: + """Normalize the OBB task and semantic palette taxonomy casing.""" + + obb = Results({}, {"task": "obb"}, [torch.zeros((0, 7))]) + semantic = Results( + {}, + {"task": "semantic_segmentation", "dataset": "CityScapes"}, + np.zeros((1, 4, 4), dtype=np.uint8), + ) + + assert obb.task == "obb" + plotted = semantic.plot(np.zeros((4, 4, 3), dtype=np.uint8)) + assert plotted is not None + assert plotted.shape == (4, 4, 3) diff --git a/tests/test_semantic_segmentation.py b/tests/test_semantic_segmentation.py new file mode 100644 index 0000000..3bb591d --- /dev/null +++ b/tests/test_semantic_segmentation.py @@ -0,0 +1,177 @@ +"""Tests for YOLO26 ADE20K semantic segmentation support.""" + +from __future__ import annotations + + +import pytest +import torch +import numpy as np +from mblt_vision.utils.evaluation.eval_ade20k import SemanticMetricAccumulator +from mblt_vision.utils.postprocess import SemanticSegPost + + +def test_semantic_postprocess_supports_logits_and_baked_maps() -> None: + """Convert logits or baked maps to input-sized integer class maps.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [4, 4]}}, + {"task": "semantic_segmentation", "dataset": "ade20k"}, + ) + logits = torch.zeros((1, 150, 2, 2)) + logits[:, 7] = 1.0 + result = post(logits) + assert isinstance(result, torch.Tensor) + assert result.shape == (1, 4, 4) + assert result.dtype == torch.int64 + assert torch.equal(result, torch.full((1, 4, 4), 7)) + + baked = post(torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])) + assert isinstance(baked, torch.Tensor) + assert baked.shape == (1, 4, 4) + assert set(baked.unique().tolist()) == {1, 2, 3, 4} + + with pytest.raises(ValueError, match=r"expects \[B, 150, H, W\]"): + post(torch.zeros((1, 19, 4, 4))) + with pytest.raises(ValueError, match=r"must be in \[0, 149\]"): + post(torch.full((1, 4, 4), 150)) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf"), float("-inf")]) +def test_semantic_postprocess_rejects_non_finite_baked_class_ids( + invalid_value: float, +) -> None: + """Reject non-finite baked class IDs before converting their dtype.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [4, 4]}}, + {"task": "semantic_segmentation", "dataset": "cityscapes"}, + ) + class_map = torch.zeros((1, 2, 2), dtype=torch.float32) + class_map[0, 0, 0] = invalid_value + + with pytest.raises(ValueError, match="must be finite"): + post(class_map) + + +def test_semantic_postprocess_rejects_fractional_baked_class_ids() -> None: + """Reject fractional baked class IDs instead of silently truncating them.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [4, 4]}}, + {"task": "semantic_segmentation", "dataset": "cityscapes"}, + ) + + with pytest.raises(ValueError, match="must be integer-valued"): + post(torch.tensor([[[0.0, 1.9], [2.0, 18.0]]])) + + +def test_semantic_postprocess_reports_semantic_letterbox_errors() -> None: + """Use semantic task labels when validating preprocessing configuration.""" + + with pytest.raises( + ValueError, + match=r"Semantic segmentation requires a LetterBox configuration in pre_cfg", + ): + SemanticSegPost({}, {"task": "semantic_segmentation", "dataset": "ade20k"}) + + +def test_semantic_postprocess_supports_mxq_hwc_and_batched_nhwc_logits() -> None: + """Convert Cityscapes MXQ channel-last logits before choosing class maps.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [4, 8]}}, + {"task": "semantic_segmentation", "dataset": "cityscapes"}, + ) + hwc_logits = torch.zeros((4, 8, 19)) + hwc_logits[..., 6] = 2.0 + hwc_result = post(hwc_logits) + assert isinstance(hwc_result, torch.Tensor) + assert hwc_result.shape == (1, 4, 8) + assert torch.equal(hwc_result, torch.full((1, 4, 8), 6)) + + nhwc_logits = torch.zeros((2, 2, 4, 19)) + nhwc_logits[0, ..., 3] = 2.0 + nhwc_logits[1, ..., 9] = 2.0 + nhwc_result = post(nhwc_logits) + assert isinstance(nhwc_result, torch.Tensor) + assert nhwc_result.shape == (2, 4, 8) + assert torch.equal(nhwc_result[0], torch.full((4, 8), 3)) + assert torch.equal(nhwc_result[1], torch.full((4, 8), 9)) + + low_resolution_hwc_logits = torch.zeros((2, 4, 19)) + low_resolution_hwc_logits[..., 12] = 2.0 + low_resolution_result = post(low_resolution_hwc_logits) + assert isinstance(low_resolution_result, torch.Tensor) + assert low_resolution_result.shape == (1, 4, 8) + assert torch.equal(low_resolution_result, torch.full((1, 4, 8), 12)) + + with pytest.raises(ValueError, match=r"expects \[B, 19, H, W\]"): + post(torch.zeros((1, 4, 8, 18))) + with pytest.raises(ValueError, match=r"expects \[H, W, 19\] MXQ logits"): + post(torch.zeros((4, 8, 18))) + + baked_width_matches_nc = torch.arange(19).reshape(1, 1, 19) + baked_result = post(baked_width_matches_nc) + assert isinstance(baked_result, torch.Tensor) + assert baked_result.shape == (1, 4, 8) + assert torch.equal(baked_result[0, 0], torch.tensor([0, 2, 4, 7, 9, 11, 14, 16])) + + +@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf")]) +def test_semantic_postprocess_rejects_nonfinite_logits(invalid_value: float) -> None: + """Reject invalid local logits before they can be converted by argmax.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [4, 4]}}, + {"task": "semantic_segmentation", "dataset": "cityscapes"}, + ) + logits = torch.zeros((1, 19, 2, 2)) + logits[0, 0, 0, 0] = invalid_value + + with pytest.raises(ValueError, match="logits must contain only finite values"): + post(logits) + + +@pytest.mark.parametrize("invalid_target", [1.5, -1.0, float("nan"), float("inf")]) +def test_semantic_metrics_reject_invalid_target_ids(invalid_target: float) -> None: + """Do not treat corrupted semantic ground truth as ignored pixels.""" + + target = torch.tensor([[invalid_target]], dtype=torch.float32).numpy() + with pytest.raises(ValueError, match="Semantic targets must be finite class IDs"): + SemanticMetricAccumulator(nc=2).update(np.zeros((1, 1)), target) + + +def test_semantic_postprocess_restores_letterbox_padding() -> None: + """Crop padding before nearest-restoring a semantic map.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [8, 8]}}, + {"task": "semantic_segmentation", "dataset": "ade20k"}, + ) + output = torch.zeros((1, 8, 8)) + output[:, 2:6, :] = 5 + restored = post(output, img0_shape=(2, 4), ratio_pad=((2.0, 2.0), (0.0, 2.0))) + assert isinstance(restored, torch.Tensor) + assert restored.shape == (2, 4) + assert torch.equal(restored, torch.full((2, 4), 5)) + + +def test_semantic_logits_restore_before_argmax_and_support_batches() -> None: + """Bilinearly restore logits before choosing classes for each original shape.""" + + post = SemanticSegPost( + {"LetterBox": {"img_size": [4, 8]}}, + {"task": "semantic_segmentation", "dataset": "cityscapes"}, + ) + logits = torch.zeros((2, 19, 2, 4)) + logits[:, 0] = 1.0 + logits[:, 1, :, 2:] = 2.0 + restored = post( + logits, + img0_shape=[(2, 8), (1, 4)], + ratio_pad=[((1.0, 1.0), (0.0, 1.0)), ((2.0, 2.0), (0.0, 1.0))], + ) + + assert isinstance(restored, list) + assert [tuple(item.shape) for item in restored] == [(2, 8), (1, 4)] + assert set(restored[0].unique().tolist()) == {0, 1} diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py new file mode 100644 index 0000000..5e671e9 --- /dev/null +++ b/tests/test_wrapper.py @@ -0,0 +1,3106 @@ +"""Tests for vision wrapper MXQ path resolution.""" + +from __future__ import annotations + +import gc +import os +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pytest +import torch +from mblt_vision.utils.postprocess import build_postprocess +from mblt_vision.utils.postprocess.base import YOLODetectionPostBase +from mblt_vision.utils.postprocess.common import ( + crop_mask, + dual_topk, + nmsout2eval, + normalize_image_shapes, + normalize_ratio_pads, + process_mask_upsample, + scale_coords, + scale_masks, +) +from mblt_vision.utils.postprocess.yolo_anchorless_post import ( + AnchorlessOutputLayout, + YOLOAnchorlessDetectionPost, + YOLOAnchorlessOBBPost, + YOLOAnchorlessPosePost, + _AnchorlessNMSInput, +) + +import mblt_vision.wrapper as wrapper +from mblt_vision._compat import create_model_class +from mblt_vision.utils.letterbox import resolve_ratio_pad +from mblt_vision.utils.preprocess import build_preprocess +from mblt_vision.utils.results import Results +from mblt_vision.utils.types import ListTensorLike +from mblt_vision.wrapper import MBLT_Engine + + +@pytest.mark.parametrize( + ("model_name", "output_shape", "expected_shape", "expected_class"), + [ + ("yolo26m-depth", (1, 192, 192), (1, 768, 768), None), + ("yolo26m-sem", (1024, 2048, 19), (1, 1024, 2048), 8), + ], +) +def test_local_mxq_dense_pipeline_uses_normalized_results( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + model_name: str, + output_shape: tuple[int, ...], + expected_shape: tuple[int, int, int], + expected_class: int | None, +) -> None: + """Exercise preprocessing, MXQ output normalization, and ``Results`` without an NPU.""" + + mxq_path = tmp_path / f"{model_name}.mxq" + mxq_path.write_bytes(b"mxq") + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + assert kwargs["mxq_path"] == str(mxq_path) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def __call__(self, input_value: np.ndarray) -> np.ndarray: + assert input_value.ndim == 3 + output = np.zeros(output_shape, dtype=np.float32) + if expected_class is not None: + output[..., expected_class] = 1.0 + else: + output[...] = 1.0 + return output + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + engine = MBLT_Engine(model_name, model_path=str(mxq_path)) + try: + preprocessed = engine.preprocess(np.zeros((20, 40, 3), dtype=np.uint8)) + raw_output = engine(preprocessed) + result = engine.postprocess(raw_output) + if expected_class is None: + assert isinstance(result.depth, torch.Tensor) + assert tuple(result.depth.shape) == expected_shape + assert torch.isfinite(result.depth).all() + else: + assert isinstance(result.semantic_mask, torch.Tensor) + assert tuple(result.semantic_mask.shape) == expected_shape + assert set(result.semantic_mask.unique().tolist()) == {expected_class} + finally: + engine.dispose() + + +def test_default_cache_dir_uses_stable_private_fallback( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Reuse a private fallback cache when the preferred cache cannot be created.""" + + def _fail_mkdir(self: Path, *args: object, **kwargs: object) -> None: + del self, args, kwargs + raise OSError("home cache is unavailable") + + monkeypatch.setattr(Path, "mkdir", _fail_mkdir) + monkeypatch.setattr( + "mblt_vision.wrapper.tempfile.gettempdir", lambda: str(tmp_path) + ) + + cache_dir = Path(wrapper._default_cache_dir()) + + assert cache_dir == tmp_path / f"mblt_model_zoo-{os.getuid()}" + assert wrapper._default_cache_dir() == str(cache_dir) + assert cache_dir.stat().st_mode & 0o777 == 0o700 + + +def test_default_cache_dir_rejects_unsafe_existing_fallback( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not reuse a fallback cache that other users can write to.""" + + unsafe_fallback = tmp_path / f"mblt_model_zoo-{os.getuid()}" + unsafe_fallback.mkdir(mode=0o700) + unsafe_fallback.chmod(0o777) + + def _fail_mkdir(self: Path, *args: object, **kwargs: object) -> None: + del self, args, kwargs + raise OSError("home cache is unavailable") + + monkeypatch.setattr(Path, "mkdir", _fail_mkdir) + monkeypatch.setattr( + "mblt_vision.wrapper.tempfile.gettempdir", lambda: str(tmp_path) + ) + + with pytest.raises(RuntimeError, match="not a private directory"): + wrapper._default_cache_dir() + + +def test_default_cache_dir_preserves_existing_probe_named_file( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Never remove a caller-owned file while checking cache writability.""" + + preferred = tmp_path / "mblt_model_zoo" + preferred.mkdir() + preferred.chmod(0o700) + marker = preferred / ".write_test" + marker.write_bytes(b"caller-owned") + monkeypatch.setattr( + "mblt_vision.wrapper.os.path.expanduser", lambda _: str(preferred) + ) + + assert wrapper._default_cache_dir() == str(preferred) + assert marker.read_bytes() == b"caller-owned" + + +def test_default_cache_dir_rejects_unsafe_preferred_directory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Do not use a preferred cache other users can write to.""" + + preferred = tmp_path / "mblt_model_zoo" + preferred.mkdir(mode=0o700) + preferred.chmod(0o777) + fallback = tmp_path / "fallback" + monkeypatch.setattr( + "mblt_vision.wrapper.os.path.expanduser", lambda _: str(preferred) + ) + monkeypatch.setattr(wrapper, "_fallback_cache_dir", lambda: str(fallback)) + + assert wrapper._default_cache_dir() == str(fallback) + + +@pytest.mark.parametrize( + ("key", "value", "expected_suffix"), + [("mxq_path", "artifact.onnx", ".mxq"), ("onnx_path", "artifact.mxq", ".onnx")], +) +def test_engine_rejects_wrong_suffix_in_direct_file_config( + key: str, value: str, expected_suffix: str +) -> None: + """Validate compatibility path aliases provided through a direct config mapping.""" + + with pytest.raises( + ValueError, match=rf"file_cfg\.{key} must end in '{expected_suffix}'" + ): + MBLT_Engine( + {"file_cfg": {key: value}, "pre_cfg": {}, "post_cfg": {}}, + ) + + +def test_cache_directory_is_resolved_lazily(monkeypatch: pytest.MonkeyPatch) -> None: + """Do not create or probe the artifact cache until an operation needs it.""" + + calls: list[str] = [] + + def resolve_cache_dir() -> str: + calls.append("resolved") + return "/tmp/mblt-model-zoo-cache" + + monkeypatch.setattr(wrapper, "_resolved_cache_dir", None) + monkeypatch.setattr(wrapper, "_default_cache_dir", resolve_cache_dir) + + assert calls == [] + assert wrapper.get_mobilint_cache_dir() == "/tmp/mblt-model-zoo-cache" + assert wrapper.get_mobilint_cache_dir() == "/tmp/mblt-model-zoo-cache" + assert calls == ["resolved"] + + +def test_onnx_runtime_defaults_to_cpu_provider() -> None: + """Avoid accelerator provider probing unless callers explicitly opt in.""" + + class _FakeOrt: + @staticmethod + def get_available_providers() -> list[str]: + return [ + "TensorrtExecutionProvider", + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + + assert wrapper._resolve_onnx_providers(_FakeOrt()) == ["CPUExecutionProvider"] + assert wrapper._resolve_onnx_providers(_FakeOrt(), ["CUDAExecutionProvider"]) == [ + "CUDAExecutionProvider" + ] + + +def test_dual_topk_logits_matches_sigmoid_scores() -> None: + """Select the same NMS-free detections before converting logits to probabilities.""" + torch.manual_seed(0) + boxes = torch.rand(24, 4) + logits = torch.rand(24, 5).mul(10.0).sub(5.0) + extra = torch.rand(24, 2) + logits_input = torch.cat([boxes, logits, extra], dim=1) + probability_input = torch.cat([boxes, logits.sigmoid(), extra], dim=1) + + actual = dual_topk( + logits_input, nc=5, n_extra=2, max_det=8, conf_thres=0.25, score_is_logits=True + ) + expected = dual_topk(probability_input, nc=5, n_extra=2, max_det=8, conf_thres=0.25) + + torch.testing.assert_close(actual, expected) + + +def test_file_config_cleansing_prefers_existing_mxq_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Use an existing local MXQ path without attempting a Hub download.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + + def _unexpected_download(**kwargs: Any) -> str: + raise AssertionError("hf_hub_download should not be called") + + monkeypatch.setattr(wrapper, "hf_hub_download", _unexpected_download) + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.file_cfg = { + "mxq_path": str(mxq_path), + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + "core_mode": "global8", + } + + engine.file_config_cleansing() + + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert "repo_id" not in engine.file_cfg + assert "filename" not in engine.file_cfg + assert "revision" not in engine.file_cfg + + +def test_model_path_defaults_to_local_mxq_for_mxq_framework(tmp_path: Path) -> None: + """Map ``model_path`` to ``mxq_path`` when MXQ inference is requested.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.framework = "mxq" + engine.file_cfg = {"model_path": str(mxq_path)} + + engine.file_config_cleansing() + + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert "onnx_path" not in engine.file_cfg or not engine.file_cfg["onnx_path"] + + +def test_model_path_defaults_to_local_onnx_for_onnx_framework(tmp_path: Path) -> None: + """Map ``model_path`` to ``onnx_path`` when ONNX inference is requested.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.framework = "onnx" + engine.file_cfg = {"model_path": str(onnx_path)} + + engine.file_config_cleansing() + + assert engine.file_cfg["onnx_path"] == str(onnx_path) + assert "mxq_path" not in engine.file_cfg or not engine.file_cfg["mxq_path"] + + +@pytest.mark.parametrize("path_argument", ["model_path", "mxq_path"]) +def test_engine_init_rejects_nonexistent_explicit_mxq_path( + tmp_path: Path, path_argument: str +) -> None: + """Do not replace an explicit missing MXQ path with a Hub artifact.""" + + missing_path = tmp_path / "missing.mxq" + model_config = { + "file_cfg": { + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + }, + "pre_cfg": {}, + "post_cfg": {}, + } + + with pytest.raises( + FileNotFoundError, match=r"Explicit MXQ model path.*missing\.mxq" + ): + path_kwargs: dict[str, Any] = {path_argument: str(missing_path)} + MBLT_Engine(model_config, **path_kwargs) + + +@pytest.mark.parametrize("path_argument", ["model_path", "onnx_path"]) +def test_engine_init_rejects_nonexistent_explicit_onnx_path( + tmp_path: Path, path_argument: str +) -> None: + """Do not replace an explicit missing ONNX path with a Hub artifact.""" + + missing_path = tmp_path / "missing.onnx" + model_config = { + "file_cfg": { + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + }, + "pre_cfg": {}, + "post_cfg": {}, + } + + with pytest.raises( + FileNotFoundError, match=r"Explicit ONNX model path.*missing\.onnx" + ): + path_kwargs: dict[str, Any] = {path_argument: str(missing_path)} + MBLT_Engine(model_config, **path_kwargs) + + +def test_engine_init_rejects_wrong_suffix_for_mxq_path(tmp_path: Path) -> None: + """Do not route an existing ONNX file through the MXQ compatibility alias.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + with pytest.raises(ValueError, match=r"mxq_path must end in '.mxq'"): + MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, mxq_path=str(onnx_path) + ) + + +@pytest.mark.parametrize("section", ["file_cfg", "pre_cfg", "post_cfg"]) +def test_engine_init_rejects_nonmapping_direct_configuration_sections( + section: str, +) -> None: + """Report malformed direct model mappings before backend construction.""" + + config: dict[str, Any] = {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}} + config[section] = [] + with pytest.raises(ValueError, match=rf"section '{section}' must be a mapping"): + MBLT_Engine(config) + + +def test_engine_init_disposes_mxq_backend_after_postprocess_setup_failure( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Dispose an acquired NPU backend when later engine setup fails.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + disposed = False + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + del kwargs + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + nonlocal disposed + disposed = True + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda *args, **kwargs: (_ for _ in ()).throw( + ValueError("invalid postprocess") + ), + ) + + with pytest.raises(ValueError, match="invalid postprocess"): + MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + model_path=str(mxq_path), + ) + + assert disposed + + +def test_engine_init_disposes_onnx_backend_after_preprocess_setup_failure( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Dispose an acquired ONNX backend when later engine setup fails.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + disposed = False + + class _Session: + def get_inputs(self) -> list[Any]: + return [type("Input", (), {"name": "input"})()] + + def get_outputs(self) -> list[Any]: + return [type("Output", (), {"name": "output"})()] + + class _FakeONNXBackend: + def __init__(self, *args: Any, **kwargs: Any) -> None: + del args, kwargs + self.session: _Session | None = None + + def create(self) -> None: + self.session = _Session() + + def dispose(self) -> None: + nonlocal disposed + disposed = True + self.session = None + + monkeypatch.setattr(wrapper, "ONNXBackend", _FakeONNXBackend) + monkeypatch.setattr(wrapper, "_load_onnxruntime", lambda: object()) + monkeypatch.setattr( + wrapper, + "build_preprocess", + lambda config: (_ for _ in ()).throw(ValueError("invalid preprocess")), + ) + + with pytest.raises(ValueError, match="invalid preprocess"): + MBLT_Engine( + { + "file_cfg": {"onnx_path": str(onnx_path)}, + "pre_cfg": {}, + "post_cfg": {}, + }, + framework="onnx", + ) + + assert disposed + + +def test_engine_context_manager_closes_backend_once( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Context exit releases an MXQ backend and close/dispose remain idempotent.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + dispose_calls = 0 + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + del kwargs + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + nonlocal dispose_calls + dispose_calls += 1 + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr(wrapper, "build_postprocess", lambda *args, **kwargs: object()) + + with MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, model_path=str(mxq_path) + ) as engine: + assert engine is not None + + engine.close() + engine.dispose() + assert dispose_calls == 1 + with pytest.raises(RuntimeError, match="closed"): + engine(torch.zeros((1, 3, 8, 8))) + + +def test_engine_finalizer_closes_unmanaged_backend( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Best-effort finalization prevents leaked backends for unmanaged engines.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + dispose_calls = 0 + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + del kwargs + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + nonlocal dispose_calls + dispose_calls += 1 + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr(wrapper, "build_postprocess", lambda *args, **kwargs: object()) + + engine = MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, model_path=str(mxq_path) + ) + del engine + gc.collect() + + assert dispose_calls == 1 + + +def test_engine_init_accepts_local_mxq_model_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Route API ``model_path`` to the MXQ backend for local MXQ inference.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + backend_kwargs: dict[str, Any] = {} + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + backend_kwargs.update(kwargs) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + engine = MBLT_Engine( + model_cls={ + "file_cfg": {}, + "pre_cfg": {}, + "post_cfg": {}, + }, + model_path=str(mxq_path), + ) + + try: + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert backend_kwargs["mxq_path"] == str(mxq_path) + finally: + engine.dispose() + + +def test_engine_init_preserves_legacy_positional_arguments( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Keep pre-``model_path`` positional arguments bound to their original fields.""" + + mxq_path = tmp_path / "legacy.mxq" + onnx_path = tmp_path / "legacy.onnx" + mxq_path.write_bytes(b"mxq") + onnx_path.write_bytes(b"onnx") + backend_kwargs: dict[str, Any] = {} + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + backend_kwargs.update(kwargs) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + engine = MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + "DEFAULT", + str(mxq_path), + str(onnx_path), + 3, + "global8", + ) + + try: + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert engine.file_cfg["onnx_path"] == str(onnx_path) + assert backend_kwargs["dev_no"] == 3 + assert backend_kwargs["core_mode"] == "global8" + finally: + engine.dispose() + + +@pytest.mark.parametrize( + ("configured_core_mode", "core_mode"), + [(None, "unsupported"), ("unsupported", None)], +) +def test_engine_init_rejects_invalid_core_modes( + configured_core_mode: str | None, + core_mode: str | None, +) -> None: + """Validate direct caller and model-config core modes before backend setup.""" + + file_cfg: dict[str, Any] = {} + if configured_core_mode is not None: + file_cfg["core_mode"] = configured_core_mode + + with pytest.raises(ValueError, match="Invalid core mode 'unsupported'"): + MBLT_Engine( + model_cls={"file_cfg": file_cfg, "pre_cfg": {}, "post_cfg": {}}, + core_mode=core_mode, + ) + + +def test_engine_init_rejects_core_modes_not_supported_by_regulus() -> None: + """Reject an Aries-only allocation mode before creating a Regulus backend.""" + + with pytest.raises(ValueError, match="not supported by regulus-ra"): + MBLT_Engine( + model_cls={"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + core_mode="global8", + target_device="regulus-ra", + ) + + +@pytest.mark.parametrize( + ( + "explicit_target_device", + "expected_target_device", + "expected_cores", + "expected_core_mode", + ), + [ + (None, "regulus-rb", [], "single"), + ( + "aries-rb", + "aries-rb", + ["0:0", "0:1", "0:2", "0:3", "1:0", "1:1", "1:2", "1:3"], + "global8", + ), + ], +) +def test_engine_init_resolves_target_device_before_default_core_settings( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + explicit_target_device: str | None, + expected_target_device: str, + expected_cores: list[str], + expected_core_mode: str, +) -> None: + """Prefer an explicit board, otherwise retain the configured board and defaults.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + backend_kwargs: dict[str, Any] = {} + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + backend_kwargs.update(kwargs) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + engine = MBLT_Engine( + model_cls={ + "file_cfg": {"target_device": "regulus-rb", "core_mode": "global8"}, + "pre_cfg": {}, + "post_cfg": {}, + }, + model_path=str(mxq_path), + target_device=explicit_target_device, + ) + + try: + assert engine.file_cfg["target_device"] == expected_target_device + assert backend_kwargs["target_device"] == expected_target_device + assert backend_kwargs["core_mode"] == expected_core_mode + assert backend_kwargs["target_cores"] == expected_cores + assert backend_kwargs["target_clusters"] == ( + [] if expected_target_device == "regulus-rb" else [0, 1] + ) + finally: + engine.dispose() + + +@pytest.mark.parametrize( + ("target_device", "expected_core_mode"), + [("aries-rb", "global8"), ("regulus-ra", "single")], +) +def test_engine_init_uses_board_specific_default_core_mode( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + target_device: str, + expected_core_mode: str, +) -> None: + """Keep direct engine construction aligned with the CLI board defaults.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + backend_kwargs: dict[str, Any] = {} + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + backend_kwargs.update(kwargs) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + engine = MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + model_path=str(mxq_path), + target_device=target_device, + ) + + try: + assert engine.file_cfg["core_mode"] == expected_core_mode + assert backend_kwargs["core_mode"] == expected_core_mode + finally: + engine.dispose() + + +def test_engine_init_preserves_shifted_positional_mxq_runtime_arguments( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Keep runtime arguments after a positional MXQ ``model_path``.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + backend_kwargs: dict[str, Any] = {} + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + backend_kwargs.update(kwargs) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + # This intentionally exercises the legacy shifted positional layout, which + # does not match the current typed ``MBLT_Engine`` constructor signature. + engine = cast(Any, MBLT_Engine)( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + "DEFAULT", + str(mxq_path), + "", + "", + 3, + "global8", + ) + + try: + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert backend_kwargs["dev_no"] == 3 + assert backend_kwargs["core_mode"] == "global8" + finally: + engine.dispose() + + +def test_legacy_wrapper_preserves_positional_path_and_framework_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep the generated compatibility constructor's original positional order.""" + + captured_kwargs: dict[str, Any] = {} + + def _capture_engine_init(self: MBLT_Engine, **kwargs: Any) -> None: + del self + captured_kwargs.update(kwargs) + + monkeypatch.setattr(MBLT_Engine, "__init__", _capture_engine_init) + compat_cls = create_model_class( + "ResNet50", "mblt_model_zoo.vision.image_classification" + ) + + compat_cls( + None, + "DEFAULT", + "global8", + "aries", + 3, + ["0:0"], + [0], + "legacy.mxq", + "legacy.onnx", + "onnx", + ) + + assert captured_kwargs["model_path"] == "" + assert captured_kwargs["mxq_path"] == "legacy.mxq" + assert captured_kwargs["onnx_path"] == "legacy.onnx" + assert captured_kwargs["framework"] == "onnx" + + +def test_legacy_wrapper_preserves_shifted_positional_model_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep the generated positional ``model_path`` slot working for ONNX.""" + + captured_kwargs: dict[str, Any] = {} + + def _capture_engine_init(self: MBLT_Engine, **kwargs: Any) -> None: + del self + captured_kwargs.update(kwargs) + + monkeypatch.setattr(MBLT_Engine, "__init__", _capture_engine_init) + compat_cls = create_model_class( + "ResNet50", "mblt_model_zoo.vision.image_classification" + ) + + compat_cls( + None, + "DEFAULT", + "global8", + "aries", + 3, + ["0:0"], + [0], + "legacy.onnx", + None, + None, + "onnx", + ) + + assert captured_kwargs["model_path"] == "legacy.onnx" + assert captured_kwargs["mxq_path"] == "" + assert captured_kwargs["onnx_path"] == "" + assert captured_kwargs["framework"] == "onnx" + + +def test_legacy_wrapper_preserves_shifted_positional_mxq_tail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep the generated shifted positional tail working for MXQ.""" + + captured_kwargs: dict[str, Any] = {} + + def _capture_engine_init(self: MBLT_Engine, **kwargs: Any) -> None: + del self + captured_kwargs.update(kwargs) + + monkeypatch.setattr(MBLT_Engine, "__init__", _capture_engine_init) + compat_cls = create_model_class( + "ResNet50", "mblt_model_zoo.vision.image_classification" + ) + + compat_cls( + None, + "DEFAULT", + "global8", + "aries", + 3, + ["0:0"], + [0], + "legacy.mxq", + None, + None, + "mxq", + ) + + assert captured_kwargs["model_path"] == "legacy.mxq" + assert captured_kwargs["mxq_path"] == "" + assert captured_kwargs["onnx_path"] == "" + assert captured_kwargs["framework"] == "mxq" + + +def test_engine_init_accepts_obb_task( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Construct the OBB postprocessor from its canonical task name.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + + engine = MBLT_Engine( + model_cls={ + "file_cfg": {}, + "pre_cfg": {"LetterBox": {"img_size": [640, 640]}}, + "post_cfg": { + "task": "obb", + "dataset": "dotav1", + "nl": 3, + "reg_max": 16, + "n_extra": 1, + }, + }, + model_path=str(mxq_path), + ) + + try: + assert isinstance(engine.postprocessor, YOLOAnchorlessOBBPost) + assert engine.postprocessor.task == "obb" + finally: + engine.dispose() + + +def test_engine_init_auto_detects_mxq_framework_from_model_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Infer the MXQ framework from a local MXQ path when framework is omitted.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + engine = MBLT_Engine( + model_cls={ + "file_cfg": {}, + "pre_cfg": {}, + "post_cfg": {}, + }, + model_path=str(mxq_path), + ) + + try: + assert engine.framework == "mxq" + assert engine.file_cfg["mxq_path"] == str(mxq_path) + finally: + engine.dispose() + + +@pytest.mark.parametrize( + "model_path_style", ["keyword", "positional-runtime", "onnx-path"] +) +def test_engine_init_accepts_local_onnx_model_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + model_path_style: str, +) -> None: + """Route API ``model_path`` to the ONNX runtime session for local ONNX inference.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + class _FakeInput: + name = "input" + shape = [1, 3, 224, 224] + + class _FakeOutput: + name = "output" + + class _FakeSession: + def __init__(self, path: str, providers: list[str]) -> None: + self.path = path + self.providers = providers + + def get_inputs(self) -> list[_FakeInput]: + return [_FakeInput()] + + def get_outputs(self) -> list[_FakeOutput]: + return [_FakeOutput()] + + class _FakeOrt: + def __init__(self) -> None: + self.session: _FakeSession | None = None + + @staticmethod + def get_available_providers() -> list[str]: + return ["CPUExecutionProvider"] + + def InferenceSession(self, path: str, providers: list[str]) -> _FakeSession: + self.session = _FakeSession(path, providers) + return self.session + + fake_ort = _FakeOrt() + monkeypatch.setattr(wrapper, "_load_onnxruntime", lambda: fake_ort) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + model_config = { + "file_cfg": {}, + "pre_cfg": {}, + "post_cfg": {}, + } + if model_path_style == "positional-runtime": + # Keep this legacy positional-runtime invocation untyped: its argument + # order is normalized at runtime by ``MBLT_Engine``. + engine = cast(Any, MBLT_Engine)( + model_config, + "DEFAULT", + str(onnx_path), + "", + "", + 3, + "global8", + ["0:0"], + [1], + {"confidence": 0.25}, + "onnx", + ["CPUExecutionProvider"], + ) + elif model_path_style == "onnx-path": + engine = MBLT_Engine(model_cls=model_config, onnx_path=str(onnx_path)) + else: + engine = MBLT_Engine(model_cls=model_config, model_path=str(onnx_path)) + + assert engine.file_cfg["onnx_path"] == str(onnx_path) + assert fake_ort.session is not None + assert fake_ort.session.path == str(onnx_path) + assert engine.framework == "onnx" + if model_path_style == "positional-runtime": + assert engine.file_cfg["dev_no"] == 3 + assert engine.file_cfg["core_mode"] == "global8" + assert engine.file_cfg["target_cores"] == ["0:0"] + assert engine.file_cfg["target_clusters"] == [1] + assert engine.postprocess_kwargs == {"confidence": 0.25} + assert fake_ort.session.providers == ["CPUExecutionProvider"] + + +def test_engine_init_rejects_framework_conflicting_with_onnx_path_alias( + tmp_path: Path, +) -> None: + """Do not silently route an explicit ONNX alias through MXQ inference.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + with pytest.raises(ValueError, match=r"Framework `mxq` conflicts with model path"): + MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + onnx_path=str(onnx_path), + framework="mxq", + ) + + +@pytest.mark.parametrize("input_count", [0, 2]) +def test_engine_init_rejects_onnx_models_without_exactly_one_input( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, input_count: int +) -> None: + """The public tensor-only engine API cannot satisfy multi-input ONNX graphs.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + class _Input: + def __init__(self, index: int) -> None: + self.name = f"input-{index}" + + class _Session: + def get_inputs(self) -> list[_Input]: + return [_Input(index) for index in range(input_count)] + + def get_outputs(self) -> list[object]: + return [] + + class _Backend: + def __init__(self, *args: Any, **kwargs: Any) -> None: + del args, kwargs + self.session = _Session() + + def create(self) -> None: + return None + + def dispose(self) -> None: + return None + + class _Ort: + @staticmethod + def get_available_providers() -> list[str]: + return ["CPUExecutionProvider"] + + monkeypatch.setattr(wrapper, "ONNXBackend", _Backend) + monkeypatch.setattr(wrapper, "_load_onnxruntime", _Ort) + + with pytest.raises(ValueError, match=rf"got {input_count} inputs"): + MBLT_Engine( + {"file_cfg": {}, "pre_cfg": {}, "post_cfg": {}}, + onnx_path=str(onnx_path), + ) + + +def test_engine_init_rejects_conflicting_framework_and_model_path( + tmp_path: Path, +) -> None: + """Fail fast when the explicit framework conflicts with the local model suffix.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + + with pytest.raises(ValueError, match="conflicts with model path"): + MBLT_Engine( + model_cls={ + "file_cfg": {}, + "pre_cfg": {}, + "post_cfg": {}, + }, + framework="onnx", + model_path=str(mxq_path), + ) + + +def test_engine_init_auto_detects_onnx_framework_from_config_model_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Infer ONNX from ``file_cfg.model_path`` when constructor inputs omit the framework.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + class _FakeInput: + name = "input" + shape = [1, 3, 224, 224] + + class _FakeOutput: + name = "output" + + class _FakeSession: + def __init__(self, path: str, providers: list[str]) -> None: + self.path = path + self.providers = providers + + def get_inputs(self) -> list[_FakeInput]: + return [_FakeInput()] + + def get_outputs(self) -> list[_FakeOutput]: + return [_FakeOutput()] + + class _FakeOrt: + def __init__(self) -> None: + self.session: _FakeSession | None = None + + @staticmethod + def get_available_providers() -> list[str]: + return ["CPUExecutionProvider"] + + def InferenceSession(self, path: str, providers: list[str]) -> _FakeSession: + self.session = _FakeSession(path, providers) + return self.session + + fake_ort = _FakeOrt() + monkeypatch.setattr(wrapper, "_load_onnxruntime", lambda: fake_ort) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + engine = MBLT_Engine( + model_cls={ + "file_cfg": {"model_path": str(onnx_path)}, + "pre_cfg": {}, + "post_cfg": {}, + } + ) + + assert engine.file_cfg["onnx_path"] == str(onnx_path) + assert fake_ort.session is not None + assert fake_ort.session.path == str(onnx_path) + assert engine.framework == "onnx" + + +def test_engine_init_rejects_conflicting_framework_and_config_model_path( + tmp_path: Path, +) -> None: + """Fail fast when the explicit framework conflicts with ``file_cfg.model_path``.""" + + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + with pytest.raises(ValueError, match="conflicts with model path"): + MBLT_Engine( + model_cls={ + "file_cfg": {"model_path": str(onnx_path)}, + "pre_cfg": {}, + "post_cfg": {}, + }, + framework="mxq", + ) + + +def test_legacy_local_path_stays_mxq_specific_for_onnx_framework( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Keep legacy ``local_path`` semantics stable in compatibility wrappers.""" + + mxq_path = tmp_path / "resnet50.mxq" + onnx_path = tmp_path / "resnet50.onnx" + mxq_path.write_bytes(b"mxq") + onnx_path.write_bytes(b"onnx") + + class _FakeInput: + name = "input" + shape = [1, 3, 224, 224] + + class _FakeOutput: + name = "output" + + class _FakeSession: + def __init__(self, path: str, providers: list[str]) -> None: + self.path = path + self.providers = providers + + def get_inputs(self) -> list[_FakeInput]: + return [_FakeInput()] + + def get_outputs(self) -> list[_FakeOutput]: + return [_FakeOutput()] + + class _FakeOrt: + def __init__(self) -> None: + self.session: _FakeSession | None = None + + @staticmethod + def get_available_providers() -> list[str]: + return ["CPUExecutionProvider"] + + def InferenceSession(self, path: str, providers: list[str]) -> _FakeSession: + self.session = _FakeSession(path, providers) + return self.session + + fake_ort = _FakeOrt() + monkeypatch.setattr(wrapper, "_load_onnxruntime", lambda: fake_ort) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + + compat_cls = create_model_class( + "ResNet50", "mblt_model_zoo.vision.image_classification" + ) + engine = compat_cls(local_path=str(mxq_path), framework="onnx") + + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert engine.file_cfg["onnx_path"] == str(onnx_path) + assert fake_ort.session is not None + assert fake_ort.session.path == str(onnx_path) + assert engine.framework == "onnx" + + +def test_engine_init_defaults_to_mxq_without_model_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Use MXQ as the fallback framework when no model path is provided.""" + + backend_kwargs: dict[str, Any] = {} + + class _FakeBackend: + def __init__(self, **kwargs: Any) -> None: + backend_kwargs.update(kwargs) + + def create(self) -> None: + return None + + def launch(self) -> None: + return None + + def get_dtype(self) -> str: + return "DataType.Float32" + + def dispose(self) -> None: + return None + + monkeypatch.setattr(wrapper, "MobilintNPUBackend", _FakeBackend) + monkeypatch.setattr(wrapper, "build_preprocess", lambda config: config) + monkeypatch.setattr( + wrapper, + "build_postprocess", + lambda pre_cfg, post_cfg, **kwargs: (pre_cfg, post_cfg, kwargs), + ) + monkeypatch.setattr( + wrapper.MBLT_Engine, + "_download_hub_artifact", + lambda self, **kwargs: "/tmp/model.mxq", + ) + + engine = MBLT_Engine( + model_cls={ + "file_cfg": { + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + }, + "pre_cfg": {}, + "post_cfg": {}, + }, + ) + + try: + assert engine.framework == "mxq" + assert "mxq_path" in backend_kwargs + finally: + engine.dispose() + + +@pytest.mark.parametrize("mask_count", [1, 50]) +def test_crop_mask_matches_ultralytics_fractional_boundaries(mask_count: int) -> None: + """Use identical fractional crop semantics on both sides of the former CPU branch.""" + + masks = torch.ones((mask_count, 5, 5), dtype=torch.float32) + boxes = torch.tensor([[1.2, 1.8, 3.6, 4.2]], dtype=torch.float32).repeat( + mask_count, 1 + ) + + cropped = crop_mask(masks, boxes) + + expected = torch.zeros((mask_count, 5, 5), dtype=torch.float32) + expected[:, 2:5, 2:4] = 1 + assert torch.equal(cropped, expected) + + +def test_file_config_cleansing_downloads_from_exact_target_device_folder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Download only the MXQ artifact from the selected board folder.""" + + calls: list[str] = [] + + def _fake_download(**kwargs: Any) -> str: + subfolder = kwargs["subfolder"] + calls.append(subfolder) + return "/tmp/global8.mxq" + + monkeypatch.setattr(wrapper, "hf_hub_download", _fake_download) + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.framework = "mxq" + engine.file_cfg = { + "mxq_path": "", + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + "core_mode": "global8", + "target_device": "aries-rb", + } + + engine.file_config_cleansing() + + assert calls == ["aries-rb"] + assert engine.file_cfg["mxq_path"] == "/tmp/global8.mxq" + assert engine.file_cfg["onnx_filename"] == "model.onnx" + assert "onnx_path" not in engine.file_cfg + + +def test_file_config_cleansing_downloads_only_onnx_for_onnx_framework( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Download only the ONNX artifact when ONNX inference is requested.""" + + calls: list[dict[str, Any]] = [] + + def _fake_download(**kwargs: Any) -> str: + calls.append(kwargs) + assert kwargs["filename"] == "model.onnx" + assert "subfolder" not in kwargs + return "/tmp/model.onnx" + + monkeypatch.setattr(wrapper, "hf_hub_download", _fake_download) + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.framework = "onnx" + engine.file_cfg = { + "onnx_path": "", + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + "core_mode": "global8", + } + + engine.file_config_cleansing() + + assert len(calls) == 1 + assert engine.file_cfg["onnx_filename"] == "model.onnx" + assert engine.file_cfg["onnx_path"] == "/tmp/model.onnx" + assert "mxq_path" not in engine.file_cfg or not engine.file_cfg["mxq_path"] + + +def test_file_config_cleansing_resolves_local_onnx( + tmp_path: Path, +) -> None: + """Resolve ONNX file path next to local MXQ file when they exist locally.""" + + mxq_path = tmp_path / "model.mxq" + mxq_path.write_bytes(b"mxq") + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.file_cfg = { + "mxq_path": str(mxq_path), + "repo_id": "mobilint/example", + "filename": "model.mxq", + "revision": "main", + "core_mode": "global8", + } + + engine.file_config_cleansing() + + assert engine.file_cfg["mxq_path"] == str(mxq_path) + assert engine.file_cfg["onnx_filename"] == "model.onnx" + assert engine.file_cfg["onnx_path"] == str(onnx_path) + + +def test_prepare_onnx_inputs_keeps_batched_nchw_layout() -> None: + """Preserve existing NCHW batches when feeding ONNX sessions.""" + + class _FakeInput: + name = "input" + shape = [1, 3, 224, 224] + + class _FakeSession: + def get_inputs(self) -> list[_FakeInput]: + return [_FakeInput()] + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.framework = "onnx" + fake_session = _FakeSession() + engine._onnx_session = fake_session + engine.model = fake_session + engine.input_name = "input" + + batch = torch.zeros((2, 3, 224, 224), dtype=torch.float32) + + inputs = engine._prepare_onnx_inputs(batch) + + assert set(inputs) == {"input"} + assert inputs["input"].shape == (2, 3, 224, 224) + assert inputs["input"].dtype == np.float32 + + +def test_prepare_onnx_inputs_transposes_static_square_hwc_images() -> None: + """Use the static ONNX channel axis to convert square HWC images to NCHW.""" + + class _FakeInput: + name = "input" + shape = [1, 3, 224, 224] + + class _FakeSession: + def get_inputs(self) -> list[_FakeInput]: + return [_FakeInput()] + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.framework = "onnx" + fake_session = _FakeSession() + engine._onnx_session = fake_session + engine.model = fake_session + engine.input_name = "input" + + image = np.zeros((224, 224, 3), dtype=np.float32) + image[..., 0] = 1.0 + image[..., 1] = 2.0 + image[..., 2] = 3.0 + + inputs = engine._prepare_onnx_inputs(image) + + assert inputs["input"].shape == (1, 3, 224, 224) + assert inputs["input"][0, :, 0, 0].tolist() == [1.0, 2.0, 3.0] + + +def test_final_onnx_detections_apply_confidence_threshold() -> None: + """Filter confidence on already-decoded ONNX detection outputs.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "nmsfree": True, + "reg_max": 16, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + final_output = np.array( + [ + [ + [10.0, 20.0, 30.0, 40.0, 0.49, 0.0], + [11.0, 21.0, 31.0, 41.0, 0.50, 1.0], + [12.0, 22.0, 32.0, 42.0, 0.90, 2.0], + ] + ], + dtype=np.float32, + ) + + result = postprocessor(final_output) + + assert len(result) == 1 + assert result[0].shape == (1, 6) + assert torch.all(result[0][:, 4] > 0.5) + + +def test_nmsfree_postprocess_supports_multilabel_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep COCO's multi-label postprocess path compatible with NMS-free models.""" + + postprocessor = build_postprocess( + {"LetterBox": {"img_size": [640, 640]}}, + {"task": "object_detection", "nl": 3, "reg_max": 16, "nmsfree": True}, + ) + decoded = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0, 0.9, 2.0], [11.0, 21.0, 31.0, 41.0, 0.0, 1.0]]], + dtype=torch.float32, + ) + monkeypatch.setattr(postprocessor, "extract_final_outputs", lambda _: (None, None)) + monkeypatch.setattr(postprocessor, "check_input", lambda x: x) + monkeypatch.setattr(postprocessor, "_pre_process", lambda _: (decoded, None)) + + result = postprocessor([torch.empty(1)], multi_label=True) + + assert len(result) == 1 + assert torch.equal(result[0], decoded[0, :1]) + + +def test_final_onnx_detections_normalize_singleton_and_channel_first() -> None: + """Accept common ONNX final-detection layouts without decoding them again.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "nmsfree": True, + "reg_max": 16, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + final_output = torch.tensor( + [ + [ + [10.0, 20.0, 30.0, 40.0, 0.90, 2.0], + [11.0, 21.0, 31.0, 41.0, 0.40, 1.0], + ] + ], + dtype=torch.float32, + ) + + singleton_result = postprocessor(final_output[:, None]) + channel_first_result = postprocessor(final_output.transpose(1, 2)) + + assert torch.equal(singleton_result[0], final_output[0, :1]) + assert torch.equal(channel_first_result[0], final_output[0, :1]) + + +def test_anchorless_pose_nms_uses_converted_provenance_for_ambiguous_shape() -> None: + """Keep converted row-major pose tensors row-major when both dimensions match.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "pose_estimation", + "nl": 3, + "n_extra": 51, + "reg_max": 16, + "conf_thres": 0.001, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + row_major = torch.zeros((56, 56), dtype=torch.float32) + row_major[:, :4] = torch.tensor([10.0, 10.0, 20.0, 20.0]) + row_major[:, 4] = torch.linspace(0.9, 0.1, 56) + + result = postprocessor.nms(_AnchorlessNMSInput([row_major], "candidates_first")) + + assert len(result) == 1 + assert result[0].shape == (1, 57) + assert torch.allclose(result[0][0, 4], torch.tensor(0.9)) + + +def test_final_onnx_segmentation_normalizes_detections_and_proto() -> None: + """Use final segmentation detections directly while preserving prototype layout.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "dflfree": True, + "nc": 80, + "n_extra": 32, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + final_output = torch.zeros((1, 1, 2, 38), dtype=torch.float32) + final_output[0, 0, 0, :6] = torch.tensor([10.0, 20.0, 30.0, 40.0, 0.90, 2.0]) + final_output[0, 0, 1, :6] = torch.tensor([11.0, 21.0, 31.0, 41.0, 0.40, 1.0]) + proto = torch.zeros((1, 32, 160, 160), dtype=torch.float32) + + result = postprocessor([final_output, proto]) + + assert len(result) == 1 + assert result[0][0].shape == (1, 38) + assert result[0][1].shape == (1, 640, 640) + + +def test_final_onnx_pose_normalizes_detections() -> None: + """Use final pose detections directly without sending them through decode.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "pose_estimation", + "nl": 3, + "n_extra": 51, + "reg_max": 16, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + final_output = torch.zeros((1, 2, 57), dtype=torch.float32) + final_output[0, 0, :6] = torch.tensor([10.0, 20.0, 30.0, 40.0, 0.90, 0.0]) + final_output[0, 1, :6] = torch.tensor([11.0, 21.0, 31.0, 41.0, 0.40, 0.0]) + + result = postprocessor(final_output) + + assert len(result) == 1 + assert result[0].shape == (1, 57) + assert torch.equal(result[0], final_output[0, :1]) + + +def test_final_onnx_pose_rejects_invalid_keypoint_confidence() -> None: + """Reject decoded pose confidences that normal decoding would sigmoid.""" + + postprocessor = cast( + YOLODetectionPostBase, + build_postprocess( + {"LetterBox": {"img_size": [640, 640]}}, + { + "task": "pose_estimation", + "nl": 3, + "n_extra": 51, + "reg_max": 16, + "conf_thres": 0.5, + "iou_thres": 0.7, + }, + ), + ) + final_output = torch.zeros((1, 1, 57), dtype=torch.float32) + final_output[0, 0, :6] = torch.tensor([10.0, 20.0, 30.0, 40.0, 0.9, 0.0]) + final_output[0, 0, 8] = 1.1 + + with pytest.raises( + ValueError, match="keypoint confidence values must be in \\[0, 1\\]" + ): + postprocessor(final_output) + + +@pytest.mark.parametrize( + ("task", "post_cfg_extra", "converted_dim"), + [ + ("object_detection", {}, 7), + ("pose_estimation", {"n_extra": 51}, 56), + ], +) +def test_non_e2e_single_converted_outputs_follow_task_shape( + task: str, + post_cfg_extra: dict[str, int], + converted_dim: int, +) -> None: + """Route non-e2e single converted detection and pose outputs by task shape.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": task, + "nl": 3, + "reg_max": 16, + "nc": 3 if task == "object_detection" else 1, + "conf_thres": 0.5, + "iou_thres": 0.7, + "e2e": False, + **post_cfg_extra, + } + postprocessor = build_postprocess(pre_cfg, post_cfg) + converted = torch.zeros((1, 2, converted_dim), dtype=torch.float32) + + result = postprocessor(converted) + + assert isinstance(result, torch.Tensor) + assert result.shape == (1, converted_dim, 2) + + +def test_non_e2e_segmentation_uses_converted_detections_and_proto() -> None: + """Route non-e2e segmentation converted detections with prototype masks.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "reg_max": 16, + "nc": 3, + "n_extra": 32, + "conf_thres": 0.5, + "iou_thres": 0.7, + "e2e": False, + } + postprocessor = build_postprocess(pre_cfg, post_cfg) + detections = torch.zeros((1, 2, 39), dtype=torch.float32) + proto = torch.zeros((1, 16, 16, 32), dtype=torch.float32) + + result = postprocessor([detections, proto]) + + assert isinstance(result, list) + assert result[0].shape == (1, 39, 2) + assert result[1].shape == (1, 32, 16, 16) + + +def test_dflfree_detection_accepts_decode_true_mxq_parts_with_reducemax() -> None: + """Accept split decode-true DFL-free detection outputs with an extra reducemax tensor.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "dflfree": True, + "nc": 3, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.1, 0.9, 0.2], [0.2, 0.3, 0.4]]], dtype=torch.float32) + reducemax = scores.max(dim=-1, keepdim=True).values + + result = postprocessor([scores, reducemax, boxes]) + + assert len(result) == 1 + assert result[0].shape == (1, 6) + assert torch.equal(result[0][0, :4], boxes[0, 0]) + assert torch.allclose(result[0][0, 4], torch.tensor(0.9)) + assert torch.allclose(result[0][0, 5], torch.tensor(1.0)) + + +def test_dflfree_detection_accepts_batched_decode_true_mxq_parts_with_reducemax() -> ( + None +): + """Preserve batched split decode-true DFL-free detection outputs.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "dflfree": True, + "nc": 3, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [ + [[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]], + [[50.0, 60.0, 70.0, 80.0], [51.0, 61.0, 71.0, 81.0]], + ], + dtype=torch.float32, + ) + scores = torch.tensor( + [ + [[0.1, 0.9, 0.2], [0.2, 0.3, 0.4]], + [[0.8, 0.1, 0.2], [0.1, 0.2, 0.3]], + ], + dtype=torch.float32, + ) + reducemax = scores.max(dim=-1, keepdim=True).values.unsqueeze(1) + batched_scores = scores.unsqueeze(1) + batched_boxes = boxes.unsqueeze(1) + + result = postprocessor([batched_scores, reducemax, batched_boxes]) + + assert len(result) == 2 + assert result[0].shape == (1, 6) + assert result[1].shape == (1, 6) + assert torch.equal(result[0][0, :4], boxes[0, 0]) + assert torch.equal(result[1][0, :4], boxes[1, 0]) + assert torch.allclose(result[0][0, 4], torch.tensor(0.9)) + assert torch.allclose(result[1][0, 4], torch.tensor(0.8)) + assert torch.allclose(result[0][0, 5], torch.tensor(1.0)) + assert torch.allclose(result[1][0, 5], torch.tensor(0.0)) + + +def test_dflfree_detection_distinguishes_equal_width_box_and_score_parts() -> None: + """Use reducemax to distinguish boxes from scores when ``nc == 4``.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "dflfree": True, + "nc": 4, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor( + [[[0.1, 0.9, 0.2, 0.3], [0.2, 0.3, 0.4, 0.1]]], dtype=torch.float32 + ) + reducemax = scores.max(dim=-1, keepdim=True).values + + result = postprocessor([boxes, reducemax, scores]) + + assert len(result) == 1 + assert result[0].shape == (1, 6) + assert torch.equal(result[0][0, :4], boxes[0, 0]) + assert torch.allclose(result[0][0, 4], torch.tensor(0.9)) + assert torch.allclose(result[0][0, 5], torch.tensor(1.0)) + + +def test_dflfree_segmentation_accepts_decode_true_mxq_parts_with_reducemax() -> None: + """Accept split decode-true DFL-free segmentation outputs with reducemax and proto tensors.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "dflfree": True, + "nc": 3, + "n_extra": 2, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.1, 0.9, 0.2], [0.2, 0.3, 0.4]]], dtype=torch.float32) + reducemax = scores.max(dim=-1, keepdim=True).values + coeffs = torch.tensor([[[0.4, 0.6], [0.2, 0.1]]], dtype=torch.float32) + proto = torch.zeros((9, 9, 2), dtype=torch.float32) + + result = postprocessor([coeffs, scores, reducemax, proto, boxes]) + + assert len(result) == 1 + assert result[0][0].shape == (1, 8) + assert result[0][1].shape == (1, 64, 64) + + +def test_dflfree_segmentation_accepts_approximate_reducemax_scores() -> None: + """Match quantized reducemax tensors without appending them as mask coefficients.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "dflfree": True, + "nc": 3, + "n_extra": 2, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.1, 0.9, 0.2], [0.2, 0.3, 0.4]]], dtype=torch.float32) + reducemax = scores.max(dim=-1, keepdim=True).values - 0.01 + coeffs = torch.tensor([[[0.4, 0.6], [0.2, 0.1]]], dtype=torch.float32) + proto = torch.zeros((9, 9, 2), dtype=torch.float32) + + result = postprocessor([coeffs, scores, reducemax, proto, boxes]) + + assert len(result) == 1 + assert result[0][0].shape == (1, 8) + assert result[0][1].shape == (1, 64, 64) + + +def test_dflfree_segmentation_excludes_reducemax_from_proto_candidates() -> None: + """Ignore unused reducemax tensors when ``n_extra == 1`` and selecting the proto tensor.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "dflfree": True, + "nc": 3, + "n_extra": 1, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.1, 0.9, 0.2], [0.2, 0.3, 0.4]]], dtype=torch.float32) + reducemax = scores.max(dim=-1, keepdim=True).values + coeffs = torch.tensor([[[0.6], [0.1]]], dtype=torch.float32) + proto = torch.zeros((9, 9, 1), dtype=torch.float32) + + result = postprocessor([coeffs, scores, reducemax, proto, boxes]) + + assert len(result) == 1 + assert result[0][0].shape == (1, 7) + assert result[0][1].shape == (1, 64, 64) + + +def test_non_e2e_dflfree_segmentation_accepts_decode_true_mxq_parts_with_reducemax() -> ( + None +): + """Route 5-part decode-true segmentation outputs through the segmentation non-e2e path.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "dflfree": True, + "nc": 3, + "n_extra": 2, + "conf_thres": 0.5, + "iou_thres": 0.7, + "e2e": False, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.1, 0.9, 0.2], [0.2, 0.3, 0.4]]], dtype=torch.float32) + reducemax = scores.max(dim=-1, keepdim=True).values + coeffs = torch.tensor([[[0.4, 0.6], [0.2, 0.1]]], dtype=torch.float32) + proto = torch.zeros((1, 9, 9, 2), dtype=torch.float32) + + result = postprocessor([coeffs, scores, reducemax, proto, boxes]) + + assert isinstance(result, list) + assert result[0].shape == (1, 300, 8) + assert result[1].shape == (1, 2, 9, 9) + assert torch.equal(result[0][0, 0, :4], boxes[0, 0]) + assert torch.allclose(result[0][0, 0, 4], torch.tensor(0.9)) + assert torch.allclose(result[0][0, 0, 5], torch.tensor(1.0)) + assert torch.equal(result[0][0, 0, 6:], coeffs[0, 0]) + + +def test_dflfree_pose_accepts_decode_true_mxq_parts_with_reducemax() -> None: + """Accept split decode-true DFL-free pose outputs with a duplicate score max tensor.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "pose_estimation", + "nl": 3, + "dflfree": True, + "nc": 1, + "n_extra": 51, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.9], [0.4]]], dtype=torch.float32) + reducemax = scores.clone() + keypoints = torch.arange(102, dtype=torch.float32).reshape(1, 2, 51) + + result = postprocessor([reducemax, scores, boxes, keypoints]) + + assert len(result) == 1 + assert result[0].shape == (1, 57) + assert torch.equal(result[0][0, :4], boxes[0, 0]) + assert torch.equal(result[0][0, 6:], keypoints[0, 0]) + + +def test_dflfree_pose_prefers_score_tensor_over_reducemax_duplicate() -> None: + """Use the actual score tensor when MXQ exports both reducemax and score parts.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "pose_estimation", + "nl": 3, + "dflfree": True, + "nc": 1, + "n_extra": 51, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + reducemax = torch.tensor([[[0.88], [0.39]]], dtype=torch.float32) + scores = torch.tensor([[[0.9], [0.4]]], dtype=torch.float32) + keypoints = torch.arange(102, dtype=torch.float32).reshape(1, 2, 51) + + for output_order in ( + [reducemax, scores, boxes, keypoints], + [scores, reducemax, boxes, keypoints], + ): + result = postprocessor(output_order) + + assert len(result) == 1 + assert result[0].shape == (1, 57) + assert torch.allclose(result[0][0, 4], torch.tensor(0.9)) + assert torch.equal(result[0][0, 6:], keypoints[0, 0]) + + +def test_non_e2e_dflfree_pose_accepts_decode_true_mxq_parts_with_reducemax() -> None: + """Route 4-part decode-true pose outputs through the pose non-e2e path.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "pose_estimation", + "nl": 3, + "dflfree": True, + "nc": 1, + "n_extra": 51, + "conf_thres": 0.5, + "iou_thres": 0.7, + "e2e": False, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + boxes = torch.tensor( + [[[10.0, 20.0, 30.0, 40.0], [11.0, 21.0, 31.0, 41.0]]], dtype=torch.float32 + ) + scores = torch.tensor([[[0.9], [0.4]]], dtype=torch.float32) + reducemax = scores.clone() + keypoints = torch.arange(102, dtype=torch.float32).reshape(1, 2, 51) + + result = postprocessor([reducemax, scores, boxes, keypoints]) + + assert isinstance(result, torch.Tensor) + assert result.shape == (1, 300, 57) + first_detection = result[0][0] + assert torch.equal(first_detection[:4], boxes[0, 0]) + assert torch.allclose(first_detection[4], torch.tensor(0.9)) + assert torch.allclose(first_detection[5], torch.tensor(0.0)) + assert torch.equal(first_detection[6:], keypoints[0, 0]) + + +def test_non_e2e_dflfree_obb_preserves_canonical_row_width() -> None: + """Pad non-e2e DFL-free OBB converted outputs using canonical pre-NMS row widths.""" + + expected_max_det = 300 + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "obb", + "nl": 3, + "nc": 15, + "n_extra": 1, + "conf_thres": 0.8, + "iou_thres": 0.7, + "dflfree": True, + "e2e": False, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + + result = postprocessor(_make_converted_obb_parts()) + + assert isinstance(result, torch.Tensor) + assert result.shape == (1, expected_max_det, 20) + first_image = result[0] + assert torch.equal(first_image[:3], _make_converted_obb_rows()[0]) + assert torch.count_nonzero(first_image[3:]) == 0 + + +def test_raw_mxq_like_outputs_are_not_final_detections() -> None: + """Do not treat split MXQ-style head tensors as already-decoded detections.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "dflfree": True, + "nc": 80, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + raw_outputs: ListTensorLike = [ + torch.zeros((1, 80, 80, 4), dtype=torch.float32), + torch.zeros((1, 80, 80, 80), dtype=torch.float32), + ] + + detections, proto = postprocessor.extract_final_outputs(raw_outputs) + + assert detections is None + assert proto is None + + +def _make_anchorless_obb_mxq_heads() -> ListTensorLike: + """Build synthetic channel-first MXQ OBB heads for anchorless models.""" + + outputs: list[np.ndarray] = [] + scale_specs = ((8, 0, 0, 0), (4, 0, 0, 1), (2, 0, 0, 2)) + for size, y_idx, x_idx, cls_idx in scale_specs: + det = np.full((1, 64, size, size), -10.0, dtype=np.float32) + for side in range(4): + det[0, side * 16 + 2, y_idx, x_idx] = 10.0 + + cls = np.full((1, 15, size, size), -10.0, dtype=np.float32) + cls[0, cls_idx, y_idx, x_idx] = 10.0 + + angle = np.zeros((1, 1, size, size), dtype=np.float32) + outputs.extend([det, cls, angle]) + return outputs + + +def _make_dflfree_obb_mxq_heads() -> ListTensorLike: + """Build synthetic channel-first MXQ OBB heads for DFL-free models.""" + + outputs: list[np.ndarray] = [] + scale_specs = ((8, 0, 0, 0), (4, 0, 0, 1), (2, 0, 0, 2)) + for size, y_idx, x_idx, cls_idx in scale_specs: + det = np.zeros((1, 4, size, size), dtype=np.float32) + det[0, :, y_idx, x_idx] = np.array([2.0, 2.0, 2.0, 2.0], dtype=np.float32) + + cls = np.full((1, 15, size, size), -10.0, dtype=np.float32) + cls[0, cls_idx, y_idx, x_idx] = 10.0 + + angle = np.zeros((1, 1, size, size), dtype=np.float32) + outputs.extend([det, cls, angle]) + return outputs + + +def _make_converted_obb_rows() -> torch.Tensor: + """Build synthetic converted OBB rows in canonical row-major format.""" + + output = torch.zeros((1, 3, 20), dtype=torch.float32) + output[0, :, :4] = torch.tensor( + [ + [12.0, 12.0, 6.0, 4.0], + [32.0, 24.0, 8.0, 6.0], + [48.0, 48.0, 10.0, 8.0], + ], + dtype=torch.float32, + ) + output[0, 0, 4] = 0.95 + output[0, 1, 5] = 0.90 + output[0, 2, 6] = 0.85 + output[0, :, -1] = torch.tensor([0.0, 0.1, -0.2], dtype=torch.float32) + return output + + +def _make_converted_obb_parts() -> ListTensorLike: + """Build shuffled converted MXQ OBB parts for decode-true outputs.""" + + rows = _make_converted_obb_rows() + boxes = rows[:, :, :4].unsqueeze(1) + scores = rows[:, :, 4:-1].unsqueeze(1) + angle = rows[:, :, -1:].unsqueeze(1) + return [angle, boxes, scores] + + +def _make_split_converted_obb_parts(class_first: bool) -> ListTensorLike: + """Build decode-true MXQ OBB parts split into box subchannels.""" + + rows = _make_converted_obb_rows() + scores = rows[:, :, 4:-1].transpose(1, 2) + angle = rows[:, :, -1:].transpose(1, 2) + xy = rows[:, :, :2].transpose(1, 2) + width = rows[:, :, 2:3].transpose(1, 2) + height = rows[:, :, 3:4].transpose(1, 2) + if class_first: + return [scores, angle, xy, width, height] + return [angle, scores, xy, width, height] + + +@pytest.mark.parametrize("dflfree", [False, True]) +def test_obb_accepts_single_converted_output(dflfree: bool) -> None: + """Accept ONNX-style converted OBB tensors before rotated NMS.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "obb", + "nl": 3, + "nc": 15, + "n_extra": 1, + "conf_thres": 0.8, + "iou_thres": 0.7, + } + if dflfree: + post_cfg["dflfree"] = True + else: + post_cfg["reg_max"] = 16 + + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + result = postprocessor(_make_converted_obb_rows().transpose(1, 2)) + + assert len(result) == 1 + assert result[0].shape == (3, 7) + assert torch.equal(result[0][:, 5], torch.tensor([0.0, 1.0, 2.0])) + + +@pytest.mark.parametrize("dflfree", [False, True]) +@pytest.mark.parametrize("class_first", [False, True]) +def test_obb_accepts_decode_true_converted_mxq_parts( + dflfree: bool, class_first: bool +) -> None: + """Accept converted MXQ OBB box, class, and angle outputs before rotated NMS.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "obb", + "nl": 3, + "nc": 15, + "n_extra": 1, + "conf_thres": 0.8, + "iou_thres": 0.7, + } + if dflfree: + post_cfg["dflfree"] = True + else: + post_cfg["reg_max"] = 16 + + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + result = postprocessor(_make_converted_obb_parts()) + split_result = postprocessor(_make_split_converted_obb_parts(class_first)) + + assert len(result) == 1 + assert result[0].shape == (3, 7) + assert torch.equal(result[0][:, 5], torch.tensor([0.0, 1.0, 2.0])) + assert len(split_result) == 1 + assert split_result[0].shape == (3, 7) + assert torch.equal(split_result[0][:, 5], torch.tensor([0.0, 1.0, 2.0])) + + +def test_anchorless_obb_accepts_channel_first_mxq_heads_and_plots_airport( + tmp_path: Path, +) -> None: + """Accept channel-first MXQ OBB heads for YOLOv8/YOLO11-style models.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "obb", + "nl": 3, + "reg_max": 16, + "nc": 15, + "n_extra": 1, + "conf_thres": 0.8, + "iou_thres": 0.7, + } + image_path = np.zeros((64, 64, 3), dtype=np.uint8) + save_path = tmp_path / "anchorless_obb_airport.jpg" + + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + result = postprocessor(_make_anchorless_obb_mxq_heads()) + + assert len(result) == 1 + assert result[0].shape[1] == 7 + assert result[0].shape[0] >= 1 + + plotted = Results(pre_cfg, post_cfg, result).plot( + image_path, save_path=str(save_path) + ) + + assert plotted is not None + assert save_path.is_file() + + +def test_dflfree_obb_accepts_channel_first_mxq_heads_and_plots_airport( + tmp_path: Path, +) -> None: + """Accept channel-first MXQ OBB heads for YOLO26-style models.""" + + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "obb", + "nl": 3, + "dflfree": True, + "nc": 15, + "n_extra": 1, + "conf_thres": 0.8, + "iou_thres": 0.7, + } + image_path = np.zeros((64, 64, 3), dtype=np.uint8) + save_path = tmp_path / "dflfree_obb_airport.jpg" + + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + result = postprocessor(_make_dflfree_obb_mxq_heads()) + + assert len(result) == 1 + assert result[0].shape[1] == 7 + assert result[0].shape[0] >= 1 + + plotted = Results(pre_cfg, post_cfg, result).plot( + image_path, save_path=str(save_path) + ) + + assert plotted is not None + assert save_path.is_file() + + +def test_anchor_segmentation_ignores_auxiliary_onnx_heads() -> None: + """Use converted YOLOv5-seg ONNX outputs and ignore auxiliary raw heads.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "instance_segmentation", + "anchors": [ + [10, 13, 16, 30, 33, 23], + [30, 61, 62, 45, 59, 119], + [116, 90, 156, 198, 373, 326], + ], + "n_extra": 32, + "nc": 80, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = cast(YOLODetectionPostBase, build_postprocess(pre_cfg, post_cfg)) + det = torch.zeros((1, 2, 117), dtype=torch.float32) + proto = torch.zeros((1, 32, 160, 160), dtype=torch.float32) + aux_heads = [ + torch.zeros((1, 3, 80, 80, 117), dtype=torch.float32), + torch.zeros((1, 3, 40, 40, 117), dtype=torch.float32), + torch.zeros((1, 3, 20, 20, 117), dtype=torch.float32), + ] + + result = postprocessor([det, proto, *aux_heads]) + + assert len(result) == 1 + assert result[0][0].shape == (0, 38) + assert result[0][1].shape == (0, 640, 640) + + +def test_anchor_segmentation_accepts_mxq_raw_heads() -> None: + """Use raw YOLOv5-seg MXQ heads when no converted detection tensor is present.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "instance_segmentation", + "anchors": [ + [10, 13, 16, 30, 33, 23], + [30, 61, 62, 45, 59, 119], + [116, 90, 156, 198, 373, 326], + ], + "n_extra": 32, + "nc": 80, + "conf_thres": 0.5, + "iou_thres": 0.7, + } + postprocessor = build_postprocess(pre_cfg, post_cfg) + raw_outputs = [ + torch.zeros((2, 20, 20, 351), dtype=torch.float32), + torch.zeros((2, 40, 40, 351), dtype=torch.float32), + torch.zeros((2, 80, 80, 351), dtype=torch.float32), + torch.zeros((2, 160, 160, 32), dtype=torch.float32), + ] + + result = postprocessor(raw_outputs) + + assert len(result) == 2 + assert result[0][0].shape == (0, 38) + assert result[0][1].shape == (0, 640, 640) + assert result[1][0].shape == (0, 38) + assert result[1][1].shape == (0, 640, 640) + + +def test_anchorless_prediction_nms_keeps_best_class_per_box() -> None: + """Keep ordinary prediction output limited to the best class per box.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "reg_max": 16, + "nc": 3, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + decoded = torch.tensor( + [ + [ + [10.0, 50.0], + [10.0, 50.0], + [20.0, 60.0], + [20.0, 60.0], + [0.90, 0.80], + [0.10, 0.85], + [0.10, 0.70], + ] + ], + dtype=torch.float32, + ) + + result = postprocessor.nms([decoded[0]]) + + assert len(result) == 1 + assert result[0].shape == (2, 6) + assert torch.equal(result[0][:, 5], torch.tensor([0.0, 1.0])) + + +def test_anchor_prediction_nms_keeps_best_class_per_box() -> None: + """Keep ordinary anchor-model prediction output single-label per box.""" + + postprocessor = cast( + YOLODetectionPostBase, + build_postprocess( + {"LetterBox": {"img_size": [640, 640]}}, + { + "task": "object_detection", + "anchors": [[10, 13, 16, 30, 33, 23]], + "nl": 1, + "nc": 2, + "conf_thres": 0.25, + "iou_thres": 0.7, + }, + ), + ) + decoded = torch.tensor( + [[50.0, 50.0, 20.0, 20.0, 1.0, 0.9, 0.8]], dtype=torch.float32 + ) + + prediction = postprocessor.nms([decoded]) + validation = postprocessor.nms_multilabel([decoded]) + + assert prediction[0].shape == (1, 6) + assert validation[0].shape == (2, 6) + + +@pytest.mark.parametrize( + ("layout", "candidate_count"), + [ + ("channels_first", 117), + ("candidates_first", 117), + ("channels_first", 116), + ("candidates_first", 116), + ], +) +def test_anchorless_nms_normalizes_known_layout_before_suppression( + layout: AnchorlessOutputLayout, + candidate_count: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Use source provenance to normalize raw and converted segmentation outputs.""" + + pre_cfg = {"LetterBox": {"img_size": [640, 640]}} + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "reg_max": 16, + "nc": 80, + "n_extra": 32, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + canonical = torch.arange(candidate_count * 116, dtype=torch.float32).reshape( + candidate_count, 116 + ) + source = canonical.transpose(0, 1) if layout == "channels_first" else canonical + captured: list[torch.Tensor] = [] + + def capture_canonical(xi: torch.Tensor, **_: Any) -> torch.Tensor: + captured.append(xi) + return torch.empty((0, 38), dtype=torch.float32) + + monkeypatch.setattr(postprocessor, "_nms_single_legacy_rows", capture_canonical) + + postprocessor.nms(_AnchorlessNMSInput([source], layout)) + + assert captured[0].shape == (candidate_count, 116) + torch.testing.assert_close(captured[0], canonical) + + +def test_anchorless_nms_shape_fallback_prefers_raw_layout_for_square_tensor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Treat an ambiguous provenance-free square tensor as channel-first.""" + + pre_cfg = {"LetterBox": {"img_size": [640, 640]}} + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "reg_max": 16, + "nc": 80, + "n_extra": 32, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + source = torch.arange(116 * 116, dtype=torch.float32).reshape(116, 116) + captured: list[torch.Tensor] = [] + + def capture_canonical(xi: torch.Tensor, **_: Any) -> torch.Tensor: + captured.append(xi) + return torch.empty((0, 38), dtype=torch.float32) + + monkeypatch.setattr(postprocessor, "_nms_single_legacy_rows", capture_canonical) + + postprocessor.nms([source]) + + torch.testing.assert_close(captured[0], source.transpose(0, 1)) + + +def test_anchorless_segmentation_preprocess_preserves_layout_provenance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Tag decoded MXQ heads as channel-first and converted outputs as candidates-first.""" + + pre_cfg = {"LetterBox": {"img_size": [640, 640]}} + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "reg_max": 16, + "nc": 80, + "n_extra": 32, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + decoded = torch.zeros((116, 116), dtype=torch.float32) + converted = decoded.transpose(0, 1).unsqueeze(0) + proto = torch.zeros((1, 160, 160, 32), dtype=torch.float32) + rearranged = torch.empty(0) + + monkeypatch.setattr(postprocessor, "rearrange", lambda _: (rearranged, proto)) + monkeypatch.setattr( + postprocessor, "decode", lambda value: [decoded] if value is rearranged else [] + ) + raw_predictions, raw_proto = postprocessor._pre_process([torch.empty(0)] * 3) + + monkeypatch.setattr(postprocessor, "conversion", lambda _: (converted, proto)) + monkeypatch.setattr( + postprocessor, "filter_conversion", lambda _: [converted.squeeze(0)] + ) + converted_predictions, converted_proto = postprocessor._pre_process( + [torch.empty(0)] * 2 + ) + + assert raw_predictions.layout == "channels_first" + assert isinstance(raw_predictions.detections, list) + assert raw_predictions.detections[0] is decoded + assert raw_proto is proto + assert converted_predictions.layout == "candidates_first" + assert isinstance(converted_predictions.detections, list) + torch.testing.assert_close( + converted_predictions.detections[0], converted.squeeze(0) + ) + assert converted_proto is proto + + +def test_anchorless_nms_normalizes_detection_without_extra_channels( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Normalize raw detection output when no segmentation or pose extras exist.""" + + pre_cfg = {"LetterBox": {"img_size": [640, 640]}} + post_cfg = { + "task": "object_detection", + "nl": 3, + "reg_max": 16, + "nc": 80, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + source = torch.arange(84 * 91, dtype=torch.float32).reshape(84, 91) + captured: list[torch.Tensor] = [] + + def capture_canonical(xi: torch.Tensor, **_: Any) -> torch.Tensor: + captured.append(xi) + return torch.empty((0, 6), dtype=torch.float32) + + monkeypatch.setattr(postprocessor, "_nms_single_legacy_rows", capture_canonical) + + postprocessor.nms(_AnchorlessNMSInput([source], "channels_first")) + + assert postprocessor.n_extra == 0 + assert captured[0].shape == (91, 84) + torch.testing.assert_close(captured[0], source.transpose(0, 1)) + + +def test_anchorless_nonambiguous_layouts_keep_identical_nms_results() -> None: + """Keep existing suppression results after canonical layout normalization.""" + + pre_cfg = {"LetterBox": {"img_size": [640, 640]}} + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "reg_max": 16, + "nc": 80, + "n_extra": 32, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + canonical = torch.zeros((2, 116), dtype=torch.float32) + canonical[0, :4] = torch.tensor([10.0, 10.0, 20.0, 20.0]) + canonical[1, :4] = torch.tensor([40.0, 40.0, 60.0, 60.0]) + canonical[0, 4] = 0.9 + canonical[1, 5] = 0.8 + canonical[:, 84:] = torch.arange(64, dtype=torch.float32).reshape(2, 32) + + raw_result = postprocessor.nms( + _AnchorlessNMSInput([canonical.transpose(0, 1)], "channels_first") + ) + converted_result = postprocessor.nms( + _AnchorlessNMSInput([canonical], "candidates_first") + ) + + torch.testing.assert_close(raw_result[0], converted_result[0]) + + +def test_anchorless_validation_nms_keeps_multilabel_candidates() -> None: + """Retain all above-threshold classes when validation requests Ultralytics semantics.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "object_detection", + "nl": 3, + "reg_max": 16, + "nc": 3, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + decoded = torch.tensor( + [ + [10.0, 10.0, 20.0, 20.0, 0.90, 0.80, 0.10], + [50.0, 50.0, 60.0, 60.0, 0.20, 0.85, 0.70], + ], + dtype=torch.float32, + ) + + result = postprocessor.nms([decoded], multi_label=True) + + assert len(result) == 1 + assert result[0].shape == (4, 6) + assert torch.equal(result[0][:, 5], torch.tensor([0.0, 1.0, 1.0, 2.0])) + + +def test_anchorless_segmentation_validation_duplicates_mask_coefficients_per_class() -> ( + None +): + """Copy mask coefficients when validation expands a box into multiple class candidates.""" + + pre_cfg = { + "LetterBox": { + "img_size": [640, 640], + } + } + post_cfg = { + "task": "instance_segmentation", + "nl": 3, + "reg_max": 16, + "nc": 3, + "n_extra": 2, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast( + YOLOAnchorlessDetectionPost, build_postprocess(pre_cfg, post_cfg) + ) + decoded = torch.tensor( + [[10.0, 10.0, 20.0, 20.0, 0.90, 0.80, 0.10, 0.25, -0.50]], + dtype=torch.float32, + ) + + result = postprocessor.nms([decoded], multi_label=True) + + assert result[0].shape == (2, 8) + assert torch.equal(result[0][:, 5], torch.tensor([0.0, 1.0])) + assert torch.equal(result[0][:, 6:], decoded[:, 7:].repeat(2, 1)) + + +def test_anchorless_pose_single_and_batch_decode_are_equivalent() -> None: + """Keep batched pose decoding equivalent to the verified v1.5.1 per-image formula.""" + + torch.manual_seed(0) + pre_cfg = { + "LetterBox": { + "img_size": [64, 64], + } + } + post_cfg = { + "task": "pose_estimation", + "nl": 3, + "reg_max": 16, + "nc": 1, + "n_extra": 51, + "conf_thres": 0.25, + "iou_thres": 0.7, + } + postprocessor = cast(YOLOAnchorlessPosePost, build_postprocess(pre_cfg, post_cfg)) + anchor_count = postprocessor.anchors_as_tensor().shape[-1] + raw = torch.randn((2, 116, anchor_count), dtype=torch.float32) + raw[:, 64, :] = 10.0 + + batched = postprocessor.decode_batch(raw) + per_image = torch.stack([postprocessor.process_box_cls(image) for image in raw]) + + torch.testing.assert_close(batched, per_image) + + +def test_scale_coords_matches_ultralytics_rounding() -> None: + """Match upstream letterbox padding rounding for keypoint scaling.""" + + coords = torch.tensor( + [[[160.0, 100.0, 1.0], [480.0, 500.0, 1.0]]], dtype=torch.float32 + ) + + ratio_pad = resolve_ratio_pad((640, 640), (581, 640)) + scaled = scale_coords((640, 640), coords.clone(), (581, 640)) + + expected = torch.tensor( + [[[160.0, 71.0, 1.0], [480.0, 471.0, 1.0]]], dtype=torch.float32 + ) + assert ratio_pad == ((1.0, 1.0), (0, 29)) + torch.testing.assert_close(scaled, expected) + + +def test_letterbox_metadata_normalization_is_batch_aware() -> None: + """Share consistent shape and ratio-pad normalization across postprocessors.""" + + ratio_pad = ((1.0, 1.0), (0.0, 80.0)) + + assert normalize_image_shapes((481, 640), batch_size=2) == [(481, 640), (481, 640)] + assert normalize_ratio_pads(ratio_pad, batch_size=2) == [ratio_pad, ratio_pad] + assert normalize_ratio_pads((ratio_pad, ratio_pad), batch_size=2) == [ + ratio_pad, + ratio_pad, + ] + labels, boxes, scores = nmsout2eval( + [torch.zeros((0, 6)), torch.zeros((0, 6))], + (640, 640), + (481, 640), + ratio_pads=ratio_pad, + ) + assert labels == boxes == scores == [[], []] + with pytest.raises(ValueError, match="Expected 2 image shapes"): + normalize_image_shapes([(481, 640)], batch_size=2) + with pytest.raises(ValueError, match="Expected 2 ratio_pad values"): + normalize_ratio_pads([ratio_pad], batch_size=2) + + +def test_scale_masks_matches_ultralytics_rounding() -> None: + """Crop mask padding with the same rounding as upstream Ultralytics.""" + + masks = torch.zeros((1, 640, 640), dtype=torch.float32) + masks[:, 80:560, :] = 1.0 + + scaled = scale_masks(masks, (481, 640)) + + assert scaled.shape == (1, 481, 640) + assert float(scaled[:, 0, :].max()) == pytest.approx(0.0) + assert float(scaled[:, 1, :].max()) > 0.0 + assert float(scaled[:, -1, :].max()) == pytest.approx(1.0) + + +@pytest.mark.parametrize("shape", [(640, 640), (481, 640)]) +def test_roi_prototype_masking_preserves_full_mask_result( + shape: tuple[int, int], +) -> None: + """ROI masking must retain the conventional path's binary mask exactly.""" + + generator = torch.Generator().manual_seed(42) + proto = torch.randn((32, 160, 160), generator=generator) + coefficients = torch.randn((40, 32), generator=generator) + # Small, fractional boxes activate the ROI path and exercise crop bounds. + starts = torch.rand((40, 2), generator=generator) + starts[:, 0] *= shape[1] - 80 + starts[:, 1] *= shape[0] - 80 + boxes = torch.cat((starts + 0.25, starts + 48.75), dim=1) + + channels, mask_h, mask_w = proto.shape + full = (coefficients @ proto.float().view(channels, -1)).view(-1, mask_h, mask_w) + expected = crop_mask(scale_masks(full, shape), boxes).gt_(0.0) + + actual = process_mask_upsample(proto, coefficients, boxes, shape) + + assert torch.equal(actual, expected) + + +def test_preprocess_with_metadata_returns_letterbox_ratio_pad() -> None: + """Expose exact LetterBox ratio and integer padding for validation scaling.""" + + engine = MBLT_Engine.__new__(MBLT_Engine) + engine.pre_cfg = { + "Reader": {"style": "numpy"}, + "LetterBox": {"img_size": [640, 640]}, + "SetOrder": {"shape": "HWC"}, + "Normalize": {"style": "cv"}, + } + engine.preprocessor = build_preprocess(engine.pre_cfg) + image = np.zeros((481, 640, 3), dtype=np.uint8) + + processed, metadata = engine.preprocess_with_metadata(image) + + assert processed.shape == (640, 640, 3) + assert metadata["img0_shape"] == (481, 640) + assert metadata["ratio_pad"] == ((1.0, 1.0), (0, 79)) + + +def test_nmsout2eval_matches_coco_json_format_without_mutation() -> None: + """Serialize detections like Ultralytics validation without changing NMS output.""" + + nms_out = torch.tensor( + [ + [10.12345, 20.23456, 110.34567, 220.45678, 0.876543, 0.0], + ], + dtype=torch.float32, + ) + original = nms_out.clone() + + labels, boxes, scores = nmsout2eval([nms_out], (640, 640), [(640, 640)]) + + assert labels == [[1]] + assert boxes == [[[10.123, 20.235, 100.222, 200.222]]] + assert scores == [[0.87654]] + assert torch.equal(nms_out, original) + + +@pytest.mark.parametrize("class_id", [-1.0, 1.9, 80.0, float("nan"), float("inf")]) +def test_nmsout2eval_rejects_invalid_coco_class_ids(class_id: float) -> None: + """Reject malformed decoded class IDs before COCO taxonomy remapping.""" + + nms_out = torch.tensor( + [[10.0, 20.0, 110.0, 220.0, 0.9, class_id]], dtype=torch.float32 + ) + + with pytest.raises(ValueError, match="finite integral values"): + nmsout2eval([nms_out], (640, 640), [(640, 640)]) + + +def test_nmsout2eval_uses_explicit_ratio_pad() -> None: + """Use dataloader-provided LetterBox padding instead of recomputing from shape.""" + + nms_out = torch.tensor([[0.0, 79.0, 10.0, 89.0, 0.9, 0.0]], dtype=torch.float32) + + _labels, boxes, _scores = nmsout2eval( + [nms_out], + (640, 640), + [(481, 640)], + ratio_pads=[((1.0, 1.0), (0, 79))], + ) + + assert boxes == [[[0.0, 0.0, 10.0, 10.0]]]