diff --git a/.claude/skills/mblt-vision/SKILL.md b/.claude/skills/mblt-vision/SKILL.md
index c6f0ae3..4c1974b 100644
--- a/.claude/skills/mblt-vision/SKILL.md
+++ b/.claude/skills/mblt-vision/SKILL.md
@@ -28,47 +28,9 @@ 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
+- Reuse wrapper.download_hub_artifact for any model that needs more than one 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
- graph contracts pinned in _sam2_contracts.py and validated at construction: two MXQ decoder
- generations exist, identified from the artifact's declared input shapes at engine
- construction (detect_decoder_contract) -- assembled (six host-flattened inputs, four
- outputs) and bridged (prompt-encoder raw inputs, two outputs; sam_tokens/object_score are
- optional in classify_decoder_outputs); the ONNX graphs use the direct-MBLT NHWC encoder
- and six named decoder inputs with a dynamic token axis. Keep the prompt-encoding host path and
- classify_decoder_outputs framework-independent. ONNX uses fpn_from_runtime and
- prepare_decoder_tensors_bridged; MXQ uses fpn_from_runtime plus the detected decoder builder.
- 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`
+- 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.
diff --git a/AGENTS.md b/AGENTS.md
index 15c1202..813ada5 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -142,118 +142,12 @@ The current ownership boundary is deliberate:
where every other face model uses `0.7`. Do not normalize the exception away; a size change here is a
durable model-behavior change requiring the guide, both skill copies, and
`mblt_vision/README.md` to be updated in the same commit.
-- `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 runtimes have different graph contracts, pinned in
- `_sam2_contracts.py` and validated at construction. Two MXQ decoder generations exist and
- both are supported, identified from the loaded artifact's declared input shapes
- (`detect_decoder_contract`, run while the engine is built so a drifted artifact fails before
- any inference is spent): the *assembled* contract takes six host-flattened positional inputs
- (tokens pre-concatenated, `image_embeddings + dense` pre-summed; four outputs), while the
- *bridged* contract -- the SDK tutorial's legacy-parser decoder, whose MBLT carries that
- assembly as an in-graph host-bridge subgraph -- takes the prompt encoder's raw outputs
- (two outputs: `classify_decoder_outputs` keeps `masks`/`iou` required and treats
- `sam_tokens`/`object_score` as optional). The uploaded ONNX graphs use the direct-MBLT
- boundary: batched NHWC encoder input `input_image_0`, NHWC FPN outputs, and six named
- decoder inputs with a dynamic prompt axis. ONNX supplies the raw prompt tensors by name;
- MXQ supplies the same semantic tensors positionally. 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); the bridged MXQ contract is numerically verified on real hardware by
- the opt-in `test_sam2_bridged_contract_predicts_numerically` (ground-truth mask IoU across
- 1-3-point prompts, artifacts supplied via `MBLT_VISION_SAM2_BRIDGED_{ENCODER,DECODER}_MXQ`),
- 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.
+- 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.
## Python-First Architecture
@@ -274,14 +168,8 @@ 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.
+- The unified benchmark runner's `TASK_CHOICES` must only list tasks it can actually execute. Do not
+ wire a new canonical task into the runner's choices before its engine and evaluator paths exist.
- 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
@@ -355,7 +243,7 @@ The current ownership boundary is deliberate:
candidate tensors inside `non_max_suppression` is 1.8x in numpy and measured ~2x
*slower* here, because six boolean index operations per iteration cost more in torch
than one gather through a shrinking index; that measurement is recorded in the function.
-- Every `eval_*` evaluator except `eval_sav` runs through
+- Every `eval_*` evaluator runs through
`utils/evaluation/_pipeline.map_batched_inference`, which prepares batches and runs
postprocessing on two separate bounded thread pools while inference stays on the consuming
thread, following `mblt-model-ops`'s helper of the same name. Results come back in input
diff --git a/README.md b/README.md
index e65f84d..35b8590 100644
--- a/README.md
+++ b/README.md
@@ -40,8 +40,7 @@ pip install "mblt-vision-python[onnxruntime-gpu]"
## Quick start
-Most models include matching `preprocess`/`postprocess` behavior around a raw
-call (mask generation is the exception -- see below):
+Each model includes its matching preprocess and postprocess behavior:
```python
from mblt_vision import ResNet50
@@ -71,20 +70,6 @@ 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
@@ -102,24 +87,17 @@ 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, pose
-estimation, and point-prompted mask generation. The selected model determines
-its task and processing pipeline.
+object and face detection, instance and semantic segmentation, OBB, and pose
+estimation. The selected model determines its task and processing pipeline.
By default it downloads the model artifact and saves a plotted result under
`runs/vision/predict/`. Use `--output` to choose the result-image path,
`--topk` for classification labels, and `--conf-thres`/`--iou-thres` for
detection-style tasks. `--framework onnx` selects ONNX Runtime inference;
`--target-device` and `--core-mode` select the MXQ board/runtime mode.
-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/benchmark/README.md b/benchmark/README.md
index d8e2b75..d3dc3dd 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,
-Cityscapes, and SA-V require the appropriate source archives or credentials.
+Public datasets can use their default download sources. ImageNet, DOTA, and
+Cityscapes require the appropriate source archives or credentials.
```bash
python benchmark/organize_coco.py
@@ -27,44 +27,6 @@ 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 b8ed7f6..bb8c977 100644
--- a/benchmark/benchmark_vision_models.py
+++ b/benchmark/benchmark_vision_models.py
@@ -35,16 +35,7 @@
CORE_MODES: tuple[CoreMode, ...] = cast(
tuple[CoreMode, ...], core_modes_for_target_device("aries-rb")
)
-# `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
-)
+TASK_CHOICES = VISION_TASKS
SUPPORTED_TARGET_DEVICES = frozenset({"aries-rb", "regulus-ra", "regulus-rb"})
diff --git a/benchmark/organize_sav.py b/benchmark/organize_sav.py
deleted file mode 100644
index 666c6a2..0000000
--- a/benchmark/organize_sav.py
+++ /dev/null
@@ -1,72 +0,0 @@
-"""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 70092e5..45638f1 100644
--- a/mblt_vision/README.md
+++ b/mblt_vision/README.md
@@ -2,12 +2,14 @@
`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, 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).
+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.
+
+Breaking change in 0.0.5: SAM2 promptable mask generation, shipped in 0.0.3 and 0.0.4, has
+been removed. The `mask_generation` task, `mblt_vision.mask_generation.SAM2HieraLarge`, the
+`sam2-hiera-large` model, the `mblt-vision predict --point` and encoder/decoder path options,
+and SA-V validation are no longer available. Pin `mblt-vision-python<0.0.5` to keep using them.
## Loading models
@@ -445,73 +447,3 @@ 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. The uploaded ONNX pair uses the same
- direct-MBLT boundary as the SDK MXQs: batched NHWC `input_image_0`, NHWC FPN outputs,
- and six raw prompt-decoder inputs. 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
-```
-
-Local decoder artifacts (`--decoder-mxq-path` / `decoder_mxq_path`) may come from either
-compiled generation: the Hub-hosted decoder, whose host assembles the prompt tokens and
-emits four outputs, or an SDK-tutorial decoder, whose graph does that assembly itself and
-emits two (masks and IoU -- `object_score` is absent from its results). The engine detects
-which one is loaded from the artifact's declared input shapes at construction; an artifact
-matching neither fails there rather than producing wrong masks.
-
-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 34fe387..a8bf94f 100644
--- a/mblt_vision/__init__.py
+++ b/mblt_vision/__init__.py
@@ -11,7 +11,6 @@
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
@@ -20,14 +19,13 @@
from ._api import list_tasks as list_tasks
from .wrapper import MBLT_Engine as MBLT_Engine
-__version__ = "0.0.6"
+__version__ = "0.0.5"
_TASK_MODULES = (
face_detection,
depth_estimation,
image_classification,
instance_segmentation,
- mask_generation,
object_detection,
obb,
pose_estimation,
@@ -51,7 +49,6 @@
"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 64852e4..df9a462 100644
--- a/mblt_vision/_tasks.py
+++ b/mblt_vision/_tasks.py
@@ -13,7 +13,6 @@
"obb",
"pose_estimation",
"face_detection",
- "mask_generation",
)
diff --git a/mblt_vision/cli/_vision.py b/mblt_vision/cli/_vision.py
index 0f4a03a..096a2ef 100644
--- a/mblt_vision/cli/_vision.py
+++ b/mblt_vision/cli/_vision.py
@@ -3,7 +3,6 @@
from __future__ import annotations
import argparse
-import math
import os
import sys
from pathlib import Path
@@ -54,39 +53,6 @@ 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."""
@@ -309,141 +275,6 @@ 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,
*,
@@ -452,14 +283,6 @@ 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 ba37af1..739900d 100644
--- a/mblt_vision/cli/predict.py
+++ b/mblt_vision/cli/predict.py
@@ -8,7 +8,6 @@
add_e2e_arg,
add_threshold_args,
add_vision_parser,
- parse_point,
run_vision_inference,
)
@@ -30,7 +29,7 @@ def add_predict_parser(
command="predict",
help_text=(
"Run vision inference for classification, depth estimation, detection, instance or semantic "
- "segmentation, OBB, pose, face detection, and point-prompted mask generation."
+ "segmentation, OBB, pose, and face detection."
),
description=(
"Run a configured Vision model on one image. The selected model determines the task, "
@@ -38,26 +37,18 @@ def add_predict_parser(
),
epilog="""Supported tasks:
image classification, depth estimation, object and face detection, instance
- and semantic segmentation, oriented bounding boxes (OBB), pose estimation,
- and point-prompted mask generation (SAM2).
+ and semantic segmentation, oriented bounding boxes (OBB), and pose estimation.
The command downloads the default MXQ artifact when no local model path is supplied,
then writes a plotted result under runs/vision/predict/ by default. Use --output to
choose the image destination. Use --framework onnx with --model-path or --onnx-path
for ONNX Runtime inference; MXQ is the default framework.
-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,
@@ -69,46 +60,5 @@ 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 a6ecc1d..5939aae 100644
--- a/mblt_vision/cli/val.py
+++ b/mblt_vision/cli/val.py
@@ -16,12 +16,9 @@
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"]
@@ -35,7 +32,6 @@
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"]
@@ -60,9 +56,9 @@ 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.
- Never returns the organized output directory itself. Several datasets use a
+ Never returns the organized output directory itself. Some 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
+ example `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.
@@ -184,49 +180,6 @@ 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."""
@@ -315,7 +268,6 @@ def _ensure_dataset(
organize_dotav1,
organize_imagenet,
organize_nyu_depth,
- organize_sav,
organize_widerface,
)
except ImportError as exc:
@@ -357,11 +309,6 @@ 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)
@@ -397,18 +344,13 @@ 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
- 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)
+ model = create_vision_engine(args)
try:
if not getattr(getattr(model, "postprocessor", None), "e2e", True):
raise SystemExit(
@@ -448,24 +390,6 @@ 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(
@@ -642,54 +566,5 @@ 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
deleted file mode 100644
index 2e39e7f..0000000
--- a/mblt_vision/datasets/sa-v.yaml
+++ /dev/null
@@ -1,27 +0,0 @@
-# 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
deleted file mode 100644
index 6e2a88e..0000000
--- a/mblt_vision/mask_generation/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-"""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
deleted file mode 100644
index 8951be4..0000000
--- a/mblt_vision/mask_generation/_sam2_contracts.py
+++ /dev/null
@@ -1,328 +0,0 @@
-"""Fixed contracts between the SAM2 encoder/decoder artifacts and their host glue.
-
-MXQ side: two decoder generations exist, distinguished by where SAM2's token
-assembly runs, and both are supported. The loaded artifact's declared input
-shapes identify which one it is (:func:`detect_decoder_contract`):
-
-* **assembled** -- the host concatenates the output tokens and sums
- ``image_embeddings + dense_prompt_embeddings``, feeding six flattened
- tensors. 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). Emits four outputs.
-* **bridged** -- the decoder MBLT carries a host-bridge subgraph that does the
- token concat and the embedding sum itself, so the artifact takes the prompt
- encoder's raw outputs. Produced by the SDK tutorial's
- ``sam2_decoder_to_mblt.py`` (legacy-parser route). Emits two outputs
- (masks and IoU; the parse's ``output_meta`` drops the SAM tokens and
- object score).
-
-Compile and calibration are out of scope for this phase, so each contract is
-a fixed validated signature rather than a configurable MBLT-input-name
-binding map.
-
-ONNX side: the Hub exports use the same direct-MBLT interface as the SDK MXQs:
-batched NHWC `input_image_0`, three NHWC FPN outputs, and six raw prompt-decoder
-inputs. ONNX Runtime receives tensors by graph name; MXQ receives the same semantic
-tensors positionally.
-"""
-
-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 assembled-contract SAM2 Hiera-Large
-# decoder MXQ, which differs from that MBLT graph's declared input names.
-DECODER_RUNTIME_ORDER: tuple[str, ...] = (
- "hrf0_nhwc",
- "src_plus_pos_src",
- "hrf1_nhwc",
- "src",
- "pos_src",
- "tokens",
-)
-
-# Bridged-contract runtime order. Unlike the assembled artifact, this one
-# matches the MBLT input-name order, because the graph's own bridge subgraph
-# consumes the prompt encoder's raw outputs directly.
-DECODER_RUNTIME_ORDER_BRIDGED: tuple[str, ...] = (
- "image_embeddings",
- "dense_prompt_embeddings",
- "image_pe",
- "sparse_prompt_embeddings_0",
- "high_res_features0_0",
- "high_res_features1_0",
-)
-
-# Declared input signatures (batch stripped, as qbruntime reports them).
-# ``-1`` is the prompt-dependent dynamic axis. These are what
-# :func:`detect_decoder_contract` matches against.
-DECODER_MXQ_INPUT_SHAPES: dict[str, tuple[tuple[int, ...], ...]] = {
- "assembled": (
- (256, 256, 32),
- (1, 4096, 256),
- (128, 128, 64),
- (1, 4096, 256),
- (1, 4096, 256),
- (1, -1, 256),
- ),
- "bridged": (
- (256, 64, 64),
- (256, 64, 64),
- (256, 64, 64),
- (1, -1, 256),
- (256, 256, 32),
- (128, 128, 64),
- ),
-}
-
-
-def detect_decoder_contract(shapes: Sequence[Sequence[int]]) -> str:
- """Identify which decoder generation a loaded artifact is, from its shapes.
-
- The two signatures share no prefix -- input 0 is ``(256, 256, 32)`` versus
- ``(256, 64, 64)`` -- so declared shapes are sufficient. A ``-1`` in the
- signature accepts any value there (the artifact may declare the axis
- dynamic or, if compiled for a single prompt length, fixed). Anything
- matching neither raises rather than guessing, since a wrong contract feeds
- tensors whose *roles* are wrong even where shapes coincide.
- """
-
- got = [tuple(int(dim) for dim in shape) for shape in shapes]
- for name, signature in DECODER_MXQ_INPUT_SHAPES.items():
- if len(got) != len(signature):
- continue
- if all(
- len(have) == len(want)
- and all(w == -1 or h == w for h, w in zip(have, want))
- for have, want in zip(got, signature)
- ):
- return name
- raise ValueError(
- f"Decoder artifact input shapes {got} match neither known contract: "
- f"assembled {list(DECODER_MXQ_INPUT_SHAPES['assembled'])} nor "
- f"bridged {list(DECODER_MXQ_INPUT_SHAPES['bridged'])}."
- )
-
-
-# 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_0"
-ENCODER_ONNX_INPUT_SHAPE: tuple[int, ...] = (1, 1024, 1024, 3)
-DECODER_ONNX_INPUT_SHAPES: dict[str, tuple[int, ...]] = {
- "image_embeddings": (1, 256, 64, 64),
- "dense_prompt_embeddings": (1, 256, 64, 64),
- "image_pe": (1, 256, 64, 64),
- "sparse_prompt_embeddings_0": (1, 1, -1, 256),
- "high_res_features0_0": (1, 256, 256, 32),
- "high_res_features1_0": (1, 128, 128, 64),
-}
-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 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.
-
- ``masks`` and ``iou`` are required. ``sam_tokens`` and ``object_score``
- exist only on assembled-contract artifacts; a bridged-contract decoder is
- parsed with an ``output_meta`` that keeps just masks and IoU, so those two
- keys are simply absent from its result rather than an error.
- """
-
- 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, required: bool) -> np.ndarray | None:
- matches = [
- array
- for array in arrays
- if array is not mask_matches[0] and array.size == size
- ]
- if len(matches) > 1 or (required and not matches):
- raise ValueError(
- f"Expected exactly one '{label}' output of size {size}, found {len(matches)}."
- )
- return matches[0] if matches else None
-
- iou = unique("iou", num_masks, required=True)
- assert iou is not None # narrowed by required=True
- result = {"masks": masks, "iou": iou.reshape(num_masks)}
- sam_tokens = unique("sam_tokens", num_masks * 256, required=False)
- if sam_tokens is not None:
- result["sam_tokens"] = sam_tokens.reshape(num_masks, 256)
- object_score = unique("object_score", 1, required=False)
- if object_score is not None:
- result["object_score"] = object_score.reshape(1)
- return result
diff --git a/mblt_vision/mask_generation/_sam2_host.py b/mblt_vision/mask_generation/_sam2_host.py
deleted file mode 100644
index 8653379..0000000
--- a/mblt_vision/mask_generation/_sam2_host.py
+++ /dev/null
@@ -1,403 +0,0 @@
-"""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.
- """
-
- sparse, dense, image_embeddings, high_res = _prompt_encoder_outputs(
- weights, features, points, labels, original_hw
- )
- 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 _prompt_encoder_outputs(
- 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 and return its raw outputs, pre-assembly.
-
- ``(sparse, dense, image_embeddings, high_res)`` -- the stage *before*
- ``decoder_token_prep``, which is what the bridged decoder contract consumes:
- its MBLT carries a host-bridge subgraph that does the token concatenation
- and the ``image_embeddings + dense`` sum inside the graph.
- """
-
- # 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"]]
- return sparse, dense, image_embeddings, high_res
-
-
-def prepare_decoder_tensors_bridged(
- weights: dict[str, torch.Tensor],
- features: BackboneFeatures,
- points: np.ndarray,
- labels: np.ndarray,
- original_hw: Sequence[int],
-) -> dict[str, np.ndarray]:
- """Build the six bridged-contract decoder inputs from the prompt encoder.
-
- Keyed by semantic role because three of the six share the shape
- ``(1, 256, 64, 64)`` -- ``image_embeddings``, ``dense_prompt_embeddings``,
- and ``image_pe`` -- so a positional guess would silently swap them.
- """
-
- sparse, dense, image_embeddings, high_res = _prompt_encoder_outputs(
- weights, features, points, labels, original_hw
- )
- tensors = {
- "image_embeddings": image_embeddings.float(),
- "dense_prompt_embeddings": dense.float(),
- "image_pe": prompt.get_dense_pe(weights).float(),
- # (1, N+1, 256) -> (1, 1, N+1, 256); axis 2 is the dynamic prompt axis.
- "sparse_prompt_embeddings_0": sparse.float().unsqueeze(1).contiguous(),
- "high_res_features0_0": high_res[0].permute(0, 2, 3, 1).contiguous(),
- "high_res_features1_0": high_res[1].permute(0, 2, 3, 1).contiguous(),
- }
- return _as_float32_arrays(tensors)
-
-
-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
deleted file mode 100644
index 63055af..0000000
--- a/mblt_vision/mask_generation/_sam2_prompt.py
+++ /dev/null
@@ -1,223 +0,0 @@
-"""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
deleted file mode 100644
index 1bcf057..0000000
--- a/mblt_vision/mask_generation/sam2.py
+++ /dev/null
@@ -1,780 +0,0 @@
-"""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,
- DECODER_RUNTIME_ORDER_BRIDGED,
- ENCODER_ONNX_INPUT_NAME,
- ENCODER_ONNX_INPUT_SHAPE,
- build_decoder_runtime_feed,
- classify_decoder_outputs,
- detect_decoder_contract,
- strip_runtime_batch,
- validate_onnx_session_inputs,
- validate_runtime_shapes,
-)
-from ._sam2_host import (
- BackboneFeatures,
- build_backbone_features,
- fpn_from_runtime,
- load_rgb,
- postprocess_masks,
- preprocess_encoder_input,
- prepare_decoder_tensors,
- prepare_decoder_tensors_bridged,
-)
-
-_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._decoder_contract: str | 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,
- )
- # Identify the decoder generation now, not on first decode: an
- # artifact matching neither contract must fail while the engine
- # is being built (the enclosing except disposes both backends),
- # not after predict has already spent an encoder inference.
- self._decoder_contract = detect_decoder_contract(
- self._backend_input_shapes(self._decoder_backend)
- )
-
- 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 NHWC layout is the shared direct-MBLT encoder input for both frameworks.
- """
-
- 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"],
- "points": points_array,
- "point_labels": labels_array,
- "selected": selected,
- }
- # Assembled-contract decoders emit an object score; the bridged
- # contract's output_meta keeps only masks and IoU. The key is added
- # only when the output exists, so membership tests see the artifact's
- # real contract rather than a synthetic None.
- if "object_score" in decoder_outputs:
- output["object_score"] = decoder_outputs["object_score"]
- 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 shared direct-MBLT NHWC tensor returned by :meth:`preprocess`.
- """
-
- encoder_backend = self._require_encoder_backend()
- feed = np.asarray(encoder_input, dtype=np.float32)
- validate_runtime_shapes([feed], [ENCODER_ONNX_INPUT_SHAPE], "encoder")
- return fpn_from_runtime(
- encoder_backend({ENCODER_ONNX_INPUT_NAME: feed}), 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.
-
- Two decoder generations exist (see ``_sam2_contracts``): the assembled
- contract takes host-flattened tensors, the bridged contract takes the
- prompt encoder's raw outputs. Which one is loaded was identified at
- construction from the artifact's declared shapes, so a mismatched feed
- is impossible rather than merely checked.
- """
-
- decoder_backend = self._require_decoder_backend()
- if self._decoder_contract == "bridged":
- decoder_tensors = prepare_decoder_tensors_bridged(
- weights, features, points, labels, original_hw
- )
- runtime_order = DECODER_RUNTIME_ORDER_BRIDGED
- else:
- decoder_tensors = prepare_decoder_tensors(
- weights, features, points, labels, original_hw
- )
- runtime_order = DECODER_RUNTIME_ORDER
- decoder_feed = build_decoder_runtime_feed(decoder_tensors, runtime_order)
- 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 the six named direct-MBLT inputs."""
-
- decoder_tensors = prepare_decoder_tensors_bridged(
- 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.
-
- These identify which decoder contract is loaded
- (:func:`detect_decoder_contract`) and gate every feed with a fail-loud
- shape check: an artifact matching neither known contract (for example
- one compiled from a different quantizer revision than the validated
- references) 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
deleted file mode 100644
index fc31f3e..0000000
--- a/mblt_vision/models/SAM2HieraLarge.yaml
+++ /dev/null
@@ -1,18 +0,0 @@
-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 95ff8e4..b0751a8 100644
--- a/mblt_vision/utils/datasets/__init__.py
+++ b/mblt_vision/utils/datasets/__init__.py
@@ -23,7 +23,6 @@
CustomDOTAv1,
CustomImageFolder,
CustomNYUDepth,
- CustomSAV,
CustomWiderface,
CustomWiderFaceDataset,
get_ade20k_loader,
@@ -43,7 +42,6 @@
organize_dotav1,
organize_imagenet,
organize_nyu_depth,
- organize_sav,
organize_widerface,
)
@@ -67,7 +65,6 @@
"CustomDOTAv1",
"CustomImageFolder",
"CustomNYUDepth",
- "CustomSAV",
"CustomWiderFaceDataset",
"CustomWiderface",
"get_ade20k_loader",
@@ -84,6 +81,5 @@
"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 0b558e9..decea74 100644
--- a/mblt_vision/utils/datasets/dataloader.py
+++ b/mblt_vision/utils/datasets/dataloader.py
@@ -275,93 +275,6 @@ 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 f51ff68..0a827cd 100644
--- a/mblt_vision/utils/datasets/organizer.py
+++ b/mblt_vision/utils/datasets/organizer.py
@@ -44,7 +44,6 @@
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,
@@ -65,8 +64,6 @@
}
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"
)
@@ -85,10 +82,6 @@
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).
}
@@ -187,7 +180,6 @@ 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("*"):
@@ -211,8 +203,6 @@ 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:
@@ -1511,227 +1501,6 @@ 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 0b59e8b..7fd18b4 100644
--- a/mblt_vision/utils/datasets/readiness.py
+++ b/mblt_vision/utils/datasets/readiness.py
@@ -38,16 +38,6 @@
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}")
@@ -56,9 +46,6 @@
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")
@@ -1057,123 +1044,6 @@ 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.
@@ -1196,7 +1066,6 @@ 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()
@@ -1218,6 +1087,4 @@ 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 e3d5b80..21b2916 100644
--- a/mblt_vision/utils/evaluation/__init__.py
+++ b/mblt_vision/utils/evaluation/__init__.py
@@ -23,12 +23,6 @@
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] = [
@@ -52,10 +46,6 @@
"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
deleted file mode 100644
index b12c12e..0000000
--- a/mblt_vision/utils/evaluation/eval_sav.py
+++ /dev/null
@@ -1,467 +0,0 @@
-"""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/results.py b/mblt_vision/utils/results.py
index c2c7b3b..d36a4bd 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 Any, cast
+from typing import cast
import cv2
import numpy as np
@@ -47,7 +47,6 @@
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:
@@ -57,7 +56,7 @@ def __init__(
self,
pre_cfg: dict,
post_cfg: dict,
- output: TensorLike | ListTensorLike | NestedListTensorLike | dict[str, Any],
+ output: TensorLike | ListTensorLike | NestedListTensorLike,
**kwargs,
) -> None:
"""
@@ -65,7 +64,7 @@ def __init__(
Args:
pre_cfg (dict): Preprocessing configuration.
post_cfg (dict): Postprocessing configuration.
- output: Raw model output, including a dictionary for mask generation.
+ output (TensorLike | ListTensorLike | NestedListTensorLike): Raw model output.
**kwargs: Additional arguments.
"""
self.pre_cfg = pre_cfg
@@ -77,20 +76,12 @@ 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 | dict[str, Any] | None
- ) = None
+ self.output: TensorLike | ListTensorLike | NestedListTensorLike | None = None
self.labels: torch.Tensor | None = None
self.scores: torch.Tensor | None = None
self.boxes: torch.Tensor | None = None
self.rboxes: torch.Tensor | None = None
self.kpts: torch.Tensor | None = None
- self.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(
@@ -140,8 +131,7 @@ 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 | dict[str, Any],
+ self, output: TensorLike | ListTensorLike | NestedListTensorLike
) -> None:
"""
Sets variables from the raw model output based on the task.
@@ -155,12 +145,6 @@ 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(
@@ -250,26 +234,6 @@ 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."
@@ -309,8 +273,6 @@ 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."
@@ -770,63 +732,6 @@ 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/tests/test_api.py b/tests/test_api.py
index 243b56d..40b4182 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -21,7 +21,6 @@ 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 68f8e6b..1b32e09 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -86,152 +86,9 @@ 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"]),
],
)
@@ -240,7 +97,7 @@ def test_find_existing_source_never_returns_the_organized_cache(
) -> 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
+ Some 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.
"""
@@ -257,167 +114,6 @@ def test_find_existing_source_never_returns_the_organized_cache(
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 a3d0f55..fb5a7a2 100644
--- a/tests/test_dataset_organizer.py
+++ b/tests/test_dataset_organizer.py
@@ -159,7 +159,6 @@ 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(
@@ -1848,118 +1847,3 @@ 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 ab8ec73..9a28641 100644
--- a/tests/test_dataset_readiness.py
+++ b/tests/test_dataset_readiness.py
@@ -1111,169 +1111,3 @@ 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
deleted file mode 100644
index 9e502d4..0000000
--- a/tests/test_eval_sav.py
+++ /dev/null
@@ -1,326 +0,0 @@
-"""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
deleted file mode 100644
index b2411d8..0000000
--- a/tests/test_mask_generation.py
+++ /dev/null
@@ -1,1171 +0,0 @@
-"""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, ClassVar
-
-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_MXQ_INPUT_SHAPES,
- DECODER_RUNTIME_ORDER,
- build_decoder_runtime_feed,
- classify_decoder_outputs,
- detect_decoder_contract,
- 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 _bridged_decoder_output_set() -> list[np.ndarray]:
- """The two outputs a bridged-contract decoder MXQ declares: IoU and masks."""
-
- return [
- np.zeros((1, 1, 3), dtype=np.float32),
- np.zeros((1, 3, 65536), dtype=np.float32),
- ]
-
-
-def test_detect_decoder_contract_identifies_both_generations() -> None:
- assert detect_decoder_contract(DECODER_MXQ_INPUT_SHAPES["assembled"]) == "assembled"
- assert detect_decoder_contract(DECODER_MXQ_INPUT_SHAPES["bridged"]) == "bridged"
- # A single-prompt-length compile declares the dynamic axis fixed; the
- # signature's -1 must accept it.
- fixed = [
- (1, 10, 256) if shape == (1, -1, 256) else shape
- for shape in DECODER_MXQ_INPUT_SHAPES["assembled"]
- ]
- assert detect_decoder_contract(fixed) == "assembled"
-
-
-def test_detect_decoder_contract_rejects_unknown_signatures() -> None:
- with pytest.raises(ValueError, match="neither known contract"):
- detect_decoder_contract([(1, 2, 3)])
- # Same length as a known contract but different shapes must not pass either.
- with pytest.raises(ValueError, match="neither known contract"):
- detect_decoder_contract([(64, 64, 256)] * 6)
-
-
-def test_classify_decoder_outputs_accepts_bridged_two_output_form() -> None:
- classified = classify_decoder_outputs(_bridged_decoder_output_set())
- assert sorted(classified) == ["iou", "masks"]
- assert classified["masks"].shape == (3, 256, 256)
- assert classified["iou"].shape == (3,)
-
-
-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)]
- # Construction now detects the decoder contract from declared shapes, so the
- # decoder handle must present a valid signature even in wiring-only tests.
- decoder_input_shape: ClassVar[list[tuple[int, ...]]] = list(
- DECODER_MXQ_INPUT_SHAPES["assembled"]
- )
-
- def __init__(self, **kwargs: Any) -> None:
- self.kwargs = kwargs
- self.disposed = False
- self.launched = False
- self.is_decoder = kwargs["mxq_path"].endswith("decoder.mxq")
- shapes = self.decoder_input_shape if self.is_decoder else self.input_shape
- self.mxq_model = _FakeModelHandle(list(shapes))
- _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 _DualContractBackend(_FakeBackend):
- """Encoder/decoder pair for driving predict end to end without hardware.
-
- Selects declared shapes and outputs by which artifact path it was built
- for, and records the decoder feed so a test can assert what was sent.
- """
-
- decoder_input_shape: ClassVar[list[tuple[int, ...]]] = list(
- DECODER_MXQ_INPUT_SHAPES["bridged"]
- )
- decoder_outputs: staticmethod = staticmethod(_bridged_decoder_output_set)
-
- def __init__(self, **kwargs: Any) -> None:
- super().__init__(**kwargs)
- self.received_feeds: list[list[np.ndarray]] = []
-
- def __call__(self, feed: list[np.ndarray]) -> list[np.ndarray]:
- self.received_feeds.append(feed)
- if self.is_decoder:
- return self.decoder_outputs()
- return [
- np.zeros((256, 256, 32), dtype=np.float32),
- np.zeros((128, 128, 64), dtype=np.float32),
- np.zeros((64, 64, 256), dtype=np.float32),
- ]
-
-
-class _AssembledContractBackend(_DualContractBackend):
- decoder_input_shape: ClassVar[list[tuple[int, ...]]] = list(
- DECODER_MXQ_INPUT_SHAPES["assembled"]
- )
- decoder_outputs = staticmethod(_decoder_output_set)
-
-
-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 encoder input shape that does not match the fed tensor --
- simulates a resolved artifact compiled from a different signature than
- expected. The decoder keeps a valid signature so construction succeeds and
- the drift is caught at the encoder feed."""
-
- 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()
-
-
-class _UnknownDecoderContractBackend(_FakeBackend):
- """Decoder declares shapes matching neither known contract."""
-
- decoder_input_shape: ClassVar[list[tuple[int, ...]]] = [(64, 64, 256)] * 6
-
-
-def test_unknown_decoder_contract_fails_at_construction_and_disposes(
- monkeypatch: pytest.MonkeyPatch, tmp_path: Path
-) -> None:
- """A drifted decoder artifact must fail while the engine is being built.
-
- Detection at first decode would let the engine construct successfully and
- burn an NPU encoder inference before reporting the invalid signature.
- """
-
- encoder_path, decoder_path = _make_engine(
- monkeypatch, tmp_path, _UnknownDecoderContractBackend
- )
- with pytest.raises(ValueError, match="neither known contract"):
- SAM2HieraLarge(
- encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path)
- )
- assert all(
- instance.disposed for instance in _UnknownDecoderContractBackend.instances
- )
-
-
-def test_predict_preprocessed_drives_the_bridged_decoder_contract(
- monkeypatch: pytest.MonkeyPatch, tmp_path: Path
-) -> None:
- """End-to-end shape flow for a bridged-contract artifact (SDK tutorial MXQ).
-
- The engine must detect the contract from the declared shapes, feed the
- prompt encoder's raw outputs in bridged order, and classify the two-output
- result -- with ``object_score`` absent rather than an error.
- """
-
- encoder_path, decoder_path = _make_engine(
- monkeypatch, tmp_path, _DualContractBackend
- )
- 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)
- result = engine.predict_preprocessed(
- encoder_input,
- (480, 640),
- points=np.array([[10.0, 20.0], [30.0, 40.0]]),
- labels=[1, 0],
- )
- assert result.masks.shape == (3, 480, 640)
- assert result.iou_predictions.shape == (3,)
- # The bridged decoder emits no object score; the key must be absent,
- # not present with a synthetic None.
- assert "object_score" not in result.output
-
- decoder = next(b for b in _DualContractBackend.instances if b.is_decoder)
- (feed,) = decoder.received_feeds
- # Bridged order: three (256, 64, 64) prompt-encoder tensors, the dynamic
- # prompt axis (2 points + 1 pad), then the two NHWC feature maps.
- assert [tuple(t.shape) for t in feed] == [
- (256, 64, 64),
- (256, 64, 64),
- (256, 64, 64),
- (1, 3, 256),
- (256, 256, 32),
- (128, 128, 64),
- ]
- finally:
- engine.close()
-
-
-def test_predict_preprocessed_drives_the_assembled_decoder_contract(
- monkeypatch: pytest.MonkeyPatch, tmp_path: Path
-) -> None:
- """The validated assembled-contract path must keep working unchanged."""
-
- encoder_path, decoder_path = _make_engine(
- monkeypatch, tmp_path, _AssembledContractBackend
- )
- engine = SAM2HieraLarge(
- encoder_mxq_path=str(encoder_path), decoder_mxq_path=str(decoder_path)
- )
- try:
- result = engine.predict_preprocessed(
- np.zeros((1024, 1024, 3), dtype=np.float32),
- (480, 640),
- points=np.array([[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]]),
- labels=[1, 1, 0],
- )
- assert result.masks.shape == (3, 480, 640)
- assert result.output["object_score"].shape == (1,)
-
- decoder = next(b for b in _AssembledContractBackend.instances if b.is_decoder)
- (feed,) = decoder.received_feeds
- # Assembled order: flattened host-side tensors, tokens = 3 points + 7.
- assert [tuple(t.shape) for t in feed] == [
- (256, 256, 32),
- (1, 4096, 256),
- (128, 128, 64),
- (1, 4096, 256),
- (1, 4096, 256),
- (1, 10, 256),
- ]
- 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_0", [1, 1024, 1024, 3])]
-_DECODER_ONNX_INPUTS = [
- _FakeOnnxInput("image_embeddings", [1, 256, 64, 64]),
- _FakeOnnxInput("dense_prompt_embeddings", [1, 256, 64, 64]),
- _FakeOnnxInput("image_pe", [1, 256, 64, 64]),
- _FakeOnnxInput("sparse_prompt_embeddings_0", [1, 1, "num_points", 256]),
- _FakeOnnxInput("high_res_features0_0", [1, 256, 256, 32]),
- _FakeOnnxInput("high_res_features1_0", [1, 128, 128, 64]),
-]
-
-
-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, 256, 256, 32), dtype=np.float32),
- np.zeros((1, 128, 128, 64), dtype=np.float32),
- np.zeros((1, 64, 64, 256), dtype=np.float32),
- ]
- return [
- np.arange(3, dtype=np.float32).reshape(1, 1, 3),
- np.zeros((1, 3, 256 * 256), 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 [
- *_DECODER_ONNX_INPUTS[:3],
- _FakeOnnxInput("sparse_prompt_embeddings_0", [1, 1, 2, 256]),
- *_DECODER_ONNX_INPUTS[4:],
- ]
-
- 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_0": (1, 1024, 1024, 3)}]
- (decoder_call,) = decoder_backend.calls
- assert decoder_call == {
- "image_embeddings": (1, 256, 64, 64),
- "dense_prompt_embeddings": (1, 256, 64, 64),
- "image_pe": (1, 256, 64, 64),
- "sparse_prompt_embeddings_0": (1, 1, 2, 256),
- "high_res_features0_0": (1, 256, 256, 32),
- "high_res_features1_0": (1, 128, 128, 64),
- }
- 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_six_direct_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_bridged,
- )
-
- 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_bridged(
- weights, features, points, labels, (480, 640)
- )
-
- assert tuple(tensors) == DECODER_ONNX_INPUT_NAMES
- assert tensors["image_embeddings"].shape == (1, 256, 64, 64)
- assert tensors["dense_prompt_embeddings"].shape == (1, 256, 64, 64)
- assert tensors["image_pe"].shape == (1, 256, 64, 64)
- assert tensors["sparse_prompt_embeddings_0"].shape == (1, 1, num_points + 1, 256)
- assert tensors["high_res_features0_0"].shape == (1, 256, 256, 32)
- assert tensors["high_res_features1_0"].shape == (1, 128, 128, 64)
- 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
deleted file mode 100644
index 78fc13f..0000000
--- a/tests/test_mask_generation_hardware.py
+++ /dev/null
@@ -1,155 +0,0 @@
-"""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()
-
-
-_BRIDGED_ENCODER_ENV = "MBLT_VISION_SAM2_BRIDGED_ENCODER_MXQ"
-_BRIDGED_DECODER_ENV = "MBLT_VISION_SAM2_BRIDGED_DECODER_MXQ"
-
-
-def _bridged_artifact_paths() -> tuple[str, str]:
- """Explicit bridged-contract artifact paths, or skip.
-
- The Hub-hosted artifacts are the assembled contract, so bridged coverage
- needs locally compiled artifacts (the SDK tutorial's
- ``sam2_decoder_to_mblt.py`` output), pointed at via environment variables.
- """
-
- encoder = os.environ.get(_BRIDGED_ENCODER_ENV)
- decoder = os.environ.get(_BRIDGED_DECODER_ENV)
- if not encoder or not decoder:
- pytest.skip(
- f"Set {_BRIDGED_ENCODER_ENV} and {_BRIDGED_DECODER_ENV} to locally "
- "compiled bridged-contract artifacts to run this test."
- )
- return encoder, decoder
-
-
-def test_sam2_bridged_contract_predicts_numerically() -> None:
- """Numerical verification of the bridged decoder contract on real hardware.
-
- Shape checks alone cannot catch swapped semantic roles among the three
- identically shaped ``(1, 256, 64, 64)`` decoder inputs -- a swap produces
- plausible but wrong masks. A near-perfect ground-truth IoU on a synthetic
- rectangle can only happen when the roles reached the graph correctly, so
- this asserts mask *values*, not tensor plumbing. Runs 1-, 2-, and 3-point
- prompts to also prove the dynamic prompt axis on hardware.
- """
-
- encoder_path, decoder_path = _bridged_artifact_paths()
- image, (x0, y0, x1, y1) = _synthetic_rectangle_image()
- ground_truth = np.zeros((480, 640), dtype=bool)
- ground_truth[y0:y1, x0:x1] = True
- center = ((x0 + x1) // 2, (y0 + y1) // 2)
- prompts = [
- ([center], [1]),
- ([center, (x0 + 30, y0 + 30)], [1, 1]),
- ([center, (x0 + 30, y0 + 30), (20, 20)], [1, 1, 0]),
- ]
-
- engine = SAM2HieraLarge(
- encoder_mxq_path=encoder_path, decoder_mxq_path=decoder_path
- )
- try:
- assert engine._decoder_contract == "bridged"
- for points, labels in prompts:
- result = engine.predict(image, points=points, labels=labels)
- 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"{len(points)}-point prompt: expected a near-perfect mask for "
- f"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
deleted file mode 100644
index da9d401..0000000
--- a/tests/test_mask_generation_onnx.py
+++ /dev/null
@@ -1,50 +0,0 @@
-"""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
deleted file mode 100644
index 0fcebe8..0000000
--- a/tests/test_mask_generation_prompt_encoding.py
+++ /dev/null
@@ -1,161 +0,0 @@
-"""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 80b3c6e..de3fb68 100644
--- a/tests/test_model_metadata.py
+++ b/tests/test_model_metadata.py
@@ -16,7 +16,6 @@
"face_detection": {"widerface"},
"image_classification": {"imagenet"},
"instance_segmentation": {"coco"},
- "mask_generation": {"sa-v"},
"object_detection": {"coco"},
"obb": {"dotav1"},
"pose_estimation": {"coco"},