diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..1ccbb1d --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,31 @@ +name: Tests + +on: + pull_request: + push: + branches: [main] + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.13"] + steps: + - uses: actions/checkout@v7 + with: + submodules: false + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install package and test dependencies + run: python -m pip install --upgrade pip && python -m pip install -e '.[dev]' + - name: Lint + run: ruff check . + - name: Test + run: python -m pytest -q diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index ea531eb..da46147 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -14,13 +14,14 @@ jobs: name: pypi url: https://pypi.org/p/samexporter permissions: + contents: read id-token: write # Mandatory for trusted publishing (OIDC) steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v7 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.x" + python-version: "3.13" - name: Install pypa/build run: >- python -m pip install build==1.2.2 twine==6.1.0 --user diff --git a/.gitmodules b/.gitmodules index 6fa4a2d..f063019 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,9 @@ [submodule "sam3"] path = sam3 url = https://github.com/facebookresearch/sam3.git +[submodule "third_party/segment-anything"] + path = third_party/segment-anything + url = https://github.com/facebookresearch/segment-anything.git +[submodule "third_party/sam2"] + path = third_party/sam2 + url = https://github.com/facebookresearch/sam2.git diff --git a/README.md b/README.md index 6e7ba10..0249ff7 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# SAMExporter — SAM / SAM2 / SAM2.1 / SAM3 / MobileSAM → ONNX +# SAMExporter — SAM / EfficientSAM / MobileSAM / SAM2 / SAM3 → ONNX -Export [Segment Anything](https://github.com/facebookresearch/segment-anything), [MobileSAM](https://github.com/ChaoningZhang/MobileSAM), [Segment Anything 2 / 2.1](https://github.com/facebookresearch/segment-anything-2), and [Segment Anything 3](https://github.com/facebookresearch/sam3) to ONNX for easy, dependency-free deployment. +Export and run [Segment Anything](https://github.com/facebookresearch/segment-anything), [MobileSAM](https://github.com/ChaoningZhang/MobileSAM), [EfficientSAM](https://github.com/yformer/EfficientSAM), [Segment Anything 2 / 2.1](https://github.com/facebookresearch/sam2), and [Segment Anything 3](https://github.com/facebookresearch/sam3) in ONNX Runtime for portable deployment. [![PyPI version](https://badge.fury.io/py/samexporter.svg)](https://badge.fury.io/py/samexporter) [![Downloads](https://pepy.tech/badge/samexporter)](https://pepy.tech/project/samexporter) @@ -14,6 +14,7 @@ Export [Segment Anything](https://github.com/facebookresearch/segment-anything), | SAM ViT-B / ViT-L / ViT-H | Point, Rectangle | Original Meta SAM | | SAM ViT-B / ViT-L / ViT-H (quantized) | Point, Rectangle | Smaller, faster variants | | MobileSAM | Point, Rectangle | Lightweight; fast on CPU | +| EfficientSAM-Ti / S | Point, Rectangle | Apache-2.0; official split ONNX models | | SAM2 Tiny / Small / Base+ / Large | Point, Rectangle | Meta SAM 2 | | SAM2.1 Tiny / Small / Base+ / Large | Point, Rectangle | Improved SAM 2 | | SAM3 ViT-H | **Text**, Point, Rectangle | Open-vocabulary text-driven segmentation | @@ -26,12 +27,12 @@ Requires **Python 3.11+**. ```bash pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu -pip install samexporter +pip install "samexporter[runtime-cpu]" ``` > **Note — Windows users:** The optional `onnxsim` model simplifier (used during ONNX export) has no pre-built wheel for Windows. If you plan to export models and want simplification, install with: > ```bash -> pip install "samexporter[export]" +> pip install "samexporter[runtime-cpu,export]" > ``` > or enable [Windows Long Path support](https://pip.pypa.io/warnings/enable-long-paths) before installing. @@ -41,9 +42,22 @@ pip install samexporter pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu git clone --recurse-submodules https://github.com/vietanhdev/samexporter cd samexporter -pip install -e . +pip install -e ".[runtime-cpu]" ``` +The source checkout pins the official model repositories as submodules so an +upstream `main` branch cannot silently change an export: + +| Family | Pinned official revision | +|---|---| +| SAM 1 | `dca509fe793f601edb92606367a655c15ac00fdf` | +| SAM 2 / 2.1 | `2b90b9f5ceec907a1c18123530e92e794ad901a4` | +| SAM 3 | `660a5e9e1b8b4c02c0ad97229b88a09a6e4ff5b7` | + +These were the latest upstream `main` revisions checked on 2026-08-30. Export +commands prefer these checkouts and retain installed-package fallbacks for the +published wheel. + --- ## SAM / MobileSAM — Convert to ONNX @@ -106,6 +120,16 @@ bash convert_all_meta_sam.sh bash convert_mobile_sam.sh ``` +MobileSAM uses the same checkpoint for both halves of the export: + +```bash +python -m samexporter.export_decoder \ + --checkpoint original_models/mobile_sam.pt \ + --output output_models/mobile_sam/mobile_sam.decoder.onnx \ + --model-type mobile \ + --return-single-mask +``` + ### 4. Run inference ```bash @@ -134,6 +158,39 @@ python -m samexporter.inference \ --- +## EfficientSAM — Run official ONNX models + +EfficientSAM-Ti and EfficientSAM-S use a distinct lightweight architecture, so +select the `efficient_sam` runtime variant. The upstream project and models are +Apache-2.0 licensed. + +Download the Ti encoder and decoder (about 40 MB combined): + +```bash +mkdir -p output_models/efficient_sam +curl -L https://huggingface.co/yunyangx/EfficientSAM/resolve/main/efficientsam_ti_encoder.onnx \ + -o output_models/efficient_sam/efficientsam_ti_encoder.onnx +curl -L https://huggingface.co/yunyangx/EfficientSAM/resolve/main/efficientsam_ti_decoder.onnx \ + -o output_models/efficient_sam/efficientsam_ti_decoder.onnx +``` + +Run point or rectangle inference: + +```bash +python -m samexporter.inference \ + --sam_variant efficient_sam \ + --encoder_model output_models/efficient_sam/efficientsam_ti_encoder.onnx \ + --decoder_model output_models/efficient_sam/efficientsam_ti_decoder.onnx \ + --image images/truck.jpg \ + --prompt images/truck_box.json \ + --output output_images/efficient_sam_truck.png +``` + +The same runtime accepts the EfficientSAM-S split ONNX pair by changing the two +model paths to `efficientsam_s_encoder.onnx` and `efficientsam_s_decoder.onnx`. + +--- + ## SAM2 / SAM2.1 — Convert to ONNX ### 1. Download checkpoints @@ -156,12 +213,15 @@ original_models/ sam2.1_hiera_large.pt ``` -### 2. Install SAM2 PyTorch package +### 2. Install the pinned SAM2 PyTorch package ```bash -pip install git+https://github.com/facebookresearch/segment-anything-2.git +SAM2_BUILD_CUDA=0 pip install -e third_party/sam2 ``` +Omit `SAM2_BUILD_CUDA=0` if you specifically need SAM2's optional CUDA +post-processing extension. It is not required for these image ONNX exports. + ### 3. Export ```bash @@ -211,30 +271,47 @@ SAM3 exports into **three separate ONNX models**: an image encoder, a language ( ### Pre-exported ONNX models -Pre-exported models are available on HuggingFace and are downloaded automatically: +NRL.ai publishes the real-model artifacts validated by this repository at +[`nrl-ai/samexporter-onnx-models`](https://huggingface.co/nrl-ai/samexporter-onnx-models). +The repository contains SAM ViT-B, MobileSAM, EfficientSAM-Ti, SAM 2.1 Tiny, +and SAM3, with checksums, source revisions, and a license in every family +directory. +```bash +hf download nrl-ai/samexporter-onnx-models \ + --include "sam3/*" \ + --local-dir output_models ``` -vietanhdev/segment-anything-3-onnx-models - sam3_image_encoder.onnx (+ .data) - sam3_language_encoder.onnx (+ .data) - sam3_decoder.onnx (+ .data) -``` + +Keep `sam3_decoder.onnx.data` beside `sam3_decoder.onnx`. The published bundle +was exported from `facebook/sam3` revision +`3c879f39826c281e95690f02c7821c4de09afae7` and source revision +`660a5e9e1b8b4c02c0ad97229b88a09a6e4ff5b7`. ### Export from PyTorch (optional) +The pinned upstream SAM3 package supports Python 3.8 or newer and requires +NumPy `<2` for source export. SAMExporter itself supports Python 3.11 through +3.13, so use any of those versions with NumPy `>=1.26,<2` when exporting SAM3. + ```bash # Clone the SAM3 source (required for export only, not inference) git submodule update --init sam3 -# Install SAM3 dependencies -pip install osam +# Install the pinned official SAM3 export dependencies +pip install -e sam3 # Export (add --simplify for ONNX simplification, requires [export] extra on Windows) python -m samexporter.export_sam3 \ --output_dir output_models/sam3 \ - --opset 18 + --opset 18 \ + --device auto \ + --max-geometric-prompts 8 ``` +Pass `--checkpoint /path/to/sam3.pt` to use an already-downloaded official +checkpoint instead of consulting the Hugging Face cache. + ### Run inference **Text-only prompt** (detects all instances matching the text): @@ -252,7 +329,25 @@ python -m samexporter.inference \ --show ``` -**Text + rectangle prompt** (text guides detection, rectangle refines region): +For broad concepts such as `"plant"`, returning every visible matching object +is expected. Keep only the highest-confidence results when a compact preview is +more useful: + +```bash +python -m samexporter.inference \ + --sam_variant sam3 \ + --encoder_model output_models/sam3/sam3_image_encoder.onnx \ + --decoder_model output_models/sam3/sam3_decoder.onnx \ + --language_encoder_model output_models/sam3/sam3_language_encoder.onnx \ + --image images/plants.png \ + --prompt images/plants_text.json \ + --text_prompt "plant" \ + --max_instances 5 \ + --output visual_results/runs/manual/sam3_plants_text_top5.png +``` + +**Text + rectangle exemplar** (text names the concept; the rectangle supplies a +positive visual example): ```bash python -m samexporter.inference \ @@ -282,7 +377,36 @@ python -m samexporter.inference \ --show ``` -> **Note:** Always pass `--text_prompt` for SAM3. Without it the model defaults to a "visual" text token and may produce zero detections. +> **Note:** Put a text mark in the prompt JSON or pass `--text_prompt`. Without +> either, the model defaults to a generic `"visual"` token and results may be +> less predictable. + +SAM3 geometry is a positive/negative *concept exemplar*, not a crop constraint. +It can find visually matching objects outside the supplied rectangle. In the +default `--sam3_output_mode auto`, a geometric prompt is treated as a selection: +duplicate queries are removed and the best mask overlapping the positive box or +point is returned. Use `--sam3_output_mode all` to preserve every matching +concept instance. Point marks are approximated as 1%-size exemplar boxes because +the PCS decoder ignores native point prompts; prefer text plus rectangles when +quality matters. + +The default `--sam3_nms_mode mask` follows the official SAM3 mask-IoU NMS +behavior. `--sam3_nms_mode box` trades some fidelity for speed and +`--sam3_nms_mode none` exposes raw thresholded queries. Tune duplicate +suppression with `--nms_threshold` and score filtering with +`--confidence_threshold`. + +New decoders accept zero to eight padded geometric marks by default. The +capacity is fixed inside ONNX because SAM3 geometry attention traces at a fixed +token count; re-export with a larger `--max-geometric-prompts` if needed. The +runtime validates the limit and masks unused slots. + +New SAM3 exports keep raw query logits, presence score, masks, and boxes as ONNX +outputs. Confidence filtering and resize-to-original are performed at runtime, +so `--confidence_threshold`, original aspect ratio, and output dimensions are +not frozen during export. Multi-instance overlays use distinct colors, contours, +and instance numbers. The runtime also remains compatible with older +three-output exports. --- @@ -303,13 +427,61 @@ Prompts are JSON files containing a list of mark objects: --- -## Tips - -- Use **quantized** models (`*.quant.onnx`) for faster inference and smaller file size with minimal accuracy loss. -- **SAM ViT-B** is the fastest SAM1 variant; **SAM ViT-H** is the most accurate. -- **SAM2 Tiny / SAM2.1 Tiny** are good CPU-friendly choices for SAM2. -- **SAM3** is slower due to its three-model pipeline but uniquely supports natural-language object queries. -- Run the encoder once per image; the lightweight decoder handles prompt changes in real time. +## Performance recommendations + +Choose the model around the prompt and deployment constraint: + +| Need | Recommended starting point | Why | +|---|---|---| +| Small CPU package / interactive clicks | EfficientSAM-Ti | Roughly 40 MB for the split ONNX pair; lowest setup cost | +| SAM-compatible lightweight encoder | MobileSAM | Reuses the familiar SAM prompt/decoder contract | +| Best original-SAM quality | SAM ViT-H | Largest and slowest SAM1 encoder, but the strongest original checkpoint | +| Better modern image masks | SAM2.1 Tiny first, then scale up | Tiny is the practical baseline; larger Hiera encoders trade latency and memory for capacity | +| Natural-language object discovery | SAM3 | Only choose it when text prompting is required; it loads image, language, and decoder models | + +For interactive applications, encode each image once and cache the returned +embedding. Re-run only the decoder as points or boxes change. On the local +8-core Intel i9-11950H using ONNX Runtime CPU, EfficientSAM-Ti took a median +966 ms to encode the 1800×1200 truck image and 42 ms per box decode (10 measured +runs after warm-up). Treat these as a local baseline, not a portable benchmark. + +Additional guidance: + +- The runtime automatically orders every installed ONNX Runtime accelerator and + keeps CPU last as a fallback. Recognized backends include TensorRT, CUDA, + TensorRT RTX, MIGraphX/ROCm, OpenVINO, DirectML, CoreML, CANN, QNN, NNAPI, + VSINPU, WebNN, WebGPU, XNNPACK, RKNPU, Vitis AI, ACL, ArmNN, and oneDNN. + Availability depends on the ONNX Runtime package/build installed on the + target machine. +- Install the one ONNX Runtime build intended for the target. The standard + `onnxruntime` wheel is CPU; NVIDIA CUDA/TensorRT uses `onnxruntime-gpu`, Intel + OpenVINO uses `onnxruntime-openvino`, Windows DirectML uses + `onnxruntime-directml`, and Qualcomm QNN uses `onnxruntime-qnn`. ONNX Runtime + recommends keeping only one of its Python packages in an environment. Follow + its installation matrix for platform-specific and custom provider builds. +- Override provider order with `--providers tensorrt,cuda,cpu`, + `--providers openvino,cpu`, or the `SAMEXPORTER_ONNX_PROVIDERS` environment + variable. Exact ONNX Runtime provider names are also accepted, and unavailable + explicit providers fail early with the installed provider list. +- TensorRT is usable through `TensorrtExecutionProvider`. Dynamic models may + incur engine-build/cache time; for repeatable deployment, configure the + provider's engine cache and shape profiles or build fixed min/opt/max engines + as demonstrated by EfficientViT-SAM upstream. When building a SAM decoder + directly with TensorRT, `orig_im_size` is a *shape tensor*: configure its + value profile with `profile.set_shape_input("orig_im_size", min_values, + opt_values, max_values)` in addition to ordinary dynamic tensor profiles. +- Quantized `*.quant.onnx` files primarily reduce storage and memory. Benchmark + accuracy and latency on the target CPU—dynamic quantization is not guaranteed + to accelerate convolution-heavy encoders. +- Avoid ONNX simplification by default for large SAM2/SAM3 exports. It can increase + file size and peak memory; enable it only after measuring the target runtime. +- Use box prompts for deterministic QA. A single positive point can validly select + a sub-part, while a tight box better communicates the expected object extent. +- Keep image embeddings in memory, but release sessions or cached embeddings when + processing large image queues; the encoders dominate memory consumption. + +Reviewed overlays and their verification status are kept in +[`visual_results/README.md`](visual_results/README.md). --- @@ -320,6 +492,43 @@ pip install pytest pytest tests/ ``` +For end-to-end model checks, first run `bash download_all_models.sh`, then: + +```bash +bash test_comprehensive.sh +``` + +The bundled point and rectangle fixtures intentionally target visible objects in +both a landscape and a portrait image. Every generated image is retained under +`visual_results/runs/` by default. Use `SAMEXPORTER_RESULTS_DIR=/path/to/results` +to redirect a run without losing its artifacts. + +--- + +## Adding more SAM variants + +Good next candidates from the broader SAM ecosystem are: + +| Variant | License | Integration path | +|---|---|---| +| EfficientViT-SAM | Apache-2.0 | Upstream provides separate encoder/decoder ONNX exporters; add an adapter for its normalized fixed-size encoder | +| HQ-SAM / HQ-SAM 2 | Apache-2.0 | Add HQ intermediate features and its custom mask decoder contract | +| MedSAM | Apache-2.0 | Reuse the SAM prompt contract with model-specific preprocessing and checkpoints | +| EdgeSAM | NTU S-Lab License 1.0 | Technically close to SAM, but review redistribution terms before bundling code or weights | + +FastSAM is intentionally not a drop-in target here: it is detector-style rather +than an encoder/prompt-decoder SAM contract, and its AGPL licensing needs a +separate distribution decision. Prefer adapters that keep third-party model code +optional and consume upstream ONNX files; this makes new integrations faster and +keeps their licenses explicit. + +The implementation was cross-checked against the official SAM mask +resize–crop–resize path, MobileSAM and EfficientSAM ONNX exporters, official +SAM2 preprocessing/contracts, EfficientViT-SAM's ONNX/TensorRT deployment, and +the MIT-licensed `sam3-triton` raw-head/post-processing split. Reference code is +used as design evidence only unless its license and attribution permit direct +integration. + --- ## AnyLabeling diff --git a/convert_all_meta_sam.sh b/convert_all_meta_sam.sh index 47f9346..1e2e3a2 100644 --- a/convert_all_meta_sam.sh +++ b/convert_all_meta_sam.sh @@ -1,3 +1,7 @@ +#!/usr/bin/env bash + +set -euo pipefail + echo "Converting all models..." echo "Converting ViT-H models..." diff --git a/convert_all_meta_sam2.sh b/convert_all_meta_sam2.sh index 9140e4b..230c8da 100644 --- a/convert_all_meta_sam2.sh +++ b/convert_all_meta_sam2.sh @@ -1,3 +1,7 @@ +#!/usr/bin/env bash + +set -euo pipefail + echo "Converting all models..." echo "Converting SAM2-Hiera-Tiny models..." diff --git a/convert_all_models.sh b/convert_all_models.sh index c09a6bc..bbaaf05 100644 --- a/convert_all_models.sh +++ b/convert_all_models.sh @@ -1,7 +1,7 @@ #!/bin/bash # convert_all_models.sh -set -e +set -euo pipefail export KMP_DUPLICATE_LIB_OK=TRUE @@ -24,7 +24,7 @@ python -m samexporter.export_decoder --checkpoint original_models/sam_vit_b_01ec echo -e "\n=== Converting Mobile SAM ===" python -m samexporter.export_encoder --checkpoint original_models/mobile_sam.pt --output output_models/mobile_sam/mobile_sam.encoder.onnx --model-type mobile --quantize-out output_models/mobile_sam/mobile_sam.encoder.quant.onnx --use-preprocess -python -m samexporter.export_decoder --checkpoint original_models/sam_vit_h_4b8939.pth --output output_models/mobile_sam/sam_vit_h_4b8939.decoder.onnx --model-type vit_h --quantize-out output_models/mobile_sam/sam_vit_h_4b8939.decoder.quant.onnx --return-single-mask +python -m samexporter.export_decoder --checkpoint original_models/mobile_sam.pt --output output_models/mobile_sam/mobile_sam.decoder.onnx --model-type mobile --quantize-out output_models/mobile_sam/mobile_sam.decoder.quant.onnx --return-single-mask echo -e "\n=== Converting Segment Anything 2 (SAM 2) ===" diff --git a/convert_mobile_sam.sh b/convert_mobile_sam.sh index 6e2664f..28fddee 100644 --- a/convert_mobile_sam.sh +++ b/convert_mobile_sam.sh @@ -1,11 +1,15 @@ +#!/usr/bin/env bash + +set -euo pipefail + echo "Converting Mobile SAM..." python -m samexporter.export_encoder --checkpoint original_models/mobile_sam.pt \ --output output_models/mobile_sam/mobile_sam.encoder.onnx \ --model-type mobile \ --quantize-out output_models/mobile_sam/mobile_sam.encoder.quant.onnx \ --use-preprocess -python -m samexporter.export_decoder --checkpoint original_models/sam_vit_h_4b8939.pth \ - --output output_models/mobile_sam/sam_vit_h_4b8939.decoder.onnx \ - --model-type vit_h \ - --quantize-out output_models/mobile_sam/sam_vit_h_4b8939.decoder.quant.onnx \ +python -m samexporter.export_decoder --checkpoint original_models/mobile_sam.pt \ + --output output_models/mobile_sam/mobile_sam.decoder.onnx \ + --model-type mobile \ + --quantize-out output_models/mobile_sam/mobile_sam.decoder.quant.onnx \ --return-single-mask diff --git a/convert_sam3.sh b/convert_sam3.sh index 7ddc0ea..5c79d02 100755 --- a/convert_sam3.sh +++ b/convert_sam3.sh @@ -3,7 +3,7 @@ # # Requirements: # - sam3 submodule initialised: git submodule update --init sam3 -# - osam installed for CLIP tokenisation: pip install osam +# - pinned SAM3 dependencies installed: pip install -e sam3 # # Optional: pass --simplify to run onnxsim after export (reduces some # redundant ops; vision_pos_enc_0/1 may be removed from the decoder). @@ -12,19 +12,21 @@ set -euo pipefail OUTPUT_DIR="${1:-output_models/sam3}" SIMPLIFY="${SIMPLIFY:-}" +DEVICE="${SAMEXPORTER_EXPORT_DEVICE:-auto}" +MAX_PROMPTS="${SAMEXPORTER_MAX_GEOMETRIC_PROMPTS:-8}" echo "Exporting SAM3 ViT-H to ONNX → $OUTPUT_DIR" +export_args=( + --output_dir "$OUTPUT_DIR" + --opset 18 + --device "$DEVICE" + --max-geometric-prompts "$MAX_PROMPTS" +) if [ -n "$SIMPLIFY" ]; then - python -m samexporter.export_sam3 \ - --output_dir "$OUTPUT_DIR" \ - --opset 18 \ - --simplify -else - python -m samexporter.export_sam3 \ - --output_dir "$OUTPUT_DIR" \ - --opset 18 + export_args+=(--simplify) fi +python -m samexporter.export_sam3 "${export_args[@]}" echo "Done – models written to $OUTPUT_DIR/" echo " sam3_image_encoder.onnx" diff --git a/download_all_models.sh b/download_all_models.sh index 4a48569..cec16b8 100644 --- a/download_all_models.sh +++ b/download_all_models.sh @@ -1,7 +1,7 @@ #!/bin/bash # download_all_models.sh -set -e +set -euo pipefail OUT_DIR="original_models" mkdir -p "$OUT_DIR" @@ -13,7 +13,10 @@ download_file() { echo " [SKIP] Already exists: $dest" else echo " Downloading $dest ..." - curl -L "$url" -o "$dest" + mkdir -p "$(dirname "$dest")" + curl --fail --location --retry 3 --continue-at - \ + "$url" --output "$dest.part" + mv "$dest.part" "$dest" echo " [OK] $dest" fi } @@ -28,6 +31,12 @@ echo -e " === MobileSAM ===" download_file "https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt" "$OUT_DIR/mobile_sam.pt" +echo -e " +=== EfficientSAM-Ti (ONNX) ===" +mkdir -p output_models/efficient_sam +download_file "https://huggingface.co/nrl-ai/samexporter-onnx-models/resolve/main/efficient_sam_ti/efficientsam_ti_encoder.onnx" "output_models/efficient_sam/efficientsam_ti_encoder.onnx" +download_file "https://huggingface.co/nrl-ai/samexporter-onnx-models/resolve/main/efficient_sam_ti/efficientsam_ti_decoder.onnx" "output_models/efficient_sam/efficientsam_ti_decoder.onnx" + echo -e " === Segment Anything 2 (SAM 2) ===" download_file "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt" "$OUT_DIR/sam2_hiera_tiny.pt" @@ -42,16 +51,14 @@ download_file "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_h download_file "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt" "$OUT_DIR/sam2.1_hiera_large.pt" echo -e "\n=== Segment Anything 3 (SAM 3) ===" -# SAM3 is a zip file containing multiple components -if [ -f "output_models/sam3/sam3_image_encoder.onnx" ]; then - echo " [SKIP] SAM3 already exported or present in output_models/sam3" -else - # We download the zip if not already present - download_file "https://huggingface.co/vietanhdev/segment-anything-3-onnx-models/resolve/main/sam3_vit_h.zip" "$OUT_DIR/sam3_vit_h.zip" - echo " Extracting SAM3 models..." - mkdir -p output_models/sam3 - unzip -o "$OUT_DIR/sam3_vit_h.zip" -d output_models/sam3 -fi +# SAM3 uses one external-data companion file for its decoder. Check every file +# independently so an interrupted or partial prior download is repaired. +mkdir -p output_models/sam3 +SAMEXPORTER_HF="https://huggingface.co/nrl-ai/samexporter-onnx-models/resolve/main/sam3" +download_file "$SAMEXPORTER_HF/sam3_image_encoder.onnx" "output_models/sam3/sam3_image_encoder.onnx" +download_file "$SAMEXPORTER_HF/sam3_language_encoder.onnx" "output_models/sam3/sam3_language_encoder.onnx" +download_file "$SAMEXPORTER_HF/sam3_decoder.onnx" "output_models/sam3/sam3_decoder.onnx" +download_file "$SAMEXPORTER_HF/sam3_decoder.onnx.data" "output_models/sam3/sam3_decoder.onnx.data" echo -e " All downloads complete!" diff --git a/images/plants_box.json b/images/plants_box.json index cbf3760..b06b600 100644 --- a/images/plants_box.json +++ b/images/plants_box.json @@ -1,6 +1,6 @@ [ { "type": "rectangle", - "data": [200, 200, 600, 600] + "data": [325, 410, 520, 770] } ] diff --git a/images/plants_box_refined.json b/images/plants_box_refined.json new file mode 100644 index 0000000..5151d6e --- /dev/null +++ b/images/plants_box_refined.json @@ -0,0 +1,4 @@ +[ + {"type": "rectangle", "data": [350, 405, 510, 775]}, + {"type": "point", "data": [430, 690], "label": 1} +] diff --git a/images/plants_point.json b/images/plants_point.json index 012c535..750231d 100644 --- a/images/plants_point.json +++ b/images/plants_point.json @@ -1,7 +1,7 @@ [ { "type": "point", - "data": [400, 400], + "data": [410, 480], "label": 1 } ] diff --git a/images/plants_text.json b/images/plants_text.json index cbbfed4..ebe59a4 100644 --- a/images/plants_text.json +++ b/images/plants_text.json @@ -1,6 +1,6 @@ [ { "type": "text", - "data": "leaf" + "data": "plant" } ] diff --git a/images/truck_box.json b/images/truck_box.json index b79b3b3..d46832b 100644 --- a/images/truck_box.json +++ b/images/truck_box.json @@ -1,6 +1,6 @@ [ { "type": "rectangle", - "data": [74, 150, 617, 400] + "data": [120, 280, 1700, 850] } ] diff --git a/images/truck_point.json b/images/truck_point.json index ab1447e..a3cb573 100644 --- a/images/truck_point.json +++ b/images/truck_point.json @@ -1,7 +1,7 @@ [ { "type": "point", - "data": [500, 375], + "data": [1000, 550], "label": 1 } ] diff --git a/images/truck_sam3_box.json b/images/truck_sam3_box.json index b79b3b3..d46832b 100644 --- a/images/truck_sam3_box.json +++ b/images/truck_sam3_box.json @@ -1,6 +1,6 @@ [ { "type": "rectangle", - "data": [74, 150, 617, 400] + "data": [120, 280, 1700, 850] } ] diff --git a/images/truck_sam3_point.json b/images/truck_sam3_point.json index ab1447e..a3cb573 100644 --- a/images/truck_sam3_point.json +++ b/images/truck_sam3_point.json @@ -1,7 +1,7 @@ [ { "type": "point", - "data": [500, 375], + "data": [1000, 550], "label": 1 } ] diff --git a/infer_mobile_sam.sh b/infer_mobile_sam.sh index 6b6a113..d744e1f 100644 --- a/infer_mobile_sam.sh +++ b/infer_mobile_sam.sh @@ -1,7 +1,12 @@ +#!/bin/bash +set -euo pipefail + +OUT_DIR="${SAMEXPORTER_RESULTS_DIR:-visual_results/runs/manual}" +mkdir -p "$OUT_DIR" python -m samexporter.inference \ --encoder_model output_models/mobile_sam/mobile_sam.encoder.onnx \ - --decoder_model output_models/mobile_sam/sam_vit_h_4b8939.decoder.onnx \ + --decoder_model output_models/mobile_sam/mobile_sam.decoder.onnx \ --image images/plants.png \ --prompt images/plants_prompt1.json \ - --output output_images/plants_01.png \ - --show + --output "$OUT_DIR/mobile_sam_plants_01.png" \ + --show 2>&1 | tee "$OUT_DIR/mobile_sam_plants_01.log" diff --git a/infer_sam.sh b/infer_sam.sh index e03f67a..91ddd95 100644 --- a/infer_sam.sh +++ b/infer_sam.sh @@ -1,7 +1,13 @@ +#!/bin/bash +set -euo pipefail + +OUT_DIR="${SAMEXPORTER_RESULTS_DIR:-visual_results/runs/manual}" +mkdir -p "$OUT_DIR" + python -m samexporter.inference \ --encoder_model output_models/sam_vit_b_01ec64.encoder.quant.onnx \ --decoder_model output_models/sam_vit_b_01ec64.decoder.quant.onnx \ --image images/truck.jpg \ --prompt images/truck_prompt.json \ - --output output_images/truck.jpg \ - --show + --output "$OUT_DIR/sam_vit_b_truck.png" \ + --show 2>&1 | tee "$OUT_DIR/sam_vit_b_truck.log" diff --git a/infer_sam2.sh b/infer_sam2.sh index 60b5e98..1972b0f 100644 --- a/infer_sam2.sh +++ b/infer_sam2.sh @@ -1,8 +1,14 @@ +#!/bin/bash +set -euo pipefail + +OUT_DIR="${SAMEXPORTER_RESULTS_DIR:-visual_results/runs/manual}" +mkdir -p "$OUT_DIR" + python -m samexporter.inference \ --encoder_model output_models/sam2_hiera_tiny.encoder.onnx \ --decoder_model output_models/sam2_hiera_tiny.decoder.onnx \ --image images/truck.jpg \ --prompt images/truck_prompt.json \ - --output output_images/sam2_truck.png \ + --output "$OUT_DIR/sam2_truck.png" \ --sam_variant sam2 \ - --show + --show 2>&1 | tee "$OUT_DIR/sam2_truck.log" diff --git a/infer_sam3.sh b/infer_sam3.sh index 56b6c2f..56bad19 100755 --- a/infer_sam3.sh +++ b/infer_sam3.sh @@ -1,53 +1,52 @@ #!/bin/bash -# SAM3 inference examples. -# -# SAM3 supports three prompt modes: -# 1. Text only – open-vocabulary detection (no geometric hint needed) -# 2. Text + point – refine detection around a clicked pixel -# 3. Text + rectangle – constrain detection to a bounding box -# -# The --text_prompt flag drives the language encoder; always supply it for -# best results. If omitted the model falls back to "visual" (no language -# guidance) which may return fewer or no detections. +# Real SAM3 text, geometric-selection, and capped-discovery examples. +# Geometry is a concept exemplar. The default auto mode returns the best mask +# overlapping a point/rectangle; pass --sam3_output_mode all to keep every +# visually matching instance. set -euo pipefail ENC="output_models/sam3/sam3_image_encoder.onnx" DEC="output_models/sam3/sam3_decoder.onnx" LANG="output_models/sam3/sam3_language_encoder.onnx" -IMG="images/truck.jpg" - -echo "--- SAM3: text-only prompt ('truck') ---" -python -m samexporter.inference \ - --encoder_model "$ENC" \ - --decoder_model "$DEC" \ - --language_encoder_model "$LANG" \ - --image "$IMG" \ - --prompt images/truck_sam3.json \ - --text_prompt "truck" \ - --output output_images/sam3_truck_text.png \ - --sam_variant sam3 - -echo "--- SAM3: text + rectangle prompt ---" -python -m samexporter.inference \ - --encoder_model "$ENC" \ - --decoder_model "$DEC" \ - --language_encoder_model "$LANG" \ - --image "$IMG" \ - --prompt images/truck_sam3_box.json \ - --text_prompt "truck" \ - --output output_images/sam3_truck_box.png \ - --sam_variant sam3 - -echo "--- SAM3: text + point prompt ---" -python -m samexporter.inference \ - --encoder_model "$ENC" \ - --decoder_model "$DEC" \ - --language_encoder_model "$LANG" \ - --image "$IMG" \ - --prompt images/truck_sam3_point.json \ - --text_prompt "truck" \ - --output output_images/sam3_truck_point.png \ - --sam_variant sam3 - -echo "Done – outputs saved to output_images/" +OUT_DIR="${SAMEXPORTER_RESULTS_DIR:-visual_results/runs/manual/sam3}" +mkdir -p "$OUT_DIR" + +run_case() { + local name=$1 + local image=$2 + local prompt=$3 + shift 3 + python -m samexporter.inference \ + --sam_variant sam3 \ + --encoder_model "$ENC" \ + --decoder_model "$DEC" \ + --language_encoder_model "$LANG" \ + --image "$image" \ + --prompt "$prompt" \ + --output "$OUT_DIR/$name.png" \ + "$@" 2>&1 | tee "$OUT_DIR/$name.log" +} + +# A singular text prompt discovers the truck. +run_case truck_text images/truck.jpg images/truck_sam3.json \ + --text_prompt "truck" + +# Text plus geometry reliably chooses the intended matching instance. +run_case truck_text_box images/truck.jpg images/truck_sam3_box.json \ + --text_prompt "truck" +run_case truck_text_point images/truck.jpg images/truck_sam3_point.json \ + --text_prompt "truck" + +# A broad concept discovers every visible plant; cap it for previews. +run_case plants_text_all images/plants.png images/plants_text.json \ + --text_prompt "plant" +run_case plants_text_top5 images/plants.png images/plants_text.json \ + --text_prompt "plant" --max_instances 5 + +# Geometry-only auto mode uses the generic visual token, then returns the best +# prompt-overlapping match instead of every similar object in the image. +run_case plants_box images/plants.png images/plants_box.json +run_case plants_box_refined images/plants.png images/plants_box_refined.json + +echo "SAM3 images and sibling logs saved to $OUT_DIR" diff --git a/pyproject.toml b/pyproject.toml index df99afe..f6201e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,15 +4,16 @@ build-backend = "setuptools.build_meta" [project] name = "samexporter" -version = "0.4.6" -description = "Exporting Segment Anything models ONNX format" +version = "0.5.0" +description = "Export and run SAM, MobileSAM, EfficientSAM, SAM2, and SAM3 with ONNX Runtime" authors = [ {name = "Viet Anh Nguyen", email = "vietanh.dev@gmail.com"}, ] readme = "README.md" requires-python = ">=3.11" license = "MIT" -license-files = ["LICENSE"] +license-files = ["LICENSE", "third_party_licenses/*"] +keywords = ["sam", "segment-anything", "onnx", "onnxruntime", "segmentation"] classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -22,40 +23,56 @@ classifiers = [ dependencies = [ "onnx==1.20.1", - "onnxruntime==1.24.2", "opencv-python==4.11.0.86", "segment-anything==1.0", "torch==2.10.0", "torchvision==0.25.0", - "timm==0.9.2", - "numpy==1.26.4", + "timm>=0.9.16,<2", + "numpy>=1.26,<3", "onnxscript==0.6.2", - "osam", # CLIP tokeniser for SAM3 text prompts + "ftfy>=6.1,<7", + "regex>=2023.0.0", ] [project.urls] "Homepage" = "https://github.com/vietanhdev/samexporter" "Bug Tracker" = "https://github.com/vietanhdev/samexporter/issues" +"Documentation" = "https://anylearning-oss.nrl.ai/docs/samexporter" +"Models" = "https://huggingface.co/nrl-ai/samexporter-onnx-models" [project.optional-dependencies] # onnxsim has no Windows wheel; the source tarball exceeds Windows 260-char # path limits. Install with: pip install samexporter[export] export = ["onnxsim==0.5.0"] -dev = ["ruff", "pre-commit", "pytest"] +runtime-cpu = ["onnxruntime==1.24.2"] +runtime-gpu = ["onnxruntime-gpu>=1.24,<2"] +runtime-openvino = ["onnxruntime-openvino>=1.23,<2"] +runtime-directml = ["onnxruntime-directml>=1.23,<2"] +runtime-qnn = ["onnxruntime-qnn>=1.23,<2"] +dev = ["onnxruntime==1.24.2", "ruff", "pre-commit", "pytest"] [tool.setuptools] -packages = ["samexporter", "samexporter.mobile_encoder", "samexporter.sam2_configs"] +packages = [ + "samexporter", + "samexporter.assets", + "samexporter.mobile_encoder", + "samexporter.sam2_configs", +] [tool.setuptools.package-data] "samexporter.sam2_configs" = ["*.yaml", "sam2.1/*.yaml"] +"samexporter.assets" = ["*.gz"] [tool.pytest.ini_options] testpaths = ["tests"] +pythonpath = ["."] [tool.ruff] line-length = 88 target-version = "py311" -exclude = [".git", "__pycache__", "build", "dist", "*.egg-info", "sam3"] +exclude = [ + ".git", "__pycache__", "build", "dist", "*.egg-info", "sam3", "third_party", +] [tool.ruff.lint] select = ["E", "W", "F", "I", "UP"] diff --git a/sam3 b/sam3 index 2d08d73..660a5e9 160000 --- a/sam3 +++ b/sam3 @@ -1 +1 @@ -Subproject commit 2d08d7312050b6f65b28f1b458cd91f8cd57814a +Subproject commit 660a5e9e1b8b4c02c0ad97229b88a09a6e4ff5b7 diff --git a/samexporter/assets/__init__.py b/samexporter/assets/__init__.py new file mode 100644 index 0000000..da8530e --- /dev/null +++ b/samexporter/assets/__init__.py @@ -0,0 +1 @@ +"""Static resources bundled with SAMExporter.""" diff --git a/samexporter/assets/bpe_simple_vocab_16e6.txt.gz b/samexporter/assets/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000..7b5088a Binary files /dev/null and b/samexporter/assets/bpe_simple_vocab_16e6.txt.gz differ diff --git a/samexporter/clip_tokenizer.py b/samexporter/clip_tokenizer.py new file mode 100644 index 0000000..4079fca --- /dev/null +++ b/samexporter/clip_tokenizer.py @@ -0,0 +1,131 @@ +"""Small NumPy CLIP tokenizer used by the SAM3 language encoder. + +Adapted from OpenAI CLIP's MIT-licensed ``simple_tokenizer.py`` and +``tokenize`` helper. See ``third_party_licenses/OPENAI_CLIP_LICENSE``. +""" + +import gzip +import html +from functools import lru_cache +from pathlib import Path + +import ftfy +import numpy as np +import regex + + +@lru_cache +def bytes_to_unicode() -> dict[int, str]: + values = list(range(ord("!"), ord("~") + 1)) + values += list(range(ord("¡"), ord("¬") + 1)) + values += list(range(ord("®"), ord("ÿ") + 1)) + characters = values[:] + extra = 0 + for value in range(256): + if value not in values: + values.append(value) + characters.append(256 + extra) + extra += 1 + return dict(zip(values, map(chr, characters))) + + +def _pairs(word: tuple[str, ...]) -> set[tuple[str, str]]: + return set(zip(word, word[1:])) + + +class CLIPTokenizer: + def __init__(self, bpe_path: str | Path | None = None) -> None: + if bpe_path is None: + bpe_path = Path(__file__).parent / "assets" / "bpe_simple_vocab_16e6.txt.gz" + self.byte_encoder = bytes_to_unicode() + with gzip.open(bpe_path, "rt", encoding="utf-8") as bpe_file: + merges = bpe_file.read().splitlines()[1 : 49152 - 256 - 2 + 1] + merge_pairs = [tuple(merge.split()) for merge in merges] + vocab = list(self.byte_encoder.values()) + vocab += [value + "" for value in vocab] + vocab += ["".join(merge) for merge in merge_pairs] + vocab += ["<|startoftext|>", "<|endoftext|>"] + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.bpe_ranks = dict(zip(merge_pairs, range(len(merge_pairs)))) + self.cache = { + "<|startoftext|>": "<|startoftext|>", + "<|endoftext|>": "<|endoftext|>", + } + self.pattern = regex.compile( + r"<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|" + r"[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+", + regex.IGNORECASE, + ) + + def _bpe(self, token: str) -> str: + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + (token[-1] + "",) + pairs = _pairs(word) + if not pairs: + return token + "" + while True: + first, second = min( + pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")) + ) + if (first, second) not in self.bpe_ranks: + break + merged = [] + index = 0 + while index < len(word): + try: + next_index = word.index(first, index) + except ValueError: + merged.extend(word[index:]) + break + merged.extend(word[index:next_index]) + index = next_index + if index < len(word) - 1 and word[index + 1] == second: + merged.append(first + second) + index += 2 + else: + merged.append(word[index]) + index += 1 + word = tuple(merged) + if len(word) == 1: + break + pairs = _pairs(word) + encoded = " ".join(word) + self.cache[token] = encoded + return encoded + + def encode(self, text: str) -> list[int]: + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + text = regex.sub(r"\s+", " ", text).strip().lower() + result = [] + for token in regex.findall(self.pattern, text): + encoded = "".join(self.byte_encoder[value] for value in token.encode()) + result.extend( + self.encoder[piece] for piece in self._bpe(encoded).split(" ") + ) + return result + + def tokenize(self, texts: str | list[str], context_length: int = 32) -> np.ndarray: + if isinstance(texts, str): + texts = [texts] + start = self.encoder["<|startoftext|>"] + end = self.encoder["<|endoftext|>"] + result = np.zeros((len(texts), context_length), dtype=np.int64) + for index, text in enumerate(texts): + tokens = [start, *self.encode(text), end] + if len(tokens) > context_length: + raise ValueError( + f"Text prompt is too long for SAM3's {context_length}-token limit" + ) + result[index, : len(tokens)] = tokens + return result + + +@lru_cache +def _default_tokenizer() -> CLIPTokenizer: + return CLIPTokenizer() + + +def tokenize(texts: str | list[str], context_length: int = 32) -> np.ndarray: + return _default_tokenizer().tokenize(texts, context_length=context_length) diff --git a/samexporter/efficient_sam_onnx.py b/samexporter/efficient_sam_onnx.py new file mode 100644 index 0000000..5840bb8 --- /dev/null +++ b/samexporter/efficient_sam_onnx.py @@ -0,0 +1,77 @@ +from typing import Any + +import cv2 +import numpy as np +import onnxruntime + +from samexporter.prompts import geometric_prompt_arrays +from samexporter.runtime import get_onnx_providers + + +class EfficientSAMONNX: + """EfficientSAM-S/Ti inference using official split ONNX models. + + The model contract follows the Apache-2.0 EfficientSAM reference exporter: + a dynamic NCHW RGB image encoder and a prompt decoder accepting one or more + point/box marks. + """ + + def __init__( + self, encoder_model_path: str, decoder_model_path: str, providers=None + ) -> None: + providers = get_onnx_providers(providers) + self.encoder_session = onnxruntime.InferenceSession( + encoder_model_path, providers=providers + ) + self.decoder_session = onnxruntime.InferenceSession( + decoder_model_path, providers=providers + ) + self.encoder_input_name = self.encoder_session.get_inputs()[0].name + self.decoder_input_names = { + model_input.name for model_input in self.decoder_session.get_inputs() + } + + def encode(self, cv_image: np.ndarray) -> dict[str, Any]: + if cv_image is None or cv_image.ndim != 3 or cv_image.shape[2] != 3: + raise ValueError("Expected a non-empty BGR image with three channels") + rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) + input_image = rgb_image.transpose(2, 0, 1)[None].astype(np.float32) / 255.0 + image_embedding = self.encoder_session.run( + None, {self.encoder_input_name: input_image} + )[0] + return { + "image_embedding": image_embedding, + "original_size": cv_image.shape[:2], + } + + def predict_masks(self, embedding: dict[str, Any], prompt) -> np.ndarray: + points, labels = geometric_prompt_arrays(prompt) + decoder_inputs = { + "image_embeddings": embedding["image_embedding"], + "batched_point_coords": points[None, None], + "batched_point_labels": labels[None, None], + "orig_im_size": np.asarray(embedding["original_size"], dtype=np.int64), + } + missing = self.decoder_input_names - decoder_inputs.keys() + if missing: + raise ValueError( + "Unsupported EfficientSAM decoder inputs: " + ", ".join(sorted(missing)) + ) + + outputs = self.decoder_session.run( + None, + {name: decoder_inputs[name] for name in self.decoder_input_names}, + ) + masks, iou_predictions = outputs[:2] + # Official shape: masks [B, queries, candidates, H, W], IoU + # [B, queries, candidates]. Select the best candidate per query. + if masks.ndim != 5 or iou_predictions.ndim != 3: + raise ValueError( + "Unexpected EfficientSAM decoder output shapes: " + f"{masks.shape}, {iou_predictions.shape}" + ) + best_indices = np.argmax(iou_predictions[0], axis=-1) + best_masks = np.stack( + [masks[0, query, best] for query, best in enumerate(best_indices)] + ) + return best_masks[:, None] diff --git a/samexporter/export_decoder.py b/samexporter/export_decoder.py index 36fc511..b77c37b 100644 --- a/samexporter/export_decoder.py +++ b/samexporter/export_decoder.py @@ -9,8 +9,14 @@ import warnings import torch -from segment_anything import sam_model_registry -from segment_anything.utils.onnx import SamOnnxModel + +from samexporter.mobile_encoder.setup_mobile_sam import setup_model +from samexporter.upstream import prefer_pinned_upstream + +prefer_pinned_upstream("sam1") + +from segment_anything import sam_model_registry # noqa: E402 +from segment_anything.utils.onnx import SamOnnxModel # noqa: E402 try: import onnxruntime # type: ignore @@ -39,9 +45,9 @@ parser.add_argument( "--model-type", - type=str, + choices=("default", "vit_h", "vit_l", "vit_b", "mobile"), required=True, - help="In ['default', 'vit_h', 'vit_l', 'vit_b']. " + help="In ['default', 'vit_h', 'vit_l', 'vit_b', 'mobile']. " "Which type of SAM model to export.", ) @@ -113,7 +119,12 @@ def run_export( return_extra_metrics=False, ): print("Loading model...") - sam = sam_model_registry[model_type](checkpoint=checkpoint) + if model_type == "mobile": + state_dict = torch.load(checkpoint, map_location="cpu", weights_only=True) + sam = setup_model() + sam.load_state_dict(state_dict, strict=True) + else: + sam = sam_model_registry[model_type](checkpoint=checkpoint) onnx_model = SamOnnxModel( model=sam, @@ -202,6 +213,7 @@ def to_numpy(tensor): from onnxruntime.quantization.quantize import quantize_dynamic # type: ignore print(f"Quantizing model and writing to {args.quantize_out}...") + pathlib.Path(args.quantize_out).parent.mkdir(parents=True, exist_ok=True) quantize_dynamic( model_input=args.output, model_output=args.quantize_out, diff --git a/samexporter/export_encoder.py b/samexporter/export_encoder.py index 0c9ee21..8d1a863 100644 --- a/samexporter/export_encoder.py +++ b/samexporter/export_encoder.py @@ -8,10 +8,14 @@ import onnx import torch from onnx.external_data_helper import convert_model_to_external_data -from segment_anything import sam_model_registry from samexporter.mobile_encoder.setup_mobile_sam import setup_model from samexporter.onnx_utils import ImageEncoderOnnxModel +from samexporter.upstream import prefer_pinned_upstream + +prefer_pinned_upstream("sam1") + +from segment_anything import sam_model_registry # noqa: E402 parser = argparse.ArgumentParser( description="Export the SAM image encoder to an ONNX model." @@ -33,7 +37,7 @@ parser.add_argument( "--model-type", - type=str, + choices=("default", "vit_h", "vit_l", "vit_b", "mobile"), required=True, help="In ['default', 'vit_h', 'vit_l', 'vit_b', 'mobile']. " "Which type of SAM model to export.", @@ -42,7 +46,7 @@ parser.add_argument( "--use-preprocess", action="store_true", - help=("Embed pre-processing into the model",), + help="Embed pre-processing into the model", ) parser.add_argument( @@ -83,7 +87,7 @@ def run_export( ): print("Loading model...") if model_type == "mobile": - checkpoint = torch.load(checkpoint, map_location="cpu") + checkpoint = torch.load(checkpoint, map_location="cpu", weights_only=True) sam = setup_model() sam.load_state_dict(checkpoint, strict=True) else: @@ -122,6 +126,7 @@ def run_export( output_names = ["image_embeddings"] onnx_base = os.path.splitext(os.path.basename(output))[0] + pathlib.Path(output).parent.mkdir(parents=True, exist_ok=True) with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=torch.jit.TracerWarning) warnings.filterwarnings("ignore", category=UserWarning) @@ -143,7 +148,6 @@ def run_export( ) # Combine the weights into a single file - pathlib.Path(output).parent.mkdir(parents=True, exist_ok=True) onnx_model = onnx.load(tmp_model_path) convert_model_to_external_data( onnx_model, @@ -193,6 +197,7 @@ def to_numpy(tensor): from onnxruntime.quantization.quantize import quantize_dynamic # type: ignore print(f"Quantizing model and writing to {args.quantize_out}...") + pathlib.Path(args.quantize_out).parent.mkdir(parents=True, exist_ok=True) quantize_dynamic( model_input=args.output, model_output=args.quantize_out, diff --git a/samexporter/export_sam2.py b/samexporter/export_sam2.py index 6ebd1e4..e6edd25 100644 --- a/samexporter/export_sam2.py +++ b/samexporter/export_sam2.py @@ -4,10 +4,26 @@ import onnx import torch -from sam2.build_sam import build_sam2 -from sam2.modeling.sam2_base import SAM2Base from torch import nn +from samexporter.upstream import prefer_pinned_upstream + +prefer_pinned_upstream("sam2") + +from sam2.build_sam import build_sam2 # noqa: E402 +from sam2.modeling.sam2_base import SAM2Base # noqa: E402 + +MODEL_CONFIGS = { + "sam2_hiera_tiny": "sam2_hiera_t.yaml", + "sam2_hiera_small": "sam2_hiera_s.yaml", + "sam2_hiera_base_plus": "sam2_hiera_b+.yaml", + "sam2_hiera_large": "sam2_hiera_l.yaml", + "sam2.1_hiera_tiny": "sam2.1/sam2.1_hiera_t.yaml", + "sam2.1_hiera_small": "sam2.1/sam2.1_hiera_s.yaml", + "sam2.1_hiera_base_plus": "sam2.1/sam2.1_hiera_b+.yaml", + "sam2.1_hiera_large": "sam2.1/sam2.1_hiera_l.yaml", +} + class SAM2ImageEncoder(nn.Module): def __init__(self, sam_model: SAM2Base) -> None: @@ -84,7 +100,7 @@ def forward( masks = masks[:, 1:, :, :] iou_predictions = iou_predictions[:, 1:] else: - masks, iou_pred = self.mask_decoder._dynamic_multimask_via_stability( + masks, iou_predictions = self.mask_decoder._dynamic_multimask_via_stability( masks, iou_predictions ) @@ -165,7 +181,7 @@ def _embed_masks( parser.add_argument( "--model_type", - type=str, + choices=tuple(MODEL_CONFIGS), required=True, help="SAM2 model type: sam2_hiera_{tiny,small,base_plus,large} or sam2.1_hiera_{tiny,small,base_plus,large}.", ) @@ -187,25 +203,7 @@ def _embed_masks( input_size = (1024, 1024) multimask_output = True - model_type = args.model_type - if model_type == "sam2_hiera_tiny": - model_cfg = "sam2_hiera_t.yaml" - elif model_type == "sam2_hiera_small": - model_cfg = "sam2_hiera_s.yaml" - elif model_type == "sam2_hiera_base_plus": - model_cfg = "sam2_hiera_b+.yaml" - elif model_type == "sam2_hiera_large": - model_cfg = "sam2_hiera_l.yaml" - elif model_type == "sam2.1_hiera_tiny": - model_cfg = "sam2.1/sam2.1_hiera_t.yaml" - elif model_type == "sam2.1_hiera_small": - model_cfg = "sam2.1/sam2.1_hiera_s.yaml" - elif model_type == "sam2.1_hiera_base_plus": - model_cfg = "sam2.1/sam2.1_hiera_b+.yaml" - elif model_type == "sam2.1_hiera_large": - model_cfg = "sam2.1/sam2.1_hiera_l.yaml" - else: - model_cfg = "sam2_hiera_l.yaml" + model_cfg = MODEL_CONFIGS[args.model_type] # Register the config directory with Hydra import os diff --git a/samexporter/export_sam3.py b/samexporter/export_sam3.py index 5bccd90..640953d 100644 --- a/samexporter/export_sam3.py +++ b/samexporter/export_sam3.py @@ -2,17 +2,26 @@ import os import pathlib import sys -from unittest.mock import MagicMock +import types import onnx import torch from torchvision.transforms import v2 -# Mock triton for Windows – must happen before any sam3 imports. -mock_triton = MagicMock() -sys.modules["triton"] = mock_triton -sys.modules["triton.language"] = MagicMock() -sys.modules["torch._inductor.runtime.triton_helpers"] = MagicMock() +from samexporter.upstream import prefer_pinned_upstream + +# SAM3 can import optional Triton helpers on Windows, where Triton is normally +# unavailable. Never replace a real Linux Triton installation: doing so breaks +# PyTorch CUDA's lazy kernel registration. +if os.name == "nt": + from importlib.machinery import ModuleSpec + from unittest.mock import MagicMock + + mock_triton = MagicMock() + mock_triton.__spec__ = ModuleSpec("triton", loader=None) + sys.modules.setdefault("triton", mock_triton) + sys.modules.setdefault("triton.language", MagicMock()) + sys.modules.setdefault("torch._inductor.runtime.triton_helpers", MagicMock()) # Ensure sam3 is in PYTHONPATH @@ -21,11 +30,10 @@ # Submodule is at samexporter/sam3. # Package 'sam3' is at samexporter/sam3/sam3. # So we add samexporter/sam3 to sys.path. -sys.path.append(os.path.join(samexporter_root, "sam3")) +prefer_pinned_upstream("sam3") sys.path.append(samexporter_root) try: - from osam._models.yoloworld.clip import tokenize from sam3.model.sam3_image import Sam3Image from sam3.model.sam3_image_processor import Sam3Processor from sam3.model_builder import build_sam3_image_model @@ -38,25 +46,44 @@ def build_sam3_image_model(): return None - def tokenize(x): - return None - -def get_replace_freqs_cis(module: torch.nn.Module) -> None: +def prepare_rope_buffers_for_onnx(module: torch.nn.Module) -> None: """Replace complex freqs_cis buffers with separate real/imag float buffers. ONNX does not support complex-valued tensors, so the complex RoPE (rotary positional embedding) buffer must be split into its real (cosine) and imaginary (sine) components before export. """ - if hasattr(module, "freqs_cis"): - freqs_cos = module.freqs_cis.real.float() - freqs_sin = module.freqs_cis.imag.float() - module.register_buffer("freqs_cos", freqs_cos) - module.register_buffer("freqs_sin", freqs_sin) - del module.freqs_cis + # Current SAM3 already registers real/imaginary buffers. Keep freqs_cis: + # upstream _apply_rope still asserts that the complex buffer exists. + freqs_cis = getattr(module, "freqs_cis", None) + if freqs_cis is not None: + if not hasattr(module, "freqs_cis_real"): + module.register_buffer("freqs_cis_real", freqs_cis.real.float()) + module.register_buffer("freqs_cis_imag", freqs_cis.imag.float()) + if hasattr(module, "use_rope_real"): + module.use_rope_real = True for child in module.children(): - get_replace_freqs_cis(child) + prepare_rope_buffers_for_onnx(child) + + +def prepare_fused_mlps_for_onnx(module: torch.nn.Module) -> None: + """Replace SAM3's inference-only BF16 fused MLP with exportable FP32 ops.""" + + def forward(mlp, value): + value = mlp.fc1(value) + value = mlp.act(value) + value = mlp.drop1(value) + value = mlp.norm(value) + value = mlp.fc2(value) + return mlp.drop2(value) + + for child in module.modules(): + if ( + child.__class__.__name__ == "Mlp" + and child.__class__.__module__ == "sam3.model.vitdet" + ): + child.forward = types.MethodType(forward, child) class SAM3ImageEncoder(torch.nn.Module): @@ -72,7 +99,10 @@ class SAM3ImageEncoder(torch.nn.Module): def __init__(self, processor: Sam3Processor) -> None: super().__init__() - self._processor: Sam3Processor = processor + # Register the backbone as a real child module. Keeping it reachable + # only through the non-Module processor makes the exporter treat every + # trainable weight as an invalid requires-grad constant. + self._backbone = processor.model.backbone # Normalise uint8 [0,255] to float [-1, 1] – identical to the # reference sam3-onnx export (export_onnx.py). self._transform = v2.Compose( @@ -82,10 +112,11 @@ def __init__(self, processor: Sam3Processor) -> None: ] ) + @torch.no_grad() def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, ...]: # image: (3, H, W) uint8 → normalise → (1, 3, H, W) float image = self._transform(image).unsqueeze(0) - backbone_out = self._processor.model.backbone._forward_image_no_act_ckpt(image) + backbone_out = self._backbone._forward_image_no_act_ckpt(image) # Remove keys that are not needed by the decoder and would add # unnecessary overhead to the ONNX graph. backbone_out.pop("vision_features", None) @@ -98,12 +129,13 @@ def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, ...]: class SAM3LanguageEncoder(torch.nn.Module): def __init__(self, processor: Sam3Processor) -> None: super().__init__() - self._processor: Sam3Processor = processor + self._model: Sam3Image = processor.model + @torch.no_grad() def forward( self, tokens: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - model: Sam3Image = self._processor.model + model = self._model # VETextEncoder forward pass text_attention_mask = (tokens != 0).bool() @@ -126,7 +158,13 @@ def __init__(self, model: Sam3Image, processor: Sam3Processor) -> None: super().__init__() self._model = model self._processor = processor + # PCS ignores point prompts, but its geometry encoder still needs a + # well-shaped point sequence when boxes are padded to a fixed capacity. + self.register_buffer("_point_embedding", torch.zeros(1, 1, 2)) + self.register_buffer("_point_mask", torch.ones(1, 1, dtype=torch.bool)) + self.register_buffer("_point_label", torch.ones(1, 1, dtype=torch.long)) + @torch.no_grad() def forward( self, original_height: torch.Tensor, @@ -148,6 +186,9 @@ def forward( geometric_prompt.box_embeddings = box_coords geometric_prompt.box_labels = box_labels geometric_prompt.box_mask = box_masks + geometric_prompt.point_embeddings = self._point_embedding + geometric_prompt.point_labels = self._point_label + geometric_prompt.point_mask = self._point_mask state = { "original_height": original_height, "original_width": original_width, @@ -168,20 +209,56 @@ def forward( }, "geometric_prompt": geometric_prompt, } - result = self._processor._forward_grounding(state) - return result["boxes"], result["scores"], result["masks"] + # Export raw heads. Keeping thresholding and original-resolution mask + # resizing outside ONNX allows runtime confidence and image dimensions + # to remain fully dynamic, matching the official Sam3Processor logic. + result = self._model.forward_grounding( + backbone_out=state["backbone_out"], + find_input=self._processor.find_stage, + geometric_prompt=state["geometric_prompt"], + find_target=None, + ) + return ( + result["pred_boxes"], + result["pred_logits"], + result["pred_masks"], + result["presence_logit_dec"], + ) -def export_sam3(output_dir: str, opset: int = 18, simplify_model: bool = False): +def export_sam3( + output_dir: str, + opset: int = 18, + simplify_model: bool = False, + checkpoint_path: str | None = None, + device: str = "auto", + max_geometric_prompts: int = 8, +): + if max_geometric_prompts < 1: + raise ValueError("max_geometric_prompts must be at least 1") output_dir = pathlib.Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) - model = build_sam3_image_model() + if device == "auto": + if torch.cuda.is_available(): + device = "cuda" + elif torch.backends.mps.is_available(): + device = "mps" + else: + device = "cpu" + model = build_sam3_image_model( + checkpoint_path=checkpoint_path, + load_from_HF=checkpoint_path is None, + device=device, + ) # Replace complex RoPE buffers with float cos/sin – required for ONNX. - get_replace_freqs_cis(model) + prepare_rope_buffers_for_onnx(model) + # The current upstream ViT uses a CUDA BF16-only fused addmm/GELU helper. + # Standard FP32 layers preserve the same MLP semantics and are portable to + # ONNX Runtime providers. + prepare_fused_mlps_for_onnx(model) processor = Sam3Processor(model) - device = "cuda" if torch.cuda.is_available() else "cpu" - model.to(device) + model.eval().to(device) # ── Image Encoder ──────────────────────────────────────────────────────── print("Exporting Image Encoder...") @@ -231,55 +308,74 @@ def export_sam3(output_dir: str, opset: int = 18, simplify_model: bool = False): # ── Decoder ────────────────────────────────────────────────────────────── print("Exporting Decoder...") - decoder = SAM3Decoder(model, processor) + decoder = SAM3Decoder(model, processor).eval().to(device) decoder_path = output_dir / "sam3_decoder.onnx" - box_coords = torch.zeros(1, 1, 4).to(device) - box_labels = torch.ones(1, 1, dtype=torch.long).to(device) - # box_masks=True means "no real box" (dummy / masked out). - box_masks = torch.ones(1, 1, dtype=torch.bool).to(device) + # Geometry is sequence-first: [num_marks, batch, coordinates]. + box_coords = torch.zeros(max_geometric_prompts, 1, 4).to(device) + box_labels = torch.ones(max_geometric_prompts, 1, dtype=torch.long).to(device) + # Trace one real neutral slot plus right-padding. Tracing with every slot + # masked can make upstream attention hit an all-masked softmax/FPE. + box_coords[0, 0] = torch.tensor([0.5, 0.5, 0.01, 0.01], device=device) + box_masks = torch.ones(1, max_geometric_prompts, dtype=torch.bool).to(device) + box_masks[0, 0] = False orig_h = torch.tensor(1008).to(device) orig_w = torch.tensor(1008).to(device) - torch.onnx.utils.export( - decoder, - args=( - orig_h, - orig_w, - vpe0, - vpe1, - vpe2, - fpn0, - fpn1, - fpn2, - l_mask, - l_feat, - l_embed, - box_coords, - box_labels, - box_masks, - ), - f=str(decoder_path), - export_params=True, - input_names=[ - "original_height", - "original_width", - "vision_pos_enc_0", - "vision_pos_enc_1", - "vision_pos_enc_2", - "backbone_fpn_0", - "backbone_fpn_1", - "backbone_fpn_2", - "language_mask", - "language_features", - "language_embeds", - "box_coords", - "box_labels", - "box_masks", - ], - output_names=["boxes", "scores", "masks"], - opset_version=opset, - ) + # The legacy tracer crashes in torchvision's multi-box RoIAlign symbolic. + # Dynamo handles the fixed padded geometry path and produces portable ONNX. + original_pin_memory = torch.Tensor.pin_memory + original_is_dynamo_compiling = torch.compiler.is_dynamo_compiling + torch.Tensor.pin_memory = lambda tensor: tensor + torch.compiler.is_dynamo_compiling = lambda: True + try: + torch.onnx.export( + decoder, + args=( + orig_h, + orig_w, + vpe0, + vpe1, + vpe2, + fpn0, + fpn1, + fpn2, + l_mask, + l_feat, + l_embed, + box_coords, + box_labels, + box_masks, + ), + f=str(decoder_path), + input_names=[ + "original_height", + "original_width", + "vision_pos_enc_0", + "vision_pos_enc_1", + "vision_pos_enc_2", + "backbone_fpn_0", + "backbone_fpn_1", + "backbone_fpn_2", + "language_mask", + "language_features", + "language_embeds", + "box_coords", + "box_labels", + "box_masks", + ], + output_names=[ + "pred_boxes", + "pred_logits", + "pred_masks", + "presence_logit_dec", + ], + opset_version=opset, + dynamo=True, + ) + finally: + torch.Tensor.pin_memory = original_pin_memory + torch.compiler.is_dynamo_compiling = original_is_dynamo_compiling print(f"Saved Decoder to {decoder_path}") # ── Simplify models conditionally ───────────────────────────────────────── @@ -308,5 +404,29 @@ def export_sam3(output_dir: str, opset: int = 18, simplify_model: bool = False): ) parser.add_argument("--opset", type=int, default=18, help="ONNX opset version") parser.add_argument("--simplify", action="store_true", help="Simplify ONNX models") + parser.add_argument( + "--checkpoint", + default=None, + help="Local official SAM3 checkpoint; otherwise download from Hugging Face", + ) + parser.add_argument( + "--device", + choices=("auto", "cpu", "cuda", "mps"), + default="auto", + help="PyTorch device used during export", + ) + parser.add_argument( + "--max-geometric-prompts", + type=int, + default=8, + help="Fixed padded capacity for rectangle/point prompts (default: 8)", + ) args = parser.parse_args() - export_sam3(args.output_dir, args.opset, args.simplify) + export_sam3( + args.output_dir, + args.opset, + args.simplify, + args.checkpoint, + args.device, + args.max_geometric_prompts, + ) diff --git a/samexporter/inference.py b/samexporter/inference.py index 09c2663..c5ed323 100644 --- a/samexporter/inference.py +++ b/samexporter/inference.py @@ -1,155 +1,233 @@ import argparse import json import pathlib -import sys - -sys.path.append(".") import cv2 import numpy as np +from samexporter.efficient_sam_onnx import EfficientSAMONNX from samexporter.sam2_onnx import SegmentAnything2ONNX from samexporter.sam3_onnx import SegmentAnything3ONNX from samexporter.sam_onnx import SegmentAnythingONNX -def str2bool(v): - return v.lower() in ("true", "1") - - -argparser = argparse.ArgumentParser() -argparser.add_argument( - "--encoder_model", - type=str, - default="output_models/sam_vit_h_4b8939.encoder.onnx", - help="Path to the ONNX encoder model", -) -argparser.add_argument( - "--decoder_model", - type=str, - default="output_models/sam_vit_h_4b8939.decoder.onnx", - help="Path to the ONNX decoder model", -) -argparser.add_argument( - "--language_encoder_model", - type=str, - default=None, - help="Path to the ONNX language encoder model (for SAM3)", -) -argparser.add_argument( - "--text_prompt", - type=str, - default=None, - help="Text prompt for SAM3 (e.g. 'truck'). Overrides any text entry in the prompt JSON.", -) -argparser.add_argument( - "--image", - type=str, - default="images/truck.jpg", - help="Path to the image", -) -argparser.add_argument( - "--prompt", - type=str, - default="images/truck_prompt.json", - help="Path to the image", -) -argparser.add_argument( - "--output", - type=str, - default=None, - help="Path to the output image", -) -argparser.add_argument( - "--show", - action="store_true", - help="Show the result", -) -argparser.add_argument( - "--sam_variant", - type=str, - default="sam", - help="Variant of SAM model. Options: sam, sam2, sam3", -) -args = argparser.parse_args() - -model = None -if args.sam_variant == "sam": - model = SegmentAnythingONNX( - args.encoder_model, - args.decoder_model, +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run a split SAM ONNX model") + parser.add_argument( + "--encoder_model", + default="output_models/sam_vit_h_4b8939.encoder.onnx", + help="Path to the ONNX encoder model", ) -elif args.sam_variant == "sam2": - model = SegmentAnything2ONNX( - args.encoder_model, - args.decoder_model, + parser.add_argument( + "--decoder_model", + default="output_models/sam_vit_h_4b8939.decoder.onnx", + help="Path to the ONNX decoder model", + ) + parser.add_argument( + "--language_encoder_model", + default=None, + help="Path to the ONNX language encoder model (SAM3 only)", + ) + parser.add_argument( + "--text_prompt", + default=None, + help="SAM3 text prompt; overrides a text mark in the prompt JSON", + ) + parser.add_argument("--image", default="images/truck.jpg") + parser.add_argument("--prompt", default="images/truck_prompt.json") + parser.add_argument("--output", default=None) + parser.add_argument("--show", action="store_true") + parser.add_argument( + "--sam_variant", + choices=("sam", "sam2", "sam3", "efficient_sam"), + default="sam", + help="ONNX model family", + ) + parser.add_argument( + "--confidence_threshold", + type=float, + default=0.5, + help="SAM3 detection confidence threshold", + ) + parser.add_argument( + "--nms_threshold", + type=float, + default=0.7, + help="SAM3 mask- or box-IoU threshold for duplicate suppression", + ) + parser.add_argument( + "--sam3_nms_mode", + choices=("mask", "box", "none"), + default="mask", + help="SAM3 duplicate suppression: official-style mask IoU, box IoU, or none", + ) + parser.add_argument( + "--max_instances", + type=int, + default=None, + help="Maximum SAM3 instances to retain after ranking", + ) + parser.add_argument( + "--sam3_output_mode", + choices=("auto", "all"), + default="auto", + help=( + "SAM3 auto returns the best prompt-overlapping instance when geometry " + "is present; all preserves open-vocabulary instance discovery" + ), ) -elif args.sam_variant == "sam3": - model = SegmentAnything3ONNX( + parser.add_argument( + "--providers", + default=None, + help=( + "Comma-separated ONNX Runtime provider names or aliases, for example " + "tensorrt,cuda,cpu or openvino,cpu" + ), + ) + return parser + + +def load_model(args): + if args.sam_variant == "sam": + return SegmentAnythingONNX( + args.encoder_model, args.decoder_model, providers=args.providers + ) + if args.sam_variant == "sam2": + return SegmentAnything2ONNX( + args.encoder_model, args.decoder_model, providers=args.providers + ) + if args.sam_variant == "efficient_sam": + return EfficientSAMONNX( + args.encoder_model, args.decoder_model, providers=args.providers + ) + return SegmentAnything3ONNX( args.encoder_model, args.decoder_model, args.language_encoder_model, + providers=args.providers, ) -image = cv2.imread(args.image) -prompt = json.load(open(args.prompt)) -text_prompt = None -if args.sam_variant == "sam3": - # --text_prompt takes priority; fall back to any text entry in the JSON. - if args.text_prompt: - text_prompt = args.text_prompt +def get_text_prompt(prompt: list[dict], override: str | None) -> str: + if override: + return override + for mark in prompt: + if mark.get("type") == "text" and isinstance(mark.get("data"), str): + return mark["data"] + return "visual" + + +def visualize(image: np.ndarray, masks: np.ndarray, prompt, variant: str) -> np.ndarray: + visualized = image.copy() + if variant == "sam3": + palette = ( + (255, 99, 71), + (60, 179, 113), + (0, 165, 255), + (238, 130, 238), + (255, 215, 0), + (255, 144, 30), + (147, 112, 219), + (64, 224, 208), + ) + for index, instance_mask in enumerate(masks[:, 0]): + binary_mask = instance_mask.astype(bool) + color = np.asarray(palette[index % len(palette)], dtype=np.uint8) + visualized[binary_mask] = ( + visualized[binary_mask].astype(np.float32) * 0.5 + color * 0.5 + ).astype(np.uint8) + contours, _ = cv2.findContours( + binary_mask.astype(np.uint8), + cv2.RETR_EXTERNAL, + cv2.CHAIN_APPROX_SIMPLE, + ) + cv2.drawContours(visualized, contours, -1, tuple(map(int, color)), 2) + if contours: + largest = max(contours, key=cv2.contourArea) + moments = cv2.moments(largest) + if moments["m00"]: + center = ( + int(moments["m10"] / moments["m00"]), + int(moments["m01"] / moments["m00"]), + ) + cv2.putText( + visualized, + str(index + 1), + center, + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (255, 255, 255), + 2, + cv2.LINE_AA, + ) else: - for p in prompt: - if p["type"] == "text": - text_prompt = p["data"] - break - if text_prompt is None: - text_prompt = "visual" - -embedding = ( - model.encode(image, text_prompt=text_prompt) - if args.sam_variant == "sam3" - else model.encode(image) -) - -masks = model.predict_masks(embedding, prompt) - -# Merge masks -mask = np.zeros((masks.shape[2], masks.shape[3], 3), dtype=np.uint8) -if args.sam_variant == "sam3": - # SAM3 returns bool (N, 1, H, W) – render all N detected instances. - for i in range(masks.shape[0]): - m = masks[i, 0] # (H, W) bool - mask[m] = [255, 0, 0] -else: - # SAM1/SAM2 return raw logits (1, 3, H, W) – threshold at 0 (= sigmoid 0.5). - for m in masks[0, :, :, :]: - mask[m > 0.0] = [255, 0, 0] - -# Binding image and mask -visualized = cv2.addWeighted(image, 0.5, mask, 0.5, 0) - -# Draw the prompt points and rectangles. -for p in prompt: - if p["type"] == "point": - color = ( - (0, 255, 0) if p["label"] == 1 else (0, 0, 255) - ) # green for positive, red for negative - cv2.circle(visualized, (p["data"][0], p["data"][1]), 10, color, -1) - elif p["type"] == "rectangle": - cv2.rectangle( - visualized, - (p["data"][0], p["data"][1]), - (p["data"][2], p["data"][3]), - (0, 255, 0), - 2, + combined_mask = np.zeros(image.shape[:2], dtype=bool) + model_masks = masks[:, 0] if variant == "efficient_sam" else masks[0] + for mask in model_masks: + combined_mask |= mask > 0.0 + color = np.asarray([255, 0, 0], dtype=np.uint8) + visualized[combined_mask] = ( + visualized[combined_mask].astype(np.float32) * 0.5 + color * 0.5 + ).astype(np.uint8) + + for mark in prompt: + if mark.get("type") == "point": + color = (0, 255, 0) if mark.get("label") == 1 else (0, 0, 255) + cv2.circle(visualized, tuple(map(int, mark["data"])), 10, color, -1) + elif mark.get("type") == "rectangle": + x1, y1, x2, y2 = map(int, mark["data"]) + cv2.rectangle(visualized, (x1, y1), (x2, y2), (0, 255, 0), 2) + return visualized + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + image = cv2.imread(args.image) + if image is None: + raise ValueError(f"Could not read image: {args.image}") + with open(args.prompt, encoding="utf-8") as prompt_file: + prompt = json.load(prompt_file) + if not isinstance(prompt, list): + raise ValueError("Prompt JSON must contain a list of marks") + + model = load_model(args) + if args.sam_variant == "sam3": + embedding = model.encode( + image, text_prompt=get_text_prompt(prompt, args.text_prompt) ) - -if args.output is not None: - pathlib.Path(args.output).parent.mkdir(parents=True, exist_ok=True) - cv2.imwrite(args.output, visualized) - -if args.show: - cv2.imshow("Result", visualized) - cv2.waitKey(0) + has_geometry = any( + mark.get("type") in ("point", "rectangle") for mark in prompt + ) + prefer_prompted_region = args.sam3_output_mode == "auto" and has_geometry + max_instances = args.max_instances + if prefer_prompted_region and max_instances is None: + max_instances = 1 + masks = model.predict_masks( + embedding, + prompt, + confidence_threshold=args.confidence_threshold, + nms_threshold=args.nms_threshold, + nms_mode=args.sam3_nms_mode, + max_instances=max_instances, + prefer_prompted_region=prefer_prompted_region, + ) + print(f"SAM3 instances retained: {len(masks)}") + else: + embedding = model.encode(image) + masks = model.predict_masks(embedding, prompt) + + visualized = visualize(image, masks, prompt, args.sam_variant) + if args.output: + output_path = pathlib.Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + if not cv2.imwrite(str(output_path), visualized): + raise OSError(f"Could not write output image: {output_path}") + if args.show: + cv2.imshow("Result", visualized) + cv2.waitKey(0) + cv2.destroyAllWindows() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/samexporter/mobile_encoder/setup_mobile_sam.py b/samexporter/mobile_encoder/setup_mobile_sam.py index cb042b3..5455c8c 100644 --- a/samexporter/mobile_encoder/setup_mobile_sam.py +++ b/samexporter/mobile_encoder/setup_mobile_sam.py @@ -1,11 +1,15 @@ -from segment_anything.modeling import ( +from samexporter.upstream import prefer_pinned_upstream + +prefer_pinned_upstream("sam1") + +from segment_anything.modeling import ( # noqa: E402 MaskDecoder, PromptEncoder, Sam, TwoWayTransformer, ) -from samexporter.mobile_encoder.tiny_vit_sam import TinyViT +from samexporter.mobile_encoder.tiny_vit_sam import TinyViT # noqa: E402 def setup_model(): diff --git a/samexporter/onnx_utils.py b/samexporter/onnx_utils.py index 0939952..be1c52b 100644 --- a/samexporter/onnx_utils.py +++ b/samexporter/onnx_utils.py @@ -1,7 +1,12 @@ import torch import torch.nn as nn -from segment_anything.modeling import Sam -from torch.nn import functional as F + +from samexporter.upstream import prefer_pinned_upstream + +prefer_pinned_upstream("sam1") + +from segment_anything.modeling import Sam # noqa: E402 +from torch.nn import functional as F # noqa: E402 class ImageEncoderOnnxModel(nn.Module): diff --git a/samexporter/prompts.py b/samexporter/prompts.py new file mode 100644 index 0000000..e4c7899 --- /dev/null +++ b/samexporter/prompts.py @@ -0,0 +1,47 @@ +from collections.abc import Sequence + +import numpy as np + + +def geometric_prompt_arrays( + prompt: Sequence[dict], *, allow_empty: bool = False +) -> tuple[np.ndarray, np.ndarray]: + """Validate point/rectangle marks and return SAM coordinate/label arrays.""" + points: list[list[float]] = [] + labels: list[float] = [] + + for index, mark in enumerate(prompt): + if not isinstance(mark, dict): + raise ValueError(f"Prompt mark {index} must be an object") + mark_type = mark.get("type") + if mark_type == "text": + continue + data = mark.get("data") + if mark_type == "point": + if not isinstance(data, (list, tuple)) or len(data) != 2: + raise ValueError(f"Point mark {index} must contain [x, y]") + label = mark.get("label") + if label not in (0, 1): + raise ValueError(f"Point mark {index} label must be 0 or 1") + points.append([float(data[0]), float(data[1])]) + labels.append(float(label)) + elif mark_type == "rectangle": + if not isinstance(data, (list, tuple)) or len(data) != 4: + raise ValueError( + f"Rectangle mark {index} must contain [x1, y1, x2, y2]" + ) + x1, y1, x2, y2 = (float(value) for value in data) + if x2 <= x1 or y2 <= y1: + raise ValueError(f"Rectangle mark {index} must have positive area") + points.extend([[x1, y1], [x2, y2]]) + labels.extend([2.0, 3.0]) + else: + raise ValueError(f"Unsupported prompt type at mark {index}: {mark_type!r}") + + if not points and not allow_empty: + raise ValueError("At least one point or rectangle prompt is required") + + return ( + np.asarray(points, dtype=np.float32).reshape(-1, 2), + np.asarray(labels, dtype=np.float32), + ) diff --git a/samexporter/runtime.py b/samexporter/runtime.py new file mode 100644 index 0000000..aa72694 --- /dev/null +++ b/samexporter/runtime.py @@ -0,0 +1,115 @@ +import logging +import os +from collections.abc import Sequence + +import onnxruntime + +_PROVIDER_PREFERENCE = ( + "NvTensorRTRTXExecutionProvider", + "TensorrtExecutionProvider", + "CUDAExecutionProvider", + "MIGraphXExecutionProvider", + "ROCMExecutionProvider", + "OpenVINOExecutionProvider", + "DmlExecutionProvider", + "CoreMLExecutionProvider", + "CANNExecutionProvider", + "QNNExecutionProvider", + "NnapiExecutionProvider", + "VSINPUExecutionProvider", + "WebNNExecutionProvider", + "WebGpuExecutionProvider", + "XnnpackExecutionProvider", + "RknpuExecutionProvider", + "VitisAIExecutionProvider", + "ACLExecutionProvider", + "ArmNNExecutionProvider", + "DnnlExecutionProvider", + "JsExecutionProvider", + "AzureExecutionProvider", + "CPUExecutionProvider", +) + +_PROVIDER_ALIASES = { + "tensorrt_rtx": "NvTensorRTRTXExecutionProvider", + "trt_rtx": "NvTensorRTRTXExecutionProvider", + "tensorrt": "TensorrtExecutionProvider", + "cuda": "CUDAExecutionProvider", + "migraphx": "MIGraphXExecutionProvider", + "rocm": "ROCMExecutionProvider", + "openvino": "OpenVINOExecutionProvider", + "directml": "DmlExecutionProvider", + "dml": "DmlExecutionProvider", + "coreml": "CoreMLExecutionProvider", + "cann": "CANNExecutionProvider", + "qnn": "QNNExecutionProvider", + "nnapi": "NnapiExecutionProvider", + "vsinpu": "VSINPUExecutionProvider", + "webnn": "WebNNExecutionProvider", + "webgpu": "WebGpuExecutionProvider", + "xnnpack": "XnnpackExecutionProvider", + "rknpu": "RknpuExecutionProvider", + "vitisai": "VitisAIExecutionProvider", + "acl": "ACLExecutionProvider", + "armnn": "ArmNNExecutionProvider", + "dnnl": "DnnlExecutionProvider", + "onednn": "DnnlExecutionProvider", + "js": "JsExecutionProvider", + "azure": "AzureExecutionProvider", + "cpu": "CPUExecutionProvider", +} + + +def _parse_requested_providers( + requested: str | Sequence[str] | None, +) -> list[str] | None: + if requested is None: + requested = os.environ.get("SAMEXPORTER_ONNX_PROVIDERS") + if requested is None: + return None + values = requested.split(",") if isinstance(requested, str) else requested + providers = [] + for value in values: + name = value.strip() + if not name: + continue + providers.append(_PROVIDER_ALIASES.get(name.lower(), name)) + return providers + + +def get_onnx_providers( + requested: str | Sequence[str] | None = None, +) -> list[str]: + """Return available ONNX Runtime providers in preference order. + + ``requested`` accepts exact ONNX Runtime provider names or short aliases + such as ``tensorrt,cuda,cpu``. When omitted, the + ``SAMEXPORTER_ONNX_PROVIDERS`` environment variable is honored, followed by + automatic accelerator discovery. CPU is retained as the final fallback when + it is installed. + """ + available = onnxruntime.get_available_providers() + requested_providers = _parse_requested_providers(requested) + if requested_providers is not None: + unavailable = [name for name in requested_providers if name not in available] + if unavailable: + raise ValueError( + "Requested ONNX Runtime provider(s) are not installed: " + + ", ".join(unavailable) + + ". Available: " + + ", ".join(available) + ) + providers = list(dict.fromkeys(requested_providers)) + else: + preference = [ + name + for name in _PROVIDER_PREFERENCE + if name in available and name != "CPUExecutionProvider" + ] + unknown = [name for name in available if name not in _PROVIDER_PREFERENCE] + providers = preference + unknown + + if "CPUExecutionProvider" in available and "CPUExecutionProvider" not in providers: + providers.append("CPUExecutionProvider") + logging.info("ONNX Runtime providers: %s", ", ".join(providers)) + return providers diff --git a/samexporter/sam2_onnx.py b/samexporter/sam2_onnx.py index 4fd47c9..fb77bb0 100644 --- a/samexporter/sam2_onnx.py +++ b/samexporter/sam2_onnx.py @@ -6,14 +6,18 @@ import onnxruntime from numpy import ndarray +from samexporter.prompts import geometric_prompt_arrays +from samexporter.runtime import get_onnx_providers + class SegmentAnything2ONNX: """Segmentation model using Segment Anything 2 (SAM2)""" - def __init__(self, encoder_model_path, decoder_model_path) -> None: - self.encoder = SAM2ImageEncoder(encoder_model_path) + def __init__(self, encoder_model_path, decoder_model_path, providers=None) -> None: + providers = get_onnx_providers(providers) + self.encoder = SAM2ImageEncoder(encoder_model_path, providers) self.decoder = SAM2ImageDecoder( - decoder_model_path, self.encoder.input_shape[2:] + decoder_model_path, self.encoder.input_shape[2:], providers=providers ) def encode(self, cv_image: np.ndarray) -> list[np.ndarray]: @@ -27,18 +31,7 @@ def encode(self, cv_image: np.ndarray) -> list[np.ndarray]: } def predict_masks(self, embedding, prompt) -> list[np.ndarray]: - points = [] - labels = [] - for mark in prompt: - if mark["type"] == "point": - points.append(mark["data"]) - labels.append(mark["label"]) - elif mark["type"] == "rectangle": - points.append([mark["data"][0], mark["data"][1]]) # top left - points.append([mark["data"][2], mark["data"][3]]) # bottom right - labels.append(2) - labels.append(3) - points, labels = np.array(points), np.array(labels) + points, labels = geometric_prompt_arrays(prompt) image_embedding = embedding["image_embedding"] high_res_feats_0 = embedding["high_res_feats_0"] @@ -74,10 +67,10 @@ def transform_masks(self, masks, original_size, transform_matrix): class SAM2ImageEncoder: - def __init__(self, path: str) -> None: + def __init__(self, path: str, providers=None) -> None: # Initialize model self.session = onnxruntime.InferenceSession( - path, providers=onnxruntime.get_available_providers() + path, providers=get_onnx_providers(providers) ) # Get model info @@ -144,10 +137,11 @@ def __init__( encoder_input_size: tuple[int, int], orig_im_size: tuple[int, int] = None, mask_threshold: float = 0.0, + providers=None, ) -> None: # Initialize model self.session = onnxruntime.InferenceSession( - path, providers=onnxruntime.get_available_providers() + path, providers=get_onnx_providers(providers) ) self.orig_im_size = ( @@ -222,7 +216,7 @@ def prepare_inputs( ), dtype=np.float32, ) - has_mask_input = np.array([0], dtype=np.float32) + has_mask_input = np.zeros(num_labels, dtype=np.float32) return ( image_embed, @@ -241,8 +235,8 @@ def prepare_points( ) -> tuple[np.ndarray, np.ndarray]: if isinstance(point_coords, np.ndarray): - input_point_coords = point_coords[np.newaxis, ...] - input_point_labels = point_labels[np.newaxis, ...] + input_point_coords = point_coords[np.newaxis, ...].copy() + input_point_labels = np.asarray(point_labels)[np.newaxis, ...].copy() else: max_num_points = max([coords.shape[0] for coords in point_coords]) # We need to make sure that all inputs have the same number of points diff --git a/samexporter/sam3_onnx.py b/samexporter/sam3_onnx.py index b21d8f4..a1c81fb 100644 --- a/samexporter/sam3_onnx.py +++ b/samexporter/sam3_onnx.py @@ -4,6 +4,9 @@ import numpy as np import onnxruntime +from samexporter.clip_tokenizer import tokenize +from samexporter.runtime import get_onnx_providers + class SegmentAnything3ONNX: """Segmentation model using Segment Anything 3 (SAM3)""" @@ -13,12 +16,16 @@ def __init__( image_encoder_path, decoder_model_path, language_encoder_path=None, + providers=None, ) -> None: - self.image_encoder = SAM3ImageEncoder(image_encoder_path) + providers = get_onnx_providers(providers) + self.image_encoder = SAM3ImageEncoder(image_encoder_path, providers) self.language_encoder = None if language_encoder_path: - self.language_encoder = SAM3LanguageEncoder(language_encoder_path) - self.decoder = SAM3ImageDecoder(decoder_model_path) + self.language_encoder = SAM3LanguageEncoder( + language_encoder_path, providers + ) + self.decoder = SAM3ImageDecoder(decoder_model_path, providers) def encode(self, cv_image: np.ndarray, text_prompt=None) -> dict[str, Any]: """Encode an image (and optional text prompt) into an embedding dict. @@ -37,6 +44,8 @@ def encode(self, cv_image: np.ndarray, text_prompt=None) -> dict[str, Any]: original_size, vision_pos_enc_{0,1,2}, backbone_fpn_{0,1,2}, language_mask, language_features, language_embeds. """ + if cv_image is None or cv_image.ndim != 3 or cv_image.shape[2] != 3: + raise ValueError("Expected a non-empty BGR image with three channels") original_size = cv_image.shape[:2] image_encoder_outputs = self.image_encoder(cv_image) @@ -73,6 +82,10 @@ def predict_masks( embedding: dict[str, Any], prompt, confidence_threshold: float = 0.5, + nms_threshold: float | None = 0.7, + nms_mode: str = "mask", + max_instances: int | None = None, + prefer_prompted_region: bool = False, ) -> np.ndarray: """Run the decoder for the given geometric prompt. @@ -86,6 +99,17 @@ def predict_masks( confidence_threshold: Minimum score to keep a detection. Detections with score below this value are discarded. Defaults to ``0.5``. + nms_threshold: + IoU threshold used to suppress duplicate queries. Set to ``None`` + to retain threshold-only behavior. + nms_mode: + ``"mask"`` follows the official SAM3 mask-IoU NMS behavior; + ``"box"`` is faster for very large masks; ``"none"`` disables NMS. + max_instances: + Optional maximum number of masks, after ranking and NMS. + prefer_prompted_region: + Rank masks overlapping positive points/rectangles ahead of remote + concept matches. Useful when geometry is intended as a selection. Returns ------- @@ -93,36 +117,67 @@ def predict_masks( of detected objects and *H* × *W* is the original image resolution. """ original_size = embedding["original_size"] - box_coords = [0.0, 0.0, 0.0, 0.0] - box_labels = [1] - # box_masks: True → dummy / no real box - # False → a real box is provided - box_masks = [True] - - for mark in prompt: - if mark["type"] == "rectangle": - x1, y1, x2, y2 = mark["data"] + box_coords = [] + box_labels = [] + + for index, mark in enumerate(prompt): + if not isinstance(mark, dict): + raise ValueError(f"Prompt mark {index} must be an object") + mark_type = mark.get("type") + if mark_type == "text": + continue + data = mark.get("data") + if mark_type == "rectangle": + if not isinstance(data, (list, tuple)) or len(data) != 4: + raise ValueError( + f"Rectangle mark {index} must contain [x1, y1, x2, y2]" + ) + x1, y1, x2, y2 = map(float, data) + if x2 <= x1 or y2 <= y1: + raise ValueError(f"Rectangle mark {index} must have positive area") cx = (x1 + x2) / 2.0 / original_size[1] cy = (y1 + y2) / 2.0 / original_size[0] w = (x2 - x1) / original_size[1] h = (y2 - y1) / original_size[0] - box_coords = [cx, cy, w, h] - box_masks = [False] - break - elif mark["type"] == "point": - x, y = mark["data"] + box_coords.append([cx, cy, w, h]) + box_labels.append(1) + elif mark_type == "point": + if not isinstance(data, (list, tuple)) or len(data) != 2: + raise ValueError(f"Point mark {index} must contain [x, y]") + label = mark.get("label") + if label not in (0, 1): + raise ValueError(f"Point mark {index} label must be 0 or 1") + x, y = map(float, data) cx = x / original_size[1] cy = y / original_size[0] # Point is represented as a very small box (1 % of image). - box_coords = [cx, cy, 0.01, 0.01] - box_masks = [False] - break - - box_coords_np = np.array(box_coords, dtype=np.float32).reshape(1, 1, 4) - box_labels_np = np.array([box_labels], dtype=np.int64) - box_masks_np = np.array([box_masks], dtype=np.bool_) - - masks, scores, _boxes = self.decoder( + box_coords.append([cx, cy, 0.01, 0.01]) + box_labels.append(label) + else: + raise ValueError( + f"Unsupported prompt type at mark {index}: {mark_type!r}" + ) + + mark_count = len(box_coords) + capacity = self.decoder.geometric_prompt_capacity + if capacity is not None and mark_count > capacity: + raise ValueError( + f"SAM3 decoder accepts at most {capacity} geometric prompts; " + "re-export it with a larger --max-geometric-prompts value" + ) + tensor_length = capacity or max(mark_count, 1) + # SAM3's Prompt contract is sequence-first: [num_marks, batch, C]. + # ONNX traces its internal geometry attention at a fixed token count, + # so exported decoders use padded slots and mark them True in box_masks. + box_coords_np = np.zeros((tensor_length, 1, 4), dtype=np.float32) + box_labels_np = np.ones((tensor_length, 1), dtype=np.int64) + box_masks_np = np.ones((1, tensor_length), dtype=np.bool_) + if mark_count: + box_coords_np[:mark_count, 0] = np.asarray(box_coords, dtype=np.float32) + box_labels_np[:mark_count, 0] = np.asarray(box_labels, dtype=np.int64) + box_masks_np[0, :mark_count] = False + + masks, scores, boxes = self.decoder( original_size, embedding["vision_pos_enc_0"], embedding["vision_pos_enc_1"], @@ -138,16 +193,212 @@ def predict_masks( box_masks_np, ) - # Filter detections by confidence score. - if len(scores) > 0: - keep = np.where(scores > confidence_threshold)[0] - masks = ( - masks[keep] - if len(keep) > 0 - else np.zeros((0,) + masks.shape[1:], dtype=masks.dtype) + masks, _scores, _boxes = self.select_detections( + masks, + scores, + boxes, + prompt, + confidence_threshold=confidence_threshold, + nms_threshold=nms_threshold, + nms_mode=nms_mode, + max_instances=max_instances, + prefer_prompted_region=prefer_prompted_region, + ) + return masks + + @staticmethod + def select_detections( + masks: np.ndarray, + scores: np.ndarray, + boxes: np.ndarray, + prompt, + *, + confidence_threshold: float, + nms_threshold: float | None, + nms_mode: str, + max_instances: int | None, + prefer_prompted_region: bool, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Filter and rank raw SAM3 detections without model dependencies.""" + if not 0.0 <= confidence_threshold <= 1.0: + raise ValueError("confidence_threshold must be between 0 and 1") + if nms_threshold is not None and not 0.0 <= nms_threshold <= 1.0: + raise ValueError("nms_threshold must be between 0 and 1 or None") + if nms_mode not in ("mask", "box", "none"): + raise ValueError("nms_mode must be 'mask', 'box', or 'none'") + if max_instances is not None and max_instances < 1: + raise ValueError("max_instances must be at least 1 or None") + + masks = np.asarray(masks) + scores = np.asarray(scores, dtype=np.float32).reshape(-1) + boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4) + if not (len(masks) == len(scores) == len(boxes)): + raise ValueError("SAM3 masks, scores, and boxes must have equal lengths") + + keep = np.flatnonzero(np.isfinite(scores) & (scores > confidence_threshold)) + if not len(keep): + return ( + np.zeros((0,) + masks.shape[1:], dtype=masks.dtype), + np.zeros((0,), dtype=np.float32), + np.zeros((0, 4), dtype=np.float32), + ) + masks, scores, boxes = masks[keep], scores[keep], boxes[keep] + + order = np.argsort(-scores, kind="stable") + masks, scores, boxes = masks[order], scores[order], boxes[order] + if nms_mode != "none" and nms_threshold is not None and len(boxes) > 1: + if nms_mode == "mask": + keep = SegmentAnything3ONNX._mask_nms(masks, scores, nms_threshold) + else: + keep = SegmentAnything3ONNX._box_nms(boxes, scores, nms_threshold) + masks, scores, boxes = masks[keep], scores[keep], boxes[keep] + + if prefer_prompted_region and len(masks) > 1: + relevance = SegmentAnything3ONNX._prompt_relevance(masks, boxes, prompt) + order = np.lexsort((-scores, -relevance)) + masks, scores, boxes = masks[order], scores[order], boxes[order] + + if max_instances is not None: + masks = masks[:max_instances] + scores = scores[:max_instances] + boxes = boxes[:max_instances] + return masks, scores, boxes + + @staticmethod + def _box_nms(boxes: np.ndarray, scores: np.ndarray, threshold: float) -> np.ndarray: + """Return score-ordered indices after standard XYXY box NMS.""" + x1, y1, x2, y2 = boxes.T + areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1) + order = np.argsort(-scores, kind="stable") + kept: list[int] = [] + while len(order): + current = int(order[0]) + kept.append(current) + if len(order) == 1: + break + rest = order[1:] + intersection_width = np.maximum( + 0.0, + np.minimum(x2[current], x2[rest]) - np.maximum(x1[current], x1[rest]), + ) + intersection_height = np.maximum( + 0.0, + np.minimum(y2[current], y2[rest]) - np.maximum(y1[current], y1[rest]), + ) + intersection = intersection_width * intersection_height + union = areas[current] + areas[rest] - intersection + iou = np.divide( + intersection, + union, + out=np.zeros_like(intersection), + where=union > 0, ) + order = rest[iou <= threshold] + return np.asarray(kept, dtype=np.int64) - return masks + @staticmethod + def _mask_nms( + masks: np.ndarray, scores: np.ndarray, threshold: float + ) -> np.ndarray: + """Return score-ordered indices after bit-packed mask-IoU NMS. + + SAM3's official NMS uses mask overlap. Masks are sampled to at most + 288 pixels on their longest side (the native SAM3 mask-head scale), + then packed into bits so large original-resolution outputs do not make + duplicate suppression memory-bound. + """ + binary = np.asarray(masks, dtype=bool) + if binary.ndim == 4 and binary.shape[1] == 1: + binary = binary[:, 0] + if binary.ndim != 3: + raise ValueError("SAM3 masks must have shape (N, 1, H, W) or (N, H, W)") + step = max(1, int(np.ceil(max(binary.shape[-2:]) / 288))) + sampled = binary[:, ::step, ::step].reshape(len(binary), -1) + packed = np.packbits(sampled, axis=1) + bit_counts = np.unpackbits(np.arange(256, dtype=np.uint8)[:, None], axis=1).sum( + axis=1 + ) + areas = bit_counts[packed].sum(axis=1, dtype=np.int64) + + order = np.argsort(-scores, kind="stable") + kept: list[int] = [] + while len(order): + current = int(order[0]) + kept.append(current) + if len(order) == 1: + break + rest = order[1:] + intersections = bit_counts[ + np.bitwise_and(packed[rest], packed[current]) + ].sum(axis=1, dtype=np.int64) + unions = areas[current] + areas[rest] - intersections + iou = np.divide( + intersections, + unions, + out=np.zeros(len(rest), dtype=np.float64), + where=unions > 0, + ) + order = rest[iou <= threshold] + return np.asarray(kept, dtype=np.int64) + + @staticmethod + def _prompt_relevance(masks: np.ndarray, boxes: np.ndarray, prompt) -> np.ndarray: + """Score positive-point/rectangle overlap for selection-style prompting.""" + relevance = np.zeros(len(masks), dtype=np.float32) + mask_height, mask_width = masks.shape[-2:] + for mark in prompt: + if not isinstance(mark, dict) or mark.get("label", 1) != 1: + continue + data = mark.get("data") + if ( + mark.get("type") == "point" + and isinstance(data, (list, tuple)) + and len(data) == 2 + ): + x = int(np.clip(round(float(data[0])), 0, mask_width - 1)) + y = int(np.clip(round(float(data[1])), 0, mask_height - 1)) + relevance += masks[:, 0, y, x].astype(np.float32) * 2.0 + relevance += ( + (boxes[:, 0] <= x) + & (x <= boxes[:, 2]) + & (boxes[:, 1] <= y) + & (y <= boxes[:, 3]) + ).astype(np.float32) + elif ( + mark.get("type") == "rectangle" + and isinstance(data, (list, tuple)) + and len(data) == 4 + ): + prompt_box = np.asarray(data, dtype=np.float32) + ix1 = np.maximum(boxes[:, 0], prompt_box[0]) + iy1 = np.maximum(boxes[:, 1], prompt_box[1]) + ix2 = np.minimum(boxes[:, 2], prompt_box[2]) + iy2 = np.minimum(boxes[:, 3], prompt_box[3]) + intersection = np.maximum(0.0, ix2 - ix1) * np.maximum(0.0, iy2 - iy1) + box_area = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum( + 0.0, boxes[:, 3] - boxes[:, 1] + ) + prompt_area = max( + 0.0, + float(prompt_box[2] - prompt_box[0]) + * float(prompt_box[3] - prompt_box[1]), + ) + union = box_area + prompt_area - intersection + relevance += np.divide( + intersection, + union, + out=np.zeros_like(intersection), + where=union > 0, + ) + center_x = (boxes[:, 0] + boxes[:, 2]) / 2.0 + center_y = (boxes[:, 1] + boxes[:, 3]) / 2.0 + relevance += ( + (prompt_box[0] <= center_x) + & (center_x <= prompt_box[2]) + & (prompt_box[1] <= center_y) + & (center_y <= prompt_box[3]) + ).astype(np.float32) + return relevance def transform_masks(self, masks, original_size, transform_matrix): """No-op: SAM3 already outputs masks in original image resolution.""" @@ -164,9 +415,9 @@ class SAM3ImageEncoder: dtype : uint8 (the model includes normalization internally) """ - def __init__(self, path: str) -> None: + def __init__(self, path: str, providers=None) -> None: self.session = onnxruntime.InferenceSession( - path, providers=onnxruntime.get_available_providers() + path, providers=get_onnx_providers(providers) ) encoder_input = self.session.get_inputs()[0] self.input_name: str = encoder_input.name @@ -222,29 +473,13 @@ class SAM3LanguageEncoder: dtype : int64 """ - def __init__(self, path: str) -> None: + def __init__(self, path: str, providers=None) -> None: self.session = onnxruntime.InferenceSession( - path, providers=onnxruntime.get_available_providers() + path, providers=get_onnx_providers(providers) ) - try: - from osam._models.yoloworld.clip import tokenize - - self._tokenize = tokenize - except ImportError: - self._tokenize = self._fallback_tokenize - - def _fallback_tokenize(self, texts, context_length: int = 32) -> np.ndarray: - """Minimal CLIP-style tokeniser fallback (zeros = empty sequence). - - Warning: the model will produce near-random language features when - this fallback is used. Install ``osam`` for correct tokenisation. - """ - return np.zeros((len(texts), context_length), dtype=np.int64) def __call__(self, text: str) -> list[np.ndarray]: - tokens = self._tokenize([text], context_length=32) - if not isinstance(tokens, np.ndarray): - tokens = np.asarray(tokens, dtype=np.int64) + tokens = tokenize([text], context_length=32) return self.session.run(None, {"tokens": tokens}) @@ -260,11 +495,24 @@ class SAM3ImageDecoder: callers can unpack in a semantically natural order. """ - def __init__(self, path: str) -> None: + def __init__(self, path: str, providers=None) -> None: self.session = onnxruntime.InferenceSession( - path, providers=onnxruntime.get_available_providers() + path, providers=get_onnx_providers(providers) ) self.input_names: list[str] = [i.name for i in self.session.get_inputs()] + self.output_names: list[str] = [ + output.name for output in self.session.get_outputs() + ] + box_input = next( + (item for item in self.session.get_inputs() if item.name == "box_coords"), + None, + ) + first_dimension = box_input.shape[0] if box_input is not None else None + self.geometric_prompt_capacity = ( + first_dimension + if isinstance(first_dimension, int) and first_dimension > 0 + else None + ) def __call__( self, @@ -319,6 +567,46 @@ def __call__( k: v for k, v in inputs.items() if k in self.input_names and v is not None } outputs = self.session.run(None, model_inputs) - # ONNX export order: [0]=boxes, [1]=scores, [2]=masks - # Return as (masks, scores, boxes) for caller convenience. + if "pred_logits" in self.output_names: + raw = dict(zip(self.output_names, outputs)) + return self.postprocess_raw_outputs(raw, original_size) + + # Backward compatibility with older exports whose processor-based + # graph returned already-filtered [boxes, scores, masks]. return outputs[2], outputs[1], outputs[0] + + @staticmethod + def postprocess_raw_outputs(raw, original_size): + """Match the official Sam3Processor postprocessing outside ONNX.""" + logits = np.asarray(raw["pred_logits"])[0].squeeze(-1) + presence = np.asarray(raw["presence_logit_dec"]).reshape(-1)[0] + + def sigmoid(value): + return 1.0 / (1.0 + np.exp(-np.clip(value, -80, 80))) + + scores = sigmoid(logits) * sigmoid(presence) + + raw_masks = np.asarray(raw["pred_masks"])[0] + height, width = original_size + masks = np.stack( + [ + sigmoid( + cv2.resize(mask, (width, height), interpolation=cv2.INTER_LINEAR) + ) + > 0.5 + for mask in raw_masks + ] + )[:, None] + + boxes_cxcywh = np.asarray(raw["pred_boxes"])[0] + cx, cy, box_width, box_height = boxes_cxcywh.T + boxes = np.stack( + [ + (cx - box_width / 2) * width, + (cy - box_height / 2) * height, + (cx + box_width / 2) * width, + (cy + box_height / 2) * height, + ], + axis=-1, + ) + return masks, scores, boxes diff --git a/samexporter/sam_onnx.py b/samexporter/sam_onnx.py index ee1cefe..3155900 100644 --- a/samexporter/sam_onnx.py +++ b/samexporter/sam_onnx.py @@ -1,56 +1,39 @@ -import logging from copy import deepcopy import cv2 import numpy as np import onnxruntime +from PIL import Image -logging.basicConfig(level=logging.DEBUG) +from samexporter.prompts import geometric_prompt_arrays +from samexporter.runtime import get_onnx_providers class SegmentAnythingONNX: """Segmentation model using Segment Anything (SAM)""" - def __init__(self, encoder_model_path, decoder_model_path) -> None: + def __init__(self, encoder_model_path, decoder_model_path, providers=None) -> None: self.target_size = 1024 - self.input_size = (684, 1024) # Load models - providers = onnxruntime.get_available_providers() - - # Pop TensorRT Runtime due to crashing issues - # TODO: Add back when TensorRT backend is stable - providers = [p for p in providers if p != "TensorrtExecutionProvider"] - - if providers: - logging.info( - "Available providers for ONNXRuntime: %s", ", ".join(providers) - ) - else: - logging.warning("No available providers for ONNXRuntime") + providers = get_onnx_providers(providers) self.encoder_session = onnxruntime.InferenceSession( encoder_model_path, providers=providers ) - self.encoder_input_name = self.encoder_session.get_inputs()[0].name + encoder_input = self.encoder_session.get_inputs()[0] + self.encoder_input_name = encoder_input.name + self.encoder_input_shape = encoder_input.shape + self.encoder_input_rank = len(encoder_input.shape) self.decoder_session = onnxruntime.InferenceSession( decoder_model_path, providers=providers ) + self.decoder_output_names = [ + model_output.name for model_output in self.decoder_session.get_outputs() + ] def get_input_points(self, prompt): """Get input points""" - points = [] - labels = [] - for mark in prompt: - if mark["type"] == "point": - points.append(mark["data"]) - labels.append(mark["label"]) - elif mark["type"] == "rectangle": - points.append([mark["data"][0], mark["data"][1]]) # top left - points.append([mark["data"][2], mark["data"][3]]) # bottom right - labels.append(2) - labels.append(3) - points, labels = np.array(points), np.array(labels) - return points, labels + return geometric_prompt_arrays(prompt) def run_encoder(self, encoder_inputs): """Run encoder""" @@ -83,7 +66,7 @@ def apply_coords(self, coords: np.ndarray, original_size, target_length): coords[..., 1] = coords[..., 1] * (new_h / old_h) return coords - def run_decoder(self, image_embedding, original_size, transform_matrix, prompt): + def run_decoder(self, image_embedding, original_size, resized_size, prompt): """Run decoder""" input_points, input_labels = self.get_input_points(prompt) @@ -95,20 +78,9 @@ def run_decoder(self, image_embedding, original_size, transform_matrix, prompt): None, : ].astype(np.float32) onnx_coord = self.apply_coords( - onnx_coord, self.input_size, self.target_size + onnx_coord, original_size, self.target_size ).astype(np.float32) - # Apply the transformation matrix to the coordinates. - onnx_coord = np.concatenate( - [ - onnx_coord, - np.ones((1, onnx_coord.shape[1], 1), dtype=np.float32), - ], - axis=2, - ) - onnx_coord = np.matmul(onnx_coord, transform_matrix.T) - onnx_coord = onnx_coord[:, :, :2].astype(np.float32) - # Create an empty mask input and an indicator for no mask. onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32) onnx_has_mask_input = np.zeros(1, dtype=np.float32) @@ -119,17 +91,61 @@ def run_decoder(self, image_embedding, original_size, transform_matrix, prompt): "point_labels": onnx_label, "mask_input": onnx_mask_input, "has_mask_input": onnx_has_mask_input, - "orig_im_size": np.array(self.input_size, dtype=np.float32), + "orig_im_size": np.array(original_size, dtype=np.float32), } - masks, _, _ = self.decoder_session.run(None, decoder_inputs) + outputs = self.decoder_session.run(None, decoder_inputs) + + # The segment-anything 1.0 ONNX wrapper converts a tensor-derived crop + # size to Python ``int`` during tracing. That freezes the export to the + # dummy image's landscape ratio. Prefer low-resolution logits and do + # the standard resize/crop/resize here so every aspect ratio remains + # correct (including portrait and extreme panoramas). + if "low_res_masks" in self.decoder_output_names: + masks = outputs[self.decoder_output_names.index("low_res_masks")] + return self.postprocess_masks(masks, original_size, resized_size) + + masks = outputs[0] + if masks.shape[-2:] != tuple(original_size): + return self.resize_masks(masks, original_size) + return masks - # Transform the masks back to the original image size. - inv_transform_matrix = np.linalg.inv(transform_matrix) - transformed_masks = self.transform_masks( - masks, original_size, inv_transform_matrix - ) + def postprocess_masks(self, masks, original_size, resized_size): + """Apply SAM's aspect-ratio-aware mask postprocessing outside ONNX.""" + output_masks = [] + for batch in masks: + batch_masks = [] + for mask in batch: + upscaled = cv2.resize( + mask, + (self.target_size, self.target_size), + interpolation=cv2.INTER_LINEAR, + ) + cropped = upscaled[: resized_size[0], : resized_size[1]] + batch_masks.append( + cv2.resize( + cropped, + (original_size[1], original_size[0]), + interpolation=cv2.INTER_LINEAR, + ) + ) + output_masks.append(batch_masks) + return np.asarray(output_masks) - return transformed_masks + @staticmethod + def resize_masks(masks, original_size): + return np.asarray( + [ + [ + cv2.resize( + mask, + (original_size[1], original_size[0]), + interpolation=cv2.INTER_LINEAR, + ) + for mask in batch + ] + for batch in masks + ] + ) def transform_masks(self, masks, original_size, transform_matrix): """Transform the masks back to the original image size.""" @@ -152,34 +168,48 @@ def encode(self, cv_image): """ Calculate embedding and metadata for a single image. """ + if cv_image is None or cv_image.ndim != 3 or cv_image.shape[2] != 3: + raise ValueError("Expected a non-empty BGR image with three channels") original_size = cv_image.shape[:2] - - # Calculate a transformation matrix to convert to self.input_size - scale_x = self.input_size[1] / cv_image.shape[1] - scale_y = self.input_size[0] / cv_image.shape[0] - scale = min(scale_x, scale_y) - transform_matrix = np.array( - [ - [scale, 0, 0], - [0, scale, 0], - [0, 0, 1], - ] - ) - cv_image = cv2.warpAffine( - cv_image, - transform_matrix[:2], - (self.input_size[1], self.input_size[0]), - flags=cv2.INTER_LINEAR, + resized_size = self.get_preprocess_shape(*original_size, self.target_size) + rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) + # Match official ResizeLongestSide.apply_image exactly. OpenCV's + # sampler can move thin/ambiguous point-prompt boundaries enough to + # select a different mask, especially after large downscales. + rgb_image = np.asarray( + Image.fromarray(rgb_image).resize( + (resized_size[1], resized_size[0]), + resample=Image.Resampling.BILINEAR, + ) ) + if self.encoder_input_rank == 3: + # Encoder exported with --use-preprocess: dynamic HWC RGB input. + input_image = rgb_image.astype(np.float32) + elif self.encoder_input_rank == 4: + # Encoder exported without preprocessing: normalized, padded NCHW. + input_image = rgb_image.astype(np.float32) + mean = np.asarray([123.675, 116.28, 103.53], dtype=np.float32) + std = np.asarray([58.395, 57.12, 57.375], dtype=np.float32) + input_image = (input_image - mean) / std + input_image = input_image.transpose(2, 0, 1) + pad_h = self.target_size - resized_size[0] + pad_w = self.target_size - resized_size[1] + input_image = np.pad(input_image, ((0, 0), (0, pad_h), (0, pad_w))) + input_image = input_image[None].astype(np.float32) + else: + raise ValueError( + f"Unsupported SAM encoder input rank: {self.encoder_input_rank}" + ) + encoder_inputs = { - self.encoder_input_name: cv_image.astype(np.float32), + self.encoder_input_name: input_image, } image_embedding = self.run_encoder(encoder_inputs) return { "image_embedding": image_embedding, "original_size": original_size, - "transform_matrix": transform_matrix, + "resized_size": resized_size, } def predict_masks(self, embedding, prompt): @@ -189,7 +219,7 @@ def predict_masks(self, embedding, prompt): masks = self.run_decoder( embedding["image_embedding"], embedding["original_size"], - embedding["transform_matrix"], + embedding["resized_size"], prompt, ) diff --git a/samexporter/upstream.py b/samexporter/upstream.py new file mode 100644 index 0000000..0d9f241 --- /dev/null +++ b/samexporter/upstream.py @@ -0,0 +1,26 @@ +import sys +from pathlib import Path + +_UPSTREAM_PATHS = { + "sam1": "third_party/segment-anything", + "sam2": "third_party/sam2", + "sam3": "sam3", +} + + +def prefer_pinned_upstream(model_family: str) -> Path | None: + """Prefer a checked-out pinned upstream while keeping wheel fallback. + + Source-tree users get the exact submodule revision recorded by this + repository. Installed wheels do not contain the submodules and continue to + resolve their separately installed upstream package. + """ + relative_path = _UPSTREAM_PATHS.get(model_family) + if relative_path is None: + raise ValueError(f"Unknown upstream model family: {model_family}") + repository_root = Path(__file__).resolve().parent.parent + upstream_path = repository_root / relative_path + if upstream_path.is_dir() and str(upstream_path) not in sys.path: + sys.path.insert(0, str(upstream_path)) + return upstream_path + return upstream_path if upstream_path.is_dir() else None diff --git a/test_all_models.sh b/test_all_models.sh index 6f85866..0e1c785 100755 --- a/test_all_models.sh +++ b/test_all_models.sh @@ -1,17 +1,19 @@ #!/bin/bash # test_all_models.sh -set -e +set -euo pipefail -mkdir -p test_outputs +OUT_DIR="${SAMEXPORTER_RESULTS_DIR:-visual_results/runs/all_models}" +mkdir -p "$OUT_DIR" run_test() { local variant=$1 local encoder=$2 local decoder=$3 local prompt_file=$4 - local output="test_outputs/${variant}_$(basename $encoder .onnx)_result.png" - local extra_args=$5 + local output="${OUT_DIR}/${variant}_$(basename "$encoder" .onnx)_result.png" + local log="${output%.png}.log" + local extra_args=${5:-} echo "Testing $variant with $encoder ..." python -m samexporter.inference \ @@ -21,7 +23,7 @@ run_test() { --image images/truck.jpg \ --prompt "$prompt_file" \ --output "$output" \ - $extra_args + $extra_args 2>&1 | tee "$log" if [ -f "$output" ]; then echo " [OK] Saved to $output" @@ -37,7 +39,7 @@ run_test "sam" "output_models/sam_vit_l_0b3195.encoder.onnx" "output_models/sam_ run_test "sam" "output_models/sam_vit_b_01ec64.encoder.onnx" "output_models/sam_vit_b_01ec64.decoder.onnx" "images/truck_prompt.json" echo -e "\n=== Testing Mobile SAM ===" -run_test "sam" "output_models/mobile_sam/mobile_sam.encoder.onnx" "output_models/mobile_sam/sam_vit_h_4b8939.decoder.onnx" "images/truck_prompt.json" +run_test "sam" "output_models/mobile_sam/mobile_sam.encoder.onnx" "output_models/mobile_sam/mobile_sam.decoder.onnx" "images/truck_prompt.json" echo -e "\n=== Testing SAM 2 ===" run_test "sam2" "output_models/sam2_hiera_tiny.encoder.onnx" "output_models/sam2_hiera_tiny.decoder.onnx" "images/truck_prompt.json" diff --git a/test_comprehensive.sh b/test_comprehensive.sh index b8f9a15..500dca0 100755 --- a/test_comprehensive.sh +++ b/test_comprehensive.sh @@ -1,9 +1,9 @@ #!/bin/bash # test_comprehensive.sh -set -e +set -euo pipefail -OUT_DIR="test_outputs/comprehensive" +OUT_DIR="${SAMEXPORTER_RESULTS_DIR:-visual_results/runs/comprehensive}" mkdir -p "$OUT_DIR" # Parameters: variant, encoder, decoder, image, prompt, suffix, extra_args @@ -19,6 +19,7 @@ run_single_test() { local img_base=$(basename "$image" | cut -d. -f1) local enc_base=$(basename "$encoder" | cut -d. -f1) local output="${OUT_DIR}/${variant}_${enc_base}_${img_base}_${suffix}.png" + local log="${output%.png}.log" echo "Testing: $variant | Model: $enc_base | Image: $img_base | Mode: $suffix" @@ -29,7 +30,7 @@ run_single_test() { --image "$image" \ --prompt "$prompt" \ --output "$output" \ - $extra_args > /dev/null 2>&1 + $extra_args 2>&1 | tee "$log" if [ -f "$output" ]; then echo " [OK] -> $output" @@ -54,31 +55,54 @@ for img in "${IMAGES[@]}"; do run_single_test "sam" "output_models/sam_vit_h_4b8939.encoder.onnx" "output_models/sam_vit_h_4b8939.decoder.onnx" "$img" "images/${img_name}_point.json" "point" # Box run_single_test "sam" "output_models/sam_vit_h_4b8939.encoder.onnx" "output_models/sam_vit_h_4b8939.decoder.onnx" "$img" "images/${img_name}_box.json" "box" + if [ "$img_name" = "plants" ]; then + run_single_test "sam" "output_models/sam_vit_h_4b8939.encoder.onnx" "output_models/sam_vit_h_4b8939.decoder.onnx" "$img" "images/plants_box_refined.json" "refined" + fi # Mobile SAM - run_single_test "sam" "output_models/mobile_sam/mobile_sam.encoder.onnx" "output_models/mobile_sam/sam_vit_h_4b8939.decoder.onnx" "$img" "images/${img_name}_box.json" "mobile_box" + run_single_test "sam" "output_models/mobile_sam/mobile_sam.encoder.onnx" "output_models/mobile_sam/mobile_sam.decoder.onnx" "$img" "images/${img_name}_box.json" "mobile_box" +done + +# 2. Test EfficientSAM-Ti +for img in "${IMAGES[@]}"; do + img_name=$(basename "$img" | cut -d. -f1) + text_prompt="$img_name" + if [ "$img_name" = "plants" ]; then + text_prompt="plant" + fi + run_single_test "efficient_sam" "output_models/efficient_sam/efficientsam_ti_encoder.onnx" "output_models/efficient_sam/efficientsam_ti_decoder.onnx" "$img" "images/${img_name}_point.json" "point" + run_single_test "efficient_sam" "output_models/efficient_sam/efficientsam_ti_encoder.onnx" "output_models/efficient_sam/efficientsam_ti_decoder.onnx" "$img" "images/${img_name}_box.json" "box" + if [ "$img_name" = "plants" ]; then + run_single_test "efficient_sam" "output_models/efficient_sam/efficientsam_ti_encoder.onnx" "output_models/efficient_sam/efficientsam_ti_decoder.onnx" "$img" "images/plants_box_refined.json" "refined" + fi done -# 2. Test SAM 2 & 2.1 +# 3. Test SAM 2 & 2.1 for model in "${MODELS_SAM2[@]}" "${MODELS_SAM21[@]}"; do for img in "${IMAGES[@]}"; do img_name=$(basename "$img" | cut -d. -f1) run_single_test "sam2" "output_models/${model}.encoder.onnx" "output_models/${model}.decoder.onnx" "$img" "images/${img_name}_point.json" "point" run_single_test "sam2" "output_models/${model}.encoder.onnx" "output_models/${model}.decoder.onnx" "$img" "images/${img_name}_box.json" "box" + if [ "$img_name" = "plants" ]; then + run_single_test "sam2" "output_models/${model}.encoder.onnx" "output_models/${model}.decoder.onnx" "$img" "images/plants_box_refined.json" "refined" + fi done done -# 3. Test SAM 3 (Point, Box, Text) +# 4. Test SAM 3 (Point, Box, Text) for img in "${IMAGES[@]}"; do img_name=$(basename "$img" | cut -d. -f1) # Point - run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "$img" "images/${img_name}_point.json" "point" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx" + run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "$img" "images/${img_name}_point.json" "point" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx --text_prompt $text_prompt" # Box - run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "$img" "images/${img_name}_box.json" "box" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx" + run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "$img" "images/${img_name}_box.json" "box" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx --text_prompt $text_prompt" + if [ "$img_name" = "plants" ]; then + run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "$img" "images/plants_box_refined.json" "refined" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx --text_prompt plant" + fi done # SAM 3 specific Text tests run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "images/truck.jpg" "images/truck_sam3.json" "text_truck" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx" -run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "images/plants.png" "images/plants_text.json" "text_leaf" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx" +run_single_test "sam3" "output_models/sam3/sam3_image_encoder.onnx" "output_models/sam3/sam3_decoder.onnx" "images/plants.png" "images/plants_text.json" "text_plant" "--language_encoder_model output_models/sam3/sam3_language_encoder.onnx" echo -e " All comprehensive tests passed!" diff --git a/tests/test_clip_tokenizer.py b/tests/test_clip_tokenizer.py new file mode 100644 index 0000000..711fc6f --- /dev/null +++ b/tests/test_clip_tokenizer.py @@ -0,0 +1,26 @@ +import numpy as np +import pytest + +from samexporter.clip_tokenizer import tokenize + + +def test_clip_tokenizer_matches_sam3_tokens(): + tokens = tokenize(["plant", "a person wearing a red hat"], context_length=32) + assert tokens.dtype == np.int64 + assert tokens.shape == (2, 32) + assert tokens[0, :3].tolist() == [49406, 3912, 49407] + assert tokens[1, :8].tolist() == [ + 49406, + 320, + 2533, + 3309, + 320, + 736, + 3801, + 49407, + ] + + +def test_clip_tokenizer_rejects_prompt_over_context_limit(): + with pytest.raises(ValueError, match="32-token limit"): + tokenize(" ".join(["plant"] * 40), context_length=32) diff --git a/tests/test_efficient_sam.py b/tests/test_efficient_sam.py new file mode 100644 index 0000000..9c79e99 --- /dev/null +++ b/tests/test_efficient_sam.py @@ -0,0 +1,77 @@ +from unittest.mock import MagicMock, patch + +import numpy as np + +from samexporter.efficient_sam_onnx import EfficientSAMONNX + + +def model_input(name, shape): + value = MagicMock() + value.name = name + value.shape = shape + return value + + +class TestEfficientSAMONNX: + @patch("onnxruntime.InferenceSession") + def test_rgb_input_and_best_mask_selection(self, inference_session): + encoder = MagicMock() + encoder.get_inputs.return_value = [ + model_input("batched_images", ["batch", 3, "height", "width"]) + ] + encoder.run.return_value = [np.zeros((1, 256, 64, 64), np.float32)] + + decoder = MagicMock() + decoder.get_inputs.return_value = [ + model_input("image_embeddings", [1, 256, 64, 64]), + model_input("batched_point_coords", [1, 1, "points", 2]), + model_input("batched_point_labels", [1, 1, "points"]), + model_input("orig_im_size", [2]), + ] + masks = np.zeros((1, 1, 3, 2, 3), np.float32) + masks[0, 0, 2] = 1 + decoder.run.return_value = [ + masks, + np.array([[[0.1, 0.2, 0.9]]], np.float32), + ] + inference_session.side_effect = [encoder, decoder] + + model = EfficientSAMONNX("encoder.onnx", "decoder.onnx") + image = np.zeros((2, 3, 3), np.uint8) + image[..., 0] = 255 # OpenCV blue becomes RGB red. + embedding = model.encode(image) + encoded_input = encoder.run.call_args.args[1]["batched_images"] + assert encoded_input.shape == (1, 3, 2, 3) + assert encoded_input[0, 0, 0, 0] == 0 + assert encoded_input[0, 2, 0, 0] == 1 + + result = model.predict_masks( + embedding, [{"type": "rectangle", "data": [0, 0, 2, 1]}] + ) + assert result.shape == (1, 1, 2, 3) + assert result.all() + decoder_inputs = decoder.run.call_args.args[1] + np.testing.assert_array_equal( + decoder_inputs["batched_point_labels"], [[[2.0, 3.0]]] + ) + + @patch("onnxruntime.InferenceSession") + def test_rejects_unknown_decoder_contract(self, inference_session): + encoder = MagicMock() + encoder.get_inputs.return_value = [ + model_input("batched_images", [1, 3, -1, -1]) + ] + decoder = MagicMock() + decoder.get_inputs.return_value = [model_input("unknown", [1])] + inference_session.side_effect = [encoder, decoder] + model = EfficientSAMONNX("encoder.onnx", "decoder.onnx") + + try: + model.predict_masks( + {"image_embedding": np.zeros(1), "original_size": (10, 10)}, + [{"type": "point", "data": [1, 1], "label": 1}], + ) + except ValueError as error: + assert "Unsupported EfficientSAM decoder inputs" in str(error) + else: + raise AssertionError("Expected incompatible decoder to be rejected") diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..15608ae --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,31 @@ +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from samexporter.inference import visualize + + +def test_visualize_only_blends_masked_pixels(): + image = np.full((20, 20, 3), 200, dtype=np.uint8) + masks = np.zeros((1, 1, 20, 20), dtype=np.float32) + masks[0, 0, 5:10, 5:10] = 1 + + result = visualize(image, masks, [], "sam") + + np.testing.assert_array_equal(result[0, 0], image[0, 0]) + assert not np.array_equal(result[6, 6], image[6, 6]) + + +def test_visualize_sam3_instances_have_distinct_colors(): + image = np.full((40, 40, 3), 100, dtype=np.uint8) + masks = np.zeros((2, 1, 40, 40), dtype=np.bool_) + masks[0, 0, 2:12, 2:12] = True + masks[1, 0, 25:35, 25:35] = True + + result = visualize(image, masks, [], "sam3") + + np.testing.assert_array_equal(result[20, 20], image[20, 20]) + assert not np.array_equal(result[5, 5], result[30, 30]) diff --git a/tests/test_prompts.py b/tests/test_prompts.py new file mode 100644 index 0000000..6efee8d --- /dev/null +++ b/tests/test_prompts.py @@ -0,0 +1,28 @@ +import pytest + +from samexporter.prompts import geometric_prompt_arrays + + +def test_text_marks_are_ignored_when_geometric_prompt_exists(): + points, labels = geometric_prompt_arrays( + [ + {"type": "text", "data": "leaf"}, + {"type": "point", "data": [4, 5], "label": 1}, + ] + ) + assert points.tolist() == [[4.0, 5.0]] + assert labels.tolist() == [1.0] + + +@pytest.mark.parametrize( + "prompt, message", + [ + ([], "At least one"), + ([{"type": "point", "data": [1, 2], "label": 4}], "label"), + ([{"type": "rectangle", "data": [2, 2, 1, 3]}], "positive area"), + ([{"type": "polygon", "data": []}], "Unsupported"), + ], +) +def test_invalid_prompts_fail_clearly(prompt, message): + with pytest.raises(ValueError, match=message): + geometric_prompt_arrays(prompt) diff --git a/tests/test_runtime.py b/tests/test_runtime.py new file mode 100644 index 0000000..2870bb4 --- /dev/null +++ b/tests/test_runtime.py @@ -0,0 +1,42 @@ +from unittest.mock import patch + +import pytest + +from samexporter.runtime import get_onnx_providers + + +@patch("onnxruntime.get_available_providers") +def test_auto_provider_order_includes_accelerators_and_cpu(mock_available): + mock_available.return_value = [ + "CPUExecutionProvider", + "AzureExecutionProvider", + "CUDAExecutionProvider", + "TensorrtExecutionProvider", + ] + assert get_onnx_providers() == [ + "TensorrtExecutionProvider", + "CUDAExecutionProvider", + "AzureExecutionProvider", + "CPUExecutionProvider", + ] + + +@patch("onnxruntime.get_available_providers") +def test_provider_aliases_and_cpu_fallback(mock_available): + mock_available.return_value = [ + "CUDAExecutionProvider", + "OpenVINOExecutionProvider", + "CPUExecutionProvider", + ] + assert get_onnx_providers("openvino,cuda") == [ + "OpenVINOExecutionProvider", + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + + +@patch("onnxruntime.get_available_providers") +def test_unavailable_explicit_provider_has_actionable_error(mock_available): + mock_available.return_value = ["CPUExecutionProvider"] + with pytest.raises(ValueError, match="CUDAExecutionProvider.*Available"): + get_onnx_providers("cuda") diff --git a/tests/test_sam3_export.py b/tests/test_sam3_export.py new file mode 100644 index 0000000..3f1c729 --- /dev/null +++ b/tests/test_sam3_export.py @@ -0,0 +1,39 @@ +import torch + +from samexporter.export_sam3 import ( + prepare_fused_mlps_for_onnx, + prepare_rope_buffers_for_onnx, +) + + +class RopeModule(torch.nn.Module): + def __init__(self, *, prepared: bool): + super().__init__() + self.use_rope_real = True + self.register_buffer("freqs_cis", torch.ones(2, dtype=torch.complex64)) + if prepared: + self.register_buffer("freqs_cis_real", self.freqs_cis.real) + self.register_buffer("freqs_cis_imag", self.freqs_cis.imag) + + +def test_rope_preparation_keeps_required_complex_buffer(): + module = RopeModule(prepared=False) + prepare_rope_buffers_for_onnx(module) + assert hasattr(module, "freqs_cis") + assert hasattr(module, "freqs_cis_real") + assert hasattr(module, "freqs_cis_imag") + assert module.use_rope_real + + +def test_rope_preparation_is_idempotent_for_current_sam3(): + module = RopeModule(prepared=True) + real_buffer = module.freqs_cis_real + prepare_rope_buffers_for_onnx(module) + assert module.freqs_cis_real is real_buffer + + +def test_fused_mlp_preparation_ignores_unrelated_modules(): + module = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.ReLU()) + original_forward = module[0].forward + prepare_fused_mlps_for_onnx(module) + assert module[0].forward == original_forward diff --git a/tests/test_sam3_onnx.py b/tests/test_sam3_onnx.py index b26bf97..4662322 100644 --- a/tests/test_sam3_onnx.py +++ b/tests/test_sam3_onnx.py @@ -23,9 +23,9 @@ backbone_fpn_2 [1, 256, 72, 72] float language_mask [1, 32] bool language_features [32, 1, 256] float - box_coords [1, 1, 4] float - box_labels [1, 1] int64 - box_masks [1, 1] bool + box_coords [capacity, 1, 4] float + box_labels [capacity, 1] int64 + box_masks [1, capacity] bool Decoder outputs: [0] boxes (N, 4) float [1] scores (N,) float @@ -87,9 +87,9 @@ def _make_mock_session(input_specs, run_return=None): ("backbone_fpn_2", [1, 256, 72, 72], "tensor(float)"), ("language_mask", [1, 32], "tensor(bool)"), ("language_features", [32, 1, 256], "tensor(float)"), - ("box_coords", [1, 1, 4], "tensor(float)"), - ("box_labels", [1, 1], "tensor(int64)"), - ("box_masks", [1, 1], "tensor(bool)"), + ("box_coords", [8, 1, 4], "tensor(float)"), + ("box_labels", [8, 1], "tensor(int64)"), + ("box_masks", [1, 8], "tensor(bool)"), ] _ENCODER_INPUT_SPECS = [ @@ -157,12 +157,13 @@ def test_rectangle_box_coords(self, MockSession): inputs = args[1] self.assertIn("box_coords", inputs) np.testing.assert_allclose( - inputs["box_coords"], - np.array([[[0.5, 0.5, 0.5, 0.6]]], dtype=np.float32), + inputs["box_coords"][0, 0], + np.array([0.5, 0.5, 0.5, 0.6], dtype=np.float32), atol=1e-5, ) # box_masks should be False (real box provided). - self.assertFalse(inputs["box_masks"].all()) + self.assertFalse(inputs["box_masks"][0, 0]) + self.assertTrue(inputs["box_masks"][0, 1:].all()) @patch("onnxruntime.InferenceSession") def test_empty_prompt_uses_dummy_box(self, MockSession): @@ -242,11 +243,164 @@ def test_point_box_coords(self, MockSession): args, _kwargs = dec_sess.run.call_args inputs = args[1] np.testing.assert_allclose( - inputs["box_coords"], - np.array([[[0.5, 0.5, 0.01, 0.01]]], dtype=np.float32), + inputs["box_coords"][0, 0], + np.array([0.5, 0.5, 0.01, 0.01], dtype=np.float32), atol=1e-5, ) + @patch("onnxruntime.InferenceSession") + def test_multiple_geometric_prompts_are_forwarded(self, MockSession): + enc_sess = _make_mock_session(_ENCODER_INPUT_SPECS) + dec_sess = _make_mock_session( + _DECODER_INPUT_SPECS, + run_return=[ + np.zeros((0, 4), dtype=np.float32), + np.zeros((0,), dtype=np.float32), + np.zeros((0, 1, 100, 200), dtype=np.bool_), + ], + ) + MockSession.side_effect = [enc_sess, dec_sess] + model = SegmentAnything3ONNX("enc.onnx", "dec.onnx") + embedding = { + "original_size": (100, 200), + **{ + name: np.zeros(1) + for name in ( + "vision_pos_enc_0", + "vision_pos_enc_1", + "vision_pos_enc_2", + "backbone_fpn_0", + "backbone_fpn_1", + "backbone_fpn_2", + "language_embeds", + ) + }, + "language_mask": np.zeros((1, 32), dtype=np.bool_), + "language_features": np.zeros((32, 1, 256), dtype=np.float32), + } + + model.predict_masks( + embedding, + [ + {"type": "rectangle", "data": [10, 20, 50, 80]}, + {"type": "point", "data": [100, 50], "label": 0}, + ], + ) + inputs = dec_sess.run.call_args.args[1] + self.assertEqual(inputs["box_coords"].shape, (8, 1, 4)) + np.testing.assert_array_equal(inputs["box_labels"][:2], [[1], [0]]) + self.assertFalse(inputs["box_masks"][0, :2].any()) + self.assertTrue(inputs["box_masks"][0, 2:].all()) + + with self.assertRaisesRegex(ValueError, "at most 8 geometric prompts"): + model.predict_masks( + embedding, + [{"type": "point", "data": [100, 50], "label": 1} for _ in range(9)], + ) + + +# --------------------------------------------------------------------------- +# SAM3 query selection tests (pure NumPy, no model session) +# --------------------------------------------------------------------------- + + +class TestSAM3DetectionSelection(unittest.TestCase): + def setUp(self): + self.masks = np.zeros((3, 1, 100, 100), dtype=np.bool_) + self.masks[0, 0, 10:50, 10:50] = True + self.masks[1, 0, 12:49, 12:49] = True + self.masks[2, 0, 60:90, 60:90] = True + self.scores = np.array([0.8, 0.7, 0.95], dtype=np.float32) + self.boxes = np.array( + [[10, 10, 50, 50], [12, 12, 49, 49], [60, 60, 90, 90]], + dtype=np.float32, + ) + + def select(self, prompt, **overrides): + options = { + "confidence_threshold": 0.5, + "nms_threshold": 0.5, + "nms_mode": "mask", + "max_instances": None, + "prefer_prompted_region": False, + } + options.update(overrides) + return SegmentAnything3ONNX.select_detections( + self.masks, self.scores, self.boxes, prompt, **options + ) + + def test_nms_removes_duplicate_queries_and_preserves_score_order(self): + masks, scores, boxes = self.select([]) + + self.assertEqual(len(masks), 2) + np.testing.assert_allclose(scores, [0.95, 0.8]) + np.testing.assert_array_equal(boxes, [self.boxes[2], self.boxes[0]]) + + def test_box_nms_is_available_as_an_explicit_fast_mode(self): + _, scores, boxes = self.select([], nms_mode="box") + + np.testing.assert_allclose(scores, [0.95, 0.8]) + np.testing.assert_array_equal(boxes, [self.boxes[2], self.boxes[0]]) + + def test_nms_can_be_disabled(self): + _, scores, _ = self.select([], nms_mode="none") + np.testing.assert_allclose(scores, [0.95, 0.8, 0.7]) + + def test_rectangle_selection_prefers_prompted_region_before_confidence(self): + masks, scores, boxes = self.select( + [{"type": "rectangle", "data": [10, 10, 50, 50]}], + max_instances=1, + prefer_prompted_region=True, + ) + + self.assertEqual(len(masks), 1) + self.assertAlmostEqual(float(scores[0]), 0.8, places=5) + np.testing.assert_array_equal(boxes[0], self.boxes[0]) + + def test_positive_point_selection_uses_mask_membership(self): + _, scores, boxes = self.select( + [{"type": "point", "data": [20, 20], "label": 1}], + max_instances=1, + prefer_prompted_region=True, + ) + + self.assertAlmostEqual(float(scores[0]), 0.8, places=5) + np.testing.assert_array_equal(boxes[0], self.boxes[0]) + + def test_confidence_filter_and_instance_limit(self): + _, scores, _ = self.select( + [], confidence_threshold=0.75, nms_threshold=None, max_instances=1 + ) + np.testing.assert_allclose(scores, [0.95]) + + def test_empty_and_invalid_selection_inputs(self): + masks, scores, boxes = self.select([], confidence_threshold=0.99) + self.assertEqual(masks.shape, (0, 1, 100, 100)) + self.assertEqual(scores.shape, (0,)) + self.assertEqual(boxes.shape, (0, 4)) + + for options, message in ( + ({"confidence_threshold": -0.1}, "confidence_threshold"), + ({"nms_threshold": 1.1}, "nms_threshold"), + ({"nms_mode": "invalid"}, "nms_mode"), + ({"max_instances": 0}, "max_instances"), + ): + with self.assertRaisesRegex(ValueError, message): + self.select([], **options) + + with self.assertRaisesRegex(ValueError, "equal lengths"): + SegmentAnything3ONNX.select_detections( + self.masks[:2], + self.scores, + self.boxes, + [], + confidence_threshold=0.5, + nms_threshold=0.7, + nms_mode="mask", + max_instances=None, + prefer_prompted_region=False, + ) + # --------------------------------------------------------------------------- # SAM3ImageEncoder tests @@ -350,6 +504,19 @@ def test_returns_masks_scores_boxes(self, MockSession): np.testing.assert_array_equal(ret_scores, scores) np.testing.assert_array_equal(ret_boxes, boxes) + def test_raw_head_postprocessing_matches_official_contract(self): + raw = { + "pred_boxes": np.array([[[0.5, 0.5, 0.5, 0.4]]], np.float32), + "pred_logits": np.array([[[2.0]]], np.float32), + "pred_masks": np.ones((1, 1, 2, 2), np.float32), + "presence_logit_dec": np.array([[1.0]], np.float32), + } + masks, scores, boxes = SAM3ImageDecoder.postprocess_raw_outputs(raw, (100, 200)) + self.assertEqual(masks.shape, (1, 1, 100, 200)) + self.assertTrue(masks.all()) + self.assertAlmostEqual(float(scores[0]), 0.6439, places=3) + np.testing.assert_allclose(boxes, [[50, 30, 150, 70]], atol=1e-5) + @patch("onnxruntime.InferenceSession") def test_dummy_language_inputs_when_none(self, MockSession): """Decoder supplies correct dummy tensors when language inputs are None.""" diff --git a/tests/test_sam_variants.py b/tests/test_sam_variants.py index 360a491..39a46ea 100644 --- a/tests/test_sam_variants.py +++ b/tests/test_sam_variants.py @@ -22,6 +22,13 @@ def test_sam1_logic(self, mock_session): mock_input.name = "image" mock_input.shape = [1, 3, 1024, 1024] mock_sess_instance.get_inputs.return_value = [mock_input] + mock_output = MagicMock() + mock_output.name = "low_res_masks" + mock_sess_instance.get_outputs.return_value = [ + MagicMock(name="masks"), + MagicMock(name="iou_predictions"), + mock_output, + ] mock_session.return_value = mock_sess_instance model = SegmentAnythingONNX("dummy_enc.onnx", "dummy_dec.onnx") @@ -37,11 +44,59 @@ def test_sam1_logic(self, mock_session): mock_sess_instance.run.return_value = [ np.zeros((1, 1, 100, 100)), np.zeros(1), - np.zeros(1), + np.zeros((1, 1, 256, 256)), ] masks = model.predict_masks(embedding, prompt) self.assertEqual(len(masks.shape), 4) + @patch("onnxruntime.InferenceSession") + def test_sam1_embedded_preprocess_preserves_portrait_aspect_and_rgb( + self, mock_session + ): + encoder = MagicMock() + encoder_input = MagicMock() + encoder_input.name = "input_image" + encoder_input.shape = ["image_height", "image_width", 3] + encoder.get_inputs.return_value = [encoder_input] + encoder.run.return_value = [np.zeros((1, 256, 64, 64), np.float32)] + decoder = MagicMock() + mock_session.side_effect = [encoder, decoder] + + model = SegmentAnythingONNX("encoder.onnx", "decoder.onnx") + image = np.zeros((1200, 600, 3), np.uint8) + image[..., 0] = 255 + embedding = model.encode(image) + + input_image = encoder.run.call_args.args[1]["input_image"] + self.assertEqual(input_image.shape, (1024, 512, 3)) + self.assertEqual(input_image[0, 0].tolist(), [0.0, 0.0, 255.0]) + self.assertEqual(embedding["resized_size"], (1024, 512)) + + @patch("onnxruntime.InferenceSession") + def test_sam1_postprocess_uses_portrait_crop(self, mock_session): + encoder = MagicMock() + encoder_input = MagicMock() + encoder_input.name = "input_image" + encoder_input.shape = ["image_height", "image_width", 3] + encoder.get_inputs.return_value = [encoder_input] + + decoder = MagicMock() + decoder.get_outputs.return_value = [ + MagicMock(name="masks"), + MagicMock(name="iou_predictions"), + MagicMock(name="low_res_masks"), + ] + mock_session.side_effect = [encoder, decoder] + model = SegmentAnythingONNX("encoder.onnx", "decoder.onnx") + + low_res = np.zeros((1, 1, 256, 256), np.float32) + low_res[:, :, :, :128] = 1 + result = model.postprocess_masks( + low_res, original_size=(1200, 600), resized_size=(1024, 512) + ) + self.assertEqual(result.shape, (1, 1, 1200, 600)) + self.assertGreater(float(result.mean()), 0.99) + @patch("onnxruntime.InferenceSession") def test_sam2_logic(self, mock_session): # Setup mock session diff --git a/third_party/sam2 b/third_party/sam2 new file mode 160000 index 0000000..2b90b9f --- /dev/null +++ b/third_party/sam2 @@ -0,0 +1 @@ +Subproject commit 2b90b9f5ceec907a1c18123530e92e794ad901a4 diff --git a/third_party/segment-anything b/third_party/segment-anything new file mode 160000 index 0000000..dca509f --- /dev/null +++ b/third_party/segment-anything @@ -0,0 +1 @@ +Subproject commit dca509fe793f601edb92606367a655c15ac00fdf diff --git a/third_party_licenses/OPENAI_CLIP_LICENSE b/third_party_licenses/OPENAI_CLIP_LICENSE new file mode 100644 index 0000000..c123b69 --- /dev/null +++ b/third_party_licenses/OPENAI_CLIP_LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 OpenAI + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/visual_results/README.md b/visual_results/README.md new file mode 100644 index 0000000..a02fc5c --- /dev/null +++ b/visual_results/README.md @@ -0,0 +1,75 @@ +# Visual model results + +This directory contains real-model inference evidence for human review. Open the +PNGs at full size to inspect boundaries, prompt placement, holes, and leakage. +Blue is the selected mask for single-result families, green is a positive +point/rectangle, and red is a negative point. SAM3 multi-instance results use +distinct colors, contours, and numbers. Mock sessions are used only by isolated +unit tests; no mocked image is presented here as quality evidence. + +[Open the cross-model contact sheet](model_comparison.jpg) for a quick visual +comparison, then inspect the linked full-resolution PNGs below. + +## Current verification status + +| Family | Real artifacts produced in this change | Validation | +|---|---|---| +| SAM ViT-B | [truck box](sam_vit_b/truck_box.png), [truck point](sam_vit_b/truck_point.png), [plants box](sam_vit_b/plants_box.png), [plants point](sam_vit_b/plants_point.png), [refined plant](sam_vit_b/plants_box_refined.png) | Official checkpoint, pinned source export, CPU ONNX | +| MobileSAM | [truck box](mobile_sam/truck_box.png), [truck point](mobile_sam/truck_point.png), [plants box](mobile_sam/plants_box.png), [plants point](mobile_sam/plants_point.png), [refined plant](mobile_sam/plants_box_refined.png) | Official checkpoint, newly exported encoder and decoder, CPU ONNX | +| EfficientSAM-Ti | [truck box](efficient_sam_ti/truck_box.png), [truck point](efficient_sam_ti/truck_point.png), [plants box](efficient_sam_ti/plants_box.png), [plants point](efficient_sam_ti/plants_point.png), [refined plant](efficient_sam_ti/plants_box_refined.png) | Official Apache-2.0 split ONNX files, CPU ONNX | +| SAM 2.1 Tiny | [truck box](sam2_1_tiny/truck_box.png), [truck point](sam2_1_tiny/truck_point.png), [plants box](sam2_1_tiny/plants_box.png), [plants point](sam2_1_tiny/plants_point.png), [refined plant](sam2_1_tiny/plants_box_refined.png) | Official checkpoint, pinned source export, CPU ONNX | +| SAM 3 | [truck text](sam3/truck_text.png), [truck box selection](sam3/truck_box.png), [truck text + point](sam3/truck_point.png), [all 20 plants](sam3/plants_text.png), [top five plants](sam3/plants_text_top5.png), [plants box selection](sam3/plants_box.png), [two-mark selection](sam3/plants_box_refined.png) | Latest official gated checkpoint, three newly exported ONNX graphs, official-style mask NMS, CPU ONNX | + +Older checked-in references remain under `reference_sam/` and +`reference_sam2/`. They are retained for comparison but are not counted as +evidence that the current exporters ran successfully. + +## Numerical checks + +- SAM ViT-B PyTorch/ONNX mask IoU: truck box 0.9921, truck point 1.0000, + plants box 0.9858, plants point 1.0000. Encoder embeddings had approximately + `1e-7` mean absolute error and cosine similarity 1.0000. +- MobileSAM PyTorch/ONNX mask IoU: truck box 0.9961, truck point 0.9188, + plants box 0.9962, plants point 0.9945. The lower truck-point agreement is an + ambiguous single-click/TinyViT case; both results visibly select the door. +- SAM 2.1 Tiny PyTorch/ONNX mask IoU: truck box 0.99757, truck point 0.99576, + plants box 0.99681, plants point 0.99761. +- EfficientSAM-Ti mask area fractions: truck box 0.2921, truck point 0.0366, + plants box 0.0324, plants point 0.0104. Both positive points are contained; + box-mask pixels outside the prompt are 0.14% for truck and 0.02% for plants. +- SAM 3 real CPU CLI runs retain one mask for truck text/box/text-plus-point, + one mask for each plant geometry selection, 20 masks for text-only `plant`, + and five when capped with `--max_instances 5`. Complete cold-process runs took + 18.9–31.8 seconds and peaked at 8.2–8.5 GiB resident memory on this host. + +The refined portrait prompt combines a tight rectangle with a positive point on +the pot. Every tested family selects the intended whole plant/pot. SAM3 geometry +still acts as a visual concept exemplar internally; the CLI's default `auto` +mode now ranks prompt overlap and returns the best local match, while +`--sam3_output_mode all` retains the broader concept-discovery behavior. + +## Run logs + +- [SAM 3 current export](sam3/export_current.log) +- [SAM 3 text run and resource usage](sam3/truck_text.log) +- [SAM 3 plant box selection](sam3/plants_box.log) +- [SAM 3 refined plant selection](sam3/plants_box_refined.log) +- [SAM 3 all-plant text discovery](sam3/plants_text.log) +- [SAM 3 capped top-five discovery](sam3/plants_text_top5.log) +- [SAM 3 packaged-tokenizer real-model validation](sam3/tokenizer_validation.log) +- [EfficientSAM refined portrait timing](efficient_sam_ti/plants_box_refined.log) +- Earlier SAM3 failure logs remain as debugging evidence for the prompt-layout, + fixed-token, and missing-temporary-artifact failures found before the final + exporter and validation runs. + +## Reproduce + +```bash +bash download_all_models.sh +bash test_comprehensive.sh +``` + +Every future test writes both a PNG and sibling `.log` under +`visual_results/runs/` by default. Set `SAMEXPORTER_RESULTS_DIR` to another +persistent path when comparing runs. Promote only reviewed representative +outputs into a named family directory. diff --git a/visual_results/efficient_sam_ti/plants_box.log b/visual_results/efficient_sam_ti/plants_box.log new file mode 100644 index 0000000..1ad6b96 --- /dev/null +++ b/visual_results/efficient_sam_ti/plants_box.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=1.27 max_rss_kb=676400 diff --git a/visual_results/efficient_sam_ti/plants_box.png b/visual_results/efficient_sam_ti/plants_box.png new file mode 100644 index 0000000..26400da Binary files /dev/null and b/visual_results/efficient_sam_ti/plants_box.png differ diff --git a/visual_results/efficient_sam_ti/plants_box_refined.log b/visual_results/efficient_sam_ti/plants_box_refined.log new file mode 100644 index 0000000..c9323b6 --- /dev/null +++ b/visual_results/efficient_sam_ti/plants_box_refined.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=1.56 max_rss_kb=675680 diff --git a/visual_results/efficient_sam_ti/plants_box_refined.png b/visual_results/efficient_sam_ti/plants_box_refined.png new file mode 100644 index 0000000..5da7a3a Binary files /dev/null and b/visual_results/efficient_sam_ti/plants_box_refined.png differ diff --git a/visual_results/efficient_sam_ti/plants_point.log b/visual_results/efficient_sam_ti/plants_point.log new file mode 100644 index 0000000..741da98 --- /dev/null +++ b/visual_results/efficient_sam_ti/plants_point.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=1.38 max_rss_kb=676652 diff --git a/visual_results/efficient_sam_ti/plants_point.png b/visual_results/efficient_sam_ti/plants_point.png new file mode 100644 index 0000000..bcbfc3e Binary files /dev/null and b/visual_results/efficient_sam_ti/plants_point.png differ diff --git a/visual_results/efficient_sam_ti/truck_box.log b/visual_results/efficient_sam_ti/truck_box.log new file mode 100644 index 0000000..6bdb3a6 --- /dev/null +++ b/visual_results/efficient_sam_ti/truck_box.log @@ -0,0 +1 @@ +elapsed_s=1.33 max_rss_kb=733376 diff --git a/visual_results/efficient_sam_ti/truck_box.png b/visual_results/efficient_sam_ti/truck_box.png new file mode 100644 index 0000000..38b12cf Binary files /dev/null and b/visual_results/efficient_sam_ti/truck_box.png differ diff --git a/visual_results/efficient_sam_ti/truck_point.log b/visual_results/efficient_sam_ti/truck_point.log new file mode 100644 index 0000000..ea3371c --- /dev/null +++ b/visual_results/efficient_sam_ti/truck_point.log @@ -0,0 +1 @@ +elapsed_s=1.38 max_rss_kb=709656 diff --git a/visual_results/efficient_sam_ti/truck_point.png b/visual_results/efficient_sam_ti/truck_point.png new file mode 100644 index 0000000..799afbc Binary files /dev/null and b/visual_results/efficient_sam_ti/truck_point.png differ diff --git a/visual_results/mobile_sam/plants_box.log b/visual_results/mobile_sam/plants_box.log new file mode 100644 index 0000000..7be7032 --- /dev/null +++ b/visual_results/mobile_sam/plants_box.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=0.88 max_rss_kb=480648 diff --git a/visual_results/mobile_sam/plants_box.png b/visual_results/mobile_sam/plants_box.png new file mode 100644 index 0000000..c3bf34b Binary files /dev/null and b/visual_results/mobile_sam/plants_box.png differ diff --git a/visual_results/mobile_sam/plants_box_refined.log b/visual_results/mobile_sam/plants_box_refined.log new file mode 100644 index 0000000..f3e350a --- /dev/null +++ b/visual_results/mobile_sam/plants_box_refined.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=0.83 max_rss_kb=478592 diff --git a/visual_results/mobile_sam/plants_box_refined.png b/visual_results/mobile_sam/plants_box_refined.png new file mode 100644 index 0000000..afe2168 Binary files /dev/null and b/visual_results/mobile_sam/plants_box_refined.png differ diff --git a/visual_results/mobile_sam/plants_point.log b/visual_results/mobile_sam/plants_point.log new file mode 100644 index 0000000..8aca748 --- /dev/null +++ b/visual_results/mobile_sam/plants_point.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=0.88 max_rss_kb=480852 diff --git a/visual_results/mobile_sam/plants_point.png b/visual_results/mobile_sam/plants_point.png new file mode 100644 index 0000000..1101e7d Binary files /dev/null and b/visual_results/mobile_sam/plants_point.png differ diff --git a/visual_results/mobile_sam/truck_box.log b/visual_results/mobile_sam/truck_box.log new file mode 100644 index 0000000..470b5b8 --- /dev/null +++ b/visual_results/mobile_sam/truck_box.log @@ -0,0 +1 @@ +elapsed_s=0.84 max_rss_kb=509000 diff --git a/visual_results/mobile_sam/truck_box.png b/visual_results/mobile_sam/truck_box.png new file mode 100644 index 0000000..94d8d90 Binary files /dev/null and b/visual_results/mobile_sam/truck_box.png differ diff --git a/visual_results/mobile_sam/truck_point.log b/visual_results/mobile_sam/truck_point.log new file mode 100644 index 0000000..f605035 --- /dev/null +++ b/visual_results/mobile_sam/truck_point.log @@ -0,0 +1 @@ +elapsed_s=0.82 max_rss_kb=485892 diff --git a/visual_results/mobile_sam/truck_point.png b/visual_results/mobile_sam/truck_point.png new file mode 100644 index 0000000..1d44432 Binary files /dev/null and b/visual_results/mobile_sam/truck_point.png differ diff --git a/visual_results/model_comparison.jpg b/visual_results/model_comparison.jpg new file mode 100644 index 0000000..20406b1 Binary files /dev/null and b/visual_results/model_comparison.jpg differ diff --git a/visual_results/reference_sam/plants_01.png b/visual_results/reference_sam/plants_01.png new file mode 100644 index 0000000..94a10f1 Binary files /dev/null and b/visual_results/reference_sam/plants_01.png differ diff --git a/visual_results/reference_sam/plants_02.png b/visual_results/reference_sam/plants_02.png new file mode 100644 index 0000000..3d466a2 Binary files /dev/null and b/visual_results/reference_sam/plants_02.png differ diff --git a/visual_results/reference_sam/truck.png b/visual_results/reference_sam/truck.png new file mode 100644 index 0000000..a35ed80 Binary files /dev/null and b/visual_results/reference_sam/truck.png differ diff --git a/visual_results/reference_sam2/sam2_truck.png b/visual_results/reference_sam2/sam2_truck.png new file mode 100644 index 0000000..76af94a Binary files /dev/null and b/visual_results/reference_sam2/sam2_truck.png differ diff --git a/visual_results/sam2_1_tiny/plants_box.log b/visual_results/sam2_1_tiny/plants_box.log new file mode 100644 index 0000000..028586e --- /dev/null +++ b/visual_results/sam2_1_tiny/plants_box.log @@ -0,0 +1,5 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:44:14.648523847 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/image_encoder/trunk/Concat_3_output_0' source:{4} target:{5}. Falling back to lenient merge. +infer time: 863.88 ms +infer time: 32.90 ms +elapsed_s=2.45 max_rss_kb=1020564 diff --git a/visual_results/sam2_1_tiny/plants_box.png b/visual_results/sam2_1_tiny/plants_box.png new file mode 100644 index 0000000..56be928 Binary files /dev/null and b/visual_results/sam2_1_tiny/plants_box.png differ diff --git a/visual_results/sam2_1_tiny/plants_box_refined.log b/visual_results/sam2_1_tiny/plants_box_refined.log new file mode 100644 index 0000000..c72a46a --- /dev/null +++ b/visual_results/sam2_1_tiny/plants_box_refined.log @@ -0,0 +1,5 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:44:19.444272219 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/image_encoder/trunk/Concat_3_output_0' source:{4} target:{5}. Falling back to lenient merge. +infer time: 1431.26 ms +infer time: 66.03 ms +elapsed_s=3.21 max_rss_kb=1034072 diff --git a/visual_results/sam2_1_tiny/plants_box_refined.png b/visual_results/sam2_1_tiny/plants_box_refined.png new file mode 100644 index 0000000..50d37cf Binary files /dev/null and b/visual_results/sam2_1_tiny/plants_box_refined.png differ diff --git a/visual_results/sam2_1_tiny/plants_point.log b/visual_results/sam2_1_tiny/plants_point.log new file mode 100644 index 0000000..a2f3dc2 --- /dev/null +++ b/visual_results/sam2_1_tiny/plants_point.log @@ -0,0 +1,5 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:44:17.028948848 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/image_encoder/trunk/Concat_3_output_0' source:{4} target:{5}. Falling back to lenient merge. +infer time: 813.78 ms +infer time: 16.51 ms +elapsed_s=2.26 max_rss_kb=1061128 diff --git a/visual_results/sam2_1_tiny/plants_point.png b/visual_results/sam2_1_tiny/plants_point.png new file mode 100644 index 0000000..7999728 Binary files /dev/null and b/visual_results/sam2_1_tiny/plants_point.png differ diff --git a/visual_results/sam2_1_tiny/truck_box.log b/visual_results/sam2_1_tiny/truck_box.log new file mode 100644 index 0000000..ba8b7c1 --- /dev/null +++ b/visual_results/sam2_1_tiny/truck_box.log @@ -0,0 +1,4 @@ +2026-08-30 21:44:09.454098394 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/image_encoder/trunk/Concat_3_output_0' source:{4} target:{5}. Falling back to lenient merge. +infer time: 913.04 ms +infer time: 17.56 ms +elapsed_s=2.34 max_rss_kb=1066720 diff --git a/visual_results/sam2_1_tiny/truck_box.png b/visual_results/sam2_1_tiny/truck_box.png new file mode 100644 index 0000000..9c9e4e9 Binary files /dev/null and b/visual_results/sam2_1_tiny/truck_box.png differ diff --git a/visual_results/sam2_1_tiny/truck_point.log b/visual_results/sam2_1_tiny/truck_point.log new file mode 100644 index 0000000..d9c3d54 --- /dev/null +++ b/visual_results/sam2_1_tiny/truck_point.log @@ -0,0 +1,4 @@ +2026-08-30 21:44:11.802843379 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/image_encoder/trunk/Concat_3_output_0' source:{4} target:{5}. Falling back to lenient merge. +infer time: 1222.32 ms +infer time: 43.61 ms +elapsed_s=2.72 max_rss_kb=1145708 diff --git a/visual_results/sam2_1_tiny/truck_point.png b/visual_results/sam2_1_tiny/truck_point.png new file mode 100644 index 0000000..34da42a Binary files /dev/null and b/visual_results/sam2_1_tiny/truck_point.png differ diff --git a/visual_results/sam3/export_current.log b/visual_results/sam3/export_current.log new file mode 100644 index 0000000..9a8823a --- /dev/null +++ b/visual_results/sam3/export_current.log @@ -0,0 +1,106 @@ +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/model_builder.py:8: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81. + import pkg_resources +/home/vietanhdev/miniconda3/lib/python3.13/site-packages/timm/models/layers/__init__.py:49: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers + warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.layers", FutureWarning) +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/model/vitdet.py:251: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + size = int(math.sqrt(xy_num)) +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/model/vitdet.py:252: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + assert size * size == xy_num +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/model/vitdet.py:254: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + if size != h or size != w: +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/model/vitdet.py:154: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + if pad_h > 0 or pad_w > 0: +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/sam/rope.py:100: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + assert xk.shape[-2] != 0 +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/sam/rope.py:53: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + assert freqs_cis.shape == (x.shape[-2], x.shape[-1]) +/home/vietanhdev/Workspaces/samexporter/sam3/sam3/model/vitdet.py:184: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + if Hp > H or Wp > W: +/home/vietanhdev/Workspaces/samexporter/samexporter/export_sam3.py:145: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! + assert text_memory.shape[1] == inputs_embeds.shape[1] +/home/vietanhdev/miniconda3/lib/python3.13/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead. + return cls.__new__(cls, *args) +/home/vietanhdev/miniconda3/lib/python3.13/site-packages/onnxscript/function_libs/torch_lib/ops/vision.py:69: UserWarning: ONNX export for RoIAlign with a non-zero sampling_ratio is not supported. The model will be exported with a sampling_ratio of 0. + sampling_ratio = _process_sampling_ratio_for_roi_align(sampling_ratio) +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +WARNING:onnxscript.optimizer._constant_folding:Skipping constant folding for op Split with multiple outputs. +Exporting Image Encoder... +Saved Image Encoder to /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx +Exporting Language Encoder... +Saved Language Encoder to /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx +Exporting Decoder... +[torch.onnx] Obtain model graph for `SAM3Decoder([...]` with `torch.export.export(..., strict=False)`... +[torch.onnx] Obtain model graph for `SAM3Decoder([...]` with `torch.export.export(..., strict=False)`... ✅ +[torch.onnx] Run decompositions... +[torch.onnx] Run decompositions... ✅ +[torch.onnx] Translate the graph into ONNX... +[torch.onnx] Translate the graph into ONNX... ✅ +[torch.onnx] Optimize the ONNX graph... +[torch.onnx] Optimize the ONNX graph... ✅ +Saved Decoder to /tmp/samexporter_sam3/onnx/sam3_decoder.onnx diff --git a/visual_results/sam3/fixed_capacity_validation.log b/visual_results/sam3/fixed_capacity_validation.log new file mode 100644 index 0000000..556b324 --- /dev/null +++ b/visual_results/sam3/fixed_capacity_validation.log @@ -0,0 +1,33 @@ +2026-08-30 18:37:50.498534337 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +libpng warning: iCCP: known incorrect sRGB profile +images/truck.jpg encode_s 23.911 +truck_text shape [1, 1, 1200, 1800] areas [0.294588] decode_s 3.515 +truck_box shape [1, 1, 1200, 1800] areas [0.291464] decode_s 3.71 +truck_point shape [1, 1, 1200, 1800] areas [0.292979] decode_s 4.255 +images/plants.png encode_s 23.265 +plants_text shape [20, 1, 1481, 987] areas [0.013404, 0.006604, 0.003406, 0.010622, 0.004746, 0.055561, 0.003844, 0.005373, 0.006701, 0.018207, 0.011179, 0.011217, 0.00285, 0.008543, 0.004404, 0.011023, 0.000697, 0.006018, 0.016192, 0.001779] decode_s 2.973 +plants_box shape [26, 1, 1481, 987] areas [0.016781, 0.003432, 0.01105, 0.000718, 0.008653, 0.091354, 0.006821, 0.006648, 0.002176, 0.00389, 0.005366, 0.007269, 0.024027, 0.004747, 0.020257, 0.019277, 0.008439, 0.011315, 0.007912, 0.001778, 0.034442, 0.016682, 0.017514, 0.00432, 0.002894, 0.011067] decode_s 3.374 +plants_box_refined shape [23, 1, 1481, 987] areas [0.013437, 0.006615, 0.003421, 0.004759, 0.010686, 0.056078, 0.002851, 0.003856, 0.001775, 0.005398, 0.006742, 0.018263, 0.011327, 0.011244, 0.008429, 0.008019, 0.004508, 0.01827, 0.000704, 0.006055, 0.016303, 0.004328, 0.011018] decode_s 3.886 + Command being timed: "python -" + User time (seconds): 500.39 + System time (seconds): 24.30 + Percent of CPU this job got: 671% + Elapsed (wall clock) time (h:mm:ss or m:ss): 1:18.19 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 12615640 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 211 + Minor (reclaiming a frame) page faults: 4505567 + Voluntary context switches: 17684 + Involuntary context switches: 4170515 + Swaps: 0 + File system inputs: 29784 + File system outputs: 22400 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/geometry_only_validation.log b/visual_results/sam3/geometry_only_validation.log new file mode 100644 index 0000000..2f45d0c --- /dev/null +++ b/visual_results/sam3/geometry_only_validation.log @@ -0,0 +1,42 @@ +Traceback (most recent call last): + File "", line 7, in + File "/home/vietanhdev/Workspaces/samexporter/samexporter/sam3_onnx.py", line 21, in __init__ + self.image_encoder = SAM3ImageEncoder(image_encoder_path, providers) + ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/home/vietanhdev/Workspaces/samexporter/samexporter/sam3_onnx.py", line 207, in __init__ + self.session = onnxruntime.InferenceSession( + ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^ + path, providers=get_onnx_providers(providers) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + ) + ^ + File "/home/vietanhdev/miniconda3/lib/python3.13/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 534, in __init__ + self._create_inference_session(providers, provider_options, disabled_optimizers) + ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/home/vietanhdev/miniconda3/lib/python3.13/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 629, in _create_inference_session + sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model) +onnxruntime.capi.onnxruntime_pybind11_state.NoSuchFile: [ONNXRuntimeError] : 3 : NO_SUCHFILE : Load model from /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx failed:Load model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx failed. File doesn't exist +Command exited with non-zero status 1 + Command being timed: "python -" + User time (seconds): 1.62 + System time (seconds): 0.04 + Percent of CPU this job got: 736% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:00.22 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 75524 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 266 + Minor (reclaiming a frame) page faults: 8793 + Voluntary context switches: 462 + Involuntary context switches: 312 + Swaps: 0 + File system inputs: 64792 + File system outputs: 248 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 1 diff --git a/visual_results/sam3/multi_prompt_validation.log b/visual_results/sam3/multi_prompt_validation.log new file mode 100644 index 0000000..b7de75a --- /dev/null +++ b/visual_results/sam3/multi_prompt_validation.log @@ -0,0 +1,50 @@ +2026-08-30 18:28:33.032410084 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +libpng warning: iCCP: known incorrect sRGB profile +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 18:29:27.499919761 [E:onnxruntime:, sequential_executor.cc:671 ExecuteKernel] Non-zero status code returned while running Split node. Name:'/geometry_encoder/Split_1' Status Message: Cannot split using values in 'split' attribute. Axis=0 Input shape={2,1,4} NumOutputs=1 Num entries in 'split' (must equal number of outputs) was 1 Sum of sizes in 'split' (must equal size of selected axis) was 1 +truck_text shape [1, 1, 1200, 1800] areas [0.29458842592592593] encode_s 19.043 decode_s 3.672 +truck_box shape [1, 1, 1200, 1800] areas [0.2890041666666667] encode_s 0.0 decode_s 3.153 +truck_point shape [0, 1, 1200, 1800] areas [] encode_s 0.0 decode_s 3.707 +plants_text shape [20, 1, 1481, 987] areas [0.013403824328012987, 0.006603741960818117, 0.0034061981998252775, 0.010622221218856615, 0.004746375398752315, 0.05556091444005016, 0.003844030464916295, 0.005373022828163834, 0.006700885994635187, 0.018206981098644295, 0.011179089131019252, 0.011217399454214717, 0.002850014400576844, 0.008543202072588485, 0.004404318941649958, 0.011023111386580579, 0.000697111059574605, 0.0060181413062588805, 0.016191584453397204, 0.0017793776898464646] encode_s 16.843 decode_s 3.233 +Traceback (most recent call last): + File "", line 28, in + File "/home/vietanhdev/Workspaces/samexporter/samexporter/sam3_onnx.py", line 155, in predict_masks + masks, scores, _boxes = self.decoder( + ~~~~~~~~~~~~^ + original_size, + ^^^^^^^^^^^^^^ + ...<11 lines>... + box_masks_np, + ^^^^^^^^^^^^^ + ) + ^ + File "/home/vietanhdev/Workspaces/samexporter/samexporter/sam3_onnx.py", line 348, in __call__ + outputs = self.session.run(None, model_inputs) + File "/home/vietanhdev/miniconda3/lib/python3.13/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 326, in run + return self._sess.run(output_names, input_feed, run_options) + ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Split node. Name:'/geometry_encoder/Split_1' Status Message: Cannot split using values in 'split' attribute. Axis=0 Input shape={2,1,4} NumOutputs=1 Num entries in 'split' (must equal number of outputs) was 1 Sum of sizes in 'split' (must equal size of selected axis) was 1 +Command exited with non-zero status 1 + Command being timed: "python -" + User time (seconds): 360.17 + System time (seconds): 18.96 + Percent of CPU this job got: 663% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:57.11 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 12572016 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 8 + Minor (reclaiming a frame) page faults: 4351383 + Voluntary context switches: 75914 + Involuntary context switches: 1404254 + Swaps: 0 + File system inputs: 1392 + File system outputs: 17448 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 1 diff --git a/visual_results/sam3/multiple_prompt_validation.log b/visual_results/sam3/multiple_prompt_validation.log new file mode 100644 index 0000000..883ea9a --- /dev/null +++ b/visual_results/sam3/multiple_prompt_validation.log @@ -0,0 +1,49 @@ +2026-08-30 18:31:40.618050965 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 18:32:03.728663740 [E:onnxruntime:, sequential_executor.cc:671 ExecuteKernel] Non-zero status code returned while running Reshape node. Name:'/geometry_encoder/encode.0/self_attn/Reshape_4' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:91 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == requested_shape_size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{3,1,256}, requested shape:{2,8,32} + +encode_s 16.192 +plants_box shape [26, 1, 1481, 987] areas [0.016781, 0.003432, 0.01105, 0.000718, 0.008653, 0.091354, 0.006821, 0.006648, 0.002176, 0.00389, 0.005366, 0.007269, 0.024027, 0.004747, 0.020257, 0.019277, 0.008439, 0.011315, 0.007912, 0.001778, 0.034442, 0.016682, 0.017514, 0.00432, 0.002894, 0.011067] decode_s 2.775 +Traceback (most recent call last): + File "", line 12, in + File "/home/vietanhdev/Workspaces/samexporter/samexporter/sam3_onnx.py", line 159, in predict_masks + masks, scores, _boxes = self.decoder( + ~~~~~~~~~~~~^ + original_size, + ^^^^^^^^^^^^^^ + ...<11 lines>... + box_masks_np, + ^^^^^^^^^^^^^ + ) + ^ + File "/home/vietanhdev/Workspaces/samexporter/samexporter/sam3_onnx.py", line 352, in __call__ + outputs = self.session.run(None, model_inputs) + File "/home/vietanhdev/miniconda3/lib/python3.13/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 326, in run + return self._sess.run(output_names, input_feed, run_options) + ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Reshape node. Name:'/geometry_encoder/encode.0/self_attn/Reshape_4' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:91 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == requested_shape_size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{3,1,256}, requested shape:{2,8,32} + +Command exited with non-zero status 1 + Command being timed: "python -" + User time (seconds): 153.57 + System time (seconds): 7.07 + Percent of CPU this job got: 631% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:25.45 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8606456 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 5 + Minor (reclaiming a frame) page faults: 3021034 + Voluntary context switches: 3350 + Involuntary context switches: 646347 + Swaps: 0 + File system inputs: 5904 + File system outputs: 3136 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 1 diff --git a/visual_results/sam3/plants_box.log b/visual_results/sam3/plants_box.log new file mode 100644 index 0000000..1682a44 --- /dev/null +++ b/visual_results/sam3/plants_box.log @@ -0,0 +1,26 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:32:40.335235339 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 1 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/plants.png --prompt images/plants_box.json --output visual_results/sam3/plants_box.png --providers cpu" + User time (seconds): 152.05 + System time (seconds): 6.78 + Percent of CPU this job got: 598% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:26.52 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8695072 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3098596 + Voluntary context switches: 3452 + Involuntary context switches: 432940 + Swaps: 0 + File system inputs: 0 + File system outputs: 3256 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/plants_box.png b/visual_results/sam3/plants_box.png new file mode 100644 index 0000000..6f8101b Binary files /dev/null and b/visual_results/sam3/plants_box.png differ diff --git a/visual_results/sam3/plants_box_refined.log b/visual_results/sam3/plants_box_refined.log new file mode 100644 index 0000000..de2729d --- /dev/null +++ b/visual_results/sam3/plants_box_refined.log @@ -0,0 +1,26 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:33:40.108249095 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 1 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/plants.png --prompt images/plants_box_refined.json --output visual_results/sam3/plants_box_refined.png --providers cpu" + User time (seconds): 146.13 + System time (seconds): 7.31 + Percent of CPU this job got: 602% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:25.48 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8770200 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3177431 + Voluntary context switches: 3382 + Involuntary context switches: 599409 + Swaps: 0 + File system inputs: 0 + File system outputs: 4952 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/plants_box_refined.png b/visual_results/sam3/plants_box_refined.png new file mode 100644 index 0000000..edb5e4b Binary files /dev/null and b/visual_results/sam3/plants_box_refined.png differ diff --git a/visual_results/sam3/plants_text.log b/visual_results/sam3/plants_text.log new file mode 100644 index 0000000..5afb13f --- /dev/null +++ b/visual_results/sam3/plants_text.log @@ -0,0 +1,26 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:32:21.168426511 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 20 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/plants.png --prompt images/plants_text.json --output visual_results/sam3/plants_text.png --providers cpu" + User time (seconds): 106.98 + System time (seconds): 5.76 + Percent of CPU this job got: 597% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:18.87 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8812232 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3128166 + Voluntary context switches: 3674 + Involuntary context switches: 422637 + Swaps: 0 + File system inputs: 0 + File system outputs: 4432 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/plants_text.png b/visual_results/sam3/plants_text.png new file mode 100644 index 0000000..da0e5b5 Binary files /dev/null and b/visual_results/sam3/plants_text.png differ diff --git a/visual_results/sam3/plants_text_top5.log b/visual_results/sam3/plants_text_top5.log new file mode 100644 index 0000000..98e1f36 --- /dev/null +++ b/visual_results/sam3/plants_text_top5.log @@ -0,0 +1,26 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 21:34:04.506122554 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 5 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/plants.png --prompt images/plants_text.json --max_instances 5 --output visual_results/sam3/plants_text_top5.png --providers cpu" + User time (seconds): 192.97 + System time (seconds): 6.88 + Percent of CPU this job got: 628% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:31.81 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8589760 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3134895 + Voluntary context switches: 3436 + Involuntary context switches: 354438 + Swaps: 0 + File system inputs: 0 + File system outputs: 3472 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/plants_text_top5.png b/visual_results/sam3/plants_text_top5.png new file mode 100644 index 0000000..f255c1d Binary files /dev/null and b/visual_results/sam3/plants_text_top5.png differ diff --git a/visual_results/sam3/tokenizer_validation.log b/visual_results/sam3/tokenizer_validation.log new file mode 100644 index 0000000..438edbc --- /dev/null +++ b/visual_results/sam3/tokenizer_validation.log @@ -0,0 +1,26 @@ +libpng warning: iCCP: known incorrect sRGB profile +2026-08-30 22:12:08.890136120 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 5 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --image images/plants.png --prompt images/plants_text.json --text_prompt plant --max_instances 5 --output /tmp/sam3-tokenizer-real.png" + User time (seconds): 152.65 + System time (seconds): 6.39 + Percent of CPU this job got: 645% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:24.62 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8518624 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3126576 + Voluntary context switches: 3520 + Involuntary context switches: 699005 + Swaps: 0 + File system inputs: 0 + File system outputs: 13024 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/truck_box.log b/visual_results/sam3/truck_box.log new file mode 100644 index 0000000..d1e42f7 --- /dev/null +++ b/visual_results/sam3/truck_box.log @@ -0,0 +1,25 @@ +2026-08-30 21:27:49.134694851 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 1 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/truck.jpg --prompt images/truck_sam3_box.json --output visual_results/sam3/truck_box.png --providers cpu" + User time (seconds): 162.50 + System time (seconds): 6.66 + Percent of CPU this job got: 622% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:27.15 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8827780 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3179111 + Voluntary context switches: 3561 + Involuntary context switches: 402909 + Swaps: 0 + File system inputs: 0 + File system outputs: 5680 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/truck_box.png b/visual_results/sam3/truck_box.png new file mode 100644 index 0000000..9e3b6bc Binary files /dev/null and b/visual_results/sam3/truck_box.png differ diff --git a/visual_results/sam3/truck_point.log b/visual_results/sam3/truck_point.log new file mode 100644 index 0000000..adcacf8 --- /dev/null +++ b/visual_results/sam3/truck_point.log @@ -0,0 +1,25 @@ +2026-08-30 21:29:19.599597233 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 1 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/truck.jpg --prompt images/truck_sam3_point.json --text_prompt truck --output visual_results/sam3/truck_point.png --providers cpu" + User time (seconds): 134.12 + System time (seconds): 6.06 + Percent of CPU this job got: 623% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:22.46 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8889924 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 0 + Minor (reclaiming a frame) page faults: 3114933 + Voluntary context switches: 3445 + Involuntary context switches: 270814 + Swaps: 0 + File system inputs: 0 + File system outputs: 5792 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/truck_point.png b/visual_results/sam3/truck_point.png new file mode 100644 index 0000000..d83b5b1 Binary files /dev/null and b/visual_results/sam3/truck_point.png differ diff --git a/visual_results/sam3/truck_text.log b/visual_results/sam3/truck_text.log new file mode 100644 index 0000000..3608ec2 --- /dev/null +++ b/visual_results/sam3/truck_text.log @@ -0,0 +1,25 @@ +2026-08-30 21:27:27.239983726 [W:onnxruntime:, graph.cc:124 MergeShapeInfo] Error merging shape info for output. '/trunk/Concat_output_0' source:{4} target:{5}. Falling back to lenient merge. +SAM3 instances retained: 1 + Command being timed: "python -m samexporter.inference --sam_variant sam3 --encoder_model /tmp/samexporter_sam3/onnx/sam3_image_encoder.onnx --language_encoder_model /tmp/samexporter_sam3/onnx/sam3_language_encoder.onnx --decoder_model /tmp/samexporter_sam3/onnx/sam3_decoder.onnx --image images/truck.jpg --prompt images/truck_sam3.json --output visual_results/sam3/truck_text.png --providers cpu" + User time (seconds): 128.44 + System time (seconds): 6.02 + Percent of CPU this job got: 618% + Elapsed (wall clock) time (h:mm:ss or m:ss): 0:21.72 + Average shared text size (kbytes): 0 + Average unshared data size (kbytes): 0 + Average stack size (kbytes): 0 + Average total size (kbytes): 0 + Maximum resident set size (kbytes): 8843048 + Average resident set size (kbytes): 0 + Major (requiring I/O) page faults: 12 + Minor (reclaiming a frame) page faults: 3172872 + Voluntary context switches: 3509 + Involuntary context switches: 359504 + Swaps: 0 + File system inputs: 952 + File system outputs: 5808 + Socket messages sent: 0 + Socket messages received: 0 + Signals delivered: 0 + Page size (bytes): 4096 + Exit status: 0 diff --git a/visual_results/sam3/truck_text.png b/visual_results/sam3/truck_text.png new file mode 100644 index 0000000..208605c Binary files /dev/null and b/visual_results/sam3/truck_text.png differ diff --git a/visual_results/sam_vit_b/plants_box.log b/visual_results/sam_vit_b/plants_box.log new file mode 100644 index 0000000..0214ae2 --- /dev/null +++ b/visual_results/sam_vit_b/plants_box.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=3.68 max_rss_kb=3896888 diff --git a/visual_results/sam_vit_b/plants_box.png b/visual_results/sam_vit_b/plants_box.png new file mode 100644 index 0000000..92efc1e Binary files /dev/null and b/visual_results/sam_vit_b/plants_box.png differ diff --git a/visual_results/sam_vit_b/plants_box_refined.log b/visual_results/sam_vit_b/plants_box_refined.log new file mode 100644 index 0000000..b6f984c --- /dev/null +++ b/visual_results/sam_vit_b/plants_box_refined.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=5.85 max_rss_kb=3919092 diff --git a/visual_results/sam_vit_b/plants_box_refined.png b/visual_results/sam_vit_b/plants_box_refined.png new file mode 100644 index 0000000..8acf3bb Binary files /dev/null and b/visual_results/sam_vit_b/plants_box_refined.png differ diff --git a/visual_results/sam_vit_b/plants_point.log b/visual_results/sam_vit_b/plants_point.log new file mode 100644 index 0000000..6c68074 --- /dev/null +++ b/visual_results/sam_vit_b/plants_point.log @@ -0,0 +1,2 @@ +libpng warning: iCCP: known incorrect sRGB profile +elapsed_s=4.84 max_rss_kb=3878844 diff --git a/visual_results/sam_vit_b/plants_point.png b/visual_results/sam_vit_b/plants_point.png new file mode 100644 index 0000000..ed709f5 Binary files /dev/null and b/visual_results/sam_vit_b/plants_point.png differ diff --git a/visual_results/sam_vit_b/truck_box.log b/visual_results/sam_vit_b/truck_box.log new file mode 100644 index 0000000..78bb98c --- /dev/null +++ b/visual_results/sam_vit_b/truck_box.log @@ -0,0 +1 @@ +elapsed_s=4.50 max_rss_kb=3955100 diff --git a/visual_results/sam_vit_b/truck_box.png b/visual_results/sam_vit_b/truck_box.png new file mode 100644 index 0000000..cb3e21b Binary files /dev/null and b/visual_results/sam_vit_b/truck_box.png differ diff --git a/visual_results/sam_vit_b/truck_point.log b/visual_results/sam_vit_b/truck_point.log new file mode 100644 index 0000000..3bacdd6 --- /dev/null +++ b/visual_results/sam_vit_b/truck_point.log @@ -0,0 +1 @@ +elapsed_s=3.62 max_rss_kb=3922240 diff --git a/visual_results/sam_vit_b/truck_point.png b/visual_results/sam_vit_b/truck_point.png new file mode 100644 index 0000000..9645189 Binary files /dev/null and b/visual_results/sam_vit_b/truck_point.png differ