diff --git a/.agents/skills/mblt-vision/SKILL.md b/.agents/skills/mblt-vision/SKILL.md index 1a19dd8..6bc345e 100644 --- a/.agents/skills/mblt-vision/SKILL.md +++ b/.agents/skills/mblt-vision/SKILL.md @@ -28,6 +28,23 @@ description: >- 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. +- A promptable or multi-artifact model (see mask_generation/SAM2HieraLarge) bypasses + MBLT_Engine.__init__, build_preprocess/build_postprocess, and create_model_class entirely, + implementing its own preprocess/predict methods; it still subclasses MBLT_Engine only for + list_models() discovery. Reuse wrapper.download_hub_artifact for any additional Hub artifact + rather than duplicating Hub-resolution logic. +- Never depend on the PyPI `sam2` package (unofficial third-party mirror) or a manually cloned + facebookresearch/sam2 checkout. mask_generation's host-side prompt encoding is a from-scratch + port verified bit-for-bit against the real predictor, backed by a small Hub-hosted weights + bundle, not package data. +- mask_generation validates on SA-V val via datasets/sa-v.yaml. SA-V is NOT auto-downloaded and + must never be mirrored on Mobilint infrastructure: Meta gates it behind a download form, so the + user supplies the official sav_val.tar (or its extracted directory) with + --annotation-dir/--image-dir, like Cityscapes. Not sha256-pinned (user-supplied, not fetched); + identity comes from the readiness inventory. CC BY 4.0 by Meta AI. Registering a new dataset + requires readiness (`_*_ready` + `dataset_ready` map) before the organizer, since staged + validation calls `dataset_ready`; also register the organizer in the + test_dataset_organizer.py parametrize lists. ## Processing and Results diff --git a/.claude/skills/mblt-vision-readme/SKILL.md b/.claude/skills/mblt-vision-readme/SKILL.md index ea6a148..2970e51 100644 --- a/.claude/skills/mblt-vision-readme/SKILL.md +++ b/.claude/skills/mblt-vision-readme/SKILL.md @@ -1,9 +1,41 @@ --- name: mblt-vision-readme -description: Write and maintain documentation for the standalone Mobilint Vision Python package. +description: >- + Write and maintain mblt-vision-python README documentation, API examples, model references, + and migration notes. --- # Mobilint Vision README Writing -Read and follow the canonical skill at -../../../.agents/skills/mblt-vision-readme/SKILL.md. +## 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/.claude/skills/mblt-vision/SKILL.md b/.claude/skills/mblt-vision/SKILL.md index 50d1f21..92f154b 100644 --- a/.claude/skills/mblt-vision/SKILL.md +++ b/.claude/skills/mblt-vision/SKILL.md @@ -1,12 +1,124 @@ --- name: mblt-vision -description: Work effectively on the standalone Mobilint Vision Python API and model registry. +description: >- + Work on the standalone Mobilint Vision Python API, model registry, preprocessing, + postprocessing, results, runtime integration, and package compatibility contracts. --- # Mobilint Vision Python -Read and follow the canonical skill at -../../../.agents/skills/mblt-vision/SKILL.md. -That canonical workflow defines WiderFace Hard AP as primary and Medium/Easy AP -as secondary metrics without mean AP, plus NYU Depth delta1 as primary with -abs_rel/RMSE secondary metrics. +## 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. +- A promptable or multi-artifact model (see mask_generation/SAM2HieraLarge) bypasses + MBLT_Engine.__init__, build_preprocess/build_postprocess, and create_model_class entirely, + implementing its own preprocess/predict methods; it still subclasses MBLT_Engine only for + list_models() discovery. Reuse wrapper.download_hub_artifact for any additional Hub artifact + rather than duplicating Hub-resolution logic. +- mask_generation supports framework="mxq" (default) and framework="onnx" with + MBLT_Engine-style inference/conflict semantics (explicit encoder/decoder path suffixes infer + the framework; NPU-only arguments are ignored for ONNX). The ONNX exports are same-stem + sam2_hiera_large_{encoder,decoder}.onnx at the Hub repo root (board-agnostic). Keep the two + graph contracts pinned in _sam2_contracts.py and validated at construction: MXQ takes six + flattened positional NHWC inputs; the ONNX graphs are NCHW with five named decoder inputs + and a dynamic token axis. Keep the prompt-encoding host path and classify_decoder_outputs + framework-independent; only fpn_from_onnx/prepare_decoder_tensors_onnx vs + fpn_from_runtime/prepare_decoder_tensors differ. eval_sav works unchanged for both. + Load the optional runtime before downloading artifacts; validate explicit artifact paths + (prompt weights included) with FileNotFoundError before any download; dispose a backend that + fails after create() inside its builder (the caller assigns it only on success) while + suppressing dispose failures so they cannot mask the original error; treat -1 in ONNX graph + validation as "must be dynamic", not a wildcard; validate point labels as + exactly 1/0; build host prompt tensors on the weights' device so device="cuda" works; and + reject single-artifact path options in every CLI command that builds the engine, not just + predict. +- Never depend on the PyPI `sam2` package (unofficial third-party mirror) or a manually cloned + facebookresearch/sam2 checkout. mask_generation's host-side prompt encoding is a from-scratch + port verified bit-for-bit against the real predictor, backed by a small Hub-hosted weights + bundle, not package data. +- mask_generation validates on SA-V val via datasets/sa-v.yaml. SA-V is NOT auto-downloaded and + must never be mirrored on Mobilint infrastructure: Meta gates it behind a download form, so the + user supplies the official sav_val.tar (or its extracted directory) with + --annotation-dir/--image-dir, like Cityscapes. Not sha256-pinned (user-supplied, not fetched); + identity comes from the readiness inventory. CC BY 4.0 by Meta AI. Readiness pins all three + inventory counts (155 videos / 293 masklets / 31967 masks); video and masklet totals alone + accept a truncated source. Both the organizer and readiness require every non-zero mask value + to be one object ID, since `{1, 2}` survives a unique-value count but `> 0` binarization makes + it all foreground. Readiness validates every mask (not just the first per video); bilevel masks + are validated from the header, since a 1-bit PNG cannot hold a non-zero background. Reject + `sav_val.txt` ids that fail SAV_VIDEO_ID_PATTERN or escape staging, before any write. Registering a new dataset requires readiness (`_*_ready` + `dataset_ready` + map) before the organizer, since staged validation calls `dataset_ready`; also register the + organizer in the test_dataset_organizer.py parametrize lists. + +## 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. +- Rank WiderFace evaluation by Hard-set AP. Expose Medium-set then Easy-set AP + as secondary metrics, and do not compute mean AP across difficulty splits. +- Rank NYU Depth evaluation by delta1. Expose abs_rel then RMSE (m) as + secondary metrics, with median-aligned metrics averaged per image. + +## 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. `assets/` holds development-only sample images: tracked in git by + deliberate exception, pruned from distribution via MANIFEST.in, never grown for one-off inputs. +- 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/.gitignore b/.gitignore index 2c98cff..004b95e 100644 --- a/.gitignore +++ b/.gitignore @@ -43,5 +43,9 @@ compile_commands.json Thumbs.db *.zip +*.tar *.tar.gz -*.egg \ No newline at end of file +*.tgz +*.egg + +runs/ \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md index 8c66d48..e870396 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -75,9 +75,110 @@ The current ownership boundary is deliberate: must remain lazy and report the appropriate package extra when unavailable. - For WiderFace evaluation, rank results by Hard-set AP and retain Medium-set then Easy-set AP as secondary metrics. Do not compute a mean across splits. +- `eval_sav` requires already-binarized candidate masks, enumerated per dtype (bool, integer + `{0, 1}`/`{0, 255}`, float `{0.0, 1.0}`). A weaker "single positive value" rule still admits a + probability map such as `{0.0, 0.5}`, or a uniform `0.5` candidate that `astype(bool)` turns + entirely into foreground. This is deliberately stricter than the SA-V ground-truth mask check, + where any single positive value is a legitimate object ID. - For NYU Depth evaluation, rank results by delta1 and retain abs_rel then RMSE (m) as secondary metrics. Median-align each image and average every metric per image, following Ultralytics' depth-validation convention. +- `mask_generation` (`SAM2HieraLarge`) is the precedent for a promptable, multi-artifact model: + it loads two backend instances (encoder + decoder) and takes point prompts, so it + bypasses `MBLT_Engine.__init__`, `build_preprocess`/`build_postprocess`, and + `create_model_class`'s single-artifact legacy constructor entirely, implementing its own + `preprocess`/`predict`/`predict_preprocessed` instead. Its `models/SAM2HieraLarge.yaml` is + documentation/test-metadata only, not actually loaded. Reuse `wrapper.download_hub_artifact` + (extracted from `MBLT_Engine._download_hub_artifact`) for any future model needing more than + one Hub artifact, rather than duplicating Hub-resolution logic. +- `mask_generation` supports both frameworks with `MBLT_Engine`-style semantics: + `framework="mxq"` (default, two `MobilintNPUBackend` MXQs from the board folder) or + `framework="onnx"` (two `ONNXBackend` sessions over the same-stem + `sam2_hiera_large_{encoder,decoder}.onnx` exports at the Hub repo root, board-agnostic). + Framework is inferred from explicit `encoder_onnx_path`/`decoder_onnx_path` vs + `encoder_mxq_path`/`decoder_mxq_path` suffixes and conflicts fail fast; NPU-only arguments are + ignored for ONNX. The two runtimes have different graph contracts, pinned in + `_sam2_contracts.py` and validated at construction: MXQ takes the six flattened positional + inputs (NHWC image and features), while the exported ONNX graphs are NCHW with five named + decoder inputs (`src_plus_pos_src` stays inside the graph) and a dynamic token axis. The + shared prompt-encoding host path and `classify_decoder_outputs` are framework-independent; + only the encoder-feed layout and decoder-feed builders differ + (`fpn_from_onnx`/`prepare_decoder_tensors_onnx` vs + `fpn_from_runtime`/`prepare_decoder_tensors`). The ONNX pipeline is numerically verified + against the official `facebookresearch/sam2` fp32 predictor (identical binary masks; + opt-in `tests/test_mask_generation_onnx.py` covers it end-to-end without NPU hardware), + and `eval_sav` works unchanged for both frameworks. Resolve the optional runtime before + downloading any artifact, so a missing `onnxruntime` reports the package extra rather than a + network failure. Build each backend so it disposes itself when `launch()`/graph validation + fails after `create()`: the caller only assigns it on success, so the constructor's cleanup + cannot otherwise reach it. Suppress disposal failures while unwinding (both in those builders + and in the constructor, which uses `_close(suppress_errors=True)` like the base engine) so a + failing `dispose()` cannot replace the original construction error. Validate every explicitly + supplied artifact path -- prompt weights included -- with a fail-fast `FileNotFoundError` after + the argument-coherence checks but before any download, and normalize the MXQ `target_device` + there too so an unknown board reports as such rather than as a Hub failure. Lowercase an + explicit `framework` before validating it, matching `_model_paths.resolve_framework`, so + `framework="ONNX"` behaves as it does for every other model. Mask-generation-only CLI overrides + (`--encoder-*-path`, `--decoder-*-path`, `--prompt-weights-path`) must be rejected by every + command that builds a non-mask engine, since the generic engine never receives them and would + silently run the downloaded default instead. ONNX graph validation treats `-1` as + "this axis must be declared dynamic" (ONNX Runtime reports a dynamic axis as a `str`), not as a + wildcard: a decoder frozen at one token count would otherwise pass construction and fail inside + ONNX Runtime for two- and three-point prompts. Point prompts are validated as 1-3 points with finite + coordinates and labels of exactly `1`/`0` (checked before the integer cast, which would + otherwise truncate `0.5` into a valid label) before any backend call, and host prompt tensors + are built on the weights' device so `device="cuda"` works. `original_hw` is likewise validated + as exactly two positive whole numbers, and the source image as numeric and finite, before + either backend runs: interpolation preserves NaN, and a zero dimension yields infinite prompt + coordinates. Both FPN converters pin each level's complete `(C, H, W)` rather than only its + channel count, since `build_backbone_features` `view()`s them and a same-element-count geometry + would be silently rearranged into a corrupted FPN that still passes every downstream check. `classify_decoder_outputs` requires + exactly three mask candidates and rejects non-finite decoder outputs, so a NaN cannot reach + `argmax` over the IoU scores or the `> 0` mask threshold. Mask generation models reject `--model-path`/`--mxq-path`/`--onnx-path` + in every CLI command that builds one, never only in `predict`. +- Never add the PyPI `sam2` package as a dependency of `mblt_vision` (it is an unofficial + third-party mirror, not Meta's) and never require a manually cloned + `facebookresearch/sam2` checkout. Host-side prompt encoding is instead a from-scratch, + dependency-free port of the official prompt encoder/mask-decoder token setup in + `mask_generation/_sam2_prompt.py` and `_sam2_host.py`, backed by a small (~16KB) bundle of + weight tensors extracted from the official checkpoint and hosted at + `mobilint/sam2-hiera-large`'s Hub repo root (`sam2_hiera_large_prompt_weights.pt`) -- + downloaded the same way as the encoder/decoder MXQ artifacts, not shipped as package data. + Any change to that port must stay numerically verified against the real + `facebookresearch/sam2` predictor (`tests/test_mask_generation_prompt_encoding.py`, opt-in, + skips without a real `sam2` install). No added dependency: the input resize/normalize step + is plain `torch` (`F.interpolate` bilinear with `antialias=True` reproduces torchvision's + tensor `Resize` bit-for-bit), so mask_generation runs on the package's existing dependencies. +- `mask_generation` evaluates on the SA-V validation split (155 videos, 293 masklets, 31967 + annotated masks; JPEG frames + per-object binary PNG masks, binarized as `> 0`). The registry + entry `datasets/sa-v.yaml` does not download it: Meta distributes SA-V through a form-gated + portal, and it must never be mirrored on Mobilint infrastructure. Users supply the official + `sav_val.tar` (or its extracted directory) via `--annotation-dir`/`--image-dir`, exactly as + Cityscapes requires its manual archives; `_resolve_sav_source` fails with the portal link and + the SAM 2 `sav_dataset/README.md` layout reference rather than falling back to a URL. The + archive is deliberately absent from `PINNED_ARCHIVE_SHA256` because it is user-supplied rather + than fetched from a URL this package controls, so identity is enforced on content by the + readiness inventory instead. SA-V is CC BY 4.0 by Meta AI. The + organizer keeps only annotated frames. Readiness pins all three counts + (`SAV_VALIDATION_MASK_COUNT` alongside the video/masklet counts), since video and masklet totals + alone accept a source truncated to a few annotated frames per masklet and would silently + evaluate a different corpus. Both the organizer's staged-mask check and the readiness check + require every non-zero mask value to be one object ID: counting unique values alone accepts a + `{1, 2}` mask that `> 0` binarization turns entirely into foreground. Readiness checks geometry + and values for *every* cached mask, not the first per video. The official split is 1-bit + bilevel, a format that cannot encode a non-zero background, so those masks are validated from + the header and only non-bilevel masks are decoded -- complete validation at header cost + (~3s for 31967 masks; decoding them all would be ~380s on every `val`). +- SA-V video ids come from `sav_val.txt` file contents rather than a directory listing, so + `construct_sav` must reject ids failing `SAV_VIDEO_ID_PATTERN` and must confirm every staged + path resolves inside the staging tree, both *before* any `makedirs`/`copy`. COCO applies the + equivalent guard to its JSON-declared file names; organizers that build paths only from + directory listings do not need it. The evaluation protocol (`eval_sav`) is ported from the + validated `sam2-mxq-pipeline` reference: seed-deterministic area-balanced sampling, synthetic + point prompts from the GT mask (distance-transform peak, dilated-mask negatives), and + own-selection mean IoU as the primary metric with best-of-3 secondary. Numbers measured on + SA-V val are a different protocol from that reference's sav_train-sampled 0.7757 and are not + directly comparable. ## Python-First Architecture @@ -98,6 +199,14 @@ The current ownership boundary is deliberate: ## Benchmark and Compilation Tooling +- The unified benchmark runner's `TASK_CHOICES` is the subset of `VISION_TASKS` it can actually + execute. `mask_generation` is excluded because `_run_target` builds a generic `MBLT_Engine` and + `_evaluate` has no `eval_sav` branch; benchmark it with `mblt-vision val` instead. Do not wire a + new canonical task into the runner's choices before its engine and evaluator paths exist. +- `classify_decoder_outputs` pins the decoder mask layout to `(N, 65536)` or `(N, 256, 256)`, + optionally batched, before reshaping. A channels-last `(1, 256, 256, 3)` output has a matching + element count and candidate count, so only the layout check catches it; without it the reshape + interleaves the candidates into plausible but corrupted masks. - 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 @@ -126,7 +235,12 @@ The current ownership boundary is deliberate: 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. + Do not add model weights, caches, or compiled build artifacts to source control. The one + deliberate exception is `assets/`: a small fixed set of sample images kept in git as inputs for + manual QA and the documented CLI examples. They are development-only and must never reach a + distributed artifact -- only `mblt_vision*` packages are built, and `MANIFEST.in` prunes + `assets` to pin that intent. Do not grow this directory for new one-off inputs, and do not + reintroduce downloaded datasets, weights, or generated outputs under it. ## Compatibility Migration and Tests diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..5c3f2b6 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,6 @@ +# Sample images under assets/ are development-only inputs for manual QA and +# documented CLI examples. They are deliberately tracked in git but must never +# reach a distributed artifact. setuptools already excludes them (only +# `mblt_vision*` packages are built), so this prunes them explicitly to pin the +# intent against future packaging changes rather than to fix a current leak. +prune assets diff --git a/README.md b/README.md index 35b8590..e65f84d 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,8 @@ pip install "mblt-vision-python[onnxruntime-gpu]" ## Quick start -Each model includes its matching preprocess and postprocess behavior: +Most models include matching `preprocess`/`postprocess` behavior around a raw +call (mask generation is the exception -- see below): ```python from mblt_vision import ResNet50 @@ -70,6 +71,20 @@ code should use the task subpackages (for example, `obb` is the canonical oriented-bounding-box task name. +Mask generation is promptable, so `SAM2HieraLarge` takes point prompts and +returns candidate masks from a single `predict()` call instead of the separate +`preprocess`/raw-call/`postprocess` steps above: + +```python +from mblt_vision.mask_generation import SAM2HieraLarge + +model = SAM2HieraLarge() # framework="onnx" for ONNX Runtime inference +result = model.predict("image.jpg", points=[[320, 240]], labels=[1]) +``` + +See the [mask generation section](mblt_vision/README.md#mask-generation) for the +prompt contract, the two-artifact layout, and SA-V validation. + ## Model Zoo migration Vision is now maintained in this package. `mblt-model-zoo` retains @@ -87,17 +102,24 @@ 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. +object and face detection, instance and semantic segmentation, OBB, pose +estimation, and point-prompted mask generation. 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. +Mask generation models require 1-3 point prompts (`--point X,Y,LABEL`, where +LABEL is 1 for positive and 0 for negative) and load two artifacts, overridden +with `--encoder-mxq-path`/`--decoder-mxq-path` (or the `--encoder-onnx-path`/ +`--decoder-onnx-path` pair) rather than `--model-path`/`--mxq-path`/`--onnx-path`. + ```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 +mblt-vision predict --source image.jpg --model sam2-hiera-large --point 320,240,1 ``` The corresponding `mblt-model-zoo` commands use the same standalone handlers for diff --git a/assets/airport.jpg b/assets/airport.jpg new file mode 100644 index 0000000..7aa310e Binary files /dev/null and b/assets/airport.jpg differ diff --git a/assets/bus.jpg b/assets/bus.jpg new file mode 100644 index 0000000..40eaaf5 Binary files /dev/null and b/assets/bus.jpg differ diff --git a/assets/cr7.jpg b/assets/cr7.jpg new file mode 100644 index 0000000..8b590b6 Binary files /dev/null and b/assets/cr7.jpg differ diff --git a/assets/frankfurt.png b/assets/frankfurt.png new file mode 100644 index 0000000..980bff1 Binary files /dev/null and b/assets/frankfurt.png differ diff --git a/assets/street.jpg b/assets/street.jpg new file mode 100644 index 0000000..bc9886e Binary files /dev/null and b/assets/street.jpg differ diff --git a/assets/volcano.jpg b/assets/volcano.jpg new file mode 100644 index 0000000..429eb58 Binary files /dev/null and b/assets/volcano.jpg differ diff --git a/benchmark/README.md b/benchmark/README.md index d3dc3dd..d8e2b75 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -10,8 +10,8 @@ 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. +Public datasets can use their default download sources. ImageNet, DOTA, +Cityscapes, and SA-V require the appropriate source archives or credentials. ```bash python benchmark/organize_coco.py @@ -27,6 +27,44 @@ python benchmark/organize_cityscapes.py \ --annotation-dir path/to/gtFine_trainvaltest.zip ``` +### SA-V (mask generation) + +Meta distributes SA-V through a +[form-gated download portal](https://ai.meta.com/datasets/segment-anything-video-downloads/) +and this package does not mirror it, so `--dataset-path` is required. Download +`sav_val.tar` (about 15 GB), then: + +```bash +python benchmark/organize_sav.py --dataset-path path/to/sav_val.tar +``` + +An already-extracted `sav_val` directory is accepted in place of the archive. +The expected source layout is documented in the +[SAM 2 `sav_dataset` README](https://github.com/facebookresearch/sam2/blob/main/sav_dataset/README.md): + +```text +sav_val +|-- sav_val.txt # 155 video ids +|-- JPEGImages_24fps/{video_id}/{frame:05d}.jpg # frames at 24fps +`-- Annotations_6fps/{video_id}/{object_id:03d}/{frame:05d}.png # masks at 6fps +``` + +The organizer installs only the annotated frames, since annotations exist at +6fps while frames are extracted at 24fps and evaluation can only use annotated +frames. The result is validated before it replaces any existing cache: 155 +videos, 293 masklets, and 31967 annotated masks, with every mask checked for +matching geometry and single-object values. + +Organizing needs roughly 20 GB of free space for the unpacked archive plus the +installed copy; the archive itself can be deleted afterwards. Raw dataset +archives are gitignored — never commit them. + +Once organized, validation reuses the cache and needs no further flags: + +```bash +mblt-vision val --model sam2-hiera-large +``` + ## Run a benchmark Use the unified runner for every Vision task. It chooses the evaluator from diff --git a/benchmark/benchmark_vision_models.py b/benchmark/benchmark_vision_models.py index 24fff74..b8ed7f6 100644 --- a/benchmark/benchmark_vision_models.py +++ b/benchmark/benchmark_vision_models.py @@ -35,7 +35,16 @@ CORE_MODES: tuple[CoreMode, ...] = cast( tuple[CoreMode, ...], core_modes_for_target_device("aries-rb") ) -TASK_CHOICES = VISION_TASKS +# `mask_generation` is a canonical Vision task but the unified runner cannot +# execute it: it needs SAM2HieraLarge's two-artifact engine and point prompts +# rather than the generic MBLT_Engine `_run_target` builds, and `eval_sav` +# rather than the generic evaluators `_evaluate` dispatches. Offering it as a +# choice would only produce an error row for every model. Benchmark it with +# `mblt-vision val --model sam2-hiera-large` until the runner grows that path. +UNSUPPORTED_BENCHMARK_TASKS: tuple[str, ...] = ("mask_generation",) +TASK_CHOICES = tuple( + task for task in VISION_TASKS if task not in UNSUPPORTED_BENCHMARK_TASKS +) SUPPORTED_TARGET_DEVICES = frozenset({"aries-rb", "regulus-ra", "regulus-rb"}) @@ -69,7 +78,9 @@ def _parse_task(value: str) -> str: """Normalize a benchmark task for argparse.""" try: - return normalize_vision_task(value) + # Validated against the runner's supported tasks, not every canonical + # Vision task, so an unsupported one reports what can actually be run. + return normalize_vision_task(value, supported=TASK_CHOICES) except (TypeError, ValueError) as exc: raise argparse.ArgumentTypeError(str(exc)) from exc diff --git a/benchmark/organize_sav.py b/benchmark/organize_sav.py new file mode 100644 index 0000000..666c6a2 --- /dev/null +++ b/benchmark/organize_sav.py @@ -0,0 +1,72 @@ +"""Organize the SA-V validation dataset for local use. + +SA-V is distributed by Meta through a form-gated download portal and is not +mirrored by this package, so `--dataset-path` is required: download +`sav_val.tar` yourself and pass it here, exactly as Cityscapes requires its +official archives. Both the archive and an already-extracted `sav_val` +directory are accepted. + +The official source layout is documented in the SAM 2 repository: +https://github.com/facebookresearch/sam2/blob/main/sav_dataset/README.md + + sav_val + |-- sav_val.txt # video ids in the split + |-- JPEGImages_24fps/{video_id}/{frame:05d}.jpg + `-- Annotations_6fps/{video_id}/{object_id:03d}/{frame:05d}.png + +Only the annotated frames are installed: annotations exist at 6fps while +frames are extracted at 24fps, and evaluation can only use annotated frames. +""" + +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.datasets import get_dataset_config +from mblt_vision.utils.datasets import organize_sav + +SAV_DOWNLOAD_CONFIG = get_dataset_config("sa-v")["download"] + + +def main() -> None: + """Parse organizer options and materialize SA-V validation data.""" + + parser = argparse.ArgumentParser( + description=( + "Organize SA-V validation data from the official sav_val.tar archive. " + f"Download it at {SAV_DOWNLOAD_CONFIG['source']} " + f"(layout: {SAV_DOWNLOAD_CONFIG['documentation']})." + ) + ) + parser.add_argument( + "--dataset-path", + required=True, + help=( + f"Path to the manually downloaded {SAV_DOWNLOAD_CONFIG['archive']} " + "or its extracted sav_val directory" + ), + ) + parser.add_argument( + "--output-dir", + default=None, + help=( + "Destination for the organized images/, annotations/, and " + "video_ids.txt (defaults to cache)" + ), + ) + args = parser.parse_args() + organize_sav( + dataset_path=args.dataset_path, + output_dir=args.output_dir, + ) + + +if __name__ == "__main__": + main() diff --git a/mblt_vision/README.md b/mblt_vision/README.md index 4eccff5..97a998f 100644 --- a/mblt_vision/README.md +++ b/mblt_vision/README.md @@ -2,9 +2,12 @@ `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. +boxes (OBB), instance and semantic segmentation, pose estimation, and promptable mask +generation (SAM2). Each model configuration includes the artifact, preprocessing, output +taxonomy, and postprocessing contract needed to produce task-specific results. Mask +generation is the one exception: `SAM2HieraLarge` loads two NPU artifacts and takes point +prompts rather than an image alone, so it has its own constructor -- see +[Mask Generation](#mask-generation). ## Loading models @@ -373,3 +376,64 @@ validation split, following Ultralytics' depth-validation convention. rotated mAP50 as the secondary metric. + +### Mask Generation + +Promptable segmentation: given an image and 1-3 point prompts (positive/negative), returns 3 +candidate masks with IoU scores. Point prompts only for now -- no box prompts and no automatic +"segment everything" mode. Unlike every other model here, `SAM2HieraLarge` downloads three +artifacts from the Hub instead of one: an image encoder, a prompt-conditioned mask decoder +(MXQ by default, or the ONNX exports with `framework="onnx"`), and a small (~16KB) bundle of +host-side prompt-encoder weights. Host-side prompt encoding +is a from-scratch, dependency-free port of the official +[`facebookresearch/sam2`](https://github.com/facebookresearch/sam2) prompt encoder and mask-decoder +token setup (numerically verified bit-for-bit against the real implementation) -- no `sam2` +package, no `torchvision`, no manually cloned repository, and no extra to install: everything +runs on the package's existing `torch` dependency. See `mblt_vision/mask_generation/`. + +```python +from mblt_vision.mask_generation import SAM2HieraLarge + +model = SAM2HieraLarge() # framework="onnx" for ONNX Runtime inference +result = model.predict("image.jpg", points=[[320, 240]], labels=[1]) +result.plot("image.jpg", save_path="result.jpg") +``` + +```bash +mblt-vision predict --source image.jpg --model sam2-hiera-large --point 320,240,1 +mblt-vision predict --source image.jpg --model sam2-hiera-large --point 320,240,1 --framework onnx +``` + +Validate with the built-in SA-V evaluation. Unlike the auto-downloading datasets, +SA-V must be obtained manually: Meta distributes it through a +[gated download form](https://ai.meta.com/datasets/segment-anything-video-downloads/), +and this package does not mirror it. Download `sav_val.tar`, then pass it (or its +extracted `sav_val` directory) on the first run; it is organized into the dataset +cache and reused afterwards. The official layout is documented in the +[SAM 2 `sav_dataset` README](https://github.com/facebookresearch/sam2/blob/main/sav_dataset/README.md). + +```bash +mblt-vision val --model sam2-hiera-large --annotation-dir /path/to/sav_val.tar +mblt-vision val --model sam2-hiera-large --framework onnx +``` + +Local artifact overrides follow the framework: `--encoder-mxq-path`/`--decoder-mxq-path` for +MXQ and `--encoder-onnx-path`/`--decoder-onnx-path` for ONNX. The shared NPU options apply to +both MXQ backends and are ignored for ONNX inference, as everywhere else in the CLI. + +| Model | Input Size
(H,W,C) | Source | Note | +| --- | --- | --- | --- | +| SAM2HieraLarge | (1024,1024,3) | [Link](https://ai.meta.com/sam2) | Point prompts only | + +
+Mask Generation (SA-V val) + +- Validation reports the own-selection mean IoU (`argmax` of the 3 predicted IoU scores) of the + predicted mask against ground truth as the primary metric, with the best-of-3 oracle IoU as + the secondary metric: 200 deterministic single-point-prompt samples by default (seed 0, + area-balanced) from the official + [SA-V](https://ai.meta.com/datasets/segment-anything-video) validation split + (155 videos, 293 masklets). Prompts are synthesized from the ground-truth mask + (distance-transform peak; `--num-points 2/3` adds negative and second positive points). + +
diff --git a/mblt_vision/__init__.py b/mblt_vision/__init__.py index 1b042e0..6d221b7 100644 --- a/mblt_vision/__init__.py +++ b/mblt_vision/__init__.py @@ -11,6 +11,7 @@ from . import face_detection as face_detection from . import image_classification as image_classification from . import instance_segmentation as instance_segmentation +from . import mask_generation as mask_generation from . import obb as obb from . import object_detection as object_detection from . import pose_estimation as pose_estimation @@ -19,13 +20,14 @@ from ._api import list_tasks as list_tasks from .wrapper import MBLT_Engine as MBLT_Engine -__version__ = "0.0.2" +__version__ = "0.0.3" _TASK_MODULES = ( face_detection, depth_estimation, image_classification, instance_segmentation, + mask_generation, object_detection, obb, pose_estimation, @@ -49,6 +51,7 @@ "depth_estimation", "image_classification", "instance_segmentation", + "mask_generation", "object_detection", "obb", "pose_estimation", diff --git a/mblt_vision/_tasks.py b/mblt_vision/_tasks.py index df9a462..64852e4 100644 --- a/mblt_vision/_tasks.py +++ b/mblt_vision/_tasks.py @@ -13,6 +13,7 @@ "obb", "pose_estimation", "face_detection", + "mask_generation", ) diff --git a/mblt_vision/cli/_vision.py b/mblt_vision/cli/_vision.py index 096a2ef..0f4a03a 100644 --- a/mblt_vision/cli/_vision.py +++ b/mblt_vision/cli/_vision.py @@ -3,6 +3,7 @@ from __future__ import annotations import argparse +import math import os import sys from pathlib import Path @@ -53,6 +54,39 @@ def parse_target_clusters(value: str | None) -> list[int] | None: return clusters or None +def parse_point(value: str) -> tuple[float, float, int]: + """Parses an `X,Y,LABEL` point prompt for mask generation models.""" + + parts = value.split(",") + if len(parts) != 3: + raise argparse.ArgumentTypeError("expected a point as X,Y,LABEL") + try: + x, y = float(parts[0]), float(parts[1]) + label = int(parts[2]) + except ValueError as exc: + raise argparse.ArgumentTypeError( + "expected numeric X,Y coordinates and an integer LABEL" + ) from exc + # `float()` accepts "nan"/"inf", which would otherwise reach Fourier prompt + # encoding and contaminate the resulting tokens, masks, and IoU scores. + if not (math.isfinite(x) and math.isfinite(y)): + raise argparse.ArgumentTypeError("expected finite X,Y coordinates") + if label not in (0, 1): + raise argparse.ArgumentTypeError( + "point LABEL must be 1 (positive) or 0 (negative)" + ) + return (x, y, label) + + +def resolve_cli_task(args: argparse.Namespace) -> str: + """Resolves the selected model's task without constructing a runtime.""" + + from mblt_vision.wrapper import resolve_model_config + + config = resolve_model_config(args.model, args.model_type) + return normalize_vision_task(config["post_cfg"]["task"]) + + def add_common_vision_args(parser: argparse.ArgumentParser) -> None: """Adds arguments shared by all vision inference commands.""" @@ -275,6 +309,141 @@ def create_vision_engine(args: argparse.Namespace) -> Any: ) +def reject_single_artifact_paths(args: argparse.Namespace) -> None: + """Rejects single-artifact model paths for two-artifact mask generation models. + + Raises: + SystemExit: If `--model-path`, `--mxq-path`, or `--onnx-path` is set. + """ + + if ( + getattr(args, "model_path", "") + or getattr(args, "mxq_path", "") + or getattr(args, "onnx_path", "") + ): + raise SystemExit( + "Mask generation models load two artifacts; use " + "`--encoder-mxq-path`/`--decoder-mxq-path` (or " + "`--encoder-onnx-path`/`--decoder-onnx-path` with `--framework onnx`) " + "instead of `--model-path`/`--mxq-path`/`--onnx-path`." + ) + + +MASK_GENERATION_ONLY_OPTIONS: tuple[tuple[str, str], ...] = ( + ("encoder_mxq_path", "--encoder-mxq-path"), + ("decoder_mxq_path", "--decoder-mxq-path"), + ("encoder_onnx_path", "--encoder-onnx-path"), + ("decoder_onnx_path", "--decoder-onnx-path"), + ("prompt_weights_path", "--prompt-weights-path"), +) + + +def reject_mask_generation_only_options(args: argparse.Namespace) -> None: + """Rejects two-artifact overrides for models that do not accept them. + + The generic engine never receives these paths, so accepting them silently + downloads and runs the default single artifact instead of the requested + local one. + + Raises: + SystemExit: If any mask-generation-only artifact override is set. + """ + + supplied = [ + flag + for attribute, flag in MASK_GENERATION_ONLY_OPTIONS + if getattr(args, attribute, "") + ] + if supplied: + raise SystemExit( + f"{', '.join(supplied)} is only supported for mask generation models " + "such as SAM2HieraLarge. Use `--model-path`, `--mxq-path`, or " + "`--onnx-path` for this model." + ) + + +def create_mask_generation_engine(args: argparse.Namespace) -> Any: + """Creates a promptable mask generation engine from shared CLI model options. + + Mask generation models load two artifacts (encoder + decoder): MXQ by + default, or ONNX with `--framework onnx`. The shared NPU options + (`--dev-no`, `--core-mode`, `--target-cores`, `--target-clusters`) apply + to both MXQ backends and are ignored for ONNX inference. + + The single-artifact path options are rejected here rather than in one + command's handler, so every command that builds a mask generation engine + fails loudly instead of silently evaluating the downloaded default in + place of an explicitly requested local artifact. + """ + + reject_single_artifact_paths(args) + + try: + import mblt_vision.mask_generation as mask_generation_module + from mblt_vision.wrapper import _model_name_aliasing + except ImportError as exc: + print(f"Missing dependencies for vision CLI: {exc}", file=sys.stderr) + raise SystemExit(2) from exc + + class_name = Path(_model_name_aliasing(args.model)).stem + model_class = getattr(mask_generation_module, class_name, None) + if model_class is None: + raise SystemExit( + f"Mask generation model class '{class_name}' is not exported by " + "mblt_vision.mask_generation." + ) + return model_class( + encoder_mxq_path=getattr(args, "encoder_mxq_path", None) or None, + decoder_mxq_path=getattr(args, "decoder_mxq_path", None) or None, + prompt_weights_path=getattr(args, "prompt_weights_path", None) or None, + encoder_dev_no=args.dev_no, + decoder_dev_no=args.dev_no, + encoder_core_mode=args.core_mode, + decoder_core_mode=args.core_mode, + encoder_target_cores=args.target_cores, + decoder_target_cores=args.target_cores, + encoder_target_clusters=args.target_clusters, + decoder_target_clusters=args.target_clusters, + target_device=args.target_device, + framework=args.framework, + encoder_onnx_path=getattr(args, "encoder_onnx_path", None) or None, + decoder_onnx_path=getattr(args, "decoder_onnx_path", None) or None, + ) + + +def run_mask_generation_inference( + args: argparse.Namespace, + *, + command: str, +) -> Any: + """Runs point-prompted mask generation for a CLI command.""" + + point_prompts = getattr(args, "points", None) or [] + if not 1 <= len(point_prompts) <= 3: + raise SystemExit( + "Mask generation requires 1 to 3 point prompts; pass `--point X,Y,LABEL` " + "(LABEL 1 positive, 0 negative) up to three times." + ) + reject_single_artifact_paths(args) + + points = [[x, y] for x, y, _ in point_prompts] + labels = [label for _, _, label in point_prompts] + model = create_mask_generation_engine(args) + try: + result = model.predict(args.source, points, labels) + save_path = resolve_output_path(args.output, command, args.source, args.model) + result.plot(source_path=args.source, save_path=save_path) + iou_text = ", ".join(f"{float(value):.4f}" for value in result.iou_predictions) + print( + f"Predicted IoU per mask candidate: [{iou_text}]; " + f"selected mask index: {result.selected}" + ) + print(f"Saved result to {os.path.relpath(save_path)}") + return result + finally: + model.dispose() + + def run_vision_inference( args: argparse.Namespace, *, @@ -283,6 +452,14 @@ def run_vision_inference( """Runs a complete vision inference pipeline for a CLI command.""" require_source_file(args.source) + if resolve_cli_task(args) == "mask_generation": + return run_mask_generation_inference(args, command=command) + if getattr(args, "points", None): + raise SystemExit( + "`--point` is only supported for mask generation models such as " + "SAM2HieraLarge." + ) + reject_mask_generation_only_options(args) model = create_vision_engine(args) try: actual_task = normalize_vision_task(model.post_cfg.get("task", "")) diff --git a/mblt_vision/cli/predict.py b/mblt_vision/cli/predict.py index 739900d..ba37af1 100644 --- a/mblt_vision/cli/predict.py +++ b/mblt_vision/cli/predict.py @@ -8,6 +8,7 @@ add_e2e_arg, add_threshold_args, add_vision_parser, + parse_point, run_vision_inference, ) @@ -29,7 +30,7 @@ def add_predict_parser( command="predict", help_text=( "Run vision inference for classification, depth estimation, detection, instance or semantic " - "segmentation, OBB, pose, and face detection." + "segmentation, OBB, pose, face detection, and point-prompted mask generation." ), description=( "Run a configured Vision model on one image. The selected model determines the task, " @@ -37,18 +38,26 @@ def add_predict_parser( ), epilog="""Supported tasks: image classification, depth estimation, object and face detection, instance - and semantic segmentation, oriented bounding boxes (OBB), and pose estimation. + and semantic segmentation, oriented bounding boxes (OBB), pose estimation, + and point-prompted mask generation (SAM2). 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. +Mask generation models take 1-3 point prompts (--point X,Y,LABEL; LABEL 1 positive, +0 negative) and load two artifacts: MXQ by default (overridable with +--encoder-mxq-path/--decoder-mxq-path) or ONNX with --framework onnx (overridable +with --encoder-onnx-path/--decoder-onnx-path) instead of --model-path/--mxq-path. + 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 + mblt-vision predict --source image.jpg --model sam2-hiera-large --point 320,240,1 + mblt-vision predict --source image.jpg --model sam2-hiera-large --point 320,240,1 --framework onnx For export-style YOLO output, use --e2e false and optionally save it with --raw-output.""", handler=_cmd_predict, @@ -60,5 +69,46 @@ def add_predict_parser( "--raw-output", help="Path to save raw export-style output with `--e2e false`.", ) + parser.add_argument( + "--point", + action="append", + dest="points", + type=parse_point, + metavar="X,Y,LABEL", + help=( + "Point prompt for mask generation models: pixel X,Y in the source image " + "and LABEL (1 positive, 0 negative). Repeat up to 3 times." + ), + ) + parser.add_argument( + "--encoder-mxq-path", + dest="encoder_mxq_path", + default="", + help="Optional local encoder MXQ path for mask generation models.", + ) + parser.add_argument( + "--decoder-mxq-path", + dest="decoder_mxq_path", + default="", + help="Optional local decoder MXQ path for mask generation models.", + ) + parser.add_argument( + "--encoder-onnx-path", + dest="encoder_onnx_path", + default="", + help="Optional local encoder ONNX path for mask generation models.", + ) + parser.add_argument( + "--decoder-onnx-path", + dest="decoder_onnx_path", + default="", + help="Optional local decoder ONNX path for mask generation models.", + ) + parser.add_argument( + "--prompt-weights-path", + dest="prompt_weights_path", + default="", + help="Optional local prompt-encoder weights path for mask generation models.", + ) 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 index cdab6ee..a6ecc1d 100644 --- a/mblt_vision/cli/val.py +++ b/mblt_vision/cli/val.py @@ -16,9 +16,12 @@ from ._vision import ( add_e2e_arg, add_threshold_args, + create_mask_generation_engine, create_vision_engine, parse_target_clusters, parse_target_cores, + reject_mask_generation_only_options, + resolve_cli_task, ) DEFAULT_IMAGENET_IMAGE_SOURCE = get_dataset_config("imagenet")["download"]["images"] @@ -32,6 +35,7 @@ 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"] +SAV_DOWNLOAD_CONFIG = get_dataset_config("sa-v")["download"] CITYSCAPES_DOWNLOAD_CONFIG = get_dataset_config("cityscapes")["download"] CITYSCAPES_IMAGE_ARCHIVE = CITYSCAPES_DOWNLOAD_CONFIG["images_archive"] CITYSCAPES_ANNOTATION_ARCHIVE = CITYSCAPES_DOWNLOAD_CONFIG["annotations_archive"] @@ -54,12 +58,21 @@ def _candidate_search_roots(data_path: str) -> list[Path]: def _find_existing_source(data_path: str, candidate_names: list[str]) -> str | None: - """Finds a nearby raw archive or extracted dataset directory.""" + """Finds a nearby raw archive or extracted dataset directory. + + Never returns the organized output directory itself. Several datasets use a + cache directory whose name also appears in their source-candidate list (for + example `sa-v` and `nyu-depth`), so `data_path.parent / name` can resolve + back to `data_path`. Handing that to an organizer as a raw source makes + source resolution fail on the incomplete cache instead of downloading the + default archive and repairing it. + """ + organized_root = Path(data_path).expanduser().resolve() for root in _candidate_search_roots(data_path): for name in candidate_names: candidate = root / name - if candidate.exists(): + if candidate.exists() and candidate.resolve() != organized_root: return str(candidate) return None @@ -171,6 +184,49 @@ def _resolve_nyu_depth_source(args: argparse.Namespace, data_path: str) -> str: return dataset_path or DEFAULT_NYU_DEPTH_SOURCE +def _resolve_sav_source(args: argparse.Namespace, data_path: str) -> str: + """Resolve the manually downloaded SA-V validation archive or directory. + + SA-V is distributed through Meta's form-gated portal and is not mirrored by + this package, so unlike the auto-downloading datasets there is no default + source to fall back to. + + Args: + args: Parsed validation CLI arguments. + data_path: Organized SA-V output path used as a discovery anchor. + + Returns: + Path to the ``sav_val.tar`` archive or its extracted directory. + + Raises: + SystemExit: If the archive or extracted directory cannot be found. + """ + + # Discovery runs even under --force-organize, unlike the auto-downloading + # datasets. There the flag skips discovery so a rebuild re-fetches from the + # canonical URL; SA-V has no such fallback, so skipping it would fail on a + # manual archive sitting in the documented discovery location -- the very + # place the error below tells the user to put it. `_find_existing_source` + # already refuses to return the organized output, so this cannot resolve to + # the stale cache that --force-organize exists to rebuild. + dataset_path = ( + args.annotation_dir + or args.image_dir + or _find_existing_source(data_path, [SAV_DOWNLOAD_CONFIG["archive"], "sav_val"]) + ) + if not dataset_path: + raise SystemExit( + "SA-V organization requires the official validation archive, which Meta " + "distributes through a gated download form and this package does not mirror.\n" + f" Download it at {SAV_DOWNLOAD_CONFIG['source']}\n" + f" Layout reference: {SAV_DOWNLOAD_CONFIG['documentation']}\n" + f"Pass the resulting {SAV_DOWNLOAD_CONFIG['archive']} (or its extracted " + "`sav_val` directory) with --annotation-dir or --image-dir, or place it " + "near the dataset path." + ) + return dataset_path + + def _resolve_ade20k_source(args: argparse.Namespace, data_path: str) -> str: """Resolve a local archive, extracted directory, or URL for ADE20K organization.""" @@ -259,6 +315,7 @@ def _ensure_dataset( organize_dotav1, organize_imagenet, organize_nyu_depth, + organize_sav, organize_widerface, ) except ImportError as exc: @@ -300,6 +357,11 @@ def _ensure_dataset( dataset_path=_resolve_nyu_depth_source(args, data_path), output_dir=data_path, ) + elif task == "mask_generation": + organize_sav( + dataset_path=_resolve_sav_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) @@ -335,13 +397,18 @@ def _run_validation(args: argparse.Namespace) -> float: eval_dota, eval_imagenet_metrics, eval_nyu_depth, + eval_sav, 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) + if resolve_cli_task(args) == "mask_generation": + model = create_mask_generation_engine(args) + else: + reject_mask_generation_only_options(args) + model = create_vision_engine(args) try: if not getattr(getattr(model, "postprocessor", None), "e2e", True): raise SystemExit( @@ -381,6 +448,24 @@ def _run_validation(args: argparse.Namespace) -> float: ) return depth_result.primary_score + if task == "mask_generation": + sav_result = eval_sav( + model=model, + data_path=data_path, + num_samples=args.num_samples, + num_points=args.num_points, + seed=args.seed, + ) + print( + "Validation score " + f"(mIoU): {sav_result.miou:.5f} " + f"(+-95%CI {sav_result.miou_ci95:.5f}), " + f"(mIoU best-of-3): {sav_result.miou_best_of_3:.5f}, " + f"samples: {sav_result.num_samples}, " + f"videos: {sav_result.distinct_videos}" + ) + return sav_result.primary_score + if task == "semantic_segmentation": if taxonomy == "cityscapes": semantic_result = eval_cityscapes( @@ -557,5 +642,54 @@ def add_val_parser( "Local archive path or download URL for dataset annotations. Cityscapes requires gtFine_trainvaltest.zip." ), ) + parser.add_argument( + "--num-samples", + type=parse_positive_int, + default=200, + help="Mask generation only: number of prompted SA-V samples to evaluate.", + ) + parser.add_argument( + "--num-points", + type=int, + default=1, + choices=[1, 2, 3], + help="Mask generation only: points per synthetic prompt.", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="Mask generation only: sampling and prompt-synthesis seed.", + ) + parser.add_argument( + "--encoder-mxq-path", + dest="encoder_mxq_path", + default="", + help="Optional local encoder MXQ path for mask generation models.", + ) + parser.add_argument( + "--decoder-mxq-path", + dest="decoder_mxq_path", + default="", + help="Optional local decoder MXQ path for mask generation models.", + ) + parser.add_argument( + "--encoder-onnx-path", + dest="encoder_onnx_path", + default="", + help="Optional local encoder ONNX path for mask generation models.", + ) + parser.add_argument( + "--decoder-onnx-path", + dest="decoder_onnx_path", + default="", + help="Optional local decoder ONNX path for mask generation models.", + ) + parser.add_argument( + "--prompt-weights-path", + dest="prompt_weights_path", + default="", + help="Optional local prompt-encoder weights path for mask generation models.", + ) add_threshold_args(parser, conf_default=None, iou_default=None) add_e2e_arg(parser) diff --git a/mblt_vision/datasets/sa-v.yaml b/mblt_vision/datasets/sa-v.yaml new file mode 100644 index 0000000..2e39e7f --- /dev/null +++ b/mblt_vision/datasets/sa-v.yaml @@ -0,0 +1,27 @@ +# SA-V (Segment Anything Video) validation split, used for point-prompted +# mask generation evaluation. 155 videos / 293 manually annotated masklets, +# distributed as extracted JPEG frames (24fps) and per-object binary PNG +# masks (6fps). Organized paths are relative to `path`. +# +# SA-V is distributed by Meta behind a form-gated download portal and is not +# mirrored on Mobilint infrastructure, so it cannot be fetched automatically: +# download `sav_val.tar` yourself and pass it (or its extracted `sav_val/` +# directory) to the organizer. The official layout this package expects is +# documented in the SAM 2 repository, linked below. SA-V is licensed +# CC BY 4.0 by Meta AI. +# +# The archive is not sha256-pinned because it is user-supplied rather than +# fetched from a URL this package controls. Dataset identity is enforced on +# content instead: readiness requires 155 videos, 293 masklets and 31967 +# annotated masks, and validates every mask's geometry and single-object +# values. +name: sa-v +path: ~/.mblt_model_zoo/datasets/sa-v +val: images +tasks: + - mask_generation +download: + type: manual + source: https://ai.meta.com/datasets/segment-anything-video-downloads/ + documentation: https://github.com/facebookresearch/sam2/blob/main/sav_dataset/README.md + archive: sav_val.tar diff --git a/mblt_vision/mask_generation/__init__.py b/mblt_vision/mask_generation/__init__.py new file mode 100644 index 0000000..6e2a88e --- /dev/null +++ b/mblt_vision/mask_generation/__init__.py @@ -0,0 +1,7 @@ +"""Promptable mask generation models (SAM2).""" + +from __future__ import annotations + +from .sam2 import SAM2HieraLarge + +__all__ = ["SAM2HieraLarge"] diff --git a/mblt_vision/mask_generation/_sam2_contracts.py b/mblt_vision/mask_generation/_sam2_contracts.py new file mode 100644 index 0000000..bd4ec15 --- /dev/null +++ b/mblt_vision/mask_generation/_sam2_contracts.py @@ -0,0 +1,243 @@ +"""Fixed contracts between the SAM2 encoder/decoder artifacts and their host glue. + +MXQ side: ported from the validated ``sam2-mxq-pipeline`` reference (real Aries2 +SA-V-200 accuracy: FP32 mIoU 0.7750 vs MXQ mIoU 0.7757, mask agreement 0.983). +Compile and calibration are out of scope for this phase, so the decoder's +compiled runtime input order is a single validated default rather than a +configurable MBLT-input-name binding map. + +ONNX side: the graph interface written by the SDK tutorial's +``sam2_export_onnx.py`` (``Sam2ImageEncoderWrapper``/``Sam2MaskDecoderWrapper`` +traces, verified numerically against the official ``facebookresearch/sam2`` +predictor). Unlike the compiled MXQ artifacts, the ONNX graphs are NCHW, keep +the pre-flattening decoder tensor shapes, and take five named decoder inputs -- +``src_plus_pos_src`` stays inside the graph instead of being a sixth input. +""" + +from __future__ import annotations + +from typing import Any, Mapping, Sequence + +import numpy as np + +# Decoder runtime input order as fed to ``MobilintNPUBackend`` / ``qbruntime.Model.infer``. +# This is the positional signature of the compiled SAM2 Hiera-Large decoder MXQ, +# which differs from the MBLT graph's declared input names. +DECODER_RUNTIME_ORDER: tuple[str, ...] = ( + "hrf0_nhwc", + "src_plus_pos_src", + "hrf1_nhwc", + "src", + "pos_src", + "tokens", +) + +# Exported ONNX graph interface. ``-1`` marks the prompt-count-dependent +# dynamic token axis (``6 output tokens + N points + 1 pad``). +ENCODER_ONNX_INPUT_NAME = "input_image" +ENCODER_ONNX_INPUT_SHAPE: tuple[int, ...] = (1, 3, 1024, 1024) +DECODER_ONNX_INPUT_SHAPES: dict[str, tuple[int, ...]] = { + "tokens": (1, -1, 256), + "src": (1, 256, 64, 64), + "pos_src": (1, 256, 64, 64), + "high_res_features_0": (1, 32, 256, 256), + "high_res_features_1": (1, 64, 128, 128), +} +DECODER_ONNX_INPUT_NAMES: tuple[str, ...] = tuple(DECODER_ONNX_INPUT_SHAPES) + +MASK_SIDE = 256 +MASK_AREA = MASK_SIDE * MASK_SIDE +# SAM2's multimask output (4 mask tokens minus the single-mask token), baked +# into both compiled artifacts. `eval_sav.CANDIDATES_PER_PROMPT` is the same +# fixed contract seen from the evaluation side. +MASK_CANDIDATE_COUNT = 3 + + +def strip_runtime_batch(value: np.ndarray) -> np.ndarray: + """Remove the outer model batch that qbruntime omits from buffer shapes.""" + + array = np.asarray(value) + if array.ndim >= 4 and array.shape[0] == 1: + array = array[0] + return np.ascontiguousarray(array, dtype=np.float32) + + +def build_decoder_runtime_feed( + tensors: Mapping[str, np.ndarray], order: Sequence[str] = DECODER_RUNTIME_ORDER +) -> list[np.ndarray]: + """Order the named decoder tensors into the compiled artifact's positional feed.""" + + missing = [role for role in order if role not in tensors] + if missing: + raise ValueError(f"Decoder tensors are missing role(s): {missing}.") + return [strip_runtime_batch(tensors[role]) for role in order] + + +def validate_runtime_shapes( + actual: Sequence[np.ndarray], expected: Sequence[Sequence[int]], label: str +) -> None: + """Fail loudly at construction time if a resolved artifact's shapes drift. + + ``-1`` in ``expected`` marks a wildcard/dynamic dimension (for example the + decoder's point-count-dependent token axis). + """ + + if len(actual) != len(expected): + raise ValueError( + f"{label} input count mismatch: feed={len(actual)}, runtime={len(expected)}." + ) + for index, (array, shape) in enumerate(zip(actual, expected)): + shape = tuple(int(dim) for dim in shape) + got = tuple(int(dim) for dim in array.shape) + if len(got) != len(shape) or any( + want != -1 and have != want for have, want in zip(got, shape) + ): + raise ValueError( + f"{label} input {index} shape mismatch: feed={got}, runtime={shape}." + ) + + +def normalize_onnx_dims(shape: Sequence[Any]) -> tuple[int, ...]: + """Map ONNX Runtime dims to ints, with symbolic/dynamic dims as ``-1``.""" + + return tuple(int(dim) if isinstance(dim, int) else -1 for dim in shape) + + +def validate_onnx_session_inputs( + session_inputs: Sequence[Any], expected: Mapping[str, Sequence[int]], label: str +) -> None: + """Fail loudly at construction time if a resolved ONNX artifact's graph drifts. + + ``session_inputs`` is ONNX Runtime input metadata (``session.get_inputs()``). + An expected static dimension must be declared statically by the graph, and + ``-1`` marks an axis the graph must declare *dynamic* -- not a wildcard that + accepts anything. ONNX Runtime reports a dynamic axis as its symbolic name + (a ``str``), so a decoder exported with a frozen token dimension such as + ``(1, 8, 256)`` is rejected here rather than working for one-point prompts + and failing inside ONNX Runtime for the advertised two- and three-point + prompts. + """ + + raw_shapes = {item.name: tuple(item.shape) for item in session_inputs} + if set(raw_shapes) != set(expected): + raise ValueError( + f"{label} ONNX input names mismatch: graph={sorted(raw_shapes)}, " + f"expected={sorted(expected)}." + ) + for name, shape in expected.items(): + want = tuple(int(dim) for dim in shape) + raw = raw_shapes[name] + got = normalize_onnx_dims(raw) + if len(raw) != len(want): + raise ValueError( + f"{label} ONNX input '{name}' shape mismatch: graph={got}, " + f"expected={want}." + ) + for axis, (raw_dim, want_dim) in enumerate(zip(raw, want)): + if want_dim == -1: + if not isinstance(raw_dim, str): + raise ValueError( + f"{label} ONNX input '{name}' axis {axis} must be dynamic " + f"to accept a varying prompt-token count, but the graph " + f"declares it as {raw_dim!r} (graph={got})." + ) + elif not isinstance(raw_dim, int) or raw_dim != want_dim: + raise ValueError( + f"{label} ONNX input '{name}' shape mismatch: graph={got}, " + f"expected={want}." + ) + + +def classify_decoder_outputs(outputs: Sequence[np.ndarray]) -> dict[str, np.ndarray]: + """Name the four Hiera decoder outputs by their unambiguous element counts. + + qbruntime does not guarantee that the runtime output order matches the + compiled graph's declared order, so each output is identified by its + unique flattened size instead of position: ``masks`` is the only output + whose size is a multiple of ``256*256``; among the rest, ``iou`` has size + ``num_masks``, ``sam_tokens`` has size ``num_masks*256``, and + ``object_score`` has size 1. + """ + + arrays = [ + np.ascontiguousarray(np.asarray(value), dtype=np.float32) for value in outputs + ] + # A NaN from a numerical/runtime failure would otherwise reach argmax over + # the IoU scores (silently selecting the wrong candidate) and the `> 0` + # mask threshold (turning non-finite logits into plausible booleans), + # corrupting predictions and SA-V metrics instead of reporting the failure. + for index, array in enumerate(arrays): + if not bool(np.isfinite(array).all()): + raise ValueError( + f"Decoded mask generation outputs must be finite; output {index} " + f"with shape {array.shape} contains NaN or infinity." + ) + mask_matches = [ + array + for array in arrays + if array.size >= MASK_AREA and array.size % MASK_AREA == 0 + ] + if len(mask_matches) != 1: + raise ValueError( + f"Expected exactly one mask output, found {len(mask_matches)} " + f"among shapes {[array.shape for array in arrays]}." + ) + # Reshaping blindly would silently interleave the candidates' pixels for a + # decoder that emits NHWC `(1, 256, 256, 3)` instead of the expected layout: + # that size is also a multiple of MASK_AREA and also yields three + # candidates, so neither the size match above nor the count check below can + # see it, and the result is plausible but corrupted masks. Pin the layout to + # the two the supported artifacts actually produce -- the MXQ runtime's + # flattened `(3, 65536)` and the ONNX graph's `(3, 256, 256)` -- ignoring + # leading batch axes. Checked here rather than against ONNX session output + # metadata so the MXQ path is covered by the same guard. + mask_shape = tuple(int(dim) for dim in mask_matches[0].shape) + unbatched_layout = ( + tuple(dim for dim in mask_shape[:-2] if dim != 1) + mask_shape[-2:] + ) + # Structure only -- the candidate count is checked separately below, so a + # decoder emitting the right layout with the wrong number of candidates + # still reports that rather than a layout error. + is_flattened = len(unbatched_layout) == 2 and unbatched_layout[-1] == MASK_AREA + is_spatial = len(unbatched_layout) == 3 and unbatched_layout[-2:] == ( + MASK_SIDE, + MASK_SIDE, + ) + if not (is_flattened or is_spatial): + raise ValueError( + f"Decoder mask output has an unsupported layout {mask_shape}; expected " + f"candidates as (N, {MASK_AREA}) or (N, {MASK_SIDE}, {MASK_SIDE}), " + "optionally batched. A channels-last layout would interleave the " + "candidates into corrupted masks." + ) + masks = mask_matches[0].reshape(-1, MASK_SIDE, MASK_SIDE) + num_masks = masks.shape[0] + # The compiled artifacts bake in SAM2's multimask slice, so a stale or + # differently exported decoder emitting 2 or 4 masks (with consistently + # sized iou/token outputs) would otherwise be accepted here and reach the + # caller as a Results.masks shape that violates the documented fixed + # three-candidate contract. Mirrors eval_sav.CANDIDATES_PER_PROMPT. + if num_masks != MASK_CANDIDATE_COUNT: + raise ValueError( + f"Expected exactly {MASK_CANDIDATE_COUNT} decoder mask candidates, " + f"got {num_masks} (mask output shape {mask_matches[0].shape})." + ) + + def unique(label: str, size: int) -> np.ndarray: + matches = [ + array + for array in arrays + if array is not mask_matches[0] and array.size == size + ] + if len(matches) != 1: + raise ValueError( + f"Expected exactly one '{label}' output of size {size}, found {len(matches)}." + ) + return matches[0] + + return { + "masks": masks, + "iou": unique("iou", num_masks).reshape(num_masks), + "sam_tokens": unique("sam_tokens", num_masks * 256).reshape(num_masks, 256), + "object_score": unique("object_score", 1).reshape(1), + } diff --git a/mblt_vision/mask_generation/_sam2_host.py b/mblt_vision/mask_generation/_sam2_host.py new file mode 100644 index 0000000..33b21ca --- /dev/null +++ b/mblt_vision/mask_generation/_sam2_host.py @@ -0,0 +1,354 @@ +"""Host-side SAM2 glue: feature-map bookkeeping, prompt encoding, and mask resizing. + +Only the two heavy backbone/mask-decoder networks run on a backend -- the NPU +via ``MobilintNPUBackend`` or ONNX Runtime via ``ONNXBackend`` (see +``sam2.py``). Everything here runs on the host with +plain ``torch`` and the tiny prompt-encoder weights in ``_sam2_prompt.py`` -- +no dependency on the ``sam2`` package, no ``torchvision``, no manually cloned +repository, and no ~900MB full checkpoint download. Ported from the validated +``sam2-mxq-pipeline`` reference, restructured to take plain tensors (feature +maps, weights) instead of mutating a live ``SAM2ImagePredictor``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Sequence, TypedDict + +import numpy as np +import torch +import torch.nn.functional as functional + +from . import _sam2_prompt as prompt + + +class BackboneFeatures(TypedDict): + """``build_backbone_features``'s two differently-shaped outputs. + + ``image_embed`` is a single deepest-level feature tensor; ``high_res_feats`` + is the list of shallower skip-connection tensors -- a plain + ``dict[str, list[Tensor]]`` cannot express that the two keys hold different + types. + """ + + image_embed: torch.Tensor + high_res_feats: list[torch.Tensor] + + +_FPN_CHANNELS = (32, 64, 256) +# Each FPN level's full (channels, height, width). `build_backbone_features` +# `view()`s these into `prompt.BB_FEAT_SIZES`, so a channel-correct level with +# the wrong spatial geometry but the same element count -- (1, 32, 128, 512) +# for (1, 32, 256, 256) -- would be silently rearranged into a plausible but +# corrupted FPN, after which the decoder feed still has the expected shape and +# passes its own validation. Pin the complete shape, not just the channel count. +_FPN_LEVEL_SHAPES: dict[int, tuple[int, int]] = dict( + zip(_FPN_CHANNELS, prompt.BB_FEAT_SIZES) +) + +_NORMALIZE_MEAN = torch.tensor(prompt.NORMALIZE_MEAN, dtype=torch.float32).view( + -1, 1, 1 +) +_NORMALIZE_STD = torch.tensor(prompt.NORMALIZE_STD, dtype=torch.float32).view(-1, 1, 1) + + +def load_rgb(path: str | Path) -> np.ndarray: + """Load an image file as an RGB ``uint8`` array.""" + + from PIL import Image + + return np.asarray(Image.open(path).convert("RGB")) + + +def preprocess_encoder_input(image: np.ndarray) -> np.ndarray: + """Resize/normalize an RGB image into the encoder MXQ's NHWC input. + + Ported from ``SAM2Transforms.__call__`` (``ToTensor`` then + ``Resize((1024, 1024))`` + ImageNet normalize) in plain torch, verified + bit-for-bit against the torchvision pipeline it replaces: ``ToTensor`` + scales only ``uint8`` inputs by 1/255, and tensor ``Resize`` is bilinear + with ``antialias=True``. + """ + + array = np.asarray(image) + if array.ndim != 3 or array.shape[2] != 3: + raise ValueError(f"Expected an HWC RGB image, got shape {array.shape}.") + # Interpolation and normalization preserve NaN/infinity, so a non-finite + # source image would reach the backend as an invalid encoder tensor and + # fail in a backend-specific way -- by the time the decoder-output + # finiteness check fires, the source is no longer identifiable. + if not ( + np.issubdtype(array.dtype, np.floating) + or np.issubdtype(array.dtype, np.integer) + or array.dtype == np.bool_ + ): + raise ValueError(f"Expected a numeric RGB image, got dtype {array.dtype}.") + if np.issubdtype(array.dtype, np.floating) and not bool(np.isfinite(array).all()): + raise ValueError("Source image must be finite; got NaN or infinity.") + # HWC -> CHW as a fresh contiguous copy; PIL-backed arrays are read-only + # and torch.from_numpy warns on non-writable memory. + chw = np.ascontiguousarray(array.transpose(2, 0, 1)) + if not chw.flags.writeable: + chw = chw.copy() + tensor = torch.from_numpy(chw) + if tensor.dtype == torch.uint8: + tensor = tensor.float().div(255.0) + else: + tensor = tensor.float() + tensor = functional.interpolate( + tensor[None], + size=prompt.INPUT_IMAGE_SIZE, + mode="bilinear", + align_corners=False, + antialias=True, + ) + tensor = (tensor - _NORMALIZE_MEAN) / _NORMALIZE_STD + return np.ascontiguousarray( + tensor.permute(0, 2, 3, 1).float().cpu().numpy(), dtype=np.float32 + ) + + +def fpn_from_runtime( + outputs: Sequence[np.ndarray], device: torch.device +) -> list[torch.Tensor]: + """Convert the three raw encoder-MXQ outputs into ordered NCHW FPN levels. + + Returns ``[feat_32ch, feat_64ch, feat_256ch]`` -- the two high-resolution + skip-connection levels followed by the deepest ``image_embed`` level -- + identified by channel count rather than position, since qbruntime does + not guarantee the runtime output order. + """ + + features: dict[int, torch.Tensor] = {} + for output in outputs: + array = np.asarray(output, dtype=np.float32) + if array.ndim == 4 and array.shape[0] == 1: + array = array[0] + if array.ndim != 3: + continue + if array.shape[-1] in _FPN_CHANNELS: + channel = int(array.shape[-1]) + tensor = torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1)[ + None + ] + elif array.shape[0] in _FPN_CHANNELS: + channel = int(array.shape[0]) + tensor = torch.from_numpy(np.ascontiguousarray(array))[None] + else: + continue + if channel in features: + raise ValueError(f"Duplicate encoder output with {channel} channels.") + expected_spatial = _FPN_LEVEL_SHAPES[channel] + if tuple(tensor.shape[2:]) != expected_spatial: + raise ValueError( + f"Encoder output with {channel} channels must be " + f"{expected_spatial} spatially, got {tuple(tensor.shape[2:])}. " + "A same-element-count geometry would be silently rearranged into " + "a corrupted FPN." + ) + features[channel] = tensor.to(device) + + missing = [channel for channel in _FPN_CHANNELS if channel not in features] + if missing: + raise ValueError( + f"Encoder outputs are missing FPN channel count(s) {missing}; " + f"got shapes {[np.asarray(o).shape for o in outputs]}." + ) + return [features[32], features[64], features[256]] + + +def fpn_from_onnx( + outputs: Sequence[np.ndarray], device: torch.device +) -> list[torch.Tensor]: + """Convert the three batched-NCHW encoder-ONNX outputs into ordered FPN levels. + + Returns ``[feat_32ch, feat_64ch, feat_256ch]`` like :func:`fpn_from_runtime`. + The exported graph declares batched NCHW outputs, so the channel axis is + fixed at axis 1 -- unlike the MXQ runtime path, a square spatial size can + never be mistaken for a channel count, and any other layout is rejected. + """ + + features: dict[int, torch.Tensor] = {} + shapes = [tuple(np.asarray(output).shape) for output in outputs] + for output in outputs: + array = np.asarray(output, dtype=np.float32) + if ( + array.ndim != 4 + or array.shape[0] != 1 + or array.shape[1] not in _FPN_CHANNELS + ): + raise ValueError( + f"Unexpected encoder ONNX output shape {array.shape}; expected " + f"(1, C, H, W) with C in {_FPN_CHANNELS}. Got shapes {shapes}." + ) + channel = int(array.shape[1]) + expected_spatial = _FPN_LEVEL_SHAPES[channel] + if tuple(array.shape[2:]) != expected_spatial: + raise ValueError( + f"Encoder ONNX output with {channel} channels must be " + f"{expected_spatial} spatially, got {tuple(array.shape[2:])}. " + "A same-element-count geometry would be silently rearranged into " + "a corrupted FPN." + ) + if channel in features: + raise ValueError(f"Duplicate encoder output with {channel} channels.") + features[channel] = torch.from_numpy(np.ascontiguousarray(array)).to(device) + + missing = [channel for channel in _FPN_CHANNELS if channel not in features] + if missing: + raise ValueError( + f"Encoder ONNX outputs are missing FPN channel count(s) {missing}; " + f"got shapes {shapes}." + ) + return [features[32], features[64], features[256]] + + +def build_backbone_features( + weights: dict[str, torch.Tensor], feature_maps: Sequence[torch.Tensor] +) -> BackboneFeatures: + """Ported from SAM2's ``_prepare_backbone_features`` for a single image + (no video memory: ``directly_add_no_mem_embed`` is always applied).""" + + vision_features = [value.flatten(2).permute(2, 0, 1) for value in feature_maps] + if prompt.DIRECTLY_ADD_NO_MEM_EMBED: + vision_features[-1] = vision_features[-1] + weights["no_mem_embed"] + features = [ + feature.permute(1, 2, 0).view(1, -1, *feature_size) + for feature, feature_size in zip( + vision_features[::-1], prompt.BB_FEAT_SIZES[::-1] + ) + ][::-1] + return {"image_embed": features[-1], "high_res_feats": features[:-1]} + + +def _as_float32_arrays(tensors: dict[str, torch.Tensor]) -> dict[str, np.ndarray]: + """Convert named torch tensors into contiguous float32 numpy arrays.""" + + return { + name: np.ascontiguousarray( + value.detach().float().cpu().numpy(), dtype=np.float32 + ) + for name, value in tensors.items() + } + + +def _decoder_prompt_tensors( + weights: dict[str, torch.Tensor], + features: BackboneFeatures, + points: np.ndarray, + labels: np.ndarray, + original_hw: Sequence[int], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]: + """Run the host prompt encoder shared by the MXQ and ONNX decoder feeds. + + Returns ``(tokens, src, pos_src, high_res)`` in the pre-flattening shapes + the mask decoder was traced with: ``tokens (1, N+7, 256)``, + ``src``/``pos_src`` NCHW ``(1, 256, 64, 64)``, and the two NCHW + high-resolution feature maps. + """ + + # Prompts arrive as host numpy arrays, so place them on the same device the + # weights live on: the prompt encoder combines them with the Gaussian + # position-encoding matrix and the learned embeddings, which torch requires + # to share a device once the engine was constructed with (or moved to) CUDA. + device = weights["positional_encoding_gaussian_matrix"].device + points_tensor = torch.as_tensor(points, dtype=torch.float32, device=device)[ + None, ... + ] + labels_tensor = torch.as_tensor(labels, dtype=torch.int64, device=device)[None, ...] + height, width = original_hw + unnorm_coords = prompt.transform_points(points_tensor, (int(height), int(width))) + + sparse = prompt.embed_points(weights, unnorm_coords, labels_tensor) + dense = prompt.dense_embeddings_for_no_mask(weights, batch_size=sparse.size(0)) + + image_embeddings = features["image_embed"][-1].unsqueeze(0) + high_res = [value[-1].unsqueeze(0) for value in features["high_res_feats"]] + tokens, src, pos_src = prompt.decoder_token_prep( + weights, + image_embeddings=image_embeddings.float(), + dense_prompt_embeddings=dense.float(), + sparse_prompt_embeddings=sparse.float(), + ) + return tokens, src, pos_src, high_res + + +def prepare_decoder_tensors( + weights: dict[str, torch.Tensor], + features: BackboneFeatures, + points: np.ndarray, + labels: np.ndarray, + original_hw: Sequence[int], +) -> dict[str, np.ndarray]: + """Run only the host prompt encoder and build the six compiled decoder inputs.""" + + tokens, src, pos_src, high_res = _decoder_prompt_tensors( + weights, features, points, labels, original_hw + ) + + def sequence(value: torch.Tensor) -> torch.Tensor: + return ( + value.flatten(2) + .transpose(1, 2) + .reshape(1, 1, -1, value.shape[1]) + .contiguous() + ) + + src_sequence = sequence(src) + pos_sequence = sequence(pos_src) + tensors = { + "hrf0_nhwc": high_res[0].permute(0, 2, 3, 1).contiguous(), + "hrf1_nhwc": high_res[1].permute(0, 2, 3, 1).contiguous(), + "src": src_sequence, + "tokens": tokens.reshape(1, 1, -1, prompt.EMBED_DIM).contiguous(), + "pos_src": pos_sequence, + "src_plus_pos_src": src_sequence + pos_sequence, + } + return _as_float32_arrays(tensors) + + +def prepare_decoder_tensors_onnx( + weights: dict[str, torch.Tensor], + features: BackboneFeatures, + points: np.ndarray, + labels: np.ndarray, + original_hw: Sequence[int], +) -> dict[str, np.ndarray]: + """Run the host prompt encoder and build the five named ONNX decoder inputs. + + The exported decoder consumes the pre-flattening NCHW tensors the graph was + traced with; ``src + pos_src`` stays inside the graph, unlike the compiled + MXQ artifact's flattened six-input runtime signature. + """ + + tokens, src, pos_src, high_res = _decoder_prompt_tensors( + weights, features, points, labels, original_hw + ) + return _as_float32_arrays( + { + "tokens": tokens.contiguous(), + "src": src.contiguous(), + "pos_src": pos_src.contiguous(), + "high_res_features_0": high_res[0].contiguous(), + "high_res_features_1": high_res[1].contiguous(), + } + ) + + +def postprocess_masks( + low_resolution_masks: np.ndarray, original_hw: Sequence[int] +) -> np.ndarray: + """Resize low-res decoder mask logits back to the original image size. + + Ported from ``SAM2Transforms.postprocess_masks`` (``max_hole_area`` / + ``max_sprinkle_area`` are 0 by default upstream, so that branch is a + no-op and is not ported). + """ + + tensor = torch.from_numpy( + np.ascontiguousarray(low_resolution_masks, dtype=np.float32) + )[None] + masks = functional.interpolate( + tensor, tuple(original_hw), mode="bilinear", align_corners=False + )[0] + return masks.detach().float().cpu().numpy() diff --git a/mblt_vision/mask_generation/_sam2_prompt.py b/mblt_vision/mask_generation/_sam2_prompt.py new file mode 100644 index 0000000..63055af --- /dev/null +++ b/mblt_vision/mask_generation/_sam2_prompt.py @@ -0,0 +1,223 @@ +"""Self-contained SAM2 Hiera-Large host-side prompt encoding. + +Ported line-for-line from the official ``facebookresearch/sam2`` modules +(``modeling/position_encoding.py::PositionEmbeddingRandom``, +``modeling/sam/prompt_encoder.py::PromptEncoder``, +``modeling/sam/mask_decoder.py``'s token embeddings, and +``sam2_image_predictor.py``'s ``_prep_prompts``/``_bb_feat_sizes``), restricted +to the point-only prompt path this package supports (no box prompts, no mask +prompts). Depends on nothing from the ``sam2`` package -- the only weights +needed for this path are the ~3k floats extracted from the official +``facebook/sam2-hiera-large`` checkpoint into ``sam2_hiera_large_prompt_weights.pt`` +(point/"not-a-point"/"no-mask" embeddings, the random Fourier position-encoding +matrix, the decoder's IoU/mask/object-score tokens, and the no-memory-embedding +parameter) -- not the ~900MB full checkpoint, which also carries the Hiera +backbone and video-memory modules this package never runs on the host (the +backbone runs on the NPU; there is no video memory in single-image inference). + +That small bundle lives at ``mobilint/sam2-hiera-large``'s Hub repo root as +``sam2_hiera_large_prompt_weights.pt``, downloaded the same way as the +encoder/decoder MXQ artifacts (see ``sam2.py``) -- not shipped as package +data, consistent with every other model's artifacts living on the Hub rather +than in the wheel. + +Numerically verified against the real ``facebookresearch/sam2`` predictor's +``sam_prompt_encoder``/``sam_mask_decoder`` on the same weights. +""" + +from __future__ import annotations + +import math +from pathlib import Path + +import torch + +# Fixed SAM2 Hiera-Large configuration (facebookresearch/sam2, sam2_hiera_l.yaml +# and modeling/sam2_base.py::_build_sam_heads); none of these are learned. +EMBED_DIM = 256 +IMAGE_EMBEDDING_SIZE = (64, 64) # image_size // backbone_stride == 1024 // 16 +INPUT_IMAGE_SIZE = (1024, 1024) +BB_FEAT_SIZES = ((256, 256), (128, 128), (64, 64)) +NUM_MASK_TOKENS = 4 # num_multimask_outputs (3) + 1 +USE_MULTIMASK_TOKEN_FOR_OBJ_PTR = True +PRED_OBJ_SCORES = True +DIRECTLY_ADD_NO_MEM_EMBED = True +MASK_THRESHOLD = 0.0 + +# SAM2's ImageNet-style input normalization (sam2/utils/transforms.py). +NORMALIZE_MEAN = (0.485, 0.456, 0.406) +NORMALIZE_STD = (0.229, 0.224, 0.225) + + +def load_prompt_weights(path: str | Path) -> dict[str, torch.Tensor]: + """Load the point/label/token embeddings extracted from the official + ``facebook/sam2-hiera-large`` checkpoint.""" + + return torch.load(path, map_location="cpu", weights_only=True) + + +def positional_encoding_for_grid( + gaussian_matrix: torch.Tensor, size: tuple[int, int] +) -> torch.Tensor: + """Dense positional encoding for an ``size`` grid. + + Ported from ``PositionEmbeddingRandom.forward``. Used once per model + (``get_dense_pe()``); ``size`` is always ``IMAGE_EMBEDDING_SIZE`` here. + """ + + h, w = size + device = gaussian_matrix.device + grid = torch.ones((h, w), device=device, dtype=torch.float32) + y_embed = (grid.cumsum(dim=0) - 0.5) / h + x_embed = (grid.cumsum(dim=1) - 0.5) / w + coords = torch.stack([x_embed, y_embed], dim=-1) + pe = _pe_encoding(coords, gaussian_matrix) + return pe.permute(2, 0, 1) # C x H x W + + +def _pe_encoding(coords: torch.Tensor, gaussian_matrix: torch.Tensor) -> torch.Tensor: + """Ported from ``PositionEmbeddingRandom._pe_encoding``. ``coords`` in [0, 1].""" + + coords = 2 * coords - 1 + coords = coords @ gaussian_matrix + coords = 2 * math.pi * coords + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + +def _forward_with_coords( + coords_input: torch.Tensor, + image_size: tuple[int, int], + gaussian_matrix: torch.Tensor, +) -> torch.Tensor: + """Ported from ``PositionEmbeddingRandom.forward_with_coords``.""" + + coords = coords_input.clone() + coords[..., 0] = coords[..., 0] / image_size[1] + coords[..., 1] = coords[..., 1] / image_size[0] + return _pe_encoding(coords.to(torch.float32), gaussian_matrix) + + +def embed_points( + weights: dict[str, torch.Tensor], points: torch.Tensor, labels: torch.Tensor +) -> torch.Tensor: + """Sparse point-prompt embeddings. + + Ported from ``PromptEncoder._embed_points`` and ``PromptEncoder.forward``'s + points branch, restricted to the point-only path (``pad=True``, no boxes): + always pads with one "not a point" token, matching this package's prompts + (there is never a box prompt to pad against instead). + + Args: + points: ``(B, N, 2)`` pixel coordinates in the 1024x1024 encoder input + space (already produced by :func:`transform_points`). + labels: ``(B, N)`` point labels (``1`` positive, ``0`` negative). + + Returns: + ``(B, N + 1, embed_dim)`` sparse embeddings (the ``+1`` is the padding + "not a point" token SAM2 always appends when there is no box prompt). + """ + + points = points + 0.5 # shift to center of pixel + # Built on the incoming tensors' own device (not the implicit CPU default), + # so the concatenation below still works when the engine was constructed + # with device="cuda". Numerically identical to the upstream port. + padding_point = torch.zeros( + (points.shape[0], 1, 2), dtype=points.dtype, device=points.device + ) + padding_label = -torch.ones( + (labels.shape[0], 1), dtype=labels.dtype, device=labels.device + ) + points = torch.cat([points, padding_point], dim=1) + labels = torch.cat([labels, padding_label], dim=1) + + point_embedding = _forward_with_coords( + points, INPUT_IMAGE_SIZE, weights["positional_encoding_gaussian_matrix"] + ) + is_padding = (labels == -1).unsqueeze(-1) + point_embedding = torch.where( + is_padding, + torch.zeros_like(point_embedding) + weights["not_a_point_embed_weight"], + point_embedding, + ) + is_negative = (labels == 0).unsqueeze(-1) + point_embedding = torch.where( + is_negative, + point_embedding + weights["point_embedding_negative"], + point_embedding, + ) + is_positive = (labels == 1).unsqueeze(-1) + point_embedding = torch.where( + is_positive, + point_embedding + weights["point_embedding_positive"], + point_embedding, + ) + return point_embedding + + +def dense_embeddings_for_no_mask( + weights: dict[str, torch.Tensor], batch_size: int +) -> torch.Tensor: + """Dense embeddings when no mask prompt is given. + + Ported from ``PromptEncoder.forward``'s ``masks is None`` branch. + """ + + no_mask_embed = weights["no_mask_embed_weight"] + return no_mask_embed.reshape(1, -1, 1, 1).expand( + batch_size, -1, IMAGE_EMBEDDING_SIZE[0], IMAGE_EMBEDDING_SIZE[1] + ) + + +def get_dense_pe(weights: dict[str, torch.Tensor]) -> torch.Tensor: + """Ported from ``PromptEncoder.get_dense_pe``.""" + + return positional_encoding_for_grid( + weights["positional_encoding_gaussian_matrix"], IMAGE_EMBEDDING_SIZE + ).unsqueeze(0) + + +def transform_points( + points: torch.Tensor, original_hw: tuple[int, int] +) -> torch.Tensor: + """Normalize original-image-pixel point coordinates into the encoder's + 1024x1024 input space. + + Ported from ``SAM2Transforms.transform_coords`` with ``normalize=True``. + """ + + height, width = original_hw + coords = points.clone() + coords[..., 0] = coords[..., 0] / width + coords[..., 1] = coords[..., 1] / height + return coords * INPUT_IMAGE_SIZE[0] + + +def decoder_token_prep( + weights: dict[str, torch.Tensor], + image_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Host-side token concatenation and image/dense embedding sum used by the + compiled decoder MXQ. + + Ported from the reference ``sam2-mxq-pipeline``'s ``sam2_decoder_prep``, + itself ported from the mask decoder's ``forward``/``predict_masks`` token + setup, restricted to this package's fixed, validated decoder contract + (``pred_obj_scores=True``, ``use_multimask_token_for_obj_ptr=True``). + """ + + output_tokens = torch.cat( + [ + weights["obj_score_token_weight"], + weights["iou_token_weight"], + weights["mask_tokens_weight"], + ], + dim=0, + ) + output_tokens = output_tokens.unsqueeze(0).expand( + sparse_prompt_embeddings.size(0), -1, -1 + ) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + image_pe = get_dense_pe(weights) + return tokens, image_embeddings + dense_prompt_embeddings, image_pe diff --git a/mblt_vision/mask_generation/sam2.py b/mblt_vision/mask_generation/sam2.py new file mode 100644 index 0000000..c523e11 --- /dev/null +++ b/mblt_vision/mask_generation/sam2.py @@ -0,0 +1,762 @@ +"""SAM2 (Segment Anything 2) promptable mask generation. + +Ported from the validated ``sam2-mxq-pipeline`` reference (real Aries2 +SA-V-200 accuracy: FP32 mIoU 0.7750 vs MXQ mIoU 0.7757, mask agreement 0.983). +Only point prompts (1-3 points, positive/negative) are supported, matching +that reference; box prompts and automatic "segment everything" grid mode are +out of scope. + +Unlike every other Vision model, SAM2 needs three independently downloaded +artifacts -- an image encoder MXQ, a prompt-conditioned mask decoder MXQ, and +a small (~16KB) bundle of host-side prompt-encoder weights extracted from the +official checkpoint (see ``_sam2_prompt.py``) -- and takes point prompts +rather than running end-to-end on an image alone. So, unlike a +``create_model_class``-generated model, ``SAM2HieraLarge`` does not go through +``MBLT_Engine.__init__``, ``build_preprocess``/``build_postprocess``, or +``file_config_cleansing`` -- it subclasses ``MBLT_Engine`` only so +``mblt_vision.list_models()``'s ``issubclass(obj, MBLT_Engine)`` filter +discovers it, and fully owns its own init/inference/cleanup. + +No dependency on the ``sam2`` package (the PyPI ``sam2`` is an unofficial +third-party mirror, not Meta's) and no manually cloned repository: host-side +prompt encoding is reimplemented from the official source in ``_sam2_prompt.py`` +and ``_sam2_host.py``, numerically verified bit-for-bit against the real +``facebookresearch/sam2`` predictor (see ``tests/test_mask_generation.py``). +""" + +from __future__ import annotations + +import math +from contextlib import suppress +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch + +from mblt_npu import MobilintNPUBackend, ONNXBackend, normalize_target_device + +from ..utils.results import Results +from ..wrapper import ( + CoreMode, + MBLT_Engine, + _load_onnxruntime, + _resolve_onnx_providers, + download_hub_artifact, +) +from . import _sam2_prompt as prompt +from ._sam2_contracts import ( + DECODER_ONNX_INPUT_NAMES, + DECODER_ONNX_INPUT_SHAPES, + DECODER_RUNTIME_ORDER, + ENCODER_ONNX_INPUT_NAME, + ENCODER_ONNX_INPUT_SHAPE, + build_decoder_runtime_feed, + classify_decoder_outputs, + strip_runtime_batch, + validate_onnx_session_inputs, + validate_runtime_shapes, +) +from ._sam2_host import ( + BackboneFeatures, + build_backbone_features, + fpn_from_onnx, + fpn_from_runtime, + load_rgb, + postprocess_masks, + preprocess_encoder_input, + prepare_decoder_tensors, + prepare_decoder_tensors_onnx, +) + +_REPO_ID = "mobilint/sam2-hiera-large" +_ENCODER_FILENAME = "sam2_hiera_large_encoder.mxq" +_DECODER_FILENAME = "sam2_hiera_large_decoder.mxq" +# The ONNX exports follow the package-wide same-stem convention and live at +# the Hub repo root (board-agnostic), unlike the board-folder MXQ artifacts. +_ENCODER_ONNX_FILENAME = "sam2_hiera_large_encoder.onnx" +_DECODER_ONNX_FILENAME = "sam2_hiera_large_decoder.onnx" +_PROMPT_WEIGHTS_FILENAME = "sam2_hiera_large_prompt_weights.pt" + + +class SAM2HieraLarge(MBLT_Engine): + """Promptable mask generation with SAM2 Hiera-Large. + + Args: + encoder_mxq_path: Explicit local path to the encoder MXQ artifact. + When omitted, downloaded from ``mobilint/sam2-hiera-large``. + decoder_mxq_path: Explicit local path to the decoder MXQ artifact. + When omitted, downloaded from ``mobilint/sam2-hiera-large``. + encoder_dev_no: Accelerator device number for the encoder backend. + decoder_dev_no: Accelerator device number for the decoder backend. + encoder_core_mode: NPU core mode for the encoder backend. Defaults to + ``"single"``, matching the validated reference configuration. + decoder_core_mode: NPU core mode for the decoder backend. Defaults to + ``"single"``. + encoder_target_cores: Optional explicit core selection for the encoder. + decoder_target_cores: Optional explicit core selection for the decoder. + encoder_target_clusters: Optional explicit cluster selection for the encoder. + decoder_target_clusters: Optional explicit cluster selection for the decoder. + target_device: NPU board identifier shared by both MXQ artifacts. + Defaults to ``"aries-rb"`` -- the only board this port has been + validated on. + revision: Hugging Face Hub revision for all artifacts. + prompt_weights_path: Explicit local path to the host-side + prompt-encoder weights bundle. When omitted, downloaded from + ``mobilint/sam2-hiera-large`` (repo root, not board-specific). + device: Torch device for the host-side prompt encoding (a handful of + small embeddings/lookups). The encoder/decoder artifacts always + run on their selected backend regardless of this setting. + framework: ``"mxq"`` (NPU, default) or ``"onnx"`` (ONNX Runtime). + When omitted, inferred from explicit local artifact paths, + matching ``MBLT_Engine`` semantics; a path whose suffix conflicts + with an explicitly selected framework fails fast. The NPU-only + arguments (``*_dev_no``, ``*_core_mode``, ``*_target_cores``, + ``*_target_clusters``, ``target_device``) are ignored for ONNX, + as in ``MBLT_Engine``. + encoder_onnx_path: Explicit local path to the encoder ONNX artifact. + When omitted with ``framework="onnx"``, downloaded from + ``mobilint/sam2-hiera-large`` (repo root, not board-specific). + decoder_onnx_path: Explicit local path to the decoder ONNX artifact. + When omitted with ``framework="onnx"``, downloaded from + ``mobilint/sam2-hiera-large`` (repo root, not board-specific). + onnx_providers: Optional ONNX Runtime execution provider order. + Defaults to CPU execution. + """ + + def __init__( + self, + encoder_mxq_path: str | None = None, + decoder_mxq_path: str | None = None, + encoder_dev_no: int | None = None, + decoder_dev_no: int | None = None, + encoder_core_mode: CoreMode | None = None, + decoder_core_mode: CoreMode | None = None, + encoder_target_cores: Sequence[str] | None = None, + decoder_target_cores: Sequence[str] | None = None, + encoder_target_clusters: Sequence[int] | None = None, + decoder_target_clusters: Sequence[int] | None = None, + target_device: str | None = None, + revision: str | None = None, + prompt_weights_path: str | None = None, + device: str = "cpu", + framework: str | None = None, + encoder_onnx_path: str | None = None, + decoder_onnx_path: str | None = None, + onnx_providers: Sequence[str] | None = None, + ) -> None: + self.pre_cfg: dict[str, Any] = {} + self.post_cfg: dict[str, Any] = {"task": "mask_generation", "dataset": "sa-v"} + self.device = torch.device(device) + self._closed = False + self._encoder_backend: MobilintNPUBackend | ONNXBackend | None = None + self._decoder_backend: MobilintNPUBackend | ONNXBackend | None = None + self.weights: dict[str, torch.Tensor] | None = None + + for label, path, suffix in ( + ("encoder_mxq_path", encoder_mxq_path, ".mxq"), + ("decoder_mxq_path", decoder_mxq_path, ".mxq"), + ("encoder_onnx_path", encoder_onnx_path, ".onnx"), + ("decoder_onnx_path", decoder_onnx_path, ".onnx"), + ): + if path and Path(path).suffix.lower() != suffix: + raise ValueError( + f"Explicit {label} must end in '{suffix}', got {path!r}." + ) + + self.framework = self._resolve_framework( + framework, + mxq_path_passed=bool(encoder_mxq_path or decoder_mxq_path), + onnx_path_passed=bool(encoder_onnx_path or decoder_onnx_path), + ) + # Every explicitly supplied artifact -- the prompt-weights bundle + # included -- is checked after the argument-coherence checks above but + # before any download, so an invalid local path is reported as such + # rather than surfacing as a Hub/network failure offline, or costing a + # pointless download online. Mirrors MBLT_Engine's fail-fast + # FileNotFoundError for explicit paths. + for label, path in ( + ("encoder_mxq_path", encoder_mxq_path), + ("decoder_mxq_path", decoder_mxq_path), + ("encoder_onnx_path", encoder_onnx_path), + ("decoder_onnx_path", decoder_onnx_path), + ("prompt_weights_path", prompt_weights_path), + ): + if path and not Path(path).is_file(): + raise FileNotFoundError( + f"Explicit {label} does not exist: {path}. Remove it to " + "download the configured artifact." + ) + + resolved_revision = revision or "main" + + try: + # Resolve the optional runtime before any download: a missing + # onnxruntime must surface as the documented package-extra install + # error rather than as a network/cache failure from an artifact + # fetch that would have been useless anyway. + ort = _load_onnxruntime() if self.framework == "onnx" else None + # Same reasoning for the MXQ board: an unknown target must report as + # such rather than as a network/cache failure from a download that + # would have been useless. Still ignored for ONNX, which is + # board-agnostic. + resolved_target_device = ( + normalize_target_device(target_device or "aries-rb") + if self.framework == "mxq" + else "" + ) + + resolved_prompt_weights_path = prompt_weights_path or download_hub_artifact( + repo_id=_REPO_ID, + filename=_PROMPT_WEIGHTS_FILENAME, + revision=resolved_revision, + ) + + if self.framework == "onnx": + providers = _resolve_onnx_providers(ort, onnx_providers) + resolved_encoder_path = encoder_onnx_path or download_hub_artifact( + repo_id=_REPO_ID, + filename=_ENCODER_ONNX_FILENAME, + revision=resolved_revision, + ) + resolved_decoder_path = decoder_onnx_path or download_hub_artifact( + repo_id=_REPO_ID, + filename=_DECODER_ONNX_FILENAME, + revision=resolved_revision, + ) + self._encoder_backend = self._build_onnx_backend( + onnx_path=resolved_encoder_path, + providers=providers, + ort_module=ort, + expected_inputs={ENCODER_ONNX_INPUT_NAME: ENCODER_ONNX_INPUT_SHAPE}, + label="encoder", + ) + self._decoder_backend = self._build_onnx_backend( + onnx_path=resolved_decoder_path, + providers=providers, + ort_module=ort, + expected_inputs=DECODER_ONNX_INPUT_SHAPES, + label="decoder", + ) + else: + resolved_encoder_path = encoder_mxq_path or download_hub_artifact( + repo_id=_REPO_ID, + filename=_ENCODER_FILENAME, + revision=resolved_revision, + subfolders=[resolved_target_device], + ) + resolved_decoder_path = decoder_mxq_path or download_hub_artifact( + repo_id=_REPO_ID, + filename=_DECODER_FILENAME, + revision=resolved_revision, + subfolders=[resolved_target_device], + ) + self._encoder_backend = self._build_backend( + mxq_path=resolved_encoder_path, + dev_no=encoder_dev_no, + core_mode=encoder_core_mode, + target_cores=encoder_target_cores, + target_clusters=encoder_target_clusters, + target_device=resolved_target_device, + ) + self._decoder_backend = self._build_backend( + mxq_path=resolved_decoder_path, + dev_no=decoder_dev_no, + core_mode=decoder_core_mode, + target_cores=decoder_target_cores, + target_clusters=decoder_target_clusters, + target_device=resolved_target_device, + ) + + weights = prompt.load_prompt_weights(resolved_prompt_weights_path) + self.weights = { + name: tensor.to(self.device) for name, tensor in weights.items() + } + except Exception: + # Suppress disposal failures while unwinding, as MBLT_Engine does: + # a backend that also raises from dispose() would otherwise replace + # the actionable original construction error. + self._close(suppress_errors=True) + raise + + @staticmethod + def _resolve_framework( + framework: str | None, *, mxq_path_passed: bool, onnx_path_passed: bool + ) -> str: + """Resolve the execution framework, mirroring ``MBLT_Engine`` semantics. + + Explicit local artifact paths select the framework when it is omitted, + and conflict loudly with an explicitly selected opposite framework. + """ + + # Lowercased before validation and before the path-conflict checks + # below, matching `_model_paths.resolve_framework`, so `framework="ONNX"` + # behaves the same here as it does for every other model. + if framework is not None: + if not isinstance(framework, str): + raise ValueError( + f"framework must be 'mxq' or 'onnx', got {framework!r}." + ) + framework = framework.lower() + if framework is not None and framework not in ("mxq", "onnx"): + raise ValueError(f"framework must be 'mxq' or 'onnx', got {framework!r}.") + if framework == "mxq" and onnx_path_passed: + raise ValueError( + "framework='mxq' conflicts with explicit encoder_onnx_path/" + "decoder_onnx_path; pass encoder_mxq_path/decoder_mxq_path instead." + ) + if framework == "onnx" and mxq_path_passed: + raise ValueError( + "framework='onnx' conflicts with explicit encoder_mxq_path/" + "decoder_mxq_path; pass encoder_onnx_path/decoder_onnx_path instead." + ) + if framework is not None: + return framework + if mxq_path_passed and onnx_path_passed: + raise ValueError( + "Both MXQ and ONNX artifact paths were passed without an explicit " + "framework; pass framework='mxq' or framework='onnx' with only the " + "matching artifact paths." + ) + return "onnx" if onnx_path_passed else "mxq" + + @staticmethod + def _build_backend( + *, + mxq_path: str, + dev_no: int | None, + core_mode: CoreMode | None, + target_cores: Sequence[str] | None, + target_clusters: Sequence[int] | None, + target_device: str, + ) -> MobilintNPUBackend: + """Build and launch one MXQ backend with SAM2's validated single-core defaults.""" + + backend = MobilintNPUBackend( + mxq_path=mxq_path, + dev_no=dev_no if dev_no is not None else 0, + core_mode=core_mode or "single", + 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=target_device, + ) + # Dispose here rather than leaving it to the constructor's cleanup: the + # caller only assigns the return value to self._*_backend on success, so + # a backend that fails after create() is unreachable from close(). + try: + backend.create() + backend.launch() + except Exception: + # Never let a dispose() failure replace the load/launch error that + # actually explains what went wrong. + with suppress(Exception): + backend.dispose() + raise + return backend + + @staticmethod + def _build_onnx_backend( + *, + onnx_path: str, + providers: Sequence[str], + ort_module: Any, + expected_inputs: dict[str, tuple[int, ...]], + label: str, + ) -> ONNXBackend: + """Build one ONNX Runtime backend and pin its graph interface. + + A resolved ONNX artifact whose input names or shapes drift from the + exported-graph contract (for example a re-export with different + wrapper input names) must fail here rather than silently produce + wrong masks, mirroring the MXQ path's construction-time shape check. + """ + + backend = ONNXBackend( + onnx_path, providers=list(providers), ort_module=ort_module + ) + # Dispose here rather than leaving it to the constructor's cleanup: the + # caller only assigns the return value to self._*_backend on success, so + # a session that fails validation is unreachable from close(). + try: + backend.create() + validate_onnx_session_inputs(backend.get_inputs(), expected_inputs, label) + except Exception: + # Never let a dispose() failure replace the session/validation error + # that actually explains what went wrong. + with suppress(Exception): + backend.dispose() + raise + return backend + + def preprocess(self, x: Any, **kwargs: Any) -> np.ndarray: + """Resize/pad/normalize an image into the canonical NHWC encoder input. + + The returned layout matches the encoder MXQ's runtime input and is + framework-independent: the ONNX path transposes to the exported + graph's NCHW layout internally. + """ + + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError( + f"SAM2HieraLarge.preprocess() does not support keyword arguments: " + f"{unexpected}" + ) + + image_array = load_rgb(x) if isinstance(x, (str, Path)) else np.asarray(x) + return preprocess_encoder_input(image_array) + + def __call__(self, x: Any) -> Any: + """Raw ``__call__`` is not meaningful for a promptable model.""" + + del x + raise NotImplementedError( + "SAM2HieraLarge requires point prompts; call predict(image, points, labels) " + "or predict_preprocessed(encoder_input, original_hw, points, labels) instead " + "of raw __call__()." + ) + + def predict_preprocessed( + self, + encoder_input: np.ndarray, + original_hw: Sequence[int], + points: Any, + labels: Any, + ) -> Results: + """Run the encoder/decoder MXQs on an already-preprocessed image and prompts. + + Args: + encoder_input: Output of :meth:`preprocess`. + original_hw: ``(height, width)`` of the original, un-preprocessed image. + points: ``(N, 2)`` array of ``(x, y)`` point-prompt coordinates in + original-image pixel space, ``1 <= N <= 3``. + labels: ``(N,)`` array of point labels (``1`` positive, ``0`` negative). + + Returns: + A :class:`Results` with ``task == "mask_generation"``. + """ + + self._ensure_open() + # Checked before either backend runs: a zero dimension makes + # transform_points produce infinite coordinates, and a fractional or + # wrong-length size fails only in the final mask resize, after both + # backends have already executed, as an opaque runtime error. + original_hw_values = list(np.asarray(original_hw).reshape(-1)) + if len(original_hw_values) != 2 or np.asarray(original_hw).ndim != 1: + raise ValueError( + f"Expected original_hw as (height, width), got {original_hw!r}." + ) + for axis, value in zip(("height", "width"), original_hw_values): + number = float(value) + if not math.isfinite(number) or number != int(number) or number <= 0: + raise ValueError( + f"original_hw {axis} must be a positive whole number, " + f"got {value!r}." + ) + original_hw = (int(original_hw_values[0]), int(original_hw_values[1])) + + points_array = np.asarray(points, dtype=np.float32) + # Read labels without an integer cast first: casting to int64 truncates, + # so 0.5/1.9/-0.1 would silently become valid 0/1 labels and pass the + # exact-label check below. + raw_labels = np.asarray(labels) + if points_array.ndim != 2 or points_array.shape[1] != 2: + raise ValueError( + f"Expected points shaped (N, 2), got {points_array.shape}." + ) + if not (1 <= len(points_array) <= 3): + raise ValueError(f"Expected 1 to 3 point prompts, got {len(points_array)}.") + # float32 preserves NaN/Inf, so this catches them before Fourier prompt + # encoding turns them into contaminated tokens, masks, and IoU scores. + if not bool(np.isfinite(points_array).all()): + raise ValueError( + f"Point coordinates must be finite, got {points_array.tolist()}." + ) + if raw_labels.ndim != 1: + raise ValueError(f"Expected labels shaped (N,), got {raw_labels.shape}.") + if len(raw_labels) != len(points_array): + raise ValueError("points and labels must have the same length.") + if not np.issubdtype(raw_labels.dtype, np.integer): + if not np.issubdtype(raw_labels.dtype, np.floating): + raise ValueError( + f"Point labels must be integers, got dtype {raw_labels.dtype}." + ) + if not bool(np.isfinite(raw_labels).all()): + raise ValueError( + f"Point labels must be finite, got {raw_labels.tolist()}." + ) + if not bool((raw_labels == np.floor(raw_labels)).all()): + raise ValueError( + f"Point labels must be whole numbers, got {raw_labels.tolist()}." + ) + labels_array = raw_labels.astype(np.int64) + # Any other value silently receives neither the positive nor the + # negative learned embedding in embed_points, which would return a + # plausible but semantically meaningless mask instead of an error. + invalid_labels = sorted(set(labels_array.tolist()) - {0, 1}) + if invalid_labels: + raise ValueError( + "Point labels must be 1 (positive) or 0 (negative), got " + f"{invalid_labels}." + ) + + weights = self._require_weights() + if self.framework == "onnx": + feature_maps = self._encode_image_onnx(encoder_input) + else: + feature_maps = self._encode_image_mxq(encoder_input) + features = build_backbone_features(weights, feature_maps) + + if self.framework == "onnx": + raw_decoder_outputs = self._decode_prompts_onnx( + weights, features, points_array, labels_array, original_hw + ) + else: + raw_decoder_outputs = self._decode_prompts_mxq( + weights, features, points_array, labels_array, original_hw + ) + decoder_outputs = classify_decoder_outputs(raw_decoder_outputs) + + full_logits = postprocess_masks(decoder_outputs["masks"], original_hw) + selected = int(np.argmax(decoder_outputs["iou"])) + binary_masks = full_logits > prompt.MASK_THRESHOLD + + output = { + "masks": binary_masks, + "low_res_masks": decoder_outputs["masks"], + "full_logits": full_logits, + "iou_predictions": decoder_outputs["iou"], + "object_score": decoder_outputs["object_score"], + "points": points_array, + "point_labels": labels_array, + "selected": selected, + } + return Results(self.pre_cfg, self.post_cfg, output) + + def predict( + self, image: str | Path | np.ndarray, points: Any, labels: Any + ) -> Results: + """Run end-to-end promptable mask generation on a raw image. + + Args: + image: Image path, or an HWC RGB array. + points: ``(N, 2)`` array of ``(x, y)`` point-prompt coordinates in + original-image pixel space, ``1 <= N <= 3``. + labels: ``(N,)`` array of point labels (``1`` positive, ``0`` negative). + + Returns: + A :class:`Results` with ``task == "mask_generation"``. + """ + + image_array = ( + load_rgb(image) if isinstance(image, (str, Path)) else np.asarray(image) + ) + encoder_input = self.preprocess(image_array) + original_hw = (int(image_array.shape[0]), int(image_array.shape[1])) + return self.predict_preprocessed(encoder_input, original_hw, points, labels) + + def _encode_image_mxq(self, encoder_input: np.ndarray) -> list[torch.Tensor]: + """Run the encoder MXQ on the NHWC input and return ordered FPN levels.""" + + encoder_backend = self._require_encoder_backend() + encoder_feed = [strip_runtime_batch(encoder_input)] + validate_runtime_shapes( + encoder_feed, self._backend_input_shapes(encoder_backend), "encoder" + ) + return fpn_from_runtime(encoder_backend(encoder_feed), self.device) + + def _encode_image_onnx(self, encoder_input: np.ndarray) -> list[torch.Tensor]: + """Run the encoder ONNX graph and return ordered FPN levels. + + ``encoder_input`` is the canonical NHWC array :meth:`preprocess` + returns; the exported graph was traced NCHW, so the transpose happens + here rather than changing the framework-independent preprocess + contract. + """ + + encoder_backend = self._require_encoder_backend() + nhwc = strip_runtime_batch(encoder_input) + if nhwc.ndim != 3: + raise ValueError( + f"Expected an NHWC encoder input with a batch of one, got shape " + f"{np.asarray(encoder_input).shape}." + ) + nchw = np.ascontiguousarray(nhwc.transpose(2, 0, 1))[None] + validate_runtime_shapes([nchw], [ENCODER_ONNX_INPUT_SHAPE], "encoder") + outputs = encoder_backend({ENCODER_ONNX_INPUT_NAME: nchw}) + return fpn_from_onnx(outputs, self.device) + + def _decode_prompts_mxq( + self, + weights: dict[str, torch.Tensor], + features: BackboneFeatures, + points: np.ndarray, + labels: np.ndarray, + original_hw: Sequence[int], + ) -> list[np.ndarray]: + """Run the decoder MXQ on the compiled artifact's positional feed.""" + + decoder_tensors = prepare_decoder_tensors( + weights, features, points, labels, original_hw + ) + decoder_feed = build_decoder_runtime_feed( + decoder_tensors, DECODER_RUNTIME_ORDER + ) + decoder_backend = self._require_decoder_backend() + validate_runtime_shapes( + decoder_feed, self._backend_input_shapes(decoder_backend), "decoder" + ) + return decoder_backend(decoder_feed) + + def _decode_prompts_onnx( + self, + weights: dict[str, torch.Tensor], + features: BackboneFeatures, + points: np.ndarray, + labels: np.ndarray, + original_hw: Sequence[int], + ) -> list[np.ndarray]: + """Run the decoder ONNX graph on its five named pre-flattening inputs.""" + + decoder_tensors = prepare_decoder_tensors_onnx( + weights, features, points, labels, original_hw + ) + validate_runtime_shapes( + [decoder_tensors[name] for name in DECODER_ONNX_INPUT_NAMES], + [DECODER_ONNX_INPUT_SHAPES[name] for name in DECODER_ONNX_INPUT_NAMES], + "decoder", + ) + decoder_backend = self._require_decoder_backend() + return decoder_backend(decoder_tensors) + + def postprocess(self, x: Any, **kwargs: Any) -> Results: + """Not supported: ``predict``/``predict_preprocessed`` already return ``Results``.""" + + del x, kwargs + raise NotImplementedError( + "SAM2HieraLarge.predict(...) and predict_preprocessed(...) already return " + "Results; there is no separate postprocess() step." + ) + + def preprocess_with_metadata(self, x: Any) -> Any: + del x + raise NotImplementedError("SAM2HieraLarge does not use letterbox metadata.") + + def set_postprocess_thresholds( + self, conf_thres: float | None = None, iou_thres: float | None = None + ) -> None: + del conf_thres, iou_thres + raise NotImplementedError( + "SAM2HieraLarge does not support configurable postprocess thresholds; " + "mask selection is argmax(iou_predictions)." + ) + + def launch(self) -> None: + """No-op: both backends are already launched during ``__init__``.""" + + self._ensure_open() + + def to(self, device: str | torch.device) -> None: + 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)}.") + if self.weights is not None: + self.weights = { + name: tensor.to(self.device) for name, tensor in self.weights.items() + } + + def cpu(self) -> None: + self.to(device="cpu") + + def gpu(self) -> None: + self.to(device="cuda") + + def cuda(self, device: str | int = 0) -> None: + if isinstance(device, int): + device = f"cuda:{device}" + elif isinstance(device, str) and 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 _require_encoder_backend(self) -> MobilintNPUBackend | ONNXBackend: + if self._encoder_backend is None: + raise RuntimeError("SAM2HieraLarge encoder backend is not initialized.") + return self._encoder_backend + + def _require_decoder_backend(self) -> MobilintNPUBackend | ONNXBackend: + if self._decoder_backend is None: + raise RuntimeError("SAM2HieraLarge decoder backend is not initialized.") + return self._decoder_backend + + def _require_weights(self) -> dict[str, torch.Tensor]: + if self.weights is None: + raise RuntimeError("SAM2HieraLarge prompt-encoder weights are not loaded.") + return self.weights + + @staticmethod + def _backend_input_shapes(backend: MobilintNPUBackend) -> list[tuple[int, ...]]: + """Read the loaded artifact's declared input shapes for a fail-loud shape check. + + A resolved encoder/decoder artifact that does not match + ``DECODER_RUNTIME_ORDER`` (for example one compiled from a different + quantizer revision than the validated reference) must fail here rather + than silently produce wrong masks. ``backend.mxq_model`` is the + slot-zero compatibility handle every mblt_npu backend preserves. + """ + + return [ + tuple(int(dim) for dim in shape) + for shape in backend.mxq_model.get_model_input_shape() + ] + + def _ensure_open(self) -> None: + if getattr(self, "_closed", False): + raise RuntimeError("SAM2HieraLarge is closed.") + + def close(self) -> None: + """Release both backends. Safe to call more than once.""" + + self._close(suppress_errors=False) + + def dispose(self) -> None: + """Compatibility alias for :meth:`close`.""" + + self.close() + + def _close(self, *, suppress_errors: bool) -> None: + if getattr(self, "_closed", False): + return + self._closed = True + first_error: Exception | None = None + for backend in (self._decoder_backend, self._encoder_backend): + if backend is None: + continue + try: + backend.dispose() + except Exception as exc: + if first_error is None: + first_error = exc + if first_error is not None and not suppress_errors: + raise first_error + + def __enter__(self) -> "SAM2HieraLarge": + self._ensure_open() + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool: + del exc_value, traceback + self._close(suppress_errors=exc_type is not None) + return False + + def __del__(self) -> None: + try: + self._close(suppress_errors=True) + except Exception: + pass diff --git a/mblt_vision/models/SAM2HieraLarge.yaml b/mblt_vision/models/SAM2HieraLarge.yaml new file mode 100644 index 0000000..fc31f3e --- /dev/null +++ b/mblt_vision/models/SAM2HieraLarge.yaml @@ -0,0 +1,18 @@ +DEFAULT: HIERA_LARGE +HIERA_LARGE: + file_cfg: + # Documentation only -- SAM2HieraLarge resolves both artifacts itself + # (mblt_vision/mask_generation/sam2.py) because it needs two independent + # MXQ files (encoder + decoder), not the single-artifact file_cfg this + # section otherwise encodes for every other model. + repo_id: mobilint/sam2-hiera-large + encoder_filename: sam2_hiera_large_encoder.mxq + decoder_filename: sam2_hiera_large_decoder.mxq + # ONNX artifacts derive the package-wide same-stem convention + # (sam2_hiera_large_{encoder,decoder}.onnx) and live at the Hub repo + # root (board-agnostic), unlike the board-folder MXQ artifacts above. + target_device: aries-rb + pre_cfg: {} + post_cfg: + task: mask_generation + dataset: sa-v diff --git a/mblt_vision/utils/datasets/__init__.py b/mblt_vision/utils/datasets/__init__.py index b0751a8..95ff8e4 100644 --- a/mblt_vision/utils/datasets/__init__.py +++ b/mblt_vision/utils/datasets/__init__.py @@ -23,6 +23,7 @@ CustomDOTAv1, CustomImageFolder, CustomNYUDepth, + CustomSAV, CustomWiderface, CustomWiderFaceDataset, get_ade20k_loader, @@ -42,6 +43,7 @@ organize_dotav1, organize_imagenet, organize_nyu_depth, + organize_sav, organize_widerface, ) @@ -65,6 +67,7 @@ "CustomDOTAv1", "CustomImageFolder", "CustomNYUDepth", + "CustomSAV", "CustomWiderFaceDataset", "CustomWiderface", "get_ade20k_loader", @@ -81,5 +84,6 @@ "organize_dotav1", "organize_imagenet", "organize_nyu_depth", + "organize_sav", "organize_widerface", ] diff --git a/mblt_vision/utils/datasets/dataloader.py b/mblt_vision/utils/datasets/dataloader.py index d23d719..07caf88 100644 --- a/mblt_vision/utils/datasets/dataloader.py +++ b/mblt_vision/utils/datasets/dataloader.py @@ -274,6 +274,93 @@ def __len__(self) -> int: return len(self.samples) +class CustomSAV(torch.utils.data.Dataset[tuple[np.ndarray, np.ndarray, str, str, str]]): + """SA-V validation dataset of per-object annotated frames. + + One sample is one ``(video_id, object_id, frame_stem)`` triple: an RGB + frame plus that object's boolean ground-truth mask. There is deliberately + no batched loader/collate companion: mask-generation evaluation is + prompt-conditioned and sample-at-a-time with globally area-balanced + sampling, so a DataLoader batch dimension buys nothing. + """ + + def __init__(self, root: str) -> None: + """Validate the organizer's ``images/``/``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"SA-V requires images/ and annotations/ directories under: {root}" + ) + samples: list[tuple[str, str, str, str, str]] = [] + for video_id in sorted(os.listdir(annotation_root)): + video_annotation_dir = os.path.join(annotation_root, video_id) + video_image_dir = os.path.join(image_root, video_id) + if not os.path.isdir(video_annotation_dir): + continue + if not os.path.isdir(video_image_dir): + raise ValueError( + f"SA-V video '{video_id}' has annotations but no frames: {root}" + ) + for object_id in sorted(os.listdir(video_annotation_dir)): + object_dir = os.path.join(video_annotation_dir, object_id) + if not os.path.isdir(object_dir): + continue + for mask_name in sorted(os.listdir(object_dir)): + if not mask_name.endswith(".png"): + continue + stem = mask_name[: -len(".png")] + frame_path = os.path.join(video_image_dir, f"{stem}.jpg") + if not os.path.isfile(frame_path): + raise ValueError( + f"SA-V mask {video_id}/{object_id}/{mask_name} has no " + f"matching frame: {frame_path}" + ) + samples.append( + ( + frame_path, + os.path.join(object_dir, mask_name), + video_id, + object_id, + stem, + ) + ) + if not samples: + raise ValueError(f"SA-V contains no annotated samples: {root}") + self.samples = samples + + def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray, str, str, str]: + """Load an RGB frame and the boolean mask for one annotated object.""" + + frame_path, mask_path, video_id, object_id, stem = self.samples[index] + frame = cv2.imread(frame_path) + if frame is None: + raise FileNotFoundError(f"SA-V frame not found: {frame_path}") + mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + if mask_image is None: + raise FileNotFoundError(f"SA-V mask not found: {mask_path}") + if mask_image.shape != frame.shape[:2]: + raise ValueError( + "SA-V frame and mask shapes must match for " + f"{video_id}/{object_id}/{stem}: frame {frame.shape[:2]}, " + f"mask {mask_image.shape}." + ) + return ( + cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), + mask_image > 0, + video_id, + object_id, + stem, + ) + + def __len__(self) -> int: + """Return the number of annotated (video, object, frame) samples.""" + + return len(self.samples) + + def get_nyu_depth_loader( dataset: CustomNYUDepth, batch_size: int, diff --git a/mblt_vision/utils/datasets/organizer.py b/mblt_vision/utils/datasets/organizer.py index f8c3484..9d7d895 100644 --- a/mblt_vision/utils/datasets/organizer.py +++ b/mblt_vision/utils/datasets/organizer.py @@ -40,6 +40,7 @@ DOTAV1_VALIDATION_SAMPLE_COUNT, IMAGE_SUFFIXES, NYU_DEPTH_VALIDATION_SAMPLE_COUNT, + SAV_VIDEO_ID_PATTERN, _canonicalize_quadrilateral, _path_has_symlink_component, _polygon_has_positive_image_overlap, @@ -60,6 +61,8 @@ } COCO_DOWNLOAD_CONFIG = get_dataset_config("coco")["download"] ADE20K_DOWNLOAD_CONFIG = get_dataset_config("ade20k")["download"] +SAV_DOWNLOAD_CONFIG = get_dataset_config("sa-v")["download"] +SAV_ARCHIVE = SAV_DOWNLOAD_CONFIG["archive"] NYU_DEPTH_URL = ( "https://github.com/ultralytics/assets/releases/download/v0.0.0/nyu-depth.zip" ) @@ -74,6 +77,10 @@ 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"], + # SA-V is deliberately absent: it is user-supplied from Meta's gated + # portal rather than fetched from a URL this package controls, so there is + # no download to pin. Its identity is enforced on content by readiness + # (video/masklet/mask counts plus per-mask geometry and values). } @@ -172,6 +179,7 @@ def _validate_staged_payloads(staged_root: Path, dataset: str) -> None: "dotav1": (staged_root / "images",), "ade20k": (staged_root / "images",), "cityscapes": (staged_root / "images",), + "sa-v": (staged_root / "images",), } for image_root in image_roots.get(dataset, ()): for image_path in image_root.rglob("*"): @@ -195,6 +203,8 @@ def _validate_staged_payloads(staged_root: Path, dataset: str) -> None: _validate_staged_semantic_masks(staged_root, dataset) elif dataset == "dotav1": _validate_staged_dotav1_labels(staged_root) + elif dataset == "sa-v": + _validate_staged_sav_masks(staged_root) def _validate_staged_coco_image_geometry(staged_root: Path) -> None: @@ -1356,6 +1366,227 @@ def organize_nyu_depth( construct_nyu_depth(local_dataset_path, output_dir) +def _resolve_sav_validation_root(dataset_dir: str) -> str: + """Resolves the extracted SA-V validation root containing the official layout. + + Args: + dataset_dir: Directory containing the SA-V validation root or its parent. + + Returns: + Directory containing ``sav_val.txt``, ``JPEGImages_24fps``, and + ``Annotations_6fps``. + + Raises: + ValueError: If no directory with the official SA-V layout is found. + """ + + root = Path(dataset_dir) + candidates = [root, root / "sav_val"] + if root.is_dir(): + candidates.extend(sorted(entry for entry in root.iterdir() if entry.is_dir())) + for candidate in candidates: + if ( + (candidate / "sav_val.txt").is_file() + and (candidate / "JPEGImages_24fps").is_dir() + and (candidate / "Annotations_6fps").is_dir() + ): + return str(candidate) + raise ValueError( + "Unable to locate the SA-V validation layout (sav_val.txt, " + f"JPEGImages_24fps, Annotations_6fps) under: {dataset_dir}." + ) + + +def construct_sav(dataset_dir: str, output_dir: str) -> None: + """Constructs the SA-V validation layout from an extracted dataset directory. + + Keeps only the annotated frames (annotations exist at 6fps, every fourth + 24fps frame), matching the validation-only trimming used by the other + organizers; evaluation can only use annotated frames. + + Args: + dataset_dir: Directory containing the SA-V validation root or its parent. + output_dir: Directory where the organized dataset will be stored. + """ + + output_dir = _validate_dense_output_root( + output_dir, "SA-V", ("images", "annotations", "video_ids.txt") + ) + dataset_root = Path(_resolve_sav_validation_root(dataset_dir)) + video_ids = [ + line.strip() + for line in (dataset_root / "sav_val.txt") + .read_text(encoding="utf-8") + .splitlines() + if line.strip() + ] + if not video_ids or len(set(video_ids)) != len(video_ids): + raise ValueError( + f"SA-V id list is empty or contains duplicates: {dataset_root / 'sav_val.txt'}." + ) + # These ids come from file contents, not a directory listing, so an entry + # such as `../../escape` would otherwise resolve outside the staging tree + # and copy attacker-controlled files there before organization failed. + # Rejected up front, before any makedirs/copy, and mirroring the same + # containment guard COCO applies to its JSON-declared file names. + invalid_ids = [ + video_id + for video_id in video_ids + if SAV_VIDEO_ID_PATTERN.fullmatch(video_id) is None + ] + if invalid_ids: + raise ValueError( + f"SA-V id list contains unsupported video ids: {sorted(invalid_ids)[:5]}. " + "Ids must match the official `sav_<6 digits>` format." + ) + print(f"Constructing SA-V validation dataset from {dataset_dir} to {output_dir}") + + output_parent_dir = os.path.dirname(output_dir) + os.makedirs(output_parent_dir, exist_ok=True) + total_masks = 0 + with TemporaryDirectory( + dir=output_parent_dir, prefix=".sa-v-staging-" + ) as staging_dir: + 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) + staging_root = Path(staging_dir).resolve() + for video_id in sorted(video_ids): + # Belt-and-braces containment on top of the id-pattern check above: + # every staged path must resolve inside the staging tree, so no + # future relaxation of that pattern can reintroduce an escape. + for staged_root in (staged_image_dir, staged_annotation_dir): + staged_video_path = Path(staged_root, video_id).resolve() + if not staged_video_path.is_relative_to(staging_root): + raise ValueError( + f"SA-V video id '{video_id}' escapes the staging directory." + ) + source_frame_dir = dataset_root / "JPEGImages_24fps" / video_id + source_annotation_dir = dataset_root / "Annotations_6fps" / video_id + if not source_frame_dir.is_dir() or not source_annotation_dir.is_dir(): + raise ValueError( + f"SA-V video '{video_id}' is missing its frame or annotation " + f"directory under {dataset_root}." + ) + annotated_stems: set[str] = set() + for object_dir in sorted(source_annotation_dir.iterdir()): + if not object_dir.is_dir(): + continue + staged_object_dir = os.path.join( + staged_annotation_dir, video_id, object_dir.name + ) + os.makedirs(staged_object_dir) + for mask_path in sorted(object_dir.glob("*.png")): + shutil.copy2( + mask_path, os.path.join(staged_object_dir, mask_path.name) + ) + annotated_stems.add(mask_path.stem) + total_masks += 1 + if not annotated_stems: + raise ValueError( + f"SA-V video '{video_id}' has no annotation masks under " + f"{source_annotation_dir}." + ) + staged_video_image_dir = os.path.join(staged_image_dir, video_id) + os.makedirs(staged_video_image_dir) + for stem in sorted(annotated_stems): + frame_path = source_frame_dir / f"{stem}.jpg" + if not frame_path.is_file(): + raise ValueError( + f"SA-V video '{video_id}' is missing annotated frame " + f"{stem}.jpg under {source_frame_dir}." + ) + shutil.copy2( + frame_path, os.path.join(staged_video_image_dir, frame_path.name) + ) + + staged_ids_path = os.path.join(staging_dir, "video_ids.txt") + with open(staged_ids_path, "w", encoding="utf-8") as ids_file: + ids_file.write("\n".join(sorted(video_ids)) + "\n") + + _validate_staged_dataset(staging_dir, "sa-v", ("mask_generation",)) + + replacements = ( + (staged_image_dir, os.path.join(output_dir, "images")), + (staged_annotation_dir, os.path.join(output_dir, "annotations")), + (staged_ids_path, os.path.join(output_dir, "video_ids.txt")), + ) + _replace_staged_directories(replacements, output_parent_dir, ".sa-v-backup-") + print( + f"Constructed SA-V validation dataset with {len(video_ids)} videos " + f"and {total_masks} annotation masks" + ) + + +def _validate_staged_sav_masks(staged_root: Path) -> None: + """Decode every staged SA-V mask before a structurally valid cache is replaced.""" + + annotation_root = staged_root / "annotations" + frame_shapes: dict[tuple[str, str], tuple[int, int]] = {} + for mask_path in sorted(annotation_root.rglob("*.png")): + video_id = mask_path.parent.parent.name + try: + with Image.open(mask_path) as mask_image: + mask = np.asarray(mask_image) + except OSError as exc: + raise ValueError( + f"Staged SA-V mask is unreadable: {mask_path}: {exc}." + ) from exc + # Counting unique values alone would accept a two-valued mask such as + # {1, 2}, which CustomSAV's `> 0` binarization turns entirely into + # foreground; require every non-zero value to be a single object ID. + if mask.ndim != 2 or len(set(np.unique(mask).tolist()) - {0}) > 1: + raise ValueError( + f"Staged SA-V mask must be a single-object binary map: {mask_path}." + ) + key = (video_id, mask_path.stem) + if key not in frame_shapes: + frame_path = staged_root / "images" / video_id / f"{mask_path.stem}.jpg" + with Image.open(frame_path) as frame_image: + frame_shapes[key] = (frame_image.height, frame_image.width) + if mask.shape != frame_shapes[key]: + raise ValueError( + "Staged SA-V mask and frame shapes must match: " + f"mask {mask.shape}, frame {frame_shapes[key]}: {mask_path}." + ) + + +def organize_sav( + dataset_path: str, + output_dir: str | None = None, +) -> None: + """Organizes SA-V validation from a manually downloaded archive or directory. + + ``dataset_path`` is required because SA-V is distributed through Meta's + form-gated portal and is not mirrored by this package, so there is no + default source to download (unlike the other organizers). + + Args: + dataset_path: Path to the manually downloaded ``sav_val.tar`` archive + or its extracted dataset directory. See + https://github.com/facebookresearch/sam2/blob/main/sav_dataset/README.md + for the official layout. + output_dir: Directory to store the organized dataset. Defaults to the + resolved Mobilint cache directory. + """ + + output_dir = _resolve_organizer_output_dir(output_dir, "sa-v") + output_dir = _validate_dense_output_root( + output_dir, "SA-V", ("images", "annotations", "video_ids.txt") + ) + with TemporaryDirectory() as temp_dir: + local_dataset_path = _resolve_source(dataset_path, temp_dir) + if local_dataset_path.endswith((".tar", ".tar.gz", ".tgz")): + print("Unpacking SA-V files to temporary directory...") + _safe_unpack_archive(local_dataset_path, temp_dir) + print("Unpacking completed") + construct_sav(temp_dir, output_dir) + return + + construct_sav(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.""" diff --git a/mblt_vision/utils/datasets/readiness.py b/mblt_vision/utils/datasets/readiness.py index 5e62977..ea85728 100644 --- a/mblt_vision/utils/datasets/readiness.py +++ b/mblt_vision/utils/datasets/readiness.py @@ -31,6 +31,16 @@ NYU_DEPTH_VALIDATION_SAMPLE_COUNT = 654 ADE20K_VALIDATION_SAMPLE_COUNT = 2000 CITYSCAPES_VALIDATION_SAMPLE_COUNT = 500 +SAV_VALIDATION_VIDEO_COUNT = 155 +SAV_VALIDATION_MASKLET_COUNT = 293 +# Total per-object annotated masks across every masklet, measured from the +# organized official `sav_val.tar` (see `datasets/sa-v.yaml`). The archive is +# user-supplied from Meta's gated portal rather than downloaded from a URL this +# package controls, so these counts -- not an archive checksum -- are what +# establish dataset identity. Masklet and video counts alone accept a truncated +# source that keeps every masklet but only a few of its annotated frames, which +# would silently change the evaluation corpus; this pins the inventory. +SAV_VALIDATION_MASK_COUNT = 31967 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}") @@ -39,6 +49,9 @@ CITYSCAPES_SAMPLE_ID_PATTERN = re.compile( r"^(?P[A-Za-z][A-Za-z0-9-]*)_\d{6}_\d{6}$" ) +SAV_VIDEO_ID_PATTERN = re.compile(r"^sav_\d{6}$") +SAV_OBJECT_ID_PATTERN = re.compile(r"^\d{3}$") +SAV_FRAME_STEM_PATTERN = re.compile(r"^\d{5}$") CITYSCAPES_VALIDATION_CITY_COUNTS = {"frankfurt": 267, "lindau": 59, "munster": 174} IMAGENET_SYNSET_ORDER = tuple( files("mblt_vision.datasets") @@ -1011,6 +1024,123 @@ def dense_dataset_ready(data_path: str | Path, dataset: str) -> bool: return False +def _sav_mask_is_valid(mask_path: Path, frame_shape: tuple[int, int]) -> bool: + """Return whether one cached SA-V mask matches its frame and is single-object. + + Both properties are checked for every mask. The value constraint is + discharged from the image mode when the file is bilevel: a 1-bit PNG can + only encode two values with zero as background, so it cannot hold the + ``{1, 2}`` case that ``CustomSAV``'s ``> 0`` binarization would turn + entirely into foreground. Any other mode is decoded and checked directly. + The official split is bilevel throughout, so this reads headers rather + than decoding ~32k full-resolution masks on every readiness call. + """ + + try: + with Image.open(mask_path) as mask_image: + if (mask_image.height, mask_image.width) != frame_shape: + return False + if mask_image.mode == "1": + return True + mask = np.asarray(mask_image) + except OSError: + return False + return mask.ndim == 2 and len(set(np.unique(mask).tolist()) - {0}) <= 1 + + +def _sav_ready(root: Path) -> bool: + """Return whether an organized SA-V validation split is complete. + + Validates the id list, per-video image/annotation directory identity, the + total masklet count, and mask/frame stem consistency. Decodes one + frame/mask pair per video (155 decodes) rather than the full multi- + thousand-file corpus, which would make this readiness probe too slow for + its per-`val`-invocation call site; the organizer's staged payload + validation decodes every mask once at install time instead. + """ + + if _path_has_symlink_component(root) or not root.is_dir(): + return False + video_ids_path = root / "video_ids.txt" + if video_ids_path.is_symlink() or not video_ids_path.is_file(): + return False + try: + video_ids = [ + line.strip() + for line in video_ids_path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + except (OSError, UnicodeDecodeError): + return False + if len(video_ids) != SAV_VALIDATION_VIDEO_COUNT or len(set(video_ids)) != len( + video_ids + ): + return False + if any(SAV_VIDEO_ID_PATTERN.fullmatch(video_id) is None for video_id in video_ids): + return False + + image_root = root / "images" + annotation_root = root / "annotations" + for directory in (image_root, annotation_root): + if directory.is_symlink() or not directory.is_dir(): + return False + entries = list(directory.iterdir()) + if any(entry.is_symlink() for entry in entries): + return False + if {entry.name for entry in entries if entry.is_dir()} != set(video_ids): + return False + + total_masklets = 0 + total_masks = 0 + for video_id in video_ids: + images = _files_by_stem(image_root / video_id, {".jpg"}, reject_symlinks=True) + if not images: + return False + if any(SAV_FRAME_STEM_PATTERN.fullmatch(stem) is None for stem in images): + return False + + annotation_dir = annotation_root / video_id + object_dirs = sorted( + entry for entry in annotation_dir.iterdir() if entry.is_dir() + ) + if not object_dirs or any(entry.is_symlink() for entry in object_dirs): + return False + if any( + SAV_OBJECT_ID_PATTERN.fullmatch(entry.name) is None for entry in object_dirs + ): + return False + total_masklets += len(object_dirs) + + # Frame geometry for every annotated stem, read from image headers so + # the per-mask comparison below covers the whole video rather than only + # its first mask. + try: + frame_shapes: dict[str, tuple[int, int]] = {} + for stem, image_path in images.items(): + with Image.open(image_path) as image: + frame_shapes[stem] = (image.height, image.width) + # One full decode per video still catches a truncated JPEG, which a + # header-only read cannot see. + with Image.open(images[sorted(images)[0]]) as image: + image.load() + except OSError: + return False + + for object_dir in object_dirs: + masks = _files_by_stem(object_dir, {".png"}, reject_symlinks=True) + if not masks or not set(masks) <= set(images): + return False + total_masks += len(masks) + for stem, mask_path in masks.items(): + if not _sav_mask_is_valid(mask_path, frame_shapes[stem]): + return False + + return ( + total_masklets == SAV_VALIDATION_MASKLET_COUNT + and total_masks == SAV_VALIDATION_MASK_COUNT + ) + + 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. @@ -1033,6 +1163,7 @@ def dataset_ready(data_path: str | Path, task: str, dataset: str | None = None) "face_detection": "widerface", "obb": "dotav1", "depth_estimation": "nyu-depth", + "mask_generation": "sa-v", }.get(normalized_task) normalized_dataset = (dataset or expected_dataset or "").lower() @@ -1054,4 +1185,6 @@ def dataset_ready(data_path: str | Path, task: str, dataset: str | None = None) return _dotav1_ready(root) if normalized_task == "depth_estimation": return dense_dataset_ready(root, normalized_dataset) + if normalized_task == "mask_generation": + return _sav_ready(root) return False diff --git a/mblt_vision/utils/evaluation/__init__.py b/mblt_vision/utils/evaluation/__init__.py index 21b2916..e3d5b80 100644 --- a/mblt_vision/utils/evaluation/__init__.py +++ b/mblt_vision/utils/evaluation/__init__.py @@ -23,6 +23,12 @@ calculate_nyu_depth_metrics, eval_nyu_depth, ) +from .eval_sav import ( + SAVMetricAccumulator, + SAVResult, + calculate_sav_sample_ious, + eval_sav, +) from .eval_widerface import WiderFaceResult, eval_widerface __all__: list[str] = [ @@ -46,6 +52,10 @@ "NYUDepthMetricAccumulator", "calculate_nyu_depth_metrics", "eval_nyu_depth", + "SAVResult", + "SAVMetricAccumulator", + "calculate_sav_sample_ious", + "eval_sav", "WiderFaceResult", "eval_widerface", ] diff --git a/mblt_vision/utils/evaluation/eval_sav.py b/mblt_vision/utils/evaluation/eval_sav.py new file mode 100644 index 0000000..b12c12e --- /dev/null +++ b/mblt_vision/utils/evaluation/eval_sav.py @@ -0,0 +1,467 @@ +"""SA-V point-prompted mask generation evaluation. + +Protocol ported from the validated ``sam2-mxq-pipeline`` reference (its +Aries2 run measured MXQ mIoU 0.7757 on 200 sav_train samples): deterministic +area-balanced sampling of (video, object, frame) masklets, synthetic point +prompts derived from the ground-truth mask (distance-transform peak plus +optional negative/second-positive points), and per-candidate mask IoU against +the ground truth. The primary metric is the mean IoU of each sample's +own-selected candidate (``argmax`` of the model's predicted IoUs). +""" + +from __future__ import annotations + +import math +import random +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Any, Protocol + +import cv2 +import numpy as np +from tqdm import tqdm + +from ..datasets import CustomSAV + + +class PromptedPrediction(Protocol): + """The two fields ``eval_sav`` actually reads off a ``predict()`` result. + + A real ``SAM2HieraLarge.predict()`` call returns a full ``Results``, but + only ``masks`` and ``selected`` are used here -- narrower than ``Results`` + so the lightweight test doubles in ``tests/test_eval_sav.py`` (structural + stand-ins, not real ``Results`` instances) satisfy it too. + """ + + # Read-only (never written), so declared as properties rather than plain + # attributes: a Protocol's plain attributes are matched invariantly, which + # would reject a real implementer's narrower concrete type (e.g. a mock's + # `selected: int` against `int | None`) even though it satisfies every + # actual read here. + @property + def masks(self) -> Any: ... + @property + def selected(self) -> int | None: ... + + +class PointPromptedEngine(Protocol): + """Structural contract ``eval_sav`` needs from a mask generation engine. + + ``SAM2HieraLarge`` satisfies this, but so do the lightweight test doubles + in ``tests/test_eval_sav.py`` (no real backend or download); the concrete + class is deliberately not required here. + """ + + post_cfg: dict[str, Any] + + def predict(self, image: Any, points: Any, labels: Any, /) -> PromptedPrediction: + """Positional-only: every real caller here invokes this positionally, + and implementations use varying parameter names (``image`` vs. + ``frame``).""" + ... + + +# Relative-mask-area bins used to balance sampling (reference dataset.py). +AREA_BINS = ((0.0, 0.005), (0.005, 0.02), (0.02, 0.08), (0.08, 1.01)) +CANDIDATES_PER_PROMPT = 3 +FRAMES_CONSIDERED_PER_MASKLET = 6 + + +@dataclass(frozen=True) +class SAVResult: + """Point-prompted SA-V mask generation metrics.""" + + miou: float + miou_ci95: float + miou_best_of_3: float + num_samples: int + distinct_videos: int + + @property + def primary_score(self) -> float: + """Return the own-selection mean IoU.""" + + return self.miou + + @property + def secondary_score(self) -> float: + """Return the best-of-3 (oracle) mean IoU.""" + + return self.miou_best_of_3 + + +def mask_iou(left: np.ndarray, right: np.ndarray) -> float: + """Return the IoU of two boolean masks; both-empty counts as 1.0.""" + + intersection = np.logical_and(left, right).sum(dtype=np.int64) + union = np.logical_or(left, right).sum(dtype=np.int64) + return float(intersection) / float(union) if union else 1.0 + + +def calculate_sav_sample_ious( + candidate_masks: np.ndarray, gt_mask: np.ndarray +) -> list[float]: + """Return per-candidate IoUs of predicted binary masks against ground truth.""" + + candidates = np.asarray(candidate_masks) + gt = np.asarray(gt_mask).astype(bool) + if candidates.ndim != 3: + raise ValueError( + f"Expected candidate masks shaped (N, H, W), got {candidates.shape}." + ) + if candidates.shape[1:] != gt.shape: + raise ValueError( + "Candidate masks and ground truth shapes must match: " + f"candidates {candidates.shape[1:]}, ground truth {gt.shape}." + ) + # `astype(bool)` would treat logits, probabilities, and NaN as foreground + # and report a plausible but meaningless IoU, so candidates must already be + # binarized in one of the documented encodings. The permitted values are + # enumerated per dtype rather than merely required to be "a single positive + # value": that weaker rule still admits a probability map such as + # {0.0, 0.5}, or a degenerate uniform 0.5 candidate that silently becomes + # all-foreground. This is deliberately stricter than the SA-V ground-truth + # mask check, where any single positive value is a legitimate object ID. + for index, candidate in enumerate(candidates): + if candidate.dtype == np.bool_: + continue + if np.issubdtype(candidate.dtype, np.floating): + if not bool(np.isfinite(candidate).all()): + raise ValueError( + f"Candidate mask {index} must be finite; got NaN or infinity." + ) + permitted: tuple[set[float], ...] = ({0.0, 1.0},) + elif np.issubdtype(candidate.dtype, np.integer): + permitted = ({0.0, 1.0}, {0.0, 255.0}) + else: + raise ValueError( + f"Candidate mask {index} has unsupported dtype {candidate.dtype}; " + "predict() must return a boolean, integer, or floating binary map." + ) + values = {float(value) for value in np.unique(candidate).tolist()} + if not any(values <= allowed for allowed in permitted): + raise ValueError( + f"Candidate mask {index} must be a binary map encoded as bool, " + f"{{0, 1}}, or {{0, 255}} for integers, or {{0.0, 1.0}} for floats; " + f"got values {sorted(values)[:5]}. predict() must return binarized " + "masks, not logits or probabilities." + ) + return [mask_iou(candidate.astype(bool), gt) for candidate in candidates] + + +def _ci95(values: Sequence[float]) -> float: + """Return the 95% confidence half-width of the mean.""" + + if len(values) < 2: + return 0.0 + array = np.asarray(values, dtype=np.float64) + return float(1.96 * array.std(ddof=1) / math.sqrt(array.size)) + + +def _distance_peak(mask: np.ndarray) -> tuple[float, float] | None: + """Return the interior point of the mask farthest from its boundary.""" + + distance = cv2.distanceTransform(mask.astype(np.uint8) * 255, cv2.DIST_L2, 5) + if float(distance.max()) < 1.0: + return None + y, x = np.unravel_index(int(distance.argmax()), distance.shape) + return float(x), float(y) + + +def _second_positive( + mask: np.ndarray, first: tuple[float, float], rng: random.Random +) -> tuple[float, float] | None: + """Sample a second positive point far from the first.""" + + ys, xs = np.where(mask) + if not len(xs): + return None + best: tuple[float, float] | None = None + best_distance = -1.0 + for _ in range(64): + index = rng.randint(0, len(xs) - 1) + point = float(xs[index]), float(ys[index]) + distance = (point[0] - first[0]) ** 2 + (point[1] - first[1]) ** 2 + if distance > best_distance: + best, best_distance = point, distance + return best if best_distance > 9 else None + + +def _negative_point( + mask: np.ndarray, height: int, width: int, rng: random.Random +) -> tuple[float, float] | None: + """Sample a background point outside the dilated mask, near it when possible.""" + + ys, xs = np.where(mask) + y0, y1, x0, x1 = int(ys.min()), int(ys.max()) + 1, int(xs.min()), int(xs.max()) + 1 + dilated = cv2.dilate( + mask.astype(np.uint8), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9)) + ).astype(bool) + pad = max(8, int(0.15 * max(y1 - y0, x1 - x0))) + bounds = ( + max(0, y0 - pad), + min(height, y1 + pad), + max(0, x0 - pad), + min(width, x1 + pad), + ) + for global_search in (False, True): + for _ in range(200): + if global_search: + x, y = rng.uniform(0, width - 1), rng.uniform(0, height - 1) + else: + y_min, y_max, x_min, x_max = bounds + x, y = rng.uniform(x_min, x_max - 1), rng.uniform(y_min, y_max - 1) + if not dilated[int(y), int(x)]: + return x, y + return None + + +def build_prompt( + mask: np.ndarray, rng: random.Random, num_points: int +) -> tuple[np.ndarray, np.ndarray] | None: + """Build a deterministic point prompt from a ground-truth mask. + + 1 point: positive distance-transform peak. 2 points: peak + negative. + 3 points: peak + far second positive + negative. Returns ``None`` when a + required point cannot be constructed for this mask. + """ + + height, width = mask.shape + first = _distance_peak(mask) + if first is None: + return None + if num_points == 1: + return np.asarray([first], np.float32), np.asarray([1], np.int64) + negative = _negative_point(mask, height, width, rng) + if negative is None: + return None + if num_points == 2: + return np.asarray([first, negative], np.float32), np.asarray([1, 0], np.int64) + if num_points != 3: + raise ValueError(f"num_points must be 1, 2, or 3; got {num_points}") + second = _second_positive(mask, first, rng) + if second is None: + return None + return ( + np.asarray([first, second, negative], np.float32), + np.asarray([1, 1, 0], np.int64), + ) + + +def _area_bin(mask: np.ndarray) -> int: + """Return the relative-area bin index of a boolean mask.""" + + fraction = int(mask.sum()) / float(mask.size) + return next( + ( + index + for index, (low, high) in enumerate(AREA_BINS) + if low <= fraction < high + ), + len(AREA_BINS) - 1, + ) + + +def _load_gt_mask(dataset: CustomSAV, index: int) -> np.ndarray | None: + """Load only the boolean ground-truth mask for one dataset sample.""" + + mask_path = dataset.samples[index][1] + mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + if mask_image is None: + return None + return mask_image > 0 + + +def iter_selected_samples( + dataset: CustomSAV, + *, + seed: int, + per_video: int, + min_mask_area: int, +) -> Iterator[int]: + """Yield dataset indices with deterministic, globally area-balanced sampling. + + Videos are visited in seed-shuffled order; per (video, object) masklet up + to six evenly spaced annotated frames are considered; per video up to + ``per_video`` samples are greedily chosen preferring the globally + least-represented relative-area bin (reference ``_select_masks``). + """ + + by_video: dict[str, dict[str, list[int]]] = {} + for index, (_, _, video_id, object_id, _) in enumerate(dataset.samples): + by_video.setdefault(video_id, {}).setdefault(object_id, []).append(index) + + video_ids = sorted(by_video) + random.Random(seed).shuffle(video_ids) + rng = random.Random(seed + 7) + bin_counts = [0] * len(AREA_BINS) + + for video_id in video_ids: + candidates: list[tuple[int, int]] = [] + for object_id in sorted(by_video[video_id]): + frame_indices = by_video[video_id][object_id] + count = len(frame_indices) + positions = sorted( + { + int( + round( + position * (count - 1) / (FRAMES_CONSIDERED_PER_MASKLET - 1) + ) + ) + for position in range(FRAMES_CONSIDERED_PER_MASKLET) + } + ) + for position in positions: + index = frame_indices[position] + mask = _load_gt_mask(dataset, index) + if mask is None or int(mask.sum()) < min_mask_area: + continue + candidates.append((index, _area_bin(mask))) + rng.shuffle(candidates) + selected: list[int] = [] + used: set[int] = set() + while len(selected) < per_video: + remaining = [item for item in candidates if item[0] not in used] + if not remaining: + break + target_bin = min( + {item[1] for item in remaining}, key=lambda value: bin_counts[value] + ) + index, chosen_bin = next( + item for item in remaining if item[1] == target_bin + ) + used.add(index) + bin_counts[chosen_bin] += 1 + selected.append(index) + yield from selected + + +class SAVMetricAccumulator: + """Accumulates per-sample candidate IoUs into SA-V summary metrics.""" + + def __init__(self) -> None: + self._own_selection_ious: list[float] = [] + self._best_ious: list[float] = [] + self._videos: set[str] = set() + + @property + def count(self) -> int: + """Return the number of accumulated samples.""" + + return len(self._own_selection_ious) + + def update( + self, candidate_ious: Sequence[float], selected: int, video_id: str + ) -> None: + """Record one sample's per-candidate IoUs and the model's selection.""" + + if len(candidate_ious) != CANDIDATES_PER_PROMPT: + raise ValueError( + f"Expected {CANDIDATES_PER_PROMPT} candidate IoUs, " + f"got {len(candidate_ious)}." + ) + if not 0 <= selected < len(candidate_ious): + raise ValueError( + f"Selected candidate index {selected} is out of range for " + f"{len(candidate_ious)} candidates." + ) + self._own_selection_ious.append(float(candidate_ious[selected])) + self._best_ious.append(float(max(candidate_ious))) + self._videos.add(video_id) + + def result(self) -> SAVResult: + """Return the accumulated SA-V metrics.""" + + if not self._own_selection_ious: + raise ValueError("SA-V evaluation received no valid samples.") + return SAVResult( + miou=float(np.mean(self._own_selection_ious)), + miou_ci95=_ci95(self._own_selection_ious), + miou_best_of_3=float(np.mean(self._best_ious)), + num_samples=len(self._own_selection_ious), + distinct_videos=len(self._videos), + ) + + +def eval_sav( + model: PointPromptedEngine, + data_path: str, + num_samples: int = 200, + num_points: int = 1, + seed: int = 0, + per_video: int = 4, + min_mask_area: int = 400, +) -> SAVResult: + """Evaluate a point-prompted mask generation model on organized SA-V val. + + Args: + model: Loaded mask generation engine exposing ``predict``. + data_path: Organized SA-V dataset root. + num_samples: Number of prompted samples to evaluate. + num_points: Points per prompt (1, 2, or 3). + seed: Sampling and prompt-synthesis seed. + per_video: Maximum samples drawn from one video. + min_mask_area: Minimum ground-truth mask area in pixels. + + Returns: + Accumulated :class:`SAVResult` metrics. + + Raises: + ValueError: If the model taxonomy mismatches, a prediction violates + the three-candidate contract, or too few valid samples exist. + """ + + dataset_name = model.post_cfg.get("dataset") + if not isinstance(dataset_name, str) or dataset_name.lower() != "sa-v": + raise ValueError( + "SA-V evaluation requires model post_cfg.dataset to be 'sa-v', " + f"got {dataset_name!r}." + ) + if num_points not in (1, 2, 3): + raise ValueError(f"num_points must be 1, 2, or 3; got {num_points}.") + if num_samples < 1: + raise ValueError(f"num_samples must be positive, got {num_samples}.") + + dataset = CustomSAV(data_path) + prompt_rng = random.Random(seed + 29) + accumulator = SAVMetricAccumulator() + progress = tqdm(total=num_samples, desc="Evaluating SA-V") + try: + for index in iter_selected_samples( + dataset, seed=seed, per_video=per_video, min_mask_area=min_mask_area + ): + if accumulator.count >= num_samples: + break + frame, gt_mask, video_id, _, _ = dataset[index] + prompt = build_prompt(gt_mask, prompt_rng, num_points) + if prompt is None: + continue + points, labels = prompt + result = model.predict(frame, points, labels) + candidate_ious = calculate_sav_sample_ious( + np.asarray(result.masks), gt_mask + ) + # Results.selected is Optional at the shared-class level (most tasks + # never set it), but a mask generation engine must populate it. Raise + # rather than assert: this is a public evaluator documented to report + # prediction-contract violations as ValueError, and an assert would + # vanish under `python -O` and resurface as an opaque comparison + # TypeError inside accumulator.update(). + if result.selected is None: + raise ValueError( + "Mask generation prediction is missing a selected mask index; " + "predict() must report which candidate it chose." + ) + accumulator.update(candidate_ious, result.selected, video_id) + progress.update(1) + finally: + progress.close() + + outcome = accumulator.result() + if outcome.num_samples < num_samples: + raise ValueError( + f"Requested {num_samples} SA-V evaluation samples but only " + f"{outcome.num_samples} valid samples were available; lower " + "--num-samples or relax the sampling constraints." + ) + return outcome diff --git a/mblt_vision/utils/postprocess/yolo_anchorless_post.py b/mblt_vision/utils/postprocess/yolo_anchorless_post.py index 8b90d31..7011dfe 100644 --- a/mblt_vision/utils/postprocess/yolo_anchorless_post.py +++ b/mblt_vision/utils/postprocess/yolo_anchorless_post.py @@ -566,7 +566,7 @@ class YOLOAnchorlessPosePost(YOLOPosePostMixin, YOLOAnchorlessDetectionPost): def extract_final_outputs( self, x: TensorLike | ListTensorLike - ) -> tuple[list[torch.Tensor] | None, torch.Tensor | None]: + ) -> tuple[list[torch.Tensor] | torch.Tensor | None, torch.Tensor | None]: """Accept QBCompiler's decode-enabled candidate-first pose output. Decode-enabled MXQs emit ``(B, anchors, 5 + keypoints)`` containing @@ -574,6 +574,11 @@ def extract_final_outputs( single-class label column and convert the boxes once. QBCompiler leaves visibility as near-zero logits in this layout, so normalize them as the split-head pose decoder does. + + When this input shape doesn't match, falls through to the base class, + whose own already-decoded-detections path can return a bare tensor + instead of a per-image list -- hence the wider return type than this + override's own ``list[torch.Tensor] | None`` branch above. """ if self.e2e and isinstance(x, (list, tuple)) and len(x) == 1: value = x[0] diff --git a/mblt_vision/utils/postprocess/yolo_dflfree_post.py b/mblt_vision/utils/postprocess/yolo_dflfree_post.py index 8cc52f2..0ba18d3 100644 --- a/mblt_vision/utils/postprocess/yolo_dflfree_post.py +++ b/mblt_vision/utils/postprocess/yolo_dflfree_post.py @@ -558,8 +558,14 @@ class YOLODFLFreePosePost(YOLOPosePostMixin, YOLODFLFreeDetectionPost): def extract_final_outputs( self, x: TensorLike | ListTensorLike - ) -> tuple[list[torch.Tensor] | None, torch.Tensor | None]: - """Accept YOLO26's decode-enabled score, xyxy, and keypoint outputs.""" + ) -> tuple[list[torch.Tensor] | torch.Tensor | None, torch.Tensor | None]: + """Accept YOLO26's decode-enabled score, xyxy, and keypoint outputs. + + When this input shape doesn't match, falls through to the base class, + whose own already-decoded-detections path can return a bare tensor + instead of a per-image list -- hence the wider return type than this + override's own ``list[torch.Tensor] | None`` branch above. + """ if self.e2e and isinstance(x, (list, tuple)) and len(x) == 4: tensors = [ value if isinstance(value, torch.Tensor) else torch.as_tensor(value) diff --git a/mblt_vision/utils/preprocess/letterbox.py b/mblt_vision/utils/preprocess/letterbox.py index 4d2c1fe..4c7e7ea 100644 --- a/mblt_vision/utils/preprocess/letterbox.py +++ b/mblt_vision/utils/preprocess/letterbox.py @@ -37,6 +37,17 @@ def _apply_letterbox( image, (resized_width, resized_height), interpolation=interpolation ) top, bottom, left, right = geometry.borders + # cv2's border value fills only the first channel and zeros the rest when + # given a bare scalar on a multi-channel image (cv::Scalar's single-value + # constructor), so broadcast a scalar to one value per channel rather than + # pass it through as-is -- a no-op for the already-per-channel tuple caller + # and for the single-channel semantic-mask caller. + channels = image.shape[2] if image.ndim == 3 else 1 + border_value: cv2.typing.Scalar = ( + (float(padding_value),) * channels + if isinstance(padding_value, int) + else tuple(float(component) for component in padding_value) + ) image = cv2.copyMakeBorder( image, top, @@ -44,7 +55,7 @@ def _apply_letterbox( left, right, cv2.BORDER_CONSTANT, - value=padding_value, + value=border_value, ) return image, geometry.ratio_pad diff --git a/mblt_vision/utils/results.py b/mblt_vision/utils/results.py index 884a333..f838401 100644 --- a/mblt_vision/utils/results.py +++ b/mblt_vision/utils/results.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from pathlib import Path -from typing import cast +from typing import Any, cast import cv2 import numpy as np @@ -42,6 +42,7 @@ RADIUS = 5 # circle radius ALPHA = 0.3 # alpha for overlay DENSE_OVERLAY_ALPHA = 0.6 +MASK_GENERATION_COLOR = (255, 144, 30) # BGR dodgerblue, matches the reference overlay class Results: @@ -51,7 +52,7 @@ def __init__( self, pre_cfg: dict, post_cfg: dict, - output: TensorLike | ListTensorLike | NestedListTensorLike, + output: TensorLike | ListTensorLike | NestedListTensorLike | dict[str, Any], **kwargs, ) -> None: """ @@ -59,7 +60,7 @@ def __init__( Args: pre_cfg (dict): Preprocessing configuration. post_cfg (dict): Postprocessing configuration. - output (TensorLike | ListTensorLike | NestedListTensorLike): Raw model output. + output: Raw model output, including a dictionary for mask generation. **kwargs: Additional arguments. """ self.pre_cfg = pre_cfg @@ -71,12 +72,20 @@ def __init__( 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.output: ( + TensorLike | ListTensorLike | NestedListTensorLike | dict[str, Any] | 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.masks: np.ndarray | None = None + self.iou_predictions: np.ndarray | None = None + self.low_res_masks: np.ndarray | None = None + self.points: np.ndarray | None = None + self.point_labels: np.ndarray | None = None + self.selected: int | None = None self.set_output(output) def _read_image( @@ -126,7 +135,8 @@ def _save_image(save_path: str | Path, image: np.ndarray) -> None: raise OSError(f"Failed to write result image: {path}") def set_output( - self, output: TensorLike | ListTensorLike | NestedListTensorLike + self, + output: TensorLike | ListTensorLike | NestedListTensorLike | dict[str, Any], ) -> None: """ Sets variables from the raw model output based on the task. @@ -140,6 +150,12 @@ def set_output( self.mask = None self.depth = None self.semantic_mask = None + self.masks = None + self.iou_predictions = None + self.low_res_masks = None + self.points = None + self.point_labels = None + self.selected = None if self.task == "image_classification": if not isinstance(output, (np.ndarray, torch.Tensor)): raise TypeError( @@ -229,6 +245,26 @@ def set_output( raise TypeError( f"Expected tensor semantic output for task {self.task}, got {type(output)}." ) + elif self.task == "mask_generation": + if not isinstance(output, dict): + raise TypeError( + f"Expected dict output for task {self.task}, got {type(output).__name__}." + ) + required_keys = {"masks", "iou_predictions"} + missing_keys = required_keys - output.keys() + if missing_keys: + raise ValueError( + f"mask_generation output is missing key(s): {sorted(missing_keys)}." + ) + # mask_generation always hands numpy arrays here (SAM2HieraLarge's + # output dict), unlike the torch/numpy-either TensorLike fields + # above -- cast to what self.masks etc. are actually declared as. + self.masks = cast(np.ndarray, output["masks"]) + self.iou_predictions = cast(np.ndarray, output["iou_predictions"]) + self.low_res_masks = cast(np.ndarray | None, output.get("low_res_masks")) + self.points = cast(np.ndarray | None, output.get("points")) + self.point_labels = cast(np.ndarray | None, output.get("point_labels")) + self.selected = output.get("selected") else: raise NotImplementedError( f"Task {self.task} is not supported for plotting results." @@ -268,6 +304,8 @@ def plot( return self._plot_pose_estimation(source_path, save_path, **kwargs) elif self.task == "obb": return self._plot_obb(source_path, save_path, **kwargs) + elif self.task == "mask_generation": + return self._plot_mask_generation(source_path, save_path, **kwargs) else: raise NotImplementedError( f"Task {self.task} is not supported for plotting results." @@ -694,6 +732,63 @@ def _plot_obb( self._save_image(save_path, img) return img + def _plot_mask_generation( + self, + source_path: str | Path | np.ndarray | Image.Image, + save_path: str | Path | None = None, + **kwargs, + ) -> np.ndarray: + """Overlay the selected mask and echoed point prompts on the source image.""" + + del kwargs + if self.masks is None: + raise ValueError("No mask_generation output found.") + img = self._read_image(source_path) + masks = ( + self.masks.detach().cpu().numpy() + if isinstance(self.masks, torch.Tensor) + else np.asarray(self.masks) + ) + if masks.ndim != 3: + raise ValueError(f"Expected masks shaped (N, H, W), got {masks.shape}.") + index = self.selected if self.selected is not None else 0 + mask = masks[index] > 0 + if tuple(mask.shape) != (img.shape[0], img.shape[1]): + raise ValueError( + f"Mask shape {mask.shape} does not match image shape {img.shape[:2]}." + ) + overlay = np.zeros_like(img, dtype=np.uint8) + overlay[mask] = MASK_GENERATION_COLOR + blended = cv2.addWeighted( + img, 1.0 - DENSE_OVERLAY_ALPHA, overlay, DENSE_OVERLAY_ALPHA, 0 + ) + result = img.copy() + result[mask] = blended[mask] + if self.points is not None and self.point_labels is not None: + points = ( + self.points.detach().cpu().numpy() + if isinstance(self.points, torch.Tensor) + else np.asarray(self.points) + ) + labels = ( + self.point_labels.detach().cpu().numpy() + if isinstance(self.point_labels, torch.Tensor) + else np.asarray(self.point_labels) + ) + for (x, y), label in zip(points, labels): + color = (0, 255, 0) if int(label) == 1 else (0, 0, 255) + cv2.circle( + result, + (int(x), int(y)), + RADIUS + 1, + color, + -1, + lineType=cv2.LINE_AA, + ) + if save_path is not None: + self._save_image(save_path, result) + return result + def _box_cls_tensor(self) -> torch.Tensor: """Returns detection output as a torch tensor.""" if self.box_cls is None: diff --git a/mblt_vision/wrapper.py b/mblt_vision/wrapper.py index 73552a3..2472c50 100644 --- a/mblt_vision/wrapper.py +++ b/mblt_vision/wrapper.py @@ -83,12 +83,52 @@ def normalize_core_mode( "core_modes_for_target_device", "MOBILINT_CACHE_DIR", "get_mobilint_cache_dir", + "download_hub_artifact", "normalize_core_mode", "resolve_model_config", "MBLT_Engine", ] +def download_hub_artifact( + *, + 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. + + Shared by :meth:`MBLT_Engine._download_hub_artifact` and any engine that + resolves Hub artifacts outside the single-artifact ``file_cfg`` flow (for + example a model that loads more than one compiled artifact). + """ + + 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 _derive_onnx_filename(file_cfg: dict[str, Any]) -> str | None: """Return the configured or MXQ-derived ONNX artifact filename. @@ -683,29 +723,12 @@ def _download_hub_artifact( ) -> 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 + return download_hub_artifact( + repo_id=repo_id, + filename=filename, + revision=revision, + subfolders=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``.""" diff --git a/tests/test_api.py b/tests/test_api.py index 40b4182..243b56d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -21,6 +21,7 @@ def test_public_discovery_exposes_all_supported_tasks() -> None: "obb", "pose_estimation", "face_detection", + "mask_generation", ] assert list_models("obb")["obb"] diff --git a/tests/test_cli.py b/tests/test_cli.py index 62250fe..68f8e6b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -86,6 +86,338 @@ def test_predict_help_explains_supported_workflows( assert "--target-device regulus-ra" in help_text +def test_predict_parses_point_prompts_for_mask_generation() -> None: + """Accept repeated `--point X,Y,LABEL` prompts and mask-generation path overrides.""" + + args = build_parser().parse_args( + [ + "predict", + "--source", + "image.jpg", + "--model", + "sam2-hiera-large", + "--point", + "320,240,1", + "--point", + "10.5,20.5,0", + "--encoder-mxq-path", + "encoder.mxq", + "--decoder-mxq-path", + "decoder.mxq", + "--encoder-onnx-path", + "encoder.onnx", + "--decoder-onnx-path", + "decoder.onnx", + ] + ) + assert args.points == [(320.0, 240.0, 1), (10.5, 20.5, 0)] + assert args.encoder_mxq_path == "encoder.mxq" + assert args.decoder_mxq_path == "decoder.mxq" + assert args.encoder_onnx_path == "encoder.onnx" + assert args.decoder_onnx_path == "decoder.onnx" + + +@pytest.mark.parametrize( + "bad_point", + # nan/inf parse fine as floats, so they must be rejected explicitly: + # they would otherwise contaminate Fourier prompt encoding downstream. + ["320,240", "320,240,2", "x,240,1", "nan,240,1", "320,inf,1", "-inf,240,0"], +) +def test_predict_rejects_malformed_point_prompts(bad_point: str) -> None: + """Fail argument parsing on malformed or out-of-range point prompts.""" + + with pytest.raises(SystemExit): + build_parser().parse_args( + [ + "predict", + "--source", + "image.jpg", + "--model", + "sam2-hiera-large", + "--point", + bad_point, + ] + ) + + +def test_predict_rejects_points_for_non_mask_generation_models( + synthetic_image_path: Path, +) -> None: + """Reject `--point` before constructing an engine for a non-promptable model.""" + + from mblt_vision.cli._vision import run_vision_inference + + args = build_parser().parse_args( + [ + "predict", + "--source", + str(synthetic_image_path), + "--model", + "resnet50", + "--point", + "320,240,1", + ] + ) + with pytest.raises(SystemExit, match="only supported for mask generation"): + run_vision_inference(args, command="predict") + + +@pytest.mark.parametrize( + ("extra_args", "match"), + [ + ([], "1 to 3 point prompts"), + ( + [ + "--point", + "1,1,1", + "--point", + "2,2,1", + "--point", + "3,3,1", + "--point", + "4,4,0", + ], + "1 to 3 point prompts", + ), + (["--point", "1,1,1", "--mxq-path", "model.mxq"], "encoder-mxq-path"), + (["--point", "1,1,1", "--onnx-path", "model.onnx"], "encoder-onnx-path"), + ], +) +def test_mask_generation_prompt_validation_fails_before_engine_construction( + synthetic_image_path: Path, extra_args: list[str], match: str +) -> None: + """Reject invalid mask-generation invocations without loading any backend.""" + + from mblt_vision.cli._vision import run_vision_inference + + args = build_parser().parse_args( + [ + "predict", + "--source", + str(synthetic_image_path), + "--model", + "sam2-hiera-large", + *extra_args, + ] + ) + with pytest.raises(SystemExit, match=match): + run_vision_inference(args, command="predict") + + +@pytest.mark.parametrize( + ("path_arg", "match"), + [ + (["--model-path", "model.mxq"], "encoder-mxq-path"), + (["--mxq-path", "model.mxq"], "encoder-mxq-path"), + (["--onnx-path", "model.onnx"], "encoder-onnx-path"), + ], +) +def test_val_rejects_single_artifact_paths_for_mask_generation( + path_arg: list[str], match: str +) -> None: + """Refuse to silently evaluate the downloaded default instead of the request. + + Validation accepting and ignoring a single-artifact path would invalidate + experiment results, so it must fail the same way prediction does. + """ + + from mblt_vision.cli._vision import create_mask_generation_engine + + args = build_parser().parse_args(["val", "--model", "sam2-hiera-large", *path_arg]) + with pytest.raises(SystemExit, match=match): + create_mask_generation_engine(args) + + +@pytest.mark.parametrize( + ("cache_name", "candidates"), + [ + ("sa-v", ["sav_val.tar", "sa-v", "sav_val"]), + ("nyu-depth", ["nyu-depth.zip", "nyu-depth"]), + ], +) +def test_find_existing_source_never_returns_the_organized_cache( + tmp_path: Path, cache_name: str, candidates: list[str] +) -> None: + """An incomplete cache must not be handed back as its own raw source. + + Several datasets use a cache directory whose name is also a source + candidate, so `data_path.parent / name` resolves back to `data_path`; + organizing from it fails instead of downloading the default archive. + """ + + from mblt_vision.cli.val import _find_existing_source + + data_path = tmp_path / "datasets" / cache_name + data_path.mkdir(parents=True) + assert _find_existing_source(str(data_path), candidates) is None + + # A genuine sibling source is still discovered. + real_source = data_path.parent / candidates[0] + real_source.write_bytes(b"archive") + assert _find_existing_source(str(data_path), candidates) == str(real_source) + + +def test_val_requires_a_manually_downloaded_sav_archive(tmp_path: Path) -> None: + """SA-V is gated by Meta and not mirrored, so there is no default source. + + The error must name the portal, the layout reference, and the flags that + accept the archive, rather than falling back to a URL. + """ + + from mblt_vision.cli.val import _resolve_sav_source + + args = build_parser().parse_args(["val", "--model", "sam2-hiera-large"]) + data_path = tmp_path / "datasets" / "sa-v" + data_path.mkdir(parents=True) + + with pytest.raises(SystemExit) as excinfo: + _resolve_sav_source(args, str(data_path)) + message = str(excinfo.value) + assert "sav_val.tar" in message + assert "ai.meta.com" in message + assert "sav_dataset" in message + assert "--annotation-dir" in message + + # A manually downloaded archive beside the dataset path is accepted. + archive = data_path.parent / "sav_val.tar" + archive.write_bytes(b"archive") + assert _resolve_sav_source(args, str(data_path)) == str(archive) + + +def test_val_finds_the_sav_archive_even_under_force_organize(tmp_path: Path) -> None: + """--force-organize rebuilds the dataset, it does not ignore the source. + + SA-V has no fallback download URL, so skipping discovery would fail on a + manual archive sitting in the very location the error message recommends. + """ + + from mblt_vision.cli.val import _resolve_sav_source + + data_path = tmp_path / "datasets" / "sa-v" + data_path.mkdir(parents=True) + archive = data_path.parent / "sav_val.tar" + archive.write_bytes(b"archive") + + args = build_parser().parse_args( + ["val", "--model", "sam2-hiera-large", "--force-organize"] + ) + assert args.force_organize is True + assert _resolve_sav_source(args, str(data_path)) == str(archive) + + +def test_benchmark_runner_excludes_unsupported_mask_generation() -> None: + """The unified runner cannot build SAM2's engine or dispatch eval_sav. + + Offering the task would accept `--task mask_generation` and then produce an + error row for every model. + """ + + from benchmark import benchmark_vision_models + from mblt_vision._tasks import VISION_TASKS + + assert "mask_generation" in VISION_TASKS # still a canonical Vision task + assert "mask_generation" not in benchmark_vision_models.TASK_CHOICES + assert set(benchmark_vision_models.TASK_CHOICES) == set(VISION_TASKS) - { + "mask_generation" + } + + with pytest.raises(SystemExit): + benchmark_vision_models._parse_args( + ["--models", "ResNet50", "--task", "mask_generation"] + ) + + +@pytest.mark.parametrize( + ("flag", "value"), + [ + ("--encoder-mxq-path", "encoder.mxq"), + ("--decoder-mxq-path", "decoder.mxq"), + ("--encoder-onnx-path", "encoder.onnx"), + ("--decoder-onnx-path", "decoder.onnx"), + ("--prompt-weights-path", "weights.pt"), + ], +) +@pytest.mark.parametrize("command", ["predict", "val"]) +def test_mask_only_overrides_rejected_for_other_tasks( + synthetic_image_path: Path, command: str, flag: str, value: str +) -> None: + """The generic engine never receives these paths. + + Accepting them would silently download and run the default single artifact + instead of the explicitly requested local one. + """ + + from mblt_vision.cli._vision import reject_mask_generation_only_options + + argv = ["predict", "--source", str(synthetic_image_path), "--model", "resnet50"] + if command == "val": + argv = ["val", "--model", "resnet50"] + args = build_parser().parse_args([*argv, flag, value]) + + with pytest.raises(SystemExit, match="only supported for mask generation"): + reject_mask_generation_only_options(args) + + +def test_mask_only_override_guard_is_a_noop_without_those_options() -> None: + """The guard only fires on the mask-generation-only flags. + + It is invoked solely from the non-mask branches, so mask generation keeps + accepting these overrides -- see the parsing tests below. + """ + + from mblt_vision.cli._vision import reject_mask_generation_only_options + + for argv in ( + ["val", "--model", "resnet50"], + ["val", "--model", "resnet50", "--mxq-path", "model.mxq"], + ): + assert ( + reject_mask_generation_only_options(build_parser().parse_args(argv)) is None + ) + + +def test_val_parses_mask_generation_options() -> None: + """Expose SA-V evaluation protocol knobs and artifact-path overrides on val.""" + + args = build_parser().parse_args( + [ + "val", + "--model", + "sam2-hiera-large", + "--num-samples", + "20", + "--num-points", + "2", + "--seed", + "7", + "--encoder-mxq-path", + "encoder.mxq", + "--decoder-mxq-path", + "decoder.mxq", + "--encoder-onnx-path", + "encoder.onnx", + "--decoder-onnx-path", + "decoder.onnx", + ] + ) + assert args.num_samples == 20 + assert args.num_points == 2 + assert args.seed == 7 + assert args.encoder_mxq_path == "encoder.mxq" + assert args.decoder_mxq_path == "decoder.mxq" + assert args.encoder_onnx_path == "encoder.onnx" + assert args.decoder_onnx_path == "decoder.onnx" + + +def test_val_defaults_match_the_reference_protocol() -> None: + """Default to the reference-validated 200-sample single-point protocol.""" + + args = build_parser().parse_args(["val", "--model", "sam2-hiera-large"]) + assert args.num_samples == 200 + assert args.num_points == 1 + assert args.seed == 0 + + def test_validation_default_dataset_path_uses_resolved_cache_root( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: diff --git a/tests/test_dataset_organizer.py b/tests/test_dataset_organizer.py index 94a4098..728528f 100644 --- a/tests/test_dataset_organizer.py +++ b/tests/test_dataset_organizer.py @@ -158,6 +158,7 @@ def test_dotav1_stage_requires_a_non_difficult_target(tmp_path: Path) -> None: (organizer.organize_ade20k, "ade20k"), (organizer.organize_cityscapes, "cityscapes"), (organizer.organize_dotav1, "dotav1"), + (organizer.organize_sav, "sa-v"), ], ) def test_organizer_defaults_use_the_lazy_cache_resolver( @@ -1680,3 +1681,118 @@ def _fail_install_and_rollback(source: str, destination: str) -> None: assert ( backup_dirs[0] / "annotations" / "keep.png" ).read_bytes() == b"old annotation" + + +def _write_sav_fixture_tree(root: Path, *, video_id: str = "sav_000001") -> None: + """Write a tiny official-layout SA-V validation tree with real payloads.""" + + frame = Image.new("RGB", (16, 12), color=(30, 60, 90)) + mask = Image.new("L", (16, 12), color=0) + mask.paste(255, (4, 3, 12, 9)) + (root / "sav_val.txt").write_text(f"{video_id}\n", encoding="utf-8") + frame_dir = root / "JPEGImages_24fps" / video_id + object_dir = root / "Annotations_6fps" / video_id / "000" + frame_dir.mkdir(parents=True) + object_dir.mkdir(parents=True) + for stem in ("00000", "00001", "00002", "00003", "00004"): + frame.save(frame_dir / f"{stem}.jpg") + for stem in ("00000", "00004"): + mask.save(object_dir / f"{stem}.png") + + +def test_organize_sav_keeps_only_annotated_frames( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Install only annotated SA-V frames and masks through staged validation.""" + + monkeypatch.setattr(readiness_module, "SAV_VALIDATION_VIDEO_COUNT", 1) + monkeypatch.setattr(readiness_module, "SAV_VALIDATION_MASKLET_COUNT", 1) + # _write_sav_fixture_tree annotates 2 of its 5 frames for the one masklet. + monkeypatch.setattr(readiness_module, "SAV_VALIDATION_MASK_COUNT", 2) + source_root = tmp_path / "source" / "sav_val" + source_root.mkdir(parents=True) + _write_sav_fixture_tree(source_root) + archive_path = tmp_path / "sav_val.tar" + with tarfile.open(archive_path, "w") as archive: + archive.add(source_root, arcname="sav_val") + + output_dir = tmp_path / "organized" + organizer.organize_sav(str(archive_path), str(output_dir)) + + assert archive_path.is_file() + video_images = output_dir / "images" / "sav_000001" + assert sorted(path.name for path in video_images.iterdir()) == [ + "00000.jpg", + "00004.jpg", + ] + masks = output_dir / "annotations" / "sav_000001" / "000" + assert sorted(path.name for path in masks.iterdir()) == ["00000.png", "00004.png"] + assert (output_dir / "video_ids.txt").read_text(encoding="utf-8") == "sav_000001\n" + assert readiness_module.dataset_ready(output_dir, "mask_generation", "sa-v") + + +@pytest.mark.parametrize( + "malicious_id", ["../../escape", "sav_000001/../../escape", "/abs/escape"] +) +def test_construct_sav_rejects_traversal_video_ids( + tmp_path: Path, malicious_id: str +) -> None: + """Reject ids that would resolve outside staging, before any file is written. + + `sav_val.txt` is file content rather than a directory listing, so an entry + like `../../escape` could otherwise copy attacker-controlled PNGs into the + output parent before organization failed. + """ + + source_root = tmp_path / "source" / "sav_val" + source_root.mkdir(parents=True) + _write_sav_fixture_tree(source_root) + (source_root / "sav_val.txt").write_text(f"{malicious_id}\n", encoding="utf-8") + escaped = source_root / "JPEGImages_24fps" / malicious_id + output_dir = tmp_path / "organized" + before = sorted(p.name for p in tmp_path.iterdir()) + + with pytest.raises(ValueError, match="unsupported video ids|escapes the staging"): + organizer.construct_sav(str(source_root), str(output_dir)) + + del escaped + # Nothing was written outside the intended output directory. + assert sorted(p.name for p in tmp_path.iterdir()) == before + + +def test_construct_sav_rejects_missing_annotated_frame( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Fail organization when an annotated frame has no matching JPEG.""" + + monkeypatch.setattr(readiness_module, "SAV_VALIDATION_VIDEO_COUNT", 1) + monkeypatch.setattr(readiness_module, "SAV_VALIDATION_MASKLET_COUNT", 1) + # _write_sav_fixture_tree annotates 2 of its 5 frames for the one masklet. + monkeypatch.setattr(readiness_module, "SAV_VALIDATION_MASK_COUNT", 2) + source_root = tmp_path / "sav_val" + source_root.mkdir(parents=True) + _write_sav_fixture_tree(source_root) + (source_root / "JPEGImages_24fps" / "sav_000001" / "00004.jpg").unlink() + + with pytest.raises(ValueError, match="missing annotated frame 00004.jpg"): + organizer.construct_sav(str(source_root.parent), str(tmp_path / "organized")) + + +def test_validate_staged_sav_masks_rejects_non_binary_masks(tmp_path: Path) -> None: + """Reject a staged SA-V mask carrying more than one object id.""" + + staged_root = tmp_path / "staged" + image_dir = staged_root / "images" / "sav_000001" + object_dir = staged_root / "annotations" / "sav_000001" / "000" + image_dir.mkdir(parents=True) + object_dir.mkdir(parents=True) + Image.new("RGB", (8, 8)).save(image_dir / "00000.jpg") + multi_object = Image.new("L", (8, 8), color=0) + multi_object.paste(1, (0, 0, 2, 2)) + multi_object.paste(2, (4, 4, 6, 6)) + multi_object.save(object_dir / "00000.png") + + with pytest.raises(ValueError, match="single-object binary map"): + organizer._validate_staged_sav_masks(staged_root) diff --git a/tests/test_dataset_readiness.py b/tests/test_dataset_readiness.py index 36e7b59..d007e55 100644 --- a/tests/test_dataset_readiness.py +++ b/tests/test_dataset_readiness.py @@ -1013,3 +1013,169 @@ def test_dense_readiness_rejects_symlinked_root_ancestors( assert not readiness.dataset_ready( symlink_traversed_parent / "ade20k", "semantic_segmentation", "ade20k" ) + + +def _write_sav_layout( + root: Path, *, video_id: str = "sav_000001", mask_mode: str = "L" +) -> None: + """Create a minimal organized SA-V validation layout with one masklet.""" + + (root / "video_ids.txt").parent.mkdir(parents=True, exist_ok=True) + (root / "video_ids.txt").write_text(f"{video_id}\n", encoding="utf-8") + image_dir = root / "images" / video_id + object_dir = root / "annotations" / video_id / "000" + image_dir.mkdir(parents=True) + object_dir.mkdir(parents=True) + frame = Image.new("RGB", (16, 12)) + mask = Image.new(mask_mode, (16, 12), color=0) + mask.paste(255 if mask_mode == "L" else 1, (4, 3, 12, 9)) + for stem in ("00000", "00004"): + frame.save(image_dir / f"{stem}.jpg") + mask.save(object_dir / f"{stem}.png") + + +@pytest.fixture +def sav_counts(monkeypatch: pytest.MonkeyPatch) -> None: + """Scale the SA-V completeness constants down to the one-video fixture.""" + + monkeypatch.setattr(readiness, "SAV_VALIDATION_VIDEO_COUNT", 1) + monkeypatch.setattr(readiness, "SAV_VALIDATION_MASKLET_COUNT", 1) + # _write_sav_layout writes two annotated frames for its single masklet. + monkeypatch.setattr(readiness, "SAV_VALIDATION_MASK_COUNT", 2) + + +def test_sav_ready_accepts_complete_layout(sav_counts: None, tmp_path: Path) -> None: + """Accept an organized SA-V layout with matching ids, frames, and masks.""" + + _write_sav_layout(tmp_path) + assert readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_video_count_mismatch( + sav_counts: None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Reject a layout whose id list disagrees with the expected video count.""" + + _write_sav_layout(tmp_path) + monkeypatch.setattr(readiness, "SAV_VALIDATION_VIDEO_COUNT", 2) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_masklet_count_mismatch( + sav_counts: None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Reject a layout whose object directories disagree with the masklet count.""" + + _write_sav_layout(tmp_path) + monkeypatch.setattr(readiness, "SAV_VALIDATION_MASKLET_COUNT", 2) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_truncated_annotated_frames( + sav_counts: None, tmp_path: Path +) -> None: + """Reject a source keeping every masklet but only some annotated frames. + + Video and masklet counts alone cannot see this truncation, which would + silently evaluate a different corpus than the pinned official split. + """ + + _write_sav_layout(tmp_path) + (tmp_path / "annotations" / "sav_000001" / "000" / "00004.png").unlink() + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_validates_every_mask_not_only_the_first( + sav_counts: None, tmp_path: Path +) -> None: + """A later mask must be checked too, not just the first pair per video. + + A cache whose first mask is valid but whose second is {1, 2} would + otherwise pass and be turned entirely to foreground by CustomSAV's `> 0`. + """ + + _write_sav_layout(tmp_path) + later_mask = tmp_path / "annotations" / "sav_000001" / "000" / "00004.png" + corrupted = np.full((12, 16), 1, dtype=np.uint8) + corrupted[3:9, 4:12] = 2 + Image.fromarray(corrupted).save(later_mask) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_later_mask_with_wrong_geometry( + sav_counts: None, tmp_path: Path +) -> None: + """Geometry is compared per mask against that mask's own frame.""" + + _write_sav_layout(tmp_path) + later_mask = tmp_path / "annotations" / "sav_000001" / "000" / "00004.png" + Image.new("L", (8, 6), color=0).save(later_mask) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_accepts_bilevel_masks(sav_counts: None, tmp_path: Path) -> None: + """The official split is 1-bit bilevel, where the format itself rules out + a non-zero background, so those masks validate from the header.""" + + _write_sav_layout(tmp_path, mask_mode="1") + assert readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_bilevel_mask_with_wrong_geometry( + sav_counts: None, tmp_path: Path +) -> None: + """The header path still enforces geometry, not only the value constraint.""" + + _write_sav_layout(tmp_path, mask_mode="1") + later_mask = tmp_path / "annotations" / "sav_000001" / "000" / "00004.png" + Image.new("1", (8, 6), color=0).save(later_mask) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_two_valued_mask_without_background( + sav_counts: None, tmp_path: Path +) -> None: + """Reject a {1, 2} mask that `> 0` binarization would make all-foreground.""" + + _write_sav_layout(tmp_path) + mask_path = tmp_path / "annotations" / "sav_000001" / "000" / "00000.png" + corrupted = np.full((12, 16), 1, dtype=np.uint8) + corrupted[3:9, 4:12] = 2 + Image.fromarray(corrupted).save(mask_path) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_orphan_mask_stem(sav_counts: None, tmp_path: Path) -> None: + """Reject a mask annotating a frame that has no matching JPEG.""" + + _write_sav_layout(tmp_path) + (tmp_path / "images" / "sav_000001" / "00004.jpg").unlink() + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_symlinked_video_directory( + sav_counts: None, tmp_path: Path +) -> None: + """Reject symlinked entries inside the managed SA-V layout.""" + + _write_sav_layout(tmp_path) + real_dir = tmp_path / "images" / "sav_000001" + moved = tmp_path / "moved" + real_dir.rename(moved) + real_dir.symlink_to(moved) + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_missing_id_list(sav_counts: None, tmp_path: Path) -> None: + """Reject a layout with no video_ids.txt identity file.""" + + _write_sav_layout(tmp_path) + (tmp_path / "video_ids.txt").unlink() + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") + + +def test_sav_ready_rejects_malformed_video_id(sav_counts: None, tmp_path: Path) -> None: + """Reject video ids that do not match the official sav_ pattern.""" + + _write_sav_layout(tmp_path, video_id="video01") + assert not readiness.dataset_ready(tmp_path, "mask_generation", "sa-v") diff --git a/tests/test_eval_sav.py b/tests/test_eval_sav.py new file mode 100644 index 0000000..9e502d4 --- /dev/null +++ b/tests/test_eval_sav.py @@ -0,0 +1,326 @@ +"""Tests for SA-V point-prompted mask generation evaluation.""" + +from __future__ import annotations + +import importlib +import random +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np +import pytest +from PIL import Image + +from mblt_vision.utils.datasets import CustomSAV +from mblt_vision.utils.evaluation.eval_sav import ( + SAVMetricAccumulator, + build_prompt, + calculate_sav_sample_ious, + eval_sav, + iter_selected_samples, + mask_iou, +) + +eval_sav_module = importlib.import_module("mblt_vision.utils.evaluation.eval_sav") + + +def _disk_mask(height: int = 96, width: int = 128, radius: int = 20) -> np.ndarray: + mask = np.zeros((height, width), dtype=bool) + yy, xx = np.mgrid[:height, :width] + mask[(yy - height // 2) ** 2 + (xx - width // 2) ** 2 <= radius**2] = True + return mask + + +def test_mask_iou_matches_hand_computed_overlap() -> None: + """Score partial overlap exactly and treat two empty masks as identical.""" + + left = np.zeros((4, 4), dtype=bool) + right = np.zeros((4, 4), dtype=bool) + left[:2, :2] = True + right[:2, :] = True + assert mask_iou(left, right) == pytest.approx(4 / 8) + assert mask_iou(np.zeros((4, 4), bool), np.zeros((4, 4), bool)) == 1.0 + + +def test_calculate_sav_sample_ious_rejects_shape_mismatch() -> None: + """Never score candidates against a ground truth of a different geometry.""" + + with pytest.raises(ValueError, match="shapes must match"): + calculate_sav_sample_ious(np.zeros((3, 4, 4), bool), np.zeros((5, 5), bool)) + with pytest.raises(ValueError, match=r"shaped \(N, H, W\)"): + calculate_sav_sample_ious(np.zeros((4, 4), bool), np.zeros((4, 4), bool)) + + +def test_calculate_sav_sample_ious_rejects_non_binary_candidates() -> None: + """Logits, probabilities, and NaN must be reported, not silently binarized. + + `astype(bool)` treats every nonzero value as foreground, including negative + logits, which would yield a plausible but meaningless IoU. + """ + + gt = np.zeros((4, 4), dtype=bool) + gt[:2, :2] = True + + logits = np.full((1, 4, 4), -3.5, dtype=np.float32) + logits[0, :2, :2] = 4.0 + with pytest.raises(ValueError, match="must be a binary map"): + calculate_sav_sample_ious(logits, gt) + + probabilities = np.linspace(0.0, 1.0, 16, dtype=np.float32).reshape(1, 4, 4) + with pytest.raises(ValueError, match="must be a binary map"): + calculate_sav_sample_ious(probabilities, gt) + + non_finite = np.zeros((1, 4, 4), dtype=np.float32) + non_finite[0, 0, 0] = np.nan + with pytest.raises(ValueError, match="must be finite"): + calculate_sav_sample_ious(non_finite, gt) + + # A single positive value is not sufficient: a thresholded probability map + # and a degenerate uniform one both have exactly one, and the latter would + # silently become an all-foreground mask. + thresholded = np.where(gt, 0.5, 0.0).astype(np.float32)[None] + with pytest.raises(ValueError, match="must be a binary map"): + calculate_sav_sample_ious(thresholded, gt) + + uniform = np.full((1, 4, 4), 0.5, dtype=np.float32) + with pytest.raises(ValueError, match="must be a binary map"): + calculate_sav_sample_ious(uniform, gt) + + with pytest.raises(ValueError, match="unsupported dtype"): + calculate_sav_sample_ious(np.array([[["a", "b"], ["c", "d"]]]), gt[:2, :2]) + + +@pytest.mark.parametrize( + "encoding", + [ + np.array([[[True, False], [False, False]]]), + np.array([[[1, 0], [0, 0]]], dtype=np.uint8), + np.array([[[255, 0], [0, 0]]], dtype=np.uint8), + np.array([[[1.0, 0.0], [0.0, 0.0]]], dtype=np.float32), + ], +) +def test_calculate_sav_sample_ious_accepts_conventional_binary_encodings( + encoding: np.ndarray, +) -> None: + """bool, {0, 1}, {0, 255} and float 0/1 all denote the same binary mask.""" + + gt = np.array([[True, False], [False, False]]) + assert calculate_sav_sample_ious(encoding, gt) == [pytest.approx(1.0)] + + +@pytest.mark.parametrize( + ("degenerate", "expected_iou"), + [ + (np.zeros((1, 2, 2), dtype=np.uint8), 0.0), + (np.ones((1, 2, 2), dtype=np.uint8), 0.25), + ], +) +def test_calculate_sav_sample_ious_accepts_empty_and_full_masks( + degenerate: np.ndarray, expected_iou: float +) -> None: + """An all-background or all-foreground prediction is valid, just wrong.""" + + gt = np.array([[True, False], [False, False]]) + assert calculate_sav_sample_ious(degenerate, gt) == [pytest.approx(expected_iou)] + + +@pytest.mark.parametrize( + ("num_points", "expected_labels"), + [(1, [1]), (2, [1, 0]), (3, [1, 1, 0])], +) +def test_build_prompt_points_lie_on_the_correct_side( + num_points: int, expected_labels: list[int] +) -> None: + """Place positive points inside the mask and negative points outside it.""" + + mask = _disk_mask() + prompt = build_prompt(mask, random.Random(0), num_points) + assert prompt is not None + points, labels = prompt + assert labels.tolist() == expected_labels + for (x, y), label in zip(points, labels): + assert bool(mask[int(y), int(x)]) == bool(label) + + +def test_build_prompt_returns_none_for_empty_mask() -> None: + """Skip masks that cannot anchor a positive point.""" + + assert build_prompt(np.zeros((32, 32), bool), random.Random(0), 1) is None + + +def test_negative_point_avoids_the_dilated_mask() -> None: + """Sample negatives outside the safety-dilated foreground, not merely outside it.""" + + mask = _disk_mask() + dilated = cv2.dilate( + mask.astype(np.uint8), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9)) + ).astype(bool) + for seed in range(5): + point = eval_sav_module._negative_point( + mask, mask.shape[0], mask.shape[1], random.Random(seed) + ) + assert point is not None + assert not dilated[int(point[1]), int(point[0])] + + +def _write_sav_root(root: Path, videos: int = 3, frames_per_video: int = 4) -> None: + """Write a tiny organized SA-V root with one masklet per video.""" + + for video_index in range(videos): + video_id = f"sav_{video_index:06d}" + image_dir = root / "images" / video_id + object_dir = root / "annotations" / video_id / "000" + image_dir.mkdir(parents=True) + object_dir.mkdir(parents=True) + for frame_index in range(frames_per_video): + stem = f"{frame_index * 4:05d}" + Image.new("RGB", (64, 48), color=(video_index, 0, 0)).save( + image_dir / f"{stem}.jpg" + ) + mask = Image.new("L", (64, 48), color=0) + mask.paste(255, (8, 8, 40, 40)) + mask.save(object_dir / f"{stem}.png") + (root / "video_ids.txt").write_text( + "\n".join(f"sav_{index:06d}" for index in range(videos)) + "\n", + encoding="utf-8", + ) + + +def test_iter_selected_samples_is_deterministic_per_seed(tmp_path: Path) -> None: + """Reproduce the exact sample selection for a fixed seed.""" + + _write_sav_root(tmp_path) + dataset = CustomSAV(str(tmp_path)) + first = list(iter_selected_samples(dataset, seed=3, per_video=2, min_mask_area=10)) + second = list(iter_selected_samples(dataset, seed=3, per_video=2, min_mask_area=10)) + other_seed = list( + iter_selected_samples(dataset, seed=4, per_video=2, min_mask_area=10) + ) + assert first and first == second + assert first != other_seed + + +def test_iter_selected_samples_caps_per_video(tmp_path: Path) -> None: + """Never draw more than per_video samples from one video.""" + + _write_sav_root(tmp_path, videos=2, frames_per_video=6) + dataset = CustomSAV(str(tmp_path)) + selected = list( + iter_selected_samples(dataset, seed=0, per_video=2, min_mask_area=10) + ) + videos = [dataset.samples[index][2] for index in selected] + for video_id in set(videos): + assert videos.count(video_id) <= 2 + + +def test_iter_selected_samples_skips_small_masks(tmp_path: Path) -> None: + """Exclude ground-truth masks below the minimum area.""" + + _write_sav_root(tmp_path, videos=1) + dataset = CustomSAV(str(tmp_path)) + assert not list( + iter_selected_samples(dataset, seed=0, per_video=4, min_mask_area=10_000) + ) + + +def test_accumulator_computes_selection_and_oracle_means() -> None: + """Aggregate own-selection and best-of-3 IoUs independently.""" + + accumulator = SAVMetricAccumulator() + accumulator.update([0.2, 0.9, 0.5], selected=1, video_id="a") + accumulator.update([0.8, 0.1, 0.4], selected=2, video_id="b") + result = accumulator.result() + assert result.miou == pytest.approx((0.9 + 0.4) / 2) + assert result.miou_best_of_3 == pytest.approx((0.9 + 0.8) / 2) + assert result.num_samples == 2 + assert result.distinct_videos == 2 + assert result.primary_score == result.miou + assert result.secondary_score == result.miou_best_of_3 + + +def test_accumulator_rejects_wrong_candidate_count() -> None: + """Enforce the three-candidate SAM2 output contract.""" + + with pytest.raises(ValueError, match="Expected 3 candidate IoUs"): + SAVMetricAccumulator().update([0.5, 0.5], selected=0, video_id="a") + + +def test_eval_sav_rejects_wrong_taxonomy() -> None: + """Refuse to score a model configured for a different dataset.""" + + # _PerfectModel (defined below) gives a real predict() the taxonomy check + # never reaches, unlike SimpleNamespace's synthesized-but-statically-invisible + # attributes. + model = _PerfectModel() + model.post_cfg = {"task": "mask_generation", "dataset": "coco"} + with pytest.raises(ValueError, match="post_cfg.dataset to be 'sa-v'"): + eval_sav(model, "/dataset") + + +@dataclass +class _MockPrediction: + """Structurally matches ``eval_sav``'s ``PromptedPrediction`` -- unlike + ``SimpleNamespace``, a dataclass's fields are visible to static typing.""" + + masks: np.ndarray + selected: int | None + + +class _PerfectModel: + """Mock engine returning the ground truth as its best-selected candidate.""" + + post_cfg = {"task": "mask_generation", "dataset": "sa-v"} + + def __init__(self) -> None: + self.calls = 0 + + def predict(self, frame, points, labels): + del points, labels + self.calls += 1 + height, width = frame.shape[:2] + gt = np.zeros((height, width), dtype=bool) + gt[8:40, 8:40] = True + masks = np.stack([np.zeros_like(gt), gt, np.ones_like(gt)]) + return _MockPrediction(masks=masks, selected=1) + + +def test_eval_sav_scores_a_mocked_engine_exactly(tmp_path: Path) -> None: + """Produce exact metrics for crafted candidates against the fixture masks.""" + + _write_sav_root(tmp_path, videos=3, frames_per_video=4) + model = _PerfectModel() + result = eval_sav( + model, str(tmp_path), num_samples=4, num_points=1, seed=0, per_video=2 + ) + assert result.miou == pytest.approx(1.0) + assert result.miou_best_of_3 == pytest.approx(1.0) + assert result.num_samples == 4 + assert model.calls == 4 + + +def test_eval_sav_rejects_prediction_without_selected_index(tmp_path: Path) -> None: + """Report a missing selection as ValueError, not a bare assert. + + The PromptedPrediction protocol permits `selected=None`, and an assert + would vanish under `python -O` and resurface as an opaque comparison + TypeError inside the accumulator. + """ + + _write_sav_root(tmp_path, videos=3, frames_per_video=4) + + class _NoSelectionModel(_PerfectModel): + def predict(self, frame, points, labels): + prediction = super().predict(frame, points, labels) + return _MockPrediction(masks=prediction.masks, selected=None) + + with pytest.raises(ValueError, match="missing a selected mask index"): + eval_sav(_NoSelectionModel(), str(tmp_path), num_samples=1, per_video=2) + + +def test_eval_sav_raises_when_samples_run_out(tmp_path: Path) -> None: + """Fail loudly instead of silently reporting fewer samples than requested.""" + + _write_sav_root(tmp_path, videos=1, frames_per_video=2) + with pytest.raises(ValueError, match="only .* valid samples were available"): + eval_sav(_PerfectModel(), str(tmp_path), num_samples=50, per_video=2) diff --git a/tests/test_mask_generation.py b/tests/test_mask_generation.py new file mode 100644 index 0000000..96662a0 --- /dev/null +++ b/tests/test_mask_generation.py @@ -0,0 +1,979 @@ +"""Tests for the mask_generation task and the SAM2HieraLarge model. + +These exercise the NPU-facing contract logic (batch stripping, decoder +output classification, two-backend wiring/cleanup) with mocked backends and +structurally-shaped (not numerically real) prompt weights, matching the +``_FakeBackend`` monkeypatch pattern used throughout ``tests/test_wrapper.py``. +The actual prompt-encoding math is verified bit-for-bit against the real +``facebookresearch/sam2`` predictor separately (not part of this repo's +default test run; see the module docstring in ``_sam2_prompt.py``), and +exercised end-to-end on real hardware only by the opt-in +``tests/test_mask_generation_hardware.py``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch + +import mblt_vision.mask_generation.sam2 as sam2_module +from mblt_vision.mask_generation._sam2_contracts import ( + DECODER_RUNTIME_ORDER, + build_decoder_runtime_feed, + classify_decoder_outputs, + strip_runtime_batch, + validate_runtime_shapes, +) +from mblt_vision.mask_generation.sam2 import SAM2HieraLarge + + +def _decoder_output_set(num_masks: int = 3) -> list[np.ndarray]: + """Synthetic decoder outputs shaped like the real Hiera-Large decoder.""" + + return [ + np.zeros((num_masks, 256, 256), dtype=np.float32), + np.arange(num_masks, dtype=np.float32), + np.ones((num_masks, 256), dtype=np.float32), + np.array([0.5], dtype=np.float32), + ] + + +def test_classify_decoder_outputs_identifies_by_shape_not_position() -> None: + """Order-independent classification, since qbruntime output order is not guaranteed.""" + + masks, iou, sam_tokens, object_score = _decoder_output_set() + for shuffled in ( + [masks, iou, sam_tokens, object_score], + [object_score, sam_tokens, iou, masks], + [iou, masks, object_score, sam_tokens], + ): + classified = classify_decoder_outputs(shuffled) + assert classified["masks"].shape == (3, 256, 256) + assert classified["iou"].shape == (3,) + assert classified["sam_tokens"].shape == (3, 256) + assert classified["object_score"].shape == (1,) + assert np.array_equal(classified["iou"], iou) + + +def test_classify_decoder_outputs_rejects_wrong_candidate_count() -> None: + """A stale/differently exported decoder must not reach the caller. + + Two or four masks with consistently sized iou/token outputs would otherwise + classify cleanly and return a Results.masks shape violating the documented + fixed three-candidate contract. + """ + + for num_masks in (2, 4): + with pytest.raises(ValueError, match="3 decoder mask candidates"): + classify_decoder_outputs(_decoder_output_set(num_masks)) + + +@pytest.mark.parametrize("bad", [np.nan, np.inf, -np.inf]) +def test_classify_decoder_outputs_rejects_non_finite_values(bad: float) -> None: + """A numerical/runtime failure must be reported, not silently absorbed. + + A NaN IoU would participate in argmax and select the wrong candidate, and + non-finite mask logits become plausible booleans under the later `> 0`. + """ + + for index in range(4): + outputs = _decoder_output_set() + outputs[index] = outputs[index].copy() + outputs[index].flat[0] = bad + with pytest.raises(ValueError, match="must be finite"): + classify_decoder_outputs(outputs) + + +@pytest.mark.parametrize( + ("shape", "accepted"), + [ + ((1, 3, 256 * 256), True), # MXQ runtime, flattened + ((3, 256, 256), True), # MXQ runtime, batch already stripped + ((1, 3, 256, 256), True), # exported ONNX graph, NCHW + ((1, 256, 256, 3), False), # channels-last re-export + ((256, 256, 3), False), + ], +) +def test_classify_decoder_outputs_pins_the_mask_layout( + shape: tuple[int, ...], accepted: bool +) -> None: + """A channels-last mask output must not be reshaped into interleaved masks. + + Its size is also a multiple of 256*256 and it also yields three candidates, + so neither the size match nor the candidate-count check can catch it. + """ + + outputs = _decoder_output_set() + outputs[0] = np.zeros(shape, dtype=np.float32) + if accepted: + assert classify_decoder_outputs(outputs)["masks"].shape == (3, 256, 256) + else: + with pytest.raises(ValueError, match="unsupported layout"): + classify_decoder_outputs(outputs) + + +def test_classify_decoder_outputs_rejects_ambiguous_mask_candidates() -> None: + """Two same-sized mask-shaped outputs cannot be told apart -- fail loudly.""" + + masks, iou, sam_tokens, object_score = _decoder_output_set() + with pytest.raises(ValueError, match="Expected exactly one mask output"): + classify_decoder_outputs([masks, masks, iou, sam_tokens, object_score]) + + +def test_strip_runtime_batch_drops_leading_batch_of_one() -> None: + """qbruntime omits the outer model batch from buffer shapes.""" + + batched = np.zeros((1, 4, 4, 3), dtype=np.float64) + stripped = strip_runtime_batch(batched) + assert stripped.shape == (4, 4, 3) + assert stripped.dtype == np.float32 + assert stripped.flags["C_CONTIGUOUS"] + + unbatched = np.zeros((4, 4, 3), dtype=np.float32) + assert strip_runtime_batch(unbatched).shape == (4, 4, 3) + + +def test_build_decoder_runtime_feed_orders_and_strips_batch() -> None: + """Feed tensors are reordered per the compiled artifact's positional signature.""" + + tensors = { + role: np.full((1, 2, 2, 3), fill_value=index, dtype=np.float32) + for index, role in enumerate(DECODER_RUNTIME_ORDER) + } + feed = build_decoder_runtime_feed(tensors) + assert len(feed) == len(DECODER_RUNTIME_ORDER) + for role, array in zip(DECODER_RUNTIME_ORDER, feed): + assert array.shape == (2, 2, 3) # leading batch-of-1 dim stripped + assert np.all(array == tensors[role][0]) + + +def test_build_decoder_runtime_feed_rejects_missing_role() -> None: + tensors = {role: np.zeros((1, 1)) for role in DECODER_RUNTIME_ORDER[:-1]} + with pytest.raises(ValueError, match="missing role"): + build_decoder_runtime_feed(tensors) + + +def test_validate_runtime_shapes_accepts_wildcard_dynamic_dim() -> None: + """``-1`` marks the decoder's point-count-dependent token axis as dynamic.""" + + actual = [np.zeros((1, 1, 9, 256), dtype=np.float32)] + validate_runtime_shapes(actual, [(1, 1, -1, 256)], "decoder") + + +def test_validate_runtime_shapes_rejects_static_dim_mismatch() -> None: + actual = [np.zeros((1, 1, 9, 256), dtype=np.float32)] + with pytest.raises(ValueError, match="shape mismatch"): + validate_runtime_shapes(actual, [(1, 1, -1, 128)], "decoder") + + +class _FakeModelHandle: + """Mirrors the ``qbruntime.Model`` slot-zero compatibility handle + ``MobilintNPUBackend.mxq_model`` exposes.""" + + def __init__(self, input_shape: list[tuple[int, ...]]) -> None: + self._input_shape = input_shape + + def get_model_input_shape(self) -> list[tuple[int, ...]]: + return self._input_shape + + +class _FakeBackend: + """Mirrors MobilintNPUBackend's minimal interface (see tests/test_wrapper.py).""" + + instances: list["_FakeBackend"] = [] + input_shape: list[tuple[int, ...]] = [(1024, 1024, 3)] + + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self.disposed = False + self.launched = False + self.mxq_model = _FakeModelHandle(self.input_shape) + _FakeBackend.instances.append(self) + + def create(self) -> None: + return None + + def launch(self) -> None: + self.launched = True + + def get_dtype(self) -> str: + return "DataType.Float32" + + def __call__(self, feed: list[np.ndarray]) -> list[np.ndarray]: + del feed + return _decoder_output_set() + + def dispose(self) -> None: + self.disposed = True + + +class _FailingDecoderBackend(_FakeBackend): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + if kwargs["mxq_path"].endswith("decoder.mxq"): + raise RuntimeError("simulated decoder load failure") + + +class _WrongShapeBackend(_FakeBackend): + """Reports an input shape that does not match the fed tensor -- simulates a + resolved artifact compiled from a different signature than expected.""" + + input_shape = [(1, 2, 3)] + + +def _fake_prompt_weights() -> dict[str, torch.Tensor]: + """Structurally-correct (shape-only, not numerically real) prompt weights.""" + + return { + "no_mem_embed": torch.zeros(1, 1, 256), + "iou_token_weight": torch.zeros(1, 256), + "mask_tokens_weight": torch.zeros(4, 256), + "obj_score_token_weight": torch.zeros(1, 256), + "positional_encoding_gaussian_matrix": torch.zeros(2, 128), + "point_embedding_negative": torch.zeros(1, 256), + "point_embedding_positive": torch.zeros(1, 256), + "not_a_point_embed_weight": torch.zeros(1, 256), + "no_mask_embed_weight": torch.zeros(1, 256), + } + + +@pytest.fixture(autouse=True) +def _reset_fake_backend_instances() -> None: + _FakeBackend.instances.clear() + + +def _make_engine( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + backend_cls: type[_FakeBackend] = _FakeBackend, +) -> tuple[Path, Path]: + encoder_path = tmp_path / "encoder.mxq" + decoder_path = tmp_path / "decoder.mxq" + encoder_path.write_bytes(b"mxq") + decoder_path.write_bytes(b"mxq") + weights_path = tmp_path / "prompt_weights.pt" + torch.save(_fake_prompt_weights(), weights_path) + monkeypatch.setattr(sam2_module, "MobilintNPUBackend", backend_cls) + # Every test below passes explicit encoder_mxq_path/decoder_mxq_path, so this + # is only ever reached for the prompt-weights resolution (no explicit + # prompt_weights_path given). + monkeypatch.setattr( + sam2_module, "download_hub_artifact", lambda **kwargs: str(weights_path) + ) + return encoder_path, decoder_path + + +def test_sam2_hiera_large_loads_two_backends_with_single_core_default( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Explicit local paths skip Hub download; both backends default to single-core.""" + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + assert len(_FakeBackend.instances) == 2 + encoder_backend, decoder_backend = _FakeBackend.instances + assert encoder_backend.kwargs["mxq_path"] == str(encoder_path) + assert decoder_backend.kwargs["mxq_path"] == str(decoder_path) + assert encoder_backend.kwargs["core_mode"] == "single" + assert decoder_backend.kwargs["core_mode"] == "single" + assert encoder_backend.launched and decoder_backend.launched + finally: + engine.close() + + assert all(instance.disposed for instance in _FakeBackend.instances) + + +def test_sam2_hiera_large_close_is_idempotent_and_context_manager_disposes( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + + with SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) as engine: + pass + assert all(instance.disposed for instance in _FakeBackend.instances) + engine.close() # second close must not raise + + +def test_sam2_hiera_large_disposes_encoder_when_decoder_construction_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Mirrors the reference pipeline's exception-safe encoder/decoder construction.""" + + encoder_path, decoder_path = _make_engine( + monkeypatch, tmp_path, _FailingDecoderBackend + ) + + with pytest.raises(RuntimeError, match="simulated decoder load failure"): + SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + + # The encoder backend was fully constructed before the decoder load failed; + # SAM2HieraLarge.__init__'s except-clause must dispose it rather than leak it. + assert len(_FakeBackend.instances) == 2 + encoder_backend = _FakeBackend.instances[0] + assert encoder_backend.kwargs["mxq_path"] == str(encoder_path) + assert encoder_backend.disposed + + +def test_predict_preprocessed_rejects_out_of_range_point_counts( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Point-count validation happens before any backend call, per the reference's 1-3 point contract.""" + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + encoder_input = np.zeros((1024, 1024, 3), dtype=np.float32) + with pytest.raises(ValueError, match="1 to 3 point prompts"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=np.zeros((0, 2)), labels=[] + ) + with pytest.raises(ValueError, match="1 to 3 point prompts"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=np.zeros((4, 2)), labels=[1, 1, 1, 1] + ) + finally: + engine.close() + + +def test_predict_preprocessed_rejects_unsupported_point_labels( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Only 1 (positive) and 0 (negative) carry a learned embedding. + + Any other value would receive neither and silently produce a plausible but + semantically meaningless mask, so reject it before any backend call. + """ + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + encoder_input = np.zeros((1024, 1024, 3), dtype=np.float32) + with pytest.raises(ValueError, match=r"must be 1 \(positive\) or 0"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=[[10.0, 10.0]], labels=[2] + ) + with pytest.raises(ValueError, match=r"must be 1 \(positive\) or 0"): + engine.predict_preprocessed( + encoder_input, + (100, 100), + points=[[10.0, 10.0], [20.0, 20.0]], + labels=[1, -1], + ) + with pytest.raises(ValueError, match=r"labels shaped \(N,\)"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=[[10.0, 10.0]], labels=[[1]] + ) + # Casting to int64 first would truncate these into valid 0/1 labels. + for truncating in ([0.5], [1.9], [-0.1]): + with pytest.raises(ValueError, match="must be whole numbers"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=[[10.0, 10.0]], labels=truncating + ) + with pytest.raises(ValueError, match="labels must be finite"): + engine.predict_preprocessed( + encoder_input, + (100, 100), + points=[[10.0, 10.0]], + labels=[float("nan")], + ) + finally: + engine.close() + + +@pytest.mark.parametrize( + "bad_hw", + [(0, 640), (640, 0), (-1, 480), (640.5, 480), (640, 480, 3), (640,)], +) +def test_predict_preprocessed_rejects_malformed_original_hw( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bad_hw: tuple[int, ...] +) -> None: + """Geometry is validated before either backend runs. + + A zero dimension makes transform_points produce infinite coordinates, and a + fractional or wrong-length size otherwise fails only in the final mask + resize, after both backends have already executed. + """ + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + with pytest.raises(ValueError, match="original_hw"): + engine.predict_preprocessed( + np.zeros((1024, 1024, 3), dtype=np.float32), + bad_hw, + points=[[10.0, 10.0]], + labels=[1], + ) + finally: + engine.close() + + +def test_preprocess_rejects_non_finite_and_non_numeric_images() -> None: + """Interpolation and normalization preserve NaN, so catch it at the source. + + Otherwise the backend fails in a runtime-specific way and the decoder-output + finiteness check fires too late to name the offending input. + """ + + from mblt_vision.mask_generation._sam2_host import preprocess_encoder_input + + non_finite = np.zeros((32, 32, 3), dtype=np.float32) + non_finite[0, 0, 0] = np.nan + with pytest.raises(ValueError, match="must be finite"): + preprocess_encoder_input(non_finite) + + infinite = np.zeros((32, 32, 3), dtype=np.float32) + infinite[1, 1, 1] = np.inf + with pytest.raises(ValueError, match="must be finite"): + preprocess_encoder_input(infinite) + + with pytest.raises(ValueError, match="numeric RGB image"): + preprocess_encoder_input(np.full((4, 4, 3), "x")) + + # uint8 and float sources both remain supported. + for supported in ( + np.zeros((32, 32, 3), dtype=np.uint8), + np.zeros((32, 32, 3), dtype=np.float32), + ): + assert preprocess_encoder_input(supported).shape == (1, 1024, 1024, 3) + + +@pytest.mark.parametrize("converter", ["runtime", "onnx"]) +def test_fpn_converters_pin_the_complete_level_shape(converter: str) -> None: + """A channel-correct level with wrong geometry must not be reshaped. + + (1, 32, 128, 512) has the same element count as (1, 32, 256, 256), so + build_backbone_features would view it into a plausible but corrupted FPN + that then passes every downstream decoder-feed check. + """ + + from mblt_vision.mask_generation._sam2_host import fpn_from_onnx, fpn_from_runtime + + if converter == "onnx": + convert = fpn_from_onnx + good = [ + np.zeros((1, 32, 256, 256), dtype=np.float32), + np.zeros((1, 64, 128, 128), dtype=np.float32), + np.zeros((1, 256, 64, 64), dtype=np.float32), + ] + wrong_geometry = np.zeros((1, 32, 128, 512), dtype=np.float32) + else: + convert = fpn_from_runtime + good = [ + np.zeros((256, 256, 32), dtype=np.float32), + np.zeros((128, 128, 64), dtype=np.float32), + np.zeros((64, 64, 256), dtype=np.float32), + ] + wrong_geometry = np.zeros((128, 512, 32), dtype=np.float32) + + assert len(convert(good, torch.device("cpu"))) == 3 + with pytest.raises(ValueError, match="spatially"): + convert([wrong_geometry, *good[1:]], torch.device("cpu")) + + +def test_predict_preprocessed_rejects_non_finite_point_coordinates( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """NaN/Inf coordinates would contaminate tokens, masks, and IoU scores.""" + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + encoder_input = np.zeros((1024, 1024, 3), dtype=np.float32) + for bad in ([[float("nan"), 10.0]], [[10.0, float("inf")]]): + with pytest.raises(ValueError, match="coordinates must be finite"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=bad, labels=[1] + ) + finally: + engine.close() + + +def test_predict_preprocessed_rejects_encoder_artifact_shape_drift( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A drifted artifact's declared input shape must fail loudly, not silently produce wrong masks.""" + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path, _WrongShapeBackend) + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + encoder_input = np.zeros((1024, 1024, 3), dtype=np.float32) + with pytest.raises(ValueError, match="encoder input 0 shape mismatch"): + engine.predict_preprocessed( + encoder_input, (100, 100), points=np.array([[1.0, 1.0]]), labels=[1] + ) + finally: + engine.close() + + +def test_raw_call_and_generic_postprocess_are_not_supported( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """SAM2HieraLarge needs prompts; the generic single-image MBLT_Engine API does not apply.""" + + encoder_path, decoder_path = _make_engine(monkeypatch, tmp_path) + engine = SAM2HieraLarge( + encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path) + ) + try: + with pytest.raises(NotImplementedError): + engine(np.zeros((1, 1))) + with pytest.raises(NotImplementedError): + engine.postprocess(np.zeros((1, 1))) + with pytest.raises(NotImplementedError): + engine.set_postprocess_thresholds() + with pytest.raises(NotImplementedError): + engine.preprocess_with_metadata(np.zeros((1, 1))) + finally: + engine.close() + + +def test_sam2_hiera_large_is_discoverable_via_list_models() -> None: + import mblt_vision + + assert ( + "SAM2HieraLarge" + in mblt_vision.list_models("mask_generation")["mask_generation"] + ) + assert mblt_vision.mask_generation.SAM2HieraLarge is SAM2HieraLarge + + +class _FakeOnnxInput: + """Mirrors ONNX Runtime's ``NodeArg`` (name + shape with symbolic dims).""" + + def __init__(self, name: str, shape: list[Any]) -> None: + self.name = name + self.shape = shape + + +_ENCODER_ONNX_INPUTS = [_FakeOnnxInput("input_image", [1, 3, 1024, 1024])] +_DECODER_ONNX_INPUTS = [ + _FakeOnnxInput("tokens", [1, "num_tokens", 256]), + _FakeOnnxInput("src", [1, 256, 64, 64]), + _FakeOnnxInput("pos_src", [1, 256, 64, 64]), + _FakeOnnxInput("high_res_features_0", [1, 32, 256, 256]), + _FakeOnnxInput("high_res_features_1", [1, 64, 128, 128]), +] + + +class _FakeONNXBackend: + """Mirrors ``mblt_npu.ONNXBackend``'s minimal interface (dict-fed sessions).""" + + instances: list["_FakeONNXBackend"] = [] + + def __init__( + self, model_path: str, *, providers: Any = None, ort_module: Any = None + ) -> None: + self.model_path = model_path + self.providers = providers + self.ort_module = ort_module + self.created = False + self.disposed = False + self.calls: list[dict[str, tuple[int, ...]]] = [] + _FakeONNXBackend.instances.append(self) + + @property + def _is_encoder(self) -> bool: + return self.model_path.endswith("encoder.onnx") + + def create(self) -> None: + self.created = True + + def get_inputs(self) -> list[_FakeOnnxInput]: + return _ENCODER_ONNX_INPUTS if self._is_encoder else _DECODER_ONNX_INPUTS + + def __call__(self, feed: dict[str, np.ndarray]) -> list[np.ndarray]: + self.calls.append({name: np.asarray(x).shape for name, x in feed.items()}) + if self._is_encoder: + # Batched NCHW FPN levels, exactly as the exported graph declares. + return [ + np.zeros((1, 32, 256, 256), dtype=np.float32), + np.zeros((1, 64, 128, 128), dtype=np.float32), + np.zeros((1, 256, 64, 64), dtype=np.float32), + ] + return [ + np.zeros((1, 3, 256, 256), dtype=np.float32), + np.arange(3, dtype=np.float32).reshape(1, 3), + np.ones((1, 3, 256), dtype=np.float32), + np.array([[0.5]], dtype=np.float32), + ] + + def dispose(self) -> None: + self.disposed = True + + +class _DriftedDecoderONNXBackend(_FakeONNXBackend): + """Reports decoder input names that do not match the exported-graph contract.""" + + def get_inputs(self) -> list[_FakeOnnxInput]: + if self._is_encoder: + return _ENCODER_ONNX_INPUTS + return [_FakeOnnxInput("renamed_tokens", [1, "num_tokens", 256])] + + +@pytest.fixture(autouse=True) +def _reset_fake_onnx_backend_instances() -> None: + _FakeONNXBackend.instances.clear() + + +def _make_onnx_engine( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + backend_cls: type[_FakeONNXBackend] = _FakeONNXBackend, +) -> tuple[Path, Path]: + encoder_path = tmp_path / "encoder.onnx" + decoder_path = tmp_path / "decoder.onnx" + encoder_path.write_bytes(b"onnx") + decoder_path.write_bytes(b"onnx") + weights_path = tmp_path / "prompt_weights.pt" + torch.save(_fake_prompt_weights(), weights_path) + monkeypatch.setattr(sam2_module, "ONNXBackend", backend_cls) + # A structurally-complete stand-in module: _resolve_onnx_providers only + # reads it when an explicit provider list is requested. + monkeypatch.setattr(sam2_module, "_load_onnxruntime", lambda: object()) + monkeypatch.setattr( + sam2_module, "download_hub_artifact", lambda **kwargs: str(weights_path) + ) + return encoder_path, decoder_path + + +def test_sam2_onnx_framework_is_inferred_from_paths_and_disposes( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Explicit .onnx paths select ONNX inference without an explicit framework.""" + + encoder_path, decoder_path = _make_onnx_engine(monkeypatch, tmp_path) + + engine = SAM2HieraLarge( + encoder_onnx_path=str(encoder_path), decoder_onnx_path=str(decoder_path) + ) + try: + assert engine.framework == "onnx" + assert len(_FakeONNXBackend.instances) == 2 + encoder_backend, decoder_backend = _FakeONNXBackend.instances + assert encoder_backend.model_path == str(encoder_path) + assert decoder_backend.model_path == str(decoder_path) + assert encoder_backend.created and decoder_backend.created + assert encoder_backend.providers == ["CPUExecutionProvider"] + finally: + engine.close() + + assert all(instance.disposed for instance in _FakeONNXBackend.instances) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ( + {"framework": "mxq", "encoder_onnx_path": "encoder.onnx"}, + "conflicts with explicit encoder_onnx_path", + ), + ( + {"framework": "onnx", "encoder_mxq_path": "encoder.mxq"}, + "conflicts with explicit encoder_mxq_path", + ), + ( + { + "encoder_mxq_path": "encoder.mxq", + "decoder_onnx_path": "decoder.onnx", + }, + "without an explicit framework", + ), + ({"framework": "tflite"}, "must be 'mxq' or 'onnx'"), + ({"encoder_onnx_path": "encoder.mxq"}, "must end in '.onnx'"), + ({"decoder_mxq_path": "decoder.onnx"}, "must end in '.mxq'"), + ], +) +def test_sam2_framework_and_path_conflicts_fail_fast( + kwargs: dict[str, Any], match: str +) -> None: + """Suffix and framework conflicts fail before any download or backend load.""" + + with pytest.raises(ValueError, match=match): + SAM2HieraLarge(**kwargs) + + +@pytest.mark.parametrize("framework", ["ONNX", "Onnx", "MXQ", "Mxq"]) +def test_sam2_framework_name_is_case_insensitive( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, framework: str +) -> None: + """Match `_model_paths.resolve_framework`, which lowercases explicit names.""" + + assert ( + SAM2HieraLarge._resolve_framework( + framework, mxq_path_passed=False, onnx_path_passed=False + ) + == framework.lower() + ) + + # The conflict checks run on the normalized value too. + with pytest.raises(ValueError, match="conflicts with explicit"): + SAM2HieraLarge._resolve_framework( + "ONNX", mxq_path_passed=True, onnx_path_passed=False + ) + + +def test_sam2_invalid_target_device_reported_before_any_download( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """An unknown board must not surface as a Hub/network failure.""" + + _make_engine(monkeypatch, tmp_path) + downloads: list[str] = [] + + def _record_download(**kwargs: Any) -> str: + downloads.append(str(kwargs.get("filename"))) + raise AssertionError("no artifact may be downloaded before target validation") + + monkeypatch.setattr(sam2_module, "download_hub_artifact", _record_download) + with pytest.raises(ValueError, match="target_device"): + SAM2HieraLarge(target_device="not-a-board") + assert downloads == [] + + +@pytest.mark.parametrize( + "kwargs", + [ + {"encoder_mxq_path": "missing-encoder.mxq"}, + {"decoder_onnx_path": "missing-decoder.onnx"}, + {"prompt_weights_path": "missing-weights.pt"}, + ], +) +def test_sam2_missing_explicit_artifact_fails_before_any_download( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, kwargs: dict[str, Any] +) -> None: + """An invalid local path must be reported as such, not as a Hub failure. + + Offline this would otherwise surface as a network error; online it would + cost a download before failing later in backend creation. + """ + + _make_onnx_engine(monkeypatch, tmp_path) + downloads: list[str] = [] + + def _record_download(**download_kwargs: Any) -> str: + downloads.append(str(download_kwargs.get("filename"))) + raise AssertionError("no artifact may be downloaded before path validation") + + monkeypatch.setattr(sam2_module, "download_hub_artifact", _record_download) + with pytest.raises(FileNotFoundError, match="does not exist"): + SAM2HieraLarge(**kwargs) + assert downloads == [] + + +def test_sam2_construction_error_survives_a_failing_dispose( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A dispose() failure must not replace the original construction error.""" + + class _UndisposableDriftedBackend(_DriftedDecoderONNXBackend): + def dispose(self) -> None: + raise RuntimeError("dispose exploded") + + encoder_path, decoder_path = _make_onnx_engine( + monkeypatch, tmp_path, _UndisposableDriftedBackend + ) + + # The graph-drift ValueError is what tells the caller what is wrong; the + # dispose RuntimeError must not surface in its place. + with pytest.raises(ValueError, match="decoder ONNX input names mismatch"): + SAM2HieraLarge( + encoder_onnx_path=str(encoder_path), decoder_onnx_path=str(decoder_path) + ) + + +def test_sam2_onnx_rejects_a_frozen_prompt_token_axis( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A decoder exported with a fixed token dimension must fail construction. + + It would otherwise work for one-point prompts and fail inside ONNX Runtime + for the advertised two- and three-point prompts. + """ + + class _FrozenTokenAxisBackend(_FakeONNXBackend): + def get_inputs(self) -> list[_FakeOnnxInput]: + if self._is_encoder: + return _ENCODER_ONNX_INPUTS + return [ + _FakeOnnxInput("tokens", [1, 8, 256]), # frozen, not symbolic + *_DECODER_ONNX_INPUTS[1:], + ] + + encoder_path, decoder_path = _make_onnx_engine( + monkeypatch, tmp_path, _FrozenTokenAxisBackend + ) + + with pytest.raises(ValueError, match="must be dynamic"): + SAM2HieraLarge( + encoder_onnx_path=str(encoder_path), decoder_onnx_path=str(decoder_path) + ) + + +def test_sam2_onnx_predict_preprocessed_uses_the_exported_graph_contract( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """NHWC preprocess output is transposed to NCHW and fed by graph input name.""" + + encoder_path, decoder_path = _make_onnx_engine(monkeypatch, tmp_path) + + with SAM2HieraLarge( + encoder_onnx_path=str(encoder_path), decoder_onnx_path=str(decoder_path) + ) as engine: + result = engine.predict_preprocessed( + np.zeros((1, 1024, 1024, 3), dtype=np.float32), + original_hw=(480, 640), + points=[[100.0, 200.0]], + labels=[1], + ) + encoder_backend, decoder_backend = _FakeONNXBackend.instances + assert encoder_backend.calls == [{"input_image": (1, 3, 1024, 1024)}] + (decoder_call,) = decoder_backend.calls + assert decoder_call == { + "tokens": (1, 8, 256), # 6 output tokens + 1 point + 1 pad + "src": (1, 256, 64, 64), + "pos_src": (1, 256, 64, 64), + "high_res_features_0": (1, 32, 256, 256), + "high_res_features_1": (1, 64, 128, 128), + } + assert result.task == "mask_generation" + assert result.masks is not None + assert result.masks.shape == (3, 480, 640) + assert result.masks.dtype == np.bool_ + # The fake decoder's iou_pred is arange(3): argmax selection is index 2. + assert result.selected == 2 + + +def test_sam2_onnx_rejects_session_interface_drift_and_disposes( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A re-exported decoder with different graph inputs must fail at construction.""" + + encoder_path, decoder_path = _make_onnx_engine( + monkeypatch, tmp_path, _DriftedDecoderONNXBackend + ) + + with pytest.raises(ValueError, match="decoder ONNX input names mismatch"): + SAM2HieraLarge( + encoder_onnx_path=str(encoder_path), decoder_onnx_path=str(decoder_path) + ) + + # The encoder session was fully constructed before the decoder validation + # failed; the constructor's except-clause must dispose it rather than leak it. + assert len(_FakeONNXBackend.instances) == 2 + assert _FakeONNXBackend.instances[0].disposed + # The decoder session is created before validation rejects it, and is never + # assigned to the engine -- so only _build_onnx_backend itself can dispose + # it. Without that, the constructor's cleanup cannot reach it and it leaks. + assert _FakeONNXBackend.instances[1].disposed + + +def test_sam2_onnx_missing_onnxruntime_raises_before_any_backend( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The optional-dependency error surfaces before any session or download. + + In an offline environment a Hub fetch attempted first would surface a + network/cache failure instead of the documented package-extra install + error -- and the download would have been useless anyway. + """ + + _make_onnx_engine(monkeypatch, tmp_path) + + def _raise() -> Any: + raise ImportError("onnxruntime is not installed") + + downloads: list[str] = [] + + def _record_download(**kwargs: Any) -> str: + downloads.append(str(kwargs.get("filename"))) + raise AssertionError("no artifact may be downloaded before onnxruntime loads") + + monkeypatch.setattr(sam2_module, "_load_onnxruntime", _raise) + monkeypatch.setattr(sam2_module, "download_hub_artifact", _record_download) + with pytest.raises(ImportError, match="onnxruntime is not installed"): + SAM2HieraLarge(framework="onnx") + assert _FakeONNXBackend.instances == [] + assert downloads == [] + + +def test_fpn_from_onnx_orders_levels_by_channel_and_rejects_other_layouts() -> None: + """Batched NCHW outputs are ordered 32/64/256 regardless of runtime order.""" + + from mblt_vision.mask_generation._sam2_host import fpn_from_onnx + + outputs = [ + np.zeros((1, 256, 64, 64), dtype=np.float32), + np.zeros((1, 32, 256, 256), dtype=np.float32), + np.zeros((1, 64, 128, 128), dtype=np.float32), + ] + levels = fpn_from_onnx(outputs, torch.device("cpu")) + assert [tuple(level.shape) for level in levels] == [ + (1, 32, 256, 256), + (1, 64, 128, 128), + (1, 256, 64, 64), + ] + + with pytest.raises(ValueError, match="Duplicate encoder output"): + fpn_from_onnx([outputs[0], outputs[0], outputs[1]], torch.device("cpu")) + # The MXQ runtime's batchless NHWC layout must be rejected, not guessed at. + with pytest.raises(ValueError, match="Unexpected encoder ONNX output shape"): + fpn_from_onnx([np.zeros((256, 256, 32), dtype=np.float32)], torch.device("cpu")) + + +@pytest.mark.parametrize(("num_points", "num_tokens"), [(1, 8), (2, 9), (3, 10)]) +def test_prepare_decoder_tensors_onnx_builds_the_five_named_inputs( + num_points: int, num_tokens: int +) -> None: + """The ONNX decoder feed keeps the traced pre-flattening shapes.""" + + from mblt_vision.mask_generation._sam2_contracts import DECODER_ONNX_INPUT_NAMES + from mblt_vision.mask_generation._sam2_host import ( + build_backbone_features, + prepare_decoder_tensors_onnx, + ) + + weights = _fake_prompt_weights() + features = build_backbone_features( + weights, + [ + torch.zeros(1, 32, 256, 256), + torch.zeros(1, 64, 128, 128), + torch.zeros(1, 256, 64, 64), + ], + ) + points = np.arange(num_points * 2, dtype=np.float32).reshape(num_points, 2) + labels = np.ones(num_points, dtype=np.int64) + tensors = prepare_decoder_tensors_onnx( + weights, features, points, labels, (480, 640) + ) + + assert tuple(tensors) == DECODER_ONNX_INPUT_NAMES + assert tensors["tokens"].shape == (1, num_tokens, 256) + assert tensors["src"].shape == (1, 256, 64, 64) + assert tensors["pos_src"].shape == (1, 256, 64, 64) + assert tensors["high_res_features_0"].shape == (1, 32, 256, 256) + assert tensors["high_res_features_1"].shape == (1, 64, 128, 128) + assert all(array.dtype == np.float32 for array in tensors.values()) diff --git a/tests/test_mask_generation_hardware.py b/tests/test_mask_generation_hardware.py new file mode 100644 index 0000000..952e132 --- /dev/null +++ b/tests/test_mask_generation_hardware.py @@ -0,0 +1,92 @@ +"""Opt-in real-hardware coverage for SAM2HieraLarge. + +Requires real Aries NPU hardware and network access (to download the two MXQ +artifacts and the small prompt-encoder weights bundle from +``mobilint/sam2-hiera-large``, unless overridden with explicit local paths). +No external ``sam2`` package or manually cloned repository is needed -- see +``mblt_vision/mask_generation/_sam2_host.py``. Skipped by default, matching +"do not require hardware ... for ordinary unit tests" (AGENTS.md). + +Override artifact paths with ``--encoder-mxq-path`` / ``--decoder-mxq-path`` +(shared NPU pytest options, see ``mblt_npu.pytest_plugin``); omitted paths +fall back to downloading from the ``mobilint/sam2-hiera-large`` Hugging Face +repository. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from mblt_npu import MobilintNPUBackend +from mblt_npu.pytest_plugin import NpuParams +from mblt_vision.mask_generation import SAM2HieraLarge + +pytestmark = [pytest.mark.requires_network, pytest.mark.requires_npu] + + +def _synthetic_rectangle_image() -> tuple[np.ndarray, tuple[int, int, int, int]]: + """A simple synthetic image with one solid rectangle to segment.""" + + image = np.full((480, 640, 3), 30, dtype=np.uint8) + box = (150, 100, 400, 300) # x0, y0, x1, y1 + x0, y0, x1, y1 = box + image[y0:y1, x0:x1] = (200, 60, 60) + return image, box + + +def test_sam2_predicts_a_precise_mask_for_a_synthetic_rectangle( + npu_params: NpuParams, +) -> None: + """A single point inside a flat-colored rectangle should segment it precisely.""" + + image, (x0, y0, x1, y1) = _synthetic_rectangle_image() + + engine = SAM2HieraLarge(**npu_params.encoder, **npu_params.decoder) + try: + result = engine.predict( + image, points=[[(x0 + x1) // 2, (y0 + y1) // 2]], labels=[1] + ) + assert result.task == "mask_generation" + assert result.masks is not None + assert result.iou_predictions is not None + assert result.masks.shape == (3, 480, 640) + assert result.selected == int(np.argmax(result.iou_predictions)) + + ground_truth = np.zeros((480, 640), dtype=bool) + ground_truth[y0:y1, x0:x1] = True + predicted = result.masks[result.selected] + intersection = np.logical_and(predicted, ground_truth).sum() + union = np.logical_or(predicted, ground_truth).sum() + iou = intersection / union + assert ( + iou > 0.95 + ), f"Expected a near-perfect mask for a flat rectangle, got IoU={iou:.4f}." + finally: + engine.close() + + +def test_sam2_resolves_artifacts_from_huggingface_hub(npu_params: NpuParams) -> None: + """Construct with no explicit paths and confirm all three artifacts resolve from Hub.""" + + if npu_params.encoder or npu_params.decoder: + pytest.skip( + "This test specifically exercises Hub resolution with no explicit paths." + ) + + engine = SAM2HieraLarge() + try: + # This test constructs the default (mxq) framework, so both backends + # are MobilintNPUBackend -- narrow past the ONNXBackend/None union + # `_encoder_backend`/`_decoder_backend` carry for the mxq_path access. + assert isinstance(engine._encoder_backend, MobilintNPUBackend) + assert isinstance(engine._decoder_backend, MobilintNPUBackend) + assert engine._encoder_backend.mxq_path.endswith("sam2_hiera_large_encoder.mxq") + assert engine._decoder_backend.mxq_path.endswith("sam2_hiera_large_decoder.mxq") + assert os.path.isfile(engine._encoder_backend.mxq_path) + assert os.path.isfile(engine._decoder_backend.mxq_path) + assert engine.weights is not None + finally: + engine.close() diff --git a/tests/test_mask_generation_onnx.py b/tests/test_mask_generation_onnx.py new file mode 100644 index 0000000..da9d401 --- /dev/null +++ b/tests/test_mask_generation_onnx.py @@ -0,0 +1,50 @@ +"""Opt-in ONNX Runtime coverage for SAM2HieraLarge. + +Requires network access (to download the two ONNX exports and the small +prompt-encoder weights bundle from ``mobilint/sam2-hiera-large``) and the +``onnxruntime`` optional extra -- but no NPU hardware, unlike +``tests/test_mask_generation_hardware.py``. Skipped by default, matching +"do not require ... downloads ... for ordinary unit tests" (AGENTS.md). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from mblt_vision.mask_generation import SAM2HieraLarge + +pytestmark = pytest.mark.requires_network + + +def _synthetic_rectangle_image() -> tuple[np.ndarray, tuple[int, int, int, int]]: + """A simple synthetic image with one solid rectangle to segment.""" + + image = np.full((480, 640, 3), 30, dtype=np.uint8) + box = (150, 100, 400, 300) # x0, y0, x1, y1 + x0, y0, x1, y1 = box + image[y0:y1, x0:x1] = (200, 60, 60) + return image, box + + +def test_sam2_onnx_predicts_a_precise_mask_for_a_synthetic_rectangle() -> None: + """A single point inside a flat-colored rectangle should segment it precisely.""" + + image, (x0, y0, x1, y1) = _synthetic_rectangle_image() + + with SAM2HieraLarge(framework="onnx") as engine: + result = engine.predict( + image, points=[[(x0 + x1) // 2, (y0 + y1) // 2]], labels=[1] + ) + assert result.task == "mask_generation" + assert result.masks is not None + assert result.iou_predictions is not None + assert result.masks.shape == (3, 480, 640) + assert result.selected == int(np.argmax(result.iou_predictions)) + + ground_truth = np.zeros((480, 640), dtype=bool) + ground_truth[y0:y1, x0:x1] = True + predicted = result.masks[result.selected] + intersection = np.logical_and(predicted, ground_truth).sum() + union = np.logical_or(predicted, ground_truth).sum() + assert intersection / union > 0.9 diff --git a/tests/test_mask_generation_prompt_encoding.py b/tests/test_mask_generation_prompt_encoding.py new file mode 100644 index 0000000..0fcebe8 --- /dev/null +++ b/tests/test_mask_generation_prompt_encoding.py @@ -0,0 +1,161 @@ +"""Differential test: our from-scratch prompt encoding vs. the real +``facebookresearch/sam2`` predictor, on the same extracted weights. + +Opt-in and network-dependent (downloads the official checkpoint from Hugging +Face and needs a real ``sam2`` install available on ``sys.path`` -- the +official package, never the unofficial PyPI mirror; see +``mblt_vision/mask_generation/_sam2_host.py``). Skips cleanly when ``sam2`` +is not importable, since this is a fidelity check for maintainers changing +``_sam2_prompt.py``/``_sam2_host.py``, not a normal unit test. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +sam2_image_predictor = pytest.importorskip("sam2.sam2_image_predictor") + +from mblt_vision.mask_generation import _sam2_host as host # noqa: E402 +from mblt_vision.mask_generation import _sam2_prompt as prompt # noqa: E402 + +pytestmark = pytest.mark.requires_network + + +@pytest.fixture(scope="module") +def real_predictor(): + predictor = sam2_image_predictor.SAM2ImagePredictor.from_pretrained( + "facebook/sam2-hiera-large" + ) + predictor.model.to(torch.device("cpu")).eval() + return predictor + + +@pytest.fixture(scope="module") +def our_weights(real_predictor): + encoder = real_predictor.model.sam_prompt_encoder + decoder = real_predictor.model.sam_mask_decoder + return { + "no_mem_embed": real_predictor.model.no_mem_embed.detach().clone(), + "iou_token_weight": decoder.iou_token.weight.detach().clone(), + "mask_tokens_weight": decoder.mask_tokens.weight.detach().clone(), + "obj_score_token_weight": decoder.obj_score_token.weight.detach().clone(), + "positional_encoding_gaussian_matrix": encoder.pe_layer.positional_encoding_gaussian_matrix.detach().clone(), + "point_embedding_negative": encoder.point_embeddings[0].weight.detach().clone(), + "point_embedding_positive": encoder.point_embeddings[1].weight.detach().clone(), + "not_a_point_embed_weight": encoder.not_a_point_embed.weight.detach().clone(), + "no_mask_embed_weight": encoder.no_mask_embed.weight.detach().clone(), + } + + +def test_config_constants_match_the_real_model(real_predictor) -> None: + encoder = real_predictor.model.sam_prompt_encoder + decoder = real_predictor.model.sam_mask_decoder + assert encoder.embed_dim == prompt.EMBED_DIM + assert tuple(encoder.image_embedding_size) == prompt.IMAGE_EMBEDDING_SIZE + assert tuple(encoder.input_image_size) == prompt.INPUT_IMAGE_SIZE + assert ( + tuple(tuple(size) for size in real_predictor._bb_feat_sizes) + == prompt.BB_FEAT_SIZES + ) + assert decoder.num_mask_tokens == prompt.NUM_MASK_TOKENS + assert ( + decoder.use_multimask_token_for_obj_ptr + == prompt.USE_MULTIMASK_TOKEN_FOR_OBJ_PTR + ) + assert decoder.pred_obj_scores == prompt.PRED_OBJ_SCORES + assert ( + real_predictor.model.directly_add_no_mem_embed + == prompt.DIRECTLY_ADD_NO_MEM_EMBED + ) + assert real_predictor.mask_threshold == prompt.MASK_THRESHOLD + + +@pytest.mark.parametrize("num_points", [1, 2, 3]) +def test_sparse_and_dense_embeddings_match_bit_for_bit( + real_predictor, our_weights, num_points: int +) -> None: + rng = np.random.default_rng(num_points) + original_hw = (480, 640) + points_np = np.stack( + [rng.uniform(0, 640, size=num_points), rng.uniform(0, 480, size=num_points)], + axis=-1, + ).astype(np.float32) + labels_np = rng.integers(0, 2, size=num_points).astype(np.int64) + + real_encoder = real_predictor.model.sam_prompt_encoder + point_coords = torch.as_tensor(points_np, dtype=torch.float32)[None, ...] + real_unnorm = real_predictor._transforms.transform_coords( + point_coords, normalize=True, orig_hw=original_hw + ) + real_labels = torch.as_tensor(labels_np, dtype=torch.int32)[None, ...] + real_sparse, real_dense = real_encoder( + points=(real_unnorm, real_labels), boxes=None, masks=None + ) + real_dense_pe = real_encoder.get_dense_pe() + + our_unnorm = prompt.transform_points(point_coords.clone(), original_hw) + our_labels = torch.as_tensor(labels_np, dtype=torch.int64)[None, ...] + our_sparse = prompt.embed_points(our_weights, our_unnorm, our_labels) + our_dense = prompt.dense_embeddings_for_no_mask(our_weights, batch_size=1) + our_dense_pe = prompt.get_dense_pe(our_weights) + + assert torch.equal(real_unnorm, our_unnorm) + assert torch.equal(real_sparse, our_sparse) + assert torch.equal(real_dense, our_dense) + assert torch.equal(real_dense_pe, our_dense_pe) + + +def test_decoder_token_prep_matches_bit_for_bit(real_predictor, our_weights) -> None: + decoder = real_predictor.model.sam_mask_decoder + sparse = torch.zeros(1, 2, prompt.EMBED_DIM) + dense = torch.zeros(1, prompt.EMBED_DIM, 64, 64) + image_embeddings = torch.randn(1, prompt.EMBED_DIM, 64, 64) + + real_output_tokens = ( + torch.cat( + [ + decoder.obj_score_token.weight, + decoder.iou_token.weight, + decoder.mask_tokens.weight, + ], + dim=0, + ) + .unsqueeze(0) + .expand(sparse.size(0), -1, -1) + ) + real_tokens = torch.cat((real_output_tokens, sparse), dim=1) + real_src = image_embeddings + dense + + our_tokens, our_src, our_pe = prompt.decoder_token_prep( + our_weights, + image_embeddings=image_embeddings, + dense_prompt_embeddings=dense, + sparse_prompt_embeddings=sparse, + ) + + assert torch.equal(real_tokens, our_tokens) + assert torch.equal(real_src, our_src) + assert torch.equal(real_predictor.model.sam_prompt_encoder.get_dense_pe(), our_pe) + + +def test_preprocess_and_postprocess_match_bit_for_bit(real_predictor) -> None: + rng = np.random.default_rng(7) + image = rng.integers(0, 255, size=(480, 640, 3), dtype=np.uint8) + + real_pre = real_predictor._transforms(np.ascontiguousarray(image))[None, ...] + real_pre = real_pre.permute(0, 2, 3, 1).float().cpu().numpy() + our_pre = host.preprocess_encoder_input(image) + assert np.array_equal(real_pre, our_pre) + + low_res = rng.standard_normal((3, 256, 256)).astype(np.float32) + real_post = ( + real_predictor._transforms.postprocess_masks( + torch.from_numpy(low_res)[None], (480, 640) + )[0] + .detach() + .numpy() + ) + our_post = host.postprocess_masks(low_res, (480, 640)) + assert np.array_equal(real_post, our_post) diff --git a/tests/test_model_metadata.py b/tests/test_model_metadata.py index de3fb68..80b3c6e 100644 --- a/tests/test_model_metadata.py +++ b/tests/test_model_metadata.py @@ -16,6 +16,7 @@ "face_detection": {"widerface"}, "image_classification": {"imagenet"}, "instance_segmentation": {"coco"}, + "mask_generation": {"sa-v"}, "object_detection": {"coco"}, "obb": {"dotav1"}, "pose_estimation": {"coco"},