diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 53f8242..516bf22 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -3,4 +3,5 @@ updates: - package-ecosystem: "cargo" directory: "/" schedule: - interval: "weekly" + interval: "monthly" + open-pull-requests-limit: 3 \ No newline at end of file diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index c35fadf..273c163 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -1,83 +1,102 @@ -name: Rust-CI +name: CI on: push: - branches: [ "main", "dev", "develop", "x", "xy" , "xyz" ] + branches: [ "main", "dev", "x" ] pull_request: branches: [ "main" ] env: CARGO_TERM_COLOR: always - jobs: - check: - name: Check - runs-on: ${{ matrix.os }} - strategy: - matrix: - os: [ubuntu-latest, macOS-latest, windows-latest] - rust: [stable] + lints: + name: Rustfmt & Clippy + runs-on: ubuntu-latest + container: jrottenberg/ffmpeg:7.1-ubuntu + steps: - - uses: actions/checkout@v2 - - uses: actions-rs/toolchain@v1 + - name: Checkout + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + DEBIAN_FRONTEND=noninteractive apt-get update + DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler + + - name: Setup Rust + uses: dtolnay/rust-toolchain@stable with: - profile: minimal - toolchain: ${{ matrix.rust }} - override: true - - uses: actions-rs/cargo@v1 - with: - command: check - args: --all + components: rustfmt, clippy + + - name: Rustfmt + run: cargo fmt --all -- --check + + - name: Clippy + run: cargo clippy --all-features --all-targets -- -D warnings + + check: + name: cargo-check + runs-on: ubuntu-latest + container: jrottenberg/ffmpeg:7.1-ubuntu + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + DEBIAN_FRONTEND=noninteractive apt-get update + DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler + + - name: Setup Rust + uses: dtolnay/rust-toolchain@stable + + - name: Check + run: cargo check --all-features --all-targets test: - name: Test - runs-on: ${{ matrix.os }} + name: cargo-test + runs-on: ubuntu-latest + container: jrottenberg/ffmpeg:7.1-ubuntu + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + DEBIAN_FRONTEND=noninteractive apt-get update + DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler + + - name: Setup Rust + uses: dtolnay/rust-toolchain@nightly + + - name: Test + run: cargo +nightly test --all-features --all-targets + + build-linux: + needs: test + name: cargo build / linux / ffmpeg ${{ matrix.ffmpeg_version }} + runs-on: ubuntu-latest + container: jrottenberg/ffmpeg:${{ matrix.ffmpeg_version }}-ubuntu + strategy: matrix: - os: [ubuntu-latest, macOS-latest, windows-latest] - rust: [stable] - steps: - - uses: actions/checkout@v2 - - uses: actions-rs/toolchain@v1 - with: - profile: minimal - toolchain: ${{ matrix.rust }} - override: true - - uses: actions-rs/cargo@v1 - with: - command: test - args: --all + ffmpeg_version: [ "5.0", "5.1", "6.0", "6.1", "7.0", "7.1" ] # "4.3", "4.4" + fail-fast: false - fmt: - name: Rustfmt - 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 rustfmt - - uses: actions-rs/cargo@v1 - with: - command: fmt - args: --all -- --check + - name: Checkout + uses: actions/checkout@v4 - 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: Install dependencies + run: | + DEBIAN_FRONTEND=noninteractive apt-get update + DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler + - name: Setup Rust + uses: dtolnay/rust-toolchain@stable + + - name: Build + run: cargo build --all-features \ No newline at end of file diff --git a/.gitignore b/.gitignore index e1a526e..8db4ec5 100644 --- a/.gitignore +++ b/.gitignore @@ -3,8 +3,6 @@ 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 @@ -18,4 +16,4 @@ Cargo.lock .debug .vscode runs/ -.DS_Store +**/*.DS_Store diff --git a/Cargo.toml b/Cargo.toml index 5b2d0c2..1cd2b59 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,67 +1,62 @@ [package] name = "usls" -version = "0.1.0" -rust-version = "1.82" edition = "2021" +version = "0.1.0-beta.1" +rust-version = "1.82" description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models." repository = "https://github.com/jamjamjon/usls" authors = ["Jamjamjon "] license = "MIT" readme = "README.md" -exclude = ["assets/*", "examples/*", "runs/*", "benches/*"] +exclude = ["assets/*", "examples/*", "runs/*", "benches/*", "tests/*"] [dependencies] -aksr = { version = "0.0.2" } +anyhow = { version = "1" } +aksr = { version = "0.0.3" } +ab_glyph = { version = "0.2.29" } image = { version = "0.25" } imageproc = { version = "0.25" } ndarray = { version = "0.16.1", features = ["rayon", "serde"] } -rayon = { version = "1.10.0" } -anyhow = { version = "1.0" } -regex = { version = "1.11.1" } +indicatif = { version = "0.17.11" } +log = "0.4.26" +minifb = { version = "0.28.0" } rand = { version = "0.8.5" } -chrono = { version = "0.4.30" } -tokenizers = { version = "0.21.0" } -log = { version = "0.4.22" } -indicatif = "0.17.8" -serde_json = "1.0" +ureq = { version = "2", default-features = true, features = [ "socks-proxy" ] } serde = { version = "1.0", features = ["derive"] } -ort = { version = "2.0.0-rc.9", default-features = false} -prost = "0.12.6" -ab_glyph = "0.2.29" -tempfile = "3.19.1" +serde_json = "1.0" +rayon = { version = "1.10.0" } +glob = "0.3.2" dirs = "6.0.0" -geo = "0.30.0" natord = "1.0.9" -half = { version = "2.3.1" } -ureq = { version = "2.12.1", default-features = false, features = [ "tls" ] } -fast_image_resize = { version = "5.1.2", features = ["image"]} -video-rs = { version = "0.10.3", features = ["ndarray"], optional = true } -minifb = { version = "0.28.0", optional = true } +geo = "0.30.0" +chrono = "0.4.40" +regex = "1.11.1" sha2 = "0.10.8" -ndarray-npy = "0.9.1" +tempfile = "3.19.1" +video-rs = { version = "0.10.3", features = ["ndarray"], optional = true } +fast_image_resize = { version = "5.1.2", features = ["image"]} +ndarray-npy = "0.9.1" +half = { version = "2.3.1" } +prost = "0.13.5" +ort = { version = "2.0.0-rc.9", default-features = false, optional = true , features = [ + "ndarray", + "copy-dylibs", + "half" +]} +tokenizers = { version = "0.21.1" } + +[build-dependencies] +prost-build = "0.13.5" [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/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" ] - -[profile.release] -lto = true -strip = true -panic = "abort" +default = ["ort-download-binaries"] +ort-download-binaries = ["ort", "ort/download-binaries"] +ort-load-dynamic = ["ort", "ort/load-dynamic"] +cuda = ["ort/cuda"] +trt = ["ort/tensorrt"] +mps = ["ort/coreml"] +video = ["dep:video-rs"] diff --git a/README.md b/README.md index 18ead33..dda3d31 100644 --- a/README.md +++ b/README.md @@ -1,47 +1,72 @@

usls

-

- - Rust Continuous Integration Badge - - - usls Version - + Rust MSRV + ONNXRuntime MSRV + - CUDA MSRV + CUDA MSRV + + + cuDNN MSRV + + TensorRT MSRV - - Crates.io Total Downloads -

+ Examples + - usls documentation + Documentation

+

+ + + Rust CI + + + Crates.io Version + + + + Crates.io Downloads + +

+

+ ⭐️ Star if helpful! ⭐️ +

-**usls** is a Rust library integrated with **ONNXRuntime**, offering a suite of advanced models for **Computer Vision** and **Vision-Language** tasks, including: +**usls** is an evolving Rust library focused on inference for advanced **vision** and **vision-language** models, along with practical vision utilities. + +- **SOTA Model Inference:** Supports a wide range of state-of-the-art vision and multi-modal models (typically with fewer than 1B parameters). +- **Multi-backend Acceleration:** Supports CPU, CUDA, TensorRT, and CoreML. +- **Easy Data Handling:** Easily read images, video streams, and folders with iterator support. +- **Rich Result Types:** Built-in containers for common vision outputs like bounding boxes (Hbb, Obb), polygons, masks, etc. +- **Annotation & Visualization:** Draw and display inference results directly, similar to OpenCV's `imshow()`. + + +## 🧩 Supported Models - **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), [YOLOv12](https://github.com/sunsmarterjie/yolov12) - **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**: [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), [Moondream2](https://github.com/vikhyat/moondream/tree/main) -- **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) +- **OCR-Related 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)
-👉 More Supported Models +Full list of supported models (click to expand) | Model | Task / Description | Example | CoreML | CUDA
FP32 | CUDA
FP16 | TensorRT
FP32 | TensorRT
FP16 | | -------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | ---------------------------- | ------ | -------------- | -------------- | ------------------ | ------------------ | @@ -94,73 +119,149 @@
-## ⛳️ Cargo Features -By default, **none of the following features are enabled**. You can enable them as needed: -- **`auto`**: Automatically downloads prebuilt ONNXRuntime binaries from Pyke’s CDN for supported platforms. +## 🛠️ Installation +**Note:** It is recommended to use the GitHub repository as the source, since the crates.io version may not be up-to-date. - - 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 Linux or macOS Users - - - 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 - ``` - -
-- **`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 - ``` - -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 -F cuda -``` - -Or use a specific commit: - -```Toml +```toml [dependencies] -usls = { git = "https://github.com/jamjamjon/usls", rev = "commit-sha" } +usls = { git = "https://github.com/jamjamjon/usls" } + +# crates.io version +usls = "latest-version" ``` -## 🥳 If you find this helpful, please give it a star ⭐ +## ⚡ Cargo Features +- **ONNXRuntime-related features (enabled by default)**, provide model inference and model zoo support: + - **`ort-download-binaries`** (**default**): Automatically downloads prebuilt `ONNXRuntime` binaries for supported platforms. Provides core model loading and inference capabilities using the `CPU` execution provider. + - **`ort-load-dynamic `** Dynamic linking. 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. [See the guide here](https://ort.pyke.io/setup/linking#dynamic-linking). + + - **`cuda`**: Enables the NVIDIA `CUDA` provider. Requires `CUDA` toolkit and `cuDNN` installed. + - **`trt`**: Enables the NVIDIA `TensorRT` provider. Requires `TensorRT` libraries installed. + - **`mps`**: Enables the Apple `CoreML` provider for macOS. -## 📌 License +- **If you only need basic features** (such as image/video reading, result visualization, etc.), you can disable the default features to minimize dependencies: + ```shell + usls = { git = "https://github.com/jamjamjon/usls", default-features = false } + ``` + - **`video`** : Enable video stream reading, and video writing.(Note: Powered by [video-rs](https://github.com/oddity-ai/video-rs) and [minifb](https://github.com/emoon/rust_minifb). Check their repositories for potential issues.) + +## ✨ Example + +- Model Inference + ```shell + cargo run -r --example yolo # CPU + cargo run -r -F cuda --example yolo -- --device cuda:0 # GPU + ``` + +- Reading Images + ```rust + // Read a single image + let image = DataLoader::try_read_one("./assets/bus.jpg")?; + + // Read multiple images + let images = DataLoader::try_read_n(&["./assets/bus.jpg", "./assets/cat.png"])?; + + // Read all images in a folder + let images = DataLoader::try_read_folder("./assets")?; + + // Read images matching a pattern (glob) + let images = DataLoader::try_read_pattern("./assets/*.Jpg")?; + + // Load images and iterate + let dl = DataLoader::new("./assets")?.with_batch(2).build()?; + for images in dl.iter() { + // Code here + } + ``` + +- Reading Video + ```rust + let dl = DataLoader::new("http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4")? + .with_batch(1) + .with_nf_skip(2) + .with_progress_bar(true) + .build()?; + for images in dl.iter() { + // Code here + } + ``` + +- Annotate + ```rust + let annotator = Annotator::default(); + let image = DataLoader::try_read_one("./assets/bus.jpg")?; + // hbb + let hbb = Hbb::default() + .with_xyxy(669.5233, 395.4491, 809.0367, 878.81226) + .with_id(0) + .with_name("person") + .with_confidence(0.87094545); + let _ = annotator.annotate(&image, &hbb)?; + + // keypoints + let keypoints: Vec = vec![ + Keypoint::default() + .with_xy(139.35767, 443.43655) + .with_id(0) + .with_name("nose") + .with_confidence(0.9739332), + Keypoint::default() + .with_xy(147.38545, 434.34055) + .with_id(1) + .with_name("left_eye") + .with_confidence(0.9098319), + Keypoint::default() + .with_xy(128.5701, 434.07516) + .with_id(2) + .with_name("right_eye") + .with_confidence(0.9320564), + ]; + let _ = annotator.annotate(&image, &keypoints)?; + ``` + + +- Visualizing Inference Results and Exporting Video + ```rust + let dl = DataLoader::new(args.source.as_str())?.build()?; + let mut viewer = Viewer::default().with_window_scale(0.5); + + for images in &dl { + // Check if the window exists and is open + if viewer.is_window_exist() && !viewer.is_window_open() { + break; + } + + // Show image in window + viewer.imshow(&images[0])?; + + // Handle key events and delay + if let Some(key) = viewer.wait_key(1) { + if key == usls::Key::Escape { + break; + } + } + + // Your custom code here + + // Write video frame (requires video feature) + // if args.save_video { + // viewer.write_video_frame(&images[0])?; + // } + } + ``` + +**All examples are located in the [examples](./examples/) directory.** + +## ❓ FAQ +See issues or open a new discussion. + +## 🤝 Contributing + +Contributions are welcome! If you have suggestions, bug reports, or want to add new features or models, feel free to open an issue or submit a pull request. + + +## 📜 License This project is licensed under [LICENSE](LICENSE). diff --git a/assets/cat.png b/assets/cat.png new file mode 100644 index 0000000..67850ac Binary files /dev/null and b/assets/cat.png differ diff --git a/assets/dog.jpg b/assets/dog.jpg new file mode 100644 index 0000000..583f69e Binary files /dev/null and b/assets/dog.jpg differ diff --git a/build.rs b/build.rs index b1335e6..4231433 100644 --- a/build.rs +++ b/build.rs @@ -1,5 +1,10 @@ -fn main() { - // Need this for CoreML. See: https://ort.pyke.io/perf/execution-providers#coreml - #[cfg(target_os = "macos")] +use std::io::Result; + +fn main() -> Result<()> { + prost_build::compile_protos(&["src/utils/onnx.proto3"], &["src"])?; + + #[cfg(any(target_os = "macos", target_os = "ios", target_os = "tvos"))] println!("cargo:rustc-link-arg=-fapple-link-rtlib"); + + Ok(()) } diff --git a/examples/annotate/README.md b/examples/annotate/README.md new file mode 100644 index 0000000..c428d7a --- /dev/null +++ b/examples/annotate/README.md @@ -0,0 +1,5 @@ +## Quick Start + +```shell +cargo run -r --example annotate +``` diff --git a/examples/annotate/main.rs b/examples/annotate/main.rs new file mode 100644 index 0000000..ce77b25 --- /dev/null +++ b/examples/annotate/main.rs @@ -0,0 +1,734 @@ +use usls::{Annotator, DataLoader, Hbb, Keypoint, Polygon, Prob, Style, SKELETON_COCO_19, Y}; + +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(); + + // load images + let image = DataLoader::try_read_one("./assets/bus.jpg")?; + println!("Read 1 images: {:?}.", image.dimensions()); + + let hbbs = vec![ + Hbb::default() + .with_xyxy(20.81192, 229.65482, 795.1383, 751.0504) + .with_id(5) + .with_name("bus") + .with_confidence(0.8815875) + .with_style( + // individual setting + Style::hbb() + .with_thickness(5) + .with_draw_fill(true) + .with_visible(true) + .with_text_visible(true) + .show_confidence(true) + .show_id(true) + .show_name(true) + .with_text_loc(usls::TextLoc::Center) + .with_color( + usls::StyleColors::default() + .with_outline(usls::Color::white()) + .with_fill(usls::Color::black().with_alpha(100)) + .with_text(usls::Color::black()) + .with_text_bg(usls::Color::white()), + ), + ), + Hbb::default() + .with_xyxy(669.5233, 395.4491, 809.0367, 878.81226) + .with_id(0) + .with_name("person") + .with_confidence(0.87094545), + Hbb::default() + .with_xyxy(48.03354, 398.6103, 245.06848, 902.5964) + .with_id(0) + .with_name("person") + .with_confidence(0.8625425), + Hbb::default() + .with_xyxy(221.26727, 405.51895, 345.14288, 857.61865) + .with_id(0) + .with_name("person") + .with_confidence(0.81437635), + Hbb::default() + .with_xyxy(0.08129883, 254.67389, 32.30627, 324.9663) + .with_id(11) + .with_name("stop sign") + .with_confidence(0.30021638), + ]; + + let keypoints: Vec = vec![ + Keypoint::default() + .with_xy(139.35767, 443.43655) + .with_id(0) + .with_name("nose") + .with_confidence(0.9739332), + Keypoint::default() + .with_xy(147.38545, 434.34055) + .with_id(1) + .with_name("left_eye") + .with_confidence(0.9098319), + Keypoint::default() + .with_xy(128.5701, 434.07516) + .with_id(2) + .with_name("right_eye") + .with_confidence(0.9320564), + Keypoint::default() + .with_xy(153.24237, 442.4857) + .with_id(3) + .with_name("left_ear") + .with_confidence(0.5992247), + Keypoint::default() + .with_xy(105.74312, 441.05765) + .with_id(4) + .with_name("right_ear") + .with_confidence(0.7259705), + Keypoint::default() + .with_xy(166.55661, 498.17484) + .with_id(5) + .with_name("left_shoulder") + .with_confidence(0.9862031), + Keypoint::default() + .with_xy(89.40589, 497.6169) + .with_id(6) + .with_name("right_shoulder") + .with_confidence(0.9879458), + Keypoint::default() + .with_xy(190.7351, 575.00226) + .with_id(7) + .with_name("left_elbow") + .with_confidence(0.9521556), + Keypoint::default() + .with_xy(116.3187, 570.6441) + .with_id(8) + .with_name("right_elbow") + .with_confidence(0.9619827), + Keypoint::default() + .with_xy(140.43465, 575.80994) + .with_id(9) + .with_name("left_wrist") + .with_confidence(0.9329945), + Keypoint::default() + .with_xy(174.73381, 558.4027) + .with_id(10) + .with_name("right_wrist") + .with_confidence(0.93989426), + Keypoint::default() + .with_xy(159.16801, 652.35846) + .with_id(11) + .with_name("left_hip") + .with_confidence(0.9849887), + Keypoint::default() + .with_xy(99.27675, 653.01874) + .with_id(12) + .with_name("right_hip") + .with_confidence(0.9861814), + Keypoint::default() + .with_xy(180.95883, 759.8797) + .with_id(13) + .with_name("left_knee") + .with_confidence(0.95086014), + Keypoint::default() + .with_xy(87.09352, 762.6029) + .with_id(14) + .with_name("right_knee") + .with_confidence(0.9532267), + Keypoint::default() + .with_xy(194.39137, 860.7901) + .with_id(15) + .with_name("left_ankle") + .with_confidence(0.7986185), + Keypoint::default() + .with_xy(70.85685, 862.53253) + .with_id(16) + .with_name("right_ankle") + .with_confidence(0.79832363), + ]; + + let probs = vec![ + Prob::default() + .with_id(654) + .with_name("minibus") + .with_confidence(0.666985), + Prob::default() + .with_id(734) + .with_name("police_van") + .with_confidence(0.20067203), + Prob::default() + .with_id(874) + .with_name("trolleybus") + .with_confidence(0.024672432), + Prob::default() + .with_id(656) + .with_name("minivan") + .with_confidence(0.02395765), + Prob::default() + .with_id(757) + .with_name("recreational_vehicle") + .with_confidence(0.012205753), + ]; + + let polygons = vec![ + Polygon::from_xys(&[ + [13.0, 251.0], + [12.0, 251.0], + [11.0, 251.0], + [10.0, 251.0], + [9.0, 251.0], + [8.0, 251.0], + [7.0, 251.0], + [6.0, 251.0], + [5.0, 251.0], + [4.0, 251.0], + [3.0, 251.0], + [2.0, 251.0], + [1.0, 251.0], + [0.0, 251.0], + [0.0, 252.0], + [0.0, 253.0], + [0.0, 254.0], + [0.0, 255.0], + [0.0, 256.0], + [0.0, 257.0], + [0.0, 258.0], + [0.0, 259.0], + [0.0, 260.0], + [0.0, 261.0], + [0.0, 262.0], + [0.0, 263.0], + [0.0, 264.0], + [0.0, 265.0], + [0.0, 266.0], + [0.0, 267.0], + [0.0, 268.0], + [0.0, 269.0], + [0.0, 270.0], + [0.0, 271.0], + [0.0, 272.0], + [0.0, 273.0], + [0.0, 274.0], + [0.0, 275.0], + [0.0, 276.0], + [0.0, 277.0], + [0.0, 278.0], + [0.0, 279.0], + [0.0, 280.0], + [0.0, 281.0], + [0.0, 282.0], + [0.0, 283.0], + [0.0, 284.0], + [0.0, 285.0], + [0.0, 286.0], + [0.0, 287.0], + [0.0, 288.0], + [0.0, 289.0], + [0.0, 290.0], + [0.0, 291.0], + [0.0, 292.0], + [0.0, 293.0], + [0.0, 294.0], + [0.0, 295.0], + [0.0, 296.0], + [0.0, 297.0], + [0.0, 298.0], + [0.0, 299.0], + [0.0, 300.0], + [0.0, 301.0], + [0.0, 302.0], + [0.0, 303.0], + [0.0, 304.0], + [0.0, 305.0], + [0.0, 306.0], + [0.0, 307.0], + [0.0, 308.0], + [0.0, 309.0], + [0.0, 310.0], + [0.0, 311.0], + [0.0, 312.0], + [0.0, 313.0], + [0.0, 314.0], + [0.0, 315.0], + [0.0, 316.0], + [0.0, 317.0], + [0.0, 318.0], + [0.0, 319.0], + [0.0, 320.0], + [0.0, 321.0], + [0.0, 322.0], + [0.0, 323.0], + [0.0, 324.0], + [0.0, 325.0], + [1.0, 325.0], + [2.0, 325.0], + [3.0, 325.0], + [4.0, 325.0], + [5.0, 325.0], + [6.0, 325.0], + [7.0, 325.0], + [8.0, 325.0], + [9.0, 325.0], + [10.0, 325.0], + [11.0, 325.0], + [12.0, 324.0], + [13.0, 324.0], + [14.0, 324.0], + [15.0, 323.0], + [16.0, 323.0], + [17.0, 322.0], + [18.0, 321.0], + [19.0, 321.0], + [20.0, 320.0], + [20.0, 319.0], + [21.0, 318.0], + [22.0, 317.0], + [23.0, 316.0], + [24.0, 315.0], + [24.0, 314.0], + [25.0, 313.0], + [26.0, 312.0], + [27.0, 311.0], + [28.0, 310.0], + [29.0, 309.0], + [30.0, 308.0], + [30.0, 307.0], + [31.0, 306.0], + [31.0, 305.0], + [31.0, 304.0], + [32.0, 303.0], + [32.0, 302.0], + [32.0, 301.0], + [33.0, 300.0], + [33.0, 299.0], + [33.0, 298.0], + [33.0, 297.0], + [33.0, 296.0], + [33.0, 295.0], + [33.0, 294.0], + [33.0, 293.0], + [33.0, 292.0], + [33.0, 291.0], + [33.0, 290.0], + [33.0, 289.0], + [33.0, 288.0], + [33.0, 287.0], + [33.0, 286.0], + [33.0, 285.0], + [33.0, 284.0], + [33.0, 283.0], + [33.0, 282.0], + [33.0, 281.0], + [33.0, 280.0], + [32.0, 279.0], + [32.0, 278.0], + [32.0, 277.0], + [31.0, 276.0], + [31.0, 275.0], + [31.0, 274.0], + [30.0, 273.0], + [30.0, 272.0], + [29.0, 271.0], + [28.0, 270.0], + [28.0, 269.0], + [27.0, 268.0], + [27.0, 267.0], + [26.0, 266.0], + [25.0, 265.0], + [25.0, 264.0], + [24.0, 263.0], + [23.0, 262.0], + [22.0, 261.0], + [21.0, 260.0], + [20.0, 259.0], + [20.0, 258.0], + [19.0, 257.0], + [18.0, 256.0], + [17.0, 255.0], + [16.0, 254.0], + [15.0, 254.0], + [14.0, 253.0], + [13.0, 252.0], + [13.0, 251.0], + ]) + .with_id(11) + .with_name("stop sign") + .with_confidence(0.5555), + Polygon::from_xys(&[ + [485.0, 149.0], + [484.0, 150.0], + [484.0, 151.0], + [483.0, 152.0], + [482.0, 153.0], + [481.0, 153.0], + [480.0, 153.0], + [479.0, 153.0], + [478.0, 153.0], + [477.0, 154.0], + [476.0, 154.0], + [475.0, 154.0], + [474.0, 154.0], + [473.0, 154.0], + [472.0, 154.0], + [471.0, 154.0], + [470.0, 154.0], + [469.0, 154.0], + [468.0, 155.0], + [467.0, 155.0], + [466.0, 155.0], + [465.0, 155.0], + [464.0, 155.0], + [463.0, 155.0], + [462.0, 156.0], + [461.0, 156.0], + [460.0, 156.0], + [459.0, 156.0], + [458.0, 156.0], + [457.0, 157.0], + [456.0, 157.0], + [455.0, 157.0], + [454.0, 157.0], + [453.0, 158.0], + [452.0, 158.0], + [451.0, 158.0], + [450.0, 158.0], + [449.0, 159.0], + [448.0, 159.0], + [447.0, 159.0], + [446.0, 159.0], + [445.0, 160.0], + [444.0, 160.0], + [443.0, 160.0], + [442.0, 160.0], + [441.0, 160.0], + [440.0, 161.0], + [439.0, 161.0], + [438.0, 161.0], + [437.0, 161.0], + [436.0, 161.0], + [435.0, 162.0], + [434.0, 162.0], + [433.0, 162.0], + [432.0, 162.0], + [431.0, 162.0], + [430.0, 162.0], + [429.0, 163.0], + [428.0, 163.0], + [427.0, 163.0], + [427.0, 164.0], + [427.0, 165.0], + [427.0, 166.0], + [427.0, 167.0], + [427.0, 168.0], + [427.0, 169.0], + [427.0, 170.0], + [427.0, 171.0], + [427.0, 172.0], + [427.0, 173.0], + [427.0, 174.0], + [427.0, 175.0], + [427.0, 176.0], + [427.0, 177.0], + [427.0, 178.0], + [427.0, 179.0], + [427.0, 180.0], + [427.0, 181.0], + [427.0, 182.0], + [427.0, 183.0], + [427.0, 184.0], + [427.0, 185.0], + [427.0, 186.0], + [427.0, 187.0], + [427.0, 188.0], + [427.0, 189.0], + [427.0, 190.0], + [428.0, 190.0], + [429.0, 191.0], + [430.0, 191.0], + [431.0, 191.0], + [432.0, 191.0], + [433.0, 191.0], + [434.0, 191.0], + [435.0, 191.0], + [436.0, 191.0], + [437.0, 191.0], + [438.0, 190.0], + [439.0, 190.0], + [440.0, 190.0], + [441.0, 190.0], + [442.0, 190.0], + [443.0, 190.0], + [444.0, 190.0], + [445.0, 189.0], + [446.0, 189.0], + [447.0, 189.0], + [448.0, 189.0], + [449.0, 189.0], + [450.0, 189.0], + [451.0, 188.0], + [452.0, 188.0], + [453.0, 188.0], + [454.0, 188.0], + [455.0, 188.0], + [456.0, 188.0], + [457.0, 187.0], + [458.0, 187.0], + [459.0, 187.0], + [460.0, 187.0], + [461.0, 186.0], + [462.0, 186.0], + [463.0, 187.0], + [464.0, 188.0], + [465.0, 189.0], + [466.0, 190.0], + [467.0, 191.0], + [467.0, 192.0], + [468.0, 193.0], + [469.0, 193.0], + [470.0, 193.0], + [471.0, 193.0], + [472.0, 193.0], + [473.0, 193.0], + [474.0, 193.0], + [475.0, 193.0], + [476.0, 193.0], + [477.0, 193.0], + [478.0, 192.0], + [479.0, 191.0], + [480.0, 190.0], + [481.0, 190.0], + [482.0, 189.0], + [483.0, 189.0], + [484.0, 189.0], + [485.0, 188.0], + [486.0, 188.0], + [487.0, 188.0], + [488.0, 188.0], + [489.0, 188.0], + [490.0, 188.0], + [491.0, 188.0], + [492.0, 188.0], + [493.0, 187.0], + [494.0, 187.0], + [495.0, 187.0], + [496.0, 187.0], + [497.0, 187.0], + [498.0, 187.0], + [499.0, 187.0], + [500.0, 186.0], + [501.0, 186.0], + [502.0, 186.0], + [503.0, 186.0], + [504.0, 185.0], + [505.0, 185.0], + [506.0, 185.0], + [507.0, 184.0], + [508.0, 184.0], + [509.0, 183.0], + [510.0, 183.0], + [511.0, 183.0], + [512.0, 182.0], + [513.0, 182.0], + [514.0, 182.0], + [515.0, 181.0], + [516.0, 181.0], + [517.0, 181.0], + [518.0, 180.0], + [519.0, 180.0], + [520.0, 180.0], + [521.0, 179.0], + [522.0, 179.0], + [523.0, 178.0], + [524.0, 178.0], + [525.0, 177.0], + [526.0, 176.0], + [527.0, 175.0], + [528.0, 174.0], + [529.0, 173.0], + [530.0, 172.0], + [531.0, 172.0], + [531.0, 171.0], + [531.0, 170.0], + [531.0, 169.0], + [531.0, 168.0], + [531.0, 167.0], + [531.0, 166.0], + [531.0, 165.0], + [531.0, 164.0], + [531.0, 163.0], + [531.0, 162.0], + [531.0, 161.0], + [531.0, 160.0], + [531.0, 159.0], + [531.0, 158.0], + [531.0, 157.0], + [531.0, 156.0], + [530.0, 155.0], + [530.0, 154.0], + [529.0, 154.0], + [528.0, 153.0], + [527.0, 152.0], + [526.0, 151.0], + [525.0, 150.0], + [524.0, 149.0], + [523.0, 149.0], + [522.0, 149.0], + [521.0, 149.0], + [520.0, 149.0], + [519.0, 149.0], + [518.0, 149.0], + [517.0, 149.0], + [516.0, 149.0], + [515.0, 149.0], + [514.0, 149.0], + [513.0, 149.0], + [512.0, 149.0], + [511.0, 149.0], + [510.0, 149.0], + [509.0, 149.0], + [508.0, 149.0], + [507.0, 149.0], + [506.0, 149.0], + [505.0, 149.0], + [504.0, 149.0], + [503.0, 149.0], + [502.0, 149.0], + [501.0, 149.0], + [500.0, 149.0], + [499.0, 149.0], + [498.0, 149.0], + [497.0, 149.0], + [496.0, 149.0], + [495.0, 149.0], + [494.0, 149.0], + [493.0, 149.0], + [492.0, 149.0], + [491.0, 149.0], + [490.0, 149.0], + [489.0, 149.0], + [488.0, 149.0], + [487.0, 149.0], + [486.0, 149.0], + [485.0, 149.0], + ]) + .with_id(9) + .with_name("traffic light") + .with_confidence(0.777777), + ]; + + // Build annotator + let annotator = Annotator::default() + .with_prob_style(Style::prob().with_text_loc(usls::TextLoc::InnerTopLeft)) + .with_hbb_style(Style::hbb().with_thickness(5).with_draw_fill(true)) + .with_keypoint_style( + Style::keypoint() + .with_skeleton(SKELETON_COCO_19.into()) + .with_radius(4) + .with_text_visible(true) + .show_confidence(false) + .show_id(true) + .show_name(false), + ) + .with_polygon_style( + Style::polygon() + .with_text_visible(true) + .show_confidence(true) + .show_id(true) + .show_name(true), + ); + + // Annotate Y + let y = Y::default() + .with_probs(&probs) + .with_hbbs(&hbbs) + .with_keypoints(&keypoints) + // .with_keypointss(&[keypoints.clone()]) + .with_polygons(&polygons); + annotator.annotate(&image, &y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Y"])? + .join(usls::timestamp(None)) + .display(), + ))?; + + // Annotate Probs + annotator.annotate(&image, &probs)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Probs"])? + .join(usls::timestamp(None)) + .display(), + ))?; + + // Annotate Prob + for prob in &probs { + annotator.annotate(&image, prob)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Prob"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } + + // Annotate Hbbs + annotator.annotate(&image, &hbbs)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Hbbs"])? + .join(usls::timestamp(None)) + .display(), + ))?; + + // Annotate Hbb + for hbb in &hbbs { + annotator.annotate(&image, hbb)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Hbb"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } + + // Annotate A set of Keypoint + annotator.annotate(&image, &keypoints)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Keypoints"])? + .join(usls::timestamp(None)) + .display(), + ))?; + + // Annotate Keypoint + for keypoint in &keypoints { + annotator.annotate(&image, keypoint)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Keypoint"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } + + // Annotate Polygons + annotator.annotate(&image, &polygons)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Polygons"])? + .join(usls::timestamp(None)) + .display(), + ))?; + + // Annotate Polygon + for polygon in &polygons { + annotator.annotate(&image, polygon)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Annotate", "Polygon"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } + + Ok(()) +} diff --git a/examples/beit/README.md b/examples/beit/README.md deleted file mode 100644 index d9eddd8..0000000 --- a/examples/beit/README.md +++ /dev/null @@ -1,6 +0,0 @@ -## Quick Start - -```shell -cargo run -r -F cuda --example beit -- --device cuda --dtype fp16 -``` - diff --git a/examples/beit/main.rs b/examples/beit/main.rs deleted file mode 100644 index aad67bd..0000000 --- a/examples/beit/main.rs +++ /dev/null @@ -1,52 +0,0 @@ -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, -} - -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(()) -} diff --git a/examples/blip/main.rs b/examples/blip/main.rs index d5e3e21..a04bfb8 100644 --- a/examples/blip/main.rs +++ b/examples/blip/main.rs @@ -17,7 +17,6 @@ fn main() -> anyhow::Result<()> { .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 @@ -30,7 +29,7 @@ fn main() -> anyhow::Result<()> { let mut model = Blip::new(options_visual, options_textual)?; // image caption - let xs = DataLoader::try_read_batch(&args.source)?; + let xs = DataLoader::try_read_n(&args.source)?; // unconditional caption let ys = model.forward(&xs, None)?; diff --git a/examples/classifier/README.md b/examples/classifier/README.md new file mode 100644 index 0000000..6f2357e --- /dev/null +++ b/examples/classifier/README.md @@ -0,0 +1,6 @@ +## Quick Start + +```shell +cargo run -r -F cuda --example classifier -- --device cuda --dtype fp16 --model beit # convnext, fastvit, deit, mobileone +``` + diff --git a/examples/fastvit/main.rs b/examples/classifier/main.rs similarity index 55% rename from examples/fastvit/main.rs rename to examples/classifier/main.rs index cb93886..fe9536e 100644 --- a/examples/fastvit/main.rs +++ b/examples/classifier/main.rs @@ -11,6 +11,10 @@ struct Args { #[argh(option, default = "String::from(\"cpu:0\")")] device: String, + /// model name + #[argh(option, default = "String::from(\"beit\")")] + model: String, + /// source image #[argh( option, @@ -25,33 +29,46 @@ struct Args { 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()) + .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339()) .init(); let args: Args = argh::from_env(); // build model - let options = Options::fastvit_t8_distill() + let options = match args.model.to_lowercase().as_str() { + "beit" => Options::beit_base(), + "convnext" => Options::convnext_v2_atto(), + "deit" => Options::deit_tiny_distill(), + "fastvit" => Options::fastvit_t8_distill(), + "mobileone" => Options::mobileone_s0(), + _ => anyhow::bail!("Unsupported model: {}", args.model), + }; + + let options = options .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)?; + let xs = DataLoader::try_read_n(&args.source)?; // run let ys = model.forward(&xs)?; - - // results - for (i, y) in ys.iter().enumerate() { - println!("{}: {:?}", i, y); - } + println!("{:?}", ys); // annotate - let annotator = Annotator::default().with_saveout(model.spec()); - annotator.annotate(&xs, &ys); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Image-Classifier", &args.model])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/clip/main.rs b/examples/clip/main.rs index e213c31..a600913 100644 --- a/examples/clip/main.rs +++ b/examples/clip/main.rs @@ -43,13 +43,13 @@ fn main() -> Result<()> { let dl = DataLoader::new("./examples/clip/images")?.build()?; // run - for (images, paths) in dl { + for images in dl { let feats_image = model.encode_images(&images)?; // use image to query texts let matrix = Ops::dot2(&feats_image, &feats_text)?; - for i in 0..paths.len() { + for i in 0..images.len() { let probs = &matrix[i]; let (id, &score) = probs .iter() @@ -58,9 +58,9 @@ fn main() -> Result<()> { .unwrap(); println!( - "({:?}%) {} => {} ", + "({:?}%) {:?} => {} ", score * 100.0, - paths[i].display(), + images[i].source(), &texts[id] ); } diff --git a/examples/convnext/README.md b/examples/convnext/README.md deleted file mode 100644 index fe6d945..0000000 --- a/examples/convnext/README.md +++ /dev/null @@ -1,6 +0,0 @@ -## Quick Start - -```shell -cargo run -r -F cuda --example convnext -- --device cuda --dtype fp16 -``` - diff --git a/examples/convnext/main.rs b/examples/convnext/main.rs deleted file mode 100644 index 6480a07..0000000 --- a/examples/convnext/main.rs +++ /dev/null @@ -1,52 +0,0 @@ -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, -} - -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(()) -} diff --git a/examples/d-fine/main.rs b/examples/d-fine/main.rs index 2726232..b74e3dc 100644 --- a/examples/d-fine/main.rs +++ b/examples/d-fine/main.rs @@ -12,17 +12,23 @@ fn main() -> Result<()> { let mut model = RTDETR::new(options)?; // load - let x = [DataLoader::try_read("./assets/bus.jpg")?]; + let xs = DataLoader::try_read_n(&["./assets/bus.jpg"])?; // run - let y = model.forward(&x)?; - println!("{:?}", y); + let ys = model.forward(&xs)?; + println!("{:?}", ys); // annotate - let annotator = Annotator::default() - .with_bboxes_thickness(3) - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/dataloader/README.md b/examples/dataloader/README.md deleted file mode 100644 index 29d81b9..0000000 --- a/examples/dataloader/README.md +++ /dev/null @@ -1,5 +0,0 @@ -## Quick Start - -```shell -cargo run -r --example dataloader -``` diff --git a/examples/dataloader/main.rs b/examples/dataloader/main.rs deleted file mode 100644 index eacb3b4..0000000 --- a/examples/dataloader/main.rs +++ /dev/null @@ -1,45 +0,0 @@ -use usls::DataLoader; - -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. 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 - )? - .with_batch(1) - .with_progress_bar(true) - .build()?; - - for (_xs, _paths) in dl { - println!("Paths: {:?}", _paths); - } - - // 2. read one image - let image = DataLoader::try_read("./assets/bus.jpg")?; - println!( - "Read one image. Height: {}, Width: {}", - image.height(), - image.width() - ); - - // 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(()) -} diff --git a/examples/db/main.rs b/examples/db/main.rs index 8f797e7..e205c35 100644 --- a/examples/db/main.rs +++ b/examples/db/main.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use usls::{models::DB, Annotator, DataLoader, Options}; +use usls::{models::DB, Annotator, DataLoader, Options, Style}; #[derive(argh::FromArgs)] /// Example @@ -16,21 +16,21 @@ struct Args { #[argh(option, default = "String::from(\"auto\")")] dtype: String, - /// show bboxes + /// show hbbs #[argh(option, default = "false")] - show_bboxes: bool, + show_hbbs: bool, - /// show mbrs + /// show obbs #[argh(option, default = "false")] - show_mbrs: bool, + show_obbs: bool, /// show bboxes confidence #[argh(option, default = "false")] - show_bboxes_conf: bool, + show_hbbs_conf: bool, /// show mbrs confidence #[argh(option, default = "false")] - show_mbrs_conf: bool, + show_obbs_conf: bool, } fn main() -> Result<()> { @@ -52,7 +52,7 @@ fn main() -> Result<()> { )?; // load image - let x = DataLoader::try_read_batch(&[ + let xs = DataLoader::try_read_n(&[ "images/db.png", "images/table.png", "images/table-ch.jpg", @@ -61,20 +61,46 @@ fn main() -> Result<()> { ])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // annotate let annotator = Annotator::default() - .without_bboxes(!args.show_bboxes) - .without_mbrs(!args.show_mbrs) - .without_bboxes_name(true) - .without_mbrs_name(true) - .without_bboxes_conf(!args.show_bboxes_conf) - .without_mbrs_conf(!args.show_mbrs_conf) - .with_polygons_alpha(60) - .with_contours_color([255, 105, 180, 255]) - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + .with_polygon_style( + Style::polygon() + .with_visible(true) + .with_text_visible(false) + .show_confidence(true) + .show_id(true) + .show_name(true) + .with_color(usls::StyleColors::default().with_outline([255, 105, 180, 255].into())), + ) + .with_hbb_style( + Style::hbb() + .with_visible(args.show_hbbs) + .with_text_visible(false) + .with_thickness(1) + .show_confidence(args.show_hbbs_conf) + .show_id(false) + .show_name(false), + ) + .with_obb_style( + Style::obb() + .with_visible(args.show_obbs) + .with_text_visible(false) + .show_confidence(args.show_obbs_conf) + .show_id(false) + .show_name(false), + ); + + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } // summary model.summary(); diff --git a/examples/deim/main.rs b/examples/deim/main.rs index cf8d4e5..b253394 100644 --- a/examples/deim/main.rs +++ b/examples/deim/main.rs @@ -12,17 +12,23 @@ fn main() -> Result<()> { let mut model = RTDETR::new(options)?; // load - let x = [DataLoader::try_read("./assets/bus.jpg")?]; + let xs = DataLoader::try_read_n(&["./assets/bus.jpg"])?; // run - let y = model.forward(&x)?; - println!("{:?}", y); + let ys = model.forward(&xs)?; + println!("{:?}", ys); // annotate - let annotator = Annotator::default() - .with_bboxes_thickness(3) - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/deit/README.md b/examples/deit/README.md deleted file mode 100644 index 962781f..0000000 --- a/examples/deit/README.md +++ /dev/null @@ -1,7 +0,0 @@ -## Quick Start - -```shell -cargo run -r -F cuda --example deit -- --device cuda --dtype fp16 -``` - - diff --git a/examples/deit/main.rs b/examples/deit/main.rs deleted file mode 100644 index 98d7c12..0000000 --- a/examples/deit/main.rs +++ /dev/null @@ -1,52 +0,0 @@ -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, -} - -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(()) -} diff --git a/examples/depth-anything/main.rs b/examples/depth-anything/main.rs index f1deeea..f5b1d3e 100644 --- a/examples/depth-anything/main.rs +++ b/examples/depth-anything/main.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use usls::{models::DepthAnything, Annotator, DataLoader, Options}; +use usls::{models::DepthAnything, Annotator, DataLoader, Options, Style}; fn main() -> Result<()> { tracing_subscriber::fmt() @@ -12,16 +12,23 @@ fn main() -> Result<()> { let mut model = DepthAnything::new(options)?; // load - let x = [DataLoader::try_read("images/street.jpg")?]; + let xs = DataLoader::try_read_n(&["images/street.jpg"])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default() - .with_colormap("Turbo") - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + let annotator = + Annotator::default().with_mask_style(Style::mask().with_colormap256("turbo".into())); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/depth-pro/main.rs b/examples/depth-pro/main.rs index 8919f93..929fa53 100644 --- a/examples/depth-pro/main.rs +++ b/examples/depth-pro/main.rs @@ -1,8 +1,9 @@ use anyhow::Result; -use usls::{models::DepthPro, Annotator, DataLoader, Options}; +use usls::DataLoader; +use usls::{models::DepthPro, Annotator, Options, Style}; #[derive(argh::FromArgs)] -/// BLIP Example +/// Example struct Args { /// device #[argh(option, default = "String::from(\"cpu:0\")")] @@ -11,10 +12,6 @@ struct Args { /// dtype #[argh(option, default = "String::from(\"q4f16\")")] dtype: String, - - /// source image - #[argh(option, default = "String::from(\"images/street.jpg\")")] - source: String, } fn main() -> Result<()> { @@ -33,16 +30,23 @@ fn main() -> Result<()> { let mut model = DepthPro::new(options)?; // load - let x = [DataLoader::try_read(&args.source)?]; + let xs = DataLoader::try_read_n(&["images/street.jpg"])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default() - .with_colormap("Turbo") - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + let annotator = + Annotator::default().with_mask_style(Style::mask().with_colormap256("turbo".into())); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/dinov2/main.rs b/examples/dinov2/main.rs index 5168785..e39427b 100644 --- a/examples/dinov2/main.rs +++ b/examples/dinov2/main.rs @@ -8,10 +8,7 @@ fn main() -> Result<()> { .init(); // images - let xs = [ - DataLoader::try_read("./assets/bus.jpg")?, - DataLoader::try_read("./assets/bus.jpg")?, - ]; + let xs = DataLoader::try_read_n(&["./assets/bus.jpg", "./assets/bus.jpg"])?; // model let options = Options::dinov2_small().with_batch_size(xs.len()).commit()?; diff --git a/examples/doclayout-yolo/main.rs b/examples/doclayout-yolo/main.rs index 99a945b..b6ab0d8 100644 --- a/examples/doclayout-yolo/main.rs +++ b/examples/doclayout-yolo/main.rs @@ -24,17 +24,22 @@ fn main() -> Result<()> { let mut model = YOLO::new(config)?; // load images - let xs = [DataLoader::try_read("images/academic.jpg")?]; + let xs = DataLoader::try_read_n(&["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); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "doclayout-yolo"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } model.summary(); diff --git a/examples/fast/main.rs b/examples/fast/main.rs index 08c8001..9304c9d 100644 --- a/examples/fast/main.rs +++ b/examples/fast/main.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use usls::{models::DB, Annotator, DataLoader, Options, Scale}; +use usls::{models::DB, Annotator, DataLoader, Options, Scale, Style}; #[derive(argh::FromArgs)] /// Example @@ -40,7 +40,7 @@ fn main() -> Result<()> { )?; // load image - let x = DataLoader::try_read_batch(&[ + let xs = DataLoader::try_read_n(&[ "images/db.png", "images/table.png", "images/table-ch.jpg", @@ -49,16 +49,46 @@ fn main() -> Result<()> { ])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // 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); + .with_polygon_style( + Style::polygon() + .with_visible(true) + .with_text_visible(false) + .show_confidence(true) + .show_id(true) + .show_name(true) + .with_color(usls::StyleColors::default().with_outline([255, 105, 180, 255].into())), + ) + .with_hbb_style( + Style::hbb() + .with_visible(false) + .with_text_visible(false) + .with_thickness(1) + .show_confidence(false) + .show_id(false) + .show_name(false), + ) + .with_obb_style( + Style::obb() + .with_visible(false) + .with_text_visible(false) + .show_confidence(false) + .show_id(false) + .show_name(false), + ); + + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/fastsam/main.rs b/examples/fastsam/main.rs index 0050fda..29e2c1e 100644 --- a/examples/fastsam/main.rs +++ b/examples/fastsam/main.rs @@ -29,17 +29,27 @@ fn main() -> Result<()> { let mut model = YOLO::new(config)?; // load images - let xs = DataLoader::try_read_batch(&["./assets/bus.jpg"])?; + let xs = DataLoader::try_read_n(&["./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); + let annotator = Annotator::default().with_hbb_style( + usls::Style::hbb() + .show_confidence(true) + .show_id(false) + .show_name(false), + ); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "FastSAM"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/fastvit/README.md b/examples/fastvit/README.md deleted file mode 100644 index ca00fdf..0000000 --- a/examples/fastvit/README.md +++ /dev/null @@ -1,13 +0,0 @@ -## 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"))] } } - -``` diff --git a/examples/florence2/main.rs b/examples/florence2/main.rs index 52f7673..4d6bde0 100644 --- a/examples/florence2/main.rs +++ b/examples/florence2/main.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use usls::{models::Florence2, Annotator, DataLoader, Options, Scale, Task}; +use usls::{models::Florence2, Annotator, DataLoader, Options, Scale, Style, Task}; #[derive(argh::FromArgs)] /// Example @@ -26,10 +26,7 @@ fn main() -> Result<()> { let args: Args = argh::from_env(); // load images - let xs = [ - DataLoader::try_read("images/green-car.jpg")?, - DataLoader::try_read("assets/bus.jpg")?, - ]; + let xs = DataLoader::try_read_n(&["images/green-car.jpg", "assets/bus.jpg"])?; // build model let ( @@ -109,12 +106,6 @@ fn main() -> Result<()> { ), ]; - // annotator - let annotator = Annotator::new() - .without_bboxes_conf(true) - .with_bboxes_thickness(3) - .with_saveout_subs(&["Florence2"]); - // inference for task in tasks.iter() { let ys = model.forward(&xs, task)?; @@ -128,44 +119,122 @@ fn main() -> Result<()> { println!("Task: {:?}\n{:?}\n", task, &ys) } Task::DenseRegionCaption => { - let annotator = annotator.clone().with_saveout("Dense-Region-Caption"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default() + .with_hbb_style(Style::hbb().show_confidence(false)) + .annotate(x, y)? + .save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Florence2", "Dense-Region-Caption"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::RegionProposal => { - let annotator = annotator - .clone() - .without_bboxes_name(false) - .with_saveout("Region-Proposal"); - - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default() + .with_hbb_style(Style::hbb().show_confidence(false).show_name(false)) + .annotate(x, y)? + .save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Florence2", "Region-Proposal"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::ObjectDetection => { - let annotator = annotator.clone().with_saveout("Object-Detection"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default().annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Florence2", "Object-Detection"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::OpenSetDetection(_) => { - let annotator = annotator.clone().with_saveout("Open-Set-Detection"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default().annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Florence2", "Open-Object-Detection"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::CaptionToPhraseGrounding(_) => { - let annotator = annotator - .clone() - .with_saveout("Caption-To-Phrase-Grounding"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default() + .with_hbb_style(Style::hbb().show_confidence(false)) + .annotate(x, y)? + .save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&[ + "runs", + "Florence2", + "Caption-To-Phrase-Grounding" + ])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::ReferringExpressionSegmentation(_) => { - let annotator = annotator - .clone() - .with_saveout("Referring-Expression-Segmentation"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default() + .with_hbb_style(Style::hbb().show_confidence(false)) + .annotate(x, y)? + .save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&[ + "runs", + "Florence2", + "Referring-Expression-Segmentation" + ])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::RegionToSegmentation(..) => { - let annotator = annotator.clone().with_saveout("Region-To-Segmentation"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default() + .with_hbb_style(Style::hbb().show_confidence(false)) + .annotate(x, y)? + .save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&[ + "runs", + "Florence2", + "Region-To-Segmentation", + ])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } Task::OcrWithRegion => { - let annotator = annotator.clone().with_saveout("Ocr-With-Region"); - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + Annotator::default() + .with_hbb_style(Style::hbb().show_confidence(false)) + .annotate(x, y)? + .save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "Florence2", "Ocr-With-Region",])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } _ => (), diff --git a/examples/grounding-dino/main.rs b/examples/grounding-dino/main.rs index 389ad89..e80083f 100644 --- a/examples/grounding-dino/main.rs +++ b/examples/grounding-dino/main.rs @@ -56,16 +56,22 @@ fn main() -> Result<()> { let mut model = GroundingDINO::new(options)?; // load images - let xs = DataLoader::try_read_batch(&args.source)?; + let xs = DataLoader::try_read_n(&args.source)?; // run let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default() - .with_bboxes_thickness(4) - .with_saveout(model.spec()); - annotator.annotate(&xs, &ys); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } // summary model.summary(); diff --git a/examples/hub/README.md b/examples/hub/README.md deleted file mode 100644 index 7cddfbc..0000000 --- a/examples/hub/README.md +++ /dev/null @@ -1,5 +0,0 @@ -## Quick Start - -```shell -RUST_LOG=usls=info cargo run -r --example hub -``` diff --git a/examples/hub/main.rs b/examples/hub/main.rs deleted file mode 100644 index 33f62e3..0000000 --- a/examples/hub/main.rs +++ /dev/null @@ -1,26 +0,0 @@ -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 (i, tag) in hub.tags().iter().enumerate() { - let files = hub.files(tag); - println!("{} :: {} => {:?}", i, tag, files); // Should be empty - } - - Ok(()) -} diff --git a/examples/imshow.rs b/examples/imshow.rs new file mode 100644 index 0000000..2c7f647 --- /dev/null +++ b/examples/imshow.rs @@ -0,0 +1,51 @@ +use usls::{DataLoader, Viewer}; + +#[derive(argh::FromArgs)] +/// Example +struct Args { + /// source + #[argh(option, default = "String::from(\"./assets\")")] + source: String, + // /// record video and save + // #[argh(option, default = "false")] + // save_video: bool, +} + +fn main() -> anyhow::Result<()> { + let args: Args = argh::from_env(); + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339()) + .init(); + + let dl = DataLoader::new(args.source.as_str())?.build()?; + let mut viewer = Viewer::default().with_window_scale(1.); + + for images in &dl { + // check out window + if viewer.is_window_exist() && !viewer.is_window_open() { + break; + } + + viewer.imshow(&images[0])?; + + // check out key event + if let Some(key) = viewer.wait_key(1) { + if key == usls::Key::Escape { + break; + } + } + + // image info + for image in &images { + println!("## {:?}", image); + } + + // // write video, need video feature + // if args.save_video { + // viewer.write_video_frame(&images[0])?; + // } + } + + Ok(()) +} diff --git a/examples/linknet/README.md b/examples/linknet/README.md index 89227df..053bb5a 100644 --- a/examples/linknet/README.md +++ b/examples/linknet/README.md @@ -1,6 +1,6 @@ ## Quick Start ```shell -cargo run -r --example fast +cargo run -r --example linknet ``` diff --git a/examples/linknet/main.rs b/examples/linknet/main.rs index 4fc3841..3dfd21a 100644 --- a/examples/linknet/main.rs +++ b/examples/linknet/main.rs @@ -1,5 +1,6 @@ use anyhow::Result; -use usls::{models::DB, Annotator, DataLoader, Options, Scale}; +use usls::DataLoader; +use usls::{models::DB, Annotator, Options, Scale, Style}; #[derive(argh::FromArgs)] /// Example @@ -40,7 +41,7 @@ fn main() -> Result<()> { )?; // load image - let x = DataLoader::try_read_batch(&[ + let xs = DataLoader::try_read_n(&[ "images/table.png", "images/table1.jpg", "images/table2.png", @@ -50,16 +51,46 @@ fn main() -> Result<()> { ])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // 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); + .with_polygon_style( + Style::polygon() + .with_visible(true) + .with_text_visible(false) + .show_confidence(true) + .show_id(true) + .show_name(true) + .with_color(usls::StyleColors::default().with_outline([255, 105, 180, 255].into())), + ) + .with_hbb_style( + Style::hbb() + .with_visible(false) + .with_text_visible(false) + .with_thickness(1) + .show_confidence(false) + .show_id(false) + .show_name(false), + ) + .with_obb_style( + Style::obb() + .with_visible(false) + .with_text_visible(false) + .show_confidence(false) + .show_id(false) + .show_name(false), + ); + + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/mobileone/README.md b/examples/mobileone/README.md deleted file mode 100644 index ca00fdf..0000000 --- a/examples/mobileone/README.md +++ /dev/null @@ -1,13 +0,0 @@ -## 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"))] } } - -``` diff --git a/examples/mobileone/main.rs b/examples/mobileone/main.rs deleted file mode 100644 index 36238c2..0000000 --- a/examples/mobileone/main.rs +++ /dev/null @@ -1,57 +0,0 @@ -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, -} - -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(()) -} diff --git a/examples/modnet/main.rs b/examples/modnet/main.rs index 39691b3..b0ee231 100644 --- a/examples/modnet/main.rs +++ b/examples/modnet/main.rs @@ -11,14 +11,22 @@ fn main() -> anyhow::Result<()> { let mut model = MODNet::new(options)?; // load image - let xs = [DataLoader::try_read("images/liuyifei.png")?]; + let xs = DataLoader::try_read_n(&["images/liuyifei.png"])?; // run let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default().with_saveout(model.spec()); - annotator.annotate(&xs, &ys); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/moondream2/main.rs b/examples/moondream2/main.rs index 299f590..963197a 100644 --- a/examples/moondream2/main.rs +++ b/examples/moondream2/main.rs @@ -116,7 +116,7 @@ fn main() -> Result<()> { )?; // load images - let xs = DataLoader::try_read_batch(&args.source)?; + let xs = DataLoader::try_read_n(&args.source)?; // run with task let task: Task = args.task.as_str().try_into()?; @@ -142,13 +142,37 @@ fn main() -> Result<()> { } Task::OpenSetDetection(_) | Task::OpenSetKeypointsDetection(_) => { println!("{:?}", ys); + // let annotator = Annotator::default() + // .with_bboxes_thickness(4) + // .without_bboxes_conf(true) + // .with_keypoints_radius(6) + // .with_keypoints_name(true) + // .with_saveout("moondream2"); + // annotator.annotate(&xs, &ys); + + // annotate let annotator = Annotator::default() - .with_bboxes_thickness(4) - .without_bboxes_conf(true) - .with_keypoints_radius(6) - .with_keypoints_name(true) - .with_saveout("moondream2"); - annotator.annotate(&xs, &ys); + .with_hbb_style( + usls::Style::hbb() + .with_draw_fill(true) + .show_confidence(false), + ) + .with_keypoint_style( + usls::Style::keypoint() + .show_confidence(false) + .show_id(true) + .show_name(false), + ); + + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "moondream2"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } _ => unimplemented!("Unsupported moondream2 task."), } diff --git a/examples/owlv2/main.rs b/examples/owlv2/main.rs index 577ae8a..f717e00 100644 --- a/examples/owlv2/main.rs +++ b/examples/owlv2/main.rs @@ -1,5 +1,6 @@ use anyhow::Result; -use usls::{models::OWLv2, Annotator, DataLoader, Options}; +use usls::DataLoader; +use usls::{models::OWLv2, Annotator, Options}; #[derive(argh::FromArgs)] /// Example @@ -20,19 +21,19 @@ struct Args { #[argh( option, default = "vec![ - String::from(\"person\"), - String::from(\"hand\"), - String::from(\"shoes\"), - String::from(\"bus\"), - String::from(\"car\"), - String::from(\"dog\"), - String::from(\"cat\"), - String::from(\"sign\"), - String::from(\"tie\"), - String::from(\"monitor\"), - String::from(\"glasses\"), - String::from(\"tree\"), - String::from(\"head\"), + String::from(\"person\"), + String::from(\"hand\"), + String::from(\"shoes\"), + String::from(\"bus\"), + String::from(\"car\"), + 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, @@ -55,16 +56,22 @@ fn main() -> Result<()> { let mut model = OWLv2::new(options)?; // load - let xs = DataLoader::try_read_batch(&args.source)?; + let xs = DataLoader::try_read_n(&args.source)?; // run let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default() - .with_bboxes_thickness(3) - .with_saveout(model.spec()); - annotator.annotate(&xs, &ys); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/picodet-layout/main.rs b/examples/picodet-layout/main.rs index fca0bcb..baa9bf4 100644 --- a/examples/picodet-layout/main.rs +++ b/examples/picodet-layout/main.rs @@ -1,5 +1,6 @@ use anyhow::Result; -use usls::{models::PicoDet, Annotator, DataLoader, Options}; +use usls::DataLoader; +use usls::{models::PicoDet, Annotator, Options}; fn main() -> Result<()> { tracing_subscriber::fmt() @@ -15,17 +16,23 @@ fn main() -> Result<()> { 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()); + let xs = DataLoader::try_read_n(&["images/academic.jpg"])?; // run let ys = model.forward(&xs)?; println!("{:?}", ys); - annotator.annotate(&xs, &ys); + + // annotate + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/read_images.rs b/examples/read_images.rs new file mode 100644 index 0000000..d397419 --- /dev/null +++ b/examples/read_images.rs @@ -0,0 +1,58 @@ +use usls::{DataLoader, Image, ImageVecExt}; + +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. Read one image + let image = Image::try_read("./assets/bus.jpg")?; + println!("Image::try_read(): {:?}", image); + // image.save("kkk.png")?; + + // => To Rgba8 + let _image_rgba = image.to_rgba8(); + + // 2. Read one image with DataLoader + let image = DataLoader::try_read_one("./assets/bus.jpg")?; + println!("DataLoader::try_read_one(): {:?}", image); + + // 3. Read N images with DataLoader + let images = DataLoader::try_read_n(&["./assets/bus.jpg", "./assets/cat.png"])?; + println!("DataLoader::try_read_n():"); + for image in images { + println!(" - {:?}", image); + } + + // 4. Read image folder with DataLoader + let images = DataLoader::try_read_folder("./assets")?; + println!("DataLoader::try_read_folder():"); + for image in images { + println!(" - {:?}", image); + } + + // 5. Glob and read image folder with DataLoader + // let images = DataLoader::try_read_pattern("./assets/*.Jpg")?; + let images = DataLoader::try_read_pattern_case_insensitive("./assets/*.Jpg")?; + println!("DataLoader::try_read_pattern_case_insensitive():"); + + for image in images { + println!(" - {:?}", image); + } + + // 6. Load images with DataLoader + let dl = DataLoader::new("./assets")?.with_batch(2).build()?; + + // iterate over the dataloader + for (i, images) in dl.iter().enumerate() { + println!("## Batch-{}: {:?}", i + 1, images); + } + + // 7. Vec <-> Vec + let images = DataLoader::try_read_n(&["./assets/bus.jpg", "./assets/cat.png"])?; + let dyn_images = images.into_dyns(); + let _images = dyn_images.into_images(); + + Ok(()) +} diff --git a/examples/viewer/main.rs b/examples/read_video.rs similarity index 50% rename from examples/viewer/main.rs rename to examples/read_video.rs index 8279204..bc6f998 100644 --- a/examples/viewer/main.rs +++ b/examples/read_video.rs @@ -1,4 +1,4 @@ -use usls::{DataLoader, Key, Viewer}; +use usls::DataLoader; #[derive(argh::FromArgs)] /// Example @@ -12,32 +12,25 @@ struct Args { } fn main() -> anyhow::Result<()> { + let args: Args = argh::from_env(); 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()?; + // load images or video stream + let dl = DataLoader::new(args.source.as_str())? + .with_batch(1) + // .with_nf_skip(1) + // .with_progress_bar(true) + .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; + // iterate over the dataloader + for images in &dl { + for image in &images { + println!("## {:?}", image); } - - // write video - viewer.write_batch(&xs)? } - // finish video write - viewer.finish_write()?; - Ok(()) } diff --git a/examples/rfdetr/README.md b/examples/rfdetr/README.md index a196caf..c3be155 100644 --- a/examples/rfdetr/README.md +++ b/examples/rfdetr/README.md @@ -7,11 +7,11 @@ cargo run -r --example rfdetr ## Results ``` -[Bboxes]: Found 6 objects -0: Bbox { xyxy: [221.55753, 408.0652, 345.23325, 860.2527], class_id: 1, name: Some("person"), confidence: 0.93212366 } -1: Bbox { xyxy: [44.967827, 397.84177, 246.13187, 905.7567], class_id: 1, name: Some("person"), confidence: 0.93540853 } -2: Bbox { xyxy: [6.2678833, 233.208, 801.6806, 737.4714], class_id: 6, name: Some("bus"), confidence: 0.93637216 } -3: Bbox { xyxy: [0.0, 555.167, 77.74801, 870.2772], class_id: 1, name: Some("person"), confidence: 0.85163206 } -4: Bbox { xyxy: [133.94543, 473.6574, 149.62558, 507.99875], class_id: 32, name: Some("tie"), confidence: 0.2992424 } -5: Bbox { xyxy: [669.81836, 395.28635, 813.44855, 879.9562], class_id: 1, name: Some("person"), confidence: 0.83661026 } +[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 } + ``` diff --git a/examples/rfdetr/main.rs b/examples/rfdetr/main.rs index c86ec96..e034e45 100644 --- a/examples/rfdetr/main.rs +++ b/examples/rfdetr/main.rs @@ -12,26 +12,25 @@ fn main() -> Result<()> { let mut model = RFDETR::new(options)?; // load - let xs = [DataLoader::try_read("./assets/bus.jpg")?]; + let xs = DataLoader::try_read_n(&["./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) - } - } - } + println!("{:?}", ys); // annotate - let annotator = Annotator::default() - .with_bboxes_thickness(3) - .with_saveout(model.spec()); - annotator.annotate(&xs, &ys); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/rtdetr/README.md b/examples/rtdetr/README.md index 35f9c32..711c097 100644 --- a/examples/rtdetr/README.md +++ b/examples/rtdetr/README.md @@ -13,4 +13,5 @@ cargo run -r --example rtdetr 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 } + ``` diff --git a/examples/rtdetr/main.rs b/examples/rtdetr/main.rs index 590b218..1d3aa5e 100644 --- a/examples/rtdetr/main.rs +++ b/examples/rtdetr/main.rs @@ -18,26 +18,23 @@ fn main() -> Result<()> { let mut model = RTDETR::new(options)?; // load - let xs = [DataLoader::try_read("./assets/bus.jpg")?]; + let xs = DataLoader::try_read_n(&["./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) - } - } - } + println!("{:?}", ys); // annotate - let annotator = Annotator::default() - .with_bboxes_thickness(3) - .with_saveout(model.spec()); - annotator.annotate(&xs, &ys); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/rtmo/main.rs b/examples/rtmo/main.rs index efe198a..314e6fc 100644 --- a/examples/rtmo/main.rs +++ b/examples/rtmo/main.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use usls::{models::RTMO, Annotator, DataLoader, Options, COCO_SKELETONS_16}; +use usls::{models::RTMO, Annotator, DataLoader, Options, Style, SKELETON_COCO_19}; fn main() -> Result<()> { tracing_subscriber::fmt() @@ -11,16 +11,31 @@ fn main() -> Result<()> { let mut model = RTMO::new(Options::rtmo_s().commit()?)?; // load image - let xs = [DataLoader::try_read("images/bus.jpg")?]; + let xs = DataLoader::try_read_n(&["./assets/bus.jpg"])?; // run let ys = model.forward(&xs)?; + println!("ys: {:?}", ys); // annotate let annotator = Annotator::default() - .with_saveout(model.spec()) - .with_skeletons(&COCO_SKELETONS_16); - annotator.annotate(&xs, &ys); + .with_hbb_style(Style::hbb().with_draw_fill(true)) + .with_keypoint_style( + Style::keypoint() + .with_skeleton(SKELETON_COCO_19.into()) + .show_confidence(false) + .show_id(true) + .show_name(false), + ); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/sam/main.rs b/examples/sam/main.rs index ca009c7..29607b7 100644 --- a/examples/sam/main.rs +++ b/examples/sam/main.rs @@ -64,10 +64,7 @@ fn main() -> Result<()> { let mut model = SAM::new(options_encoder, options_decoder)?; // Load image - let xs = [DataLoader::try_read("images/truck.jpg")?]; - - // Build annotator - let annotator = Annotator::default().with_saveout(model.spec()); + let xs = DataLoader::try_read_n(&["images/truck.jpg"])?; // Prompt let prompts = vec![ @@ -79,7 +76,18 @@ fn main() -> Result<()> { // Run & Annotate let ys = model.forward(&xs, &prompts)?; - annotator.annotate(&xs, &ys); + + // annotate + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/sapiens/main.rs b/examples/sapiens/main.rs index 08d3167..a209e27 100644 --- a/examples/sapiens/main.rs +++ b/examples/sapiens/main.rs @@ -23,17 +23,22 @@ fn main() -> Result<()> { let mut model = Sapiens::new(options)?; // load - let x = [DataLoader::try_read("images/paul-george.jpg")?]; + let xs = DataLoader::try_read_n(&["images/paul-george.jpg"])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default() - .without_masks(true) - .with_polygons_name(true) - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/slanet/main.rs b/examples/slanet/main.rs index fe4cef1..362671c 100644 --- a/examples/slanet/main.rs +++ b/examples/slanet/main.rs @@ -1,5 +1,5 @@ use anyhow::Result; -use usls::{models::SLANet, Annotator, DataLoader, Options}; +use usls::{models::SLANet, Annotator, Color, DataLoader, Options}; #[derive(argh::FromArgs)] /// Example @@ -33,18 +33,34 @@ fn main() -> Result<()> { let mut model = SLANet::new(options)?; // load - let xs = DataLoader::try_read_batch(&[args.source])?; + let xs = DataLoader::try_read_n(&[args.source])?; // run let ys = model.forward(&xs)?; - // println!("{:?}", ys); + 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); + let annotator = Annotator::default().with_keypoint_style( + usls::Style::keypoint() + .with_text_visible(false) + .with_skeleton( + ( + [(0, 1), (1, 2), (2, 3), (3, 0)], + [Color::black(), Color::red(), Color::green(), Color::blue()], + ) + .into(), + ), + ); + + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } // summary model.summary(); diff --git a/examples/smolvlm/main.rs b/examples/smolvlm/main.rs index a6bf1e5..3bc87fb 100644 --- a/examples/smolvlm/main.rs +++ b/examples/smolvlm/main.rs @@ -57,7 +57,7 @@ fn main() -> Result<()> { )?; // load images - let xs = DataLoader::try_read_batch(&args.source)?; + let xs = DataLoader::try_read_n(&args.source)?; // run let ys = model.forward(&xs, &args.prompt)?; diff --git a/examples/svtr/main.rs b/examples/svtr/main.rs index d53a2ef..bd92036 100644 --- a/examples/svtr/main.rs +++ b/examples/svtr/main.rs @@ -37,9 +37,9 @@ fn main() -> Result<()> { .build()?; // run - for (xs, paths) in dl { + for xs in &dl { let ys = model.forward(&xs)?; - println!("{paths:?}: {:?}", ys) + println!("ys: {:?}", ys); } // summary diff --git a/examples/trocr/main.rs b/examples/trocr/main.rs index 3b7d8ea..f64b72e 100644 --- a/examples/trocr/main.rs +++ b/examples/trocr/main.rs @@ -32,7 +32,7 @@ fn main() -> anyhow::Result<()> { let args: Args = argh::from_env(); // load images - let xs = DataLoader::try_read_batch(&[ + let xs = DataLoader::try_read_n(&[ "images/text-en-dark.png", "images/text-hello-rust-handwritten.png", ])?; diff --git a/examples/viewer/README.md b/examples/viewer/README.md deleted file mode 100644 index 0cfe0e0..0000000 --- a/examples/viewer/README.md +++ /dev/null @@ -1,5 +0,0 @@ -## Quick Start - -```shell -RUST_LOG=usls=info cargo run -F ffmpeg -r --example viewer -``` diff --git a/examples/yolo-sam/main.rs b/examples/yolo-sam/main.rs index b66fb63..74c35d6 100644 --- a/examples/yolo-sam/main.rs +++ b/examples/yolo-sam/main.rs @@ -1,7 +1,7 @@ use anyhow::Result; use usls::{ models::{SamPrompt, SAM, YOLO}, - Annotator, DataLoader, Options, Scale, + Annotator, DataLoader, Options, Scale, Style, }; #[derive(argh::FromArgs)] @@ -30,37 +30,41 @@ fn main() -> Result<()> { // build YOLOv8 let options_yolo = Options::yolo_detect() .with_model_scale(Scale::N) - .with_model_version(8.0.into()) + .with_model_version(8.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_batch(&["images/dog.jpg"])?; + let xs = DataLoader::try_read_n(&["images/dog.jpg"])?; // build annotator - let annotator = Annotator::default() - .with_bboxes_thickness(7) - .without_bboxes_name(true) - .without_bboxes_conf(true) - .without_mbrs(true) - .with_saveout("YOLO-SAM"); + let annotator = Annotator::default().with_hbb_style(Style::hbb().with_draw_fill(true)); // run & annotate let ys_det = yolo.forward(&xs)?; for y_det in ys_det.iter() { - if let Some(bboxes) = y_det.bboxes() { - for bbox in bboxes { + if let Some(hbbs) = y_det.hbbs() { + for hbb in hbbs { let ys_sam = sam.forward( &xs, &[SamPrompt::default().with_bbox( - bbox.xmin(), - bbox.ymin(), - bbox.xmax(), - bbox.ymax(), + hbb.xmin(), + hbb.ymin(), + hbb.xmax(), + hbb.ymax(), )], )?; - annotator.annotate(&xs, &ys_sam); + // annotator.annotate(&xs, &ys_sam); + for (x, y) in xs.iter().zip(ys_sam.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "YOLO-SAM"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } } } diff --git a/examples/yolo/main.rs b/examples/yolo/main.rs index 71ec5fb..86111bd 100644 --- a/examples/yolo/main.rs +++ b/examples/yolo/main.rs @@ -1,7 +1,7 @@ use anyhow::Result; use usls::{ - models::YOLO, Annotator, DataLoader, Options, COCO_CLASS_NAMES_80, COCO_SKELETONS_16, - IMAGENET_NAMES_1K, + models::YOLO, Annotator, DataLoader, Options, Style, NAMES_COCO_80, NAMES_COCO_KEYPOINTS_17, + NAMES_IMAGENET_1K, SKELETON_COCO_19, SKELETON_COLOR_COCO_19, }; #[derive(argh::FromArgs, Debug)] @@ -39,10 +39,6 @@ struct Args { #[argh(option, default = "true")] trt_fp16: bool, - /// find_contours - #[argh(option, default = "true")] - find_contours: bool, - /// batch_size #[argh(option, default = "1")] batch_size: usize, @@ -91,6 +87,10 @@ struct Args { #[argh(switch)] use_coco_80_classes: bool, + /// use_coco_17_keypoints_classes + #[argh(switch)] + use_coco_17_keypoints_classes: bool, + /// use_imagenet_1k_classes #[argh(switch)] use_imagenet_1k_classes: bool, @@ -118,6 +118,10 @@ struct Args { /// keypoint_names #[argh(option)] keypoint_names: Vec, + + /// topk + #[argh(option, default = "5")] + topk: usize, } fn main() -> Result<()> { @@ -131,7 +135,7 @@ fn main() -> Result<()> { 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_version(args.ver.try_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()?) @@ -166,16 +170,20 @@ fn main() -> Result<()> { } else { &args.keypoint_confs }) - .with_find_contours(args.find_contours) + .with_topk(args.topk) .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); + options = options.with_class_names(&NAMES_COCO_80); + } + + if args.use_coco_17_keypoints_classes { + options = options.with_keypoint_names(&NAMES_COCO_KEYPOINTS_17); } if args.use_imagenet_1k_classes { - options = options.with_class_names(&IMAGENET_NAMES_1K); + options = options.with_class_names(&NAMES_IMAGENET_1K); } if let Some(nc) = args.num_classes { @@ -216,26 +224,35 @@ fn main() -> Result<()> { // build annotator let annotator = Annotator::default() - .with_skeletons(&COCO_SKELETONS_16) - .without_masks(true) - .with_bboxes_thickness(3) - .with_saveout(model.spec()); + .with_obb_style(Style::obb().with_draw_fill(true)) + .with_hbb_style( + Style::hbb() + .with_draw_fill(true) + .with_palette(&usls::Color::palette_coco_80()), + ) + .with_keypoint_style( + Style::keypoint() + .with_skeleton((SKELETON_COCO_19, SKELETON_COLOR_COCO_19).into()) + .show_confidence(false) + .show_id(true) + .show_name(false), + ) + .with_mask_style(Style::mask().with_draw_mask_polygon_largest(true)); // run & annotate - for (xs, _paths) in dl { + for xs in &dl { 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) - // } - // } - // } + println!("ys: {:?}", ys); - // plot - annotator.annotate(&xs, &ys); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } } model.summary(); diff --git a/examples/yolop/main.rs b/examples/yolop/main.rs index ed8283d..7f062db 100644 --- a/examples/yolop/main.rs +++ b/examples/yolop/main.rs @@ -12,16 +12,22 @@ fn main() -> Result<()> { let mut model = YOLOPv2::new(options)?; // load image - let x = [DataLoader::try_read("images/car-view.jpg")?]; + let xs = DataLoader::try_read_n(&["images/car-view.jpg"])?; // run - let y = model.forward(&x)?; + let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default() - .with_polygons_name(true) - .with_saveout(model.spec()); - annotator.annotate(&x, &y); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/examples/yolov8-rtdetr/main.rs b/examples/yolov8-rtdetr/main.rs index 87f611f..3b96e8e 100644 --- a/examples/yolov8-rtdetr/main.rs +++ b/examples/yolov8-rtdetr/main.rs @@ -29,17 +29,23 @@ fn main() -> Result<()> { let mut model = YOLO::new(config)?; // load images - let xs = DataLoader::try_read_batch(&["./assets/bus.jpg"])?; + let xs = DataLoader::try_read_n(&["./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); + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", "YOLOv8-RT-DETR"])? + .join(usls::timestamp(None)) + .display(), + ))?; + } Ok(()) } diff --git a/src/misc/engine.rs b/src/inference/engine.rs similarity index 90% rename from src/misc/engine.rs rename to src/inference/engine.rs index 1c6543b..ad6cdf0 100644 --- a/src/misc/engine.rs +++ b/src/inference/engine.rs @@ -3,13 +3,9 @@ use anyhow::Result; use half::{bf16, f16}; use log::{debug, info, warn}; use ndarray::{Array, IxDyn}; -#[allow(unused_imports)] use ort::{ execution_providers::ExecutionProvider, - session::{ - builder::{GraphOptimizationLevel, SessionBuilder}, - Session, SessionInputValue, - }, + session::{builder::GraphOptimizationLevel, Session, SessionInputValue}, tensor::TensorElementType, value::{DynValue, Value}, }; @@ -17,9 +13,31 @@ use prost::Message; use std::collections::HashSet; use crate::{ - build_progress_bar, elapsed, human_bytes, onnx, DType, Device, Iiix, MinOptMax, Ops, Ts, Xs, X, + build_progress_bar, elapsed, human_bytes_binary, onnx, DType, Device, Iiix, MinOptMax, Ops, Ts, + Xs, PROGRESS_BAR_STYLE_CYAN_2, PROGRESS_BAR_STYLE_FINISH, X, }; +impl From for DType { + fn from(dtype: TensorElementType) -> Self { + match dtype { + TensorElementType::Int8 => Self::Int8, + TensorElementType::Int16 => Self::Int16, + TensorElementType::Int32 => Self::Int32, + TensorElementType::Int64 => Self::Int64, + TensorElementType::Uint8 => Self::Uint8, + TensorElementType::Uint16 => Self::Uint16, + TensorElementType::Uint32 => Self::Uint32, + TensorElementType::Uint64 => Self::Uint64, + TensorElementType::Float16 => Self::Fp16, + TensorElementType::Float32 => Self::Fp32, + TensorElementType::Float64 => Self::Fp64, + TensorElementType::Bfloat16 => Self::Bf16, + TensorElementType::String => Self::String, + TensorElementType::Bool => Self::Bool, + } + } +} + /// A struct for tensor attrs composed of the names, the dtypes, and the dimensions. #[derive(Builder, Debug, Clone)] pub struct OrtTensorAttr { @@ -42,16 +60,17 @@ pub struct Engine { pub spec: String, pub device: Device, pub trt_fp16: bool, - #[args(inc = true)] + #[args(inc)] pub iiixs: Vec, - #[args(alias = "parameters")] + #[args(aka = "parameters")] pub params: Option, - #[args(alias = "memory")] + #[args(aka = "memory")] pub wbmems: Option, pub inputs_minoptmax: Vec>, pub onnx: Option, pub ts: Ts, pub num_dry_run: usize, + pub graph_opt_level: Option, } impl Default for Engine { @@ -68,6 +87,7 @@ impl Default for Engine { inputs_minoptmax: vec![], onnx: None, ts: Ts::default(), + graph_opt_level: None, } } } @@ -151,7 +171,7 @@ impl Engine { self.num_dry_run as u64, "DryRun", Some(self.spec()), - crate::PROGRESS_BAR_STYLE_CYAN_2, + PROGRESS_BAR_STYLE_CYAN_2, )?; // dummy @@ -181,14 +201,14 @@ impl Engine { self.spec, match self.params { Some(bytes) if bytes != 0 => { - human_bytes(bytes as f64, true) + human_bytes_binary(bytes as f64, 2) } _ => "Unknown".to_string(), }, self.device, )); pb.set_style(indicatif::ProgressStyle::with_template( - crate::PROGRESS_BAR_STYLE_FINISH, + PROGRESS_BAR_STYLE_FINISH, )?); pb.finish(); } @@ -349,7 +369,7 @@ impl Engine { spec_max += &s_max; } - let p = crate::Dir::Cache.path_with_subs(&["trt-cache"])?; + let p = crate::Dir::Cache.crate_dir_default_with_subs(&["trt-cache"])?; let ep = ort::execution_providers::TensorRTExecutionProvider::default() .with_device_id(id as i32) .with_fp16(self.trt_fp16) @@ -430,8 +450,14 @@ impl Engine { } // session + let graph_opt_level = match self.graph_opt_level { + Some(0) => GraphOptimizationLevel::Disable, + Some(1) => GraphOptimizationLevel::Level1, + Some(2) => GraphOptimizationLevel::Level2, + _ => GraphOptimizationLevel::Level3, + }; let session = builder - .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_optimization_level(graph_opt_level)? .with_intra_threads(std::thread::available_parallelism()?.get())? .commit_from_file(self.file())?; @@ -602,6 +628,24 @@ impl Engine { }) } + // pub fn to_ort(&self) -> TensorElementType { + // match self { + // Self::Int8 => TensorElementType::Int8, + // Self::Int16 => TensorElementType::Int16, + // Self::Int32 => TensorElementType::Int32, + // Self::Int64 => TensorElementType::Int64, + // Self::Uint8 => TensorElementType::Uint8, + // Self::Uint16 => TensorElementType::Uint16, + // Self::Uint32 => TensorElementType::Uint32, + // Self::Uint64 => TensorElementType::Uint64, + // Self::Fp16 => TensorElementType::Float16, + // Self::Fp32 => TensorElementType::Float32, + // Self::Fp64 => TensorElementType::Float64, + // Self::Bf16 => TensorElementType::Bfloat16, + // _ => todo!(), + // } + // } + pub fn load_onnx>(p: P) -> Result { let f = std::fs::read(p.as_ref())?; onnx::ModelProto::decode(f.as_slice()).map_err(|err| { @@ -692,7 +736,7 @@ impl Engine { x.inputs .dtypes() .iter() - .map(DType::from_ort) + .map(|x| DType::from(*x)) .collect::>() .into() }) @@ -715,7 +759,7 @@ impl Engine { x.outputs .dtypes() .iter() - .map(DType::from_ort) + .map(|x| DType::from(*x)) .collect::>() .into() }) @@ -733,13 +777,13 @@ impl Engine { self.device, match self.params { Some(bytes) if bytes != 0 => { - human_bytes(bytes as f64, true) + human_bytes_binary(bytes as f64, 2) } _ => "Unknown".to_string(), }, match self.wbmems { Some(bytes) if bytes != 0 => { - human_bytes(bytes as f64, true) + human_bytes_binary(bytes as f64, 2) } _ => "Unknown".to_string(), }, diff --git a/src/inference/hbb.rs b/src/inference/hbb.rs new file mode 100644 index 0000000..63402c0 --- /dev/null +++ b/src/inference/hbb.rs @@ -0,0 +1,290 @@ +use aksr::Builder; + +use crate::{InstanceMeta, Keypoint, Style}; + +#[derive(Builder, Clone, Default)] +pub struct Hbb { + x: f32, + y: f32, + w: f32, + h: f32, + meta: InstanceMeta, + style: Option