Skip to content

Add promptable YOLOE ONNX support#27

Draft
vlordier wants to merge 10 commits into
pan93412:masterfrom
vlordier:feat/yoloe-closed-set-support
Draft

Add promptable YOLOE ONNX support#27
vlordier wants to merge 10 commits into
pan93412:masterfrom
vlordier:feat/yoloe-closed-set-support

Conversation

@vlordier

@vlordier vlordier commented Apr 17, 2026

Copy link
Copy Markdown

Summary

This PR adds promptable YOLOE ONNX support while keeping the currently validated scope explicit after the docs cleanup.

Validated scope

  • Exported closed-set YOLOE ONNX support is supported when the matching labels are supplied at runtime.
  • The currently validated promptable path uses a custom two-input YOLOE detector plus a JSON bundle of raw text prompt features.
  • Rust-side ONNX prompt-encoder integration exists, but it is not yet covered by the committed parity fixture in this branch.
  • The ignored local parity test relies on the untracked exported-promptable-yoloe-26n/ bundle, so that reviewer fixture is still local-only.

What changed

  • generalize model input/output handling to read ONNX session metadata instead of hardcoded names
  • support caller-provided labels for exported closed-set YOLOE models
  • decode segmentation-style exports, including mask reconstruction from output0 coefficients plus output1 prototypes
  • add YOLO-style letterbox preprocessing, including Ultralytics-style rect mode for promptable parity work
  • add promptable two-input YOLOE inference with runtime prompt embeddings loaded from JSON raw text prompt features
  • add Rust-side prompt encoding with ONNX Runtime plus tokenizers
  • add CLI support for --label, --labels-file, --prompt, --prompt-encoder-model, --prompt-tokenizer, --rect, and --no-display
  • add a helper export script for promptable YOLOE detector plus prompt-feature assets
  • document both the closed-set export flow and the currently validated promptable runtime flow
  • add regression coverage for rect preprocessing, postprocessed decode handling, and promptable parity

Validation

  • cargo test --lib
  • cargo test promptable_rect_png_fixture_matches_python_baseline --test promptable_parity -- --ignored

Notes

  • Stock Ultralytics ONNX exports remain single-input closed-set graphs, so true runtime prompting still requires a custom promptable export.
  • The committed parity fixture validates the JSON-bundle promptable path; it does not yet cover the Rust ONNX prompt-encoder path.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request expands the library's capabilities to support YOLOE models and custom label configurations, enabling users to specify class labels through CLI arguments or external files. The core inference engine was updated to accommodate dynamic input/output names and flexible tensor layouts, specifically addressing models with additional detection channels. Review feedback highlights a logic issue where providing a subset of labels could lead to misclassification, identifies redundant error handling in the label retrieval process, and notes an unreachable error variant that should be removed for cleanliness.

Comment thread src/lib.rs Outdated
Comment thread src/lib.rs Outdated
Comment thread src/error.rs
@vlordier vlordier changed the title Add closed-set YOLOE ONNX support Add promptable YOLOE ONNX support Apr 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant