mii/usls
1
0
Fork 0
mirror of https://github.com/mii443/usls.git synced 2025-08-22 15:45:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

🐍 v0.1.0 (#53)

Browse Source
This commit is contained in:
Jamjamjon
2025-01-12 16:59:57 +08:00
committed by GitHub
parent 4e932c4910
commit 0f2d84b8c5
256 changed files with 12485 additions and 9088 deletions
Download Patch File Download Diff File Expand all files Collapse all files

224
.github/workflows/rust-ci.yml vendored
View File

@ -11,193 +11,73 @@ env:
jobs:
build-on-linux:
name: build / linux / ffmpeg ${{ matrix.ffmpeg_version }}
runs-on: ubuntu-latest
container: jrottenberg/ffmpeg:${{ matrix.ffmpeg_version }}-ubuntu
check:
name: Check
runs-on: ${{ matrix.os }}
strategy:
matrix:
ffmpeg_version: ["4.3", "4.4", "5.0", "5.1", "6.0", "6.1", "7.0"]
fail-fast: false
os: [ubuntu-latest, macOS-latest, windows-latest]
rust: [stable]
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
apt update
apt install -y --no-install-recommends clang curl pkg-config
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
- uses: actions/checkout@v2
- uses: actions-rs/toolchain@v1
with:
toolchain: stable
- name: Build
run: cargo build
build-on-macos:
name: build / macos / ffmpeg latest
runs-on: macos-latest
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
brew install ffmpeg pkg-config
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
profile: minimal
toolchain: ${{ matrix.rust }}
override: true
- uses: actions-rs/cargo@v1
with:
toolchain: stable
- name: Build
run: cargo build
build-on-windows:
name: build / windows / ffmpeg latest
runs-on: windows-latest
env:
FFMPEG_DOWNLOAD_URL: https://www.gyan.dev/ffmpeg/builds/ffmpeg-release-full-shared.7z
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
$VCINSTALLDIR = $(& "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" -latest -property installationPath)
Add-Content $env:GITHUB_ENV "LIBCLANG_PATH=${VCINSTALLDIR}\VC\Tools\LLVM\x64\bin`n"
Invoke-WebRequest "${env:FFMPEG_DOWNLOAD_URL}" -OutFile ffmpeg-release-full-shared.7z
7z x ffmpeg-release-full-shared.7z
mkdir ffmpeg
mv ffmpeg-*/* ffmpeg/
Add-Content $env:GITHUB_ENV "FFMPEG_DIR=${pwd}\ffmpeg`n"
Add-Content $env:GITHUB_PATH "${pwd}\ffmpeg\bin`n"
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
with:
toolchain: stable
- name: Build
run: cargo build
test-on-linux:
name: test / linux / ffmpeg ${{ matrix.ffmpeg_version }}
runs-on: ubuntu-latest
container: jrottenberg/ffmpeg:${{ matrix.ffmpeg_version }}-ubuntu
command: check
args: --all
test:
name: Test
runs-on: ${{ matrix.os }}
strategy:
matrix:
ffmpeg_version: ["4.3", "4.4", "5.0", "5.1", "6.0", "6.1", "7.0"]
fail-fast: false
os: [ubuntu-latest, macOS-latest, windows-latest]
rust: [stable]
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
apt update
apt install -y --no-install-recommends clang curl pkg-config
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
- uses: actions/checkout@v2
- uses: actions-rs/toolchain@v1
with:
toolchain: stable
- name: Run Tests with All Features
run: cargo test --all-features
- name: Run Tests in Release Mode
run: cargo test --release
test-on-macos:
name: test / macos / ffmpeg latest
runs-on: macos-latest
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
brew install ffmpeg pkg-config
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
profile: minimal
toolchain: ${{ matrix.rust }}
override: true
- uses: actions-rs/cargo@v1
with:
toolchain: stable
command: test
args: --all
- name: Run Tests with All Features
run: cargo test --all-features
- name: Run Tests in Release Mode
run: cargo test --release
test-on-windows:
name: test / windows / ffmpeg latest
runs-on: windows-latest
env:
FFMPEG_DOWNLOAD_URL: https://www.gyan.dev/ffmpeg/builds/ffmpeg-release-full-shared.7z
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
$VCINSTALLDIR = $(& "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" -latest -property installationPath)
Add-Content $env:GITHUB_ENV "LIBCLANG_PATH=${VCINSTALLDIR}\VC\Tools\LLVM\x64\bin`n"
Invoke-WebRequest "${env:FFMPEG_DOWNLOAD_URL}" -OutFile ffmpeg-release-full-shared.7z
7z x ffmpeg-release-full-shared.7z
mkdir ffmpeg
mv ffmpeg-*/* ffmpeg/
Add-Content $env:GITHUB_ENV "FFMPEG_DIR=${pwd}\ffmpeg`n"
Add-Content $env:GITHUB_PATH "${pwd}\ffmpeg\bin`n"
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
with:
toolchain: stable
- name: Run Tests with All Features
run: cargo test --all-features
- name: Run Tests in Release Mode
run: cargo test --release
lints:
fmt:
name: Rustfmt
runs-on: ubuntu-latest
container: jrottenberg/ffmpeg:6-ubuntu
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Install dependencies
run: |
apt update
apt install -y --no-install-recommends clang curl pkg-config
- name: Setup Rust
uses: dtolnay/rust-toolchain@v1
- uses: actions/checkout@v2
- uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: stable
components: rustfmt, clippy
override: true
- run: rustup component add rustfmt
- uses: actions-rs/cargo@v1
with:
command: fmt
args: --all -- --check
- name: Rustfmt
run: cargo fmt --all -- --check
clippy:
name: Clippy
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: stable
override: true
- run: rustup component add clippy
- uses: actions-rs/cargo@v1
with:
command: clippy
args: --all --all-targets -- -D warnings
- name: Clippy
run: cargo clippy --all --all-targets --all-features -- -D warnings

3
.gitignore vendored
View File

@ -3,6 +3,8 @@
debug/
target/
**/*.DS_Store
# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
Cargo.lock
@ -13,7 +15,6 @@ Cargo.lock
# MSVC Windows builds of rustc generate these, which store debugging information
*.pdb
.debug
.vscode
runs/

74
Cargo.toml
View File

@ -1,64 +1,66 @@
[package]
name = "usls"
version = "0.0.20"
version = "0.1.0"
rust-version = "1.79"
edition = "2021"
description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models."
repository = "https://github.com/jamjamjon/usls"
authors = ["Jamjamjon <xxyydzml@outlook.com>"]
license = "MIT"
readme = "README.md"
exclude = ["assets/*", "examples/*", "scripts/*", "runs/*"]
exclude = ["assets/*", "examples/*", "runs/*", "benches/*"]
[dependencies]
clap = { version = "4.2.4", features = ["derive"] }
aksr = { version = "0.0.2" }
image = { version = "0.25.2" }
imageproc = { version = "0.24" }
ndarray = { version = "0.16.1", features = ["rayon"] }
ort = { version = "2.0.0-rc.9", default-features = false }
rayon = { version = "1.10.0" }
anyhow = { version = "1.0.75" }
regex = { version = "1.5.4" }
rand = { version = "0.8.5" }
chrono = { version = "0.4.30" }
half = { version = "2.3.1" }
dirs = { version = "5.0.1" }
ureq = { version = "2.9.1", default-features = true, features = [
"socks-proxy",
] }
tokenizers = { version = "0.15.2" }
rayon = "1.10.0"
log = { version = "0.4.22" }
indicatif = "0.17.8"
image = "0.25.2"
imageproc = { version = "0.24" }
ab_glyph = "0.2.23"
geo = "0.28.0"
prost = "0.12.4"
fast_image_resize = { version = "4.2.1", features = ["image"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde = { version = "1.0", features = ["derive"] }
ort = { version = "2.0.0-rc.9", default-features = false}
prost = "0.12.4"
ab_glyph = "0.2.23"
dirs = { version = "5.0.1" }
tempfile = "3.12.0"
video-rs = { version = "0.9.0", features = ["ndarray"] }
geo = "0.28.0"
half = { version = "2.3.1" }
ureq = { version = "2.12.1", default-features = false, features = [ "tls" ] }
fast_image_resize = { version = "4.2.1", features = ["image"]}
natord = "1.0.9"
tracing = "0.1.40"
tracing-subscriber = "0.3.18"
minifb = "0.27.0"
video-rs = { version = "0.10.0", features = ["ndarray"], optional = true }
minifb = { version = "0.27.0", optional = true }
sha2 = "0.10.8"
[dev-dependencies]
argh = "0.1.13"
tracing-subscriber = { version = "0.3.18", features = ["env-filter", "chrono"] }
[[example]]
name = "viewer"
required-features = ["ffmpeg"]
[features]
default = [
"ort/load-dynamic",
"ort/copy-dylibs",
"ort/half",
"ort/ndarray",
"ort/cuda",
"ort/tensorrt",
"ort/coreml",
"ort/copy-dylibs",
"ort/load-dynamic",
"ort/half",
]
auto = ["ort/download-binaries"]
ffmpeg = ["dep:video-rs", "dep:minifb"]
cuda = [ "ort/cuda" ]
trt = [ "ort/tensorrt" ]
mps = [ "ort/coreml" ]
[dev-dependencies]
criterion = "0.5.1"
[[bench]]
name = "yolo"
harness = false
[lib]
bench = false
[profile.release]
# lto = true
strip = true
panic = "abort"

282
README.md
View File

@ -1,221 +1,161 @@
<p align="center">
<h2 align="center">usls</h2>
</p>
<p align="center">
<a href="https://docs.rs/usls"><strong>Documentation</strong></a>
<br>
<br>
<a href="https://github.com/jamjamjon/usls/actions/workflows/rust-ci.yml">
<img src="https://github.com/jamjamjon/usls/actions/workflows/rust-ci.yml/badge.svg" alt="Rust Continuous Integration Badge">
</a>
<a href='https://crates.io/crates/usls'>
<img src='https://img.shields.io/crates/v/usls.svg' alt='usls Version'>
</a>
<a href='https://crates.io/crates/usls'>
<img src='https://img.shields.io/crates/msrv/usls-yellow?' alt='Rust MSRV'>
</a>
<a href='https://github.com/microsoft/onnxruntime/releases'>
<img src='https://img.shields.io/badge/ONNXRuntime-v1.19.x-239DFF?style=for-the-badge&logo=onnx' alt='ONNXRuntime Release Page'>
<img src='https://img.shields.io/badge/onnxruntime-%3E%3D%201.19.0-3399FF' alt='ONNXRuntime MSRV'>
</a>
<a href='https://developer.nvidia.com/cuda-toolkit-archive'>
<img src='https://img.shields.io/badge/CUDA-12.x-76B900?style=for-the-badge&logo=nvidia' alt='CUDA Toolkit Page'>
<img src='https://img.shields.io/badge/cuda-%3E%3D%2012.0-green' alt='CUDA MSRV'>
</a>
<a href='https://developer.nvidia.com/tensorrt'>
<img src='https://img.shields.io/badge/TensorRT-10.x.x.x-76B900?style=for-the-badge&logo=nvidia' alt='TensorRT Page'>
<img src='https://img.shields.io/badge/TensorRT-%3E%3D%2012.0-0ABF53' alt='TensorRT MSRV'>
</a>
<a href="https://crates.io/crates/usls">
<img alt="Crates.io Total Downloads" src="https://img.shields.io/crates/d/usls?&color=946CE6">
</a>
</p>
<p align="center">
<a href='https://crates.io/crates/usls'>
<img src='https://img.shields.io/crates/v/usls.svg?style=for-the-badge&logo=rust' alt='Crates Page'>
<a href="./examples">
<img src="https://img.shields.io/badge/Examples-1A86FD?&logo=anki" alt="Examples">
</a>
<!-- Documentation Badge -->
<!-- <a href="https://docs.rs/usls">
<img src='https://img.shields.io/badge/Documents-usls-000000?style=for-the-badge&logo=docs.rs' alt='Documentation'>
</a> -->
<!-- Downloads Badge -->
<a href="">
<img alt="Crates.io Total Downloads" src="https://img.shields.io/crates/d/usls?style=for-the-badge&color=3ECC5F">
<a href='https://docs.rs/usls'>
<img src='https://img.shields.io/badge/Docs-usls-yellow?&logo=docs.rs&color=FFA200' alt='usls documentation'>
</a>
</p>
**`usls`** is a Rust library integrated with **ONNXRuntime** that provides a collection of state-of-the-art models for **Computer Vision** and **Vision-Language** tasks, including:
**usls** is a Rust library integrated with **ONNXRuntime**, offering a suite of advanced models for **Computer Vision** and **Vision-Language** tasks, including:
- **YOLO Models**: [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [YOLOv7](https://github.com/WongKinYiu/yolov7), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [YOLOv10](https://github.com/THU-MIG/yolov10), [YOLOv11](https://github.com/ultralytics/ultralytics)
- **YOLO Models**: [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [YOLOv7](https://github.com/WongKinYiu/yolov7), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [YOLOv10](https://github.com/THU-MIG/yolov10), [YOLO11](https://github.com/ultralytics/ultralytics)
- **SAM Models**: [SAM](https://github.com/facebookresearch/segment-anything), [SAM2](https://github.com/facebookresearch/segment-anything-2), [MobileSAM](https://github.com/ChaoningZhang/MobileSAM), [EdgeSAM](https://github.com/chongzhou96/EdgeSAM), [SAM-HQ](https://github.com/SysCV/sam-hq), [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM)
- **Vision Models**: [RTDETR](https://arxiv.org/abs/2304.08069), [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo), [DB](https://arxiv.org/abs/1911.08947), [SVTR](https://arxiv.org/abs/2205.00159), [Depth-Anything-v1-v2](https://github.com/LiheYoung/Depth-Anything), [DINOv2](https://github.com/facebookresearch/dinov2), [MODNet](https://github.com/ZHKKKe/MODNet), [Sapiens](https://arxiv.org/abs/2408.12569), [DepthPro](https://github.com/apple/ml-depth-pro)
- **Vision-Language Models**: [CLIP](https://github.com/openai/CLIP), [BLIP](https://arxiv.org/abs/2201.12086), [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [Florence2](https://arxiv.org/abs/2311.06242)
- **Vision Models**: [RT-DETR](https://arxiv.org/abs/2304.08069), [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo), [Depth-Anything](https://github.com/LiheYoung/Depth-Anything), [DINOv2](https://github.com/facebookresearch/dinov2), [MODNet](https://github.com/ZHKKKe/MODNet), [Sapiens](https://arxiv.org/abs/2408.12569), [DepthPro](https://github.com/apple/ml-depth-pro), [FastViT](https://github.com/apple/ml-fastvit), [BEiT](https://github.com/microsoft/unilm/tree/master/beit), [MobileOne](https://github.com/apple/ml-mobileone)
- **Vision-Language Models**: [CLIP](https://github.com/openai/CLIP), [jina-clip-v1](https://huggingface.co/jinaai/jina-clip-v1), [BLIP](https://arxiv.org/abs/2201.12086), [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [Florence2](https://arxiv.org/abs/2311.06242)
- **OCR Models**: [FAST](https://github.com/czczup/FAST), [DB(PaddleOCR-Det)](https://arxiv.org/abs/1911.08947), [SVTR(PaddleOCR-Rec)](https://arxiv.org/abs/2205.00159), [SLANet](https://paddlepaddle.github.io/PaddleOCR/latest/algorithm/table_recognition/algorithm_table_slanet.html), [TrOCR](https://huggingface.co/microsoft/trocr-base-printed), [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO)
<details>
<summary>Click to expand Supported Models</summary>
## Supported Models
| Model | Task / Type | Example | CUDA f32 | CUDA f16 | TensorRT f32 | TensorRT f16 |
|---------------------------------------------------------------------|----------------------------------------------------------------------------------------------|----------------------------|----------|----------|--------------|--------------|
| [YOLOv5](https://github.com/ultralytics/yolov5) | Classification<br>Object Detection<br>Instance Segmentation | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [YOLOv6](https://github.com/meituan/YOLOv6) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [YOLOv7](https://github.com/WongKinYiu/yolov7) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [YOLOv8](https://github.com/ultralytics/ultralytics) | Object Detection<br>Instance Segmentation<br>Classification<br>Oriented Object Detection<br>Keypoint Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [YOLOv9](https://github.com/WongKinYiu/yolov9) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [YOLOv10](https://github.com/THU-MIG/yolov10) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [YOLOv11](https://github.com/ultralytics/ultralytics) | Object Detection<br>Instance Segmentation<br>Classification<br>Oriented Object Detection<br>Keypoint Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [RTDETR](https://arxiv.org/abs/2304.08069) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM) | Instance Segmentation | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [SAM](https://github.com/facebookresearch/segment-anything) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | | |
| [SAM2](https://github.com/facebookresearch/segment-anything-2) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | | |
| [MobileSAM](https://github.com/ChaoningZhang/MobileSAM) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | | |
| [EdgeSAM](https://github.com/chongzhou96/EdgeSAM) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | | |
| [SAM-HQ](https://github.com/SysCV/sam-hq) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | | |
| [YOLO-World](https://github.com/AILab-CVC/YOLO-World) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ |
| [DINOv2](https://github.com/facebookresearch/dinov2) | Vision-Self-Supervised | [demo](examples/dinov2) | ✅ | ✅ | ✅ | ✅ |
| [CLIP](https://github.com/openai/CLIP) | Vision-Language | [demo](examples/clip) | ✅ | ✅ | ✅ Visual<br>❌ Textual | ✅ Visual<br>❌ Textual |
| [BLIP](https://github.com/salesforce/BLIP) | Vision-Language | [demo](examples/blip) | ✅ | ✅ | ✅ Visual<br>❌ Textual | ✅ Visual<br>❌ Textual |
| [DB](https://arxiv.org/abs/1911.08947) | Text Detection | [demo](examples/db) | ✅ | ✅ | ✅ | ✅ |
| [SVTR](https://arxiv.org/abs/2205.00159) | Text Recognition | [demo](examples/svtr) | ✅ | ✅ | ✅ | ✅ |
| [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo) | Keypoint Detection | [demo](examples/rtmo) | ✅ | ✅ | ❌ | ❌ |
| [YOLOPv2](https://arxiv.org/abs/2208.11434) | Panoptic Driving Perception | [demo](examples/yolop) | ✅ | ✅ | ✅ | ✅ |
| [Depth-Anything v1 & v2](https://github.com/LiheYoung/Depth-Anything) | Monocular Depth Estimation | [demo](examples/depth-anything) | ✅ | ✅ | ❌ | ❌ |
| [MODNet](https://github.com/ZHKKKe/MODNet) | Image Matting | [demo](examples/modnet) | ✅ | ✅ | ✅ | ✅ |
| [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO) | Open-Set Detection With Language | [demo](examples/grounding-dino) | ✅ | ✅ | | |
| [Sapiens](https://github.com/facebookresearch/sapiens/tree/main) | Body Part Segmentation | [demo](examples/sapiens) | ✅ | ✅ | | |
| [Florence2](https://arxiv.org/abs/2311.06242) | a Variety of Vision Tasks | [demo](examples/florence2) | ✅ | ✅ | | |
| [DepthPro](https://github.com/apple/ml-depth-pro) | Monocular Depth Estimation | [demo](examples/depth-pro) | ✅ | ✅ | | |
<summary>👉 More Supported Models</summary>
| Model | Task / Description | Example | CoreML | CUDA<br />FP32 | CUDA<br />FP16 | TensorRT<br />FP32 | TensorRT<br />FP16 |
| -------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | ---------------------------- | ------ | -------------- | -------------- | ------------------ | ------------------ |
| [BEiT](https://github.com/microsoft/unilm/tree/master/beit) | Image Classification | [demo](examples/beit) | ✅ | ✅ | ✅ | | |
| [ConvNeXt](https://github.com/facebookresearch/ConvNeXt) | Image Classification | [demo](examples/convnext) | ✅ | ✅ | ✅ | | |
| [FastViT](https://github.com/apple/ml-fastvit) | Image Classification | [demo](examples/fastvit) | ✅ | ✅ | ✅ | | |
| [MobileOne](https://github.com/apple/ml-mobileone) | Image Classification | [demo](examples/mobileone) | ✅ | ✅ | ✅ | | |
| [DeiT](https://github.com/facebookresearch/deit) | Image Classification | [demo](examples/deit) | ✅ | ✅ | ✅ | | |
| [DINOv2](https://github.com/facebookresearch/dinov2) | Vision Embedding | [demo](examples/dinov2) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLOv5](https://github.com/ultralytics/yolov5) | Image Classification<br />Object Detection<br />Instance Segmentation | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLOv6](https://github.com/meituan/YOLOv6) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLOv7](https://github.com/WongKinYiu/yolov7) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLOv8<br />YOLO11](https://github.com/ultralytics/ultralytics) | Object Detection<br />Instance Segmentation<br />Image Classification<br />Oriented Object Detection<br />Keypoint Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLOv9](https://github.com/WongKinYiu/yolov9) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLOv10](https://github.com/THU-MIG/yolov10) | Object Detection | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [RT-DETR](https://github.com/lyuwenyu/RT-DETR) | Object Detection | [demo](examples/rtdetr) | ✅ | ✅ | ✅ | | |
| [PP-PicoDet](https://github.com/PaddlePaddle/PaddleDetection/tree/release/2.8/configs/picodet) | Object Detection | [demo](examples/picodet-layout) | ✅ | ✅ | ✅ | | |
| [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO) | Object Detection | [demo](examples/picodet-layout) | ✅ | ✅ | ✅ | | |
| [D-FINE](https://github.com/manhbd-22022602/D-FINE) | Object Detection | [demo](examples/d-fine) | ✅ | ✅ | ✅ | | |
| [DEIM](https://github.com/ShihuaHuang95/DEIM) | Object Detection | [demo](examples/deim) | ✅ | ✅ | ✅ | | |
| [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo) | Keypoint Detection | [demo](examples/rtmo) | ✅ | ✅ | ✅ | ❌ | ❌ |
| [SAM](https://github.com/facebookresearch/segment-anything) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | ✅ | | |
| [SAM2](https://github.com/facebookresearch/segment-anything-2) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | ✅ | | |
| [MobileSAM](https://github.com/ChaoningZhang/MobileSAM) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | ✅ | | |
| [EdgeSAM](https://github.com/chongzhou96/EdgeSAM) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | ✅ | | |
| [SAM-HQ](https://github.com/SysCV/sam-hq) | Segment Anything | [demo](examples/sam) | ✅ | ✅ | ✅ | | |
| [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM) | Instance Segmentation | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [YOLO-World](https://github.com/AILab-CVC/YOLO-World) | Open-Set Detection With Language | [demo](examples/yolo) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO) | Open-Set Detection With Language | [demo](examples/grounding-dino) | ✅ | ✅ | ✅ | | |
| [CLIP](https://github.com/openai/CLIP) | Vision-Language Embedding | [demo](examples/clip) | ✅ | ✅ | ✅ | ❌ | ❌ |
| [jina-clip-v1](https://huggingface.co/jinaai/jina-clip-v1) | Vision-Language Embedding | [demo](examples/clip) | ✅ | ✅ | ✅ | ❌ | ❌ |
| [BLIP](https://github.com/salesforce/BLIP) | Image Captioning | [demo](examples/blip) | ✅ | ✅ | ✅ | ❌ | ❌ |
| [DB(PaddleOCR-Det)](https://arxiv.org/abs/1911.08947) | Text Detection | [demo](examples/db) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [FAST](https://github.com/czczup/FAST) | Text Detection | [demo](examples/fast) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [LinkNet](https://arxiv.org/abs/1707.03718) | Text Detection | [demo](examples/linknet) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [SVTR(PaddleOCR-Rec)](https://arxiv.org/abs/2205.00159) | Text Recognition | [demo](examples/svtr) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [SLANet](https://paddlepaddle.github.io/PaddleOCR/latest/algorithm/table_recognition/algorithm_table_slanet.html) | Tabel Recognition | [demo](examples/slanet) | ✅ | ✅ | ✅ | | |
| [TrOCR](https://huggingface.co/microsoft/trocr-base-printed) | Text Recognition | [demo](examples/trocr) | ✅ | ✅ | ✅ | | |
| [YOLOPv2](https://arxiv.org/abs/2208.11434) | Panoptic Driving Perception | [demo](examples/yolop) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [DepthAnything v1<br />DepthAnything v2](https://github.com/LiheYoung/Depth-Anything) | Monocular Depth Estimation | [demo](examples/depth-anything) | ✅ | ✅ | ✅ | ❌ | ❌ |
| [DepthPro](https://github.com/apple/ml-depth-pro) | Monocular Depth Estimation | [demo](examples/depth-pro) | ✅ | ✅ | ✅ | | |
| [MODNet](https://github.com/ZHKKKe/MODNet) | Image Matting | [demo](examples/modnet) | ✅ | ✅ | ✅ | ✅ | ✅ |
| [Sapiens](https://github.com/facebookresearch/sapiens/tree/main) | Foundation for Human Vision Models | [demo](examples/sapiens) | ✅ | ✅ | ✅ | | |
| [Florence2](https://arxiv.org/abs/2311.06242) | a Variety of Vision Tasks | [demo](examples/florence2) | ✅ | ✅ | ✅ | | |
</details>
## ⛳️ Cargo Features
## ⛳️ ONNXRuntime Linking
By default, **none of the following features are enabled**. You can enable them as needed:
- **`auto`**: Automatically downloads prebuilt ONNXRuntime binaries from Pyke’s CDN for supported platforms.
- If disabled, you'll need to [compile `ONNXRuntime` from source](https://github.com/microsoft/onnxruntime) or [download a precompiled package](https://github.com/microsoft/onnxruntime/releases), and then [link it manually](https://ort.pyke.io/setup/linking).
<details>
<summary>You have two options to link the ONNXRuntime library</summary>
<summary>👉 For Linux or macOS Users</summary>
- ### Option 1: Manual Linking
- #### For detailed setup instructions, refer to the [ORT documentation](https://ort.pyke.io/setup/linking).
- #### For Linux or macOS Users:
- Download the ONNX Runtime package from the [Releases page](https://github.com/microsoft/onnxruntime/releases).
- Download from the [Releases page](https://github.com/microsoft/onnxruntime/releases).
- Set up the library path by exporting the `ORT_DYLIB_PATH` environment variable:
```shell
export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.19.0
```
- ### Option 2: Automatic Download
Just use `--features auto`
```shell
cargo run -r --example yolo --features auto
export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.20.1
```
</details>
- **`ffmpeg`**: Adds support for video streams, real-time frame visualization, and video export.
## 🎈 Demo
- Powered by [video-rs](https://github.com/oddity-ai/video-rs) and [minifb](https://github.com/emoon/rust_minifb). For any issues related to `ffmpeg` features, please refer to the issues of these two crates.
- **`cuda`**: Enables the NVIDIA TensorRT provider.
- **`trt`**: Enables the NVIDIA TensorRT provider.
- **`mps`**: Enables the Apple CoreML provider.
## 🎈 Example
* **Using `CUDA`**
```Shell
cargo run -r --example yolo # blip, clip, yolop, svtr, db, ...
```
cargo run -r -F cuda --example yolo -- --device cuda:0
```
* **Using Apple `CoreML`**
```
cargo run -r -F mps --example yolo -- --device mps
```
* **Using `TensorRT`**
```
cargo run -r -F trt --example yolo -- --device trt
```
* **Using `CPU`**
```
cargo run -r --example yolo
```
All examples are located in the [examples](./examples/) directory.
## 🥂 Integrate Into Your Own Project
- #### Add `usls` as a dependency to your project's `Cargo.toml`
Add `usls` as a dependency to your project's `Cargo.toml`
```Shell
cargo add usls
cargo add usls -F cuda
```
Or use a specific commit:
```Toml
[dependencies]
usls = { git = "https://github.com/jamjamjon/usls", rev = "commit-sha" }
```
- #### Follow the pipeline
- Build model with the provided `models` and `Options`
- Load images, video and stream with `DataLoader`
- Do inference
- Retrieve inference results from `Vec<Y>`
- Annotate inference results with `Annotator`
- Display images and write them to video with `Viewer`
<br/>
<details>
<summary>example code</summary>
```rust
use usls::{models::YOLO, Annotator, DataLoader, Nms, Options, Vision, YOLOTask, YOLOVersion};
fn main() -> anyhow::Result<()> {
// Build model with Options
let options = Options::new()
.with_trt(0)
.with_model("yolo/v8-m-dyn.onnx")?
.with_yolo_version(YOLOVersion::V8) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
.with_yolo_task(YOLOTask::Detect) // YOLOTask: Classify, Detect, Pose, Segment, Obb
.with_ixx(0, 0, (1, 2, 4).into())
.with_ixx(0, 2, (0, 640, 640).into())
.with_ixx(0, 3, (0, 640, 640).into())
.with_confs(&[0.2]);
let mut model = YOLO::new(options)?;
// Build DataLoader to load image(s), video, stream
let dl = DataLoader::new(
// "./assets/bus.jpg", // local image
// "images/bus.jpg", // remote image
// "../images-folder", // local images (from folder)
// "../demo.mp4", // local video
// "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // online video
"rtsp://admin:kkasd1234@192.168.2.217:554/h264/ch1/", // stream
)?
.with_batch(2) // iterate with batch_size = 2
.build()?;
// Build annotator
let annotator = Annotator::new()
.with_bboxes_thickness(4)
.with_saveout("YOLO-DataLoader");
// Build viewer
let mut viewer = Viewer::new().with_delay(10).with_scale(1.).resizable(true);
// Run and annotate results
for (xs, _) in dl {
let ys = model.forward(&xs, false)?;
// annotator.annotate(&xs, &ys);
let images_plotted = annotator.plot(&xs, &ys, false)?;
// show image
viewer.imshow(&images_plotted)?;
// check out window and key event
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) {
break;
}
// write video
viewer.write_batch(&images_plotted)?;
// Retrieve inference results
for y in ys {
// bboxes
if let Some(bboxes) = y.bboxes() {
for bbox in bboxes {
println!(
"Bbox: {}, {}, {}, {}, {}, {}",
bbox.xmin(),
bbox.ymin(),
bbox.xmax(),
bbox.ymax(),
bbox.confidence(),
bbox.id(),
);
}
}
}
}
// finish video write
viewer.finish_write()?;
Ok(())
}
```
</details>
</br>
## 🥳 If you find this helpful, please give it a star ⭐
## 📌 License
This project is licensed under [LICENSE](LICENSE).

94
benches/yolo.rs
View File

@ -1,94 +0,0 @@
use anyhow::Result;
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use usls::{models::YOLO, DataLoader, Options, Vision, YOLOTask, YOLOVersion};
enum Stage {
Pre,
Run,
Post,
Pipeline,
}
fn yolo_stage_bench(
model: &mut YOLO,
x: &[image::DynamicImage],
stage: Stage,
n: u64,
) -> std::time::Duration {
let mut t_pre = std::time::Duration::new(0, 0);
let mut t_run = std::time::Duration::new(0, 0);
let mut t_post = std::time::Duration::new(0, 0);
let mut t_pipeline = std::time::Duration::new(0, 0);
for _ in 0..n {
let t0 = std::time::Instant::now();
let xs = model.preprocess(x).unwrap();
t_pre += t0.elapsed();
let t = std::time::Instant::now();
let xs = model.inference(xs).unwrap();
t_run += t.elapsed();
let t = std::time::Instant::now();
let _ys = black_box(model.postprocess(xs, x).unwrap());
t_post += t.elapsed();
t_pipeline += t0.elapsed();
}
match stage {
Stage::Pre => t_pre,
Stage::Run => t_run,
Stage::Post => t_post,
Stage::Pipeline => t_pipeline,
}
}
pub fn benchmark_cuda(c: &mut Criterion, h: isize, w: isize) -> Result<()> {
let mut group = c.benchmark_group(format!("YOLO ({}-{})", w, h));
group
.significance_level(0.05)
.sample_size(80)
.measurement_time(std::time::Duration::new(20, 0));
let options = Options::default()
.with_yolo_version(YOLOVersion::V8) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
.with_yolo_task(YOLOTask::Detect) // YOLOTask: Classify, Detect, Pose, Segment, Obb
.with_model("yolo/v8-m-dyn.onnx")?
.with_cuda(0)
// .with_cpu()
.with_dry_run(0)
.with_ixx(0, 2, (320, h, 1280).into())
.with_ixx(0, 3, (320, w, 1280).into())
.with_confs(&[0.2, 0.15]);
let mut model = YOLO::new(options)?;
let xs = [DataLoader::try_read("./assets/bus.jpg")?];
group.bench_function("pre-process", |b| {
b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Pre, n))
});
group.bench_function("run", |b| {
b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Run, n))
});
group.bench_function("post-process", |b| {
b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Post, n))
});
group.bench_function("pipeline", |b| {
b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Pipeline, n))
});
group.finish();
Ok(())
}
pub fn criterion_benchmark(c: &mut Criterion) {
// benchmark_cuda(c, 416, 416).unwrap();
benchmark_cuda(c, 640, 640).unwrap();
benchmark_cuda(c, 448, 768).unwrap();
// benchmark_cuda(c, 800, 800).unwrap();
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);

6
examples/beit/README.md Normal file
View File

@ -0,0 +1,6 @@
## Quick Start
```shell
cargo run -r -F cuda --example beit -- --device cuda --dtype fp16
```

52
examples/beit/main.rs Normal file
View File

@ -0,0 +1,52 @@
use usls::{models::ImageClassifier, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// source image
#[argh(
option,
default = "vec![
String::from(\"images/dog.jpg\"),
String::from(\"images/siamese.png\"),
String::from(\"images/ailurus-fulgens.jpg\"),
]"
)]
source: Vec<String>,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::beit_base()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = ImageClassifier::try_from(options)?;
// load images
let xs = DataLoader::try_read_batch(&args.source)?;
// run
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default().with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

14
examples/blip/README.md
View File

@ -3,20 +3,12 @@ This demo shows how to use [BLIP](https://arxiv.org/abs/2201.12086) to do condit
## Quick Start
```shell
cargo run -r --example blip
cargo run -r -F cuda --example blip -- --device cuda:0 --source images/dog.jpg --source ./assets/bus.jpg --source images/green-car.jpg
```
## Results
```shell
[Unconditional]: a group of people walking around a bus
[Conditional]: three man walking in front of a bus
Some(["three man walking in front of a bus"])
Unconditional: Ys([Y { Texts: [Text("a dog running through a field of grass")] }, Y { Texts: [Text("a group of people walking around a bus")] }, Y { Texts: [Text("a green volkswagen beetle parked in front of a yellow building")] }])
Conditional: Ys([Y { Texts: [Text("this image depicting a dog running in a field")] }, Y { Texts: [Text("this image depict a bus in barcelona")] }, Y { Texts: [Text("this image depict a blue volkswagen beetle parked in a street in havana, cuba")] }])
```
## TODO
* [ ] Multi-batch inference for image caption
* [ ] VQA
* [ ] Retrival
* [ ] TensorRT support for textual model

52
examples/blip/main.rs
View File

@ -1,28 +1,44 @@
use usls::{models::Blip, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// visual
let options_visual = Options::default()
.with_model("blip/visual-base.onnx")?
// .with_ixx(0, 2, 384.into())
// .with_ixx(0, 3, 384.into())
.with_profile(false);
#[derive(argh::FromArgs)]
/// BLIP Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
// textual
let options_textual = Options::default()
.with_model("blip/textual-base.onnx")?
.with_tokenizer("blip/tokenizer.json")?
.with_profile(false);
/// source image
#[argh(option, default = "vec![String::from(\"./assets/bus.jpg\")]")]
source: Vec<String>,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options_visual = Options::blip_v1_base_caption_visual()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let options_textual = Options::blip_v1_base_caption_textual()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = Blip::new(options_visual, options_textual)?;
// image caption (this demo use batch_size=1)
let xs = [DataLoader::try_read("images/bus.jpg")?];
let image_embeddings = model.encode_images(&xs)?;
let _y = model.caption(&image_embeddings, None, true)?; // unconditional
let y = model.caption(&image_embeddings, Some("three man"), true)?; // conditional
println!("{:?}", y[0].texts());
// image caption
let xs = DataLoader::try_read_batch(&args.source)?;
// unconditional caption
let ys = model.forward(&xs, None)?;
println!("Unconditional: {:?}", ys);
// conditional caption
let ys = model.forward(&xs, Some("this image depict"))?;
println!("Conditional: {:?}", ys);
Ok(())
}

13
examples/clip/README.md
View File

@ -3,18 +3,13 @@ This demo showcases how to use [CLIP](https://github.com/openai/CLIP) to compute
## Quick Start
```shell
cargo run -r --example clip
cargo run -r -F cuda --example clip -- --device cuda:0
```
## Results
```shell
(90.11472%) ./examples/clip/images/carrot.jpg => 几个胡萝卜
[0.04573484, 0.0048218793, 0.0011618224, 0.90114725, 0.0036694852, 0.031348046, 0.0121166315]
(94.07785%) ./examples/clip/images/peoples.jpg => Some people holding wine glasses in a restaurant
[0.050406333, 0.0011632168, 0.0019338318, 0.0013227565, 0.003916758, 0.00047858112, 0.9407785]
(86.59852%) ./examples/clip/images/doll.jpg => There is a doll with red hair and a clock on a table
[0.07032883, 0.00053773675, 0.0006372929, 0.06066096, 0.0007378078, 0.8659852, 0.0011121632]
(99.9675%) ./examples/clip/images/carrot.jpg => Some carrots
(99.93718%) ./examples/clip/images/doll.jpg => There is a doll with red hair and a clock on a table
(100.0%) ./examples/clip/images/drink.jpg => Some people holding wine glasses in a restaurant
```

0
examples/clip/images/peoples.jpg → examples/clip/images/drink.jpg
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 176 KiB

After

Width:  |  Height:  |  Size: 176 KiB

58
examples/clip/main.rs
View File

@ -1,43 +1,54 @@
use usls::{models::Clip, DataLoader, Options};
use anyhow::Result;
use usls::{models::Clip, DataLoader, Ops, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// visual
let options_visual = Options::default().with_model("clip/visual-base-dyn.onnx")?;
#[derive(argh::FromArgs)]
/// CLIP Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
// textual
let options_textual = Options::default()
.with_model("clip/textual-base-dyn.onnx")?
.with_tokenizer("clip/tokenizer.json")?;
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options_visual = Options::jina_clip_v1_visual()
// clip_vit_b32_visual()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let options_textual = Options::jina_clip_v1_textual()
// clip_vit_b32_textual()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = Clip::new(options_visual, options_textual)?;
// texts
let texts = vec![
"A photo of a dinosaur ".to_string(),
"A photo of a cat".to_string(),
"A photo of a dog".to_string(),
"几个胡萝卜".to_string(),
"There are some playing cards on a striped table cloth".to_string(),
"There is a doll with red hair and a clock on a table".to_string(),
"Some people holding wine glasses in a restaurant".to_string(),
"A photo of a dinosaur",
"A photo of a cat",
"A photo of a dog",
"Some carrots",
"There are some playing cards on a striped table cloth",
"There is a doll with red hair and a clock on a table",
"Some people holding wine glasses in a restaurant",
];
let feats_text = model.encode_texts(&texts)?; // [n, ndim]
// load image
// load images
let dl = DataLoader::new("./examples/clip/images")?.build()?;
// loop
// run
for (images, paths) in dl {
let feats_image = model.encode_images(&images).unwrap();
let feats_image = model.encode_images(&images)?;
// use image to query texts
let matrix = match feats_image.embedding() {
Some(x) => x.dot2(feats_text.embedding().unwrap())?,
None => continue,
};
let matrix = Ops::dot2(&feats_image, &feats_text)?;
// summary
for i in 0..paths.len() {
let probs = &matrix[i];
let (id, &score) = probs
@ -52,7 +63,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
paths[i].display(),
&texts[id]
);
println!("{:?}\n", probs);
}
}

6
examples/convnext/README.md Normal file
View File

@ -0,0 +1,6 @@
## Quick Start
```shell
cargo run -r -F cuda --example convnext -- --device cuda --dtype fp16
```

52
examples/convnext/main.rs Normal file
View File

@ -0,0 +1,52 @@
use usls::{models::ImageClassifier, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// source image
#[argh(
option,
default = "vec![
String::from(\"images/dog.jpg\"),
String::from(\"images/siamese.png\"),
String::from(\"images/ailurus-fulgens.jpg\"),
]"
)]
source: Vec<String>,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::convnext_v2_atto()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = ImageClassifier::try_from(options)?;
// load images
let xs = DataLoader::try_read_batch(&args.source)?;
// run
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default().with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

5
examples/d-fine/README.md Normal file
View File

@ -0,0 +1,5 @@
## Quick Start
```shell
cargo run -r --example d-fine
```

28
examples/d-fine/main.rs Normal file
View File

@ -0,0 +1,28 @@
use anyhow::Result;
use usls::{models::RTDETR, Annotator, DataLoader, Options};
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// options
let options = Options::d_fine_n_coco().commit()?;
let mut model = RTDETR::new(options)?;
// load
let x = [DataLoader::try_read("./assets/bus.jpg")?];
// run
let y = model.forward(&x)?;
println!("{:?}", y);
// annotate
let annotator = Annotator::default()
.with_bboxes_thickness(3)
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())
}

5
examples/dataloader/README.md Normal file
View File

@ -0,0 +1,5 @@
## Quick Start
```shell
cargo run -r --example dataloader
```

71
examples/dataloader/main.rs
View File

@ -1,66 +1,45 @@
use usls::{
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOTask, YOLOVersion,
};
use usls::DataLoader;
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_max_level(tracing::Level::ERROR)
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let options = Options::new()
.with_device(Device::Cuda(0))
.with_model("yolo/v8-m-dyn.onnx")?
.with_yolo_version(YOLOVersion::V8)
.with_yolo_task(YOLOTask::Detect)
.with_batch(2)
.with_ixx(0, 2, (416, 640, 800).into())
.with_ixx(0, 3, (416, 640, 800).into())
.with_confs(&[0.2]);
let mut model = YOLO::new(options)?;
// build annotator
let annotator = Annotator::new()
.with_bboxes_thickness(4)
.with_saveout("YOLO-DataLoader");
// build dataloader
let dl = DataLoader::new(
// 1. iterator
let dl = DataLoader::try_from(
// "images/bus.jpg", // remote image
// "../images", // image folder
// "../demo.mp4", // local video
// "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // remote video
// "rtsp://admin:xyz@192.168.2.217:554/h265/ch1/", // rtsp h264 stream
// "./assets/bus.jpg", // local image
"../7.mp4",
"./assets/bus.jpg", // local image
)?
.with_batch(1)
.with_progress_bar(true)
.build()?;
let mut viewer = Viewer::new().with_delay(10).with_scale(1.).resizable(true);
// iteration
for (xs, _) in dl {
// inference & annotate
let ys = model.run(&xs)?;
let images_plotted = annotator.plot(&xs, &ys, false)?;
// show image
viewer.imshow(&images_plotted)?;
// check out window and key event
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) {
break;
for (_xs, _paths) in dl {
println!("Paths: {:?}", _paths);
}
// write video
viewer.write_batch(&images_plotted)?;
}
// 2. read one image
let image = DataLoader::try_read("./assets/bus.jpg")?;
println!(
"Read one image. Height: {}, Width: {}",
image.height(),
image.width()
);
// finish video write
viewer.finish_write()?;
// images -> video
// DataLoader::is2v("runs/YOLO-DataLoader", &["runs", "is2v"], 24)?;
// 3. read several images
let images = DataLoader::try_read_batch(&[
"./assets/bus.jpg",
"./assets/bus.jpg",
"./assets/bus.jpg",
"./assets/bus.jpg",
"./assets/bus.jpg",
])?;
println!("Read {} images.", images.len());
Ok(())
}

9
examples/db/README.md
View File

@ -4,15 +4,6 @@
cargo run -r --example db
```
### Speed test
| Model | Image size | TensorRT<br />f16<br />batch=1<br />(ms) | TensorRT<br />f32<br />batch=1<br />(ms) | CUDA<br />f32<br />batch=1<br />(ms) |
| --------------- | ---------- | ---------------------------------------- | ---------------------------------------- | ------------------------------------ |
| ppocr-v3-db-dyn | 640x640 | 1.8585 | 2.5739 | 4.3314 |
| ppocr-v4-db-dyn | 640x640 | 2.0507 | 2.8264 | 6.6064 |
***Test on RTX3060***
## Results
![](https://github.com/jamjamjon/assets/releases/download/db/demo-paper.png)

49
examples/db/main.rs
View File

@ -1,35 +1,48 @@
use anyhow::Result;
use usls::{models::DB, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// build model
let options = Options::default()
.with_ixx(0, 0, (1, 4, 8).into())
.with_ixx(0, 2, (608, 960, 1280).into())
.with_ixx(0, 3, (608, 960, 1280).into())
// .with_trt(0)
.with_confs(&[0.4])
.with_min_width(5.0)
.with_min_height(12.0)
.with_model("db/ppocr-v4-db-dyn.onnx")?;
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::ppocr_det_v4_server_ch()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = DB::new(options)?;
// load image
let x = [
DataLoader::try_read("images/db.png")?,
DataLoader::try_read("images/street.jpg")?,
];
let x = DataLoader::try_read_batch(&[
"images/table.png",
"images/table1.jpg",
"images/table2.png",
"images/table-ch.jpg",
"images/db.png",
"images/street.jpg",
])?;
// run
let y = model.run(&x)?;
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.without_bboxes(true)
.without_mbrs(true)
.with_polygons_alpha(60)
.with_contours_color([255, 105, 180, 255])
.without_mbrs(true)
.with_saveout("DB");
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())

7
examples/deim/README.md Normal file
View File

@ -0,0 +1,7 @@
## Quick Start
```shell
cargo run -r --example deim
```

28
examples/deim/main.rs Normal file
View File

@ -0,0 +1,28 @@
use anyhow::Result;
use usls::{models::RTDETR, Annotator, DataLoader, Options};
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// options
let options = Options::deim_dfine_s_coco().commit()?;
let mut model = RTDETR::new(options)?;
// load
let x = [DataLoader::try_read("./assets/bus.jpg")?];
// run
let y = model.forward(&x)?;
println!("{:?}", y);
// annotate
let annotator = Annotator::default()
.with_bboxes_thickness(3)
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())
}

7
examples/deit/README.md Normal file
View File

@ -0,0 +1,7 @@
## Quick Start
```shell
cargo run -r -F cuda --example deit -- --device cuda --dtype fp16
```

52
examples/deit/main.rs Normal file
View File

@ -0,0 +1,52 @@
use usls::{models::ImageClassifier, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// source image
#[argh(
option,
default = "vec![
String::from(\"images/dog.jpg\"),
String::from(\"images/siamese.png\"),
String::from(\"images/ailurus-fulgens.jpg\"),
]"
)]
source: Vec<String>,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::deit_tiny_distill()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = ImageClassifier::try_from(options)?;
// load images
let xs = DataLoader::try_read_batch(&args.source)?;
// run
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default().with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

20
examples/depth-anything/main.rs
View File

@ -1,24 +1,26 @@
use anyhow::Result;
use usls::{models::DepthAnything, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// options
let options = Options::default()
// .with_model("depth-anything/v1-s-dyn.onnx")?
.with_model("depth-anything/v2-s.onnx")?
.with_ixx(0, 2, (384, 512, 1024).into())
.with_ixx(0, 3, (384, 512, 1024).into());
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// build model
let options = Options::depth_anything_v2_small().commit()?;
let mut model = DepthAnything::new(options)?;
// load
let x = [DataLoader::try_read("images/street.jpg")?];
// run
let y = model.run(&x)?;
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.with_colormap("Turbo")
.with_saveout("Depth-Anything");
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())

10
examples/depth-pro/README.md Normal file
View File

@ -0,0 +1,10 @@
## Quick Start
```shell
cargo run -r -F cuda --example depth-pro -- --device cuda
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/depth-pro/demo-depth-pro.png)

44
examples/depth-pro/main.rs
View File

@ -1,25 +1,47 @@
use anyhow::Result;
use usls::{models::DepthPro, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// options
let options = Options::default()
.with_model("depth-pro/q4f16.onnx")? // bnb4, f16
.with_ixx(0, 0, 1.into()) // batch. Note: now only support batch_size = 1
.with_ixx(0, 1, 3.into()) // channel
.with_ixx(0, 2, 1536.into()) // height
.with_ixx(0, 3, 1536.into()); // width
#[derive(argh::FromArgs)]
/// BLIP Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// dtype
#[argh(option, default = "String::from(\"q4f16\")")]
dtype: String,
/// source image
#[argh(option, default = "String::from(\"images/street.jpg\")")]
source: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// model
let options = Options::depth_pro()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = DepthPro::new(options)?;
// load
let x = [DataLoader::try_read("images/street.jpg")?];
let x = [DataLoader::try_read(&args.source)?];
// run
let y = model.run(&x)?;
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.with_colormap("Turbo")
.with_saveout("Depth-Pro");
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())

51
examples/dinov2/main.rs
View File

@ -1,40 +1,25 @@
use usls::{models::Dinov2, DataLoader, Options};
use anyhow::Result;
use usls::{models::DINOv2, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// build model
let options = Options::default()
.with_model("dinov2/s-dyn.onnx")?
.with_ixx(0, 2, 224.into())
.with_ixx(0, 3, 224.into());
let mut model = Dinov2::new(options)?;
let x = [DataLoader::try_read("images/bus.jpg")?];
let y = model.run(&x)?;
println!("{y:?}");
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// TODO:
// query from vector
// let ys = model.query_from_vec(
// "./assets/bus.jpg",
// &[
// "./examples/dinov2/images/bus.jpg",
// "./examples/dinov2/images/1.jpg",
// "./examples/dinov2/images/2.jpg",
// ],
// Metric::L2,
// )?;
// images
let xs = [
DataLoader::try_read("./assets/bus.jpg")?,
DataLoader::try_read("./assets/bus.jpg")?,
];
// or query from folder
// let ys = model.query_from_folder("./assets/bus.jpg", "./examples/dinov2/images", Metric::IP)?;
// model
let options = Options::dinov2_small().with_batch_size(xs.len()).commit()?;
let mut model = DINOv2::new(options)?;
// results
// for (i, y) in ys.iter().enumerate() {
// println!(
// "Top-{:<3}{:.7} {}",
// i + 1,
// y.1,
// y.2.canonicalize()?.display()
// );
// }
// encode images
let y = model.encode_images(&xs)?;
println!("Feat shape: {:?}", y.shape());
Ok(())
}

10
examples/doclayout-yolo/README.md Normal file
View File

@ -0,0 +1,10 @@
## Quick Start
```shell
cargo run -r -F cuda --example doclayout-yolo -- --device cuda
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/yolo/demo-doclayout-yolo.png)

42
examples/doclayout-yolo/main.rs Normal file
View File

@ -0,0 +1,42 @@
use anyhow::Result;
use usls::{models::YOLO, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let config = Options::doclayout_yolo_docstructbench()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = YOLO::new(config)?;
// load images
let xs = [DataLoader::try_read("images/academic.jpg")?];
// run
let ys = model.forward(&xs)?;
// println!("{:?}", ys);
// annotate
let annotator = Annotator::default()
.with_bboxes_thickness(3)
.with_saveout("doclayout-yolo");
annotator.annotate(&xs, &ys);
model.summary();
Ok(())
}

6
examples/fast/README.md Normal file
View File

@ -0,0 +1,6 @@
## Quick Start
```shell
cargo run -r --example fast
```

65
examples/fast/main.rs Normal file
View File

@ -0,0 +1,65 @@
use anyhow::Result;
use usls::{models::DB, Annotator, DataLoader, Options, Scale};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// scale
#[argh(option, default = "String::from(\"t\")")]
scale: String,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = match args.scale.as_str().try_into()? {
Scale::T => Options::fast_tiny(),
Scale::S => Options::fast_small(),
Scale::B => Options::fast_base(),
_ => unimplemented!("Unsupported model scale: {:?}. Try b, s, t.", args.scale),
};
let mut model = DB::new(
options
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?,
)?;
// load image
let x = DataLoader::try_read_batch(&[
"images/table.png",
"images/table1.jpg",
"images/table2.png",
"images/table-ch.jpg",
"images/db.png",
"images/street.jpg",
])?;
// run
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.without_bboxes(true)
.without_mbrs(true)
.with_polygons_alpha(60)
.with_contours_color([255, 105, 180, 255])
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())
}

5
examples/fastsam/README.md Normal file
View File

@ -0,0 +1,5 @@
## Quick Start
```shell
cargo run -r -F cuda --example fastsam -- --device cuda
```

45
examples/fastsam/main.rs Normal file
View File

@ -0,0 +1,45 @@
use anyhow::Result;
use usls::{models::YOLO, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"fp16\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let config = Options::fastsam_s()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = YOLO::new(config)?;
// load images
let xs = DataLoader::try_read_batch(&["./assets/bus.jpg"])?;
// run
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default()
.without_masks(true)
.with_bboxes_thickness(3)
.with_saveout("fastsam");
annotator.annotate(&xs, &ys);
Ok(())
}

13
examples/fastvit/README.md Normal file
View File

@ -0,0 +1,13 @@
## Quick Start
```shell
cargo run -r -F cuda --example mobileone -- --device cuda --dtype fp16
```
```shell
0: Y { Probs: { Top5: [(263, 0.6109131, Some("Pembroke, Pembroke Welsh corgi")), (264, 0.2062352, Some("Cardigan, Cardigan Welsh corgi")), (231, 0.028572788, Some("collie")), (273, 0.015174894, Some("dingo, warrigal, warragal, Canis dingo")), (248, 0.014367299, Some("Eskimo dog, husky"))] } }
1: Y { Probs: { Top5: [(284, 0.9907692, Some("siamese cat, Siamese")), (285, 0.0015794479, Some("Egyptian cat")), (174, 0.0015189401, Some("Norwegian elkhound, elkhound")), (225, 0.00031838714, Some("malinois")), (17, 0.00027021166, Some("jay"))] } }
2: Y { Probs: { Top5: [(387, 0.94238573, Some("lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens")), (368, 0.0029994072, Some("gibbon, Hylobates lar")), (277, 0.0016564301, Some("red fox, Vulpes vulpes")), (356, 0.0015081967, Some("weasel")), (295, 0.001427932, Some("American black bear, black bear, Ursus americanus, Euarctos americanus"))] } }
```

57
examples/fastvit/main.rs Normal file
View File

@ -0,0 +1,57 @@
use usls::{models::ImageClassifier, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// source image
#[argh(
option,
default = "vec![
String::from(\"images/dog.jpg\"),
String::from(\"images/siamese.png\"),
String::from(\"images/ailurus-fulgens.jpg\"),
]"
)]
source: Vec<String>,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::fastvit_t8_distill()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = ImageClassifier::try_from(options)?;
// load images
let xs = DataLoader::try_read_batch(&args.source)?;
// run
let ys = model.forward(&xs)?;
// results
for (i, y) in ys.iter().enumerate() {
println!("{}: {:?}", i, y);
}
// annotate
let annotator = Annotator::default().with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

30
examples/florence2/README.md Normal file
View File

@ -0,0 +1,30 @@
## Quick Start
```shell
cargo run -r -F cuda --example florence2 -- --device cuda --scale base --dtype fp16
```
```Shell
Task: Caption(0)
Ys([Y { Texts: [Text("A green car parked in front of a yellow building.")] }, Y { Texts: [Text("A group of people walking down a street next to a bus.")] }])
Task: Caption(1)
Ys([Y { Texts: [Text("The image shows a green car parked in front of a yellow building with two brown doors. The car is on the road, and the building has a wall and a tree in the background.")] }, Y { Texts: [Text("The image shows a group of people walking down a street next to a bus, with a building in the background. The bus is likely part of the World Electric Emission Bus, which is a new bus that will be launched in Madrid. The people are walking on the road, and there are trees and a sign board to the left of the bus.")] }])
Task: Caption(2)
Ys([Y { Texts: [Text("The image shows a vintage Volkswagen Beetle car parked on a cobblestone street in front of a yellow building with two wooden doors. The car is a light blue color with silver rims and appears to be in good condition. The building has a sloping roof and is painted in a bright yellow color. The sky is blue and there are trees in the background. The overall mood of the image is peaceful and serene.")] }, Y { Texts: [Text("The image shows a blue and white bus with the logo of the Brazilian football club, Cero Emisiones, on the side. The bus is parked on a street with a building in the background. There are several people walking on the sidewalk in front of the bus, some of them are carrying bags and one person is holding a camera. The sky is blue and there are trees and a traffic light visible in the top right corner of the image. The image appears to be taken during the day.")] }])
```
## Results
| Task | Demo |
| -----| ------|
|Caption-To-Phrase-Grounding | <img src='https://github.com/jamjamjon/assets/releases/download/florence2/Caption-To-Phrase-Grounding-car.png' alt=''> |
| Ocr-With-Region | <img src='https://github.com/jamjamjon/assets/releases/download/florence2/Ocr-With-Region.png' alt=''>|
| Dense-Region-Caption | <img src='https://github.com/jamjamjon/assets/releases/download/florence2/Dense-Region-Caption-car.png' alt=''>|
| Object-Detection | <img src='https://github.com/jamjamjon/assets/releases/download/florence2/Object-Detection-car.png' alt=''>|
| Region-Proposal | <img src='https://github.com/jamjamjon/assets/releases/download/florence2/Region-Proposal.png' alt=''>|
| Referring-Expression-Segmentation | <img src='https://github.com/jamjamjon/assets/releases/download/florence2/Referring-Expression-Segmentation.png' alt=''>|

173
examples/florence2/main.rs
View File

@ -1,116 +1,133 @@
use usls::{models::Florence2, Annotator, DataLoader, Options, Task};
use anyhow::Result;
use usls::{models::Florence2, Annotator, DataLoader, Options, Scale, Task};
fn main() -> Result<(), Box<dyn std::error::Error>> {
let batch_size = 3;
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
// vision encoder
let options_vision_encoder = Options::default()
.with_model("florence2/base-vision-encoder-f16.onnx")?
.with_ixx(0, 2, (512, 768, 800).into())
.with_ixx(0, 3, 768.into())
.with_ixx(0, 0, (1, batch_size as _, 8).into());
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
// text embed
let options_text_embed = Options::default()
.with_model("florence2/base-embed-tokens-f16.onnx")?
.with_tokenizer("florence2/tokenizer.json")?
.with_batch(batch_size);
/// scale
#[argh(option, default = "String::from(\"base\")")]
scale: String,
}
// transformer encoder
let options_encoder = Options::default()
.with_model("florence2/base-encoder-f16.onnx")?
.with_batch(batch_size);
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// transformer decoder
let options_decoder = Options::default()
.with_model("florence2/base-decoder-f16.onnx")?
.with_batch(batch_size);
let args: Args = argh::from_env();
// transformer decoder merged
let options_decoder_merged = Options::default()
.with_model("florence2/base-decoder-merged-f16.onnx")?
.with_batch(batch_size);
// load images
let xs = [
DataLoader::try_read("images/green-car.jpg")?,
DataLoader::try_read("assets/bus.jpg")?,
];
// build model
let mut model = Florence2::new(
let (
options_vision_encoder,
options_text_embed,
options_encoder,
options_decoder,
options_decoder_merged,
) = match args.scale.as_str().try_into()? {
Scale::B => (
Options::florence2_visual_encoder_base(),
Options::florence2_textual_embed_base(),
Options::florence2_texual_encoder_base(),
Options::florence2_texual_decoder_base(),
Options::florence2_texual_decoder_merged_base(),
),
Scale::L => todo!(),
_ => anyhow::bail!("Unsupported Florence2 scale."),
};
let mut model = Florence2::new(
options_vision_encoder
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_text_embed
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_encoder
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_decoder
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_decoder_merged
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
)?;
// load images
let xs = [
// DataLoader::try_read("florence2/car.jpg")?, // for testing region-related tasks
DataLoader::try_read("florence2/car.jpg")?,
// DataLoader::try_read("images/db.png")?,
DataLoader::try_read("assets/bus.jpg")?,
];
// region-related tasks
let quantizer = usls::Quantizer::default();
// let coords = [449., 270., 556., 372.]; // wheel
let coords = [31., 156., 581., 373.]; // car
let (width_car, height_car) = (xs[0].width(), xs[0].height());
let quantized_coords = quantizer.quantize(&coords, (width_car as _, height_car as _));
// run with tasks
let ys = model.run_with_tasks(
&xs,
&[
// w/ inputs
// tasks
let tasks = [
// w inputs
Task::Caption(0),
Task::Caption(1),
Task::Caption(2),
Task::Ocr,
Task::OcrWithRegion,
// Task::OcrWithRegion,
Task::RegionProposal,
Task::ObjectDetection,
Task::DenseRegionCaption,
// w/o inputs
Task::OpenSetDetection("a vehicle".into()),
Task::CaptionToPhraseGrounding(
"A vehicle with two wheels parked in front of a building.".into(),
),
Task::ReferringExpressionSegmentation("a vehicle".into()),
Task::OpenSetDetection("a vehicle"),
Task::CaptionToPhraseGrounding("A vehicle with two wheels parked in front of a building."),
Task::ReferringExpressionSegmentation("a vehicle"),
Task::RegionToSegmentation(
quantized_coords[0],
quantized_coords[1],
quantized_coords[2],
quantized_coords[3],
// 31, 156, 581, 373, // car
449, 270, 556, 372, // wheel
),
Task::RegionToCategory(
quantized_coords[0],
quantized_coords[1],
quantized_coords[2],
quantized_coords[3],
// 31, 156, 581, 373,
449, 270, 556, 372,
),
Task::RegionToDescription(
quantized_coords[0],
quantized_coords[1],
quantized_coords[2],
quantized_coords[3],
// 31, 156, 581, 373,
449, 270, 556, 372,
),
],
)?;
];
// annotator
let annotator = Annotator::new()
.without_bboxes_conf(true)
.with_bboxes_thickness(3)
.with_saveout_subs(&["Florence2"]);
for (task, ys_) in ys.iter() {
// inference
for task in tasks.iter() {
let ys = model.forward(&xs, task)?;
// annotate
match task {
Task::Caption(_)
| Task::Ocr
| Task::RegionToCategory(..)
| Task::RegionToDescription(..) => {
println!("Task: {:?}\n{:?}\n", task, ys_)
println!("Task: {:?}\n{:?}\n", task, &ys)
}
Task::DenseRegionCaption => {
let annotator = annotator.clone().with_saveout("Dense-Region-Caption");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::RegionProposal => {
let annotator = annotator
@ -118,40 +135,42 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
.without_bboxes_name(false)
.with_saveout("Region-Proposal");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::ObjectDetection => {
let annotator = annotator.clone().with_saveout("Object-Detection");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::OpenSetDetection(_) => {
let annotator = annotator.clone().with_saveout("Open-Set-Detection");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::CaptionToPhraseGrounding(_) => {
let annotator = annotator
.clone()
.with_saveout("Caption-To-Phrase-Grounding");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::ReferringExpressionSegmentation(_) => {
let annotator = annotator
.clone()
.with_saveout("Referring-Expression-Segmentation");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::RegionToSegmentation(..) => {
let annotator = annotator.clone().with_saveout("Region-To-Segmentation");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
Task::OcrWithRegion => {
let annotator = annotator.clone().with_saveout("Ocr-With-Region");
annotator.annotate(&xs, ys_);
annotator.annotate(&xs, &ys);
}
_ => (),
}
}
model.summary();
Ok(())
}

2
examples/grounding-dino/README.md
View File

@ -1,7 +1,7 @@
## Quick Start
```shell
cargo run -r --example grounding-dino
cargo run -r -F cuda --example grounding-dino -- --device cuda --dtype fp16
```

95
examples/grounding-dino/main.rs
View File

@ -1,41 +1,72 @@
use anyhow::Result;
use usls::{models::GroundingDINO, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
let opts = Options::default()
.with_ixx(0, 0, (1, 1, 4).into())
.with_ixx(0, 2, (640, 800, 1200).into())
.with_ixx(0, 3, (640, 1200, 1200).into())
// .with_i10((1, 1, 4).into())
// .with_i11((256, 256, 512).into())
// .with_i20((1, 1, 4).into())
// .with_i21((256, 256, 512).into())
// .with_i30((1, 1, 4).into())
// .with_i31((256, 256, 512).into())
// .with_i40((1, 1, 4).into())
// .with_i41((256, 256, 512).into())
// .with_i50((1, 1, 4).into())
// .with_i51((256, 256, 512).into())
// .with_i52((256, 256, 512).into())
.with_model("grounding-dino/swint-ogc-dyn-u8.onnx")? // TODO: current onnx model does not support bs > 1
// .with_model("grounding-dino/swint-ogc-dyn-f32.onnx")?
.with_tokenizer("grounding-dino/tokenizer.json")?
.with_confs(&[0.2])
.with_profile(false);
let mut model = GroundingDINO::new(opts)?;
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
// Load images and set class names
let x = [DataLoader::try_read("images/bus.jpg")?];
let texts = [
"person", "hand", "shoes", "bus", "dog", "cat", "sign", "tie", "monitor", "window",
"glasses", "tree", "head",
];
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
// Run and annotate
let y = model.run(&x, &texts)?;
/// source image
#[argh(option, default = "vec![String::from(\"./assets/bus.jpg\")]")]
source: Vec<String>,
/// open class names
#[argh(
option,
default = "vec![
String::from(\"person\"),
String::from(\"hand\"),
String::from(\"shoes\"),
String::from(\"bus\"),
String::from(\"dog\"),
String::from(\"cat\"),
String::from(\"sign\"),
String::from(\"tie\"),
String::from(\"monitor\"),
String::from(\"glasses\"),
String::from(\"tree\"),
String::from(\"head\"),
]"
)]
labels: Vec<String>,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
let options = Options::grounding_dino_tiny()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_text_names(&args.labels.iter().map(|x| x.as_str()).collect::<Vec<_>>())
.commit()?;
let mut model = GroundingDINO::new(options)?;
// load images
let xs = DataLoader::try_read_batch(&args.source)?;
// run
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default()
.with_bboxes_thickness(4)
.with_saveout("GroundingDINO");
annotator.annotate(&x, &y);
.with_saveout(model.spec());
annotator.annotate(&xs, &ys);
// summary
model.summary();
Ok(())
}

5
examples/hub/README.md Normal file
View File

@ -0,0 +1,5 @@
## Quick Start
```shell
RUST_LOG=usls=info cargo run -r --example hub
```

26
examples/hub/main.rs Normal file
View File

@ -0,0 +1,26 @@
use usls::Hub;
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// 1. Download from default github release
let path = Hub::default().try_fetch("images/bus.jpg")?;
println!("Fetch one image: {:?}", path);
// 2. Download from specific github release url
let path = Hub::default()
.try_fetch("https://github.com/jamjamjon/assets/releases/download/images/bus.jpg")?;
println!("Fetch one file: {:?}", path);
// 3. Fetch tags and files
let hub = Hub::default().with_owner("jamjamjon").with_repo("usls");
for tag in hub.tags().iter() {
let files = hub.files(tag);
println!("{} => {:?}", tag, files); // Should be empty
}
Ok(())
}

6
examples/linknet/README.md Normal file
View File

@ -0,0 +1,6 @@
## Quick Start
```shell
cargo run -r --example fast
```

65
examples/linknet/main.rs Normal file
View File

@ -0,0 +1,65 @@
use anyhow::Result;
use usls::{models::DB, Annotator, DataLoader, Options, Scale};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// scale
#[argh(option, default = "String::from(\"t\")")]
scale: String,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = match args.scale.as_str().try_into()? {
Scale::T => Options::linknet_r18(),
Scale::S => Options::linknet_r34(),
Scale::B => Options::linknet_r50(),
_ => unimplemented!("Unsupported model scale: {:?}. Try b, s, t.", args.scale),
};
let mut model = DB::new(
options
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?,
)?;
// load image
let x = DataLoader::try_read_batch(&[
"images/table.png",
"images/table1.jpg",
"images/table2.png",
"images/table-ch.jpg",
"images/db.png",
"images/street.jpg",
])?;
// run
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.without_bboxes(true)
.without_mbrs(true)
.with_polygons_alpha(60)
.with_contours_color([255, 105, 180, 255])
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())
}

13
examples/mobileone/README.md Normal file
View File

@ -0,0 +1,13 @@
## Quick Start
```shell
cargo run -r -F cuda --example mobileone -- --device cuda --dtype fp16
```
```shell
0: Y { Probs: { Top5: [(263, 0.6109131, Some("Pembroke, Pembroke Welsh corgi")), (264, 0.2062352, Some("Cardigan, Cardigan Welsh corgi")), (231, 0.028572788, Some("collie")), (273, 0.015174894, Some("dingo, warrigal, warragal, Canis dingo")), (248, 0.014367299, Some("Eskimo dog, husky"))] } }
1: Y { Probs: { Top5: [(284, 0.9907692, Some("siamese cat, Siamese")), (285, 0.0015794479, Some("Egyptian cat")), (174, 0.0015189401, Some("Norwegian elkhound, elkhound")), (225, 0.00031838714, Some("malinois")), (17, 0.00027021166, Some("jay"))] } }
2: Y { Probs: { Top5: [(387, 0.94238573, Some("lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens")), (368, 0.0029994072, Some("gibbon, Hylobates lar")), (277, 0.0016564301, Some("red fox, Vulpes vulpes")), (356, 0.0015081967, Some("weasel")), (295, 0.001427932, Some("American black bear, black bear, Ursus americanus, Euarctos americanus"))] } }
```

57
examples/mobileone/main.rs Normal file
View File

@ -0,0 +1,57 @@
use usls::{models::ImageClassifier, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// source image
#[argh(
option,
default = "vec![
String::from(\"images/dog.jpg\"),
String::from(\"images/siamese.png\"),
String::from(\"images/ailurus-fulgens.jpg\"),
]"
)]
source: Vec<String>,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::mobileone_s0()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = ImageClassifier::try_from(options)?;
// load images
let xs = DataLoader::try_read_batch(&args.source)?;
// run
let ys = model.forward(&xs)?;
// results
for (i, y) in ys.iter().enumerate() {
println!("{}: {:?}", i, y);
}
// annotate
let annotator = Annotator::default().with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

20
examples/modnet/main.rs
View File

@ -1,22 +1,24 @@
use usls::{models::MODNet, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// build model
let options = Options::default()
.with_model("modnet/dyn-f32.onnx")?
.with_ixx(0, 2, (416, 512, 800).into())
.with_ixx(0, 3, (416, 512, 800).into());
let options = Options::modnet_photographic().commit()?;
let mut model = MODNet::new(options)?;
// load image
let x = [DataLoader::try_read("images/liuyifei.png")?];
let xs = [DataLoader::try_read("images/liuyifei.png")?];
// run
let y = model.run(&x)?;
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default().with_saveout("MODNet");
annotator.annotate(&x, &y);
let annotator = Annotator::default().with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

10
examples/picodet-layout/README.md Normal file
View File

@ -0,0 +1,10 @@
## Quick Start
```shell
cargo run -r --example picodet-layout
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/picodet/demo-layout-1x.png)

31
examples/picodet-layout/main.rs Normal file
View File

@ -0,0 +1,31 @@
use anyhow::Result;
use usls::{models::PicoDet, Annotator, DataLoader, Options};
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// options
let options = Options::picodet_layout_1x()
// picodet_l_layout_3cls()
// picodet_l_layout_17cls()
.commit()?;
let mut model = PicoDet::new(options)?;
// load
let xs = [DataLoader::try_read("images/academic.jpg")?];
// annotator
let annotator = Annotator::default()
.with_bboxes_thickness(3)
.with_saveout(model.spec());
// run
let ys = model.forward(&xs)?;
println!("{:?}", ys);
annotator.annotate(&xs, &ys);
Ok(())
}

17
examples/rtdetr/README.md Normal file
View File

@ -0,0 +1,17 @@
## Quick Start
```shell
cargo run -r --example rtdetr
```
## Results
```
[Bboxes]: Found 5 objects
0: Bbox { xyxy: [47.969677, 397.81808, 246.22426, 904.8823], class_id: 0, name: Some("person"), confidence: 0.94432133 }
1: Bbox { xyxy: [668.0796, 399.28854, 810.3779, 880.7412], class_id: 0, name: Some("person"), confidence: 0.93386495 }
2: Bbox { xyxy: [20.852705, 229.30482, 807.43494, 729.51196], class_id: 5, name: Some("bus"), confidence: 0.9319465 }
3: Bbox { xyxy: [223.28226, 405.37265, 343.92603, 859.50366], class_id: 0, name: Some("person"), confidence: 0.9130827 }
4: Bbox { xyxy: [0.0, 552.6165, 65.99908, 868.00525], class_id: 0, name: Some("person"), confidence: 0.7910869 }
```

43
examples/rtdetr/main.rs Normal file
View File

@ -0,0 +1,43 @@
use anyhow::Result;
use usls::{models::RTDETR, Annotator, DataLoader, Options};
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// options
let options = Options::rtdetr_v2_s_coco()
// rtdetr_v1_r18vd_coco()
// rtdetr_v2_ms_coco()
// rtdetr_v2_m_coco()
// rtdetr_v2_l_coco()
// rtdetr_v2_x_coco()
.commit()?;
let mut model = RTDETR::new(options)?;
// load
let xs = [DataLoader::try_read("./assets/bus.jpg")?];
// run
let ys = model.forward(&xs)?;
// extract bboxes
for y in ys.iter() {
if let Some(bboxes) = y.bboxes() {
println!("[Bboxes]: Found {} objects", bboxes.len());
for (i, bbox) in bboxes.iter().enumerate() {
println!("{}: {:?}", i, bbox)
}
}
}
// annotate
let annotator = Annotator::default()
.with_bboxes_thickness(3)
.with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

23
examples/rtmo/main.rs
View File

@ -1,25 +1,26 @@
use anyhow::Result;
use usls::{models::RTMO, Annotator, DataLoader, Options, COCO_SKELETONS_16};
fn main() -> Result<(), Box<dyn std::error::Error>> {
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// build model
let options = Options::default()
.with_model("rtmo/s-dyn.onnx")?
.with_nk(17)
.with_confs(&[0.3])
.with_kconfs(&[0.5]);
let mut model = RTMO::new(options)?;
let mut model = RTMO::new(Options::rtmo_s().commit()?)?;
// load image
let x = [DataLoader::try_read("images/bus.jpg")?];
let xs = [DataLoader::try_read("images/bus.jpg")?];
// run
let y = model.run(&x)?;
let ys = model.forward(&xs)?;
// annotate
let annotator = Annotator::default()
.with_saveout("RTMO")
.with_saveout(model.spec())
.with_skeletons(&COCO_SKELETONS_16);
annotator.annotate(&x, &y);
annotator.annotate(&xs, &ys);
Ok(())
}

9
examples/sam/README.md
View File

@ -3,19 +3,18 @@
```Shell
# SAM
cargo run -r --example sam
cargo run -r -F cuda --example sam -- --device cuda --kind sam
# MobileSAM
cargo run -r --example sam -- --kind mobile-sam
cargo run -r -F cuda --example sam -- --device cuda --kind mobile-sam
# EdgeSAM
cargo run -r --example sam -- --kind edge-sam
cargo run -r -F cuda --example sam -- --device cuda --kind edge-sam
# SAM-HQ
cargo run -r --example sam -- --kind sam-hq
cargo run -r -F cuda --example sam -- --device cuda --kind sam-hq
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/sam/demo-car.png)

136
examples/sam/main.rs
View File

@ -1,97 +1,73 @@
use clap::Parser;
use anyhow::Result;
use usls::{
models::{SamKind, SamPrompt, SAM},
Annotator, DataLoader, Options,
Annotator, DataLoader, Options, Scale,
};
#[derive(Parser)]
#[command(author, version, about, long_about = None)]
pub struct Args {
#[arg(long, value_enum, default_value_t = SamKind::Sam)]
pub kind: SamKind,
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
#[arg(long, default_value_t = 0)]
pub device_id: usize,
/// scale
#[argh(option, default = "String::from(\"t\")")]
scale: String,
#[arg(long)]
pub use_low_res_mask: bool,
/// SAM kind
#[argh(option, default = "String::from(\"sam\")")]
kind: String,
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = Args::parse();
// Options
let (options_encoder, options_decoder, saveout) = match args.kind {
SamKind::Sam => {
let options_encoder = Options::default()
// .with_model("sam/sam-vit-b-encoder.onnx")?;
.with_model("sam/sam-vit-b-encoder-u8.onnx")?;
let options_decoder = Options::default()
.with_sam_kind(SamKind::Sam)
// .with_model("sam/sam-vit-b-decoder.onnx")?;
// .with_model("sam/sam-vit-b-decoder-singlemask.onnx")?;
.with_model("sam/sam-vit-b-decoder-u8.onnx")?;
(options_encoder, options_decoder, "SAM")
}
SamKind::Sam2 => {
let options_encoder = Options::default()
// .with_model("sam/sam2-hiera-tiny-encoder.onnx")?;
// .with_model("sam/sam2-hiera-small-encoder.onnx")?;
.with_model("sam/sam2-hiera-base-plus-encoder.onnx")?;
let options_decoder = Options::default()
.with_sam_kind(SamKind::Sam2)
// .with_model("sam/sam2-hiera-tiny-decoder.onnx")?;
// .with_model("sam/sam2-hiera-small-decoder.onnx")?;
.with_model("sam/sam2-hiera-base-plus-decoder.onnx")?;
(options_encoder, options_decoder, "SAM2")
}
SamKind::MobileSam => {
let options_encoder =
Options::default().with_model("sam/mobile-sam-vit-t-encoder.onnx")?;
let options_decoder = Options::default()
.with_sam_kind(SamKind::MobileSam)
.with_model("sam/mobile-sam-vit-t-decoder.onnx")?;
(options_encoder, options_decoder, "Mobile-SAM")
}
SamKind::SamHq => {
let options_encoder = Options::default().with_model("sam/sam-hq-vit-t-encoder.onnx")?;
let options_decoder = Options::default()
.with_sam_kind(SamKind::SamHq)
.with_model("sam/sam-hq-vit-t-decoder.onnx")?;
(options_encoder, options_decoder, "SAM-HQ")
}
SamKind::EdgeSam => {
let options_encoder = Options::default().with_model("sam/edge-sam-3x-encoder.onnx")?;
let options_decoder = Options::default()
.with_sam_kind(SamKind::EdgeSam)
.with_model("sam/edge-sam-3x-decoder.onnx")?;
(options_encoder, options_decoder, "Edge-SAM")
}
};
let options_encoder = options_encoder
.with_cuda(args.device_id)
.with_ixx(0, 2, (800, 1024, 1024).into())
.with_ixx(0, 3, (800, 1024, 1024).into());
let options_decoder = options_decoder
.with_cuda(args.device_id)
.use_low_res_mask(args.use_low_res_mask)
.with_find_contours(true);
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// Build model
let (options_encoder, options_decoder) = match args.kind.as_str().try_into()? {
SamKind::Sam => (
Options::sam_v1_base_encoder(),
Options::sam_v1_base_decoder(),
),
SamKind::Sam2 => match args.scale.as_str().try_into()? {
Scale::T => (Options::sam2_tiny_encoder(), Options::sam2_tiny_decoder()),
Scale::S => (Options::sam2_small_encoder(), Options::sam2_small_decoder()),
Scale::B => (
Options::sam2_base_plus_encoder(),
Options::sam2_base_plus_decoder(),
),
_ => unimplemented!("Unsupported model scale: {:?}. Try b, s, t.", args.scale),
},
SamKind::MobileSam => (
Options::mobile_sam_tiny_encoder(),
Options::mobile_sam_tiny_decoder(),
),
SamKind::SamHq => (
Options::sam_hq_tiny_encoder(),
Options::sam_hq_tiny_decoder(),
),
SamKind::EdgeSam => (
Options::edge_sam_3x_encoder(),
Options::edge_sam_3x_decoder(),
),
};
let options_encoder = options_encoder
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let options_decoder = options_decoder.commit()?;
let mut model = SAM::new(options_encoder, options_decoder)?;
// Load image
let xs = [
DataLoader::try_read("images/truck.jpg")?,
// DataLoader::try_read("images/dog.jpg")?,
];
let xs = [DataLoader::try_read("images/truck.jpg")?];
// Build annotator
let annotator = Annotator::default().with_saveout(saveout);
let annotator = Annotator::default().with_saveout(model.spec());
// Prompt
let prompts = vec![
@ -102,7 +78,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
];
// Run & Annotate
let ys = model.run(&xs, &prompts)?;
let ys = model.forward(&xs, &prompts)?;
annotator.annotate(&xs, &ys);
Ok(())

3
examples/sapiens/README.md
View File

@ -1,10 +1,9 @@
## Quick Start
```shell
cargo run -r --example sapiens
cargo run -r -F cuda --example sapiens -- --device cuda
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/sapiens/demo.png)

35
examples/sapiens/main.rs
View File

@ -1,27 +1,38 @@
use usls::{
models::{Sapiens, SapiensTask},
Annotator, DataLoader, Options, BODY_PARTS_28,
};
use anyhow::Result;
use usls::{models::Sapiens, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build
let options = Options::default()
.with_model("sapiens/seg-0.3b-dyn.onnx")?
.with_sapiens_task(SapiensTask::Seg)
.with_names(&BODY_PARTS_28);
let options = Options::sapiens_seg_0_3b()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = Sapiens::new(options)?;
// load
let x = [DataLoader::try_read("images/paul-george.jpg")?];
// run
let y = model.run(&x)?;
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.without_masks(true)
.with_polygons_name(false)
.with_saveout("Sapiens");
.with_polygons_name(true)
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())

9
examples/slanet/README.md Normal file
View File

@ -0,0 +1,9 @@
## Quick Start
```shell
cargo run -r -F cuda --example slanet -- --device cuda
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/slanet/demo.png)

48
examples/slanet/main.rs Normal file
View File

@ -0,0 +1,48 @@
use anyhow::Result;
use usls::{models::SLANet, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// source
#[argh(option, default = "String::from(\"images/table.png\")")]
source: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::slanet_lcnet_v2_mobile_ch()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = SLANet::new(options)?;
// load
let xs = DataLoader::try_read_batch(&[args.source])?;
// run
let ys = model.forward(&xs)?;
println!("{:?}", ys);
// annotate
let annotator = Annotator::default()
.with_keypoints_radius(2)
.with_skeletons(&[(0, 1), (1, 2), (2, 3), (3, 0)])
.with_saveout(model.spec());
annotator.annotate(&xs, &ys);
// summary
model.summary();
Ok(())
}

32
examples/svtr/README.md
View File

@ -1,29 +1,21 @@
## Quick Start
```shell
cargo run -r --example svtr
cargo run -r -F cuda --example svtr -- --device cuda
```
### Speed test
| Model | Width | TensorRT<br />f16<br />batch=1<br />(ms) | TensorRT<br />f32<br />batch=1<br />(ms) | CUDA<br />f32<br />batch=1<br />(ms) |
| --------------------------- | :---: | :--------------------------------------: | :--------------------------------------: | :----------------------------------: |
| ppocr-v4-server-svtr-ch-dyn | 1500 | 4.2116 | 13.0013 | 20.8673 |
| ppocr-v4-svtr-ch-dyn | 1500 | 2.0435 | 3.1959 | 10.1750 |
| ppocr-v3-svtr-ch-dyn | 1500 | 1.8596 | 2.9401 | 6.8210 |
***Test on RTX3060***
## Results
```shell
["./examples/svtr/images/5.png"]: Some(["are closely jointed. Some examples are illustrated in Fig.7."])
["./examples/svtr/images/6.png"]: Some(["小菊儿胡同71号"])
["./examples/svtr/images/4.png"]: Some(["我在南锣鼓捣猫呢"])
["./examples/svtr/images/1.png"]: Some(["你有这么高速运转的机械进入中国，记住我给出的原理"])
["./examples/svtr/images/2.png"]: Some(["冀B6G000"])
["./examples/svtr/images/9.png"]: Some(["from the background, but also separate text instances which"])
["./examples/svtr/images/8.png"]: Some(["110022345"])
["./examples/svtr/images/3.png"]: Some(["粤A·68688"])
["./examples/svtr/images/7.png"]: Some(["Please lower your volume"])
["./examples/svtr/images/license-ch-2.png"]: Ys([Y { Texts: [Text("粤A·68688")] }])
["./examples/svtr/images/license-ch.png"]: Ys([Y { Texts: [Text("冀B6G000")] }])
["./examples/svtr/images/sign-ch-2.png"]: Ys([Y { Texts: [Text("我在南锣鼓捣猫呢")] }])
["./examples/svtr/images/sign-ch.png"]: Ys([Y { Texts: [Text("小菊儿胡同71号")] }])
["./examples/svtr/images/text-110022345.png"]: Ys([Y { Texts: [Text("110022345")] }])
["./examples/svtr/images/text-ch.png"]: Ys([Y { Texts: [Text("你有这么高速运转的机械进入中国，记住我给出的原理")] }])
["./examples/svtr/images/text-en-2.png"]: Ys([Y { Texts: [Text("from the background, but also separate text instances which")] }])
["./examples/svtr/images/text-en-dark.png"]: Ys([Y { Texts: [Text("Please lower your volume")] }])
["./examples/svtr/images/text-en.png"]: Ys([Y { Texts: [Text("are closely jointed. Some examples are illustrated in Fig.7.")] }])
["./examples/svtr/images/text-hello-rust-handwritten.png"]: Ys([Y { Texts: [Text("HeloRuSt")] }])
```

0
examples/svtr/images/3.png → examples/svtr/images/license-ch-2.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 59 KiB

After

Width:  |  Height:  |  Size: 59 KiB

0
examples/svtr/images/2.png → examples/svtr/images/license-ch.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 13 KiB

After

Width:  |  Height:  |  Size: 13 KiB

0
examples/svtr/images/4.png → examples/svtr/images/sign-ch-2.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 15 KiB

After

Width:  |  Height:  |  Size: 15 KiB

0
examples/svtr/images/6.png → examples/svtr/images/sign-ch.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 10 KiB

After

Width:  |  Height:  |  Size: 10 KiB

0
examples/svtr/images/8.png → examples/svtr/images/text-110022345.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 24 KiB

After

Width:  |  Height:  |  Size: 24 KiB

0
examples/svtr/images/1.png → examples/svtr/images/text-ch.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 14 KiB

After

Width:  |  Height:  |  Size: 14 KiB

0
examples/svtr/images/9.png → examples/svtr/images/text-en-2.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 9.0 KiB

After

Width:  |  Height:  |  Size: 9.0 KiB

0
examples/svtr/images/7.png → examples/svtr/images/text-en-dark.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 13 KiB

After

Width:  |  Height:  |  Size: 13 KiB

0
examples/svtr/images/5.png → examples/svtr/images/text-en.png
View File

Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 17 KiB

After

Width:  |  Height:  |  Size: 17 KiB

BIN
examples/svtr/images/text-hello-rust-handwritten.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 3.7 KiB

42
examples/svtr/main.rs
View File

@ -1,24 +1,44 @@
use anyhow::Result;
use usls::{models::SVTR, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::default()
.with_ixx(0, 0, (1, 2, 8).into())
.with_ixx(0, 2, (320, 960, 1600).into())
.with_ixx(0, 3, (320, 960, 1600).into())
.with_confs(&[0.2])
.with_vocab("svtr/ppocr_rec_vocab.txt")?
.with_model("svtr/ppocr-v4-svtr-ch-dyn.onnx")?;
let options = Options::ppocr_rec_v4_ch()
// svtr_v2_teacher_ch()
// .with_batch_size(2)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = SVTR::new(options)?;
// load images
let dl = DataLoader::new("./examples/svtr/images")?.build()?;
let dl = DataLoader::new("./examples/svtr/images")?
.with_batch(model.batch() as _)
.with_progress_bar(false)
.build()?;
// run
for (xs, paths) in dl {
let ys = model.run(&xs)?;
println!("{paths:?}: {:?}", ys[0].texts())
let ys = model.forward(&xs)?;
println!("{paths:?}: {:?}", ys)
}
//summary
model.summary();
Ok(())
}

13
examples/trocr/README.md Normal file
View File

@ -0,0 +1,13 @@
## Quick Start
```shell
cargo run -r -F cuda --example trocr -- --device cuda --dtype fp16 --scale s --kind printed
cargo run -r -F cuda --example trocr -- --device cuda --dtype fp16 --scale s --kind hand-written
```
```shell
Ys([Y { Texts: [Text("PLEASE LOWER YOUR VOLUME")] }, Y { Texts: [Text("HELLO RUST")] }])
```

96
examples/trocr/main.rs Normal file
View File

@ -0,0 +1,96 @@
use usls::{
models::{TrOCR, TrOCRKind},
DataLoader, Options, Scale,
};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// scale
#[argh(option, default = "String::from(\"s\")")]
scale: String,
/// kind
#[argh(option, default = "String::from(\"printed\")")]
kind: String,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// load images
let xs = DataLoader::try_read_batch(&[
"images/text-en-dark.png",
"images/text-hello-rust-handwritten.png",
])?;
// build model
let (options_encoder, options_decoder, options_decoder_merged) =
match args.scale.as_str().try_into()? {
Scale::S => match args.kind.as_str().try_into()? {
TrOCRKind::Printed => (
Options::trocr_encoder_small_printed(),
Options::trocr_decoder_small_printed(),
Options::trocr_decoder_merged_small_printed(),
),
TrOCRKind::HandWritten => (
Options::trocr_encoder_small_handwritten(),
Options::trocr_decoder_small_handwritten(),
Options::trocr_decoder_merged_small_handwritten(),
),
},
Scale::B => match args.kind.as_str().try_into()? {
TrOCRKind::Printed => (
Options::trocr_encoder_base_printed(),
Options::trocr_decoder_base_printed(),
Options::trocr_decoder_merged_base_printed(),
),
TrOCRKind::HandWritten => (
Options::trocr_encoder_base_handwritten(),
Options::trocr_decoder_base_handwritten(),
Options::trocr_decoder_merged_base_handwritten(),
),
},
x => anyhow::bail!("Unsupported TrOCR scale: {:?}", x),
};
let mut model = TrOCR::new(
options_encoder
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_decoder
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_decoder_merged
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
)?;
// inference
let ys = model.forward(&xs)?;
println!("{:?}", ys);
// summary
model.summary();
Ok(())
}

5
examples/viewer/README.md Normal file
View File

@ -0,0 +1,5 @@
## Quick Start
```shell
RUST_LOG=usls=info cargo run -F ffmpeg -r --example viewer
```

43
examples/viewer/main.rs Normal file
View File

@ -0,0 +1,43 @@
use usls::{DataLoader, Key, Viewer};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// source
#[argh(
option,
default = "String::from(\"http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4\")"
)]
source: String,
}
fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
let dl = DataLoader::new(&args.source)?.with_batch(1).build()?;
let mut viewer = Viewer::new().with_delay(5).with_scale(1.).resizable(true);
// run & annotate
for (xs, _paths) in dl {
// show image
viewer.imshow(&xs)?;
// check out window and key event
if !viewer.is_open() || viewer.is_key_pressed(Key::Escape) {
break;
}
// write video
viewer.write_batch(&xs)?
}
// finish video write
viewer.finish_write()?;
Ok(())
}

2
examples/yolo-sam/README.md
View File

@ -1,7 +1,7 @@
## Quick Start
```shell
cargo run -r --example yolo-sam
cargo run -r -F cuda --example yolo-sam -- --device cuda
```
## Results

55
examples/yolo-sam/main.rs
View File

@ -1,31 +1,42 @@
use anyhow::Result;
use usls::{
models::{SamKind, SamPrompt, YOLOTask, YOLOVersion, SAM, YOLO},
Annotator, DataLoader, Options, Vision,
models::{SamPrompt, SAM, YOLO},
Annotator, DataLoader, Options, Scale,
};
fn main() -> Result<(), Box<dyn std::error::Error>> {
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build SAM
let options_encoder = Options::default().with_model("sam/mobile-sam-vit-t-encoder.onnx")?;
let options_decoder = Options::default()
.with_find_contours(true)
.with_sam_kind(SamKind::Sam)
.with_model("sam/mobile-sam-vit-t-decoder.onnx")?;
let (options_encoder, options_decoder) = (
Options::mobile_sam_tiny_encoder().commit()?,
Options::mobile_sam_tiny_decoder().commit()?,
);
let mut sam = SAM::new(options_encoder, options_decoder)?;
// build YOLOv8-Det
let options_yolo = Options::default()
.with_yolo_version(YOLOVersion::V8)
.with_yolo_task(YOLOTask::Detect)
.with_model("yolo/v8-m-dyn.onnx")?
.with_cuda(0)
.with_ixx(0, 2, (416, 640, 800).into())
.with_ixx(0, 3, (416, 640, 800).into())
.with_find_contours(false)
.with_confs(&[0.45]);
// build YOLOv8
let options_yolo = Options::yolo_detect()
.with_model_scale(Scale::N)
.with_model_version(8.0.into())
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut yolo = YOLO::new(options_yolo)?;
// load one image
let xs = [DataLoader::try_read("images/dog.jpg")?];
let xs = DataLoader::try_read_batch(&["images/dog.jpg"])?;
// build annotator
let annotator = Annotator::default()
@ -36,11 +47,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
.with_saveout("YOLO-SAM");
// run & annotate
let ys_det = yolo.run(&xs)?;
for y_det in ys_det {
let ys_det = yolo.forward(&xs)?;
for y_det in ys_det.iter() {
if let Some(bboxes) = y_det.bboxes() {
for bbox in bboxes {
let ys_sam = sam.run(
let ys_sam = sam.forward(
&xs,
&[SamPrompt::default().with_bbox(
bbox.xmin(),

182
examples/yolo/README.md
View File

@ -1,175 +1,65 @@
<h1 align='center'>YOLO-Series</h1>
| Detection | Instance Segmentation | Pose |
| :---------------: | :------------------------: |:---------------: |
| <img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-det.png' width="300px"> | <img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-seg.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-pose.png' width="300px"> |
| :----------------------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------------------------------: | :-----------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-det.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-seg.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-pose.png' width="300px">` |
| Classification | Obb |
| :------------------------: |:------------------------: |
|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-cls.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-obb-2.png' width="628px">
| :----------------------------------------------------------------------------------------------------: | :------------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-cls.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-obb-2.png' width="628px">` |
| Head Detection | Fall Detection | Trash Detection |
| :------------------------: |:------------------------: |:------------------------: |
|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-head.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-falldown.png' width="300px">|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-trash.png' width="300px">
| :-----------------------------------------------------------------------------------------------------: | :---------------------------------------------------------------------------------------------------------: | :------------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-head.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-falldown.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-trash.png' width="300px">` |
| YOLO-World | Face Parsing | FastSAM |
| :------------------------: |:------------------------: |:------------------------: |
|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-yolov8-world.png' width="300px"> |<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-face-parsing.png' width="300px">|<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-fastsam.png' width="300px">
| :-------------------------------------------------------------------------------------------------------------: | :-------------------------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------------------------: |
| `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-yolov8-world.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-face-parsing.png' width="300px">` | `<img src='https://github.com/jamjamjon/assets/releases/download/yolo/demo-fastsam.png' width="300px">` |
## Quick Start
```Shell
# customized
cargo run -r --example yolo -- --task detect --ver v8 --nc 6 --model xxx.onnx # YOLOv8
# Your customized YOLOv8 model
cargo run -r --example yolo -- --task detect --ver v8 --num-classes 6 --model xxx.onnx # YOLOv8
# Classify
cargo run -r --example yolo -- --task classify --ver v5 --scale s --width 224 --height 224 --nc 1000 # YOLOv5
cargo run -r --example yolo -- --task classify --ver v8 --scale n --width 224 --height 224 --nc 1000 # YOLOv8
cargo run -r --example yolo -- --task classify --ver v11 --scale n --width 224 --height 224 --nc 1000 # YOLOv11
cargo run -r --example yolo -- --task classify --ver 5 --scale s --image-width 224 --image-height 224 --num-classes 1000 --use-imagenet-1k-classes # YOLOv5
cargo run -r --example yolo -- --task classify --ver 8 --scale n --image-width 224 --image-height 224 # YOLOv8
cargo run -r --example yolo -- --task classify --ver 11 --scale n --image-width 224 --image-height 224 # YOLOv11
# Detect
cargo run -r --example yolo -- --task detect --ver v5 --scale n # YOLOv5
cargo run -r --example yolo -- --task detect --ver v6 --scale n # YOLOv6
cargo run -r --example yolo -- --task detect --ver v7 --scale t # YOLOv7
cargo run -r --example yolo -- --task detect --ver v8 --scale n # YOLOv8
cargo run -r --example yolo -- --task detect --ver v9 --scale t # YOLOv9
cargo run -r --example yolo -- --task detect --ver v10 --scale n # YOLOv10
cargo run -r --example yolo -- --task detect --ver v11 --scale n # YOLOv11
cargo run -r --example yolo -- --task detect --ver rtdetr --scale l # RTDETR
cargo run -r --example yolo -- --task detect --ver v8 --model yolo/v8-s-world-v2-shoes.onnx # YOLOv8-world
cargo run -r --example yolo -- --task detect --ver 5 --scale n --use-coco-80-classes # YOLOv5
cargo run -r --example yolo -- --task detect --ver 6 --scale n --use-coco-80-classes # YOLOv6
cargo run -r --example yolo -- --task detect --ver 7 --scale t --use-coco-80-classes # YOLOv7
cargo run -r --example yolo -- --task detect --ver 8 --scale n --use-coco-80-classes # YOLOv8
cargo run -r --example yolo -- --task detect --ver 9 --scale t --use-coco-80-classes # YOLOv9
cargo run -r --example yolo -- --task detect --ver 10 --scale n --use-coco-80-classes # YOLOv10
cargo run -r --example yolo -- --task detect --ver 11 --scale n --use-coco-80-classes # YOLOv11
cargo run -r --example yolo -- --task detect --ver 8 --model v8-s-world-v2-shoes.onnx # YOLOv8-world
# Pose
cargo run -r --example yolo -- --task pose --ver v8 --scale n # YOLOv8-Pose
cargo run -r --example yolo -- --task pose --ver v11 --scale n # YOLOv11-Pose
cargo run -r --example yolo -- --task pose --ver 8 --scale n # YOLOv8-Pose
cargo run -r --example yolo -- --task pose --ver 11 --scale n # YOLOv11-Pose
# Segment
cargo run -r --example yolo -- --task segment --ver v5 --scale n # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver v8 --scale n # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver v11 --scale n # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver v8 --model yolo/FastSAM-s-dyn-f16.onnx # FastSAM
cargo run -r --example yolo -- --task segment --ver 5 --scale n # YOLOv5-Segment
cargo run -r --example yolo -- --task segment --ver 8 --scale n # YOLOv8-Segment
cargo run -r --example yolo -- --task segment --ver 11 --scale n # YOLOv8-Segment
# Obb
cargo run -r --example yolo -- --ver v8 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv8-Obb
cargo run -r --example yolo -- --ver v11 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv11-Obb
cargo run -r --example yolo -- --ver 8 --task obb --scale n --image-width 1024 --image-height 1024 --source images/dota.png # YOLOv8-Obb
cargo run -r --example yolo -- --ver 11 --task obb --scale n --image-width 1024 --image-height 1024 --source images/dota.png # YOLOv11-Obb
```
**`cargo run -r --example yolo -- --help` for more options**
## YOLOs configs with `Options`
<details open>
<summary>Use official YOLO Models</summary>
```Rust
let options = Options::default()
.with_yolo_version(YOLOVersion::V5) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
.with_yolo_task(YOLOTask::Classify) // YOLOTask: Classify, Detect, Pose, Segment, Obb
.with_model("xxxx.onnx")?;
```
</details>
<details open>
<summary>Cutomized your own YOLO model</summary>
```Rust
// This config is for YOLOv8-Segment
use usls::{AnchorsPosition, BoxType, ClssType, YOLOPreds};
let options = Options::default()
.with_yolo_preds(
YOLOPreds {
bbox: Some(BoxType::Cxcywh),
clss: ClssType::Clss,
coefs: Some(true),
anchors: Some(AnchorsPosition::After),
..Default::default()
}
)
// .with_nc(80)
// .with_names(&COCO_CLASS_NAMES_80)
.with_model("xxxx.onnx")?;
```
</details>
## Other YOLOv8 Solution Models
| Model | Weights | Datasets|
|:---------------------: | :--------------------------: | :-------------------------------: |
| Face-Landmark Detection | [yolov8-face-dyn-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-n-face-dyn-f16.onnx) | |
| Head Detection | [yolov8-head-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-head-f16.onnx) | |
| Fall Detection | [yolov8-falldown-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-falldown-f16.onnx) | |
| Trash Detection | [yolov8-plastic-bag-f16](https://github.com/jamjamjon/assets/releases/download/yolo/v8-plastic-bag-f16.onnx) | |
| FaceParsing | [yolov8-face-parsing-dyn](https://github.com/jamjamjon/assets/releases/download/yolo/v8-face-parsing-dyn.onnx) | [CelebAMask-HQ](https://github.com/switchablenorms/CelebAMask-HQ/tree/master/face_parsing)<br />[[Processed YOLO labels]](https://github.com/jamjamjon/assets/releases/download/yolo/CelebAMask-HQ-YOLO-Labels.zip)[[Python Script]](../../scripts/CelebAMask-HQ-To-YOLO-Labels.py) |
## Export ONNX Models
<details close>
<summary>YOLOv5</summary>
[Here](https://docs.ultralytics.com/yolov5/tutorials/model_export/)
</details>
<details close>
<summary>YOLOv6</summary>
[Here](https://github.com/meituan/YOLOv6/tree/main/deploy/ONNX)
</details>
<details close>
<summary>YOLOv7</summary>
[Here](https://github.com/WongKinYiu/yolov7?tab=readme-ov-file#export)
</details>
<details close>
<summary>YOLOv8, YOLOv11</summary>
```Shell
pip install -U ultralytics
# export onnx model with dynamic shapes
yolo export model=yolov8m.pt format=onnx simplify dynamic
yolo export model=yolov8m-cls.pt format=onnx simplify dynamic
yolo export model=yolov8m-pose.pt format=onnx simplify dynamic
yolo export model=yolov8m-seg.pt format=onnx simplify dynamic
yolo export model=yolov8m-obb.pt format=onnx simplify dynamic
# export onnx model with fixed shapes
yolo export model=yolov8m.pt format=onnx simplify
yolo export model=yolov8m-cls.pt format=onnx simplify
yolo export model=yolov8m-pose.pt format=onnx simplify
yolo export model=yolov8m-seg.pt format=onnx simplify
yolo export model=yolov8m-obb.pt format=onnx simplify
```
</details>
<details close>
<summary>YOLOv9</summary>
[Here](https://github.com/WongKinYiu/yolov9/blob/main/export.py)
</details>
<details close>
<summary>YOLOv10</summary>
[Here](https://github.com/THU-MIG/yolov10#export)
</details>
| Model | Weights |
| :---------------------: | :------------------------------------------------------: |
| Face-Landmark Detection | [yolov8-n-face](https://github.com/jamjamjon/assets/releases/download/yolo/v8-n-face-fp16.onnx) |
| Head Detection | [yolov8-head](https://github.com/jamjamjon/assets/releases/download/yolo/v8-head-fp16.onnx) |
| Fall Detection | [yolov8-falldown](https://github.com/jamjamjon/assets/releases/download/yolo/v8-falldown-fp16.onnx) |
| Trash Detection | [yolov8-plastic-bag](https://github.com/jamjamjon/assets/releases/download/yolo/v8-plastic-bag-fp16.onnx) |
| FaceParsing | [yolov8-face-parsing-seg](https://github.com/jamjamjon/assets/releases/download/yolo/v8-face-parsing.onnx) |

362
examples/yolo/main.rs
View File

@ -1,171 +1,213 @@
use anyhow::Result;
use clap::Parser;
use usls::{
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOScale, YOLOTask,
YOLOVersion, COCO_SKELETONS_16,
models::YOLO, Annotator, DataLoader, Options, COCO_CLASS_NAMES_80, COCO_SKELETONS_16,
IMAGENET_NAMES_1K,
};
#[derive(Parser, Clone)]
#[command(author, version, about, long_about = None)]
pub struct Args {
/// Path to the model
#[arg(long)]
pub model: Option<String>,
#[derive(argh::FromArgs, Debug)]
/// Example
struct Args {
/// model file
#[argh(option)]
model: Option<String>,
/// Input source path
#[arg(long, default_value_t = String::from("./assets/bus.jpg"))]
pub source: String,
/// source
#[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
source: String,
/// YOLO Task
#[arg(long, value_enum, default_value_t = YOLOTask::Detect)]
pub task: YOLOTask,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// YOLO Version
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)]
pub ver: YOLOVersion,
/// task
#[argh(option, default = "String::from(\"det\")")]
task: String,
/// YOLO Scale
#[arg(long, value_enum, default_value_t = YOLOScale::N)]
pub scale: YOLOScale,
/// version
#[argh(option, default = "8.0")]
ver: f32,
/// Batch size
#[arg(long, default_value_t = 1)]
pub batch_size: usize,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// Minimum input width
#[arg(long, default_value_t = 224)]
pub width_min: isize,
/// scale
#[argh(option, default = "String::from(\"n\")")]
scale: String,
/// Input width
#[arg(long, default_value_t = 640)]
pub width: isize,
/// trt_fp16
#[argh(option, default = "true")]
trt_fp16: bool,
/// Maximum input width
#[arg(long, default_value_t = 1024)]
pub width_max: isize,
/// find_contours
#[argh(option, default = "true")]
find_contours: bool,
/// Minimum input height
#[arg(long, default_value_t = 224)]
pub height_min: isize,
/// batch_size
#[argh(option, default = "1")]
batch_size: usize,
/// Input height
#[arg(long, default_value_t = 640)]
pub height: isize,
/// min_batch_size
#[argh(option, default = "1")]
min_batch_size: usize,
/// Maximum input height
#[arg(long, default_value_t = 1024)]
pub height_max: isize,
/// max_batch_size
#[argh(option, default = "4")]
max_batch_size: usize,
/// Number of classes
#[arg(long, default_value_t = 80)]
pub nc: usize,
/// min_image_width
#[argh(option, default = "224")]
min_image_width: isize,
/// Class confidence
#[arg(long)]
pub confs: Vec<f32>,
/// image_width
#[argh(option, default = "640")]
image_width: isize,
/// Enable TensorRT support
#[arg(long)]
pub trt: bool,
/// max_image_width
#[argh(option, default = "1280")]
max_image_width: isize,
/// Enable CUDA support
#[arg(long)]
pub cuda: bool,
/// min_image_height
#[argh(option, default = "224")]
min_image_height: isize,
/// Enable CoreML support
#[arg(long)]
pub coreml: bool,
/// image_height
#[argh(option, default = "640")]
image_height: isize,
/// Use TensorRT half precision
#[arg(long)]
pub half: bool,
/// max_image_height
#[argh(option, default = "1280")]
max_image_height: isize,
/// Device ID to use
#[arg(long, default_value_t = 0)]
pub device_id: usize,
/// num_classes
#[argh(option)]
num_classes: Option<usize>,
/// Enable performance profiling
#[arg(long)]
pub profile: bool,
/// num_keypoints
#[argh(option)]
num_keypoints: Option<usize>,
/// Disable contour drawing
#[arg(long)]
pub no_contours: bool,
/// use_coco_80_classes
#[argh(switch)]
use_coco_80_classes: bool,
/// Show result
#[arg(long)]
pub view: bool,
/// use_imagenet_1k_classes
#[argh(switch)]
use_imagenet_1k_classes: bool,
/// Do not save output
#[arg(long)]
pub nosave: bool,
/// confs
#[argh(option)]
confs: Vec<f32>,
/// keypoint_confs
#[argh(option)]
keypoint_confs: Vec<f32>,
/// exclude_classes
#[argh(option)]
exclude_classes: Vec<usize>,
/// retain_classes
#[argh(option)]
retain_classes: Vec<usize>,
/// class_names
#[argh(option)]
class_names: Vec<String>,
/// keypoint_names
#[argh(option)]
keypoint_names: Vec<String>,
}
fn main() -> Result<()> {
let args = Args::parse();
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// model path
let path = match &args.model {
None => format!(
"yolo/{}-{}-{}.onnx",
args.ver.name(),
args.scale.name(),
args.task.name()
),
Some(x) => x.to_string(),
};
let args: Args = argh::from_env();
// saveout
let saveout = match &args.model {
None => format!(
"{}-{}-{}",
args.ver.name(),
args.scale.name(),
args.task.name()
),
Some(x) => {
let p = std::path::PathBuf::from(&x);
p.file_stem().unwrap().to_str().unwrap().to_string()
}
};
// device
let device = if args.cuda {
Device::Cuda(args.device_id)
} else if args.trt {
Device::Trt(args.device_id)
} else if args.coreml {
Device::CoreML(args.device_id)
} else {
Device::Cpu(args.device_id)
};
// build options
let options = Options::new()
.with_model(&path)?
.with_yolo_version(args.ver)
.with_yolo_task(args.task)
.with_device(device)
.with_trt_fp16(args.half)
.with_ixx(0, 0, (1, args.batch_size as _, 4).into())
.with_ixx(0, 2, (args.height_min, args.height, args.height_max).into())
.with_ixx(0, 3, (args.width_min, args.width, args.width_max).into())
.with_confs(if args.confs.is_empty() {
let mut options = Options::yolo()
.with_model_file(&args.model.unwrap_or_default())
.with_model_task(args.task.as_str().try_into()?)
.with_model_version(args.ver.into())
.with_model_scale(args.scale.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.with_trt_fp16(args.trt_fp16)
.with_model_ixx(
0,
0,
(args.min_batch_size, args.batch_size, args.max_batch_size).into(),
)
.with_model_ixx(
0,
2,
(
args.min_image_height,
args.image_height,
args.max_image_height,
)
.into(),
)
.with_model_ixx(
0,
3,
(args.min_image_width, args.image_width, args.max_image_width).into(),
)
.with_class_confs(if args.confs.is_empty() {
&[0.2, 0.15]
} else {
&args.confs
})
.with_nc(args.nc)
// .with_names(&COCO_CLASS_NAMES_80)
// .with_names2(&COCO_KEYPOINTS_17)
.with_find_contours(!args.no_contours) // find contours or not
.exclude_classes(&[0])
// .retain_classes(&[0, 5])
.with_profile(args.profile);
.with_keypoint_confs(if args.keypoint_confs.is_empty() {
&[0.5]
} else {
&args.keypoint_confs
})
.with_find_contours(args.find_contours)
.retain_classes(&args.retain_classes)
.exclude_classes(&args.exclude_classes);
if args.use_coco_80_classes {
options = options.with_class_names(&COCO_CLASS_NAMES_80);
}
if args.use_imagenet_1k_classes {
options = options.with_class_names(&IMAGENET_NAMES_1K);
}
if let Some(nc) = args.num_classes {
options = options.with_nc(nc);
}
if let Some(nk) = args.num_keypoints {
options = options.with_nk(nk);
}
if !args.class_names.is_empty() {
options = options.with_class_names(
&args
.class_names
.iter()
.map(|x| x.as_str())
.collect::<Vec<_>>(),
);
}
if !args.keypoint_names.is_empty() {
options = options.with_keypoint_names(
&args
.keypoint_names
.iter()
.map(|x| x.as_str())
.collect::<Vec<_>>(),
);
}
// build model
let mut model = YOLO::new(options)?;
let mut model = YOLO::try_from(options.commit()?)?;
// build dataloader
let dl = DataLoader::new(&args.source)?
@ -175,56 +217,28 @@ fn main() -> Result<()> {
// build annotator
let annotator = Annotator::default()
.with_skeletons(&COCO_SKELETONS_16)
.without_masks(true) // No masks plotting when doing segment task.
.without_masks(true)
.with_bboxes_thickness(3)
.with_keypoints_name(false) // Enable keypoints names
.with_saveout_subs(&["YOLO"])
.with_saveout(&saveout);
// build viewer
let mut viewer = if args.view {
Some(Viewer::new().with_delay(5).with_scale(1.).resizable(true))
} else {
None
};
.with_saveout(model.spec());
// run & annotate
for (xs, _paths) in dl {
// let ys = model.run(&xs)?; // way one
let ys = model.forward(&xs, args.profile)?; // way two
let images_plotted = annotator.plot(&xs, &ys, !args.nosave)?;
let ys = model.forward(&xs)?;
// extract bboxes
// for y in ys.iter() {
// if let Some(bboxes) = y.bboxes() {
// println!("[Bboxes]: Found {} objects", bboxes.len());
// for (i, bbox) in bboxes.iter().enumerate() {
// println!("{}: {:?}", i, bbox)
// }
// }
// }
// show image
match &mut viewer {
Some(viewer) => viewer.imshow(&images_plotted)?,
None => continue,
// plot
annotator.annotate(&xs, &ys);
}
// check out window and key event
match &mut viewer {
Some(viewer) => {
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) {
break;
}
}
None => continue,
}
// write video
if !args.nosave {
match &mut viewer {
Some(viewer) => viewer.write_batch(&images_plotted)?,
None => continue,
}
}
}
// finish video write
if !args.nosave {
if let Some(viewer) = &mut viewer {
viewer.finish_write()?;
}
}
model.summary();
Ok(())
}

18
examples/yolop/main.rs
View File

@ -1,22 +1,26 @@
use anyhow::Result;
use usls::{models::YOLOPv2, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
// build model
let options = Options::default()
.with_model("yolop/v2-dyn-480x800.onnx")?
.with_confs(&[0.3]);
let options = Options::yolop_v2_480x800().commit()?;
let mut model = YOLOPv2::new(options)?;
// load image
let x = [DataLoader::try_read("images/car.jpg")?];
let x = [DataLoader::try_read("images/car-view.jpg")?];
// run
let y = model.run(&x)?;
let y = model.forward(&x)?;
// annotate
let annotator = Annotator::default()
.with_polygons_name(true)
.with_saveout("YOLOPv2");
.with_saveout(model.spec());
annotator.annotate(&x, &y);
Ok(())

9
examples/yolov8-rtdetr/README.md Normal file
View File

@ -0,0 +1,9 @@
## Quick Start
```shell
cargo run -r -F cuda --example yolov8-rtdetr -- --device cuda
```
```shell
Ys([Y { BBoxes: [Bbox { xyxy: [668.71356, 395.4159, 809.01587, 879.3043], class_id: 0, name: Some("person"), confidence: 0.950527 }, Bbox { xyxy: [48.866394, 399.50665, 248.22641, 904.7525], class_id: 0, name: Some("person"), confidence: 0.9504415 }, Bbox { xyxy: [20.197449, 230.00304, 805.026, 730.3445], class_id: 5, name: Some("bus"), confidence: 0.94705224 }, Bbox { xyxy: [221.3088, 405.65436, 345.44052, 860.2628], class_id: 0, name: Some("person"), confidence: 0.93062377 }, Bbox { xyxy: [0.34117508, 549.8391, 76.50758, 868.87646], class_id: 0, name: Some("person"), confidence: 0.71064234 }, Bbox { xyxy: [282.12543, 484.14166, 296.43207, 520.96246], class_id: 27, name: Some("tie"), confidence: 0.40305245 }] }])
```

45
examples/yolov8-rtdetr/main.rs Normal file
View File

@ -0,0 +1,45 @@
use anyhow::Result;
use usls::{models::YOLO, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let config = Options::yolo_v8_rtdetr_l()
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = YOLO::new(config)?;
// load images
let xs = DataLoader::try_read_batch(&["./assets/bus.jpg"])?;
// run
let ys = model.forward(&xs)?;
println!("{:?}", ys);
// annotate
let annotator = Annotator::default()
.with_bboxes_thickness(3)
.with_saveout(model.spec());
annotator.annotate(&xs, &ys);
Ok(())
}

2
rust-toolchain.toml
View File

@ -1,2 +0,0 @@
[toolchain]
channel = "1.79"

63
scripts/CelebAMask-HQ-To-YOLO-Labels.py
View File

@ -1,63 +0,0 @@
import cv2
import numpy as np
from pathlib import Path
from tqdm import tqdm
mapping = {
'background': 0,
'skin': 1,
'nose': 2,
'eye_g': 3,
'l_eye': 4,
'r_eye': 5,
'l_brow': 6,
'r_brow': 7,
'l_ear': 8,
'r_ear': 9,
'mouth': 10,
'u_lip': 11,
'l_lip': 12,
'hair': 13,
'hat': 14,
'ear_r': 15,
'neck_l': 16,
'neck': 17,
'cloth': 18
}
def main():
saveout_dir = Path("labels")
if not saveout_dir.exists():
saveout_dir.mkdir()
else:
import shutil
shutil.rmtree(saveout_dir)
saveout_dir.mkdir()
image_list = [x for x in Path("CelebAMask-HQ-mask-anno/").rglob("*.png")]
for image_path in tqdm(image_list, total=len(image_list)):
image_gray = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
stem = image_path.stem
name, cls_ = stem.split("_", 1)
segments = cv2.findContours(image_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
saveout = saveout_dir / f"{int(name)}.txt"
with open(saveout, 'a+') as f:
for segment in segments:
line = f"{mapping[cls_]}"
segment = segment / 512
for seg in segment:
xn, yn = seg[0]
line += f" {xn} {yn}"
f.write(line + "\n")
if __name__ == "__main__":
main()

8
scripts/convert2f16.py
View File

@ -1,8 +0,0 @@
import onnx
from pathlib import Path
from onnxconverter_common import float16
model_f32 = "onnx_model.onnx"
model_f16 = float16.convert_float_to_float16(onnx.load(model_f32))
saveout = Path(model_f32).with_name(Path(model_f32).stem + "-f16.onnx")
onnx.save(model_f16, saveout)

14
src/core/device.rs
View File

@ -1,14 +0,0 @@
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum Device {
Auto(usize),
Cpu(usize),
Cuda(usize),
Trt(usize),
CoreML(usize),
// Cann(usize),
// Acl(usize),
// Rocm(usize),
// Rknpu(usize),
// Openvino(usize),
// Onednn(usize),
}

426
src/core/hub.rs
View File

@ -1,426 +0,0 @@
use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use serde::{Deserialize, Serialize};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use crate::Dir;
/// Represents a downloadable asset in a release
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Asset {
pub name: String,
pub browser_download_url: String,
pub size: u64,
}
/// Represents a GitHub release
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Release {
pub tag_name: String,
pub assets: Vec<Asset>,
}
/// Manages interactions with a GitHub repository's releases
pub struct Hub {
/// github api
_gh_api_release: String,
/// GitHub repository owner
owner: String,
/// GitHub repository name
repo: String,
/// Optional list of releases fetched from GitHub
releases: Option<Vec<Release>>,
/// Path to cache file
cache: PathBuf,
/// Optional release tag to be used
tag: Option<String>,
/// Filename for the asset, used in cache management
file_name: Option<String>,
file_size: Option<u64>,
/// Full URL constructed for downloading the asset
url: Option<String>,
/// Local path where the asset will be stored
path: PathBuf,
/// Directory to store the downloaded file
to: Dir,
/// Download timeout in seconds
timeout: u64,
/// Time to live (cache duration)
ttl: std::time::Duration,
/// Maximum attempts for downloading
max_attempts: u32,
}
impl std::fmt::Debug for Hub {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Hub")
.field("owner", &self.owner)
.field("repo", &self.repo)
.field("cache", &self.cache)
.field("path", &self.path)
.field("releases", &self.releases.as_ref().map(|x| x.len()))
.field("ttl", &self.ttl)
.field("max_attempts", &self.max_attempts)
.finish()
}
}
impl Default for Hub {
fn default() -> Self {
let owner = "jamjamjon".to_string();
let repo = "assets".to_string();
let _gh_api_release = format!("https://api.github.com/repos/{}/{}/releases", owner, repo);
Self {
owner,
repo,
_gh_api_release,
url: None,
path: PathBuf::new(),
to: Dir::Cache,
tag: None,
file_name: None,
file_size: None,
releases: None,
cache: PathBuf::new(),
timeout: 3000,
max_attempts: 3,
ttl: std::time::Duration::from_secs(10 * 60),
}
}
}
impl Hub {
pub fn new() -> Result<Self> {
let mut to = Dir::Cache;
let cache = to
.path()
.or_else(|_| {
to = Dir::Home;
to.path()
})?
.join("cache_releases");
Ok(Self {
to,
cache,
..Default::default()
})
}
pub fn with_owner(mut self, owner: &str) -> Self {
self.owner = owner.to_string();
self
}
pub fn with_repo(mut self, repo: &str) -> Self {
self.repo = repo.to_string();
self
}
pub fn with_ttl(mut self, x: u64) -> Self {
self.ttl = std::time::Duration::from_secs(x);
self
}
pub fn with_timeout(mut self, x: u64) -> Self {
self.timeout = x;
self
}
pub fn with_max_attempts(mut self, x: u32) -> Self {
self.max_attempts = x;
self
}
pub fn fetch(mut self, s: &str) -> Result<Self> {
// try to fetch from hub or local cache
let p = PathBuf::from(s);
match p.exists() {
true => self.path = p,
false => {
// check remote
match s.split_once('/') {
Some((tag, file_name)) => {
// Extract tag and file from input string
self.tag = Some(tag.to_string());
self.file_name = Some(file_name.to_string());
// Check if releases are already loaded in memory
if self.releases.is_none() {
self.releases = Some(self.connect_remote()?);
}
if let Some(releases) = &self.releases {
// Validate the tag
let tags: Vec<&str> =
releases.iter().map(|x| x.tag_name.as_str()).collect();
if !tags.contains(&tag) {
anyhow::bail!(
"Hub tag '{}' not found in releases. Available tags: {:?}",
tag,
tags
);
}
// Validate the file
if let Some(release) = releases.iter().find(|r| r.tag_name == tag) {
let files: Vec<&str> =
release.assets.iter().map(|x| x.name.as_str()).collect();
if !files.contains(&file_name) {
anyhow::bail!(
"Hub file '{}' not found in tag '{}'. Available files: {:?}",
file_name,
tag,
files
);
} else {
for f_ in release.assets.iter() {
if f_.name.as_str() == file_name {
self.url = Some(f_.browser_download_url.clone());
self.file_size = Some(f_.size);
break;
}
}
}
}
self.path = self.to.path_with_subs(&[tag])?.join(file_name);
}
}
_ => anyhow::bail!(
"Download failed due to invalid format. Expected: <tag>/<file>, got: {}",
s
),
}
}
}
Ok(self)
}
/// Fetch releases from GitHub and cache them
fn fetch_and_cache_releases(url: &str, cache_path: &Path) -> Result<String> {
let response = ureq::get(url)
.set("User-Agent", "my-app")
.call()
.context("Failed to fetch releases from remote")?;
if response.status() != 200 {
anyhow::bail!(
"Failed to fetch releases from remote ({}): status {} - {}",
url,
response.status(),
response.status_text()
);
}
let body = response
.into_string()
.context("Failed to read response body")?;
// Ensure cache directory exists
let parent_dir = cache_path
.parent()
.context("Invalid cache path; no parent directory found")?;
std::fs::create_dir_all(parent_dir)
.with_context(|| format!("Failed to create cache directory: {:?}", parent_dir))?;
// Create temporary file
let mut temp_file = tempfile::NamedTempFile::new_in(parent_dir)
.context("Failed to create temporary cache file")?;
// Write data to temporary file
temp_file
.write_all(body.as_bytes())
.context("Failed to write to temporary cache file")?;
// Persist temporary file as the cache
temp_file.persist(cache_path).with_context(|| {
format!("Failed to persist temporary cache file to {:?}", cache_path)
})?;
Ok(body)
}
pub fn tags(&mut self) -> Option<Vec<&str>> {
if self.releases.is_none() {
self.releases = self.connect_remote().ok();
}
self.releases
.as_ref()
.map(|releases| releases.iter().map(|x| x.tag_name.as_str()).collect())
}
pub fn files(&mut self, tag: &str) -> Option<Vec<&str>> {
if self.releases.is_none() {
self.releases = self.connect_remote().ok();
}
self.releases.as_ref().map(|releases| {
releases
.iter()
.find(|r| r.tag_name == tag)
.map(|a| a.assets.iter().map(|x| x.name.as_str()).collect())
})?
}
pub fn connect_remote(&mut self) -> Result<Vec<Release>> {
let span = tracing::span!(tracing::Level::INFO, "Hub-connect_remote");
let _guard = span.enter();
let should_download = if !self.cache.exists() {
tracing::info!("No cache found, fetching data from GitHub");
true
} else {
match std::fs::metadata(&self.cache)?.modified() {
Err(_) => {
tracing::info!("Cannot get file modified time, fetching new data from GitHub");
true
}
Ok(modified_time) => {
if std::time::SystemTime::now().duration_since(modified_time)? < self.ttl {
tracing::info!("Using cached data");
false
} else {
tracing::info!("Cache expired, fetching new data from GitHub");
true
}
}
}
};
let body = if should_download {
Self::fetch_and_cache_releases(&self._gh_api_release, &self.cache)?
} else {
std::fs::read_to_string(&self.cache)?
};
let releases: Vec<Release> = serde_json::from_str(&body)?;
Ok(releases)
}
/// Commit the downloaded file, downloading if necessary
pub fn commit(&self) -> Result<String> {
if let Some(url) = &self.url {
// Download if the file does not exist or if the size of file does not match
if !self.path.is_file()
|| self.path.is_file()
&& Some(std::fs::metadata(&self.path)?.len()) != self.file_size
{
let name = format!(
"{}/{}",
self.tag.as_ref().unwrap(),
self.file_name.as_ref().unwrap()
);
Self::download(
url.as_str(),
&self.path,
Some(&name),
Some(self.timeout),
Some(self.max_attempts),
)?;
}
}
self.path
.to_str()
.map(|s| s.to_string())
.with_context(|| format!("Failed to convert PathBuf: {:?} to String", self.path))
}
/// Download a file from a github release to a specified path with a progress bar
pub fn download<P: AsRef<Path> + std::fmt::Debug>(
src: &str,
dst: P,
prompt: Option<&str>,
timeout: Option<u64>,
max_attempts: Option<u32>,
) -> Result<()> {
// TODO: other url, not just github release page
let max_attempts = max_attempts.unwrap_or(2);
let timeout_duration = std::time::Duration::from_secs(timeout.unwrap_or(2000));
let agent = ureq::AgentBuilder::new().try_proxy_from_env(true).build();
for i_try in 0..max_attempts {
let resp = agent
.get(src)
.timeout(timeout_duration)
.call()
.with_context(|| {
format!(
"Failed to download file from {}, timeout: {:?}",
src, timeout_duration
)
})?;
let ntotal = resp
.header("Content-Length")
.and_then(|s| s.parse::<u64>().ok())
.context("Content-Length header is missing or invalid")?;
let pb = ProgressBar::new(ntotal);
pb.set_style(
ProgressStyle::with_template(
"{prefix:.cyan.bold} {msg} |{bar}| ({percent_precise}%, {binary_bytes}/{binary_total_bytes}, {binary_bytes_per_sec})",
)?
.progress_chars("██ "),
);
pb.set_prefix(if i_try == 0 {
" Fetching"
} else {
" Re-Fetching"
});
pb.set_message(prompt.unwrap_or_default().to_string());
let mut reader = resp.into_reader();
let mut buffer = [0; 256];
let mut downloaded_bytes = 0usize;
let mut file = std::fs::File::create(&dst)
.with_context(|| format!("Failed to create destination file: {:?}", dst))?;
loop {
let bytes_read = reader.read(&mut buffer)?;
if bytes_read == 0 {
break;
}
file.write_all(&buffer[..bytes_read])
.context("Failed to write to file")?;
downloaded_bytes += bytes_read;
pb.inc(bytes_read as u64);
}
// check size
if downloaded_bytes as u64 != ntotal {
continue;
}
// update
pb.set_prefix(" Downloaded");
pb.set_style(ProgressStyle::with_template(
crate::PROGRESS_BAR_STYLE_FINISH_3,
)?);
pb.finish();
if i_try != max_attempts {
break;
} else {
anyhow::bail!("Exceeded the maximum number of download attempts");
}
}
Ok(())
}
}

6
src/core/metric.rs
View File

@ -1,6 +0,0 @@
#[derive(Debug)]
pub enum Metric {
IP,
Cos,
L2,
}

45
src/core/mod.rs
View File

@ -1,45 +0,0 @@
mod annotator;
mod dataloader;
mod device;
mod dir;
mod dynconf;
mod hub;
mod logits_sampler;
mod media;
mod metric;
mod min_opt_max;
pub mod onnx;
pub mod ops;
mod options;
mod ort_engine;
mod task;
mod tokenizer_stream;
mod ts;
mod viewer;
mod vision;
mod x;
mod xs;
pub use annotator::Annotator;
pub use dataloader::DataLoader;
pub use device::Device;
pub use dir::Dir;
pub use dynconf::DynConf;
pub use hub::Hub;
pub use logits_sampler::LogitsSampler;
pub use media::*;
pub use metric::Metric;
pub use min_opt_max::MinOptMax;
pub use ops::Ops;
pub use options::Options;
pub use ort_engine::*;
pub use task::Task;
pub use tokenizer_stream::TokenizerStream;
pub use ts::Ts;
pub use viewer::Viewer;
pub use vision::Vision;
pub use x::X;
pub use xs::Xs;
// re-export
pub use minifb::Key;

295
src/core/options.rs
View File

@ -1,295 +0,0 @@
//! Options for build models.
use anyhow::Result;
use crate::{
models::{SamKind, SapiensTask, YOLOPreds, YOLOTask, YOLOVersion},
Device, Hub, Iiix, MinOptMax, Task,
};
/// Options for building models
#[derive(Debug, Clone)]
pub struct Options {
pub onnx_path: String,
pub task: Task,
pub device: Device,
pub batch_size: usize,
pub iiixs: Vec<Iiix>,
pub profile: bool,
pub num_dry_run: usize,
// trt related
pub trt_engine_cache_enable: bool,
pub trt_int8_enable: bool,
pub trt_fp16_enable: bool,
// options for Vision and Language models
pub nc: Option<usize>,
pub nk: Option<usize>,
pub nm: Option<usize>,
pub confs: Vec<f32>,
pub confs2: Vec<f32>,
pub confs3: Vec<f32>,
pub kconfs: Vec<f32>,
pub iou: Option<f32>,
pub tokenizer: Option<String>,
pub vocab: Option<String>,
pub context_length: Option<usize>,
pub names: Option<Vec<String>>, // names
pub names2: Option<Vec<String>>, // names2
pub names3: Option<Vec<String>>, // names3
pub min_width: Option<f32>,
pub min_height: Option<f32>,
pub unclip_ratio: f32, // DB
pub yolo_task: Option<YOLOTask>,
pub yolo_version: Option<YOLOVersion>,
pub yolo_preds: Option<YOLOPreds>,
pub find_contours: bool,
pub sam_kind: Option<SamKind>,
pub use_low_res_mask: Option<bool>,
pub sapiens_task: Option<SapiensTask>,
pub classes_excluded: Vec<isize>,
pub classes_retained: Vec<isize>,
}
impl Default for Options {
fn default() -> Self {
Self {
onnx_path: String::new(),
device: Device::Cuda(0),
profile: false,
batch_size: 1,
iiixs: vec![],
num_dry_run: 3,
trt_engine_cache_enable: true,
trt_int8_enable: false,
trt_fp16_enable: false,
nc: None,
nk: None,
nm: None,
confs: vec![0.3f32],
confs2: vec![0.3f32],
confs3: vec![0.3f32],
kconfs: vec![0.5f32],
iou: None,
tokenizer: None,
vocab: None,
context_length: None,
names: None,
names2: None,
names3: None,
min_width: None,
min_height: None,
unclip_ratio: 1.5,
yolo_task: None,
yolo_version: None,
yolo_preds: None,
find_contours: false,
sam_kind: None,
use_low_res_mask: None,
sapiens_task: None,
task: Task::Untitled,
classes_excluded: vec![],
classes_retained: vec![],
}
}
}
impl Options {
pub fn new() -> Self {
Default::default()
}
pub fn with_task(mut self, task: Task) -> Self {
self.task = task;
self
}
pub fn with_model(mut self, onnx_path: &str) -> Result<Self> {
self.onnx_path = Hub::new()?.fetch(onnx_path)?.commit()?;
Ok(self)
}
pub fn with_batch_size(mut self, n: usize) -> Self {
self.batch_size = n;
self
}
pub fn with_batch(mut self, n: usize) -> Self {
self.batch_size = n;
self
}
pub fn with_dry_run(mut self, n: usize) -> Self {
self.num_dry_run = n;
self
}
pub fn with_device(mut self, device: Device) -> Self {
self.device = device;
self
}
pub fn with_cuda(mut self, id: usize) -> Self {
self.device = Device::Cuda(id);
self
}
pub fn with_trt(mut self, id: usize) -> Self {
self.device = Device::Trt(id);
self
}
pub fn with_cpu(mut self) -> Self {
self.device = Device::Cpu(0);
self
}
pub fn with_coreml(mut self, id: usize) -> Self {
self.device = Device::CoreML(id);
self
}
pub fn with_trt_fp16(mut self, x: bool) -> Self {
self.trt_fp16_enable = x;
self
}
pub fn with_yolo_task(mut self, x: YOLOTask) -> Self {
self.yolo_task = Some(x);
self
}
pub fn with_sapiens_task(mut self, x: SapiensTask) -> Self {
self.sapiens_task = Some(x);
self
}
pub fn with_yolo_version(mut self, x: YOLOVersion) -> Self {
self.yolo_version = Some(x);
self
}
pub fn with_profile(mut self, profile: bool) -> Self {
self.profile = profile;
self
}
pub fn with_find_contours(mut self, x: bool) -> Self {
self.find_contours = x;
self
}
pub fn with_sam_kind(mut self, x: SamKind) -> Self {
self.sam_kind = Some(x);
self
}
pub fn use_low_res_mask(mut self, x: bool) -> Self {
self.use_low_res_mask = Some(x);
self
}
pub fn with_names(mut self, names: &[&str]) -> Self {
self.names = Some(names.iter().map(|x| x.to_string()).collect::<Vec<String>>());
self
}
pub fn with_names2(mut self, names: &[&str]) -> Self {
self.names2 = Some(names.iter().map(|x| x.to_string()).collect::<Vec<String>>());
self
}
pub fn with_names3(mut self, names: &[&str]) -> Self {
self.names3 = Some(names.iter().map(|x| x.to_string()).collect::<Vec<String>>());
self
}
pub fn with_vocab(mut self, vocab: &str) -> Result<Self> {
self.vocab = Some(Hub::new()?.fetch(vocab)?.commit()?);
Ok(self)
}
pub fn with_context_length(mut self, n: usize) -> Self {
self.context_length = Some(n);
self
}
pub fn with_tokenizer(mut self, tokenizer: &str) -> Result<Self> {
self.tokenizer = Some(Hub::new()?.fetch(tokenizer)?.commit()?);
Ok(self)
}
pub fn with_unclip_ratio(mut self, x: f32) -> Self {
self.unclip_ratio = x;
self
}
pub fn with_min_width(mut self, x: f32) -> Self {
self.min_width = Some(x);
self
}
pub fn with_min_height(mut self, x: f32) -> Self {
self.min_height = Some(x);
self
}
pub fn with_yolo_preds(mut self, x: YOLOPreds) -> Self {
self.yolo_preds = Some(x);
self
}
pub fn with_nc(mut self, nc: usize) -> Self {
self.nc = Some(nc);
self
}
pub fn with_nk(mut self, nk: usize) -> Self {
self.nk = Some(nk);
self
}
pub fn with_iou(mut self, x: f32) -> Self {
self.iou = Some(x);
self
}
pub fn with_confs(mut self, x: &[f32]) -> Self {
self.confs = x.to_vec();
self
}
pub fn with_confs2(mut self, x: &[f32]) -> Self {
self.confs2 = x.to_vec();
self
}
pub fn with_confs3(mut self, x: &[f32]) -> Self {
self.confs3 = x.to_vec();
self
}
pub fn with_kconfs(mut self, kconfs: &[f32]) -> Self {
self.kconfs = kconfs.to_vec();
self
}
pub fn with_ixx(mut self, i: usize, ii: usize, x: MinOptMax) -> Self {
self.iiixs.push(Iiix::from((i, ii, x)));
self
}
pub fn exclude_classes(mut self, xs: &[isize]) -> Self {
self.classes_retained.clear();
self.classes_excluded.extend_from_slice(xs);
self
}
pub fn retain_classes(mut self, xs: &[isize]) -> Self {
self.classes_excluded.clear();
self.classes_retained.extend_from_slice(xs);
self
}
}

669
src/core/ort_engine.rs
View File

@ -1,669 +0,0 @@
use anyhow::Result;
use half::f16;
use ndarray::{Array, IxDyn};
use ort::{
execution_providers::{ExecutionProvider, TensorRTExecutionProvider},
session::{builder::SessionBuilder, Session},
tensor::TensorElementType,
};
use prost::Message;
use std::collections::HashSet;
use crate::{
build_progress_bar, human_bytes, onnx, Device, Dir, MinOptMax, Ops, Options, Ts, Xs,
CHECK_MARK, CROSS_MARK, X,
};
/// A struct for input composed of the i-th input, the ii-th dimension, and the value.
#[derive(Clone, Debug, Default)]
pub struct Iiix {
pub i: usize,
pub ii: usize,
pub x: MinOptMax,
}
impl From<(usize, usize, MinOptMax)> for Iiix {
fn from((i, ii, x): (usize, usize, MinOptMax)) -> Self {
Self { i, ii, x }
}
}
/// A struct for tensor attrs composed of the names, the dtypes, and the dimensions.
#[derive(Debug)]
pub struct OrtTensorAttr {
pub names: Vec<String>,
pub dtypes: Vec<TensorElementType>,
pub dimss: Vec<Vec<usize>>,
}
/// ONNXRuntime Backend
#[derive(Debug)]
pub struct OrtEngine {
name: String,
session: Session,
device: Device,
inputs_minoptmax: Vec<Vec<MinOptMax>>,
inputs_attrs: OrtTensorAttr,
outputs_attrs: OrtTensorAttr,
profile: bool,
num_dry_run: usize,
model_proto: onnx::ModelProto,
params: usize,
wbmems: usize,
ts: Ts,
}
impl OrtEngine {
pub fn new(config: &Options) -> Result<Self> {
let span = tracing::span!(tracing::Level::INFO, "OrtEngine-new");
let _guard = span.enter();
// onnx graph
let model_proto = Self::load_onnx(&config.onnx_path)?;
let graph = match &model_proto.graph {
Some(graph) => graph,
None => anyhow::bail!("No graph found in this proto. Failed to parse ONNX model."),
};
// model params & mems
let byte_alignment = 16; // 16 for simd; 8 for most
let mut params: usize = 0;
let mut wbmems: usize = 0;
let mut initializer_names: HashSet<&str> = HashSet::new();
for tensor_proto in graph.initializer.iter() {
initializer_names.insert(&tensor_proto.name);
let param = tensor_proto.dims.iter().product::<i64>() as usize;
params += param;
// mems
let param = Ops::make_divisible(param, byte_alignment);
let n = Self::nbytes_from_onnx_dtype_id(tensor_proto.data_type as usize);
let wbmem = param * n;
wbmems += wbmem;
}
// inputs & outputs
let inputs_attrs = Self::io_from_onnx_value_info(&initializer_names, &graph.input)?;
let outputs_attrs = Self::io_from_onnx_value_info(&initializer_names, &graph.output)?;
let inputs_minoptmax =
Self::build_inputs_minoptmax(&inputs_attrs, &config.iiixs, config.batch_size)?;
// build
ort::init().commit()?;
let mut builder = Session::builder()?;
let mut device = config.device.to_owned();
match device {
Device::Trt(device_id) => {
Self::build_trt(
&inputs_attrs.names,
&inputs_minoptmax,
&mut builder,
device_id,
config.trt_int8_enable,
config.trt_fp16_enable,
config.trt_engine_cache_enable,
)?;
}
Device::Cuda(device_id) => {
Self::build_cuda(&mut builder, device_id).unwrap_or_else(|err| {
tracing::warn!("{err}, Using cpu");
device = Device::Cpu(0);
})
}
Device::CoreML(_) => Self::build_coreml(&mut builder).unwrap_or_else(|err| {
tracing::warn!("{err}, Using cpu");
device = Device::Cpu(0);
}),
Device::Cpu(_) => {
Self::build_cpu(&mut builder)?;
}
_ => todo!(),
}
let session = builder
.with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)?
.commit_from_file(&config.onnx_path)?;
// summary
tracing::info!(
"{CHECK_MARK} Backend: ONNXRuntime | Opset: {} | Device: {:?} | Params: {}",
model_proto.opset_import[0].version,
device,
human_bytes(params as f64),
);
Ok(Self {
name: config.onnx_path.to_owned(),
session,
device,
inputs_minoptmax,
inputs_attrs,
outputs_attrs,
profile: config.profile,
num_dry_run: config.num_dry_run,
model_proto,
params,
wbmems,
ts: Ts::default(),
})
}
fn build_trt(
names: &[String],
inputs_minoptmax: &[Vec<MinOptMax>],
builder: &mut SessionBuilder,
device_id: usize,
int8_enable: bool,
fp16_enable: bool,
engine_cache_enable: bool,
) -> Result<()> {
let span = tracing::span!(tracing::Level::INFO, "OrtEngine-build_trt");
let _guard = span.enter();
// auto generate shapes
let mut spec_min = String::new();
let mut spec_opt = String::new();
let mut spec_max = String::new();
for (i, name) in names.iter().enumerate() {
if i != 0 {
spec_min.push(',');
spec_opt.push(',');
spec_max.push(',');
}
let mut s_min = format!("{}:", name);
let mut s_opt = format!("{}:", name);
let mut s_max = format!("{}:", name);
for d in inputs_minoptmax[i].iter() {
let min_ = &format!("{}x", d.min());
let opt_ = &format!("{}x", d.opt());
let max_ = &format!("{}x", d.max());
s_min += min_;
s_opt += opt_;
s_max += max_;
}
s_min.pop();
s_opt.pop();
s_max.pop();
spec_min += &s_min;
spec_opt += &s_opt;
spec_max += &s_max;
}
let p = Dir::Cache.path_with_subs(&["trt-cache"])?;
let trt = TensorRTExecutionProvider::default()
.with_device_id(device_id as i32)
.with_int8(int8_enable)
.with_fp16(fp16_enable)
.with_engine_cache(engine_cache_enable)
.with_engine_cache_path(p.to_str().unwrap())
.with_timing_cache(false)
.with_profile_min_shapes(spec_min)
.with_profile_opt_shapes(spec_opt)
.with_profile_max_shapes(spec_max);
if trt.is_available()? && trt.register(builder).is_ok() {
tracing::info!("🐢 Initial model serialization with TensorRT may require a wait...\n");
Ok(())
} else {
anyhow::bail!("{CROSS_MARK} TensorRT initialization failed")
}
}
fn build_cuda(builder: &mut SessionBuilder, device_id: usize) -> Result<()> {
let ep = ort::execution_providers::CUDAExecutionProvider::default()
.with_device_id(device_id as i32);
if ep.is_available()? && ep.register(builder).is_ok() {
Ok(())
} else {
anyhow::bail!("{CROSS_MARK} CUDA initialization failed")
}
}
fn build_coreml(builder: &mut SessionBuilder) -> Result<()> {
let ep = ort::execution_providers::CoreMLExecutionProvider::default().with_subgraphs(); //.with_ane_only();
if ep.is_available()? && ep.register(builder).is_ok() {
Ok(())
} else {
anyhow::bail!("{CROSS_MARK} CoreML initialization failed")
}
}
fn build_cpu(builder: &mut SessionBuilder) -> Result<()> {
let ep = ort::execution_providers::CPUExecutionProvider::default();
if ep.is_available()? && ep.register(builder).is_ok() {
Ok(())
} else {
anyhow::bail!("{CROSS_MARK} CPU initialization failed")
}
}
pub fn dry_run(&mut self) -> Result<()> {
if self.num_dry_run > 0 {
// pb
let name = std::path::Path::new(&self.name);
let pb = build_progress_bar(
self.num_dry_run as u64,
" DryRun",
Some(
name.file_name()
.and_then(|x| x.to_str())
.unwrap_or_default(),
),
crate::PROGRESS_BAR_STYLE_CYAN_2,
)?;
// dummy inputs
let mut xs = Vec::new();
for i in self.inputs_minoptmax.iter() {
let mut x: Vec<usize> = Vec::new();
for i_ in i.iter() {
x.push(i_.opt());
}
let x: Array<f32, IxDyn> = Array::ones(x).into_dyn();
xs.push(X::from(x));
}
let xs = Xs::from(xs);
// run
for _ in 0..self.num_dry_run {
pb.inc(1);
self.run(xs.clone())?;
}
self.ts.clear();
// update
let name = std::path::Path::new(&self.name);
pb.set_message(format!(
"{} on {:?}",
name.file_name()
.and_then(|x| x.to_str())
.unwrap_or_default(),
self.device,
));
pb.set_style(indicatif::ProgressStyle::with_template(
crate::PROGRESS_BAR_STYLE_FINISH,
)?);
pb.finish();
}
Ok(())
}
pub fn run(&mut self, xs: Xs) -> Result<Xs> {
let span = tracing::span!(tracing::Level::INFO, "OrtEngine-run");
let _guard = span.enter();
// inputs dtype alignment
let mut xs_ = Vec::new();
let t_pre = std::time::Instant::now();
for (idtype, x) in self.inputs_attrs.dtypes.iter().zip(xs.into_iter()) {
let x_ = match &idtype {
TensorElementType::Float32 => ort::value::Value::from_array(x.view())?.into_dyn(),
TensorElementType::Float16 => {
ort::value::Value::from_array(x.mapv(f16::from_f32).view())?.into_dyn()
}
TensorElementType::Int32 => {
ort::value::Value::from_array(x.mapv(|x_| x_ as i32).view())?.into_dyn()
}
TensorElementType::Int64 => {
ort::value::Value::from_array(x.mapv(|x_| x_ as i64).view())?.into_dyn()
}
TensorElementType::Uint8 => {
ort::value::Value::from_array(x.mapv(|x_| x_ as u8).view())?.into_dyn()
}
TensorElementType::Int8 => {
ort::value::Value::from_array(x.mapv(|x_| x_ as i8).view())?.into_dyn()
}
TensorElementType::Bool => {
ort::value::Value::from_array(x.mapv(|x_| x_ != 0.).view())?.into_dyn()
}
_ => todo!(),
};
xs_.push(Into::<ort::session::SessionInputValue<'_>>::into(x_));
}
let t_pre = t_pre.elapsed();
self.ts.add_or_push(0, t_pre);
// inference
let t_run = std::time::Instant::now();
let outputs = self.session.run(&xs_[..])?;
let t_run = t_run.elapsed();
self.ts.add_or_push(1, t_run);
// oputput
let mut ys = Xs::new();
let t_post = std::time::Instant::now();
for (dtype, name) in self
.outputs_attrs
.dtypes
.iter()
.zip(self.outputs_attrs.names.iter())
{
let y = &outputs[name.as_str()];
let y_ = match &dtype {
TensorElementType::Float32 => match y.try_extract_tensor::<f32>() {
Err(err) => {
tracing::error!("Error: {:?}. Output name: {:?}", err, name);
Array::zeros(0).into_dyn()
}
Ok(x) => x.view().into_owned(),
},
TensorElementType::Float16 => match y.try_extract_tensor::<f16>() {
Err(err) => {
tracing::error!("Error: {:?}. Output name: {:?}", err, name);
Array::zeros(0).into_dyn()
}
Ok(x) => x.view().mapv(f16::to_f32).into_owned(),
},
TensorElementType::Int64 => match y.try_extract_tensor::<i64>() {
Err(err) => {
tracing::error!("Error: {:?}. Output name: {:?}", err, name);
Array::zeros(0).into_dyn()
}
Ok(x) => x.view().to_owned().mapv(|x| x as f32).into_owned(),
},
_ => todo!(),
};
ys.push_kv(name.as_str(), X::from(y_))?;
}
let t_post = t_post.elapsed();
self.ts.add_or_push(2, t_post);
if self.profile {
let len = 10usize;
let n = 4usize;
tracing::info!(
"[Profile] {:>len$.n$?} ({:>len$.n$?} avg) [alignment: {:>len$.n$?} ({:>len$.n$?} avg) | inference: {:>len$.n$?} ({:>len$.n$?} avg) | to_f32: {:>len$.n$?} ({:>len$.n$?} avg)]",
t_pre + t_run + t_post,
self.ts.avg(),
t_pre,
self.ts.avgi(0),
t_run,
self.ts.avgi(1),
t_post,
self.ts.avgi(2),
);
}
Ok(ys)
}
fn build_inputs_minoptmax(
inputs_attrs: &OrtTensorAttr,
iiixs: &[Iiix],
batch_size: usize,
) -> Result<Vec<Vec<MinOptMax>>> {
let span = tracing::span!(tracing::Level::INFO, "OrtEngine-build_inputs_minoptmax");
let _guard = span.enter();
// init
let mut ys: Vec<Vec<MinOptMax>> = inputs_attrs
.dimss
.iter()
.map(|dims| dims.iter().map(|&x| MinOptMax::from(x)).collect())
.collect();
// update from customized
for iiix in iiixs.iter() {
if let Some(x) = inputs_attrs
.dimss
.get(iiix.i)
.and_then(|dims| dims.get(iiix.ii))
{
// dynamic
if *x == 0 {
ys[iiix.i][iiix.ii] = iiix.x.clone();
}
} else {
anyhow::bail!(
"Cannot retrieve the {}-th dimension of the {}-th input.",
iiix.ii,
iiix.i,
);
}
}
// deal with the dynamic axis
ys.iter_mut().enumerate().for_each(|(i, xs)| {
xs.iter_mut().enumerate().for_each(|(ii, x)| {
if x.is_dyn() {
let n = if ii == 0 { batch_size } else { 1 };
let y = MinOptMax::from(n);
tracing::warn!(
"Using dynamic shapes in inputs without specifying it: the {}-th input, the {}-th dimension. \
Using {:?} by default. You should make it clear when using TensorRT.",
i + 1, ii + 1, y
);
*x = y;
}
});
});
Ok(ys)
}
#[allow(dead_code)]
fn nbytes_from_onnx_dtype_id(x: usize) -> usize {
match x {
7 | 11 | 13 => 8, // i64, f64, u64
1 | 6 | 12 => 4, // f32, i32, u32
10 | 16 | 5 | 4 => 2, // f16, bf16, i16, u16
2 | 3 | 9 => 1, // u8, i8, bool
8 => 4, // string(1~4)
_ => todo!(),
}
}
#[allow(dead_code)]
fn nbytes_from_onnx_dtype(x: &ort::tensor::TensorElementType) -> usize {
match x {
ort::tensor::TensorElementType::Float64
| ort::tensor::TensorElementType::Uint64
| ort::tensor::TensorElementType::Int64 => 8, // i64, f64, u64
ort::tensor::TensorElementType::Float32
| ort::tensor::TensorElementType::Uint32
| ort::tensor::TensorElementType::Int32
| ort::tensor::TensorElementType::String => 4, // f32, i32, u32, string(1~4)
ort::tensor::TensorElementType::Float16
| ort::tensor::TensorElementType::Bfloat16
| ort::tensor::TensorElementType::Int16
| ort::tensor::TensorElementType::Uint16 => 2, // f16, bf16, i16, u16
ort::tensor::TensorElementType::Uint8
| ort::tensor::TensorElementType::Int8
| ort::tensor::TensorElementType::Bool => 1, // u8, i8, bool
}
}
#[allow(dead_code)]
fn ort_dtype_from_onnx_dtype_id(value: i32) -> Option<ort::tensor::TensorElementType> {
match value {
0 => None,
1 => Some(ort::tensor::TensorElementType::Float32),
2 => Some(ort::tensor::TensorElementType::Uint8),
3 => Some(ort::tensor::TensorElementType::Int8),
4 => Some(ort::tensor::TensorElementType::Uint16),
5 => Some(ort::tensor::TensorElementType::Int16),
6 => Some(ort::tensor::TensorElementType::Int32),
7 => Some(ort::tensor::TensorElementType::Int64),
8 => Some(ort::tensor::TensorElementType::String),
9 => Some(ort::tensor::TensorElementType::Bool),
10 => Some(ort::tensor::TensorElementType::Float16),
11 => Some(ort::tensor::TensorElementType::Float64),
12 => Some(ort::tensor::TensorElementType::Uint32),
13 => Some(ort::tensor::TensorElementType::Uint64),
14 => None, // COMPLEX64
15 => None, // COMPLEX128
16 => Some(ort::tensor::TensorElementType::Bfloat16),
_ => None,
}
}
fn io_from_onnx_value_info(
initializer_names: &HashSet<&str>,
value_info: &[onnx::ValueInfoProto],
) -> Result<OrtTensorAttr> {
let mut dimss: Vec<Vec<usize>> = Vec::new();
let mut dtypes: Vec<ort::tensor::TensorElementType> = Vec::new();
let mut names: Vec<String> = Vec::new();
for v in value_info.iter() {
if initializer_names.contains(v.name.as_str()) {
continue;
}
names.push(v.name.to_string());
let dtype = match &v.r#type {
Some(dtype) => dtype,
None => continue,
};
let dtype = match &dtype.value {
Some(dtype) => dtype,
None => continue,
};
let tensor = match dtype {
onnx::type_proto::Value::TensorType(tensor) => tensor,
_ => continue,
};
let tensor_type = tensor.elem_type;
let tensor_type = match Self::ort_dtype_from_onnx_dtype_id(tensor_type) {
Some(dtype) => dtype,
None => continue,
};
dtypes.push(tensor_type);
let shapes = match &tensor.shape {
Some(shapes) => shapes,
None => continue,
};
let mut shape_: Vec<usize> = Vec::new();
for shape in shapes.dim.iter() {
match &shape.value {
None => continue,
Some(value) => match value {
onnx::tensor_shape_proto::dimension::Value::DimValue(x) => {
shape_.push(*x as _);
}
onnx::tensor_shape_proto::dimension::Value::DimParam(_) => {
shape_.push(0);
}
},
}
}
dimss.push(shape_);
}
Ok(OrtTensorAttr {
dimss,
dtypes,
names,
})
}
pub fn load_onnx<P: AsRef<std::path::Path>>(p: P) -> Result<onnx::ModelProto> {
let f = std::fs::read(p)?;
Ok(onnx::ModelProto::decode(f.as_slice())?)
}
pub fn oshapes(&self) -> &Vec<Vec<usize>> {
&self.outputs_attrs.dimss
}
pub fn odimss(&self) -> &Vec<Vec<usize>> {
&self.outputs_attrs.dimss
}
pub fn onames(&self) -> &Vec<String> {
&self.outputs_attrs.names
}
pub fn odtypes(&self) -> &Vec<ort::tensor::TensorElementType> {
&self.outputs_attrs.dtypes
}
pub fn ishapes(&self) -> &Vec<Vec<usize>> {
&self.inputs_attrs.dimss
}
pub fn idimss(&self) -> &Vec<Vec<usize>> {
&self.inputs_attrs.dimss
}
pub fn inames(&self) -> &Vec<String> {
&self.inputs_attrs.names
}
pub fn idtypes(&self) -> &Vec<ort::tensor::TensorElementType> {
&self.inputs_attrs.dtypes
}
pub fn device(&self) -> &Device {
&self.device
}
pub fn inputs_minoptmax(&self) -> &Vec<Vec<MinOptMax>> {
&self.inputs_minoptmax
}
pub fn batch(&self) -> &MinOptMax {
&self.inputs_minoptmax[0][0]
}
pub fn try_height(&self) -> Option<&MinOptMax> {
self.inputs_minoptmax.first().and_then(|x| x.get(2))
}
pub fn try_width(&self) -> Option<&MinOptMax> {
self.inputs_minoptmax.first().and_then(|x| x.get(3))
}
pub fn height(&self) -> &MinOptMax {
&self.inputs_minoptmax[0][2]
}
pub fn width(&self) -> &MinOptMax {
&self.inputs_minoptmax[0][3]
}
pub fn is_batch_dyn(&self) -> bool {
self.ishapes()[0][0] == 0
}
pub fn try_fetch(&self, key: &str) -> Option<String> {
match self.session.metadata() {
Err(_) => None,
Ok(metadata) => metadata.custom(key).unwrap_or_default(),
}
}
pub fn session(&self) -> &Session {
&self.session
}
pub fn ir_version(&self) -> usize {
self.model_proto.ir_version as usize
}
pub fn opset_version(&self) -> usize {
self.model_proto.opset_import[0].version as usize
}
pub fn producer_name(&self) -> String {
self.model_proto.producer_name.to_string()
}
pub fn producer_version(&self) -> String {
self.model_proto.producer_version.to_string()
}
pub fn model_version(&self) -> usize {
self.model_proto.model_version as usize
}
pub fn parameters(&self) -> usize {
self.params
}
pub fn memory_weights(&self) -> usize {
self.wbmems
}
pub fn ts(&self) -> &Ts {
&self.ts
}
}

87
src/core/tokenizer_stream.rs
View File

@ -1,87 +0,0 @@
// TODO: refactor
use anyhow::Result;
/// This is a wrapper around a tokenizer to ensure that tokens can be returned to the user in a
/// streaming way rather than having to wait for the full decoding.
#[derive(Debug)]
pub struct TokenizerStream {
tokenizer: tokenizers::Tokenizer,
tokens: Vec<u32>,
prev_index: usize,
current_index: usize,
}
impl TokenizerStream {
pub fn new(tokenizer: tokenizers::Tokenizer) -> Self {
Self {
tokenizer,
tokens: Vec::new(),
prev_index: 0,
current_index: 0,
}
}
pub fn into_inner(self) -> tokenizers::Tokenizer {
self.tokenizer
}
fn decode(&self, tokens: &[u32]) -> Result<String> {
match self.tokenizer.decode(tokens, true) {
Ok(str) => Ok(str),
Err(err) => anyhow::bail!("cannot decode: {err}"),
}
}
pub fn next_token(&mut self, token: u32) -> Result<Option<String>> {
let prev_text = if self.tokens.is_empty() {
String::new()
} else {
let tokens = &self.tokens[self.prev_index..self.current_index];
self.decode(tokens)?
};
self.tokens.push(token);
let text = self.decode(&self.tokens[self.prev_index..])?;
if text.len() > prev_text.len() {
let text = text.split_at(prev_text.len());
self.prev_index = self.current_index;
self.current_index = self.tokens.len();
Ok(Some(text.1.to_string()))
} else {
Ok(None)
}
}
pub fn decode_rest(&self) -> Result<Option<String>> {
let prev_text = if self.tokens.is_empty() {
String::new()
} else {
let tokens = &self.tokens[self.prev_index..self.current_index];
self.decode(tokens)?
};
let text = self.decode(&self.tokens[self.prev_index..])?;
if text.len() > prev_text.len() {
let text = text.split_at(prev_text.len());
Ok(Some(text.1.to_string()))
} else {
Ok(None)
}
}
pub fn decode_all(&self) -> Result<String> {
self.decode(&self.tokens)
}
pub fn get_token(&self, token_s: &str) -> Option<u32> {
self.tokenizer.get_vocab(true).get(token_s).copied()
}
pub fn tokenizer(&self) -> &tokenizers::Tokenizer {
&self.tokenizer
}
pub fn clear(&mut self) {
self.tokens.clear();
self.prev_index = 0;
self.current_index = 0;
}
}

49
src/core/ts.rs
View File

@ -1,49 +0,0 @@
use std::time::Duration;
#[derive(Debug, Default)]
pub struct Ts {
n: usize,
ts: Vec<Duration>,
}
impl Ts {
pub fn total(&self) -> Duration {
self.ts.iter().sum::<Duration>()
}
pub fn n(&self) -> usize {
self.n / self.ts.len()
}
pub fn avg(&self) -> Duration {
self.total() / self.n() as u32
}
pub fn avgi(&self, i: usize) -> Duration {
if i >= self.ts.len() {
panic!("Index out of bound");
}
self.ts[i] / self.n() as u32
}
pub fn ts(&self) -> &Vec<Duration> {
&self.ts
}
pub fn add_or_push(&mut self, i: usize, x: Duration) {
match self.ts.get_mut(i) {
Some(elem) => *elem += x,
None => {
if i >= self.ts.len() {
self.ts.push(x)
}
}
}
self.n += 1;
}
pub fn clear(&mut self) {
self.n = Default::default();
self.ts = Default::default();
}
}

51
src/core/vision.rs
View File

@ -1,51 +0,0 @@
use crate::{Options, Xs, Y};
pub trait Vision: Sized {
type Input; // DynamicImage
/// Creates a new instance of the model with the given options.
fn new(options: Options) -> anyhow::Result<Self>;
/// Preprocesses the input data.
fn preprocess(&self, xs: &[Self::Input]) -> anyhow::Result<Xs>;
/// Executes the model on the preprocessed data.
fn inference(&mut self, xs: Xs) -> anyhow::Result<Xs>;
/// Postprocesses the model's output.
fn postprocess(&self, xs: Xs, xs0: &[Self::Input]) -> anyhow::Result<Vec<Y>>;
/// Executes the full pipeline.
fn run(&mut self, xs: &[Self::Input]) -> anyhow::Result<Vec<Y>> {
let ys = self.preprocess(xs)?;
let ys = self.inference(ys)?;
let ys = self.postprocess(ys, xs)?;
Ok(ys)
}
/// Executes the full pipeline.
fn forward(&mut self, xs: &[Self::Input], profile: bool) -> anyhow::Result<Vec<Y>> {
let span = tracing::span!(tracing::Level::INFO, "Vision-forward");
let _guard = span.enter();
let t_pre = std::time::Instant::now();
let ys = self.preprocess(xs)?;
let t_pre = t_pre.elapsed();
let t_exe = std::time::Instant::now();
let ys = self.inference(ys)?;
let t_exe = t_exe.elapsed();
let t_post = std::time::Instant::now();
let ys = self.postprocess(ys, xs)?;
let t_post = t_post.elapsed();
if profile {
tracing::info!(
"> Preprocess: {t_pre:?} | Execution: {t_exe:?} | Postprocess: {t_post:?}"
);
}
Ok(ys)
}
}

279
src/lib.rs
View File

@ -1,258 +1,49 @@
//! **usls** is a Rust library integrated with **ONNXRuntime** that provides a collection of state-of-the-art models for **Computer Vision** and **Vision-Language** tasks, including:
//! **usls** is a Rust library integrated with **ONNXRuntime**, offering a suite of advanced models for **Computer Vision** and **Vision-Language** tasks, including:
//!
//! - **YOLO Models**: [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [YOLOv7](https://github.com/WongKinYiu/yolov7), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [YOLOv10](https://github.com/THU-MIG/yolov10)
//! - **YOLO Models**: [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [YOLOv7](https://github.com/WongKinYiu/yolov7), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [YOLOv10](https://github.com/THU-MIG/yolov10), [YOLO11](https://github.com/ultralytics/ultralytics)
//! - **SAM Models**: [SAM](https://github.com/facebookresearch/segment-anything), [SAM2](https://github.com/facebookresearch/segment-anything-2), [MobileSAM](https://github.com/ChaoningZhang/MobileSAM), [EdgeSAM](https://github.com/chongzhou96/EdgeSAM), [SAM-HQ](https://github.com/SysCV/sam-hq), [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM)
//! - **Vision Models**: [RTDETR](https://arxiv.org/abs/2304.08069), [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo), [DB](https://arxiv.org/abs/1911.08947), [SVTR](https://arxiv.org/abs/2205.00159), [Depth-Anything-v1-v2](https://github.com/LiheYoung/Depth-Anything), [DINOv2](https://github.com/facebookresearch/dinov2), [MODNet](https://github.com/ZHKKKe/MODNet), [Sapiens](https://arxiv.org/abs/2408.12569)
//! - **Vision-Language Models**: [CLIP](https://github.com/openai/CLIP), [BLIP](https://arxiv.org/abs/2201.12086), [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [Florence2](https://arxiv.org/abs/2311.06242)
//! - **Vision Models**: [RT-DETR](https://arxiv.org/abs/2304.08069), [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo), [Depth-Anything](https://github.com/LiheYoung/Depth-Anything), [DINOv2](https://github.com/facebookresearch/dinov2), [MODNet](https://github.com/ZHKKKe/MODNet), [Sapiens](https://arxiv.org/abs/2408.12569), [DepthPro](https://github.com/apple/ml-depth-pro), [FastViT](https://github.com/apple/ml-fastvit), [BEiT](https://github.com/microsoft/unilm/tree/master/beit), [MobileOne](https://github.com/apple/ml-mobileone)
//! - **Vision-Language Models**: [CLIP](https://github.com/openai/CLIP), [jina-clip-v1](https://huggingface.co/jinaai/jina-clip-v1), [BLIP](https://arxiv.org/abs/2201.12086), [GroundingDINO](https://github.com/IDEA-Research/GroundingDINO), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [Florence2](https://arxiv.org/abs/2311.06242)
//! - **OCR Models**: [DB](https://arxiv.org/abs/1911.08947), [FAST](https://github.com/czczup/FAST), [SVTR](https://arxiv.org/abs/2205.00159), [SLANet](https://paddlepaddle.github.io/PaddleOCR/latest/algorithm/table_recognition/algorithm_table_slanet.html), [TrOCR](https://huggingface.co/microsoft/trocr-base-printed), [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO)
//! - **And more...**
//!
//! # Examples
//! ## ⛳️ Cargo Features
//!
//! Refer to [All Runnable Demos](https://github.com/jamjamjon/usls/tree/main/examples)
//! By default, **none of the following features are enabled**. You can enable them as needed:
//!
//! # Quick Start
//! - **`auto`**: Automatically downloads prebuilt ONNXRuntime binaries from Pyke’s CDN for supported platforms.
//!
//! The following demo shows how to build a `YOLO` with [`Options`], load `image(s)`, `video` and `stream` with [`DataLoader`], and annotate the model's inference results with [`Annotator`].
//! - If disabled, you'll need to [compile `ONNXRuntime` from source](https://github.com/microsoft/onnxruntime) or [download a precompiled package](https://github.com/microsoft/onnxruntime/releases), and then [link it manually](https://ort.pyke.io/setup/linking).
//!
//! ```ignore
//! use usls::{models::YOLO, Annotator, DataLoader, Options, Vision, YOLOTask, YOLOVersion};
//! <details>
//! <summary>👉 For Linux or macOS Users</summary>
//!
//! fn main() -> anyhow::Result<()> {
//! // Build model with Options
//! let options = Options::new()
//! .with_trt(0)
//! .with_model("yolo/v8-m-dyn.onnx")?
//! .with_yolo_version(YOLOVersion::V8) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
//! .with_yolo_task(YOLOTask::Detect) // YOLOTask: Classify, Detect, Pose, Segment, Obb
//! .with_i00((1, 1, 4).into())
//! .with_i02((0, 640, 640).into())
//! .with_i03((0, 640, 640).into())
//! .with_confs(&[0.2]);
//! let mut model = YOLO::new(options)?;
//!
//! // Build DataLoader to load image(s), video, stream
//! let dl = DataLoader::new(
//! "./assets/bus.jpg", // local image
//! // "images/bus.jpg", // remote image
//! // "../set-negs", // local images (from folder)
//! // "../hall.mp4", // local video
//! // "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // remote video
//! // "rtsp://admin:kkasd1234@192.168.2.217:554/h264/ch1/", // stream
//! )?
//! .with_batch(3) // iterate with batch_size = 3
//! .build()?;
//!
//! // Build annotator
//! let annotator = Annotator::new().with_saveout("YOLO-Demo");
//!
//! // Run and Annotate images
//! for (xs, _) in dl {
//! let ys = model.forward(&xs, false)?;
//! annotator.annotate(&xs, &ys);
//! }
//!
//! Ok(())
//! }
//! - Download from the [Releases page](https://github.com/microsoft/onnxruntime/releases).
//! - Set up the library path by exporting the `ORT_DYLIB_PATH` environment variable:
//! ```shell
//! export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.20.1
//! ```
//!
//! </details>
//! - **`ffmpeg`**: Adds support for video streams, real-time frame visualization, and video export.
//!
//! - Powered by [video-rs](https://github.com/oddity-ai/video-rs) and [minifb](https://github.com/emoon/rust_minifb). For any issues related to `ffmpeg` features, please refer to the issues of these two crates.
//! - **`cuda`**: Enables the NVIDIA TensorRT provider.
//! - **`trt`**: Enables the NVIDIA TensorRT provider.
//! - **`mps`**: Enables the Apple CoreML provider.
//!
//! ## 🎈 Example
//!
//! ```Shell
//! cargo run -r -F cuda --example svtr -- --device cuda
//! ```
//!
//! All examples are located in the [examples](https://github.com/jamjamjon/usls/tree/main/examples) directory.
//! # What's More
//!
//! This guide covers the process of using provided models for inference, including how to build a model, load data, annotate results, and retrieve the outputs. Click the sections below to expand for detailed instructions.
//!
//! <details>
//! <summary>Build the Model</summary>
//!
//! To build a model, you can use the provided [models] with [Options]:
//!
//! ```ignore
//! use usls::{models::YOLO, Annotator, DataLoader, Options, Vision};
//!
//! let options = Options::default()
//! .with_yolo_version(YOLOVersion::V8) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
//! .with_yolo_task(YOLOTask::Detect) // YOLOTask: Classify, Detect, Pose, Segment, Obb
//! .with_model("xxxx.onnx")?;
//! let mut model = YOLO::new(options)?;
//! ```
//!
//! **And there're many options provided by [Options]**
//!
//! - **Choose Execution Provider:**
//! Select `CUDA` (default), `TensorRT`, or `CoreML`:
//!
//! ```ignore
//! let options = Options::default()
//! .with_cuda(0)
//! // .with_trt(0)
//! // .with_coreml(0)
//! // .with_cpu();
//! ```
//!
//! - **Dynamic Input Shapes:**
//! Specify dynamic shapes with [MinOptMax]:
//!
//! ```ignore
//! let options = Options::default()
//! .with_i00((1, 2, 4).into()) // batch(min=1, opt=2, max=4)
//! .with_i02((416, 640, 800).into()) // height(min=416, opt=640, max=800)
//! .with_i03((416, 640, 800).into()); // width(min=416, opt=640, max=800)
//! ```
//!
//! - **Set Confidence Thresholds:**
//! Adjust thresholds for each category:
//!
//! ```ignore
//! let options = Options::default()
//! .with_confs(&[0.4, 0.15]); // class_0: 0.4, others: 0.15
//! ```
//!
//! - **Set Class Names:**
//! Provide class names if needed:
//!
//! ```ignore
//! let options = Options::default()
//! .with_names(&COCO_CLASS_NAMES_80);
//! ```
//!
//! **More options are detailed in the [Options] documentation.**
//!
//!
//! </details>
//!
//! <details>
//! <summary>Load Images, Video and Stream</summary>
//!
//! - **Load a Single Image**
//! Use [DataLoader::try_read] to load an image from a local file or remote source:
//!
//! ```ignore
//! let x = DataLoader::try_read("./assets/bus.jpg")?; // from local
//! let x = DataLoader::try_read("images/bus.jpg")?; // from remote
//! ```
//!
//! Alternatively, use [image::ImageReader] directly:
//!
//! ```ignore
//! let x = image::ImageReader::open("myimage.png")?.decode()?;
//! ```
//!
//! - **Load Multiple Images, Videos, or Streams**
//! Create a [DataLoader] instance for batch processing:
//!
//! ```ignore
//! let dl = DataLoader::new(
//! "./assets/bus.jpg", // local image
//! // "images/bus.jpg", // remote image
//! // "../set-negs", // local images (from folder)
//! // "../hall.mp4", // local video
//! // "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4", // remote video
//! // "rtsp://admin:kkasd1234@192.168.2.217:554/h264/ch1/", // stream
//! )?
//! .with_batch(3) // iterate with batch_size = 3
//! .build()?;
//!
//! // Iterate through the data
//! for (xs, _) in dl {}
//! ```
//!
//! - **Convert Images to Video**
//! Use [DataLoader::is2v] to create a video from a sequence of images:
//!
//! ```ignore
//! let fps = 24;
//! let image_folder = "runs/YOLO-DataLoader";
//! let saveout = ["runs", "is2v"];
//! DataLoader::is2v(image_folder, &saveout, fps)?;
//! ```
//!
//! </details>
//!
//! <details>
//! <summary>Annotate Inference Results</summary>
//!
//! - **Create an Annotator Instance**
//!
//! ```ignore
//! let annotator = Annotator::default();
//! ```
//!
//! - **Set Saveout Name:**
//!
//! ```ignore
//! let annotator = Annotator::default()
//! .with_saveout("YOLOs");
//! ```
//!
//! - **Set Bounding Box Line Width:**
//!
//! ```ignore
//! let annotator = Annotator::default()
//! .with_bboxes_thickness(4);
//! ```
//!
//! - **Disable Mask Plotting**
//!
//! ```ignore
//! let annotator = Annotator::default()
//! .without_masks(true);
//! ```
//!
//! - **Perform Inference and nnotate the results**
//!
//! ```ignore
//! for (xs, _paths) in dl {
//! let ys = model.run(&xs)?;
//! annotator.annotate(&xs, &ys);
//! }
//! ```
//!
//! More options are detailed in the [Annotator] documentation.
//!
//! </details>
//!
//! <details>
//! <summary>Retrieve Model's Inference Results</summary>
//!
//! Retrieve the inference outputs, which are saved in a [`Vec<Y>`]:
//!
//! - **Get Detection Bounding Boxes**
//!
//! ```ignore
//! let ys = model.run(&xs)?;
//! for y in ys {
//! // bboxes
//! if let Some(bboxes) = y.bboxes() {
//! for bbox in bboxes {
//! println!(
//! "Bbox: {}, {}, {}, {}, {}, {}",
//! bbox.xmin(),
//! bbox.ymin(),
//! bbox.xmax(),
//! bbox.ymax(),
//! bbox.confidence(),
//! bbox.id(),
//! );
//! }
//! }
//! }
//! ```
//!
//! </details>
//!
//! <details>
//! <summary>Custom Model Implementation</summary>
//!
//! You can also implement your own model using [OrtEngine] and [Options]. [OrtEngine] supports ONNX model loading, metadata parsing, dry_run, inference, and other functions, with execution providers such as CUDA, TensorRT, CoreML, etc.
//!
//! For more details, refer to the [Demo: Depth-Anything](https://github.com/jamjamjon/usls/blob/main/src/models/depth_anything.rs).
//!
//! </details>
mod core;
mod misc;
pub mod models;
mod utils;
mod ys;
mod xy;
pub use core::*;
pub use misc::*;
pub use models::*;
pub use utils::*;
pub use ys::*;
pub use xy::*;

80
src/core/annotator.rs → src/misc/annotator.rs
View File

@ -1,16 +1,15 @@
use crate::{
colormap256, string_now, Bbox, Dir, Hub, Keypoint, Mask, Mbr, Polygon, Prob, CHECK_MARK,
CROSS_MARK, Y,
};
use ab_glyph::{FontArc, PxScale};
use anyhow::Result;
use image::{DynamicImage, GenericImage, Rgba, RgbaImage};
use imageproc::map::map_colors;
use crate::{
string_now, Bbox, Color, ColorMap256, Dir, Hub, Keypoint, Mask, Mbr, Polygon, Prob, Y,
};
/// Annotator for struct `Y`
#[derive(Clone)]
pub struct Annotator {
// TODO: Add lifetime
font: FontArc,
_scale: f32, // Cope with ab_glyph & imageproc=0.24.0
scale_dy: f32,
@ -18,6 +17,7 @@ pub struct Annotator {
saveout: Option<String>,
saveout_subs: Vec<String>,
decimal_places: usize,
palette: Vec<Color>,
// About mbrs
without_mbrs: bool,
@ -57,7 +57,7 @@ pub struct Annotator {
// About masks
without_masks: bool,
colormap: Option<[[u8; 3]; 256]>,
colormap: Option<[Color; 256]>,
// About probs
probs_topk: usize,
@ -73,6 +73,7 @@ impl Default for Annotator {
_scale: 6.666667,
scale_dy: 28.,
polygons_alpha: 179,
palette: Color::palette_base_20(),
saveout: None,
saveout_subs: vec![],
saveout_base: String::from("runs"),
@ -272,22 +273,8 @@ impl Annotator {
}
pub fn with_colormap(mut self, x: &str) -> Self {
let x = match x {
"turbo" | "Turbo" | "TURBO" => colormap256::TURBO,
"inferno" | "Inferno" | "INFERNO" => colormap256::INFERNO,
"plasma" | "Plasma" | "PLASMA" => colormap256::PLASMA,
"viridis" | "Viridis" | "VIRIDIS" => colormap256::VIRIDIS,
"magma" | "Magma" | "MAGMA" => colormap256::MAGMA,
"bentcoolwarm" | "BentCoolWarm" | "BENTCOOLWARM" => colormap256::BENTCOOLWARM,
"blackbody" | "BlackBody" | "BLACKBODY" => colormap256::BLACKBODY,
"extendedkindLmann" | "ExtendedKindLmann" | "EXTENDEDKINDLMANN" => {
colormap256::EXTENDEDKINDLMANN
}
"kindlmann" | "KindLmann" | "KINDLMANN" => colormap256::KINDLMANN,
"smoothcoolwarm" | "SmoothCoolWarm" | "SMOOTHCOOLWARM" => colormap256::SMOOTHCOOLWARM,
_ => todo!(),
};
self.colormap = Some(x);
let x = ColorMap256::from(x);
self.colormap = Some(x.data());
self
}
@ -355,7 +342,7 @@ impl Annotator {
}
// mkdir even no filename specified
Dir::Currnet.raw_path_with_subs(&subs)
Dir::Current.raw_path_with_subs(&subs)
}
/// Annotate images, save, and no return
@ -365,9 +352,6 @@ impl Annotator {
/// Plot images and return plotted images
pub fn plot(&self, imgs: &[DynamicImage], ys: &[Y], save: bool) -> Result<Vec<DynamicImage>> {
let span = tracing::span!(tracing::Level::INFO, "Annotator-plot");
let _guard = span.enter();
let mut vs: Vec<DynamicImage> = Vec::new();
// annotate
@ -418,9 +402,9 @@ impl Annotator {
if save {
let saveout = self.saveout()?.join(format!("{}.png", string_now("-")));
match img_rgba.save(&saveout) {
Err(err) => tracing::error!("{} Saving failed: {:?}", CROSS_MARK, err),
Err(err) => anyhow::bail!("Failed to save annotated image: {:?}", err),
Ok(_) => {
tracing::info!("{} Annotated image saved to: {:?}", CHECK_MARK, saveout);
println!("Annotated image saved to: {:?}", saveout);
}
}
}
@ -428,6 +412,7 @@ impl Annotator {
// RgbaImage -> DynamicImage
vs.push(image::DynamicImage::from(img_rgba));
}
Ok(vs)
}
@ -673,7 +658,7 @@ impl Annotator {
let luma = if let Some(colormap) = self.colormap {
let luma = map_colors(mask.mask(), |p| {
let x = p[0];
image::Rgb(colormap[x as usize])
image::Rgb(colormap[x as usize].rgb().into())
});
image::DynamicImage::from(luma)
} else {
@ -774,42 +759,15 @@ impl Annotator {
/// Load custom font
fn load_font(path: Option<&str>) -> Result<FontArc> {
let path_font = match path {
None => Hub::new()?.fetch("fonts/Arial.ttf")?.commit()?,
None => Hub::default().try_fetch("fonts/Arial.ttf")?,
Some(p) => p.into(),
};
let buffer = std::fs::read(path_font)?;
Ok(FontArc::try_from_vec(buffer.to_owned())?)
let buf = std::fs::read(path_font)?;
Ok(FontArc::try_from_vec(buf.to_owned())?)
}
/// Pick color from pallette
/// Color palette
pub fn get_color(&self, n: usize) -> (u8, u8, u8, u8) {
Self::color_palette()[n % Self::color_palette().len()]
}
/// Color pallette
fn color_palette() -> [(u8, u8, u8, u8); 20] {
// TODO: more colors
[
(0, 255, 127, 255), // spring green
(255, 105, 180, 255), // hot pink
(255, 99, 71, 255), // tomato
(255, 215, 0, 255), // glod
(188, 143, 143, 255), // rosy brown
(0, 191, 255, 255), // deep sky blue
(143, 188, 143, 255), // dark sea green
(238, 130, 238, 255), // violet
(154, 205, 50, 255), // yellow green
(205, 133, 63, 255), // peru
(30, 144, 255, 255), // dodger blue
(112, 128, 144, 255), // slate gray
(127, 255, 212, 255), // aqua marine
(51, 153, 255, 255), // blue
(0, 255, 255, 255), // cyan
(138, 43, 226, 255), // blue violet
(165, 42, 42, 255), // brown
(216, 191, 216, 255), // thistle
(240, 255, 255, 255), // azure
(95, 158, 160, 255), // cadet blue
]
self.palette[n % self.palette.len()].rgba()
}
}

171
src/misc/color.rs Normal file
View File

@ -0,0 +1,171 @@
use anyhow::Result;
use rand::Rng;
/// Color: 0xRRGGBBAA
#[derive(Copy, Clone)]
pub struct Color(u32);
impl std::fmt::Debug for Color {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Color")
.field("RGBA", &self.rgba())
.field("HEX", &self.hex())
.finish()
}
}
impl std::fmt::Display for Color {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.hex())
}
}
impl From<u32> for Color {
fn from(x: u32) -> Self {
Self(x)
}
}
impl From<(u8, u8, u8)> for Color {
fn from((r, g, b): (u8, u8, u8)) -> Self {
Self::from_rgba(r, g, b, 0xff)
}
}
impl From<[u8; 3]> for Color {
fn from(c: [u8; 3]) -> Self {
Self::from((c[0], c[1], c[2]))
}
}
impl From<(u8, u8, u8, u8)> for Color {
fn from((r, g, b, a): (u8, u8, u8, u8)) -> Self {
Self::from_rgba(r, g, b, a)
}
}
impl From<[u8; 4]> for Color {
fn from(c: [u8; 4]) -> Self {
Self::from((c[0], c[1], c[2], c[3]))
}
}
impl TryFrom<&str> for Color {
type Error = &'static str;
fn try_from(x: &str) -> Result<Self, Self::Error> {
let hex = x.trim_start_matches('#');
let hex = match hex.len() {
6 => format!("{}ff", hex),
8 => hex.to_string(),
_ => return Err("Failed to convert `Color` from str: invalid length"),
};
u32::from_str_radix(&hex, 16)
.map(Self)
.map_err(|_| "Failed to convert `Color` from str: invalid hex")
}
}
impl Color {
const fn from_rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
Self(((r as u32) << 24) | ((g as u32) << 16) | ((b as u32) << 8) | (a as u32))
}
pub fn rgba(&self) -> (u8, u8, u8, u8) {
let r = ((self.0 >> 24) & 0xff) as u8;
let g = ((self.0 >> 16) & 0xff) as u8;
let b = ((self.0 >> 8) & 0xff) as u8;
let a = (self.0 & 0xff) as u8;
(r, g, b, a)
}
pub fn rgb(&self) -> (u8, u8, u8) {
let (r, g, b, _) = self.rgba();
(r, g, b)
}
pub fn bgr(&self) -> (u8, u8, u8) {
let (r, g, b) = self.rgb();
(b, g, r)
}
pub fn hex(&self) -> String {
format!("#{:08x}", self.0)
}
pub fn create_palette<A: Into<Self> + Copy>(xs: &[A]) -> Vec<Self> {
xs.iter().copied().map(Into::into).collect()
}
pub fn try_create_palette<A: TryInto<Self> + Copy>(xs: &[A]) -> Result<Vec<Self>>
where
<A as TryInto<Self>>::Error: std::fmt::Debug,
{
xs.iter()
.copied()
.map(|x| {
x.try_into()
.map_err(|e| anyhow::anyhow!("Failed to convert: {:?}", e))
})
.collect()
}
pub fn palette_rand(n: usize) -> Vec<Self> {
let mut rng = rand::thread_rng();
let xs: Vec<(u8, u8, u8)> = (0..n)
.map(|_| {
(
rng.gen_range(0..=255),
rng.gen_range(0..=255),
rng.gen_range(0..=255),
)
})
.collect();
Self::create_palette(&xs)
}
pub fn palette_base_20() -> Vec<Self> {
Self::create_palette(&[
0x00ff7fff, // SpringGreen
0xff69b4ff, // HotPink
0xff6347ff, // Tomato
0xffd700ff, // Gold
0xbc8f8fff, // RosyBrown
0x00bfffff, // DeepSkyBlue
0x8fb88fff, // DarkSeaGreen
0xee82eeff, // Violet
0x9acd32ff, // YellowGreen
0xcd853fff, // Peru
0x1e90ffff, // DodgerBlue
0xd74a49ff, // ?
0x7fffd4ff, // AquaMarine
0x3399ffff, // Blue2
0x00ffffff, // Cyan
0x8a2befff, // BlueViolet
0xa52a2aff, // Brown
0xd8bfd8ff, // Thistle
0xf0ffffff, // Azure
0x609ea0ff, // CadetBlue
])
}
pub fn palette_cotton_candy_5() -> Vec<Self> {
Self::try_create_palette(&["#ff595e", "#ffca3a", "#8ac926", "#1982c4", "#6a4c93"])
.expect("Faild to create palette: Cotton Candy")
}
pub fn palette_tropical_sunrise_5() -> Vec<Self> {
// https://colorkit.co/palette/e12729-f37324-f8cc1b-72b043-007f4e/
Self::try_create_palette(&["#e12729", "#f37324", "#f8cc1b", "#72b043", "#007f4e"])
.expect("Faild to create palette: Tropical Sunrise")
}
pub fn palette_rainbow_10() -> Vec<Self> {
Self::create_palette(&[
0xff595eff, 0xff924cff, 0xffca3aff, 0xc5ca30ff, 0x8ac926ff, 0x52a675ff, 0x1982c4ff,
0x4267acff, 0x6a4c93ff, 0xb5a6c9ff,
])
}
}

435
src/misc/colormap256.rs Normal file
View File

@ -0,0 +1,435 @@
use crate::Color;
pub enum ColorMap256 {
Turbo,
Inferno,
Plasma,
Viridis,
Magma,
BentCoolWarm,
BlackBody,
ExtendedKindLmann,
KindLmann,
SmoothCoolWarm,
}
impl From<&str> for ColorMap256 {
fn from(s: &str) -> Self {
match s {
"turbo" | "Turbo" | "TURBO" => Self::Turbo,
"inferno" | "Inferno" | "INFERNO" => Self::Inferno,
"plasma" | "Plasma" | "PLASMA" => Self::Plasma,
"viridis" | "Viridis" | "VIRIDIS" => Self::Viridis,
"magma" | "Magma" | "MAGMA" => Self::Magma,
"bentcoolwarm" | "BentCoolWarm" | "BENTCOOLWARM" => Self::BentCoolWarm,
"blackbody" | "BlackBody" | "BLACKBODY" => Self::BlackBody,
"extendedkindLmann" | "ExtendedKindLmann" | "EXTENDEDKINDLMANN" => {
Self::ExtendedKindLmann
}
"kindlmann" | "KindLmann" | "KINDLMANN" => Self::KindLmann,
"smoothcoolwarm" | "SmoothCoolWarm" | "SMOOTHCOOLWARM" => Self::SmoothCoolWarm,
_ => todo!(),
}
}
}
impl ColorMap256 {
pub fn data(&self) -> [Color; 256] {
let xs = match self {
Self::Turbo => [
0x30123bff, 0x321543ff, 0x33184aff, 0x341b51ff, 0x351e58ff, 0x36215fff, 0x372466ff,
0x38276dff, 0x392a73ff, 0x3a2d79ff, 0x3b2f80ff, 0x3c3286ff, 0x3d358bff, 0x3e3891ff,
0x3f3b97ff, 0x3f3e9cff, 0x4040a2ff, 0x4143a7ff, 0x4146acff, 0x4249b1ff, 0x424bb5ff,
0x434ebaff, 0x4451bfff, 0x4454c3ff, 0x4456c7ff, 0x4559cbff, 0x455ccfff, 0x455ed3ff,
0x4661d6ff, 0x4664daff, 0x4666ddff, 0x4669e0ff, 0x466be3ff, 0x476ee6ff, 0x4771e9ff,
0x4773ebff, 0x4776eeff, 0x4778f0ff, 0x477bf2ff, 0x467df4ff, 0x4680f6ff, 0x4682f8ff,
0x4685faff, 0x4687fbff, 0x458afcff, 0x458cfdff, 0x448ffeff, 0x4391feff, 0x4294ffff,
0x4196ffff, 0x4099ffff, 0x3e9bfeff, 0x3d9efeff, 0x3ba0fdff, 0x3aa3fcff, 0x38a5fbff,
0x37a8faff, 0x35abf8ff, 0x33adf7ff, 0x31aff5ff, 0x2fb2f4ff, 0x2eb4f2ff, 0x2cb7f0ff,
0x2ab9eeff, 0x28bcebff, 0x27bee9ff, 0x25c0e7ff, 0x23c3e4ff, 0x22c5e2ff, 0x20c7dfff,
0x1fc9ddff, 0x1ecbdaff, 0x1ccdd8ff, 0x1bd0d5ff, 0x1ad2d2ff, 0x1ad4d0ff, 0x19d5cdff,
0x18d7caff, 0x18d9c8ff, 0x18dbc5ff, 0x18ddc2ff, 0x18dec0ff, 0x18e0bdff, 0x19e2bbff,
0x19e3b9ff, 0x1ae4b6ff, 0x1ce6b4ff, 0x1de7b2ff, 0x1fe9afff, 0x20eaacff, 0x22ebaaff,
0x25eca7ff, 0x27eea4ff, 0x2aefa1ff, 0x2cf09eff, 0x2ff19bff, 0x32f298ff, 0x35f394ff,
0x38f491ff, 0x3cf58eff, 0x3ff68aff, 0x43f787ff, 0x46f884ff, 0x4af880ff, 0x4ef97dff,
0x52fa7aff, 0x55fa76ff, 0x59fb73ff, 0x5dfc6fff, 0x61fc6cff, 0x65fd69ff, 0x69fd66ff,
0x6dfe62ff, 0x71fe5fff, 0x75fe5cff, 0x79fe59ff, 0x7dff56ff, 0x80ff53ff, 0x84ff51ff,
0x88ff4eff, 0x8bff4bff, 0x8fff49ff, 0x92ff47ff, 0x96fe44ff, 0x99fe42ff, 0x9cfe40ff,
0x9ffd3fff, 0xa1fd3dff, 0xa4fc3cff, 0xa7fc3aff, 0xa9fb39ff, 0xacfb38ff, 0xaffa37ff,
0xb1f936ff, 0xb4f836ff, 0xb7f735ff, 0xb9f635ff, 0xbcf534ff, 0xbef434ff, 0xc1f334ff,
0xc3f134ff, 0xc6f034ff, 0xc8ef34ff, 0xcbed34ff, 0xcdec34ff, 0xd0ea34ff, 0xd2e935ff,
0xd4e735ff, 0xd7e535ff, 0xd9e436ff, 0xdbe236ff, 0xdde037ff, 0xdfdf37ff, 0xe1dd37ff,
0xe3db38ff, 0xe5d938ff, 0xe7d739ff, 0xe9d539ff, 0xebd339ff, 0xecd13aff, 0xeecf3aff,
0xefcd3aff, 0xf1cb3aff, 0xf2c93aff, 0xf4c73aff, 0xf5c53aff, 0xf6c33aff, 0xf7c13aff,
0xf8be39ff, 0xf9bc39ff, 0xfaba39ff, 0xfbb838ff, 0xfbb637ff, 0xfcb336ff, 0xfcb136ff,
0xfdae35ff, 0xfdac34ff, 0xfea933ff, 0xfea732ff, 0xfea431ff, 0xfea130ff, 0xfe9e2fff,
0xfe9b2dff, 0xfe992cff, 0xfe962bff, 0xfe932aff, 0xfe9029ff, 0xfd8d27ff, 0xfd8a26ff,
0xfc8725ff, 0xfc8423ff, 0xfb8122ff, 0xfb7e21ff, 0xfa7b1fff, 0xf9781eff, 0xf9751dff,
0xf8721cff, 0xf76f1aff, 0xf66c19ff, 0xf56918ff, 0xf46617ff, 0xf36315ff, 0xf26014ff,
0xf15d13ff, 0xf05b12ff, 0xef5811ff, 0xed5510ff, 0xec530fff, 0xeb500eff, 0xea4e0dff,
0xe84b0cff, 0xe7490cff, 0xe5470bff, 0xe4450aff, 0xe2430aff, 0xe14109ff, 0xdf3f08ff,
0xdd3d08ff, 0xdc3b07ff, 0xda3907ff, 0xd83706ff, 0xd63506ff, 0xd43305ff, 0xd23105ff,
0xd02f05ff, 0xce2d04ff, 0xcc2b04ff, 0xca2a04ff, 0xc82803ff, 0xc52603ff, 0xc32503ff,
0xc12302ff, 0xbe2102ff, 0xbc2002ff, 0xb91e02ff, 0xb71d02ff, 0xb41b01ff, 0xb21a01ff,
0xaf1801ff, 0xac1701ff, 0xa91601ff, 0xa71401ff, 0xa41301ff, 0xa11201ff, 0x9e1001ff,
0x9b0f01ff, 0x980e01ff, 0x950d01ff, 0x920b01ff, 0x8e0a01ff, 0x8b0902ff, 0x880802ff,
0x850702ff, 0x810602ff, 0x7e0502ff, 0x7a0403ff,
],
Self::Inferno => [
0x000004ff, 0x010005ff, 0x010106ff, 0x010108ff, 0x02010aff, 0x02020cff, 0x02020eff,
0x030210ff, 0x040312ff, 0x040314ff, 0x050417ff, 0x060419ff, 0x07051bff, 0x08051dff,
0x09061fff, 0x0a0722ff, 0x0b0724ff, 0x0c0826ff, 0x0d0829ff, 0x0e092bff, 0x10092dff,
0x110a30ff, 0x120a32ff, 0x140b34ff, 0x150b37ff, 0x160b39ff, 0x180c3cff, 0x190c3eff,
0x1b0c41ff, 0x1c0c43ff, 0x1e0c45ff, 0x1f0c48ff, 0x210c4aff, 0x230c4cff, 0x240c4fff,
0x260c51ff, 0x280b53ff, 0x290b55ff, 0x2b0b57ff, 0x2d0b59ff, 0x2f0a5bff, 0x310a5cff,
0x320a5eff, 0x340a5fff, 0x360961ff, 0x380962ff, 0x390963ff, 0x3b0964ff, 0x3d0965ff,
0x3e0966ff, 0x400a67ff, 0x420a68ff, 0x440a68ff, 0x450a69ff, 0x470b6aff, 0x490b6aff,
0x4a0c6bff, 0x4c0c6bff, 0x4d0d6cff, 0x4f0d6cff, 0x510e6cff, 0x520e6dff, 0x540f6dff,
0x550f6dff, 0x57106eff, 0x59106eff, 0x5a116eff, 0x5c126eff, 0x5d126eff, 0x5f136eff,
0x61136eff, 0x62146eff, 0x64156eff, 0x65156eff, 0x67166eff, 0x69166eff, 0x6a176eff,
0x6c186eff, 0x6d186eff, 0x6f196eff, 0x71196eff, 0x721a6eff, 0x741a6eff, 0x751b6eff,
0x771c6dff, 0x781c6dff, 0x7a1d6dff, 0x7c1d6dff, 0x7d1e6dff, 0x7f1e6cff, 0x801f6cff,
0x82206cff, 0x84206bff, 0x85216bff, 0x87216bff, 0x88226aff, 0x8a226aff, 0x8c2369ff,
0x8d2369ff, 0x8f2469ff, 0x902568ff, 0x922568ff, 0x932667ff, 0x952667ff, 0x972766ff,
0x982766ff, 0x9a2865ff, 0x9b2964ff, 0x9d2964ff, 0x9f2a63ff, 0xa02a63ff, 0xa22b62ff,
0xa32c61ff, 0xa52c60ff, 0xa62d60ff, 0xa82e5fff, 0xa92e5eff, 0xab2f5eff, 0xad305dff,
0xae305cff, 0xb0315bff, 0xb1325aff, 0xb3325aff, 0xb43359ff, 0xb63458ff, 0xb73557ff,
0xb93556ff, 0xba3655ff, 0xbc3754ff, 0xbd3853ff, 0xbf3952ff, 0xc03a51ff, 0xc13a50ff,
0xc33b4fff, 0xc43c4eff, 0xc63d4dff, 0xc73e4cff, 0xc83f4bff, 0xca404aff, 0xcb4149ff,
0xcc4248ff, 0xce4347ff, 0xcf4446ff, 0xd04545ff, 0xd24644ff, 0xd34743ff, 0xd44842ff,
0xd54a41ff, 0xd74b3fff, 0xd84c3eff, 0xd94d3dff, 0xda4e3cff, 0xdb503bff, 0xdd513aff,
0xde5238ff, 0xdf5337ff, 0xe05536ff, 0xe15635ff, 0xe25734ff, 0xe35933ff, 0xe45a31ff,
0xe55c30ff, 0xe65d2fff, 0xe75e2eff, 0xe8602dff, 0xe9612bff, 0xea632aff, 0xeb6429ff,
0xeb6628ff, 0xec6726ff, 0xed6925ff, 0xee6a24ff, 0xef6c23ff, 0xef6e21ff, 0xf06f20ff,
0xf1711fff, 0xf1731dff, 0xf2741cff, 0xf3761bff, 0xf37819ff, 0xf47918ff, 0xf57b17ff,
0xf57d15ff, 0xf67e14ff, 0xf68013ff, 0xf78212ff, 0xf78410ff, 0xf8850fff, 0xf8870eff,
0xf8890cff, 0xf98b0bff, 0xf98c0aff, 0xf98e09ff, 0xfa9008ff, 0xfa9207ff, 0xfa9407ff,
0xfb9606ff, 0xfb9706ff, 0xfb9906ff, 0xfb9b06ff, 0xfb9d07ff, 0xfc9f07ff, 0xfca108ff,
0xfca309ff, 0xfca50aff, 0xfca60cff, 0xfca80dff, 0xfcaa0fff, 0xfcac11ff, 0xfcae12ff,
0xfcb014ff, 0xfcb216ff, 0xfcb418ff, 0xfbb61aff, 0xfbb81dff, 0xfbba1fff, 0xfbbc21ff,
0xfbbe23ff, 0xfac026ff, 0xfac228ff, 0xfac42aff, 0xfac62dff, 0xf9c72fff, 0xf9c932ff,
0xf9cb35ff, 0xf8cd37ff, 0xf8cf3aff, 0xf7d13dff, 0xf7d340ff, 0xf6d543ff, 0xf6d746ff,
0xf5d949ff, 0xf5db4cff, 0xf4dd4fff, 0xf4df53ff, 0xf4e156ff, 0xf3e35aff, 0xf3e55dff,
0xf2e661ff, 0xf2e865ff, 0xf2ea69ff, 0xf1ec6dff, 0xf1ed71ff, 0xf1ef75ff, 0xf1f179ff,
0xf2f27dff, 0xf2f482ff, 0xf3f586ff, 0xf3f68aff, 0xf4f88eff, 0xf5f992ff, 0xf6fa96ff,
0xf8fb9aff, 0xf9fc9dff, 0xfafda1ff, 0xfcffa4ff,
],
Self::Plasma => [
0x0d0887ff, 0x100788ff, 0x130789ff, 0x16078aff, 0x19068cff, 0x1b068dff, 0x1d068eff,
0x20068fff, 0x220690ff, 0x240691ff, 0x260591ff, 0x280592ff, 0x2a0593ff, 0x2c0594ff,
0x2e0595ff, 0x2f0596ff, 0x310597ff, 0x330597ff, 0x350498ff, 0x370499ff, 0x38049aff,
0x3a049aff, 0x3c049bff, 0x3e049cff, 0x3f049cff, 0x41049dff, 0x43039eff, 0x44039eff,
0x46039fff, 0x48039fff, 0x4903a0ff, 0x4b03a1ff, 0x4c02a1ff, 0x4e02a2ff, 0x5002a2ff,
0x5102a3ff, 0x5302a3ff, 0x5502a4ff, 0x5601a4ff, 0x5801a4ff, 0x5901a5ff, 0x5b01a5ff,
0x5c01a6ff, 0x5e01a6ff, 0x6001a6ff, 0x6100a7ff, 0x6300a7ff, 0x6400a7ff, 0x6600a7ff,
0x6700a8ff, 0x6900a8ff, 0x6a00a8ff, 0x6c00a8ff, 0x6e00a8ff, 0x6f00a8ff, 0x7100a8ff,
0x7201a8ff, 0x7401a8ff, 0x7501a8ff, 0x7701a8ff, 0x7801a8ff, 0x7a02a8ff, 0x7b02a8ff,
0x7d03a8ff, 0x7e03a8ff, 0x8004a8ff, 0x8104a7ff, 0x8305a7ff, 0x8405a7ff, 0x8606a6ff,
0x8707a6ff, 0x8808a6ff, 0x8a09a5ff, 0x8b0aa5ff, 0x8d0ba5ff, 0x8e0ca4ff, 0x8f0da4ff,
0x910ea3ff, 0x920fa3ff, 0x9410a2ff, 0x9511a1ff, 0x9613a1ff, 0x9814a0ff, 0x99159fff,
0x9a169fff, 0x9c179eff, 0x9d189dff, 0x9e199dff, 0xa01a9cff, 0xa11b9bff, 0xa21d9aff,
0xa31e9aff, 0xa51f99ff, 0xa62098ff, 0xa72197ff, 0xa82296ff, 0xaa2395ff, 0xab2494ff,
0xac2694ff, 0xad2793ff, 0xae2892ff, 0xb02991ff, 0xb12a90ff, 0xb22b8fff, 0xb32c8eff,
0xb42e8dff, 0xb52f8cff, 0xb6308bff, 0xb7318aff, 0xb83289ff, 0xba3388ff, 0xbb3488ff,
0xbc3587ff, 0xbd3786ff, 0xbe3885ff, 0xbf3984ff, 0xc03a83ff, 0xc13b82ff, 0xc23c81ff,
0xc33d80ff, 0xc43e7fff, 0xc5407eff, 0xc6417dff, 0xc7427cff, 0xc8437bff, 0xc9447aff,
0xca457aff, 0xcb4679ff, 0xcc4778ff, 0xcc4977ff, 0xcd4a76ff, 0xce4b75ff, 0xcf4c74ff,
0xd04d73ff, 0xd14e72ff, 0xd24f71ff, 0xd35171ff, 0xd45270ff, 0xd5536fff, 0xd5546eff,
0xd6556dff, 0xd7566cff, 0xd8576bff, 0xd9586aff, 0xda5a6aff, 0xda5b69ff, 0xdb5c68ff,
0xdc5d67ff, 0xdd5e66ff, 0xde5f65ff, 0xde6164ff, 0xdf6263ff, 0xe06363ff, 0xe16462ff,
0xe26561ff, 0xe26660ff, 0xe3685fff, 0xe4695eff, 0xe56a5dff, 0xe56b5dff, 0xe66c5cff,
0xe76e5bff, 0xe76f5aff, 0xe87059ff, 0xe97158ff, 0xe97257ff, 0xea7457ff, 0xeb7556ff,
0xeb7655ff, 0xec7754ff, 0xed7953ff, 0xed7a52ff, 0xee7b51ff, 0xef7c51ff, 0xef7e50ff,
0xf07f4fff, 0xf0804eff, 0xf1814dff, 0xf1834cff, 0xf2844bff, 0xf3854bff, 0xf3874aff,
0xf48849ff, 0xf48948ff, 0xf58b47ff, 0xf58c46ff, 0xf68d45ff, 0xf68f44ff, 0xf79044ff,
0xf79143ff, 0xf79342ff, 0xf89441ff, 0xf89540ff, 0xf9973fff, 0xf9983eff, 0xf99a3eff,
0xfa9b3dff, 0xfa9c3cff, 0xfa9e3bff, 0xfb9f3aff, 0xfba139ff, 0xfba238ff, 0xfca338ff,
0xfca537ff, 0xfca636ff, 0xfca835ff, 0xfca934ff, 0xfdab33ff, 0xfdac33ff, 0xfdae32ff,
0xfdaf31ff, 0xfdb130ff, 0xfdb22fff, 0xfdb42fff, 0xfdb52eff, 0xfeb72dff, 0xfeb82cff,
0xfeba2cff, 0xfebb2bff, 0xfebd2aff, 0xfebe2aff, 0xfec029ff, 0xfdc229ff, 0xfdc328ff,
0xfdc527ff, 0xfdc627ff, 0xfdc827ff, 0xfdca26ff, 0xfdcb26ff, 0xfccd25ff, 0xfcce25ff,
0xfcd025ff, 0xfcd225ff, 0xfbd324ff, 0xfbd524ff, 0xfbd724ff, 0xfad824ff, 0xfada24ff,
0xf9dc24ff, 0xf9dd25ff, 0xf8df25ff, 0xf8e125ff, 0xf7e225ff, 0xf7e425ff, 0xf6e626ff,
0xf6e826ff, 0xf5e926ff, 0xf5eb27ff, 0xf4ed27ff, 0xf3ee27ff, 0xf3f027ff, 0xf2f227ff,
0xf1f426ff, 0xf1f525ff, 0xf0f724ff, 0xf0f921ff,
],
Self::Viridis => [
0x440154ff, 0x440256ff, 0x450457ff, 0x450559ff, 0x46075aff, 0x46085cff, 0x460a5dff,
0x460b5eff, 0x470d60ff, 0x470e61ff, 0x471063ff, 0x471164ff, 0x471365ff, 0x481467ff,
0x481668ff, 0x481769ff, 0x48186aff, 0x481a6cff, 0x481b6dff, 0x481c6eff, 0x481d6fff,
0x481f70ff, 0x482071ff, 0x482173ff, 0x482374ff, 0x482475ff, 0x482576ff, 0x482677ff,
0x482878ff, 0x482979ff, 0x472a7aff, 0x472c7aff, 0x472d7bff, 0x472e7cff, 0x472f7dff,
0x46307eff, 0x46327eff, 0x46337fff, 0x463480ff, 0x453581ff, 0x453781ff, 0x453882ff,
0x443983ff, 0x443a83ff, 0x443b84ff, 0x433d84ff, 0x433e85ff, 0x423f85ff, 0x424086ff,
0x424186ff, 0x414287ff, 0x414487ff, 0x404588ff, 0x404688ff, 0x3f4788ff, 0x3f4889ff,
0x3e4989ff, 0x3e4a89ff, 0x3e4c8aff, 0x3d4d8aff, 0x3d4e8aff, 0x3c4f8aff, 0x3c508bff,
0x3b518bff, 0x3b528bff, 0x3a538bff, 0x3a548cff, 0x39558cff, 0x39568cff, 0x38588cff,
0x38598cff, 0x375a8cff, 0x375b8dff, 0x365c8dff, 0x365d8dff, 0x355e8dff, 0x355f8dff,
0x34608dff, 0x34618dff, 0x33628dff, 0x33638dff, 0x32648eff, 0x32658eff, 0x31668eff,
0x31678eff, 0x31688eff, 0x30698eff, 0x306a8eff, 0x2f6b8eff, 0x2f6c8eff, 0x2e6d8eff,
0x2e6e8eff, 0x2e6f8eff, 0x2d708eff, 0x2d718eff, 0x2c718eff, 0x2c728eff, 0x2c738eff,
0x2b748eff, 0x2b758eff, 0x2a768eff, 0x2a778eff, 0x2a788eff, 0x29798eff, 0x297a8eff,
0x297b8eff, 0x287c8eff, 0x287d8eff, 0x277e8eff, 0x277f8eff, 0x27808eff, 0x26818eff,
0x26828eff, 0x26828eff, 0x25838eff, 0x25848eff, 0x25858eff, 0x24868eff, 0x24878eff,
0x23888eff, 0x23898eff, 0x238a8dff, 0x228b8dff, 0x228c8dff, 0x228d8dff, 0x218e8dff,
0x218f8dff, 0x21908dff, 0x21918cff, 0x20928cff, 0x20928cff, 0x20938cff, 0x1f948cff,
0x1f958bff, 0x1f968bff, 0x1f978bff, 0x1f988bff, 0x1f998aff, 0x1f9a8aff, 0x1e9b8aff,
0x1e9c89ff, 0x1e9d89ff, 0x1f9e89ff, 0x1f9f88ff, 0x1fa088ff, 0x1fa188ff, 0x1fa187ff,
0x1fa287ff, 0x20a386ff, 0x20a486ff, 0x21a585ff, 0x21a685ff, 0x22a785ff, 0x22a884ff,
0x23a983ff, 0x24aa83ff, 0x25ab82ff, 0x25ac82ff, 0x26ad81ff, 0x27ad81ff, 0x28ae80ff,
0x29af7fff, 0x2ab07fff, 0x2cb17eff, 0x2db27dff, 0x2eb37cff, 0x2fb47cff, 0x31b57bff,
0x32b67aff, 0x34b679ff, 0x35b779ff, 0x37b878ff, 0x38b977ff, 0x3aba76ff, 0x3bbb75ff,
0x3dbc74ff, 0x3fbc73ff, 0x40bd72ff, 0x42be71ff, 0x44bf70ff, 0x46c06fff, 0x48c16eff,
0x4ac16dff, 0x4cc26cff, 0x4ec36bff, 0x50c46aff, 0x52c569ff, 0x54c568ff, 0x56c667ff,
0x58c765ff, 0x5ac864ff, 0x5cc863ff, 0x5ec962ff, 0x60ca60ff, 0x63cb5fff, 0x65cb5eff,
0x67cc5cff, 0x69cd5bff, 0x6ccd5aff, 0x6ece58ff, 0x70cf57ff, 0x73d056ff, 0x75d054ff,
0x77d153ff, 0x7ad151ff, 0x7cd250ff, 0x7fd34eff, 0x81d34dff, 0x84d44bff, 0x86d549ff,
0x89d548ff, 0x8bd646ff, 0x8ed645ff, 0x90d743ff, 0x93d741ff, 0x95d840ff, 0x98d83eff,
0x9bd93cff, 0x9dd93bff, 0xa0da39ff, 0xa2da37ff, 0xa5db36ff, 0xa8db34ff, 0xaadc32ff,
0xaddc30ff, 0xb0dd2fff, 0xb2dd2dff, 0xb5de2bff, 0xb8de29ff, 0xbade28ff, 0xbddf26ff,
0xc0df25ff, 0xc2df23ff, 0xc5e021ff, 0xc8e020ff, 0xcae11fff, 0xcde11dff, 0xd0e11cff,
0xd2e21bff, 0xd5e21aff, 0xd8e219ff, 0xdae319ff, 0xdde318ff, 0xdfe318ff, 0xe2e418ff,
0xe5e419ff, 0xe7e419ff, 0xeae51aff, 0xece51bff, 0xefe51cff, 0xf1e51dff, 0xf4e61eff,
0xf6e620ff, 0xf8e621ff, 0xfbe723ff, 0xfde725ff,
],
Self::Magma => [
0x000004ff, 0x010005ff, 0x010106ff, 0x010108ff, 0x020109ff, 0x02020bff, 0x02020dff,
0x03030fff, 0x030312ff, 0x040414ff, 0x050416ff, 0x060518ff, 0x06051aff, 0x07061cff,
0x08071eff, 0x090720ff, 0x0a0822ff, 0x0b0924ff, 0x0c0926ff, 0x0d0a29ff, 0x0e0b2bff,
0x100b2dff, 0x110c2fff, 0x120d31ff, 0x130d34ff, 0x140e36ff, 0x150e38ff, 0x160f3bff,
0x180f3dff, 0x19103fff, 0x1a1042ff, 0x1c1044ff, 0x1d1147ff, 0x1e1149ff, 0x20114bff,
0x21114eff, 0x221150ff, 0x241253ff, 0x251255ff, 0x271258ff, 0x29115aff, 0x2a115cff,
0x2c115fff, 0x2d1161ff, 0x2f1163ff, 0x311165ff, 0x331067ff, 0x341069ff, 0x36106bff,
0x38106cff, 0x390f6eff, 0x3b0f70ff, 0x3d0f71ff, 0x3f0f72ff, 0x400f74ff, 0x420f75ff,
0x440f76ff, 0x451077ff, 0x471078ff, 0x491078ff, 0x4a1079ff, 0x4c117aff, 0x4e117bff,
0x4f127bff, 0x51127cff, 0x52137cff, 0x54137dff, 0x56147dff, 0x57157eff, 0x59157eff,
0x5a167eff, 0x5c167fff, 0x5d177fff, 0x5f187fff, 0x601880ff, 0x621980ff, 0x641a80ff,
0x651a80ff, 0x671b80ff, 0x681c81ff, 0x6a1c81ff, 0x6b1d81ff, 0x6d1d81ff, 0x6e1e81ff,
0x701f81ff, 0x721f81ff, 0x732081ff, 0x752181ff, 0x762181ff, 0x782281ff, 0x792282ff,
0x7b2382ff, 0x7c2382ff, 0x7e2482ff, 0x802582ff, 0x812581ff, 0x832681ff, 0x842681ff,
0x862781ff, 0x882781ff, 0x892881ff, 0x8b2981ff, 0x8c2981ff, 0x8e2a81ff, 0x902a81ff,
0x912b81ff, 0x932b80ff, 0x942c80ff, 0x962c80ff, 0x982d80ff, 0x992d80ff, 0x9b2e7fff,
0x9c2e7fff, 0x9e2f7fff, 0xa02f7fff, 0xa1307eff, 0xa3307eff, 0xa5317eff, 0xa6317dff,
0xa8327dff, 0xaa337dff, 0xab337cff, 0xad347cff, 0xae347bff, 0xb0357bff, 0xb2357bff,
0xb3367aff, 0xb5367aff, 0xb73779ff, 0xb83779ff, 0xba3878ff, 0xbc3978ff, 0xbd3977ff,
0xbf3a77ff, 0xc03a76ff, 0xc23b75ff, 0xc43c75ff, 0xc53c74ff, 0xc73d73ff, 0xc83e73ff,
0xca3e72ff, 0xcc3f71ff, 0xcd4071ff, 0xcf4070ff, 0xd0416fff, 0xd2426fff, 0xd3436eff,
0xd5446dff, 0xd6456cff, 0xd8456cff, 0xd9466bff, 0xdb476aff, 0xdc4869ff, 0xde4968ff,
0xdf4a68ff, 0xe04c67ff, 0xe24d66ff, 0xe34e65ff, 0xe44f64ff, 0xe55064ff, 0xe75263ff,
0xe85362ff, 0xe95462ff, 0xea5661ff, 0xeb5760ff, 0xec5860ff, 0xed5a5fff, 0xee5b5eff,
0xef5d5eff, 0xf05f5eff, 0xf1605dff, 0xf2625dff, 0xf2645cff, 0xf3655cff, 0xf4675cff,
0xf4695cff, 0xf56b5cff, 0xf66c5cff, 0xf66e5cff, 0xf7705cff, 0xf7725cff, 0xf8745cff,
0xf8765cff, 0xf9785dff, 0xf9795dff, 0xf97b5dff, 0xfa7d5eff, 0xfa7f5eff, 0xfa815fff,
0xfb835fff, 0xfb8560ff, 0xfb8761ff, 0xfc8961ff, 0xfc8a62ff, 0xfc8c63ff, 0xfc8e64ff,
0xfc9065ff, 0xfd9266ff, 0xfd9467ff, 0xfd9668ff, 0xfd9869ff, 0xfd9a6aff, 0xfd9b6bff,
0xfe9d6cff, 0xfe9f6dff, 0xfea16eff, 0xfea36fff, 0xfea571ff, 0xfea772ff, 0xfea973ff,
0xfeaa74ff, 0xfeac76ff, 0xfeae77ff, 0xfeb078ff, 0xfeb27aff, 0xfeb47bff, 0xfeb67cff,
0xfeb77eff, 0xfeb97fff, 0xfebb81ff, 0xfebd82ff, 0xfebf84ff, 0xfec185ff, 0xfec287ff,
0xfec488ff, 0xfec68aff, 0xfec88cff, 0xfeca8dff, 0xfecc8fff, 0xfecd90ff, 0xfecf92ff,
0xfed194ff, 0xfed395ff, 0xfed597ff, 0xfed799ff, 0xfed89aff, 0xfdda9cff, 0xfddc9eff,
0xfddea0ff, 0xfde0a1ff, 0xfde2a3ff, 0xfde3a5ff, 0xfde5a7ff, 0xfde7a9ff, 0xfde9aaff,
0xfdebacff, 0xfcecaeff, 0xfceeb0ff, 0xfcf0b2ff, 0xfcf2b4ff, 0xfcf4b6ff, 0xfcf6b8ff,
0xfcf7b9ff, 0xfcf9bbff, 0xfcfbbdff, 0xfcfdbfff,
],
Self::BentCoolWarm => [
0x3b4cc0ff, 0x3c4ec1ff, 0x3d4fc2ff, 0x3e50c2ff, 0x3f52c3ff, 0x4053c4ff, 0x4155c4ff,
0x4256c5ff, 0x4357c5ff, 0x4459c6ff, 0x455ac7ff, 0x465bc7ff, 0x475dc8ff, 0x485ec8ff,
0x495fc9ff, 0x4a61caff, 0x4b62caff, 0x4c64cbff, 0x4d65cbff, 0x4e66ccff, 0x4f68ccff,
0x5169cdff, 0x526aceff, 0x536cceff, 0x546dcfff, 0x556ecfff, 0x5670d0ff, 0x5771d0ff,
0x5972d1ff, 0x5a74d1ff, 0x5b75d2ff, 0x5c76d2ff, 0x5d78d3ff, 0x5f79d3ff, 0x607ad4ff,
0x617cd4ff, 0x627dd5ff, 0x647ed5ff, 0x6580d6ff, 0x6681d6ff, 0x6782d6ff, 0x6983d7ff,
0x6a85d7ff, 0x6b86d8ff, 0x6d87d8ff, 0x6e89d9ff, 0x6f8ad9ff, 0x718bdaff, 0x728ddaff,
0x738edaff, 0x758fdbff, 0x7691dbff, 0x7792dcff, 0x7993dcff, 0x7a95dcff, 0x7b96ddff,
0x7d97ddff, 0x7e98deff, 0x809adeff, 0x819bdeff, 0x829cdfff, 0x849edfff, 0x859fdfff,
0x87a0e0ff, 0x88a2e0ff, 0x8aa3e1ff, 0x8ba4e1ff, 0x8ca5e1ff, 0x8ea7e2ff, 0x8fa8e2ff,
0x91a9e2ff, 0x92abe3ff, 0x94ace3ff, 0x95ade3ff, 0x97afe4ff, 0x98b0e4ff, 0x9ab1e4ff,
0x9bb2e5ff, 0x9db4e5ff, 0x9fb5e5ff, 0xa0b6e6ff, 0xa2b8e6ff, 0xa3b9e6ff, 0xa5bae6ff,
0xa6bbe7ff, 0xa8bde7ff, 0xaabee7ff, 0xabbfe8ff, 0xadc1e8ff, 0xaec2e8ff, 0xb0c3e8ff,
0xb2c4e9ff, 0xb3c6e9ff, 0xb5c7e9ff, 0xb7c8eaff, 0xb8caeaff, 0xbacbeaff, 0xbccceaff,
0xbdcdebff, 0xbfcfebff, 0xc1d0ebff, 0xc2d1ecff, 0xc4d2ecff, 0xc6d4ecff, 0xc8d5ecff,
0xc9d6edff, 0xcbd7edff, 0xcdd9edff, 0xcfdaedff, 0xd0dbeeff, 0xd2dceeff, 0xd4deeeff,
0xd6dfeeff, 0xd7e0efff, 0xd9e1efff, 0xdbe3efff, 0xdde4efff, 0xdfe5f0ff, 0xe1e6f0ff,
0xe2e8f0ff, 0xe4e9f0ff, 0xe6eaf1ff, 0xe8ebf1ff, 0xeaedf1ff, 0xeceef1ff, 0xeeeff2ff,
0xeff0f2ff, 0xf1f2f2ff, 0xf2f1f1ff, 0xf2f0efff, 0xf1eeedff, 0xf1edebff, 0xf1ebe8ff,
0xf1eae6ff, 0xf0e8e4ff, 0xf0e7e2ff, 0xf0e5e0ff, 0xefe4deff, 0xefe2dbff, 0xefe1d9ff,
0xeedfd7ff, 0xeeded5ff, 0xeedcd3ff, 0xeddbd1ff, 0xedd9cfff, 0xedd8cdff, 0xecd6cbff,
0xecd5c9ff, 0xecd3c7ff, 0xebd2c4ff, 0xebd0c2ff, 0xebcfc0ff, 0xeacdbeff, 0xeaccbcff,
0xe9cabaff, 0xe9c9b8ff, 0xe9c7b6ff, 0xe8c5b4ff, 0xe8c4b3ff, 0xe8c2b1ff, 0xe7c1afff,
0xe7bfadff, 0xe6beabff, 0xe6bca9ff, 0xe6bba7ff, 0xe5b9a5ff, 0xe5b8a3ff, 0xe4b6a1ff,
0xe4b59fff, 0xe4b39eff, 0xe3b19cff, 0xe3b09aff, 0xe2ae98ff, 0xe2ad96ff, 0xe2ab94ff,
0xe1aa93ff, 0xe1a891ff, 0xe0a78fff, 0xe0a58dff, 0xdfa38cff, 0xdfa28aff, 0xdfa088ff,
0xde9f86ff, 0xde9d85ff, 0xdd9c83ff, 0xdd9a81ff, 0xdc9880ff, 0xdc977eff, 0xdb957cff,
0xdb947bff, 0xda9279ff, 0xda9077ff, 0xd98f76ff, 0xd98d74ff, 0xd98c72ff, 0xd88a71ff,
0xd8886fff, 0xd7876eff, 0xd7856cff, 0xd6846bff, 0xd68269ff, 0xd58068ff, 0xd47f66ff,
0xd47d65ff, 0xd37b63ff, 0xd37a62ff, 0xd27860ff, 0xd2775fff, 0xd1755dff, 0xd1735cff,
0xd0725aff, 0xd07059ff, 0xcf6e58ff, 0xcf6c56ff, 0xce6b55ff, 0xcd6953ff, 0xcd6752ff,
0xcc6651ff, 0xcc644fff, 0xcb624eff, 0xcb604dff, 0xca5f4bff, 0xc95d4aff, 0xc95b49ff,
0xc85948ff, 0xc85746ff, 0xc75645ff, 0xc65444ff, 0xc65243ff, 0xc55041ff, 0xc54e40ff,
0xc44c3fff, 0xc34a3eff, 0xc3483dff, 0xc2463cff, 0xc1443aff, 0xc14239ff, 0xc04038ff,
0xc03e37ff, 0xbf3c36ff, 0xbe3a35ff, 0xbe3734ff, 0xbd3533ff, 0xbc3232ff, 0xbc3031ff,
0xbb2d30ff, 0xba2b2fff, 0xba282eff, 0xb9252dff, 0xb8222cff, 0xb81e2bff, 0xb71b2aff,
0xb61629ff, 0xb51228ff, 0xb50c27ff, 0xb40426ff,
],
Self::BlackBody => [
0x000000ff, 0x030101ff, 0x070201ff, 0x0a0302ff, 0x0d0402ff, 0x100503ff, 0x120603ff,
0x140704ff, 0x160804ff, 0x180905ff, 0x1a0a05ff, 0x1b0b06ff, 0x1d0b06ff, 0x1e0c07ff,
0x200d08ff, 0x210e08ff, 0x220f09ff, 0x240f09ff, 0x25100aff, 0x26100aff, 0x28110bff,
0x29110bff, 0x2b120cff, 0x2c120cff, 0x2e120dff, 0x2f130dff, 0x31130eff, 0x32130eff,
0x34140fff, 0x36140fff, 0x37140fff, 0x391510ff, 0x3a1510ff, 0x3c1510ff, 0x3e1611ff,
0x3f1611ff, 0x411611ff, 0x421712ff, 0x441712ff, 0x461712ff, 0x471813ff, 0x491813ff,
0x4b1813ff, 0x4c1914ff, 0x4e1914ff, 0x501914ff, 0x511914ff, 0x531a15ff, 0x551a15ff,
0x561a15ff, 0x581a15ff, 0x5a1b16ff, 0x5b1b16ff, 0x5d1b16ff, 0x5f1b16ff, 0x611c17ff,
0x621c17ff, 0x641c17ff, 0x661c17ff, 0x681d18ff, 0x691d18ff, 0x6b1d18ff, 0x6d1d18ff,
0x6f1d19ff, 0x701e19ff, 0x721e19ff, 0x741e19ff, 0x761e1aff, 0x771e1aff, 0x791f1aff,
0x7b1f1aff, 0x7d1f1bff, 0x7f1f1bff, 0x801f1bff, 0x821f1bff, 0x84201cff, 0x86201cff,
0x88201cff, 0x89201cff, 0x8b201dff, 0x8d201dff, 0x8f201dff, 0x91211dff, 0x93211eff,
0x94211eff, 0x96211eff, 0x98211fff, 0x9a211fff, 0x9c211fff, 0x9e211fff, 0xa02120ff,
0xa12220ff, 0xa32220ff, 0xa52220ff, 0xa72221ff, 0xa92221ff, 0xab2221ff, 0xad2221ff,
0xaf2222ff, 0xb12222ff, 0xb22222ff, 0xb32422ff, 0xb42622ff, 0xb52821ff, 0xb62a21ff,
0xb72c21ff, 0xb82d21ff, 0xb92f20ff, 0xba3120ff, 0xbb3220ff, 0xbc341fff, 0xbd351fff,
0xbe371fff, 0xbf381fff, 0xc03a1eff, 0xc13b1eff, 0xc23d1eff, 0xc33e1dff, 0xc4401dff,
0xc5411cff, 0xc6421cff, 0xc7441cff, 0xc8451bff, 0xc9471bff, 0xca481aff, 0xcb491aff,
0xcc4b19ff, 0xcd4c19ff, 0xce4d18ff, 0xcf4f18ff, 0xd05017ff, 0xd15217ff, 0xd25316ff,
0xd35415ff, 0xd45515ff, 0xd55714ff, 0xd65813ff, 0xd75913ff, 0xd85b12ff, 0xd95c11ff,
0xda5d10ff, 0xdb5f0fff, 0xdc600eff, 0xdd610dff, 0xde620cff, 0xdf640bff, 0xe06509ff,
0xe16608ff, 0xe26807ff, 0xe36905ff, 0xe36b05ff, 0xe36d06ff, 0xe46e07ff, 0xe47007ff,
0xe47208ff, 0xe47408ff, 0xe57609ff, 0xe5770aff, 0xe5790aff, 0xe57b0bff, 0xe57c0cff,
0xe67e0cff, 0xe6800dff, 0xe6820eff, 0xe6830eff, 0xe6850fff, 0xe6870fff, 0xe78810ff,
0xe78a11ff, 0xe78c11ff, 0xe78d12ff, 0xe78f13ff, 0xe79113ff, 0xe79214ff, 0xe89415ff,
0xe89615ff, 0xe89716ff, 0xe89916ff, 0xe89a17ff, 0xe89c18ff, 0xe89e18ff, 0xe89f19ff,
0xe8a11aff, 0xe8a21aff, 0xe9a41bff, 0xe9a61bff, 0xe9a71cff, 0xe9a91dff, 0xe9aa1dff,
0xe9ac1eff, 0xe9ae1eff, 0xe9af1fff, 0xe9b120ff, 0xe9b220ff, 0xe9b421ff, 0xe9b522ff,
0xe9b722ff, 0xe9b923ff, 0xe9ba23ff, 0xe9bc24ff, 0xe9bd25ff, 0xe9bf25ff, 0xe9c026ff,
0xe9c226ff, 0xe9c327ff, 0xe9c528ff, 0xe9c728ff, 0xe9c829ff, 0xe8ca2aff, 0xe8cb2aff,
0xe8cd2bff, 0xe8ce2bff, 0xe8d02cff, 0xe8d12dff, 0xe8d32dff, 0xe8d52eff, 0xe8d62fff,
0xe8d82fff, 0xe7d930ff, 0xe7db30ff, 0xe7dc31ff, 0xe7de32ff, 0xe7df32ff, 0xe7e133ff,
0xe6e234ff, 0xe6e434ff, 0xe6e535ff, 0xe7e73cff, 0xe9e745ff, 0xeae84eff, 0xece957ff,
0xedea5eff, 0xeeeb66ff, 0xf0ec6dff, 0xf1ec75ff, 0xf2ed7cff, 0xf3ee83ff, 0xf5ef89ff,
0xf6f090ff, 0xf7f197ff, 0xf8f19eff, 0xf9f2a4ff, 0xf9f3abff, 0xfaf4b1ff, 0xfbf5b8ff,
0xfcf6beff, 0xfcf7c5ff, 0xfdf8cbff, 0xfdf9d2ff, 0xfef9d8ff, 0xfefadfff, 0xfefbe5ff,
0xfffcecff, 0xfffdf2ff, 0xfffef9ff, 0xffffffff,
],
Self::ExtendedKindLmann => [
0x000000ff, 0x050004ff, 0x090009ff, 0x0d010dff, 0x100111ff, 0x130115ff, 0x160118ff,
0x18011bff, 0x1a011eff, 0x1b0222ff, 0x1c0226ff, 0x1d022aff, 0x1d022eff, 0x1e0232ff,
0x1e0335ff, 0x1e0339ff, 0x1e033dff, 0x1d0341ff, 0x1d0344ff, 0x1c0348ff, 0x1b044bff,
0x1b044fff, 0x1a0452ff, 0x190455ff, 0x180458ff, 0x17045cff, 0x16055fff, 0x150562ff,
0x140565ff, 0x130567ff, 0x12056aff, 0x12056dff, 0x11056fff, 0x0e0573ff, 0x080677ff,
0x060878ff, 0x060b78ff, 0x060f77ff, 0x061276ff, 0x061674ff, 0x051972ff, 0x051c70ff,
0x051f6dff, 0x05216bff, 0x052468ff, 0x052665ff, 0x052863ff, 0x052a60ff, 0x052c5eff,
0x042e5bff, 0x043059ff, 0x043157ff, 0x043355ff, 0x043453ff, 0x043651ff, 0x04374fff,
0x04394dff, 0x043a4cff, 0x043b4aff, 0x033d49ff, 0x033e47ff, 0x033f46ff, 0x034145ff,
0x034243ff, 0x034342ff, 0x034441ff, 0x034540ff, 0x03473fff, 0x03483dff, 0x04493cff,
0x044a3aff, 0x044b38ff, 0x044d37ff, 0x044e35ff, 0x044f33ff, 0x045031ff, 0x04512fff,
0x04522dff, 0x04542bff, 0x045529ff, 0x045627ff, 0x045724ff, 0x045822ff, 0x04591fff,
0x045b1dff, 0x045c1aff, 0x055d18ff, 0x055e15ff, 0x055f12ff, 0x05600fff, 0x05610dff,
0x05620aff, 0x056408ff, 0x056506ff, 0x066605ff, 0x086705ff, 0x0a6805ff, 0x0b6905ff,
0x0d6a05ff, 0x0f6b05ff, 0x116c05ff, 0x146d05ff, 0x166e05ff, 0x1a6f05ff, 0x1d7005ff,
0x207005ff, 0x247105ff, 0x287205ff, 0x2b7306ff, 0x2f7406ff, 0x337406ff, 0x377506ff,
0x3b7606ff, 0x3f7606ff, 0x437706ff, 0x477706ff, 0x4c7806ff, 0x507806ff, 0x547906ff,
0x587906ff, 0x5c7a06ff, 0x617a06ff, 0x657a06ff, 0x697b06ff, 0x6d7b06ff, 0x717b06ff,
0x767b06ff, 0x7a7b06ff, 0x7e7b06ff, 0x827b06ff, 0x877b07ff, 0x8b7b07ff, 0x907b07ff,
0x957a07ff, 0x9a7a07ff, 0xa07908ff, 0xa57808ff, 0xab7708ff, 0xb17608ff, 0xb77509ff,
0xbd7309ff, 0xc47109ff, 0xca6f0aff, 0xd16c0aff, 0xd8690aff, 0xde660bff, 0xe5620bff,
0xec5e0bff, 0xf35a0cff, 0xf45b1bff, 0xf55c25ff, 0xf55e2eff, 0xf56034ff, 0xf6623aff,
0xf6633fff, 0xf66543ff, 0xf66747ff, 0xf6694aff, 0xf66b4dff, 0xf76d4fff, 0xf76f53ff,
0xf77057ff, 0xf7725bff, 0xf77360ff, 0xf87565ff, 0xf8766aff, 0xf87870ff, 0xf87976ff,
0xf97a7bff, 0xf97b81ff, 0xf97d87ff, 0xf97e8dff, 0xf97f93ff, 0xf98199ff, 0xf9829eff,
0xf983a4ff, 0xf984a9ff, 0xf985afff, 0xf986b4ff, 0xf987baff, 0xf989bfff, 0xf98ac4ff,
0xf98bc9ff, 0xfa8cceff, 0xfa8dd3ff, 0xfa8ed8ff, 0xfa8fddff, 0xfa90e1ff, 0xfa91e6ff,
0xfa92ebff, 0xfa93efff, 0xfa94f3ff, 0xfa95f8ff, 0xf898faff, 0xf59bfaff, 0xf29ffaff,
0xefa2fbff, 0xeca5fbff, 0xeaa8fbff, 0xe8abfbff, 0xe6adfbff, 0xe5b0fbff, 0xe3b2fbff,
0xe2b4fbff, 0xe1b6fbff, 0xe0b8fcff, 0xe0bafcff, 0xdfbcfcff, 0xdfbefcff, 0xdebffcff,
0xdec1fcff, 0xdec3fcff, 0xdec4fcff, 0xdfc6fcff, 0xdfc7fcff, 0xdfc9fcff, 0xe0cafcff,
0xe0ccfdff, 0xe1cdfdff, 0xe2cffdff, 0xe2d0fdff, 0xe3d1fdff, 0xe4d3fdff, 0xe5d4fdff,
0xe5d5fdff, 0xe6d7fdff, 0xe7d8fdff, 0xe7dafdff, 0xe7dbfdff, 0xe8ddfdff, 0xe8defdff,
0xe8e0fdff, 0xe8e1feff, 0xe9e3feff, 0xe9e4feff, 0xe9e6feff, 0xe9e7feff, 0xe9e9feff,
0xe9eafeff, 0xeaecfeff, 0xeaedfeff, 0xeaeffeff, 0xebf0feff, 0xebf2feff, 0xecf3feff,
0xedf5feff, 0xedf6feff, 0xeef7feff, 0xeff9feff, 0xf0fafeff, 0xf2fbfeff, 0xf3fcfeff,
0xf5fdffff, 0xf8feffff, 0xfbffffff, 0xffffffff,
],
Self::KindLmann => [
0x000000ff, 0x050004ff, 0x090008ff, 0x0d010dff, 0x110110ff, 0x140114ff, 0x160117ff,
0x19011aff, 0x1b011dff, 0x1d0220ff, 0x1e0223ff, 0x1f0226ff, 0x20022aff, 0x21022dff,
0x220230ff, 0x230233ff, 0x240336ff, 0x250339ff, 0x25033cff, 0x26033fff, 0x260342ff,
0x260344ff, 0x270347ff, 0x27044aff, 0x27044dff, 0x270450ff, 0x270453ff, 0x270456ff,
0x270459ff, 0x27045dff, 0x270560ff, 0x270563ff, 0x260566ff, 0x26056aff, 0x25056dff,
0x250570ff, 0x240674ff, 0x230677ff, 0x22067bff, 0x21067eff, 0x200681ff, 0x200684ff,
0x1f0688ff, 0x1e078bff, 0x1d078eff, 0x1c0791ff, 0x1b0794ff, 0x1a0797ff, 0x19079aff,
0x19079dff, 0x1808a0ff, 0x1808a3ff, 0x1408a6ff, 0x0f08aaff, 0x0809aeff, 0x080cafff,
0x080fafff, 0x0813afff, 0x0816afff, 0x0819afff, 0x081caeff, 0x0820adff, 0x0823acff,
0x0826aaff, 0x0829a8ff, 0x082ba6ff, 0x082ea5ff, 0x0831a3ff, 0x0833a0ff, 0x08359eff,
0x08389cff, 0x073a9aff, 0x073c98ff, 0x073e95ff, 0x074093ff, 0x074291ff, 0x07448fff,
0x07468dff, 0x07478bff, 0x074989ff, 0x074b87ff, 0x064c85ff, 0x064e84ff, 0x065082ff,
0x065180ff, 0x06537fff, 0x06547dff, 0x06567bff, 0x06577aff, 0x065878ff, 0x065a77ff,
0x065b76ff, 0x065d74ff, 0x065e73ff, 0x055f72ff, 0x066071ff, 0x056270ff, 0x05636eff,
0x05646dff, 0x05666cff, 0x05676bff, 0x05686aff, 0x056969ff, 0x056b68ff, 0x056c67ff,
0x056d66ff, 0x056e65ff, 0x057064ff, 0x057163ff, 0x057262ff, 0x067360ff, 0x06755fff,
0x06765eff, 0x06775cff, 0x06785bff, 0x067a59ff, 0x067b58ff, 0x067c56ff, 0x067d54ff,
0x067f53ff, 0x068051ff, 0x06814fff, 0x06824dff, 0x06844bff, 0x06854aff, 0x078648ff,
0x068746ff, 0x078943ff, 0x078a41ff, 0x078b3fff, 0x078c3dff, 0x078e3bff, 0x078f38ff,
0x079036ff, 0x079134ff, 0x079331ff, 0x07942fff, 0x07952cff, 0x07962aff, 0x079727ff,
0x079925ff, 0x079a22ff, 0x089b1fff, 0x089c1dff, 0x089d1aff, 0x089f17ff, 0x08a014ff,
0x08a112ff, 0x08a20fff, 0x08a30cff, 0x08a50aff, 0x08a608ff, 0x0ca708ff, 0x0fa808ff,
0x11a908ff, 0x12aa08ff, 0x14ab08ff, 0x16ad08ff, 0x18ae08ff, 0x1aaf08ff, 0x1db008ff,
0x20b109ff, 0x23b209ff, 0x26b309ff, 0x29b409ff, 0x2db509ff, 0x30b609ff, 0x34b709ff,
0x38b809ff, 0x3bb909ff, 0x3fba09ff, 0x43bb09ff, 0x47bc09ff, 0x4bbd09ff, 0x4fbe09ff,
0x53be09ff, 0x57bf09ff, 0x5bc009ff, 0x5fc109ff, 0x63c109ff, 0x67c209ff, 0x6bc309ff,
0x6fc409ff, 0x74c409ff, 0x78c509ff, 0x7cc60aff, 0x80c60aff, 0x85c70aff, 0x89c70aff,
0x8dc80aff, 0x91c80aff, 0x96c90aff, 0x9ac90aff, 0x9eca0aff, 0xa3ca0aff, 0xa7ca0aff,
0xabcb0aff, 0xafcb0aff, 0xb4cb0aff, 0xb8cc0aff, 0xbccc0aff, 0xc1cc0aff, 0xc5cd0aff,
0xc9cd0aff, 0xcdcd0aff, 0xd1cd0aff, 0xd6cd0aff, 0xdacd0bff, 0xdfcd0bff, 0xe4cd0bff,
0xe9cd0bff, 0xedcd0bff, 0xf3cd0cff, 0xf6cc39ff, 0xf7cd56ff, 0xf8cd69ff, 0xf9ce77ff,
0xf9cf83ff, 0xfacf8dff, 0xfad095ff, 0xfad19dff, 0xfbd2a3ff, 0xfbd3a9ff, 0xfbd4aeff,
0xfbd6b3ff, 0xfcd7b8ff, 0xfcd8bcff, 0xfcd9c0ff, 0xfcdac3ff, 0xfcdcc7ff, 0xfcddcaff,
0xfddecdff, 0xfde0d0ff, 0xfde1d3ff, 0xfde2d5ff, 0xfde3d8ff, 0xfde5daff, 0xfde6ddff,
0xfde8dfff, 0xfee9e1ff, 0xfeeae3ff, 0xfeece5ff, 0xfeede7ff, 0xfeeee9ff, 0xfef0ebff,
0xfef1edff, 0xfef2efff, 0xfef4f1ff, 0xfef5f3ff, 0xfef7f5ff, 0xfff8f6ff, 0xfff9f8ff,
0xfffbfaff, 0xfffcfcff, 0xfffefdff, 0xffffffff,
],
Self::SmoothCoolWarm => [
0x3b4cc0ff, 0x3c4ec2ff, 0x3d50c3ff, 0x3e51c5ff, 0x3f53c7ff, 0x4055c8ff, 0x4257caff,
0x4358cbff, 0x445accff, 0x455cceff, 0x465ecfff, 0x485fd1ff, 0x4961d2ff, 0x4a63d4ff,
0x4b64d5ff, 0x4c66d6ff, 0x4e68d8ff, 0x4f6ad9ff, 0x506bdaff, 0x516ddbff, 0x536fddff,
0x5470deff, 0x5572dfff, 0x5674e0ff, 0x5875e2ff, 0x5977e3ff, 0x5a78e4ff, 0x5b7ae5ff,
0x5d7ce6ff, 0x5e7de7ff, 0x5f7fe8ff, 0x6181e9ff, 0x6282eaff, 0x6384ebff, 0x6585ecff,
0x6687edff, 0x6788eeff, 0x698aefff, 0x6a8cf0ff, 0x6b8df0ff, 0x6d8ff1ff, 0x6e90f2ff,
0x6f92f3ff, 0x7193f4ff, 0x7295f4ff, 0x7396f5ff, 0x7598f6ff, 0x7699f6ff, 0x779af7ff,
0x799cf8ff, 0x7a9df8ff, 0x7b9ff9ff, 0x7da0f9ff, 0x7ea2faff, 0x80a3faff, 0x81a4fbff,
0x82a6fbff, 0x84a7fcff, 0x85a8fcff, 0x86aafcff, 0x88abfdff, 0x89acfdff, 0x8baefdff,
0x8caffeff, 0x8db0feff, 0x8fb1feff, 0x90b2feff, 0x92b4feff, 0x93b5ffff, 0x94b6ffff,
0x96b7ffff, 0x97b8ffff, 0x99baffff, 0x9abbffff, 0x9bbcffff, 0x9dbdffff, 0x9ebeffff,
0x9fbfffff, 0xa1c0ffff, 0xa2c1ffff, 0xa3c2feff, 0xa5c3feff, 0xa6c4feff, 0xa8c5feff,
0xa9c6feff, 0xaac7fdff, 0xacc8fdff, 0xadc9fdff, 0xaec9fcff, 0xb0cafcff, 0xb1cbfcff,
0xb2ccfbff, 0xb4cdfbff, 0xb5cefaff, 0xb6cefaff, 0xb7cff9ff, 0xb9d0f9ff, 0xbad1f8ff,
0xbbd1f8ff, 0xbdd2f7ff, 0xbed3f6ff, 0xbfd3f6ff, 0xc0d4f5ff, 0xc1d4f4ff, 0xc3d5f4ff,
0xc4d6f3ff, 0xc5d6f2ff, 0xc6d7f1ff, 0xc8d7f1ff, 0xc9d8f0ff, 0xcad8efff, 0xcbd8eeff,
0xccd9edff, 0xcdd9ecff, 0xcedaebff, 0xd0daeaff, 0xd1dae9ff, 0xd2dbe8ff, 0xd3dbe7ff,
0xd4dbe6ff, 0xd5dbe5ff, 0xd6dce4ff, 0xd7dce3ff, 0xd8dce2ff, 0xd9dce1ff, 0xdadce0ff,
0xdbdddeff, 0xdcddddff, 0xdddcdcff, 0xdedcdbff, 0xdfdcd9ff, 0xe1dbd8ff, 0xe2dad6ff,
0xe3dad5ff, 0xe4d9d3ff, 0xe5d9d2ff, 0xe5d8d1ff, 0xe6d8cfff, 0xe7d7ceff, 0xe8d6ccff,
0xe9d6cbff, 0xead5c9ff, 0xebd4c8ff, 0xebd3c6ff, 0xecd3c5ff, 0xedd2c3ff, 0xeed1c2ff,
0xeed0c0ff, 0xefcfbfff, 0xefcebdff, 0xf0cebbff, 0xf1cdbaff, 0xf1ccb8ff, 0xf2cbb7ff,
0xf2cab5ff, 0xf3c9b4ff, 0xf3c8b2ff, 0xf4c7b1ff, 0xf4c6afff, 0xf4c5adff, 0xf5c4acff,
0xf5c3aaff, 0xf5c1a9ff, 0xf6c0a7ff, 0xf6bfa6ff, 0xf6bea4ff, 0xf6bda2ff, 0xf7bca1ff,
0xf7ba9fff, 0xf7b99eff, 0xf7b89cff, 0xf7b79bff, 0xf7b599ff, 0xf7b497ff, 0xf7b396ff,
0xf7b194ff, 0xf7b093ff, 0xf7af91ff, 0xf7ad90ff, 0xf7ac8eff, 0xf7ab8cff, 0xf7a98bff,
0xf7a889ff, 0xf7a688ff, 0xf6a586ff, 0xf6a385ff, 0xf6a283ff, 0xf6a081ff, 0xf59f80ff,
0xf59d7eff, 0xf59c7dff, 0xf49a7bff, 0xf4997aff, 0xf49778ff, 0xf39577ff, 0xf39475ff,
0xf29274ff, 0xf29072ff, 0xf18f71ff, 0xf18d6fff, 0xf08b6eff, 0xf08a6cff, 0xef886bff,
0xee8669ff, 0xee8568ff, 0xed8366ff, 0xed8165ff, 0xec7f63ff, 0xeb7d62ff, 0xea7c60ff,
0xea7a5fff, 0xe9785dff, 0xe8765cff, 0xe7745bff, 0xe67259ff, 0xe57058ff, 0xe56f56ff,
0xe46d55ff, 0xe36b54ff, 0xe26952ff, 0xe16751ff, 0xe0654fff, 0xdf634eff, 0xde614dff,
0xdd5f4bff, 0xdc5d4aff, 0xdb5b49ff, 0xda5947ff, 0xd85646ff, 0xd75445ff, 0xd65244ff,
0xd55042ff, 0xd44e41ff, 0xd34c40ff, 0xd1493eff, 0xd0473dff, 0xcf453cff, 0xce433bff,
0xcc4039ff, 0xcb3e38ff, 0xca3b37ff, 0xc83936ff, 0xc73635ff, 0xc63434ff, 0xc43132ff,
0xc32e31ff, 0xc12b30ff, 0xc0282fff, 0xbf252eff, 0xbd222dff, 0xbc1e2cff, 0xba1a2bff,
0xb91629ff, 0xb71128ff, 0xb60b27ff, 0xb40426ff,
],
};
Color::create_palette(&xs)
.try_into()
.expect("Vector length is not 256")
}
}

104
src/core/dataloader.rs → src/misc/dataloader.rs
View File

@ -1,18 +1,18 @@
use anyhow::{anyhow, Result};
use image::DynamicImage;
use indicatif::{ProgressBar, ProgressStyle};
use indicatif::ProgressBar;
use log::{info, warn};
use std::collections::VecDeque;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
#[cfg(feature = "ffmpeg")]
use video_rs::{
encode::{Encoder, Settings},
time::Time,
Decoder, Url,
};
use crate::{
build_progress_bar, string_now, Dir, Hub, Location, MediaType, CHECK_MARK, CROSS_MARK,
};
use crate::{build_progress_bar, Hub, Location, MediaType};
type TempReturnType = (Vec<DynamicImage>, Vec<PathBuf>);
@ -28,7 +28,8 @@ impl Iterator for DataLoaderIterator {
fn next(&mut self) -> Option<Self::Item> {
match &self.progress_bar {
None => self.receiver.recv().ok(),
Some(progress_bar) => match self.receiver.recv().ok() {
Some(progress_bar) => {
match self.receiver.recv().ok() {
Some(item) => {
progress_bar.inc(self.batch_size);
Some(item)
@ -36,13 +37,16 @@ impl Iterator for DataLoaderIterator {
None => {
progress_bar.set_prefix("Iterated");
progress_bar.set_style(
indicatif::ProgressStyle::with_template(crate::PROGRESS_BAR_STYLE_FINISH_2)
.unwrap(),
match indicatif::ProgressStyle::with_template(crate::PROGRESS_BAR_STYLE_FINISH_2) {
Ok(x) => x,
Err(err) => panic!("Failed to set style for progressbar in `DataLoaderIterator`: {}", err),
},
);
progress_bar.finish();
None
}
},
}
}
}
}
}
@ -93,6 +97,7 @@ pub struct DataLoader {
receiver: mpsc::Receiver<TempReturnType>,
/// Video decoder for handling video or stream data.
#[cfg(feature = "ffmpeg")]
decoder: Option<video_rs::decode::Decoder>,
/// Number of images or frames; `u64::MAX` is used for live streams (indicating no limit).
@ -102,10 +107,18 @@ pub struct DataLoader {
with_pb: bool,
}
impl TryFrom<&str> for DataLoader {
type Error = anyhow::Error;
fn try_from(str: &str) -> Result<Self, Self::Error> {
Self::new(str)
}
}
impl DataLoader {
pub fn new(source: &str) -> Result<Self> {
let span = tracing::span!(tracing::Level::INFO, "DataLoader-new");
let _guard = span.enter();
// TODO: multi-types
// Vec<&str>
// Number of frames or stream
let mut nf = 0;
@ -147,6 +160,21 @@ impl DataLoader {
}
// video decoder
#[cfg(not(feature = "ffmpeg"))]
{
match &media_type {
MediaType::Video(Location::Local)
| MediaType::Video(Location::Remote)
| MediaType::Stream => {
anyhow::bail!(
"Video processing requires the features: `ffmpeg`. \
\nConsider enabling them by passing, e.g., `--features ffmpeg`"
);
}
_ => {}
};
}
#[cfg(feature = "ffmpeg")]
let decoder = match &media_type {
MediaType::Video(Location::Local) => Some(Decoder::new(source_path)?),
MediaType::Video(Location::Remote) | MediaType::Stream => {
@ -157,6 +185,7 @@ impl DataLoader {
};
// video & stream frames
#[cfg(feature = "ffmpeg")]
if let Some(decoder) = &decoder {
nf = match decoder.frames() {
Err(_) => u64::MAX,
@ -166,7 +195,7 @@ impl DataLoader {
}
// summary
tracing::info!("{} Found {:?} x{}", CHECK_MARK, media_type, nf);
info!("Found {:?} x{}", media_type, nf);
Ok(DataLoader {
paths,
@ -174,6 +203,7 @@ impl DataLoader {
bound: 50,
receiver: mpsc::sync_channel(1).1,
batch_size: 1,
#[cfg(feature = "ffmpeg")]
decoder,
nf,
with_pb: true,
@ -190,6 +220,11 @@ impl DataLoader {
self
}
pub fn with_batch_size(mut self, x: usize) -> Self {
self.batch_size = x;
self
}
pub fn with_progress_bar(mut self, x: bool) -> Self {
self.with_pb = x;
self
@ -201,11 +236,19 @@ impl DataLoader {
let batch_size = self.batch_size;
let data = self.paths.take().unwrap_or_default();
let media_type = self.media_type.clone();
#[cfg(feature = "ffmpeg")]
let decoder = self.decoder.take();
// Spawn the producer thread
std::thread::spawn(move || {
DataLoader::producer_thread(sender, data, batch_size, media_type, decoder);
DataLoader::producer_thread(
sender,
data,
batch_size,
media_type,
#[cfg(feature = "ffmpeg")]
decoder,
);
});
Ok(self)
@ -216,10 +259,8 @@ impl DataLoader {
mut data: VecDeque<PathBuf>,
batch_size: usize,
media_type: MediaType,
mut decoder: Option<video_rs::decode::Decoder>,
#[cfg(feature = "ffmpeg")] mut decoder: Option<video_rs::decode::Decoder>,
) {
let span = tracing::span!(tracing::Level::INFO, "DataLoader-producer-thread");
let _guard = span.enter();
let mut yis: Vec<DynamicImage> = Vec::with_capacity(batch_size);
let mut yps: Vec<PathBuf> = Vec::with_capacity(batch_size);
@ -228,7 +269,7 @@ impl DataLoader {
while let Some(path) = data.pop_front() {
match Self::try_read(&path) {
Err(err) => {
tracing::warn!("{} {:?} | {:?}", CROSS_MARK, path, err);
warn!("{:?} | {:?}", path, err);
continue;
}
Ok(img) => {
@ -245,6 +286,7 @@ impl DataLoader {
}
}
}
#[cfg(feature = "ffmpeg")]
MediaType::Video(_) | MediaType::Stream => {
if let Some(decoder) = decoder.as_mut() {
let (w, h) = decoder.size();
@ -279,12 +321,12 @@ impl DataLoader {
}
}
}
_ => todo!(),
_ => unimplemented!(),
}
// Deal with remaining data
if !yis.is_empty() && sender.send((yis, yps)).is_err() {
tracing::info!("Receiver dropped, stopping production");
info!("Receiver dropped, stopping production");
}
}
@ -319,13 +361,30 @@ impl DataLoader {
// try to fetch from hub or local cache
if !path.exists() {
let p = Hub::new()?.fetch(path.to_str().unwrap())?.commit()?;
let p = Hub::default().try_fetch(path.to_str().unwrap())?;
path = PathBuf::from(&p);
}
let img = Self::read_into_rgb8(path)?;
Ok(DynamicImage::from(img))
}
pub fn try_read_batch<P: AsRef<Path> + std::fmt::Debug>(
paths: &[P],
) -> Result<Vec<DynamicImage>> {
let images = paths
.iter()
.filter_map(|path| match Self::try_read(path) {
Ok(img) => Some(img),
Err(err) => {
warn!("Failed to read from: {:?}. Error: {:?}", path, err);
None
}
})
.collect();
Ok(images)
}
fn read_into_rgb8<P: AsRef<Path>>(path: P) -> Result<image::RgbImage> {
let path = path.as_ref();
let img = image::ImageReader::open(path)
@ -357,6 +416,7 @@ impl DataLoader {
}
/// Convert images into a video
#[cfg(feature = "ffmpeg")]
pub fn is2v<P: AsRef<Path>>(source: P, subs: &[&str], fps: usize) -> Result<()> {
let paths = Self::load_from_folder(source.as_ref())?;
if paths.is_empty() {
@ -364,9 +424,9 @@ impl DataLoader {
}
let mut encoder = None;
let mut position = Time::zero();
let saveout = Dir::Currnet
let saveout = crate::Dir::Current
.raw_path_with_subs(subs)?
.join(format!("{}.mp4", string_now("-")));
.join(format!("{}.mp4", crate::string_now("-")));
let pb = build_progress_bar(
paths.len() as u64,
"Converting",
@ -406,7 +466,7 @@ impl DataLoader {
// update
pb.set_prefix("Converted");
pb.set_message(saveout.to_str().unwrap_or_default().to_string());
pb.set_style(ProgressStyle::with_template(
pb.set_style(indicatif::ProgressStyle::with_template(
crate::PROGRESS_BAR_STYLE_FINISH_4,
)?);
pb.finish();

63
src/misc/device.rs Normal file
View File

@ -0,0 +1,63 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Device {
Auto(usize),
Cpu(usize),
Cuda(usize),
TensorRT(usize),
CoreML(usize),
// Cann(usize),
// Acl(usize),
// Rocm(usize),
// Rknpu(usize),
// Openvino(usize),
// Onednn(usize),
}
impl std::fmt::Display for Device {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
let x = match self {
Self::Auto(i) => format!("auto:{}", i),
Self::Cpu(i) => format!("cpu:{}", i),
Self::Cuda(i) => format!("cuda:{}", i),
Self::TensorRT(i) => format!("tensorrt:{}", i),
Self::CoreML(i) => format!("mps:{}", i),
};
write!(f, "{}", x)
}
}
impl TryFrom<&str> for Device {
type Error = anyhow::Error;
fn try_from(s: &str) -> Result<Self, Self::Error> {
// device and its id
let d_id: Vec<&str> = s.trim().split(':').collect();
let (d, id) = match d_id.len() {
1 => (d_id[0], 0),
2 => (d_id[0], d_id[1].parse::<usize>().unwrap_or(0)),
_ => anyhow::bail!(
"Fail to parse device string: {s}. Expect: `device:device_id` or `device`. e.g. `cuda:0` or `cuda`"
),
};
// TODO: device-id checking
match d.to_lowercase().as_str() {
"cpu" => Ok(Self::Cpu(id)),
"cuda" => Ok(Self::Cuda(id)),
"trt" | "tensorrt" => Ok(Self::TensorRT(id)),
"coreml" | "mps" => Ok(Self::CoreML(id)),
_ => anyhow::bail!("Unsupported device str: {s:?}."),
}
}
}
impl Device {
pub fn id(&self) -> usize {
match self {
Device::Auto(i) => *i,
Device::Cpu(i) => *i,
Device::Cuda(i) => *i,
Device::TensorRT(i) => *i,
Device::CoreML(i) => *i,
}
}
}

6
src/core/dir.rs → src/misc/dir.rs
View File

@ -4,7 +4,7 @@ pub enum Dir {
Home,
Cache,
Config,
Currnet,
Current,
Document,
Data,
Download,
@ -15,7 +15,7 @@ pub enum Dir {
impl Dir {
pub fn saveout(subs: &[&str]) -> anyhow::Result<std::path::PathBuf> {
Self::Currnet.raw_path_with_subs(subs)
Self::Current.raw_path_with_subs(subs)
}
/// Retrieves the base path for the specified directory type, optionally appending the `usls` subdirectory.
@ -30,7 +30,7 @@ impl Dir {
Dir::Home => dirs::home_dir(),
Dir::Cache => dirs::cache_dir(),
Dir::Config => dirs::config_dir(),
Dir::Currnet => std::env::current_dir().ok(),
Dir::Current => std::env::current_dir().ok(),
_ => None,
};

Some files were not shown because too many files have changed in this diff Show More