diff --git a/.gitignore b/.gitignore index b972a7f..5cd4521 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,25 @@ devenv.local.nix # pre-commit .pre-commit-config.yaml +# Editor settings +.vscode/ + +# Local Cargo config +.cargo/ + +# Python local environment +.venv/ +__pycache__/ +*.py[cod] + +# Local model weights +/yoloe-26n-seg.pt +/yoloe-v8s-seg.pt +/yoloe-*.onnx +/yoloe-*.onnx.data +/mobileclip*.pt +/mobileclip*.ts + # Added by cargo diff --git a/Cargo.lock b/Cargo.lock index af222d4..b7fbbd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -40,11 +40,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", + "getrandom 0.3.4", "once_cell", + "serde", "version_check", "zerocopy", ] +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "aligned" version = "0.4.3" @@ -246,7 +257,7 @@ dependencies = [ "anyhow", "arrayvec", "log", - "nom", + "nom 8.0.0", "num-rational", "v_frame", ] @@ -275,6 +286,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.22.1" @@ -404,6 +421,15 @@ dependencies = [ "vec_map", ] +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.2.56" @@ -475,7 +501,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3538270d33cc669650c4b093848450d380def10c331d38c768e34cac80576e6e" dependencies = [ "termcolor", - "unicode-width", + "unicode-width 0.1.14", ] [[package]] @@ -496,6 +522,33 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf43edc576402991846b093a7ca18a3477e0ef9c588cde84964b5d3e43016642" +[[package]] +name = "compact_str" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "console" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width 0.2.2", + "windows-sys 0.61.2", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -606,6 +659,50 @@ dependencies = [ "winapi", ] +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.117", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "der" version = "0.7.10" @@ -616,6 +713,37 @@ dependencies = [ "zeroize", ] +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.117", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -643,6 +771,12 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "equator" version = "0.4.2" @@ -679,6 +813,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "euclid" version = "0.22.13" @@ -775,6 +918,12 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.1.5" @@ -1107,6 +1256,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "image" version = "0.25.9" @@ -1169,6 +1324,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" +dependencies = [ + "console", + "portable-atomic", + "unicode-width 0.2.2", + "unit-prefix", + "web-time", +] + [[package]] name = "instant" version = "0.1.13" @@ -1368,6 +1536,22 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69" +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + [[package]] name = "malloc_buf" version = "0.0.6" @@ -1436,6 +1620,12 @@ dependencies = [ "paste", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -1458,6 +1648,28 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "moxcms" version = "0.7.11" @@ -1581,6 +1793,16 @@ dependencies = [ "memoffset", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nom" version = "8.0.0" @@ -1773,6 +1995,28 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "onig" +version = "6.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0" +dependencies = [ + "bitflags 2.11.0", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f86c6eef3d6df15f23bcfb6af487cbd2fed4e5581d58d5bf1f5f8b7f6727dc" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "openssl" version = "0.10.75" @@ -2188,6 +2432,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -2225,6 +2480,35 @@ dependencies = [ "bitflags 2.11.0", ] +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + [[package]] name = "renderdoc-sys" version = "1.1.0" @@ -2286,6 +2570,12 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "schannel" version = "0.1.28" @@ -2356,6 +2646,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -2510,6 +2801,18 @@ dependencies = [ "num-traits", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom 7.1.3", + "serde", + "unicode-segmentation", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -2672,6 +2975,40 @@ dependencies = [ "strict-num", ] +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "indicatif", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "toml_datetime" version = "0.6.11" @@ -2764,25 +3101,58 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" + [[package]] name = "unicode-width" version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode-xid" version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + [[package]] name = "ureq" version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdc97a28575b85cfedf2a7e7d3cc64b3e11bd8ac766666318003abbacc7a21fc" dependencies = [ - "base64", + "base64 0.22.1", "der", "log", "native-tls", @@ -2800,7 +3170,7 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d81f9efa9df032be5934a46a068815a10a042b494b6a58cb0a1a97bb5467ed6f" dependencies = [ - "base64", + "base64 0.22.1", "http", "httparse", "log", @@ -3051,6 +3421,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki-root-certs" version = "1.0.6" @@ -3528,6 +3908,7 @@ dependencies = [ "ndarray", "ort", "thiserror 2.0.18", + "tokenizers", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 5465913..3a67a82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ image = "0.25.9" ndarray = "0.17" ort = "2.0.0-rc.11" thiserror = "2.0.18" +tokenizers = "0.22.1" [workspace] members = [ diff --git a/README.md b/README.md index 936e736..2dd1f22 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,13 @@ # YOLO inference in Rust -This project provides a Rust implementation of the YOLO v11 object detection model, enabling inference on images to identify objects along with their bounding boxes, labels, and confidence scores. It utilizes the YOLO v11 model in ONNX format and leverages the `ort` library for ONNX Runtime integration. The implementation is inspired by the [YOLOv8 example](https://github.com/pykeio/ort/tree/main/examples/yolov8) from the `ort` repository. +This project provides a Rust implementation of YOLO-style ONNX inference, enabling detection and segmentation on images with bounding boxes, labels, confidence scores, and decoded masks. It works out of the box with YOLO v8/v11 COCO exports and can also run exported closed-set YOLOE ONNX models when you provide the matching labels. The implementation is inspired by the [YOLOv8 example](https://github.com/pykeio/ort/tree/main/examples/yolov8) from the `ort` repository. See docs.rs for the [latest documentation](https://docs.rs/yolo-rs). ## Features -- **Object Detection**: Detects objects within an image and provides their bounding boxes, labels, and confidence scores. -- **ONNX Model Integration**: Employs the YOLO v11 model in ONNX format for efficient inference. +- **Object Detection and Segmentation**: Detects objects within an image and provides their bounding boxes, labels, confidence scores, and segmentation masks when the model exports them. +- **YOLO-style ONNX Integration**: Supports default YOLO v8/v11 exports and exported closed-set YOLOE models. - **Rust Implementation**: Written entirely in Rust, ensuring performance and safety. - **ONNX Runtime**: Utilizes the `ort` library for executing the ONNX model. @@ -31,6 +31,74 @@ On a MacBook Pro (2024) with M3 Max, it tooks about **57ms** to inferring an ima - [**yolo-cli**](examples/yolo-cli): Command-line interface for running YOLO inference on images. +## YOLOE support + +This crate can run exported YOLOE ONNX models when they behave like regular fixed-class detectors. In practice that means exporting from Ultralytics after configuring the classes you want to detect, then passing the same labels into `yolo-rs`. + +Validated export flow in Python: + +```python +from ultralytics import YOLOE + +model = YOLOE("yoloe-v8s-seg.pt") +embeddings = model.get_text_pe(["person", "baseball bat", "baseball glove"]) +model.set_classes(["person", "baseball bat", "baseball glove"], embeddings) +model.export(format="onnx") +``` + +That export produces a standard segmentation-style ONNX graph with: + +- input `images` shaped `[1, 3, 640, 640]` +- output `output0` shaped `[1, 39, 8400]` +- output `output1` shaped `[1, 32, 160, 160]` + +`yolo-rs` follows the same broad exported ONNX contract as Ultralytics' YOLO segmentation examples: it letterboxes the input image, decodes boxes from `output0`, and reconstructs instance masks from the trailing mask coefficients in `output0` plus the prototype tensor in `output1`. + +Example CLI usage: + +```bash +cargo run --release -p example-yolo-gui -- yoloe-v8s-seg.onnx examples/yolo-cli/data/baseball.jpg --labels-file target/yoloe_inspect_v839/yoloe.labels.txt --no-display +``` + +Current scope: + +- Closed-set YOLOE detection exported to ONNX is supported. +- Detection heads with extra mask coefficients can be decoded for boxes and classes when you provide the label list. +- Segmentation-style exports can decode instance masks from `output1` and the mask coefficients stored in `output0`. +- Preprocessing and box rescaling now follow YOLO-style letterboxing instead of stretching the image to `640x640`. +- Runtime open-vocabulary prompting is available only for custom promptable YOLOE ONNX exports that expose a prompt embedding input. Ultralytics' stock ONNX exporter still emits single-input graphs, so those default exports must keep using `set_classes(...)` before export. + +## True Runtime Prompting + +`yolo-rs` now also supports promptable YOLOE ONNX graphs with a second model input for prompt embeddings. + +That path is different from Ultralytics' stock `model.export(format="onnx")`, which still emits a single-input ONNX and bakes prompts in ahead of time. For true runtime open-vocabulary prompting you need three artifacts: + +- a promptable YOLOE detector ONNX with inputs `images` and `prompt_embeddings` +- a prompt-encoder ONNX that converts tokenized text into the detector's prompt embedding space +- a tokenizer JSON compatible with that prompt encoder + +This repository includes a helper export script at [scripts/export_promptable_yoloe.py](/Users/vincent/Work/yolo-rs/scripts/export_promptable_yoloe.py) to generate those assets from upstream YOLOE weights. + +Example CLI usage once those assets exist: + +```bash +cargo run --release -p example-yolo-gui -- \ + exported-promptable-yoloe/yoloe-promptable.onnx \ + examples/yolo-cli/data/baseball.jpg \ + --prompt person \ + --prompt "baseball bat" \ + --prompt "baseball glove" \ + --prompt-encoder-model exported-promptable-yoloe/yoloe-prompt-encoder.onnx \ + --prompt-tokenizer exported-promptable-yoloe/tokenizer.json +``` + +Current runtime-prompt scope: + +- `yolo-rs` can run promptable two-input YOLOE ONNX graphs with runtime text prompts. +- The prompt encoder is a separate ONNX model; prompt strings are tokenized in Rust and encoded at runtime. +- Stock Ultralytics ONNX exports remain single-input closed-set graphs, so they still require `set_classes(...)` before export. + ## Acknowledgements This project is inspired by the [YOLOv8 example](https://github.com/pykeio/ort/tree/main/examples/yolov8) from the `ort` repository. diff --git a/examples/yolo-cli/README.md b/examples/yolo-cli/README.md index a953e25..fde0a2d 100644 --- a/examples/yolo-cli/README.md +++ b/examples/yolo-cli/README.md @@ -14,6 +14,40 @@ where: - You can export the model according to [Ultralytics' manual](https://docs.ultralytics.com/integrations/onnx/). - `data/baseball.jpg` is the path to the image file. +For exported closed-set YOLOE models, pass the matching labels when you run the CLI: + +```bash +cargo run --release exported-yoloe.onnx image.jpg --label person --label bus +``` + +For non-interactive validation, skip the GUI window: + +```bash +cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt --no-display +``` + +For segmentation-style exports, the CLI will also decode and draw instance masks when the model exposes a second ONNX output with mask prototypes. + +Open-vocabulary runtime prompting is available only for custom promptable YOLOE ONNX exports. Ultralytics' default ONNX export still produces a single-input fixed-class graph. + +For true runtime prompting with a promptable two-input YOLOE ONNX graph, pass prompts together with a prompt encoder and tokenizer: + +```bash +cargo run --release promptable-yoloe.onnx image.jpg \ + --prompt person \ + --prompt bus \ + --prompt-encoder-model yoloe-prompt-encoder.onnx \ + --prompt-tokenizer tokenizer.json +``` + +This path expects a custom promptable export rather than Ultralytics' default single-input ONNX export. See [scripts/export_promptable_yoloe.py](/Users/vincent/Work/yolo-rs/scripts/export_promptable_yoloe.py) for a helper script. + +You can also load labels from a file: + +```bash +cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt +``` + ## Downloading models ```shell diff --git a/examples/yolo-cli/src/main.rs b/examples/yolo-cli/src/main.rs index 87c0519..6028574 100644 --- a/examples/yolo-cli/src/main.rs +++ b/examples/yolo-cli/src/main.rs @@ -1,3 +1,4 @@ +use std::fs; use std::path::PathBuf; use anyhow::{Context, Result}; @@ -5,7 +6,11 @@ use clap::Parser; use ort::execution_providers::{CUDAExecutionProvider, CoreMLExecutionProvider}; use raqote::{DrawOptions, DrawTarget, LineJoin, PathBuilder, SolidSource, Source, StrokeStyle}; use show_image::{AsImageView, WindowOptions, event}; -use yolo_rs::{YoloEntityOutput, image_to_yolo_input_tensor, inference, model}; +use yolo_rs::{ + YoloEntityOutput, YoloSegmentationOutput, image_to_yolo_input_tensor, inference, + inference_segment, inference_segment_with_prompts, inference_with_prompts, model, + prompt::YoloPromptEncoderSession, +}; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -13,11 +18,60 @@ struct Args { model_path: PathBuf, picture_path: PathBuf, + /// Override the class labels used to decode the model output. + /// + /// Pass this multiple times for exported YOLOE closed-set models, for example: + /// --label person --label bus + #[arg(long = "label")] + labels: Vec, + + /// Runtime open-vocabulary prompts for a promptable YOLOE ONNX model. + #[arg(long = "prompt")] + prompts: Vec, + + /// Path to a prompt-encoder ONNX model that converts tokenized text to prompt embeddings. + #[arg(long)] + prompt_encoder_model: Option, + + /// Path to a tokenizer JSON compatible with the prompt encoder. + #[arg(long)] + prompt_tokenizer: Option, + + /// Read class labels from a UTF-8 text file, one label per line. + #[arg(long)] + labels_file: Option, + #[arg(long)] probability_threshold: Option, #[arg(long)] iou_threshold: Option, + + /// Skip opening the GUI window and only print detections. + #[arg(long)] + no_display: bool, +} + +fn read_labels(args: &Args) -> Result>> { + let mut labels = if let Some(path) = &args.labels_file { + fs::read_to_string(path) + .with_context(|| format!("failed to read labels file {:?}", path.display()))? + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(ToOwned::to_owned) + .collect::>() + } else { + Vec::new() + }; + + labels.extend(args.labels.iter().cloned()); + + if labels.is_empty() { + Ok(None) + } else { + Ok(Some(labels)) + } } #[show_image::main] @@ -41,9 +95,20 @@ fn main() -> Result<()> { .with_context(|| format!("failed to open image {:?}", args.picture_path.display()))?; tracing::info!("Loading models {:?}…", args.model_path.display()); + let labels = read_labels(&args)?; + let prompt_mode = !args.prompts.is_empty(); let mut model = { - let mut model = model::YoloModelSession::from_filename_v8(&args.model_path) - .with_context(|| format!("failed to load model {:?}", args.model_path.display()))?; + let mut model = if prompt_mode { + model::YoloModelSession::from_filename(&args.model_path) + } else if let Some(labels) = labels { + model::YoloModelSession::from_filename_with_labels( + &args.model_path, + labels.into_iter(), + ) + } else { + model::YoloModelSession::from_filename_v8(&args.model_path) + } + .with_context(|| format!("failed to load model {:?}", args.model_path.display()))?; model.iou_threshold = args.iou_threshold; model.probability_threshold = args.probability_threshold; @@ -54,17 +119,80 @@ fn main() -> Result<()> { tracing::debug!("Converting image to tensor…"); let input = image_to_yolo_input_tensor(&original_img); + let prompt_embeddings = if prompt_mode { + let prompt_encoder_model = args + .prompt_encoder_model + .as_ref() + .context("--prompt-encoder-model is required when using --prompt")?; + let prompt_tokenizer = args + .prompt_tokenizer + .as_ref() + .context("--prompt-tokenizer is required when using --prompt")?; + tracing::info!("Encoding runtime prompts…"); + let mut prompt_encoder = YoloPromptEncoderSession::from_files( + prompt_encoder_model, + prompt_tokenizer, + ) + .with_context(|| { + format!( + "failed to load prompt encoder {:?} and tokenizer {:?}", + prompt_encoder_model.display(), + prompt_tokenizer.display() + ) + })?; + Some(prompt_encoder.encode(args.prompts.iter().map(String::as_str))?) + } else { + None + }; + // Run YOLOv8 inference tracing::info!("Running inference…"); let now = std::time::Instant::now(); - let result = inference(&mut model, input.view())?; + let has_mask_output = model.session.outputs().len() > 1; + let result = if has_mask_output { + let (img_width, img_height) = (original_img.width(), original_img.height()); + let mut dt = DrawTarget::new(img_width as _, img_height as _); + let result = if let Some(prompt_embeddings) = prompt_embeddings.as_ref() { + inference_segment_with_prompts(&mut model, input.view(), prompt_embeddings.view())? + } else { + inference_segment(&mut model, input.view())? + }; + let data = dt.get_data_mut(); + for YoloSegmentationOutput { entity, mask } in &result { + let fill = match entity.label.as_str() { + "baseball bat" => 0x40001080u32, + "baseball glove" => 0x40208040u32, + _ => 0x40801040u32, + }; + + for (x, y, pixel) in mask.enumerate_pixels() { + if pixel.0[0] > 0 { + let index = (y as usize) * (img_width as usize) + (x as usize); + data[index] = fill; + } + } + } + + (result.into_iter().map(|output| output.entity).collect::>(), dt) + } else { + let (img_width, img_height) = (original_img.width(), original_img.height()); + let dt = DrawTarget::new(img_width as _, img_height as _); + ( + if let Some(prompt_embeddings) = prompt_embeddings.as_ref() { + inference_with_prompts(&mut model, input.view(), prompt_embeddings.view())? + } else { + inference(&mut model, input.view())? + }, + dt, + ) + }; + tracing::info!("Inference took {:?}", now.elapsed()); tracing::debug!("Drawing bounding boxes…"); let (img_width, img_height) = (original_img.width(), original_img.height()); - - let mut dt = DrawTarget::new(img_width as _, img_height as _); + let (result, mut dt) = result; for YoloEntityOutput { bounding_box: bbox, @@ -117,6 +245,10 @@ fn main() -> Result<()> { ); } + if args.no_display { + return Ok(()); + } + let overlay: show_image::Image = dt.into(); tracing::info!("Displaying image…"); diff --git a/scripts/export_promptable_yoloe.py b/scripts/export_promptable_yoloe.py new file mode 100644 index 0000000..dd1aaa0 --- /dev/null +++ b/scripts/export_promptable_yoloe.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from ultralytics import YOLOE +from ultralytics.nn.modules.head import Detect +from ultralytics.nn.text_model import build_text_model + + +class PromptableDetector(torch.nn.Module): + def __init__(self, checkpoint: str) -> None: + super().__init__() + wrapper = YOLOE(checkpoint) + self.model = wrapper.model + self.model.eval().float() + + for module in self.model.modules(): + if isinstance(module, Detect): + module.dynamic = False + module.export = True + module.format = "onnx" + module.max_det = 300 + + def forward(self, images: torch.Tensor, prompt_embeddings: torch.Tensor): + return self.model.predict(images, tpe=prompt_embeddings) + + +class PromptEncoder(torch.nn.Module): + def __init__(self, checkpoint: str) -> None: + super().__init__() + wrapper = YOLOE(checkpoint) + self.model = wrapper.model + self.model.eval().float() + self.text_model = build_text_model(self.model.args.get("text_model"), device=torch.device("cpu")) + self.text_model.eval() + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + text_features = self.text_model.encode_text(input_ids, dtype=torch.float32) + text_features = text_features.reshape(1, input_ids.shape[0], text_features.shape[-1]) + head = self.model.model[-1] + return head.get_tpe(text_features) + + +def save_tokenizer(encoder: PromptEncoder, output_path: Path) -> None: + tokenizer = encoder.text_model.tokenizer + output_path.parent.mkdir(parents=True, exist_ok=True) + + if hasattr(tokenizer, "save"): + tokenizer.save(str(output_path)) + return + + if hasattr(tokenizer, "save_pretrained"): + tokenizer.save_pretrained(str(output_path.parent)) + return + + raise RuntimeError("Tokenizer does not support save() or save_pretrained(); export it manually.") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", help="Path to a YOLOE checkpoint such as yoloe-v8s-seg.pt") + parser.add_argument("--output-dir", type=Path, default=Path("exported-promptable-yoloe")) + parser.add_argument("--opset", type=int, default=17) + parser.add_argument("--height", type=int, default=640) + parser.add_argument("--width", type=int, default=640) + args = parser.parse_args() + + args.output_dir.mkdir(parents=True, exist_ok=True) + + detector = PromptableDetector(args.checkpoint) + prompt_encoder = PromptEncoder(args.checkpoint) + + sample_image = torch.zeros((1, 3, args.height, args.width), dtype=torch.float32) + sample_tokens = prompt_encoder.text_model.tokenize(["object"]).to(torch.long) + sample_prompt_embeddings = prompt_encoder(sample_tokens) + + detector_path = args.output_dir / "yoloe-promptable.onnx" + prompt_encoder_path = args.output_dir / "yoloe-prompt-encoder.onnx" + tokenizer_path = args.output_dir / "tokenizer.json" + + detector_output_names = ["output0", "output1"] if "seg" in Path(args.checkpoint).stem else ["output0"] + + torch.onnx.export( + detector, + (sample_image, sample_prompt_embeddings), + detector_path, + opset_version=args.opset, + input_names=["images", "prompt_embeddings"], + output_names=detector_output_names, + dynamic_axes={ + "images": {0: "batch"}, + "prompt_embeddings": {1: "prompt_count"}, + "output0": {0: "batch", 2: "anchors"}, + **({"output1": {0: "batch", 2: "mask_height", 3: "mask_width"}} if len(detector_output_names) > 1 else {}), + }, + ) + + torch.onnx.export( + prompt_encoder, + sample_tokens, + prompt_encoder_path, + opset_version=args.opset, + input_names=["input_ids"], + output_names=["prompt_embeddings"], + dynamic_axes={ + "input_ids": {0: "prompt_count", 1: "token_count"}, + "prompt_embeddings": {1: "prompt_count"}, + }, + ) + + save_tokenizer(prompt_encoder, tokenizer_path) + + print(f"Wrote detector ONNX to {detector_path}") + print(f"Wrote prompt encoder ONNX to {prompt_encoder_path}") + print(f"Wrote tokenizer to {tokenizer_path}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/error.rs b/src/error.rs index c8fac32..fd35fe8 100644 --- a/src/error.rs +++ b/src/error.rs @@ -10,6 +10,50 @@ pub enum YoloError { OrtInputError(ort::Error), #[error("run inference: {0}")] OrtInferenceError(ort::Error), - #[error("extract sensor: {0}")] - OrtExtractSensorError(ort::Error), + #[error("extract tensor: {0}")] + OrtExtractTensorError(ort::Error), + #[error("model has no inputs")] + MissingModelInput, + #[error("model has no outputs")] + MissingModelOutput, + #[error("model has no prompt embedding input")] + MissingPromptInput, + #[error("model has no segmentation mask output")] + MissingMaskOutput, + #[error("unsupported model output shape {0:?}; expected a 3D detection tensor")] + InvalidOutputShape(Vec), + #[error("unsupported mask output shape {0:?}; expected a 4D mask prototype tensor")] + InvalidMaskShape(Vec), + #[error("unsupported prompt embedding shape {0:?}; expected [1, prompts, embedding_dim]")] + InvalidPromptEmbeddingShape(Vec), + #[error("unsupported prompt encoder output shape {0:?}; expected [prompts, embedding_dim] or [1, prompts, embedding_dim]")] + InvalidPromptEncoderOutputShape(Vec), + #[error("prompt encoder model exposes {0} inputs, but only 1 to 3 text-encoder inputs are supported")] + UnsupportedPromptEncoderInputCount(usize), + #[error("prompt list is empty")] + EmptyPromptList, + #[error("prompt label count {labels} does not match prompt embedding count {embeddings}")] + PromptLabelEmbeddingMismatch { labels: usize, embeddings: usize }, + #[error("load tokenizer: {0}")] + TokenizerLoadError(String), + #[error("tokenize prompts: {0}")] + PromptTokenizationError(String), + #[error( + "model output has {available} detection channels, but {required} are required for 4 box coordinates plus {label_count} labels" + )] + InsufficientDetectionChannels { + available: usize, + required: usize, + label_count: usize, + }, + #[error( + "model output has {available} mask coefficients, but {required} are required for {prototype_count} mask prototypes" + )] + InsufficientMaskCoefficients { + available: usize, + required: usize, + prototype_count: usize, + }, + #[error("detected class index {class_id} is out of range for {label_count} labels")] + UnknownClassId { class_id: usize, label_count: usize }, } diff --git a/src/lib.rs b/src/lib.rs index 1d3bd6a..763c26b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,17 +1,21 @@ -//! A Rust library for the YOLO v11 object detection model. +//! A Rust library for YOLO-style ONNX detection and segmentation models. //! -//! This library provides a high-level API for running the YOLO v11 object detection model. -//! Currently, it supports only the inference. +//! This library provides a high-level API for running YOLO-style ONNX models, +//! including exported closed-set YOLOE segmentation models. pub mod error; pub mod model; +pub mod prompt; use arcstr::ArcStr; use error::YoloError; -use image::{DynamicImage, GenericImageView, Rgba, imageops::FilterType}; +use image::{DynamicImage, GenericImageView, GrayImage, Luma, Rgba, imageops::FilterType}; use model::YoloModelSession; -use ndarray::{Array4, ArrayBase, ArrayView4, Axis, s}; +use ndarray::{Array4, ArrayBase, ArrayView3, ArrayView4}; use ort::{inputs, value::TensorRef}; +use prompt::YoloPromptEmbeddingsView; + +const DEFAULT_MODEL_SIZE: u32 = 640; #[derive(Debug, Clone, Copy)] pub struct BoundingBox { @@ -26,6 +30,11 @@ pub struct YoloInput { pub tensor: Array4, // 640x640 pub raw_width: u32, pub raw_height: u32, + pub model_width: u32, + pub model_height: u32, + pub resize_ratio: f32, + pub pad_w: f32, + pub pad_h: f32, } impl YoloInput { @@ -34,6 +43,11 @@ impl YoloInput { tensor_view: self.tensor.view(), raw_width: self.raw_width, raw_height: self.raw_height, + model_width: self.model_width, + model_height: self.model_height, + resize_ratio: self.resize_ratio, + pad_w: self.pad_w, + pad_h: self.pad_h, } } } @@ -43,6 +57,11 @@ pub struct YoloInputView<'a> { pub tensor_view: ArrayView4<'a, f32>, pub raw_width: u32, pub raw_height: u32, + pub model_width: u32, + pub model_height: u32, + pub resize_ratio: f32, + pub pad_w: f32, + pub pad_h: f32, } #[derive(Debug, Clone)] @@ -57,6 +76,257 @@ pub struct YoloEntityOutput { pub confidence: f32, } +#[derive(Debug, Clone)] +pub struct YoloSegmentationOutput { + pub entity: YoloEntityOutput, + pub mask: GrayImage, +} + +fn validate_prompt_embeddings( + prompts: YoloPromptEmbeddingsView<'_>, +) -> Result, YoloError> { + if prompts.embeddings.shape().len() != 3 || prompts.embeddings.shape()[0] != 1 { + return Err(YoloError::InvalidPromptEmbeddingShape( + prompts.embeddings.shape().to_vec(), + )); + } + if prompts.labels.len() != prompts.embeddings.shape()[1] { + return Err(YoloError::PromptLabelEmbeddingMismatch { + labels: prompts.labels.len(), + embeddings: prompts.embeddings.shape()[1], + }); + } + + Ok(prompts) +} + +#[derive(Debug, Clone)] +struct DecodedDetection { + entity: YoloEntityOutput, + mask_coefficients: Vec, +} + +#[derive(Debug, Clone, Copy)] +struct LetterboxInfo { + raw_width: u32, + raw_height: u32, + model_width: u32, + model_height: u32, + resize_ratio: f32, + pad_w: f32, + pad_h: f32, +} + +fn output_layout(output: ArrayView3<'_, f32>, label_count: usize) -> Result<(usize, usize), YoloError> { + if output.shape()[0] != 1 { + return Err(YoloError::InvalidOutputShape(output.shape().to_vec())); + } + + let channels_axis = if output.shape()[1] <= output.shape()[2] { 1 } else { 2 }; + let detection_axis = if channels_axis == 1 { 2 } else { 1 }; + + let channel_count = output.shape()[channels_axis]; + let required_channels = 4 + label_count; + + if channel_count < required_channels { + return Err(YoloError::InsufficientDetectionChannels { + available: channel_count, + required: required_channels, + label_count, + }); + } + + Ok((channels_axis, output.shape()[detection_axis])) +} + +fn intersection(box1: &BoundingBox, box2: &BoundingBox) -> f32 { + (box1.x2.min(box2.x2) - box1.x1.max(box2.x1)) * (box1.y2.min(box2.y2) - box1.y1.max(box2.y1)) +} + +fn union(box1: &BoundingBox, box2: &BoundingBox) -> f32 { + ((box1.x2 - box1.x1) * (box1.y2 - box1.y1)) + ((box2.x2 - box2.x1) * (box2.y2 - box2.y1)) + - intersection(box1, box2) +} + +fn non_maximum_suppression( + mut boxes: Vec, + iou_threshold: f32, +) -> Vec { + if boxes.is_empty() { + return Vec::new(); + } + + boxes.sort_unstable_by(|a, b| b.entity.confidence.total_cmp(&a.entity.confidence)); + + let mut result = Vec::with_capacity(boxes.len()); + for current in boxes { + if result.iter().all(|selected: &DecodedDetection| { + let iou = intersection(&selected.entity.bounding_box, ¤t.entity.bounding_box) + / union(&selected.entity.bounding_box, ¤t.entity.bounding_box); + iou < iou_threshold + }) { + result.push(current); + } + } + + result.shrink_to_fit(); + result +} + +fn decode_detections( + output: ArrayView3<'_, f32>, + labels: &[ArcStr], + probability_threshold: f32, + letterbox: LetterboxInfo, + expected_mask_coefficients: Option, +) -> Result, YoloError> { + let (channels_axis, detection_count) = output_layout(output, labels.len())?; + + Ok((0..detection_count) + .filter_map(|row| { + let row = if channels_axis == 1 { + output.slice(ndarray::s![0, .., row]) + } else { + output.slice(ndarray::s![0, row, ..]) + }; + + let (class_id, prob) = row + .iter() + .skip(4) + .take(labels.len()) + .enumerate() + .map(|(index, value)| (index, *value)) + .reduce(|accum, row| if row.1 > accum.1 { row } else { accum }) + .filter(|(_, prob)| *prob >= probability_threshold)?; + + let label = labels + .get(class_id) + .cloned() + .ok_or(YoloError::UnknownClassId { + class_id, + label_count: labels.len(), + }) + .ok()?; + + let mask_coefficients = row + .iter() + .skip(4 + labels.len()) + .copied() + .collect::>(); + + if let Some(required) = expected_mask_coefficients + && mask_coefficients.len() < required + { + return Some(Err(YoloError::InsufficientMaskCoefficients { + available: mask_coefficients.len(), + required, + prototype_count: required, + })); + } + + let x1 = ((row[0_usize] - row[2_usize] / 2.0) - letterbox.pad_w) / letterbox.resize_ratio; + let y1 = ((row[1_usize] - row[3_usize] / 2.0) - letterbox.pad_h) / letterbox.resize_ratio; + let x2 = ((row[0_usize] + row[2_usize] / 2.0) - letterbox.pad_w) / letterbox.resize_ratio; + let y2 = ((row[1_usize] + row[3_usize] / 2.0) - letterbox.pad_h) / letterbox.resize_ratio; + + Some(Ok(DecodedDetection { + entity: YoloEntityOutput { + bounding_box: BoundingBox { + x1: x1.clamp(0.0, letterbox.raw_width as f32), + y1: y1.clamp(0.0, letterbox.raw_height as f32), + x2: x2.clamp(0.0, letterbox.raw_width as f32), + y2: y2.clamp(0.0, letterbox.raw_height as f32), + }, + label, + confidence: prob, + }, + mask_coefficients, + })) + }) + .collect::, _>>()?) +} + +fn decode_segmentation_masks( + detections: Vec, + prototypes: ArrayView3<'_, f32>, + letterbox: LetterboxInfo, +) -> Result, YoloError> { + let prototype_count = prototypes.shape()[0]; + let mask_height = prototypes.shape()[1]; + let mask_width = prototypes.shape()[2]; + + let mask_pad_w = letterbox.pad_w * (mask_width as f32 / letterbox.model_width as f32); + let mask_pad_h = letterbox.pad_h * (mask_height as f32 / letterbox.model_height as f32); + let mask_content_width = (mask_width as f32 - (mask_pad_w * 2.0)).max(1.0); + let mask_content_height = (mask_height as f32 - (mask_pad_h * 2.0)).max(1.0); + + let mut results = Vec::with_capacity(detections.len()); + for detection in detections { + let mut mask = GrayImage::new(letterbox.raw_width, letterbox.raw_height); + let bbox = detection.entity.bounding_box; + + let start_x = bbox.x1.floor().max(0.0) as u32; + let end_x = bbox.x2.ceil().min(letterbox.raw_width as f32) as u32; + let start_y = bbox.y1.floor().max(0.0) as u32; + let end_y = bbox.y2.ceil().min(letterbox.raw_height as f32) as u32; + + let mut proto_mask = vec![0.0f32; mask_width * mask_height]; + for proto_y in 0..mask_height { + for proto_x in 0..mask_width { + let value = detection + .mask_coefficients + .iter() + .take(prototype_count) + .enumerate() + .map(|(index, coefficient)| coefficient * prototypes[[index, proto_y, proto_x]]) + .sum::(); + proto_mask[proto_y * mask_width + proto_x] = value; + } + } + + for y in start_y..end_y { + for x in start_x..end_x { + let scaled_x = ((x as f32 + 0.5) / letterbox.raw_width as f32) * mask_content_width + mask_pad_w; + let scaled_y = ((y as f32 + 0.5) / letterbox.raw_height as f32) * mask_content_height + mask_pad_h; + let value = bilinear_sample(&proto_mask, mask_width, mask_height, scaled_x, scaled_y); + + if value > 0.5 { + mask.put_pixel(x, y, Luma([255])); + } + } + } + + results.push(YoloSegmentationOutput { + entity: detection.entity, + mask, + }); + } + + Ok(results) +} + +fn bilinear_sample(buffer: &[f32], width: usize, height: usize, x: f32, y: f32) -> f32 { + let x = x.clamp(0.0, (width.saturating_sub(1)) as f32); + let y = y.clamp(0.0, (height.saturating_sub(1)) as f32); + + let x0 = x.floor() as usize; + let y0 = y.floor() as usize; + let x1 = (x0 + 1).min(width.saturating_sub(1)); + let y1 = (y0 + 1).min(height.saturating_sub(1)); + + let dx = x - x0 as f32; + let dy = y - y0 as f32; + + let top_left = buffer[y0 * width + x0]; + let top_right = buffer[y0 * width + x1]; + let bottom_left = buffer[y1 * width + x0]; + let bottom_right = buffer[y1 * width + x1]; + + let top = top_left + dx * (top_right - top_left); + let bottom = bottom_left + dx * (bottom_right - bottom_left); + top + dy * (bottom - top) +} + /// Convert an image to a YOLO input tensor. /// /// The input image is resized to 640x640 and normalized to the range [0, 1]. @@ -65,12 +335,28 @@ pub struct YoloEntityOutput { /// You can pass the resulting tensor to the [`inference`] function. /// Note that you might need to call [`YoloInput::view`] to get a view of the tensor. pub fn image_to_yolo_input_tensor(original_image: &DynamicImage) -> YoloInput { - let mut input = ArrayBase::zeros((1, 3, 640, 640)); + let model_width = DEFAULT_MODEL_SIZE; + let model_height = DEFAULT_MODEL_SIZE; + let mut input = ArrayBase::from_elem( + (1, 3, model_height as usize, model_width as usize), + 114.0f32 / 255.0, + ); + + let raw_width = original_image.width(); + let raw_height = original_image.height(); + let resize_ratio = (model_width as f32 / raw_width as f32) + .min(model_height as f32 / raw_height as f32); + let resized_width = ((raw_width as f32) * resize_ratio).round() as u32; + let resized_height = ((raw_height as f32) * resize_ratio).round() as u32; + let pad_w = (model_width as f32 - resized_width as f32) / 2.0; + let pad_h = (model_height as f32 - resized_height as f32) / 2.0; + let left = (pad_w - 0.1).round().max(0.0) as usize; + let top = (pad_h - 0.1).round().max(0.0) as usize; - let image = original_image.resize_exact(640, 640, FilterType::CatmullRom); + let image = original_image.resize_exact(resized_width, resized_height, FilterType::CatmullRom); for (x, y, Rgba([r, g, b, _])) in image.pixels() { - let x = x as usize; - let y = y as usize; + let x = left + x as usize; + let y = top + y as usize; input[[0, 0, y, x]] = (r as f32) / 255.; input[[0, 1, y, x]] = (g as f32) / 255.; @@ -79,115 +365,283 @@ pub fn image_to_yolo_input_tensor(original_image: &DynamicImage) -> YoloInput { YoloInput { tensor: input, - raw_width: original_image.width(), - raw_height: original_image.height(), + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, } } /// Inference on the YOLO model, returning the detected entities. /// /// The input tensor should be obtained from the [`image_to_yolo_input_tensor`] function. -/// The [`YoloModelSession`] can be obtained from the [`YoloModelSession::from_filename_v8`] method. +/// The [`YoloModelSession`] can be obtained from the [`YoloModelSession::from_filename_v8`] +/// or [`YoloModelSession::from_filename_with_labels`] methods. pub fn inference( model: &mut YoloModelSession, YoloInputView { tensor_view, raw_width, raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, }: YoloInputView, ) -> Result, YoloError> { - fn intersection(box1: &BoundingBox, box2: &BoundingBox) -> f32 { - (box1.x2.min(box2.x2) - box1.x1.max(box2.x1)) - * (box1.y2.min(box2.y2) - box1.y1.max(box2.y1)) - } + // Due to the lifetime of the model, we need to clone the + // labels and thresholds early. + let iou_threshold = model.get_iou_threshold(); + let probability_threshold = model.get_probability_threshold(); + let labels = model.get_labels().to_vec(); + let input_name = model.get_input_name()?.to_owned(); + let output_name = model.get_output_name()?.to_owned(); - fn union(box1: &BoundingBox, box2: &BoundingBox) -> f32 { - ((box1.x2 - box1.x1) * (box1.y2 - box1.y1)) + ((box2.x2 - box2.x1) * (box2.y2 - box2.y1)) - - intersection(box1, box2) - } + // Run YOLO-style inference using the session's configured input/output names. + let inputs = inputs![input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?]; + let outputs = model + .as_mut() + .run(inputs) + .map_err(YoloError::OrtInferenceError)?; + let output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let output = output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(output.shape().to_vec()))?; + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; + let boxes = decode_detections( + output, + &labels, + probability_threshold, + letterbox, + None, + )?; - fn non_maximum_suppression( - mut boxes: Vec, - iou_threshold: f32, - ) -> Vec { - // Early return if no boxes are provided - if boxes.is_empty() { - return Vec::new(); - } + Ok(non_maximum_suppression(boxes, iou_threshold) + .into_iter() + .map(|decoded| decoded.entity) + .collect()) +} - // Sort boxes by confidence descending using sort_unstable_by for better performance - boxes.sort_unstable_by(|a, b| b.confidence.total_cmp(&a.confidence)); - - let mut result = Vec::with_capacity(boxes.len()); - - // Iterate through each box and select it if it doesn't overlap significantly with already selected boxes - for current in boxes.into_iter() { - // Check if the current box has a high IoU with any box in the result - // Using `iter().all()` ensures we short-circuit on the first overlap found - if result.iter().all(|selected: &YoloEntityOutput| { - let iou = intersection(&selected.bounding_box, ¤t.bounding_box) - / union(&selected.bounding_box, ¤t.bounding_box); - iou < iou_threshold - }) { - result.push(current); - } - } +pub fn inference_with_prompts( + model: &mut YoloModelSession, + YoloInputView { + tensor_view, + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }: YoloInputView, + prompts: YoloPromptEmbeddingsView<'_>, +) -> Result, YoloError> { + let prompts = validate_prompt_embeddings(prompts)?; + let iou_threshold = model.get_iou_threshold(); + let probability_threshold = model.get_probability_threshold(); + let input_name = model.get_input_name()?.to_owned(); + let prompt_input_name = model.get_prompt_input_name()?.to_owned(); + let output_name = model.get_output_name()?.to_owned(); - result.shrink_to_fit(); + let inputs = inputs![ + input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?, + prompt_input_name.as_str() => TensorRef::from_array_view(prompts.embeddings).map_err(YoloError::OrtInputError)?, + ]; + let outputs = model + .as_mut() + .run(inputs) + .map_err(YoloError::OrtInferenceError)?; + let output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let output = output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(output.shape().to_vec()))?; + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; + let boxes = decode_detections( + output, + prompts.labels, + probability_threshold, + letterbox, + None, + )?; - result - } + Ok(non_maximum_suppression(boxes, iou_threshold) + .into_iter() + .map(|decoded| decoded.entity) + .collect()) +} - // Due to the lifetime of the model, we need to clone the - // labels and thresholds early. +/// Inference on a segmentation-style YOLO model, returning the detected entities and decoded masks. +/// +/// This is intended for YOLO segmentation exports, including exported closed-set YOLOE segmentation models. +pub fn inference_segment( + model: &mut YoloModelSession, + YoloInputView { + tensor_view, + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }: YoloInputView, +) -> Result, YoloError> { let iou_threshold = model.get_iou_threshold(); let probability_threshold = model.get_probability_threshold(); let labels = model.get_labels().to_vec(); + let input_name = model.get_input_name()?.to_owned(); + let output_name = model.get_output_name()?.to_owned(); + let mask_output_name = model.get_mask_output_name()?.to_owned(); - // Run YOLOv8 inference - let inputs = inputs!["images" => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?]; + let inputs = inputs![input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?]; let outputs = model .as_mut() .run(inputs) .map_err(YoloError::OrtInferenceError)?; - let output = outputs["output0"] + + let detection_output = outputs[output_name.as_str()] .try_extract_array::() - .map_err(YoloError::OrtExtractSensorError)? - .reversed_axes(); - let output = output.slice(s![.., .., 0]); + .map_err(YoloError::OrtExtractTensorError)?; + let detection_output = detection_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(detection_output.shape().to_vec()))?; - // Turn the output tensor into bounding boxes - let boxes = output - .axis_iter(Axis(0)) - .filter_map(|row| { - let (class_id, prob) = row - .iter() - .skip(4) // skip bounding box coordinates - .enumerate() - .map(|(index, value)| (index, *value)) - .reduce(|accum, row| if row.1 > accum.1 { row } else { accum }) - .filter(|(_, prob)| *prob >= probability_threshold)?; + let mask_output = outputs[mask_output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let mask_output = mask_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidMaskShape(mask_output.shape().to_vec()))?; + if mask_output.shape()[0] != 1 { + return Err(YoloError::InvalidMaskShape(mask_output.shape().to_vec())); + } - let label = labels[class_id].clone(); + let prototype_count = mask_output.shape()[1]; + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; + let boxes = decode_detections( + detection_output, + &labels, + probability_threshold, + letterbox, + Some(prototype_count), + )?; + let boxes = non_maximum_suppression(boxes, iou_threshold); - let xc = row[0_usize] / 640. * (raw_width as f32); - let yc = row[1_usize] / 640. * (raw_height as f32); - let w = row[2_usize] / 640. * (raw_width as f32); - let h = row[3_usize] / 640. * (raw_height as f32); + decode_segmentation_masks( + boxes, + mask_output.slice(ndarray::s![0, .., .., ..]), + letterbox, + ) +} - Some(YoloEntityOutput { - bounding_box: BoundingBox { - x1: xc - w / 2., - y1: yc - h / 2., - x2: xc + w / 2., - y2: yc + h / 2., - }, - label, - confidence: prob, - }) - }) - .collect::>(); +pub fn inference_segment_with_prompts( + model: &mut YoloModelSession, + YoloInputView { + tensor_view, + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }: YoloInputView, + prompts: YoloPromptEmbeddingsView<'_>, +) -> Result, YoloError> { + let prompts = validate_prompt_embeddings(prompts)?; + let iou_threshold = model.get_iou_threshold(); + let probability_threshold = model.get_probability_threshold(); + let input_name = model.get_input_name()?.to_owned(); + let prompt_input_name = model.get_prompt_input_name()?.to_owned(); + let output_name = model.get_output_name()?.to_owned(); + let mask_output_name = model.get_mask_output_name()?.to_owned(); + + let inputs = inputs![ + input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?, + prompt_input_name.as_str() => TensorRef::from_array_view(prompts.embeddings).map_err(YoloError::OrtInputError)?, + ]; + let outputs = model + .as_mut() + .run(inputs) + .map_err(YoloError::OrtInferenceError)?; + + let detection_output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let detection_output = detection_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(detection_output.shape().to_vec()))?; + + let mask_output = outputs[mask_output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let mask_output = mask_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidMaskShape(mask_output.shape().to_vec()))?; + if mask_output.shape()[0] != 1 { + return Err(YoloError::InvalidMaskShape(mask_output.shape().to_vec())); + } + + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; + let prototype_count = mask_output.shape()[1]; + let boxes = decode_detections( + detection_output, + prompts.labels, + probability_threshold, + letterbox, + Some(prototype_count), + )?; + let boxes = non_maximum_suppression(boxes, iou_threshold); - // Perform non-maximum suppression (NMS) - Ok(non_maximum_suppression(boxes, iou_threshold)) + decode_segmentation_masks( + boxes, + mask_output.slice(ndarray::s![0, .., .., ..]), + letterbox, + ) } diff --git a/src/model.rs b/src/model.rs index 362456c..69988b3 100644 --- a/src/model.rs +++ b/src/model.rs @@ -13,12 +13,24 @@ use crate::error::YoloError; pub struct YoloModelSession { pub session: ort::session::Session, pub labels: Vec, + pub input_name: Option, + pub output_name: Option, pub probability_threshold: Option, // default = 0.5 pub iou_threshold: Option, // default = 0.7 } impl YoloModelSession { + /// Wrap a generic ONNX session to a [`YoloModelSession`] without predefined labels. + pub fn from_filename(filename: impl AsRef) -> Result { + let session = ort::session::Session::builder() + .map_err(YoloError::OrtSessionBuildError)? + .commit_from_file(filename) + .map_err(YoloError::OrtSessionLoadError)?; + + Ok(Self::new(session, std::iter::empty::())) + } + /// Wrap a ONNX session to a [`YoloModelSession`]. /// /// The `session` is the ONNX runtime session, and the `labels` are the YOLO labels. @@ -29,6 +41,25 @@ impl YoloModelSession { Self { session, labels: labels.map(Into::into).collect(), + input_name: None, + output_name: None, + probability_threshold: None, + iou_threshold: None, + } + } + + /// Wrap an ONNX session to a [`YoloModelSession`] with explicit input/output names. + pub fn new_with_io_names( + session: ort::session::Session, + labels: impl Iterator>, + input_name: impl Into, + output_name: impl Into, + ) -> Self { + Self { + session, + labels: labels.map(Into::into).collect(), + input_name: Some(input_name.into()), + output_name: Some(output_name.into()), probability_threshold: None, iou_threshold: None, } @@ -122,6 +153,8 @@ impl YoloModelSession { Self { session, labels: LABELS.to_vec(), + input_name: None, + output_name: None, probability_threshold: None, iou_threshold: None, } @@ -143,6 +176,22 @@ impl YoloModelSession { Ok(Self::new_v8(session)) } + /// Load a YOLO-style ONNX model from a filename with caller-provided labels. + /// + /// This is useful for exported YOLOE closed-set models where the class list is not + /// the default COCO vocabulary. + pub fn from_filename_with_labels( + filename: impl AsRef, + labels: impl Iterator>, + ) -> Result { + let session = ort::session::Session::builder() + .map_err(YoloError::OrtSessionBuildError)? + .commit_from_file(filename) + .map_err(YoloError::OrtSessionLoadError)?; + + Ok(Self::new(session, labels)) + } + pub fn get_labels(&self) -> &[ArcStr] { &self.labels } @@ -154,6 +203,36 @@ impl YoloModelSession { pub fn get_iou_threshold(&self) -> f32 { self.iou_threshold.unwrap_or(0.7) } + + pub fn get_input_name(&self) -> Result<&str, YoloError> { + self.input_name + .as_deref() + .or_else(|| self.session.inputs().first().map(|input| input.name())) + .ok_or(YoloError::MissingModelInput) + } + + pub fn get_output_name(&self) -> Result<&str, YoloError> { + self.output_name + .as_deref() + .or_else(|| self.session.outputs().first().map(|output| output.name())) + .ok_or(YoloError::MissingModelOutput) + } + + pub fn get_prompt_input_name(&self) -> Result<&str, YoloError> { + self.session + .inputs() + .get(1) + .map(|input| input.name()) + .ok_or(YoloError::MissingPromptInput) + } + + pub fn get_mask_output_name(&self) -> Result<&str, YoloError> { + self.session + .outputs() + .get(1) + .map(|output| output.name()) + .ok_or(YoloError::MissingMaskOutput) + } } impl AsRef for YoloModelSession { diff --git a/src/prompt.rs b/src/prompt.rs new file mode 100644 index 0000000..3f32cf6 --- /dev/null +++ b/src/prompt.rs @@ -0,0 +1,205 @@ +use std::path::Path; + +use arcstr::ArcStr; +use ndarray::{Array2, Array3, ArrayView3, Axis}; +use ort::{inputs, value::TensorRef}; +use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer, TruncationParams}; + +use crate::error::YoloError; + +#[derive(Debug, Clone)] +pub struct YoloPromptEmbeddings { + labels: Vec, + embeddings: Array3, +} + +impl YoloPromptEmbeddings { + pub fn try_new( + labels: impl IntoIterator>, + embeddings: Array3, + ) -> Result { + if embeddings.shape().len() != 3 || embeddings.shape()[0] != 1 { + return Err(YoloError::InvalidPromptEmbeddingShape(embeddings.shape().to_vec())); + } + + let labels = labels.into_iter().map(Into::into).collect::>(); + if labels.is_empty() { + return Err(YoloError::EmptyPromptList); + } + if labels.len() != embeddings.shape()[1] { + return Err(YoloError::PromptLabelEmbeddingMismatch { + labels: labels.len(), + embeddings: embeddings.shape()[1], + }); + } + + Ok(Self { labels, embeddings }) + } + + pub fn labels(&self) -> &[ArcStr] { + &self.labels + } + + pub fn embeddings(&self) -> &Array3 { + &self.embeddings + } + + pub fn view(&self) -> YoloPromptEmbeddingsView<'_> { + YoloPromptEmbeddingsView { + labels: &self.labels, + embeddings: self.embeddings.view(), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct YoloPromptEmbeddingsView<'a> { + pub labels: &'a [ArcStr], + pub embeddings: ArrayView3<'a, f32>, +} + +#[derive(Debug)] +pub struct YoloPromptEncoderSession { + session: ort::session::Session, + tokenizer: Tokenizer, + max_length: Option, +} + +impl YoloPromptEncoderSession { + pub fn from_files( + model_path: impl AsRef, + tokenizer_path: impl AsRef, + ) -> Result { + let session = ort::session::Session::builder() + .map_err(YoloError::OrtSessionBuildError)? + .commit_from_file(model_path) + .map_err(YoloError::OrtSessionLoadError)?; + + let tokenizer = Tokenizer::from_file(tokenizer_path.as_ref()) + .map_err(|error| YoloError::TokenizerLoadError(error.to_string()))?; + + Ok(Self { + session, + tokenizer, + max_length: None, + }) + } + + pub fn encode<'a>( + &mut self, + prompts: impl IntoIterator, + ) -> Result { + let labels = prompts + .into_iter() + .map(ArcStr::from) + .collect::>(); + if labels.is_empty() { + return Err(YoloError::EmptyPromptList); + } + + let prompt_texts = labels.iter().map(ToString::to_string).collect::>(); + let mut tokenizer = self.tokenizer.clone(); + if let Some(max_length) = self.max_length { + tokenizer + .with_truncation(Some(TruncationParams { + max_length, + ..Default::default() + })) + .map_err(|error| YoloError::PromptTokenizationError(error.to_string()))?; + tokenizer.with_padding(Some(PaddingParams { + strategy: PaddingStrategy::Fixed(max_length), + ..Default::default() + })); + } else { + tokenizer.with_padding(Some(PaddingParams { + strategy: PaddingStrategy::BatchLongest, + ..Default::default() + })); + } + + let encodings = tokenizer + .encode_batch(prompt_texts, true) + .map_err(|error| YoloError::PromptTokenizationError(error.to_string()))?; + let batch_size = encodings.len(); + let sequence_length = encodings.first().map(|encoding| encoding.len()).unwrap_or(0); + + let input_ids = Array2::from_shape_vec( + (batch_size, sequence_length), + encodings + .iter() + .flat_map(|encoding| encoding.get_ids().iter().map(|token| i64::from(*token))) + .collect(), + ) + .map_err(|_| YoloError::PromptTokenizationError("failed to build input_ids tensor".into()))?; + let attention_mask = Array2::from_shape_vec( + (batch_size, sequence_length), + encodings + .iter() + .flat_map(|encoding| encoding.get_attention_mask().iter().map(|token| i64::from(*token))) + .collect(), + ) + .map_err(|_| YoloError::PromptTokenizationError("failed to build attention_mask tensor".into()))?; + let token_type_ids = Array2::::zeros((batch_size, sequence_length)); + + let output_name = self + .session + .outputs() + .first() + .map(|output| output.name().to_owned()) + .ok_or(YoloError::MissingModelOutput)?; + let outputs = match self.session.inputs().len() { + 1 => { + let input_name = self.session.inputs()[0].name().to_owned(); + let inputs = inputs![input_name.as_str() => TensorRef::from_array_view(input_ids.view()).map_err(YoloError::OrtInputError)?]; + self.session.run(inputs).map_err(YoloError::OrtInferenceError)? + } + 2 => { + let input_name = self.session.inputs()[0].name().to_owned(); + let mask_name = self.session.inputs()[1].name().to_owned(); + let inputs = inputs![ + input_name.as_str() => TensorRef::from_array_view(input_ids.view()).map_err(YoloError::OrtInputError)?, + mask_name.as_str() => TensorRef::from_array_view(attention_mask.view()).map_err(YoloError::OrtInputError)?, + ]; + self.session.run(inputs).map_err(YoloError::OrtInferenceError)? + } + 3 => { + let first_name = self.session.inputs()[0].name().to_owned(); + let second_name = self.session.inputs()[1].name().to_owned(); + let third_name = self.session.inputs()[2].name().to_owned(); + let inputs = inputs![ + first_name.as_str() => TensorRef::from_array_view(input_ids.view()).map_err(YoloError::OrtInputError)?, + second_name.as_str() => TensorRef::from_array_view(attention_mask.view()).map_err(YoloError::OrtInputError)?, + third_name.as_str() => TensorRef::from_array_view(token_type_ids.view()).map_err(YoloError::OrtInputError)?, + ]; + self.session.run(inputs).map_err(YoloError::OrtInferenceError)? + } + count => return Err(YoloError::UnsupportedPromptEncoderInputCount(count)), + }; + let output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + + let mut embeddings = if let Ok(output) = output.view().into_dimensionality::() { + output.to_owned().insert_axis(Axis(0)) + } else if let Ok(output) = output.view().into_dimensionality::() { + if output.shape()[0] != 1 { + return Err(YoloError::InvalidPromptEncoderOutputShape(output.shape().to_vec())); + } + output.to_owned() + } else { + return Err(YoloError::InvalidPromptEncoderOutputShape(output.shape().to_vec())); + }; + + normalize_last_axis(&mut embeddings); + YoloPromptEmbeddings::try_new(labels, embeddings) + } +} + +fn normalize_last_axis(embeddings: &mut Array3) { + for mut prompt in embeddings.axis_iter_mut(Axis(1)) { + let norm = prompt.iter().map(|value| value * value).sum::().sqrt(); + if norm > 0.0 { + prompt.iter_mut().for_each(|value| *value /= norm); + } + } +} \ No newline at end of file