mirror of
https://github.com/mii443/usls.git
synced 2025-08-22 15:45:41 +00:00
🐍 v0.1.0 (#53)
This commit is contained in:
Jamjamjon
GitHub
224
.github/workflows/rust-ci.yml
vendored
224
.github/workflows/rust-ci.yml
vendored
@ -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
3
.gitignore
vendored
@ -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/
|
||||
|
||||
76
Cargo.toml
76
Cargo.toml
@ -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/ndarray",
|
||||
"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"
|
||||
|
||||
310
README.md
310
README.md
@ -1,221 +1,161 @@
|
||||
<p align="center">
|
||||
<h2 align="center">usls</h2>
|
||||
</p>
|
||||
<h2 align="center">usls</h2>
|
||||
|
||||
<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="./examples">
|
||||
<img src="https://img.shields.io/badge/Examples-1A86FD?&logo=anki" alt="Examples">
|
||||
</a>
|
||||
<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>
|
||||
|
||||
<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>
|
||||
<!-- 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>
|
||||
**usls** is a Rust library integrated with **ONNXRuntime**, offering a suite of advanced models for **Computer Vision** and **Vision-Language** tasks, including:
|
||||
|
||||
</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:
|
||||
|
||||
- **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:
|
||||
|
||||
<details>
|
||||
<summary>You have two options to link the ONNXRuntime library</summary>
|
||||
- **`auto`**: Automatically downloads prebuilt ONNXRuntime binaries from Pyke’s CDN for supported platforms.
|
||||
|
||||
- ### Option 1: Manual Linking
|
||||
- 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).
|
||||
|
||||
- #### For detailed setup instructions, refer to the [ORT documentation](https://ort.pyke.io/setup/linking).
|
||||
<details>
|
||||
<summary>👉 For Linux or macOS Users</summary>
|
||||
|
||||
- #### For Linux or macOS Users:
|
||||
- Download the ONNX Runtime package 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
|
||||
```
|
||||
- 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
|
||||
```
|
||||
|
||||
- ### Option 2: Automatic Download
|
||||
Just use `--features auto`
|
||||
```shell
|
||||
cargo run -r --example yolo --features auto
|
||||
</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
|
||||
|
||||
* **Using `CUDA`**
|
||||
|
||||
```
|
||||
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
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 🎈 Demo
|
||||
|
||||
```Shell
|
||||
cargo run -r --example yolo # blip, clip, yolop, svtr, db, ...
|
||||
```
|
||||
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`
|
||||
```Shell
|
||||
cargo add usls
|
||||
```
|
||||
Add `usls` as a dependency to your project's `Cargo.toml`
|
||||
|
||||
Or use a specific commit:
|
||||
```Toml
|
||||
[dependencies]
|
||||
usls = { git = "https://github.com/jamjamjon/usls", rev = "commit-sha" }
|
||||
```
|
||||
```Shell
|
||||
cargo add usls -F cuda
|
||||
```
|
||||
|
||||
- #### 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`
|
||||
Or use a specific commit:
|
||||
|
||||
<br/>
|
||||
<details>
|
||||
<summary>example code</summary>
|
||||
```Toml
|
||||
[dependencies]
|
||||
usls = { git = "https://github.com/jamjamjon/usls", rev = "commit-sha" }
|
||||
```
|
||||
|
||||
```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).
|
||||
|
||||
@ -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
6
examples/beit/README.md
Normal 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
52
examples/beit/main.rs
Normal 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(())
|
||||
}
|
||||
@ -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
|
||||
|
||||
@ -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(())
|
||||
}
|
||||
|
||||
@ -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
|
||||
```
|
||||
|
Before Width: | Height: | Size: 176 KiB After Width: | Height: | Size: 176 KiB |
@ -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
6
examples/convnext/README.md
Normal 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
52
examples/convnext/main.rs
Normal 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
5
examples/d-fine/README.md
Normal file
@ -0,0 +1,5 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example d-fine
|
||||
```
|
||||
28
examples/d-fine/main.rs
Normal file
28
examples/d-fine/main.rs
Normal 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
5
examples/dataloader/README.md
Normal file
@ -0,0 +1,5 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example dataloader
|
||||
```
|
||||
@ -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;
|
||||
}
|
||||
|
||||
// write video
|
||||
viewer.write_batch(&images_plotted)?;
|
||||
for (_xs, _paths) in dl {
|
||||
println!("Paths: {:?}", _paths);
|
||||
}
|
||||
|
||||
// finish video write
|
||||
viewer.finish_write()?;
|
||||
// 2. read one image
|
||||
let image = DataLoader::try_read("./assets/bus.jpg")?;
|
||||
println!(
|
||||
"Read one image. Height: {}, Width: {}",
|
||||
image.height(),
|
||||
image.width()
|
||||
);
|
||||
|
||||
// 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(())
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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
7
examples/deim/README.md
Normal file
@ -0,0 +1,7 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example deim
|
||||
```
|
||||
|
||||
|
||||
28
examples/deim/main.rs
Normal file
28
examples/deim/main.rs
Normal 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
7
examples/deit/README.md
Normal 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
52
examples/deit/main.rs
Normal 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(())
|
||||
}
|
||||
@ -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
10
examples/depth-pro/README.md
Normal file
@ -0,0 +1,10 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r -F cuda --example depth-pro -- --device cuda
|
||||
```
|
||||
|
||||
|
||||
## Results
|
||||
|
||||
|
||||
@ -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(())
|
||||
|
||||
@ -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
10
examples/doclayout-yolo/README.md
Normal file
@ -0,0 +1,10 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r -F cuda --example doclayout-yolo -- --device cuda
|
||||
```
|
||||
|
||||
|
||||
## Results
|
||||
|
||||
|
||||
42
examples/doclayout-yolo/main.rs
Normal file
42
examples/doclayout-yolo/main.rs
Normal 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
6
examples/fast/README.md
Normal file
@ -0,0 +1,6 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example fast
|
||||
```
|
||||
|
||||
65
examples/fast/main.rs
Normal file
65
examples/fast/main.rs
Normal 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
5
examples/fastsam/README.md
Normal file
@ -0,0 +1,5 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r -F cuda --example fastsam -- --device cuda
|
||||
```
|
||||
45
examples/fastsam/main.rs
Normal file
45
examples/fastsam/main.rs
Normal 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
13
examples/fastvit/README.md
Normal 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
57
examples/fastvit/main.rs
Normal 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
30
examples/florence2/README.md
Normal 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=''>|
|
||||
|
||||
|
||||
@ -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")?,
|
||||
// tasks
|
||||
let tasks = [
|
||||
// w inputs
|
||||
Task::Caption(0),
|
||||
Task::Caption(1),
|
||||
Task::Caption(2),
|
||||
Task::Ocr,
|
||||
// Task::OcrWithRegion,
|
||||
Task::RegionProposal,
|
||||
Task::ObjectDetection,
|
||||
Task::DenseRegionCaption,
|
||||
// w/o inputs
|
||||
Task::OpenSetDetection("a vehicle"),
|
||||
Task::CaptionToPhraseGrounding("A vehicle with two wheels parked in front of a building."),
|
||||
Task::ReferringExpressionSegmentation("a vehicle"),
|
||||
Task::RegionToSegmentation(
|
||||
// 31, 156, 581, 373, // car
|
||||
449, 270, 556, 372, // wheel
|
||||
),
|
||||
Task::RegionToCategory(
|
||||
// 31, 156, 581, 373,
|
||||
449, 270, 556, 372,
|
||||
),
|
||||
Task::RegionToDescription(
|
||||
// 31, 156, 581, 373,
|
||||
449, 270, 556, 372,
|
||||
),
|
||||
];
|
||||
|
||||
// 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
|
||||
Task::Caption(0),
|
||||
Task::Caption(1),
|
||||
Task::Caption(2),
|
||||
Task::Ocr,
|
||||
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::RegionToSegmentation(
|
||||
quantized_coords[0],
|
||||
quantized_coords[1],
|
||||
quantized_coords[2],
|
||||
quantized_coords[3],
|
||||
),
|
||||
Task::RegionToCategory(
|
||||
quantized_coords[0],
|
||||
quantized_coords[1],
|
||||
quantized_coords[2],
|
||||
quantized_coords[3],
|
||||
),
|
||||
Task::RegionToDescription(
|
||||
quantized_coords[0],
|
||||
quantized_coords[1],
|
||||
quantized_coords[2],
|
||||
quantized_coords[3],
|
||||
),
|
||||
],
|
||||
)?;
|
||||
|
||||
// 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(())
|
||||
}
|
||||
|
||||
@ -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
|
||||
```
|
||||
|
||||
|
||||
|
||||
@ -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
5
examples/hub/README.md
Normal file
@ -0,0 +1,5 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
RUST_LOG=usls=info cargo run -r --example hub
|
||||
```
|
||||
26
examples/hub/main.rs
Normal file
26
examples/hub/main.rs
Normal 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
6
examples/linknet/README.md
Normal file
@ -0,0 +1,6 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example fast
|
||||
```
|
||||
|
||||
65
examples/linknet/main.rs
Normal file
65
examples/linknet/main.rs
Normal 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
13
examples/mobileone/README.md
Normal 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
57
examples/mobileone/main.rs
Normal 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(())
|
||||
}
|
||||
@ -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
10
examples/picodet-layout/README.md
Normal file
@ -0,0 +1,10 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example picodet-layout
|
||||
```
|
||||
|
||||
|
||||
## Results
|
||||
|
||||
|
||||
31
examples/picodet-layout/main.rs
Normal file
31
examples/picodet-layout/main.rs
Normal 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
17
examples/rtdetr/README.md
Normal 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
43
examples/rtdetr/main.rs
Normal 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(())
|
||||
}
|
||||
@ -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(())
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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(())
|
||||
|
||||
@ -1,10 +1,9 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example sapiens
|
||||
cargo run -r -F cuda --example sapiens -- --device cuda
|
||||
```
|
||||
|
||||
|
||||
## Results
|
||||
|
||||
|
||||
|
||||
@ -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
9
examples/slanet/README.md
Normal file
@ -0,0 +1,9 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r -F cuda --example slanet -- --device cuda
|
||||
```
|
||||
|
||||
## Results
|
||||
|
||||
|
||||
48
examples/slanet/main.rs
Normal file
48
examples/slanet/main.rs
Normal 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(())
|
||||
}
|
||||
@ -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")] }])
|
||||
|
||||
```
|
||||
|
Before Width: | Height: | Size: 59 KiB After Width: | Height: | Size: 59 KiB |
|
Before Width: | Height: | Size: 13 KiB After Width: | Height: | Size: 13 KiB |
|
Before Width: | Height: | Size: 15 KiB After Width: | Height: | Size: 15 KiB |
|
Before Width: | Height: | Size: 10 KiB After Width: | Height: | Size: 10 KiB |
|
Before Width: | Height: | Size: 24 KiB After Width: | Height: | Size: 24 KiB |
|
Before Width: | Height: | Size: 14 KiB After Width: | Height: | Size: 14 KiB |
|
Before Width: | Height: | Size: 9.0 KiB After Width: | Height: | Size: 9.0 KiB |
|
Before Width: | Height: | Size: 13 KiB After Width: | Height: | Size: 13 KiB |
|
Before Width: | Height: | Size: 17 KiB After Width: | Height: | Size: 17 KiB |
BIN
examples/svtr/images/text-hello-rust-handwritten.png
Normal file
BIN
examples/svtr/images/text-hello-rust-handwritten.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 3.7 KiB |
@ -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
13
examples/trocr/README.md
Normal 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
96
examples/trocr/main.rs
Normal 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
5
examples/viewer/README.md
Normal 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
43
examples/viewer/main.rs
Normal 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(())
|
||||
}
|
||||
@ -1,7 +1,7 @@
|
||||
## Quick Start
|
||||
|
||||
```shell
|
||||
cargo run -r --example yolo-sam
|
||||
cargo run -r -F cuda --example yolo-sam -- --device cuda
|
||||
```
|
||||
|
||||
## Results
|
||||
|
||||
@ -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(),
|
||||
|
||||
@ -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">` |
|
||||
|
||||
| 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"> |
|
||||
|
||||
| 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">
|
||||
|
||||
| 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">
|
||||
|
||||
| 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">
|
||||
|
||||
|
||||
| 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">` |
|
||||
|
||||
| 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">` |
|
||||
|
||||
| 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">` |
|
||||
|
||||
## 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) |
|
||||
|
||||
@ -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,
|
||||
}
|
||||
|
||||
// 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,
|
||||
}
|
||||
}
|
||||
// plot
|
||||
annotator.annotate(&xs, &ys);
|
||||
}
|
||||
|
||||
// finish video write
|
||||
if !args.nosave {
|
||||
if let Some(viewer) = &mut viewer {
|
||||
viewer.finish_write()?;
|
||||
}
|
||||
}
|
||||
model.summary();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -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
9
examples/yolov8-rtdetr/README.md
Normal 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
45
examples/yolov8-rtdetr/main.rs
Normal 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(())
|
||||
}
|
||||
@ -1,2 +0,0 @@
|
||||
[toolchain]
|
||||
channel = "1.79"
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
@ -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
426
src/core/hub.rs
@ -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(())
|
||||
}
|
||||
}
|
||||
@ -1,6 +0,0 @@
|
||||
#[derive(Debug)]
|
||||
pub enum Metric {
|
||||
IP,
|
||||
Cos,
|
||||
L2,
|
||||
}
|
||||
@ -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;
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
@ -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;
|
||||
}
|
||||
}
|
||||
@ -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();
|
||||
}
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
271
src/lib.rs
271
src/lib.rs
@ -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)?;
|
||||
//! - 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
|
||||
//! ```
|
||||
//!
|
||||
//! // 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()?;
|
||||
//! </details>
|
||||
//! - **`ffmpeg`**: Adds support for video streams, real-time frame visualization, and video export.
|
||||
//!
|
||||
//! // Build annotator
|
||||
//! let annotator = Annotator::new().with_saveout("YOLO-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.
|
||||
//!
|
||||
//! // Run and Annotate images
|
||||
//! for (xs, _) in dl {
|
||||
//! let ys = model.forward(&xs, false)?;
|
||||
//! annotator.annotate(&xs, &ys);
|
||||
//! }
|
||||
//! ## 🎈 Example
|
||||
//!
|
||||
//! Ok(())
|
||||
//! }
|
||||
//! ```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::*;
|
||||
|
||||
@ -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
171
src/misc/color.rs
Normal 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
435
src/misc/colormap256.rs
Normal 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")
|
||||
}
|
||||
}
|
||||
@ -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,21 +28,25 @@ 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(item) => {
|
||||
progress_bar.inc(self.batch_size);
|
||||
Some(item)
|
||||
}
|
||||
None => {
|
||||
progress_bar.set_prefix(" Iterated");
|
||||
progress_bar.set_style(
|
||||
indicatif::ProgressStyle::with_template(crate::PROGRESS_BAR_STYLE_FINISH_2)
|
||||
.unwrap(),
|
||||
Some(progress_bar) => {
|
||||
match self.receiver.recv().ok() {
|
||||
Some(item) => {
|
||||
progress_bar.inc(self.batch_size);
|
||||
Some(item)
|
||||
}
|
||||
None => {
|
||||
progress_bar.set_prefix("Iterated");
|
||||
progress_bar.set_style(
|
||||
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
|
||||
progress_bar.finish();
|
||||
None
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -55,7 +59,7 @@ impl IntoIterator for DataLoader {
|
||||
let progress_bar = if self.with_pb {
|
||||
build_progress_bar(
|
||||
self.nf,
|
||||
" Iterating",
|
||||
"Iterating",
|
||||
Some(&format!("{:?}", self.media_type)),
|
||||
crate::PROGRESS_BAR_STYLE_CYAN_2,
|
||||
)
|
||||
@ -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,12 +424,12 @@ 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",
|
||||
"Converting",
|
||||
Some(&format!("{:?}", MediaType::Video(Location::Local))),
|
||||
crate::PROGRESS_BAR_STYLE_CYAN_2,
|
||||
)?;
|
||||
@ -404,9 +464,9 @@ impl DataLoader {
|
||||
}
|
||||
|
||||
// update
|
||||
pb.set_prefix(" Converted");
|
||||
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
63
src/misc/device.rs
Normal 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -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
Reference in New Issue
Block a user